Add a ml user friendly API for scenario sweep

README.md

@@ -186,6 +186,22 @@ V = sweep.compute(batch_size=512)
 disconnected = sweep.get_disconnected()  # which rows a topology trip islanded
 ```
 
+For PyTorch users, the same batch is available as a differentiable layer that
+takes lightsim2grid-style per-element inputs (`True = connected` statuses),
+solves every row in one GPU pass and back-propagates through it with the
+adjoint method — consecutive calls with the same batch size reuse the whole GPU
+setup (no new cuDSS analysis, refactorization only), and the transposed system
+the backward needs is built lazily on the first `backward()`:
+
+```python
+from gpusim2grid.differentiable import BatchPowerFlow
+pf = BatchPowerFlow.from_lsgrid(grid, nb_iter=8)
+V = pf(load_p=load_p, load_q=load_q, gen_p=gen_p, gen_v=gen_v,   # (n_scen, n_elem) tensors
+       line_status=line_status, trafo_status=trafo_status)        # bool, True = connected
+loss = (V.abs() - 1.0).pow(2).sum()
+loss.backward()   # gradients w.r.t. load_p, load_q, gen_p and gen_v
+```
+
 See [`examples/`](examples/):
 
 - `ieee14_basic.py` — end-to-end AC power flow on the IEEE 14-bus case.
@@ -196,6 +212,7 @@ See [`examples/`](examples/):
 - `limit_violations.py` — fused per-contingency bus voltage / branch current limit checking.
 - `distributed_slack.py` — augmented solve (distributed slack in the Jacobian) via the lightsim2grid bridge.
 - `differentiable_pf.py` — derivatives through a single power flow via the adjoint method.
+- `batch_differentiable_pf.py` — `BatchPowerFlow`: a batch of scenarios (injections, set-points, branch statuses) as one differentiable PyTorch layer, trained over a few steps.
 
 ## How it works
 
@@ -253,8 +270,9 @@ design. See the [roadmap](#roadmap) below for areas where help is especially wel
 
 Directions we plan to pursue — **any help or ideas are very welcome**:
 
-- **Extend derivatives to the injection sweep**, and later to the full contingency
-  analysis path (currently differentiation is limited to a single power flow).
+- **Extend derivatives** beyond `BatchPowerFlow`'s inputs (`load_p`, `load_q`,
+  `gen_p`, `gen_v`): generator contingencies (`gen_status`), a JAX front-end on the
+  same DLPack primitives, and the plain contingency-analysis path.
 - **Best action selector** — given one or several grid snapshots and a list of
   candidate actions, find the best action(s) to apply, by evaluating the candidates in
   batch on the GPU.

docs/api.rst

@@ -150,6 +150,123 @@ violations") for the equivalent ``ContingencyAnalysisGPU`` walkthroughs — the
 only difference on :class:`~gpusim2grid.ScenarioSweepGPU` is that each row
 also carries its own injection.
 
+.. _scenario-sweep-reuse:
+
+What ``compute()`` reuses across calls
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The GPU batch driver built by the first ``compute()`` is kept alive by the
+sweep and reused by later calls. Each ``compute()`` compares the current
+settings with the ones the live driver was built with and takes one of three
+paths:
+
+* **cold** — no driver yet, or something the driver's shape depends on
+  changed: the number of rows, ``batch_size``, ``strategy``,
+  ``refactor_period``, a cuDSS choice (``reordering_alg`` /
+  ``matching_alg`` / ``pivot_epsilon_alg``), the step-scaling knobs,
+  ``handle_disconnected_grid``, ``fixed_batch_capacity``, or the base state
+  (``set_contingency_gens`` reserving a different set of buses);
+* **warm** — same shape, but ``set_topology`` or ``set_contingency_gens`` was
+  called since the last ``compute()``;
+* **hot** — same shape and topology; only ``set_injections`` /
+  ``set_injections_from_elements`` / ``set_gen_v`` were called (or nothing).
+
+``nb_iter`` is applied to the live driver and never forces a rebuild.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 34 22 22 22
+
+   * - Phase
+     - cold
+     - warm
+     - hot
+   * - Host: per-unit Sbus build from the MW / MVAr matrices (numpy path
+       only; the device path, ``set_injections_dlpack``, skips it)
+     - yes, if injections changed
+     - yes, if injections changed
+     - yes, if injections changed
+   * - Host: topology preprocessing — Ybus patch triplets → CSR positions,
+       connectivity / masking, flat patch arrays, tripped-branch table,
+       per-row PV pins and slack weights (``ScenarioSweepBatch``)
+     - yes
+     - yes
+     - no
+   * - Device: allocation of the chunk buffers (V, Ybus values, J values, F,
+       dx, Ibus), block-diagonal CSR structure, cuSPARSE SpMV descriptor
+     - yes
+     - no
+     - no
+   * - cuDSS context creation + ANALYSIS (symbolic factorization)
+     - yes
+     - no
+     - no
+   * - H→D upload of the patch / mask / tripped-branch arrays
+     - yes
+     - yes
+     - no
+   * - Sbus rows: one H→D copy (numpy) or D→D copy (torch) into the
+       original-order buffer, then a device gather into batch order
+     - yes
+     - yes
+     - yes
+   * - ``gen_v`` rows (same copy + gather, Vm-fixed columns only)
+     - yes, if set
+     - yes, if set
+     - only if it changed
+   * - Per chunk: tile V and Ybus, apply the row's Ybus patches, reseed
+       ``gen_v``, slice Sbus
+     - yes
+     - yes
+     - yes
+   * - Newton-Raphson iterations: SpMV, fill F, fill J
+     - yes
+     - yes
+     - yes
+   * - cuDSS first FACTORIZATION (iteration 0 of the first chunk)
+     - yes
+     - no
+     - no
+   * - cuDSS REFACTORIZATION (every other iteration, per the strategy)
+     - yes
+     - yes
+     - yes
+   * - cuDSS SOLVE, voltage update, residuals, result store
+     - yes
+     - yes
+     - yes
+   * - Limits (``compute_limit_violations``): branch admittance upload /
+       per-row sentinel reset
+     - yes / yes
+     - no / yes
+     - no / yes
+   * - Reported one-time timings (``t_analysis_ms``, ``t_alloc_ms``,
+       ``t_context_init_ms``; ``t_preprocess_ms``, ``t_source_init_ms``)
+     - measured
+     - 0 ; measured
+     - 0 ; 0
+   * - Counters bumped
+     - ``driver_build_counter``, ``source_build_counter``
+     - ``source_build_counter``
+     - —
+
+The result buffer behind ``v_results_dlpack()`` is overwritten in place by a
+warm or hot call and freed (reallocated) by a cold one — clone the tensor for
+a snapshot either way.
+
+For the differentiable layer (:class:`~gpusim2grid.differentiable.BatchPowerFlow`)
+the same table applies to its forward, with one addition when gradients are
+requested: the batched Jacobian is refilled at the converged voltages after
+the iterations (one extra ``fill J``, on every path). Its backward has its own
+lazily built state: the first ``backward()`` transposes the Jacobian pattern,
+builds the J→Jᵀ position map and buffers and runs one cuDSS ANALYSIS + one
+FACTORIZATION of Jᵀ; every later backward only permutes the values with a
+kernel, REFACTORIZES (once per new forward — a second backward on the same
+forward only solves) and SOLVES. A cold forward discards that state, and the
+next backward rebuilds it. ``timings.adjoint_n_analysis`` /
+``adjoint_n_factorize`` / ``adjoint_n_refactorize`` / ``adjoint_n_solve``
+count these over the driver's life.
+
 .. automodule:: gpusim2grid.scenario_sweep
    :members:
    :undoc-members:
@@ -164,6 +281,40 @@ Single AC power flow
 Differentiable power flow (alpha)
 ---------------------------------
 
+Two entry points, both PyTorch ``autograd`` integrations of the GPU solver
+using the adjoint method (the converged Jacobian is transposed and factorized
+once, then reused for every backward):
+
+* :class:`~gpusim2grid.differentiable.BatchPowerFlow` -- a **batch** of
+  scenarios as one differentiable layer, driven by the same per-element inputs
+  as lightsim2grid's ``ScenarioSweep``: ``load_p``, ``load_q``, ``gen_p``,
+  ``gen_v`` (all differentiable) and the boolean ``line_status`` /
+  ``trafo_status`` masks (``True`` = connected). Built once from a solved
+  lightsim2grid grid; consecutive calls with the same number of rows reuse the
+  GPU batch driver (no cuDSS analysis, refactorization only), and the
+  transposed system is built lazily on the first ``backward()`` — see
+  :ref:`scenario-sweep-reuse` for the phase-by-phase table.
+* :func:`~gpusim2grid.differentiable.solve_power_flow` /
+  :class:`~gpusim2grid.differentiable.PowerFlowFunction` -- a single power flow
+  from raw ``Sbus`` tensors.
+
+.. code-block:: python
+
+    from gpusim2grid.differentiable import BatchPowerFlow
+
+    pf = BatchPowerFlow.from_lsgrid(grid, nb_iter=8)          # grid: solved lightsim2grid grid
+    V = pf(load_p=load_p, load_q=load_q, gen_p=gen_p,          # (n_scen, n_load / n_gen) MW, MVAr
+           gen_v=gen_v,                                        # (n_scen, n_gen) vm_pu, NaN = keep
+           line_status=line_status, trafo_status=trafo_status) # bool, True = connected
+    loss = ((V.abs() - 1.0) ** 2).sum()
+    loss.backward()                                            # d loss / d load_p, ..., d gen_v
+
+Rows whose branch trips island the grid come back as ``NaN`` (mask them out of
+the loss); they get a zero gradient. A forward → forward → backward(first)
+pattern is refused with a clear error unless the layer was built with
+``snapshot_jacobian=True`` (one extra device copy of the Jacobians per forward),
+which is also what ``torch.autograd.gradcheck`` needs.
+
 .. automodule:: gpusim2grid.differentiable
    :members:
    :undoc-members:

docs/conf.py

@@ -25,7 +25,7 @@
     import gpusim2grid
     release = gpusim2grid.__version__
 except Exception:  # pragma: no cover - best effort
-    release = "0.1.2"
+    release = "0.2.0.rc0"
 version = release
 
 # -- General configuration ---------------------------------------------------

docs/examples.rst

@@ -92,3 +92,16 @@ power flow, using the adjoint method via the PyTorch integration.
 
 .. literalinclude:: ../examples/differentiable_pf.py
    :language: python
+
+Batched differentiable power flow
+---------------------------------
+
+A whole batch of scenarios (random load scalings and random N-1 line trips) as
+one differentiable PyTorch layer (``BatchPowerFlow``): a few optimisation steps
+learn a generator redispatch and voltage set-points that flatten the voltage
+profile. The reuse counters printed at each step show the GPU batch driver
+being built once, the Jacobians being refactorized (never re-analysed) on
+later calls, and the transposed system being built on the first backward only.
+
+.. literalinclude:: ../examples/batch_differentiable_pf.py
+   :language: python

examples/README.md

@@ -21,6 +21,7 @@ importable.)
 | `limit_violations.py` | N-1 screen with `compute_limit_violations=True`: fused, on-device per-contingency bus voltage / branch current / divergence check, reported as a bounded per-contingency violation list. Takes an optional `<case>` argument. |
 | `distributed_slack.py` | Augmented solve: distributed slack carried in the Jacobian (via the lightsim2grid C++ bridge), matched against the CPU reference. Same path also covers HVDC droop, SVC, and remote voltage control. Needs a bridge-enabled build. |
 | `differentiable_pf.py` | Derivatives through a single power flow (adjoint method) via the PyTorch autograd integration. Requires PyTorch with CUDA. |
+| `batch_differentiable_pf.py` | `BatchPowerFlow`: a whole batch of scenarios (load / generator injections, generator set-points, line / trafo statuses) as one differentiable PyTorch layer; a few optimisation steps, with the GPU-reuse counters printed at each call. Requires PyTorch with CUDA. |
 
 `_common.py` is a shared helper (not a standalone example): it wraps
 lightsim2grid/pandapower to produce the plain NumPy/SciPy arrays the solver

examples/batch_differentiable_pf.py

@@ -0,0 +1,86 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+"""
+batch_differentiable_pf.py — a batch of power flows as a differentiable
+PyTorch layer, driven by lightsim2grid-style per-element inputs.
+
+``BatchPowerFlow`` takes, per scenario (row), the load / generator injections,
+the generator voltage set-points and the line / trafo statuses, solves the whole
+batch in ONE GPU pass, and back-propagates through it with the adjoint method.
+Here a tiny model learns a generator redispatch + voltage set-points that keep
+every bus of the IEEE 14-bus grid near 1.0 pu under random load scalings and
+random N-1 line trips.
+
+The point of the timings printed at the end: the first call builds the GPU
+batch driver (cuDSS analysis + first factorization); every later call with the
+same batch size reuses it (only the new rows move to the GPU, the Jacobians are
+refactorized), and the transposed system the backward needs is built once, on
+the first backward, then only refactorized.
+
+Requires PyTorch with CUDA. Run:
+    python examples/batch_differentiable_pf.py
+"""
+import numpy as np
+
+from _common import load_case
+from gpusim2grid.compilation_options import is_fp32
+
+
+def main():
+    import torch
+
+    if not torch.cuda.is_available():
+        raise SystemExit("This example requires a CUDA-capable PyTorch build.")
+    if is_fp32:
+        print("NOTE: gpusim2grid was built in FP32; gradients are most accurate "
+              "against an FP64 build.")
+
+    from gpusim2grid.differentiable import BatchPowerFlow
+
+    grid = load_case("case14")["grid"]
+    pf = BatchPowerFlow.from_lsgrid(grid, nb_iter=8, tol_base=1e-10)
+    dev = pf.device
+    rdt = torch.float32 if is_fp32 else torch.float64
+
+    n_scen = 64
+    rng = np.random.default_rng(0)
+    # Random load scalings and one random line trip per row (True = connected).
+    scales = torch.tensor(rng.uniform(0.8, 1.2, size=(n_scen, 1)), dtype=rdt, device=dev)
+    load_p = pf._load_p_base[None, :] * scales
+    load_q = pf._load_q_base[None, :] * scales
+    line_status = torch.ones(n_scen, pf.n_line, dtype=torch.bool, device=dev)
+    line_status[torch.arange(n_scen), torch.tensor(rng.integers(2, 12, size=n_scen))] = False
+
+    # Learnable per-generator redispatch (MW) and voltage set-points (pu),
+    # shared across the batch; the slack absorbs the mismatch.
+    dp = torch.zeros(pf.n_gen, dtype=rdt, device=dev, requires_grad=True)
+    vset = torch.full((pf.n_gen,), 1.04, dtype=rdt, device=dev, requires_grad=True)
+    opt = torch.optim.Adam([dp, vset], lr=5e-3)
+
+    for step in range(15):
+        opt.zero_grad()
+        gen_p = (pf._gen_p_base + dp)[None, :].expand(n_scen, -1)
+        gen_v = vset[None, :].expand(n_scen, -1)
+        V = pf(load_p=load_p, load_q=load_q, gen_p=gen_p, gen_v=gen_v,
+               line_status=line_status)
+        valid = torch.isfinite(V.real)                       # islanded rows are NaN
+        vm = torch.where(valid, V, torch.ones_like(V)).abs()
+        loss = ((vm - 1.0) ** 2 * valid).sum() / valid.sum() + 1e-4 * (dp ** 2).sum()
+        loss.backward()
+        opt.step()
+        t = pf.timings
+        print(f"step {step:2d}  loss {loss.item():.3e}  "
+              f"driver builds {pf.sweep.driver_build_counter}  "
+              f"analysis {t.t_analysis_ms:6.2f} ms  first-factorize {t.t_first_factorize.wall_ms:5.2f} ms  "
+              f"refactorize x{t.n_refactorize}  "
+              f"adjoint: analysis x{t.adjoint_n_analysis} factorize x{t.adjoint_n_factorize} "
+              f"refactorize x{t.adjoint_n_refactorize} solve x{t.adjoint_n_solve}")
+
+    print(f"learned redispatch (MW): {dp.detach().cpu().numpy().round(3)}")
+    print(f"learned set-points (pu): {vset.detach().cpu().numpy().round(4)}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -51,6 +51,11 @@ docs = [
   "sphinx-rtd-theme>=2",
   "myst-parser>=2",
 ]
+# gpusim2grid.differentiable (BatchPowerFlow / PowerFlowFunction) needs a
+# CUDA-enabled PyTorch; everything else works without it.
+torch = [
+  "torch",
+]
 
 [project.urls]
 Homepage = "https://github.com/Grid2Op/gpusim2grid"

src/_cpp/acpf_nr.cu

@@ -66,6 +66,7 @@
 #include <iostream>
 #include <vector>
 #include <cassert>
+#include <cmath>       // isnan / NAN (NanMaxFunctor)
 
 // ---------------------------------------------------------------------------
 // Error-checking macros — throw on failure so the AcPfNrState constructor
@@ -109,6 +110,17 @@ struct AbsFunctor {
     cuda_real_type operator()(cuda_real_type x) const { return fabs(x); }
 };
 
+// NaN-propagating max for the ‖F‖∞ reductions below. thrust::maximum is
+// `a < b ? b : a`, which drops a NaN `b`: an all-NaN F would reduce to the
+// initial 0 and read as converged. Same sticky-NaN rule as the batched
+// compute_residuals_kernel.
+struct NanMaxFunctor {
+    __host__ __device__
+    cuda_real_type operator()(cuda_real_type a, cuda_real_type b) const {
+        return (isnan(a) || isnan(b)) ? cuda_real_type(NAN) : (b > a ? b : a);
+    }
+};
+
 // =============================================================================
 // §2  CPU helper: build_J_structure
 //   Derives the J CSR sparsity skeleton and the four dS scatter maps from the
@@ -1181,7 +1193,7 @@ AcPfNrState::AcPfNrState(
             d_F.begin(), d_F.end(),
             AbsFunctor{},
             cuda_real_type(0.),
-            thrust::maximum<cuda_real_type>());
+            NanMaxFunctor{});
         timings.t_mismatch += timer.stop_ms();
 
         // Guard against tol being tuned for lightsim2grid's FP64 CPU check
@@ -1236,7 +1248,7 @@ AcPfNrState::AcPfNrState(
                 d_F.begin(), d_F.end(),
                 AbsFunctor{},
                 cuda_real_type(0.),
-                thrust::maximum<cuda_real_type>());
+                NanMaxFunctor{});
             timings.t_mismatch += timer.stop_ms();
 
             if (norm_F < static_cast<cuda_real_type>(tol)) {