manage progress and output from solver

offline_optimization_script.py

@@ -28,13 +28,26 @@
 import itertools
 import json
 import logging
+import os
 from pathlib import Path
 from typing import Literal
 
-import numpy as np
+# don't use tqdm, because we want to suppress
+# tqdm bars of optimization process (from linopy)
+import progress.bar
 
-from ptxboa import DEFAULT_CACHE_DIR, DEFAULT_DATA_DIR
-from ptxboa.api import PtxboaAPI
+# suppress most of the solver output from
+# linopy and HIGHs
+# MUST be set before importing ptxboa
+os.environ["HIGHS_OUTPUT_FLAG"] = "false"
+os.environ["TQDM_DISABLE"] = "1"
+
+
+from ptxboa import (  # noqa E402 module level import not at top
+    DEFAULT_CACHE_DIR,
+    DEFAULT_DATA_DIR,
+)
+from ptxboa.api import PtxboaAPI  # noqa E402 module level import not at top
 
 
 def product_dict(**kwargs):
@@ -47,31 +60,7 @@ def product_dict(**kwargs):
         yield dict(zip(keys, instance))
 
 
-def main(
-    cache_dir: Path = DEFAULT_CACHE_DIR,
-    out_dir=None,
-    loglevel: Literal["debug", "info", "warning", "error"] = "info",
-):
-    cache_dir = Path(cache_dir)
-    cache_dir.mkdir(exist_ok=True)
-
-    out_dir = Path(out_dir) if out_dir else cache_dir
-    out_dir.mkdir(exist_ok=True)
-
-    fmt = "[%(asctime)s %(levelname)7s] %(message)s"
-    datefmt = "%Y-%m-%d %H:%M:%S"
-    logging.basicConfig(
-        level=loglevel.upper(),
-        format=fmt,
-        datefmt=datefmt,
-        handlers=[
-            logging.FileHandler(cache_dir / "offline_optimization_script.log"),
-            logging.StreamHandler(),
-        ],
-    )
-    logging.info(f"starting offline optimization script with cache_dir: {cache_dir}")
-    api = PtxboaAPI(data_dir=DEFAULT_DATA_DIR, cache_dir=cache_dir)
-
+def generate_param_sets(api):
     # these are the parameter dimensions that are relevant for the optimization
     param_arrays = {
         "scenario": api.get_dimension("scenario").index.tolist(),
@@ -98,23 +87,55 @@ def main(
         "secproc_co2": "Specific costs",
     }
 
-    n_total = np.prod([len(x) for x in param_arrays.values()])
-    logging.info(f"Total number of parameter combinations: {n_total}")
-    one_percent = n_total // 100
-    assert len(list(product_dict(**param_arrays))) == n_total
+    param_sets = [
+        param_set | static_params for param_set in product_dict(**param_arrays)
+    ]
+    return param_sets
+
+
+def main(
+    cache_dir: Path = DEFAULT_CACHE_DIR,
+    out_dir=None,
+    loglevel: Literal["debug", "info", "warning", "error"] = "info",
+):
+    cache_dir = Path(cache_dir)
+    cache_dir.mkdir(exist_ok=True)
+
+    out_dir = Path(out_dir) if out_dir else cache_dir
+    out_dir.mkdir(exist_ok=True)
+
+    # set up logging
+    fmt = "[%(asctime)s %(levelname)7s] %(message)s"
+    datefmt = "%Y-%m-%d %H:%M:%S"
+    logging.basicConfig(
+        level=loglevel.upper(),
+        format=fmt,
+        datefmt=datefmt,
+        handlers=[
+            logging.FileHandler(cache_dir / "offline_optimization_script.log"),
+        ],
+    )
+    logging.info(f"starting offline optimization script with cache_dir: {cache_dir}")
+    api = PtxboaAPI(data_dir=DEFAULT_DATA_DIR, cache_dir=cache_dir)
+
+    param_sets = generate_param_sets(api)
 
     results = []  # save results
-    for i, param_set in enumerate(product_dict(**param_arrays)):
-        logging.info(f"parameter combination {i} of {n_total}")
-        if i % one_percent == 0:
-            p = i / n_total * 100
-            logging.info(f"{p:.2f} % of all parameter combinations calculated")
+    for params in progress.bar.Bar(
+        suffix=(
+            "%(index)s/%(max)s, "
+            "%(percent)d%%, "
+            "elapsed %(elapsed_td)s, "
+            "eta %(eta_td)s"
+        )
+    ).iter(param_sets):
 
-        params = param_set | static_params
         result = {"params": params}
         try:
             logging.info(f"calculating parameter set {params}")
+
             _df, metadata = api.calculate(optimize_flh=True, **params)
+
             result["error"] = None
             result["result"] = metadata.get("flh_opt_hash")
         except Exception as e:

requirements-dev.txt

@@ -10,3 +10,4 @@ bumpversion>=0.5
 mkdocs>=1.5
 mkdocs-material>=9.4
 mkdocstrings-python>=1.7
+progress