Update benchmark: add --include-solve flag for PR #619 solution-unpack benchmark

Author: MaykThewessen

benchmark/scripts/benchmark_matrix_gen.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 """
-Benchmark script for linopy matrix generation performance.
+Benchmark script for linopy matrix generation and solution-unpacking performance.
 
 Covers the code paths optimised by PRs #616–#619:
   - #616  cached_property on MatrixAccessor (flat_vars / flat_cons)
@@ -11,13 +11,16 @@
 Usage
 -----
     # Quick run (24 snapshots only):
-    python dev-scripts/benchmark_matrix_gen.py --quick
+    python benchmark/scripts/benchmark_matrix_gen.py --quick
 
-    # Full sweep with JSON output:
-    python dev-scripts/benchmark_matrix_gen.py -o results.json --label "after-PR-616"
+    # Full matrix-generation sweep with JSON output:
+    python benchmark/scripts/benchmark_matrix_gen.py -o results.json --label "after-PR-616"
+
+    # Include solution-unpacking benchmark (requires HiGHS solver, #619):
+    python benchmark/scripts/benchmark_matrix_gen.py --include-solve -o results.json
 
     # Compare two runs:
-    python dev-scripts/benchmark_matrix_gen.py --compare before.json after.json
+    python benchmark/scripts/benchmark_matrix_gen.py --compare before.json after.json
 """
 
 from __future__ import annotations
@@ -38,19 +41,17 @@
 # Model builders
 # ---------------------------------------------------------------------------
 
-def build_scigrid(n_snapshots: int):
-    """Return a linopy Model from PyPSA SciGrid-DE with extended snapshots."""
+def build_scigrid_network(n_snapshots: int):
+    """Return a PyPSA Network (SciGrid-DE) with extended snapshots, without building the model."""
     import pypsa
 
     n = pypsa.examples.scigrid_de()
     orig_snapshots = n.snapshots
     orig_len = len(orig_snapshots)
 
-    # Create unique snapshots by extending from the first timestamp
     new_snapshots = pd.date_range(orig_snapshots[0], periods=n_snapshots, freq="h")
     n.set_snapshots(new_snapshots)
 
-    # Tile time-varying data to fill extended snapshots
     for component_t in (n.generators_t, n.loads_t, n.storage_units_t):
         for attr in list(component_t):
             df = getattr(component_t, attr)
@@ -60,7 +61,12 @@ def build_scigrid(n_snapshots: int):
                 setattr(component_t, attr, pd.DataFrame(
                     tiled, index=new_snapshots, columns=df.columns,
                 ))
+    return n
 
+
+def build_scigrid(n_snapshots: int):
+    """Return a linopy Model from PyPSA SciGrid-DE with extended snapshots."""
+    n = build_scigrid_network(n_snapshots)
     n.optimize.create_model(include_objective_constant=False)
     return n.model
 
@@ -152,6 +158,82 @@ def do_full():
     return results
 
 
+# ---------------------------------------------------------------------------
+# Solution-unpacking benchmark (#619)
+# ---------------------------------------------------------------------------
+
+def benchmark_solution_unpack(n_snapshots: int, repeats: int = 3) -> list[dict]:
+    """
+    Benchmark the solution-assignment loop in Model.solve (PR #619).
+
+    Strategy: solve once with HiGHS to get a real solution vector, then
+    re-run only the assignment loop (sol[idx] → var.solution) repeatedly
+    without re-solving, isolating the unpacking cost from solver time.
+    """
+    import xarray as xr
+
+    n = build_scigrid_network(n_snapshots)
+    n.optimize.create_model(include_objective_constant=False)
+    model = n.model
+
+    # Solve once to populate the raw solution
+    status, _ = model.solve(solver_name="highs", io_api="direct")
+    if status != "ok":
+        print(f"  WARNING: solve failed ({status}), skipping solution-unpack benchmark")
+        return []
+
+    # Reconstruct the raw solution Series (as returned by the solver):
+    # a float-indexed Series mapping variable label → solution value.
+    nan = float("nan")
+    parts = []
+    for name, var in model.variables.items():
+        if var.solution is None:
+            continue
+        labels = np.ravel(var.labels)
+        values = np.ravel(var.solution.values)
+        parts.append(pd.Series(values, index=labels.astype(float)))
+    if not parts:
+        print("  WARNING: no solution found on variables, was solve successful?")
+        return []
+    sol_series = pd.concat(parts).drop_duplicates()
+    sol_series.loc[-1] = nan
+
+    n_vars = sum(np.ravel(model.variables[name].labels).size for name in model.variables)
+    results = []
+
+    # ----- Old path (pandas label-based, pre-#619) -----
+    def unpack_pandas():
+        for name, var in model.variables.items():
+            idx = np.ravel(var.labels).astype(float)
+            try:
+                vals = sol_series[idx].values.reshape(var.labels.shape)
+            except KeyError:
+                vals = sol_series.reindex(idx).values.reshape(var.labels.shape)
+            var.solution = xr.DataArray(vals, var.coords)
+
+    results.append(time_phase(unpack_pandas, "unpack_pandas (before)", repeats))
+
+    # ----- New path (numpy dense array, #619) -----
+    def unpack_numpy():
+        sol_max_idx = int(max(sol_series.index.max(), 0))
+        sol_arr = np.full(sol_max_idx + 1, nan)
+        mask = sol_series.index >= 0
+        valid = sol_series.index[mask].astype(int)
+        sol_arr[valid] = sol_series.values[mask]
+        for name, var in model.variables.items():
+            idx = np.ravel(var.labels)
+            safe_idx = np.clip(idx, 0, sol_max_idx)
+            vals = sol_arr[safe_idx]
+            vals[idx < 0] = nan
+            var.solution = xr.DataArray(vals.reshape(var.labels.shape), var.coords)
+
+    results.append(time_phase(unpack_numpy, "unpack_numpy (after)", repeats))
+
+    for r in results:
+        r.update(model_type="scigrid_solve", size=n_snapshots, n_vars=n_vars, n_cons=0)
+    return results
+
+
 # ---------------------------------------------------------------------------
 # Runners
 # ---------------------------------------------------------------------------
@@ -161,7 +243,7 @@ def do_full():
 SYNTHETIC_SIZES = [20, 50, 100, 200]
 
 
-def run_benchmarks(model_type: str, quick: bool, repeats: int) -> list[dict]:
+def run_benchmarks(model_type: str, quick: bool, repeats: int, include_solve: bool = False) -> list[dict]:
     """Run benchmarks across problem sizes, return flat list of results."""
     all_results = []
 
@@ -185,6 +267,18 @@ def run_benchmarks(model_type: str, quick: bool, repeats: int) -> list[dict]:
             del model
             gc.collect()
 
+    if include_solve:
+        # Solution-unpacking benchmark for PR #619 (SciGrid-DE only, small sizes)
+        solve_sizes = QUICK_SNAPSHOTS if quick else [24, 100]
+        for n_snap in solve_sizes:
+            print(f"\n{'='*60}")
+            print(f"SciGrid-DE solve + unpack  {n_snap} snapshots  (#619)")
+            print(f"{'='*60}")
+            for r in benchmark_solution_unpack(n_snap, repeats):
+                all_results.append(r)
+                print(f"  {r['phase']:30s}  {r['best_s']:.4f}s  (median {r['median_s']:.4f}s)")
+            gc.collect()
+
     if model_type in ("synthetic", "all"):
         sizes = [20, 50] if quick else SYNTHETIC_SIZES
         for n in sizes:
@@ -246,6 +340,10 @@ def main():
         "--compare", nargs=2, metavar=("BEFORE", "AFTER"),
         help="Compare two JSON result files instead of running benchmarks",
     )
+    parser.add_argument(
+        "--include-solve", action="store_true",
+        help="Also benchmark solution unpacking (PR #619); requires HiGHS solver",
+    )
     args = parser.parse_args()
 
     if args.compare:
@@ -258,7 +356,7 @@ def main():
     print(f"Python {sys.version.split()[0]}, numpy {np.__version__}, "
           f"{platform.machine()}, {platform.system()}")
 
-    results = run_benchmarks(args.model, args.quick, args.repeats)
+    results = run_benchmarks(args.model, args.quick, args.repeats, args.include_solve)
 
     if args.output:
         out = {