Refactor simulation to use dedicated class

schedgen/opt/simulate.py

+#!/usr/bin/env python3
+
+import numpy as np
+
+
+def runtime(schedule, frame, timeout):
+    result = simulate(schedule, frame, timeout=timeout)
+    total = len(result)
+    result = result.loc[result < float("inf")]
+    return (len(result), total, result.sum())
+
+
+def simulate(schedule, frame, timeout=None, mu=0, sigma=0, seed=None):
+    exps = frame.copy()
+
+    rng = np.random.default_rng(seed)
+
+    sc = "solved"
+    # Find a unique name
+    while sc in exps.columns:
+        sc = sc + "_"
+    exps[sc] = float("inf")
+
+    time = 0.0
+    for strategy_time, strategy in schedule:
+        jitter = rng.normal(mu, sigma, exps.shape[0])
+
+        solved = (exps[strategy] + jitter).loc[
+            (exps[strategy] + jitter) <= strategy_time
+        ]
+
+        solving_time = solved.clip(0.0) + time
+        solving_time.name = sc
+        time = time + strategy_time
+
+        update_idx = exps.loc[exps[sc] == float("inf")].index.intersection(solved.index)
+        exps.update(solving_time.loc[update_idx])
+
+    if timeout:
+        exps[sc].loc[exps[sc] > timeout] = float("inf")
+
+    return exps[sc]

schedgen/schedgen-optimize.py

@@ -9,6 +9,7 @@ from schedgen.opt.schedule import (
     read_from_file,
 )
 from schedgen.opt.order import OptimizeOrder
+from schedgen.opt.simulate import runtime
 from schedgen.opt.helper import message, file_or_dir_path, new_file_path, file_path
 from schedgen.parsers.gridtpt import GridTPT
 from schedgen.parsers.csvresult import CSVResult
@@ -137,19 +138,15 @@ args_parser.add_argument(
 )
 
 
-def writeScript(outFile, logics, schedule, seconds=False):
+def writeScript(outFile, logics, schedule, timeout, seconds=False):
     message(f"Writing schedule to {outFile}.", logics)
     with open(outFile, "w") as af:
         if len(logics) == 1:
             af.write(f"    {logics[0]})\n")
         else:
             af.write(f"    {logics})\n")
-        optimizer = OptimizeOrder(r.frame)
-        (_, time, solved) = optimizer.runtime(schedule)
-        num_solved = solved.sum()
-        af.write(
-            f"        # Prediction: {num_solved}/{solved.count()} in {time:.2f}s\n"
-        )
+        (solved, out_of, in_time) = runtime(schedule, r.frame, timeout)
+        af.write(f"        # Prediction: {solved}/{out_of} in {in_time:.2f}s\n")
         for i in range(len(schedule) - 1):
             (t, e) = schedule[i]
             if seconds:
@@ -171,9 +168,9 @@ def writeCSV(outFile, logics, schedule):
             af.write(f"{t:.3f};{e}\n")
 
 
-def writeSchedule(outFile, logics, full, schedule, seconds=False):
+def writeSchedule(outFile, logics, full, schedule, timeout, seconds=False):
     if full:
-        writeScript(outFile, logics, schedule, seconds)
+        writeScript(outFile, logics, schedule, timeout, seconds)
     else:
         writeCSV(outFile, logics, schedule)
 
@@ -295,12 +292,12 @@ if __name__ == "__main__":
             longest = time_slices[-1]
             message(f"  For {t:.2f}s\t'{best_opt}'", logics)
             order_optimizer = OptimizeOrder(exps)
-            _, best_time, best_solves = order_optimizer.runtime([(longest, best_opt)])
-            best_count = best_solves.sum()
+            (best_count, _, best_time) = runtime(
+                [(longest, best_opt)], r.frame, timeout
+            )
             message(f"This solves {best_count} in {best_time}s.", logics)
             for e in exps:
-                _, time, solves = order_optimizer.runtime([(longest, e)])
-                new_count = solves.sum()
+                (new_count, _, time) = runtime([(longest, e)], r.frame, timeout)
                 if new_count > best_count:
                     message(
                         "Something went wrong: optimizer returned a better schedule.",
@@ -313,7 +310,6 @@ if __name__ == "__main__":
             message("The best schedule is:", logics)
             message(f"  Run\t'{best_opt}'", logics)
             message(f"This solves {best_count} in {best_time}s.", logics)
-            after_solved = best_solves
             after_count = best_count
             after_total_time = best_time
             best_schedule = [(t, best_opt)]
@@ -324,14 +320,11 @@ if __name__ == "__main__":
 
             order_optimizer = OptimizeOrder(exps)
 
-            (
-                before_wait,
-                before_total_time,
-                before_solved,
-            ) = order_optimizer.runtime(schedule)
-            before_count = before_solved.value_counts().loc[True]
+            (before_count, total_count, before_total_time) = runtime(
+                schedule, r.frame, timeout
+            )
             message(
-                f"Current schedule solves {before_count}/{before_solved.count()} in {before_total_time:.2f}s.",
+                f"Current schedule solves {before_count}/{total_count} in {before_total_time:.2f}s.",
                 logics,
             )
 
@@ -349,12 +342,11 @@ if __name__ == "__main__":
                 message("\nTweaked schedule:", logics)
             for (t, e) in best_schedule:
                 message(f"  For {t:.2f}s\t'{e}'", logics)
-            after_wait, after_total_time, after_solved = order_optimizer.runtime(
-                best_schedule
+            (after_count, total_count, after_total_time) = runtime(
+                best_schedule, r.frame, timeout
             )
-            after_count = after_solved.value_counts().loc[True]
             message(
-                f"This schedule solves {after_count}/{after_solved.count()} in {after_total_time:.2f}s.",
+                f"This schedule solves {after_count}/{total_count} in {after_total_time:.2f}s.",
                 logics,
             )
 
@@ -365,6 +357,8 @@ if __name__ == "__main__":
         full_schedule = best_schedule
 
     if args.result:
-        writeSchedule(args.result, logics, args.full, full_schedule, args.seconds_out)
+        writeSchedule(
+            args.result, logics, args.full, full_schedule, timeout, args.seconds_out
+        )
 
     message("All done.", logics)

schedgen/schedgen-simulate.py

@@ -6,6 +6,7 @@ import numpy as np
 from datetime import datetime
 from schedgen.opt.schedule import read_from_file
 from schedgen.opt.helper import message, file_or_dir_path, new_file_path, file_path
+from schedgen.opt.simulate import simulate
 from schedgen.parsers.gridtpt import GridTPT
 from schedgen.parsers.csvresult import CSVResult
 
@@ -161,35 +162,9 @@ if __name__ == "__main__":
     pre_schedule = None
     pre_schedule = read_from_file(args.schedule)
 
-    sc = "solved"
-    # Find a unique name
-    while sc in exps.columns:
-        sc = sc + "_"
-
-    exps[sc] = float("inf")
-
-    time = 0.0
-    rng = np.random.default_rng(args.seed)
-
-    # Some pandas magic to do the simulation
-    for strategy_time, strategy in pre_schedule:
-        jitter = rng.normal(args.mu, args.sigma, exps.shape[0])
-
-        solved = (exps[strategy] + jitter).loc[
-            (exps[strategy] + jitter) <= strategy_time
-        ]
-
-        solving_time = solved.clip(0.0) + time
-        solving_time.name = sc
-        time = time + strategy_time
-
-        update_idx = exps.loc[exps[sc] == float("inf")].index.intersection(solved.index)
-        exps.update(solving_time.loc[update_idx])
-
-    if timeout:
-        exps[sc].loc[exps[sc] > timeout] = float("inf")
+    result = simulate(pre_schedule, exps, timeout, args.mu, args.sigma, args.seed)
 
     message("Writing output.", logics)
-    writeResult(args, exps[sc])
+    writeResult(args, result)
 
     message("All done.", logics)