chore: cleaning the code

engine/studies/mixed_lh.py

@@ -1,5 +1,7 @@
 """mixed design: full factorial on primary factors, latin hypercube on secondary"""
+
 import sys
+
 sys.path.insert(0, "..")
 import logging
 from itertools import product
@@ -16,6 +18,7 @@ log = logging.getLogger(__name__)
 
 LH_SAMPLES = 10
 
+
 def generate_configs(lh_samples: int = LH_SAMPLES):
     primary = [f for f in FACTORS if f.primary]
     secondary = [f for f in FACTORS if not f.primary]
@@ -28,7 +31,9 @@ def generate_configs(lh_samples: int = LH_SAMPLES):
         samples = lhs.random(n=lh_samples)
         for s in samples:
             sec_vals = {
-                secondary[i].name: secondary[i].levels[int(s[i] * len(secondary[i].levels))]
+                secondary[i].name: secondary[i].levels[
+                    int(s[i] * len(secondary[i].levels))
+                ]
                 for i in range(len(secondary))
             }
             base = {primary[i].name: p_combo[i] for i in range(len(primary))}
@@ -36,10 +41,13 @@ def generate_configs(lh_samples: int = LH_SAMPLES):
 
             for seed in range(SEEDS_PER_CONFIG):
                 cfg = {**base, "seed": seed}
-                cfg["id"] = hashlib.md5(json.dumps(cfg, sort_keys=True).encode()).hexdigest()[:8]
+                cfg["id"] = hashlib.md5(
+                    json.dumps(cfg, sort_keys=True).encode()
+                ).hexdigest()[:8]
                 configs.append(cfg)
     return configs
 
+
 def run_single(cfg: dict) -> dict:
     from engine.wrapper import PHANTOM
     import numpy as np
@@ -62,7 +70,8 @@ def run_single(cfg: dict) -> dict:
         obs, reward, term, trunc, _ = env.step(action)
         total_reward += reward
         steps += 1
-        if term: break
+        if term:
+            break
 
     env.close()
     return {
@@ -73,23 +82,33 @@ def run_single(cfg: dict) -> dict:
         "steps": steps,
     }
 
-def run_study(max_workers: int = None, output: str = "results_mixed.jsonl", lh_samples: int = LH_SAMPLES):
+
+def run_study(
+    max_workers: int = None,
+    output: str = "results_mixed.jsonl",
+    lh_samples: int = LH_SAMPLES,
+):
     configs = generate_configs(lh_samples)
     n_primary_cells = int(np.prod([len(f.levels) for f in FACTORS if f.primary]))
-    log.info(f"mixed LH: {len(configs)} configs ({n_primary_cells} primary × {lh_samples} LH × {SEEDS_PER_CONFIG} seeds)")
+    log.info(
+        f"mixed LH: {len(configs)} configs ({n_primary_cells} primary × {lh_samples} LH × {SEEDS_PER_CONFIG} seeds)"
+    )
 
     results = []
     with ProcessPoolExecutor(max_workers=max_workers) as ex:
         for i, result in enumerate(ex.map(run_single, configs)):
             results.append(result)
-            if (i+1) % 100 == 0: log.info(f"progress: {i+1}/{len(configs)}")
+            if (i + 1) % 100 == 0:
+                log.info(f"progress: {i + 1}/{len(configs)}")
 
     Path(output).write_text("\n".join(json.dumps(r) for r in results))
     log.info(f"wrote {len(results)} results to {output}")
     return results
 
+
 if __name__ == "__main__":
     import argparse
+
     p = argparse.ArgumentParser()
     p.add_argument("--workers", type=int, default=None)
     p.add_argument("--output", default="results_mixed.jsonl")
@@ -100,7 +119,9 @@ if __name__ == "__main__":
     primary = [f for f in FACTORS if f.primary]
     secondary = [f for f in FACTORS if not f.primary]
     configs = generate_configs(args.lh_samples)
-    log.info(f"design: {len(configs)} runs | primary: {[f.name for f in primary]} | secondary (LH): {[f.name for f in secondary]}")
+    log.info(
+        f"design: {len(configs)} runs | primary: {[f.name for f in primary]} | secondary (LH): {[f.name for f in secondary]}"
+    )
 
     if not args.dry_run:
         run_study(args.workers, args.output, args.lh_samples)