diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 9b09538..1b23ebc 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -9,3 +9,9 @@ updates: directory: "/" schedule: interval: "weekly" + ignore: + # The v5→v7 bump silently broke coverage uploads in the sibling + # DecisionRules.jl repo. Keep codecov-action pinned until a deliberate, + # verified migration — see the comment in .github/workflows/CI.yml. + - dependency-name: "codecov/codecov-action" + update-types: ["version-update:semver-major"] diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index d6b04cb..be1a6df 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -17,6 +17,7 @@ jobs: permissions: actions: write contents: read + id-token: write # OIDC token for tokenless Codecov uploads (see codecov step) strategy: fail-fast: false matrix: @@ -37,8 +38,22 @@ jobs: - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 - uses: julia-actions/julia-processcoverage@v1 - - uses: codecov/codecov-action@v7 + # Pinned to v5: codecov-action@v7 silently stopped uploading coverage in + # the sibling DecisionRules.jl repo (Codecov "Missing Head Commit" on + # PRs); v5 is the last version verified to upload from this workflow + # shape. Before re-bumping, migrate deliberately and confirm a commit + # appears on Codecov. + # Authentication uses OIDC (`use_oidc` + the job's `id-token: write` + # permission) because the CODECOV_TOKEN secret is not set in this repo + # ("Token length: 0" in CI) and tokenless uploads are rejected on + # protected branches. OIDC requires the Codecov GitHub App to be + # installed for the organization AND this repository to be activated on + # codecov.io (it currently is not); if uploads fail with an OIDC error, + # either install the app + activate the repo, or set the CODECOV_TOKEN + # secret and replace `use_oidc` with `token: ${{ secrets.CODECOV_TOKEN }}`. + - uses: codecov/codecov-action@v5 with: files: lcov.info - token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: false + use_oidc: true + # Fail loudly so upload breakage is visible instead of silent. + fail_ci_if_error: true diff --git a/.gitignore b/.gitignore index a432b01..8f248ab 100644 --- a/.gitignore +++ b/.gitignore @@ -42,3 +42,5 @@ logs/ # Slurm batch scripts (user-specific, not part of the package) *.sbatch *.sh +examples/HydroPowerModels/results/ +*strong.json diff --git a/Project.toml b/Project.toml index 42ea84c..f07852b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,33 +1,42 @@ name = "DecisionRulesExa" uuid = "7c3e91a4-d8f2-4b6a-9e15-a2c4f7b80d53" -authors = ["Andrew Rosemberg and contributors"] version = "0.1.0" +authors = ["Andrew Rosemberg and contributors"] [deps] -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [compat] +CSV = "0.10.16" CUDA = "6" ChainRulesCore = "1.26" ExaModels = "0.11" Flux = "0.16" +JSON = "1.6.1" MadNLP = "0.10" MadNLPGPU = "0.10" NLPModels = "0.21" +Tables = "1.12.1" Zygote = "0.7" julia = "1.10, 1.11, 1.12" [extras] +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Test", "Statistics"] diff --git a/README.md b/README.md index a06eca1..57cc681 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,16 @@ train_tsddr( ) ``` +> **Note on the uncertainty parameter**: `train_tsddr` writes the full sampled +> trajectory (length `T * nw`) into `p_uncertainty` with +> `ExaModels.set_parameter!`, which enforces an exact size match (ExaModels ≥ +> 0.11). The `p_w` built by `build_deterministic_equivalent` / +> `build_linear_tracking_problem` holds only the `(T - 1) * nw` dynamics +> entries, so for `train_tsddr` your NLP needs an uncertainty parameter of +> length `T * nw` (as the Hydro example's `p_inflow` is). See the +> `"train_tsddr open-loop smoke test"` testset in `test/runtests.jl` for a +> minimal full-length variant of the problem above. + For GPU, replace `backend = nothing` with `backend = CUDABackend()` and add `linear_solver = CUDSSSolver` to `madnlp_kwargs`. ## What you need to provide @@ -71,7 +81,59 @@ For a custom problem you need: - **An uncertainty sampler** `() -> w_flat` returning a flat `Float32`/`Float64` vector of length `T * nw`. - **A Flux policy** (LSTM or MLP) mapping `(w_t, x_{t-1})` to target `x_t` at each stage. -The package provides `build_deterministic_equivalent` for generic problems and `build_linear_tracking_problem` as a ready-made demo. For domain-specific models (power systems, robotics), build the ExaModels NLP directly — see `examples/HydroPowerModels/` for a complete AC-OPF example. +The package provides `build_deterministic_equivalent` for generic problems and `build_linear_tracking_problem` as a ready-made demo. For domain-specific models, build the ExaModels NLP directly; `examples/BatteryStorageOPF/` contains an AC-OPF battery example. + +## Strict reachable target equality + +The usual TS-DDR deterministic equivalent uses slack-penalized target +constraints, + +```text +x_t - pi_theta(w_t, x_{t-1}) = delta_t, +objective += rho * penalty(delta_t). +``` + +This is the right default for open-loop target trajectories and for cases where +the policy can request states that are not reachable from the previous realized +state. The target multipliers are then gradients of the penalized projection +problem, so their quality depends on the penalty calibration. + +For policies whose output is guaranteed to lie in a one-stage reachable state +set, a stricter formulation is possible: + +```text +x_t = pi_theta(w_t, x_{t-1}), pi_theta(w_t, x_{t-1}) in R(w_t, x_{t-1}). +``` + +In that case the deterministic equivalent does not need target slack variables +or target penalties. The multiplier on the equality is the local envelope +sensitivity of the true stage problem with respect to the policy-imposed next +state, not the sensitivity of a penalized approximation. This is useful when: + +- users can define a differentiable or piecewise differentiable map into a + subset of the one-stage reachable set; +- total recourse is guaranteed by the model for every state produced by that + map. + +There are two strict hydro paths: + +- **Embedded strict DE** evaluates the policy inside the NLP against realized + reservoir states. This is the usual way to make strict mode safe because the + policy sees the state from which its next target must be reachable. +- **Regular strict DE with reachable rollout** computes targets before solving + the NLP, but starts from the true initial state and feeds the previous target + back to the reachable policy. If + `x̂_t ∈ R(x̂_{t-1}, w_t)` and `x̂_0 = x_0`, the full target path is feasible by + induction. The strict equality then forces the realized path to equal that + reachable target path. + +Do not use strict equality for a generic open-loop target policy. For +unreachable targets, the slack-penalty formulation is the robust fallback. + +The hydro reachable policy keeps recurrence over inflows only. Optional +`combiner_layers` / `DR_HEAD_LAYERS` add a nonlinear feed-forward map from +`[encoded_inflow; reservoir_state]` to targets without adding recurrence over +the state input. ## Parallel GPU solves @@ -217,7 +279,24 @@ Choose DecisionRules.jl when: - [`examples/end_to_end_cpu.jl`](examples/end_to_end_cpu.jl) — minimal CPU demo with a linear tracking problem - [`examples/end_to_end_gpu.jl`](examples/end_to_end_gpu.jl) — same demo on GPU with CUDSS -- [`examples/HydroPowerModels/`](examples/HydroPowerModels/) — full multi-stage hydrothermal scheduling with DC and AC OPF +- [`examples/BatteryStorageOPF/`](examples/BatteryStorageOPF/) — reproducible PGLib AC-OPF cases with linear battery storage + +## Repository Map + +| Path | Purpose | +|---|---| +| `src/DecisionRulesExa.jl` | Module entrypoint and public exports | +| `src/policy.jl` | MLP, state-conditioned LSTM policies, bounded target policies, nonlinear target heads | +| `src/deterministic_equivalent.jl` | Generic open-loop deterministic-equivalent builder and solve helpers | +| `src/embedded_deterministic_equivalent.jl` | Generic embedded-policy deterministic equivalent with nonlinear oracle | +| `src/training.jl` | `train_tsddr`, embedded training, solver retry/warm-start handling | +| `src/rollout.jl` | Stage-wise rollout evaluation for ExaModels problems | +| `src/critic_control_variate.jl` | Scalar critic/control-variate helpers | +| `src/utils.jl` | Indexing and small shared utilities | +| `examples/end_to_end_cpu.jl` | Minimal CPU training demo | +| `examples/end_to_end_gpu.jl` | Minimal GPU training demo | +| `examples/BatteryStorageOPF/` | PGLib AC-OPF battery-storage example | +| `test/runtests.jl` | Unit and smoke tests | ## Citation diff --git a/examples/BatteryStorageOPF/.gitignore b/examples/BatteryStorageOPF/.gitignore new file mode 100644 index 0000000..23662d2 --- /dev/null +++ b/examples/BatteryStorageOPF/.gitignore @@ -0,0 +1,3 @@ +# Generated outputs from run_case300.jl (manifests + battery tables are +# reproducible from the seeded API, so they are not tracked). +results/ diff --git a/examples/BatteryStorageOPF/Project.toml b/examples/BatteryStorageOPF/Project.toml new file mode 100644 index 0000000..0bcfbe0 --- /dev/null +++ b/examples/BatteryStorageOPF/Project.toml @@ -0,0 +1,50 @@ +[compat] +CUDA = "6" +ChainRulesCore = "1.26" +DecisionRulesExa = "0.1" +ExaModels = "0.11" +Flux = "0.16" +Ipopt = "1.15" +JSON = "1.6" +JuMP = "1.31" +MadNLP = "0.10" +MadNLPGPU = "0.10" +NLPModels = "0.21" +NNlib = "0.9" +PGLib = "0.2.2" +PowerModels = "0.21" +StableRNGs = "1" +Zygote = "0.7" +julia = "1.10, 1.11, 1.12" + +[deps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +DecisionRulesExa = "7c3e91a4-d8f2-4b6a-9e15-a2c4f7b80d53" +ExaModels = "1037b233-b668-4ce9-9b63-f9f681f55dd2" +Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" +Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9" +JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" +JuMP = "4076af6c-e467-56ae-b986-b466b2749572" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" +MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" +NLPModels = "a4795742-8479-5a88-8948-cc11e1c8c1a6" +NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" +PGLib = "07a8691f-3d11-4330-951b-3c50f98338be" +PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce" +Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[extras] +CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" + +[sources.DecisionRulesExa] +path = "../.." diff --git a/examples/BatteryStorageOPF/README.md b/examples/BatteryStorageOPF/README.md new file mode 100644 index 0000000..4146f2b --- /dev/null +++ b/examples/BatteryStorageOPF/README.md @@ -0,0 +1,485 @@ +# Battery-Storage AC-OPF with TS-DDR + +This example builds a **stochastic AC optimal-power-flow (AC-OPF) problem with +batteries** on top of any [PGLib-OPF](https://github.com/power-grid-lib/pglib-opf) +benchmark network, entirely in [ExaModels](https://github.com/exanauts/ExaModels.jl), +and trains a **decision-rule policy (TS-DDR)** on the *true* AC-polar model with +[DecisionRulesExa](../../). It solves on CPU with [MadNLP](https://github.com/MadNLP/MadNLP.jl) +and on GPU with MadNLP + CUDSS. + +If you have never used Julia or power systems before: you can change a few numbers +(the case name, the number of batteries, a random seed), run one command, and get +a trained battery policy with a manifest and trajectory that let anyone reproduce +it exactly. + +* **Phase 1** — deterministic AC-polar foundation + reproducible PGLib battery-case + generator + manifest (no uncertainty, no training). +* **Phase 2** *(this document also covers)* — seeded demand uncertainty, a paired + scenario protocol, a battery-SoC **reachable policy**, a target-constrained + ExaModels problem (strict / soft), CPU/GPU training, rollout evaluation, + checkpointing, and compact trajectories. + +There is **no claim here that TS-DDR beats SDDP** — that comparison is not part of +this example. + +--- + +## 1. The battery model and units + +Every quantity is **per unit (pu)** on the network's `baseMVA` power base. Energy +is in **pu·hours (pu·h)** and time in **hours**. All unit conversion happens in a +single data layer (`src/network_data.jl`); the model code never re-scales. + +For each battery `b` and stage `t` (stage length `Δt` hours) the model uses +continuous, non-negative charge/discharge powers and a linear state of charge +`e` — the **battery SoC** (*not* the SOC-WR power-flow relaxation used by SDDP in +a later phase): + +$$ +0 \le p^{ch}_{t,b} \le \bar p^{ch}_b, \qquad +0 \le p^{dis}_{t,b} \le \bar p^{dis}_b, +$$ + +$$ +e_{t+1,b} = (1-\sigma_b\,\Delta t)\,e_{t,b} + + \eta^{ch}_b\,\Delta t\,p^{ch}_{t,b} + - \frac{\Delta t}{\eta^{dis}_b}\,p^{dis}_{t,b}, +\qquad +\underline e_b \le e_{t,b} \le \bar e_b, +$$ + +$$ +p^{bat}_{t,b} = p^{dis}_{t,b} - p^{ch}_{t,b}. +$$ + +`p_bat` enters the **active** power balance at the battery's bus; batteries run at +**unity power factor** (no reactive injection). A strictly-positive throughput cost +`c_cycle·(p_ch + p_dis)` keeps the formulation continuous and removes any incentive +for simultaneous charge/discharge (audited by the tests). The network model is the +standard **AC-polar OPF** (voltages, the four branch-end flow equations with taps, +shifts and charging shunts, angle-difference limits, apparent-power limits at both +ends) with **hard** active and reactive balance and **no reactive slack**. These +equations live in ONE shared implementation (`src/acp_core.jl`) used by both the +deterministic Phase-1 builder and the Phase-2 operational builder. + +The Phase-1 deterministic foundation has **no active recourse** — it is the +base-ACP parity artifact (hard balance). The Phase-2 **operational** model adds +the two-sided active recourse described in §3. + +**Defaults & units** (all configurable in `make_battery_case`): power/energy in pu +/ pu·h on `baseMVA`; `duration_hours = 4.0` (a "4-hour battery"); `initial_soc = +0.5`; `charge_efficiency = discharge_efficiency = 0.95`; `fleet_power_fraction = +0.25` (fleet power = 0.25·Σload, split equally, never derived from generator +prices); `self_discharge_rate = 0.0` per hour; `cycle_cost_per_mwh = 2.0` USD/MWh. +`stage_hours` (Δt) defaults to `1.0` h at model-build time. + +--- + +## 2. Demand uncertainty and what the policy observes + +The only stochastic driver is **demand** (no renewables, no battery inflow). The +process is a **pure seeded function** (`src/demand_process.jl`): + +* a deterministic hourly **base shape** `base_shape[t]` (time-of-day, mean 1); +* a small finite set of **joint atoms**, each an explicit + `(L, R_1, …, R_R)` — a system-wide load factor `L` and per-region factors — with + **explicit probabilities**; +* buses are split into `R` **regions** by a deterministic, price-free topology rule + (nearest graph-anchor by hop distance, `assign_regions`), so scarcity can move + spatially between atoms. + +At stage `t`, given the atom `a_t`, the realized per-stage uncertainty vector is + +``` +w_t = [ s_t , r_{1,t} , … , r_{R,t} ], s_t = base_shape[t]·L^{a_t}, r_{k,t} = R_k^{a_t} +``` + +(length `nw = 1 + R`). The realized demand at bus `p` is +`pd_p·s_t·r_{region(p),t}` and `qd_p·s_t·r_{region(p),t}` — active and reactive +scaled by the **same** factor, so every bus keeps its **power factor**. The policy +**observes `w_t` and the current battery SoC**, and chooses the next SoC target; it +**never sees future atoms**. Atoms are drawn i.i.d. across stages (a small, +SDDP-friendly support). Default: `nregion = 3`, `period = 24`, `1 + nregion` +atoms (a calm atom plus one per-region scarcity atom). Training and evaluation +use **distinct declared seeds** (`train_seed`, `eval_seed`). + +### Demand calibration presets + +Demand amplitude is a **case-design** parameter: too heavy a process makes the +realized demand unservable and forces nonzero active recourse (deficit `d⁺` or +surplus `d⁻`), which invalidates a scientific run. `DEMAND_PRESETS` is a fixed +ladder, strongest first: + +| preset | `base_amplitude` | high system factor | scarce regional | off-region | +|:--|:--|:--|:--|:--| +| `:D0` | 0.05 | 1.03 | 1.05 | 0.99 | +| `:D1` | 0.04 | 1.02 | 1.04 | 0.99 | +| `:D2` | 0.03 | 1.01 | 1.03 | 0.995 | + +The public default (`DEFAULT_DEMAND_PRESET = :D2`) is a **conservative tutorial +default**, chosen so the documented one-command examples run — **not** a +case300 gate-passing benchmark. `:D2` passes the current case14 +zero-active-recourse checks but **does not** pass them on case300, and **the +current case300 battery placement/configuration is not yet a scientifically +accepted candidate** (screening it is Phase-3 work). Every scientific experiment +must record its preset explicitly and independently pass the zero-active-recourse +gates for its own case — **both** directions (`d⁺` and `d⁻`) numerically zero +within tolerance over the fixed stored evaluation paths, with every declared atom +checked individually. Rejected (stronger) presets remain available as named +experimental presets via `make_load_process(case; preset = :D0)`. +`demand_multiplier_summary(process)` reports the peak system and per-bus demand +multipliers. + +**Targetless feasibility diagnostic.** `build_targetless_diagnostic_de` builds +the full-horizon ACP with batteries freely optimized and the two-sided active +recourse as the only valve. It is **genuinely targetless**: the model contains +no target variables, constraints, slacks, or penalties (not a zero-weight target +— no target data at all). Interpretation is deliberately asymmetric: a positive +active recourse at the returned *local* solution does **not** prove the demand is +physically infeasible (ACP is nonconvex; it is a local optimum); a successful +**recourse-disabled** solve (`allow_active_recourse = false`) proves a +zero-recourse feasible point *was found*; a **failed** recourse-disabled solve is +inconclusive and proves nothing. This separates demand-design signals from policy +failures without over-claiming. + +### Paired protocol + +`scenario_index_matrix(process, horizon, paths; seed)` produces a **stage-major** +`horizon × paths` matrix of atom indices. `write_scenario_protocol` serializes it +(with all process parameters, seeds, units, and hashes); `materialize_scenario` +turns one column into `w_flat`. **Every method shares the same stored index +matrix** — no method re-draws a differently-shaped random array. +`reconstruct_scenario_protocol` rebuilds and **verifies both** the process hash and +the index-matrix hash. + +```julia +train_mat = scenario_index_matrix(process, T, 64; seed = process.train_seed) +eval_mat = scenario_index_matrix(process, T, 64; seed = process.eval_seed) # fixed held-out +write_scenario_protocol("results/eval_protocol.json", process, eval_mat; kind="eval", seed=process.eval_seed) +``` + +--- + +### Artifacts and exact reconstruction + +`run_tiny_training.jl` writes the artifact set that defines an experiment, and +`evaluate_checkpoint.jl` reconstructs from those artifacts **alone** — it never +regenerates the demand process from defaults or environment variables: + +| artifact | contents | +|:--|:--| +| `stochastic_manifest_.json` | the case, the **exact** demand process (base-shape vector, `region_of_bus`, anchors, ordered atoms + probabilities, period, seeds, preset), horizons, stage duration, target mode, active-recourse price, penalty coefficients, policy architecture + activation, and five hashes | +| `train_protocol_.json` / `eval_protocol_.json` | the stage-major scenario-index matrices | +| `checkpoint_.jls` | policy state + architecture + case/source/process hashes | +| `trajectory_.json` | compact per-stage records | + +Five **distinct** hashes are stored, none reused for another role: the +load-process content hash, the train and eval **index-matrix** hashes, and the +train and eval **protocol-file** (exact-bytes) hashes. +`reconstruct_stochastic_manifest` rebuilds the process field-for-field and +verifies its hash; `verify_protocol_file` checks the file bytes. Tampering with +the base shape, an atom, the scenario-index order, or a protocol file is +detected (see `test/test_artifacts.jl`, which also runs a **fresh-Julia-process +round trip** reproducing indices, statuses, costs, and active recourse). + +## 3. Two-sided active recourse — the complete-recourse slack (always present) + +The physical model carries a **two-sided active nodal recourse** — always present, +never a mode, and *not* a target-tracking device. At **every** bus there are two +nonnegative, **UNBOUNDED** variables entering the **active** balance: + +``` +d⁺[t,i] ≥ 0 (active deficit / injection — covers an active shortfall) +d⁻[t,i] ≥ 0 (active surplus / absorption — absorbs an active excess) +active balance: p_d − d⁺ + d⁻ + gs·vm² − Σpg − Σ(p_dis−p_ch) + Σp_fr + Σp_to = 0 +``` + +This is the classical multistage recourse device that gives **relatively +complete recourse**: because `d⁺` can inject arbitrary local power and `d⁻` can +absorb it, the stage subproblem is **feasible for every incoming SoC and every +dynamically reachable battery target**, in both the charging and discharging +directions. Concretely: a target that forces a battery to **charge** at a +network-constrained bus is served by local `d⁺`; a target that forces it to +**discharge** into a bus whose outgoing branches are saturated is absorbed by +local `d⁻`. They are an **artificial active-balance recourse, not curtailed +customer load**: `d⁺` may exceed local demand and may be positive at a bus with +`p_d = 0` — they exist at buses with no load too, which is exactly what makes the +recourse complete and why they are never reported as a per-load fraction. + +The recourse is **active-only**: it does **not** touch reactive power, so reactive +KCL stays a **hard equality** with no reactive slack. Batteries are unity-power- +factor, so the battery target only moves active injection; reactive feasibility is +a property of the base network and the (feasible) demand process, independent of +the target. + +Both directions are priced at `active_recourse_cost_per_mwh = 10 000` USD/MWh +(default), stage cost `cost · baseMVA · Δt · Σ(d⁺ + d⁻)`, **included in physical +operating cost**. They are a safety valve: **an accepted run leaves both +directions at ~0 within tolerance** — a nonzero deficit/surplus on an otherwise +sensible target is a case-design signal (the target is not network-deliverable), +not a solver failure. The solve **always succeeds**; the cost tells you whether +the target was deliverable. A scientific candidate path requires **both** `d⁺` +and `d⁻` numerically zero within tolerance. + +## 4. Strict vs soft target mode + +The interstage state is battery SoC. The policy outputs a **target next SoC** +`ê_{t+1}`; the operational problem (`build_battery_tsddr_de`) ties it to the +realized SoC in one of two **clearly separated** modes (orthogonal to active recourse): + +* **strict** (default, **PRIMARY**) — `ê_{t+1,b} − e_{t+1,b} = 0`; **no target + slack, no target penalty**. This is the production operational and training + target: its multiplier is an economic shadow price uncontaminated by a penalty. + Backed by the two-sided active recourse (§3), strict has **complete recourse** + — for every supported PGLib case and every dynamically reachable target + (including the exact reachable endpoints), the strict stage NLP solves and + reproduces the target to `≤ 1e-5`. +* **soft** (diagnostic/fallback only) — `ê_{t+1,b} − e_{t+1,b} − δ⁺ + δ⁻ = 0` with + `δ⁺,δ⁻ ≥ 0` and a documented **training-only** penalty + `ρ1·Σ(δ⁺+δ⁻) + (ρ2/2)·Σ((δ⁺)²+(δ⁻)²)`. + +The target penalty is **never** counted as physical operating cost, and target +violations are reported separately. Target slacks (`δ±`) and the active +recourse (`d±`) are distinct mechanisms and never share a name. + +The default is **strict** — it is the primary mode. Soft is retained only as a +diagnostic/fallback. + +### Strict start sequence (deterministic warm-start helper) + +Strict feasibility is guaranteed by the two-sided active recourse (§3), not by the +start. This helper only makes the accepted solve **cheaper** (fewer iterations): +`solve_stage_with_starts` tries a **fixed** deterministic sequence and accepts the +**first** solver-accepted result (never the cheapest): + +1. flat start (`vm = 1`); +2. **target-consistent battery start** — with `Δe = ê − (1−σΔt)·e_prev`, + `Δe ≥ 0 ⇒ p_charge = Δe/(η_ch·Δt), p_discharge = 0`; otherwise + `p_charge = 0, p_discharge = −Δe·η_dis/Δt`, clipped to the power bounds; +3. seed from the corresponding solved **targetless-diagnostic** ACP point; +4. the **previous stage's** accepted solution. + +Only the starting point varies — no tolerance change, no iteration-limit +manipulation, no equation/bound/generator/network change, no dropped constraint. +Every attempted start and its solver status is logged. + +### Reachable target policy + +Every target lies in the one-stage physical reachability interval + +``` +ℓ = max(e_min, (1−σΔt)·e_t − (Δt/η_dis)·p̄_dis) +u = min(e_max, (1−σΔt)·e_t + η_ch·Δt·p̄_ch) +``` + +(battery dynamics only — **not** a network-feasibility proof), via +`ê = ℓ + (u−ℓ)·y`. The canonical default activation is + +``` +stretchedsigmoid(z) = clamp((sigmoid(z) − 0.03)/0.94, 0, 1 − 1e-3) +``` + +which attains the **lower** edge exactly at finite weights while keeping a small +margin **below** the exact upper edge: an exact `y = 1` (store-max) target under a +strict equality drives the stage NLP onto a measure-zero set that interior-point +solvers cannot converge into. `hardsigmoidsafe(z) = clamp(0.5 + 0.5z, 0, 1 − 1e-3)` +is a documented option with the same safe margin. Reachability bounds are physical +projection **data**: gradients stop through `ℓ` and `u` and flow only through the +normalized output. Recurrent state is reset at every scenario boundary. + +--- + +## 5. Physical cost vs training-only penalty, reporting vs look-ahead + +`decompose_costs(prob, result)` reports, recomputed independently from the primal +solution: + +| quantity | meaning | +|---|---| +| `total_solver_objective` | the raw solver objective | +| `generator_cost` | `Δt·Σ(c2·pg² + c1·pg + c0)` (original PGLib costs, `c0` scaled too) | +| `battery_throughput_cost` | `baseMVA·Δt·Σ c_cycle(p_ch + p_dis)` | +| `active_recourse_cost` | `cost·baseMVA·Δt·Σ(d⁺ + d⁻)` — **physical** | +| `active_deficit_pu`, `active_surplus_pu`, `active_deficit_energy_mwh`, `active_surplus_energy_mwh`, `total_active_recourse_energy_mwh`, `max_active_deficit_pu`, `max_active_surplus_pu` | recourse audit, by direction (no per-load fraction) | +| **`physical_operating_cost`** | generator + throughput + **active-recourse** cost | +| `target_penalty`, `target_violation` | soft mode only — **training only, never physical** | +| `reporting_physical_cost` | physical cost over stages `1:reporting_horizon` | +| `lookahead_physical_cost` | physical cost over the look-ahead buffer | + +and `total_solver_objective ≈ physical_operating_cost + target_penalty` (checked in +the tests). **Improvement is always judged on reporting-window physical cost, +never the total objective.** + +The horizon is a **reporting horizon** followed by a **look-ahead buffer**: +`T = reporting_horizon + lookahead`. Both are recorded in the manifest and the +checkpoint and are identical for every method that uses the case. Only the reported +(reporting-horizon) physical cost is used to compare policies. + +--- + +## 6. Environment setup + +This example has its own environment and uses the parent `DecisionRulesExa` package +through a relative `[sources]` path in `Project.toml`. + +```bash +module load julia # Julia 1.12.x (matches the parent package) +cd examples/BatteryStorageOPF +julia --pkgimages=no --project=. setup_env.jl +``` + +> **Why `--pkgimages=no`?** On this cluster the system Julia cannot build the +> native precompile image for the `Pkg` stdlib (a MadNLP dependency). Disabling +> native package images sidesteps that; the model and solver are unaffected. **Use +> `--pkgimages=no` on every command below.** The first run of each command JITs +> from source and may take a few minutes. + +--- + +## 7. Run tiny CPU training (one command) + +```bash +julia --pkgimages=no --project=. run_tiny_training.jl +``` + +This builds a battery case + demand process, writes the paired **train** and +**eval** protocols and the stochastic manifest to `results/`, evaluates a freshly +initialized policy on the fixed held-out eval scenarios (the "before" physical +cost), trains the reachable policy on the full-horizon target-constrained DE with +deterministic scenario replay, re-evaluates (the "after" physical cost), and saves +a checkpoint and a compact trajectory. It prints the initial and final held-out +**physical operating cost** and the improvement. + +Override defaults via environment variables, e.g. a different case / battery count +/ seed / horizon / mode: + +```bash +BAT_CASE=case14_ieee BAT_NBAT=3 BAT_SEED=1 BAT_NREGION=3 \ +BAT_REPORT=4 BAT_LOOKAHEAD=1 BAT_MODE=strict BAT_BATCHES=40 \ + julia --pkgimages=no --project=. run_tiny_training.jl +``` + +## 8. Evaluate a saved checkpoint (no retraining) + +```bash +BAT_CKPT=results/checkpoint_case14_ieee.jls \ + julia --pkgimages=no --project=. evaluate_checkpoint.jl +``` + +Reloads the policy (verifying the case-manifest, MATPOWER-source, and load-process +hashes), runs the non-anticipative stage-wise rollout on the **fixed** eval +scenarios, prints the held-out physical cost, and writes a trajectory. + +## 9. Run on a GPU (when available) + +```bash +# On a GPU node (see the GPU sbatch recipe), with a functional CUDA device: +julia --pkgimages=no --project=. run_gpu_training.jl +``` + +Builds the ExaModels model with a `CUDABackend()` and solves with MadNLP's CUDSS +GPU linear solver via `MadNLPGPU`. The CPU and GPU builders represent the **same** +mathematical problem (a structural-parity test asserts equal variable/constraint +counts and target-multiplier slice). + +--- + +## 10. Change the case, battery count, and seed (Julia API) + +```julia +include("src/BatteryStorageOPF.jl"); using .BatteryStorageOPF +using DecisionRulesExa, Flux, MadNLP, Random + +case = make_battery_case("case300_ieee"; number_of_batteries = 20, seed = 20260722) +process = make_load_process(case; nregion = 3) +T = 5; report = 4; lookahead = 1 +de = build_battery_tsddr_de(case, process; reporting_horizon = report, + lookahead = lookahead, mode = :strict, stage_hours = 1.0) +stage = build_battery_stage_problem(case, process; mode = :strict, stage_hours = 1.0) +Random.seed!(1) +policy = battery_reachable_policy(case, process; dt = 1.0, layers = [64, 64], combiner_layers = [64]) + +train_mat = scenario_index_matrix(process, T, 64; seed = process.train_seed) +eval_mat = scenario_index_matrix(process, T, 64; seed = process.eval_seed) +before = evaluate_paired(policy, stage, process, eval_mat; reporting_horizon = report) +train_battery_tsddr(policy, de, process, train_mat; num_batches = 100, num_train_per_batch = 16) +after = evaluate_paired(policy, stage, process, eval_mat; reporting_horizon = report) +``` + +`available_pglib_cases()` lists every benchmark; names resolve leniently. The same +API works for **any** PGLib case — only the name changes. + +--- + +## 11. File inventory + +``` +examples/BatteryStorageOPF/ +├── Project.toml # example env (DecisionRulesExa via [sources]) +├── setup_env.jl # resolve + instantiate + precompile +├── run_case300.jl # Phase-1 deterministic CPU smoke runner +├── run_tiny_training.jl # Phase-2 tiny CPU TS-DDR training + eval +├── evaluate_checkpoint.jl # Phase-2 checkpoint evaluation (no retraining) +├── run_gpu_training.jl # Phase-2 GPU training smoke +├── src/ +│ ├── network_data.jl # typed per-unit PGLib parsing + stable id maps +│ ├── battery_data.jl # battery structs + make_battery_case + validation +│ ├── manifest.jl # Phase-1 manifest write / hash / reconstruct +│ ├── acp_core.jl # SHARED AC-polar equations (single source of truth) +│ ├── battery_opf_exa.jl # Phase-1 deterministic model (uses acp_core) +│ ├── reference_powermodels.jl # PowerModels/Ipopt ACP parity check +│ ├── demand_process.jl # seeded finite-support demand + paired protocol +│ ├── battery_tsddr.jl # target-constrained DE (strict/soft) + cost decomposition +│ ├── battery_policy.jl # battery-SoC reachable policy (edge-reaching map) +│ ├── battery_training.jl # train/rollout entrypoints + checkpoint + trajectory +│ └── stochastic_manifest.jl# Phase-2 manifest (case + process + horizons + seeds + hashes) +└── test/ + ├── runtests.jl # Phase-1 suite + ├── runtests_phase2.jl # Phase-2 suite (process/policy/DE/gradients/checkpoint/parity) + ├── test_artifacts.jl # exact reconstruction, 5 hashes, tampering, round trip + └── test_e2e_training.jl # tiny fixed-seed end-to-end training test +``` + +--- + +## 12. Tests + +```bash +julia --pkgimages=no --project=. test/runtests.jl # Phase 1 +julia --pkgimages=no --project=. test/runtests_phase2.jl # Phase 2 +julia --pkgimages=no --project=. test/test_artifacts.jl # artifacts + round trip +julia --pkgimages=no --project=. test/test_e2e_training.jl # tiny end-to-end training +``` + +The Phase-2 suite checks exact demand replay, distinct train/eval protocols, +probability/input validation, deterministic region assignment, power-factor +preservation, protocol serialization/hashing/reconstruction, stochastic-manifest +reconstruction, reachability bounds and target containment, recurrent reset and +determinism, finite/finite-difference-checked policy gradients, the target +multiplier's finite-difference sign and magnitude, strict target equality, the +soft physical/penalty decomposition, zero active recourse in both directions and battery-balance residuals, +absence of simultaneous charge/discharge (soft mode), exact checkpoint reload, and +CPU/GPU structural parity (GPU checks are gated on `CUDA.functional()` and reported +as skipped on a CPU node). + +--- + +## 13. Attribution and reproducibility + +**Network data & license.** Networks come from the *Power Grid Library for +Benchmarking AC Optimal Power Flow Algorithms* (PGLib-OPF), release **23.07**, +distributed via `PGLib.jl`, licensed **CC BY 4.0** +(). Please cite +S. Babaeinejadsarookolaee *et al.*, arXiv:1908.02788. The manifest records the +exact MATPOWER filename, its SHA-256, the upstream release, the license, and the +package/Julia versions. + +**Reproducibility.** Battery placement and scenario generation are seeded +(`StableRNGs`, stable across Julia versions). `make_battery_case` and +`make_load_process` record every parameter, seed, and hash; `write_stochastic_manifest` +serializes the case content hash (incl. the MATPOWER source hash), the process +hash, the horizon / reporting-horizon / look-ahead treatment, the target mode, the +train/eval seeds, and the protocol hashes. `reconstruct_stochastic_manifest` +rebuilds and verifies all of them. Checkpoints identify the Flux state, +architecture, target mode, horizons, seeds, and the case/source/process hashes, and +reload to reproduce policy outputs **exactly** on a fixed input. diff --git a/examples/BatteryStorageOPF/evaluate_checkpoint.jl b/examples/BatteryStorageOPF/evaluate_checkpoint.jl new file mode 100644 index 0000000..4cbbf58 --- /dev/null +++ b/examples/BatteryStorageOPF/evaluate_checkpoint.jl @@ -0,0 +1,90 @@ +#!/usr/bin/env julia +# evaluate_checkpoint.jl +# +# Evaluate a saved TS-DDR checkpoint using ONLY the saved artifacts. It does NOT +# regenerate the demand process from defaults or environment variables: +# +# 1. reconstruct the case + demand process FIELD-FOR-FIELD from the stochastic +# manifest and verify the load-process content hash; +# 2. verify the evaluation protocol FILE bytes against the manifest hash, then +# read the stored evaluation scenario-index matrix and verify ITS hash; +# 3. reload the policy from the checkpoint (case/source/process hashes checked); +# 4. run the non-anticipative stage-wise rollout on exactly those scenarios. +# +# Run (from examples/BatteryStorageOPF): +# module load julia +# BAT_MANIFEST=results/stochastic_manifest_case14_ieee.json \ +# BAT_EVAL_PROTOCOL=results/eval_protocol_case14_ieee.json \ +# BAT_CKPT=results/checkpoint_case14_ieee.jls \ +# julia --pkgimages=no --project=. evaluate_checkpoint.jl + +include(joinpath(@__DIR__, "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using MadNLP +using Flux +using Printf + +const RES = joinpath(@__DIR__, "results") +const CASE = get(ENV, "BAT_CASE", "case14_ieee") +const MANIFEST = get(ENV, "BAT_MANIFEST", joinpath(RES, "stochastic_manifest_$(CASE).json")) +const EVALPROT = get(ENV, "BAT_EVAL_PROTOCOL", joinpath(RES, "eval_protocol_$(CASE).json")) +const CKPT = get(ENV, "BAT_CKPT", joinpath(RES, "checkpoint_$(CASE).jls")) +const OUTDIR = get(ENV, "BAT_OUTDIR", RES) + +function main() + for f in (MANIFEST, EVALPROT, CKPT) + isfile(f) || error("missing artifact: $f (run run_tiny_training.jl first)") + end + + # 1. Experiment definition, reconstructed exactly and hash-verified. + case, process, meta = reconstruct_stochastic_manifest(MANIFEST) + @printf("Reconstructed experiment from %s\n", MANIFEST) + @printf(" load-process hash verified : %s\n", process_hash(process)) + @printf(" preset=%s target_mode=%s report=%d lookahead=%d dt=%.3g\n", + process.preset, meta.mode, meta.reporting_horizon, meta.lookahead, meta.stage_hours) + @printf(" active recourse=%.1f USD/MWh rho1=%.3g rho2=%.3g activation=%s\n", + meta.active_recourse_cost_per_mwh, meta.rho1, meta.rho2, meta.activation) + + # 2. Evaluation protocol: verify FILE bytes, then the index-matrix content. + file_hash = verify_protocol_file(EVALPROT, meta.eval_protocol_file_sha256) + @printf(" eval protocol file hash : %s (verified)\n", file_hash) + proto_process, eval_mat, pmeta = reconstruct_scenario_protocol(EVALPROT) + process_hash(proto_process) == process_hash(process) || + error("evaluation protocol describes a different demand process than the manifest") + if meta.eval_index_matrix_hash !== nothing + index_matrix_hash(eval_mat) == String(meta.eval_index_matrix_hash) || + error("evaluation index-matrix hash mismatch vs manifest") + end + @printf(" eval index-matrix hash : %s (verified) paths=%d horizon=%d\n", + index_matrix_hash(eval_mat), pmeta.paths, pmeta.horizon) + + # 3. Policy from the checkpoint. + policy, ck = load_checkpoint(CKPT, case, process) + @printf(" checkpoint target_mode=%s activation=%s\n", + ck["architecture"]["target_mode"], ck["architecture"]["activation"]) + + # 4. Rollout on exactly the stored scenarios. + stage = build_battery_stage_problem(case, process; mode = meta.mode, + rho1 = meta.rho1, rho2 = meta.rho2, + active_recourse_cost_per_mwh = meta.active_recourse_cost_per_mwh, + stage_hours = meta.stage_hours) + ev = evaluate_paired(policy, stage, process, eval_mat; + reporting_horizon = meta.reporting_horizon, keep_trajectories = true) + @printf("\n mean held-out physical cost = %.6f (n_ok=%d, n_failed=%d)\n", + ev.mean_reporting_physical_cost, ev.n_ok, ev.n_failed) + @printf(" active deficit = %.6g MWh active surplus = %.6g MWh (max deficit %.6g / surplus %.6g pu)\n", + ev.total_active_deficit_energy_mwh, ev.total_active_surplus_energy_mwh, + ev.max_active_deficit_pu, ev.max_active_surplus_pu) + for (p, c) in enumerate(ev.reporting_physical_costs) + @printf(" path %-3d reporting physical cost = %.6f\n", p, c) + end + + traj = joinpath(OUTDIR, "eval_trajectory_$(CASE).json") + write_trajectory(traj, ev.trajectories; + meta = Dict("case" => CASE, "target_mode" => String(meta.mode), + "reconstructed_from" => MANIFEST)) + @printf(" wrote %s\n", traj) + return nothing +end + +main() diff --git a/examples/BatteryStorageOPF/run_case300.jl b/examples/BatteryStorageOPF/run_case300.jl new file mode 100644 index 0000000..6bd1f29 --- /dev/null +++ b/examples/BatteryStorageOPF/run_case300.jl @@ -0,0 +1,92 @@ +#!/usr/bin/env julia +# run_case300.jl +# +# CPU smoke runner for the deterministic battery-storage AC-OPF foundation. +# Builds a reproducible battery case from a PGLib benchmark, writes its manifest +# + human-readable battery file, builds a short-horizon ExaModels AC-polar +# deterministic equivalent, solves it on the CPU with MadNLP, and reports the +# solver status, objective, residuals, chosen battery buses, and manifest hash. +# +# Run (from examples/BatteryStorageOPF): +# module load julia +# julia --pkgimages=no --project=. run_case300.jl +# +# Change the case / battery count / seed / horizon via environment variables: +# BAT_CASE=case14_ieee BAT_NBAT=3 BAT_SEED=1 BAT_HORIZON=4 \ +# julia --pkgimages=no --project=. run_case300.jl + +include(joinpath(@__DIR__, "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using MadNLP +using Printf + +const CASE = get(ENV, "BAT_CASE", "case300_ieee") +const NBAT = parse(Int, get(ENV, "BAT_NBAT", "20")) +const SEED = parse(Int, get(ENV, "BAT_SEED", "20260722")) +const HORIZON = parse(Int, get(ENV, "BAT_HORIZON", "4")) +const OUTDIR = get(ENV, "BAT_OUTDIR", joinpath(@__DIR__, "results")) + +function main() + mkpath(OUTDIR) + @printf("Building battery case \"%s\" (batteries=%d, seed=%d)\n", CASE, NBAT, SEED) + case = make_battery_case(CASE; number_of_batteries = NBAT, seed = SEED) + nd = case.network + + @printf(" network: %d buses, %d gens, %d branches, %d loads, baseMVA=%.1f\n", + nbus(nd), ngen(nd), nbranch(nd), nload(nd), nd.baseMVA) + @printf(" total load = %.4f pu ; eligible load buses = %d\n", + case.total_load_pu, length(case.eligible_bus_ids)) + @printf(" battery buses (stable order) = %s\n", string(case.selected_bus_ids)) + if !isempty(case.batteries) + b = case.batteries[1] + @printf(" per-battery: p̄=%.4f pu, e_max=%.4f pu·h, e_init=%.4f pu·h, η=%.2f/%.2f\n", + b.p_charge_max, b.e_max, b.e_init, b.eta_ch, b.eta_dis) + end + + hash = manifest_hash(case) + man_path = joinpath(OUTDIR, "manifest_$(CASE)_seed$(SEED).json") + bat_path = joinpath(OUTDIR, "batteries_$(CASE)_seed$(SEED).csv") + write_battery_file(case, bat_path) + write_manifest(case, man_path; extra_files = [bat_path]) + @printf(" manifest hash = %s\n", hash) + @printf(" wrote %s\n %s\n", man_path, bat_path) + + # Short horizon with a gentle time-of-day load shape (±3%) so the batteries + # cycle (peak/off-peak arbitrage) — exercises the state equation — while + # staying within the base case's feasible region. + profile = HORIZON == 4 ? [1.00, 1.03, 1.00, 0.97] : + [1.0 + 0.03 * sinpi(2 * (t - 1) / HORIZON) for t in 1:HORIZON] + @printf("\nBuilding ExaModels AC-polar DE (T=%d, Δt=1.0 h) ...\n", HORIZON) + prob = build_battery_de(case, HORIZON; stage_hours = 1.0, demand_profile = profile) + + @printf("Solving on CPU with MadNLP ...\n") + t0 = time() + result = solve_de!(prob; print_level = MadNLP.ERROR, tol = 1e-6, max_iter = 1000) + dt = time() - t0 + + sol = battery_solution(prob, result) + prim = max_primal_residual(prob, result) + balres = battery_balance_residuals(prob, sol) + simult = simultaneous_charge_discharge_power(sol) + + println("\n── Results ─────────────────────────────────────────────") + @printf("status : %s\n", string(result.status)) + @printf("accepted : %s\n", solve_succeeded(result.status)) + @printf("objective (USD) : %.6f\n", result.objective) + @printf("solve time (s) : %.2f\n", dt) + @printf("max primal residual : %.3e\n", prim) + if !isempty(case.batteries) + @printf("max |battery balance|: %.3e\n", maximum(abs, balres)) + @printf("max simultaneous charge/discharge power : %.3e pu\n", maximum(simult)) + @printf("SoC[:,1] (init) : %s\n", string(round.(sol.soc[:, 1], digits = 4))) + @printf("SoC[:,end] (final) : %s\n", string(round.(sol.soc[:, end], digits = 4))) + @printf("Σ discharge (pu) : %.4f ; Σ charge (pu) : %.4f\n", + sum(sol.p_dis), sum(sol.p_ch)) + end + println("────────────────────────────────────────────────────────") + + solve_succeeded(result.status) || error("smoke solve did not reach an accepted status") + return nothing +end + +main() diff --git a/examples/BatteryStorageOPF/run_gpu_training.jl b/examples/BatteryStorageOPF/run_gpu_training.jl new file mode 100644 index 0000000..94df82f --- /dev/null +++ b/examples/BatteryStorageOPF/run_gpu_training.jl @@ -0,0 +1,69 @@ +#!/usr/bin/env julia +# run_gpu_training.jl +# +# GPU TS-DDR training smoke for the battery-storage example (Phase-2 prompt §5), +# using the existing DecisionRulesExa + MadNLPGPU path: the ExaModels model is +# built with a CUDA backend and MadNLP solves it with the CUDSS GPU linear solver. +# +# Run on a GPU node (see the GPU sbatch recipe in the README): +# module load julia +# julia --pkgimages=no --project=. run_gpu_training.jl +# +# Requires a functional CUDA device; it errors clearly otherwise. Same defaults / +# env-var overrides as run_tiny_training.jl. + +include(joinpath(@__DIR__, "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using DecisionRulesExa +using ExaModels +using CUDA +using MadNLPGPU +using MadNLP +using Flux +using Random +using Printf + +CUDA.functional() || error("CUDA is not functional on this machine; run on a GPU node.") +CUDA.allowscalar(false) + +const CASE = get(ENV, "BAT_CASE", "case300_ieee") +const NBAT = parse(Int, get(ENV, "BAT_NBAT", "20")) +const SEED = parse(Int, get(ENV, "BAT_SEED", "20260722")) +const NREGION = parse(Int, get(ENV, "BAT_NREGION", "3")) +const REPORT = parse(Int, get(ENV, "BAT_REPORT", "4")) +const LOOKAH = parse(Int, get(ENV, "BAT_LOOKAHEAD", "1")) +const MODE = Symbol(get(ENV, "BAT_MODE", "strict")) +const BATCHES = parse(Int, get(ENV, "BAT_BATCHES", "10")) +const NTRAIN = parse(Int, get(ENV, "BAT_NTRAIN", "4")) +const POLSEED = parse(Int, get(ENV, "BAT_POLICY_SEED", "1234")) + +function main() + T = REPORT + LOOKAH + @printf("GPU device: %s\n", CUDA.name(CUDA.device())) + @printf("Case %s: %d batteries, %d regions, mode=%s, T=%d\n", CASE, NBAT, NREGION, MODE, T) + + case = make_battery_case(CASE; number_of_batteries = NBAT, seed = SEED) + process = make_load_process(case; nregion = NREGION, period = max(T, 4)) + train_mat = scenario_index_matrix(process, T, NTRAIN; seed = process.train_seed) + + # Build the target-constrained DE on the GPU (CUDABackend); the policy lives + # on the GPU too so the rollout produces device targets. + de = build_battery_tsddr_de(case, process; reporting_horizon = REPORT, lookahead = LOOKAH, + mode = MODE, stage_hours = 1.0, backend = CUDABackend()) + Random.seed!(POLSEED) + policy = battery_reachable_policy(case, process; dt = 1.0, layers = [64, 64], + combiner_layers = [64]) |> Flux.gpu + x0 = policy_initial_state(case) |> Flux.gpu + + sampler, _ = make_replay_sampler(process, train_mat) + @printf("\nTraining on GPU (%d batches × %d scenarios) ...\n", BATCHES, NTRAIN) + train_tsddr(policy, x0, de, de.p_x0, de.p_target, de.p_w, sampler; + num_batches = BATCHES, num_train_per_batch = NTRAIN, + optimizer = Flux.Adam(1f-3), + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-4, + linear_solver = MadNLPGPU.CUDSSSolver)) + println("GPU training smoke completed.") + return nothing +end + +main() diff --git a/examples/BatteryStorageOPF/run_tiny_training.jl b/examples/BatteryStorageOPF/run_tiny_training.jl new file mode 100644 index 0000000..2df798f --- /dev/null +++ b/examples/BatteryStorageOPF/run_tiny_training.jl @@ -0,0 +1,137 @@ +#!/usr/bin/env julia +# run_tiny_training.jl +# +# Beginner-usable CPU TS-DDR training that PRODUCES THE ARTIFACT SET the rest of +# the workflow consumes: +# +# results/stochastic_manifest_.json — the experiment definition (exact +# demand process, horizons, target mode, VOLL/penalties, policy arch, and +# five distinct hashes) +# results/train_protocol_.json — training scenario-index matrix +# results/eval_protocol_.json — FIXED held-out evaluation matrix +# results/checkpoint_.jls — policy checkpoint +# results/trajectory_.json — compact per-stage trajectory +# +# `evaluate_checkpoint.jl` reconstructs the experiment from these artifacts alone +# — it never regenerates the demand process from defaults or env vars. +# +# Run (from examples/BatteryStorageOPF): +# module load julia +# julia --pkgimages=no --project=. run_tiny_training.jl +# +# Env overrides: BAT_CASE, BAT_NBAT, BAT_SEED, BAT_NREGION, BAT_PRESET, +# BAT_REPORT, BAT_LOOKAHEAD, BAT_MODE (soft|strict), BAT_BATCHES, BAT_NTRAIN, +# BAT_NEVAL, BAT_LR, BAT_POLICY_SEED, BAT_OUTDIR. + +include(joinpath(@__DIR__, "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using MadNLP +using Flux +using Random +using Printf +using SHA + +const CASE = get(ENV, "BAT_CASE", "case14_ieee") +const NBAT = parse(Int, get(ENV, "BAT_NBAT", "3")) +const SEED = parse(Int, get(ENV, "BAT_SEED", "20260722")) +const NREGION = parse(Int, get(ENV, "BAT_NREGION", "3")) +const PRESET = Symbol(get(ENV, "BAT_PRESET", string(DEFAULT_DEMAND_PRESET))) +const REPORT = parse(Int, get(ENV, "BAT_REPORT", "4")) +const LOOKAH = parse(Int, get(ENV, "BAT_LOOKAHEAD", "1")) +const MODE = Symbol(get(ENV, "BAT_MODE", "strict")) +const BATCHES = parse(Int, get(ENV, "BAT_BATCHES", "40")) +const NTRAIN = parse(Int, get(ENV, "BAT_NTRAIN", "8")) +const NEVAL = parse(Int, get(ENV, "BAT_NEVAL", "8")) +const LR = parse(Float64, get(ENV, "BAT_LR", "0.01")) +const OUTDIR = get(ENV, "BAT_OUTDIR", joinpath(@__DIR__, "results")) +const POLSEED = parse(Int, get(ENV, "BAT_POLICY_SEED", "1234")) +const LAYERS = [32, 32] +const COMBINER = [32] + +function main() + mkpath(OUTDIR) + T = REPORT + LOOKAH + @printf("Case %s: %d batteries, %d regions, preset=%s, mode=%s, T=%d (report=%d + look=%d)\n", + CASE, NBAT, NREGION, PRESET, MODE, T, REPORT, LOOKAH) + + case = make_battery_case(CASE; number_of_batteries = NBAT, seed = SEED) + process = make_load_process(case; preset = PRESET, nregion = NREGION, period = max(T, 4)) + mult = demand_multiplier_summary(process) + @printf(" demand multipliers: max system %.4f, max bus %.4f\n", + mult.max_system_multiplier, mult.max_bus_multiplier) + + # Paired protocols: distinct declared seeds; eval is the fixed held-out set. + train_mat = scenario_index_matrix(process, T, NTRAIN; seed = process.train_seed) + eval_mat = scenario_index_matrix(process, T, NEVAL; seed = process.eval_seed) + train_path = joinpath(OUTDIR, "train_protocol_$(CASE).json") + eval_path = joinpath(OUTDIR, "eval_protocol_$(CASE).json") + write_scenario_protocol(train_path, process, train_mat; kind = "train", seed = process.train_seed) + write_scenario_protocol(eval_path, process, eval_mat; kind = "eval", seed = process.eval_seed) + + # Five DISTINCT hashes: process content, the two index matrices, and the two + # protocol FILES (exact bytes). + man_path = joinpath(OUTDIR, "stochastic_manifest_$(CASE).json") + write_stochastic_manifest(man_path, case, process; + reporting_horizon = REPORT, lookahead = LOOKAH, mode = MODE, stage_hours = 1.0, + active_recourse_cost_per_mwh = DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH, rho1 = 0.0, rho2 = 0.0, + activation = string(stretchedsigmoid), safe_upper_margin = 1e-3, + policy_layers = LAYERS, policy_combiner_layers = COMBINER, policy_seed = POLSEED, + train_index_matrix_hash = index_matrix_hash(train_mat), + eval_index_matrix_hash = index_matrix_hash(eval_mat), + train_protocol_file_sha256 = bytes2hex(open(sha256, train_path)), + eval_protocol_file_sha256 = bytes2hex(open(sha256, eval_path)), + train_paths = NTRAIN, eval_paths = NEVAL) + @printf(" artifacts → %s\n process hash = %s\n", OUTDIR, process_hash(process)) + + de = build_battery_tsddr_de(case, process; reporting_horizon = REPORT, + lookahead = LOOKAH, mode = MODE, stage_hours = 1.0) + stage = build_battery_stage_problem(case, process; mode = MODE, stage_hours = 1.0) + + Random.seed!(POLSEED) # reproducible initialization + policy = battery_reachable_policy(case, process; dt = 1.0, + layers = LAYERS, combiner_layers = COMBINER) + + @printf("\nEvaluating INITIAL policy on %d held-out paired scenarios ...\n", NEVAL) + ev0 = evaluate_paired(policy, stage, process, eval_mat; reporting_horizon = REPORT) + @printf(" initial mean held-out physical cost = %.4f (n_ok=%d, n_failed=%d, deficit=%.3g / surplus=%.3g MWh)\n", + ev0.mean_reporting_physical_cost, ev0.n_ok, ev0.n_failed, + ev0.total_active_deficit_energy_mwh, ev0.total_active_surplus_energy_mwh) + + @printf("\nTraining (%d batches × %d scenarios) ...\n", BATCHES, NTRAIN) + tr = train_battery_tsddr(policy, de, process, train_mat; + num_batches = BATCHES, num_train_per_batch = NTRAIN, + optimizer = Flux.Adam(Float32(LR)), + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-6)) + @printf(" solves: n_ok=%d / n_total=%d (failed=%d)\n", tr.n_ok, tr.n_total, tr.n_failed) + isempty(tr.failure_counts) || @printf(" failure_counts = %s\n", string(tr.failure_counts)) + + @printf("\nEvaluating TRAINED policy on the SAME held-out scenarios ...\n") + ev1 = evaluate_paired(policy, stage, process, eval_mat; reporting_horizon = REPORT, + keep_trajectories = true) + @printf(" trained mean held-out physical cost = %.4f (n_ok=%d, n_failed=%d, deficit=%.3g / surplus=%.3g MWh)\n", + ev1.mean_reporting_physical_cost, ev1.n_ok, ev1.n_failed, + ev1.total_active_deficit_energy_mwh, ev1.total_active_surplus_energy_mwh) + + ckpt = joinpath(OUTDIR, "checkpoint_$(CASE).jls") + save_checkpoint(ckpt, policy, de; case = case, process = process, + extra = Dict("initial_mean_physical" => ev0.mean_reporting_physical_cost, + "final_mean_physical" => ev1.mean_reporting_physical_cost, + "manifest_path" => man_path, + "eval_protocol_path" => eval_path)) + traj = joinpath(OUTDIR, "trajectory_$(CASE).json") + write_trajectory(traj, ev1.trajectories; + meta = Dict("case" => CASE, "target_mode" => String(MODE), + "reporting_horizon" => REPORT, "lookahead" => LOOKAH)) + + Δ = ev0.mean_reporting_physical_cost - ev1.mean_reporting_physical_cost + println("\n── Summary ─────────────────────────────────────────────") + @printf("initial held-out physical cost : %.4f\n", ev0.mean_reporting_physical_cost) + @printf("trained held-out physical cost : %.4f\n", ev1.mean_reporting_physical_cost) + @printf("improvement (initial − trained): %.4f (%.2f%%)\n", + Δ, 100Δ / ev0.mean_reporting_physical_cost) + @printf("manifest : %s\ncheckpoint : %s\ntrajectory : %s\n", man_path, ckpt, traj) + println("────────────────────────────────────────────────────────") + return nothing +end + +main() diff --git a/examples/BatteryStorageOPF/setup_env.jl b/examples/BatteryStorageOPF/setup_env.jl new file mode 100644 index 0000000..c0df3fd --- /dev/null +++ b/examples/BatteryStorageOPF/setup_env.jl @@ -0,0 +1,33 @@ +# setup_env.jl +# +# One-shot helper that resolves the BatteryStorageOPF example environment. It +# activates THIS directory and instantiates the dependency set declared in +# Project.toml — including the parent DecisionRulesExa package, brought in through +# the relative `[sources]` path — letting Pkg fetch a mutually compatible set of +# versions and precompile them. Run once (or after deleting Manifest.toml): +# +# module load julia # Julia 1.12.x (matches the parent package) +# julia --pkgimages=no --project=examples/BatteryStorageOPF \ +# examples/BatteryStorageOPF/setup_env.jl +# +# The depot lives in /tmp per project policy; do not redirect it elsewhere. On a +# fresh node-local depot the General registry is installed first. + +import Pkg +Pkg.activate(@__DIR__) + +# A fresh /tmp depot has no registry; install General before resolving. +try + Pkg.Registry.add("General") +catch err + @info "General registry already present or add skipped" err +end + +# Resolve + install + precompile everything declared in Project.toml (deps + +# [sources] DecisionRulesExa). This pulls the TS-DDR stack (Flux, Zygote, CUDA, +# MadNLP/MadNLPGPU) plus the PGLib/PowerModels/ExaModels modeling stack. +Pkg.resolve() +Pkg.instantiate() +Pkg.precompile() + +@info "BatteryStorageOPF environment resolved" project = Base.active_project() diff --git a/examples/BatteryStorageOPF/src/BatteryStorageOPF.jl b/examples/BatteryStorageOPF/src/BatteryStorageOPF.jl new file mode 100644 index 0000000..bb38773 --- /dev/null +++ b/examples/BatteryStorageOPF/src/BatteryStorageOPF.jl @@ -0,0 +1,106 @@ +# BatteryStorageOPF.jl +# +# Self-contained example module for a stochastic battery-storage AC-OPF problem +# built on any PGLib-OPF case, with a full TS-DDR path (policy → target-constrained +# ExaModels projection → envelope-theorem gradient) trained/evaluated on the true +# AC-polar model. +# +# Phase 1: deterministic AC-polar foundation + reproducible PGLib battery-case +# generator + manifest (network_data, battery_data, manifest, +# battery_opf_exa, reference_powermodels). +# Phase 2: seeded finite-support demand process + paired protocol, battery-SoC +# reachable policy, target-constrained DE (strict/soft) with cost +# decomposition and reporting/look-ahead horizons, CPU/GPU training and +# rollout entrypoints, checkpointing, and the stochastic manifest +# (demand_process, battery_tsddr, battery_policy, battery_training, +# stochastic_manifest). +# +# Load it from the example environment (see README): +# +# module load julia +# julia --pkgimages=no --project=. -e 'include("src/BatteryStorageOPF.jl")' +# +# `--pkgimages=no` is required on this cluster (the system Julia cannot build the +# native precompile image for the `Pkg` stdlib, a MadNLP dependency). + +module BatteryStorageOPF + +using TOML + +# Order matters: data layer → case construction → manifest → physical model → +# reference → demand process → target-constrained DE → policy → training → +# stochastic manifest. The Phase-2 files depend on the parent DecisionRulesExa +# and Flux, declared in this example's Project.toml (see [sources]). +include("network_data.jl") +include("battery_data.jl") +include("manifest.jl") +include("acp_core.jl") # SHARED AC-polar equations (single source of truth) +include("battery_opf_exa.jl") +include("reference_powermodels.jl") +include("demand_process.jl") +include("battery_tsddr.jl") +include("battery_policy.jl") +include("battery_training.jl") +include("stochastic_manifest.jl") + +# ── Public API ──────────────────────────────────────────────────────────────── +# Network / case construction +export NetworkData, BusData, GenData, BranchData, LoadData +export nbus, ngen, nbranch, nload +export resolve_pglib_case, available_pglib_cases, load_pglib_network, parse_network +export BatteryData, BatteryCase, nbattery +export make_battery_case, validate_battery_case, eligible_load_bus_ids + +# Phase-1 manifest +export battery_manifest, manifest_hash, canonical_content +export write_manifest, write_battery_file, reconstruct_case + +# Phase-1 ExaModels physical model +export BatteryExaProblem, build_battery_de, solve_de!, solve_succeeded +export set_demand!, set_initial_soc! +export battery_solution, battery_balance_residuals +export simultaneous_charge_discharge_power, max_primal_residual + +# Reference / parity +export reference_ac_opf, exa_base_objective, check_base_acp_parity + +# Phase-2 demand process + paired protocol +export LoadAtom, LoadProcess, make_load_process, assign_regions +export default_base_shape, default_load_atoms +export n_uncertainty, natom +export scenario_index_matrix, materialize_scenario, materialize_all +export process_canonical_content, process_hash, index_matrix_hash +export DEMAND_PRESETS, DEFAULT_DEMAND_PRESET, load_process_from_fields +export demand_multiplier_summary +export write_scenario_protocol, reconstruct_scenario_protocol + +# Shared AC-polar blocks (single source of truth) +export add_acp_variables!, add_battery_variables!, add_acp_network_constraints! +export add_nodal_balance!, add_battery_dynamics! +export add_generator_cost!, add_cycle_cost! +export acp_constraint_count, rated_branch_positions, validate_stage_hours + +# Phase-2 operational (target-constrained) problem + cost decomposition +export BatteryTSDDRProblem, build_battery_tsddr_de, build_battery_stage_problem +export set_tsddr_uncertainty!, set_tsddr_initial_soc!, set_tsddr_targets! +export set_realized_demand!, DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH, ACTIVE_RECOURSE_LB_TOL_PU +export build_targetless_diagnostic_de, is_targetless, variable_offsets, target_consistent_start! +export seed_start_from_solution!, reset_flat_start!, solve_stage_with_starts +export target_multipliers, tsddr_solution, decompose_costs +export tsddr_balance_residuals, tsddr_max_primal_residual + +# Phase-2 reachable policy +export BatteryReachablePolicy, battery_reachable_policy +export battery_reachable_bounds, stretchedsigmoid, hardsigmoidsafe +export BOUNDED_TARGET_ACTIVATIONS, policy_initial_state, load_battery_policy! + +# Phase-2 training / rollout / checkpoint / trajectory +export make_replay_sampler, train_battery_tsddr +export evaluate_battery_policy, evaluate_paired +export battery_checkpoint, save_checkpoint, load_checkpoint, write_trajectory + +# Phase-2 stochastic manifest +export stochastic_manifest, write_stochastic_manifest, reconstruct_stochastic_manifest +export verify_protocol_file + +end # module diff --git a/examples/BatteryStorageOPF/src/acp_core.jl b/examples/BatteryStorageOPF/src/acp_core.jl new file mode 100644 index 0000000..da6f426 --- /dev/null +++ b/examples/BatteryStorageOPF/src/acp_core.jl @@ -0,0 +1,405 @@ +# acp_core.jl +# +# SHARED AC-polar (ACP) construction — the single source of truth for the battery +# AC-OPF equations (canonical spec: docs/src/casestudies/battery_storage_opf.md, +# "Implementation invariants": *one shared ACP constraint implementation underlies +# deterministic, stochastic, training, and evaluation builders*). +# +# Both builders call these blocks: +# * `build_battery_de` (Phase-1 deterministic foundation, no active +# recourse, no targets — accepted base-ACP parity +# artifact); +# * `build_battery_tsddr_de` (Phase-2 operational/stochastic model: adds the +# two-sided absolute active nodal recourse (deficit +# d⁺ + surplus d⁻) and the target equations). +# +# Nothing here changes generator prices, pmin/pmax, qmin/qmax, branch limits, +# voltage limits, admittances, or topology: every quantity is read from the +# accepted `NetworkData`/`BatteryCase` produced by network_data.jl/battery_data.jl. +# +# ── Variable creation order (shared; extra blocks appended by the caller) ────── +# va, vm, pg, qg, p_fr, q_fr, p_to, q_to, p_ch, p_dis, e, +# [active_deficit, active_surplus] (Phase 2 only) +# [slack_pos, slack_neg] (Phase 2 soft target mode only) +# +# ── Constraint creation order (shared) ──────────────────────────────────────── +# 1 reference angle T·nRef +# 2 from-end active flow T·nBranch +# 3 from-end reactive flow T·nBranch +# 4 to-end active flow T·nBranch +# 5 to-end reactive flow T·nBranch +# 6 angle-difference limits T·nBranch +# 7 thermal limit (from end) T·nRated +# 8 thermal limit (to end) T·nRated +# 9 active nodal balance T·nBus +# 10 reactive nodal balance T·nBus (HARD equality: no reactive slack) +# 11 battery initial SoC nBat +# 12 battery energy balance T·nBat +# ... then the caller appends TARGET constraints LAST (Phase 2). +# +# `acp_constraint_count` reproduces that total so the caller can compute the +# contiguous target-multiplier slice exactly. + +using ExaModels + +# Flat stage-major index helpers, (t,i) → (t-1)*n + i. +@inline _bidx(nB, t, b) = (t - 1) * nB + b # bus (t = 1..T) +@inline _gidx(nG, t, g) = (t - 1) * nG + g # generator (t = 1..T) +@inline _bridx(nBR, t, r) = (t - 1) * nBR + r # branch (t = 1..T) +@inline _kidx(nK, t, k) = (t - 1) * nK + k # battery power (t = 1..T) +@inline _eidx(nK, t, k) = (t - 1) * nK + k # battery SoC (t = 1..T+1) + +""" + _ac_branch_coeffs(br, T) -> NamedTuple + +ExaModelsPower/PowerModels-compatible AC-polar branch coefficients `c1..c8`, +retaining the transformer tap and phase shift and both line-charging shunts: + +``` +tr = tap·cos(shift), ti = tap·sin(shift), ttm = tr² + ti² +g = br_r/(br_r²+br_x²), b = −br_x/(br_r²+br_x²) +c1 = (−g·tr − b·ti)/ttm c2 = (−b·tr + g·ti)/ttm +c3 = (−g·tr + b·ti)/ttm c4 = (−b·tr − g·ti)/ttm +c5 = (g + g_fr)/ttm c6 = (b + b_fr)/ttm +c7 = g + g_to c8 = b + b_to +``` +""" +function _ac_branch_coeffs(br::BranchData, ::Type{T}) where {T} + r2x2 = br.br_r^2 + br.br_x^2 + g = r2x2 > 0 ? T(br.br_r / r2x2) : zero(T) + b = r2x2 > 0 ? T(-br.br_x / r2x2) : zero(T) + tap = T(br.tap); sh = T(br.shift) + tr = tap * cos(sh); ti = tap * sin(sh) + ttm = tr^2 + ti^2 + ttm = ttm > 0 ? ttm : one(T) + return ( + c1 = (-g * tr - b * ti) / ttm, + c2 = (-b * tr + g * ti) / ttm, + c3 = (-g * tr + b * ti) / ttm, + c4 = (-b * tr - g * ti) / ttm, + c5 = (g + T(br.g_fr)) / ttm, + c6 = (b + T(br.b_fr)) / ttm, + c7 = g + T(br.g_to), + c8 = b + T(br.b_to), + ) +end + +""" + rated_branch_positions(nd) -> Vector{Int} + +Array positions of branches carrying a finite PGLib apparent-power limit. An +absent limit adds no artificial finite bound (canonical spec, Appendix A). +""" +rated_branch_positions(nd::NetworkData) = + [bp for bp in 1:nbranch(nd) if isfinite(nd.branches[bp].rate_a)] + +""" + acp_constraint_count(nd, case, T) -> Int + +Number of constraints the shared ACP + battery blocks add, in creation order +(see the file header). The Phase-2 builder appends its target constraints after +these, so its contiguous target-multiplier slice starts at this count + 1. +""" +function acp_constraint_count(nd::NetworkData, case::BatteryCase, T::Int) + nBus = nbus(nd); nBranch = nbranch(nd); nBat = length(case.batteries) + nRef = length(nd.ref_bus_positions) + nRated = length(rated_branch_positions(nd)) + return T * nRef + 5 * (T * nBranch) + 2 * (T * nRated) + + 2 * (T * nBus) + nBat + (T * nBat) +end + +# ── Variables ───────────────────────────────────────────────────────────────── + +""" + add_acp_variables!(core, nd, T, ft) -> (core, vars) + +Add the shared AC-polar variables in canonical order and return them as a +NamedTuple `(va, vm, pg, qg, p_fr, q_fr, p_to, q_to)`. + +Bounds come straight from the PGLib data: voltage magnitude limits, generator +active/reactive limits (non-finite reactive limits fall back to ±1e4 pu), and +branch-flow box bounds from `rate_a` (±1e4 pu when unlimited). +""" +function add_acp_variables!(core, nd::NetworkData, T::Int, ft::Type{<:AbstractFloat}) + nBus = nbus(nd); nGen = ngen(nd); nBranch = nbranch(nd) + core, va = ExaModels.add_var(core, T * nBus) + core, vm = ExaModels.add_var(core, T * nBus; + lvar = ft.(repeat([b.vmin for b in nd.buses], T)), + uvar = ft.(repeat([b.vmax for b in nd.buses], T)), + start = ones(ft, T * nBus)) + core, pg = ExaModels.add_var(core, T * nGen; + lvar = ft.(repeat([g.pmin for g in nd.gens], T)), + uvar = ft.(repeat([g.pmax for g in nd.gens], T))) + core, qg = ExaModels.add_var(core, T * nGen; + lvar = ft.(repeat([isfinite(g.qmin) ? g.qmin : -1e4 for g in nd.gens], T)), + uvar = ft.(repeat([isfinite(g.qmax) ? g.qmax : 1e4 for g in nd.gens], T))) + fr_lb = ft.(repeat([isfinite(b.rate_a) ? -b.rate_a : -1e4 for b in nd.branches], T)) + fr_ub = ft.(repeat([isfinite(b.rate_a) ? b.rate_a : 1e4 for b in nd.branches], T)) + core, p_fr = ExaModels.add_var(core, T * nBranch; lvar = fr_lb, uvar = fr_ub) + core, q_fr = ExaModels.add_var(core, T * nBranch; lvar = fr_lb, uvar = fr_ub) + core, p_to = ExaModels.add_var(core, T * nBranch; lvar = fr_lb, uvar = fr_ub) + core, q_to = ExaModels.add_var(core, T * nBranch; lvar = fr_lb, uvar = fr_ub) + return core, (va = va, vm = vm, pg = pg, qg = qg, + p_fr = p_fr, q_fr = q_fr, p_to = p_to, q_to = q_to) +end + +""" + add_battery_variables!(core, case, T, ft) -> (core, bat) + +Add battery charge/discharge power (`0 ≤ p ≤ p̄`, per battery) and the `(T+1)` +energy states bounded by `[e_min, e_max]` and started at `e_init`. Returns +`(p_ch, p_dis, e)`. +""" +function add_battery_variables!(core, case::BatteryCase, T::Int, ft::Type{<:AbstractFloat}) + nBat = length(case.batteries) + pch_ub = ft.(repeat([b.p_charge_max for b in case.batteries], T)) + pdis_ub = ft.(repeat([b.p_discharge_max for b in case.batteries], T)) + core, p_ch = ExaModels.add_var(core, T * nBat; lvar = ft(0), uvar = pch_ub) + core, p_dis = ExaModels.add_var(core, T * nBat; lvar = ft(0), uvar = pdis_ub) + core, e = ExaModels.add_var(core, (T + 1) * nBat; + lvar = ft.(repeat([b.e_min for b in case.batteries], T + 1)), + uvar = ft.(repeat([b.e_max for b in case.batteries], T + 1)), + start = ft.(repeat([b.e_init for b in case.batteries], T + 1))) + return core, (p_ch = p_ch, p_dis = p_dis, e = e) +end + +# ── Objective blocks ────────────────────────────────────────────────────────── + +""" + add_generator_cost!(core, v, nd, T, dt, ft) -> core + +Original PGLib polynomial generator cost, duration-scaled: +`Δt·(c2·pg² + c1·pg + c0)` — every term, including the constant `c0`, is +multiplied by `Δt` (canonical spec, "Stage objective and cost accounting"). +""" +function add_generator_cost!(core, v, nd::NetworkData, T::Int, dt::Float64, ft::Type{<:AbstractFloat}) + nGen = ngen(nd) + items = [(t = t, g = gp, c2 = ft(g.cost2 * dt), c1 = ft(g.cost1 * dt), c0 = ft(g.cost0 * dt)) + for t in 1:T for (gp, g) in enumerate(nd.gens)] + core, _ = ExaModels.add_obj(core, + it.c2 * v.pg[_gidx(nGen, it.t, it.g)]^2 + it.c1 * v.pg[_gidx(nGen, it.t, it.g)] + it.c0 + for it in items) + return core +end + +""" + add_cycle_cost!(core, bat, case, nd, T, dt, ft) -> core + +Battery throughput/degradation cost `S^base·Δt·c^cycle_b·(p_ch + p_dis)`, +computed per battery (the fleet need not share one price). Skipped when every +battery has a zero cycle price. +""" +function add_cycle_cost!(core, bat, case::BatteryCase, nd::NetworkData, T::Int, + dt::Float64, ft::Type{<:AbstractFloat}) + nBat = length(case.batteries) + nBat == 0 && return core + items = [(idx = _kidx(nBat, t, k), c = ft(b.cycle_cost_per_mwh * nd.baseMVA * dt)) + for t in 1:T for (k, b) in enumerate(case.batteries)] + any(it -> it.c > 0, items) || return core + core, _ = ExaModels.add_obj(core, it.c * (bat.p_ch[it.idx] + bat.p_dis[it.idx]) for it in items) + return core +end + +# ── Constraint blocks ───────────────────────────────────────────────────────── + +""" + add_acp_network_constraints!(core, v, nd, T, ft) -> core + +Blocks 1–8: reference angle; the four AC-polar branch-end flow equalities (with +taps, shifts, and both line-charging shunts); angle-difference limits; and +apparent-power limits at BOTH ends for rate-limited branches only. +""" +function add_acp_network_constraints!(core, v, nd::NetworkData, T::Int, ft::Type{<:AbstractFloat}) + nBus = nbus(nd); nBranch = nbranch(nd) + br_ac = [_ac_branch_coeffs(br, ft) for br in nd.branches] + + # 1. Reference angle: va[t, ref] = 0 + ref_items = [(t = t, ref = ref) for t in 1:T for ref in nd.ref_bus_positions] + core, _ = ExaModels.add_con(core, v.va[_bidx(nBus, it.t, it.ref)] for it in ref_items) + + # 2. From-end active flow + pfr_items = [(t = t, f = br.f_pos, tb = br.t_pos, br = bp, + c3 = br_ac[bp].c3, c4 = br_ac[bp].c4, c5 = br_ac[bp].c5) + for t in 1:T for (bp, br) in enumerate(nd.branches)] + core, _ = ExaModels.add_con(core, + v.p_fr[_bridx(nBranch, it.t, it.br)] + - it.c5 * v.vm[_bidx(nBus, it.t, it.f)]^2 + - it.c3 * v.vm[_bidx(nBus, it.t, it.f)] * v.vm[_bidx(nBus, it.t, it.tb)] + * cos(v.va[_bidx(nBus, it.t, it.f)] - v.va[_bidx(nBus, it.t, it.tb)]) + - it.c4 * v.vm[_bidx(nBus, it.t, it.f)] * v.vm[_bidx(nBus, it.t, it.tb)] + * sin(v.va[_bidx(nBus, it.t, it.f)] - v.va[_bidx(nBus, it.t, it.tb)]) + for it in pfr_items) + + # 3. From-end reactive flow + qfr_items = [(t = t, f = br.f_pos, tb = br.t_pos, br = bp, + c3 = br_ac[bp].c3, c4 = br_ac[bp].c4, c6 = br_ac[bp].c6) + for t in 1:T for (bp, br) in enumerate(nd.branches)] + core, _ = ExaModels.add_con(core, + v.q_fr[_bridx(nBranch, it.t, it.br)] + + it.c6 * v.vm[_bidx(nBus, it.t, it.f)]^2 + + it.c4 * v.vm[_bidx(nBus, it.t, it.f)] * v.vm[_bidx(nBus, it.t, it.tb)] + * cos(v.va[_bidx(nBus, it.t, it.f)] - v.va[_bidx(nBus, it.t, it.tb)]) + - it.c3 * v.vm[_bidx(nBus, it.t, it.f)] * v.vm[_bidx(nBus, it.t, it.tb)] + * sin(v.va[_bidx(nBus, it.t, it.f)] - v.va[_bidx(nBus, it.t, it.tb)]) + for it in qfr_items) + + # 4. To-end active flow + pto_items = [(t = t, f = br.f_pos, tb = br.t_pos, br = bp, + c1 = br_ac[bp].c1, c2 = br_ac[bp].c2, c7 = br_ac[bp].c7) + for t in 1:T for (bp, br) in enumerate(nd.branches)] + core, _ = ExaModels.add_con(core, + v.p_to[_bridx(nBranch, it.t, it.br)] + - it.c7 * v.vm[_bidx(nBus, it.t, it.tb)]^2 + - it.c1 * v.vm[_bidx(nBus, it.t, it.tb)] * v.vm[_bidx(nBus, it.t, it.f)] + * cos(v.va[_bidx(nBus, it.t, it.tb)] - v.va[_bidx(nBus, it.t, it.f)]) + - it.c2 * v.vm[_bidx(nBus, it.t, it.tb)] * v.vm[_bidx(nBus, it.t, it.f)] + * sin(v.va[_bidx(nBus, it.t, it.tb)] - v.va[_bidx(nBus, it.t, it.f)]) + for it in pto_items) + + # 5. To-end reactive flow + qto_items = [(t = t, f = br.f_pos, tb = br.t_pos, br = bp, + c1 = br_ac[bp].c1, c2 = br_ac[bp].c2, c8 = br_ac[bp].c8) + for t in 1:T for (bp, br) in enumerate(nd.branches)] + core, _ = ExaModels.add_con(core, + v.q_to[_bridx(nBranch, it.t, it.br)] + + it.c8 * v.vm[_bidx(nBus, it.t, it.tb)]^2 + + it.c2 * v.vm[_bidx(nBus, it.t, it.tb)] * v.vm[_bidx(nBus, it.t, it.f)] + * cos(v.va[_bidx(nBus, it.t, it.tb)] - v.va[_bidx(nBus, it.t, it.f)]) + - it.c1 * v.vm[_bidx(nBus, it.t, it.tb)] * v.vm[_bidx(nBus, it.t, it.f)] + * sin(v.va[_bidx(nBus, it.t, it.tb)] - v.va[_bidx(nBus, it.t, it.f)]) + for it in qto_items) + + # 6. Angle-difference limits + ang_lb = ft.(repeat([br.angmin for br in nd.branches], T)) + ang_ub = ft.(repeat([br.angmax for br in nd.branches], T)) + ang_items = [(t = t, f = br.f_pos, tb = br.t_pos) for t in 1:T for br in nd.branches] + core, _ = ExaModels.add_con(core, + v.va[_bidx(nBus, it.t, it.f)] - v.va[_bidx(nBus, it.t, it.tb)] + for it in ang_items; lcon = ang_lb, ucon = ang_ub) + + # 7/8. Apparent-power thermal limits at both ends (rate-limited branches only). + rated = rated_branch_positions(nd) + if !isempty(rated) + th_items = [(t = t, br = bp) for t in 1:T for bp in rated] + th_ub = ft.([nd.branches[it.br].rate_a^2 for it in th_items]) + th_lb = fill(ft(-Inf), length(th_items)) + core, _ = ExaModels.add_con(core, + v.p_fr[_bridx(nBranch, it.t, it.br)]^2 + v.q_fr[_bridx(nBranch, it.t, it.br)]^2 + for it in th_items; lcon = th_lb, ucon = th_ub) + core, _ = ExaModels.add_con(core, + v.p_to[_bridx(nBranch, it.t, it.br)]^2 + v.q_to[_bridx(nBranch, it.t, it.br)]^2 + for it in th_items; lcon = th_lb, ucon = th_ub) + end + return core +end + +""" + add_nodal_balance!(core, v, bat, nd, case, T, ft, p_pd, p_qd; active_deficit=nothing) -> core + +Blocks 9–10: active and reactive nodal balance. + +Active: `p^d − d⁺ + d⁻ + gs·vm² − Σpg − Σ(p_dis−p_ch) + Σp_fr + Σp_to = 0` +Reactive: `q^d − bs·vm² − Σqg + Σq_fr + Σq_to = 0` + +`active_deficit` (`d⁺ ≥ 0`) and `active_surplus` (`d⁻ ≥ 0`) are the per-bus +**two-sided active nodal recourse** (Phase 2), or both `nothing` (Phase 1, hard +balance). They are UNBOUNDED nonnegative variables at EVERY bus: + +* `d⁺` (active deficit / injection) covers an active-power SHORTFALL — e.g. the + power to CHARGE a battery at a network-constrained bus; +* `d⁻` (active surplus / absorption) absorbs an active-power EXCESS — e.g. the + power a battery is FORCED to DISCHARGE into a bus whose outgoing branches are + thermally saturated. + +They are an artificial active-balance recourse, NOT curtailed customer load: `d⁺` +may exceed local demand and may be positive at a bus with `p^d = 0`. Together they +give relatively complete recourse: the stage subproblem is feasible for EVERY +incoming state and EVERY dynamically reachable battery target, in both the +charging and discharging directions. Both are priced at VOLL, so an accepted +solution leaves both at ~0. They do NOT touch reactive power: the reactive balance +stays a HARD equality with no slack (canonical spec, "Active load deficit" and +"Implementation invariants"). +""" +function add_nodal_balance!(core, v, bat, nd::NetworkData, case::BatteryCase, + T::Int, ft::Type{<:AbstractFloat}, p_pd, p_qd; + active_deficit = nothing, active_surplus = nothing) + nBus = nbus(nd); nGen = ngen(nd); nBranch = nbranch(nd) + nBat = length(case.batteries) + + gen_bus_items = [(brow = _bidx(nBus, t, g.bus_pos), gcol = _gidx(nGen, t, gp)) + for t in 1:T for (gp, g) in enumerate(nd.gens)] + fr_bus_items = [(brow = _bidx(nBus, t, br.f_pos), bcol = _bridx(nBranch, t, bp)) + for t in 1:T for (bp, br) in enumerate(nd.branches)] + to_bus_items = [(brow = _bidx(nBus, t, br.t_pos), bcol = _bridx(nBranch, t, bp)) + for t in 1:T for (bp, br) in enumerate(nd.branches)] + bat_bus_items = [(brow = _bidx(nBus, t, b.bus_pos), kcol = _kidx(nBat, t, k)) + for t in 1:T for (k, b) in enumerate(case.batteries)] + + # ── Active balance: p^d − d⁺ + d⁻ (both ≥ 0, unbounded above) ────────────── + kcl_p_init = [(t = t, b = b, gs = ft(nd.buses[b].gs)) for t in 1:T for b in 1:nBus] + core, c_kcl_p = if active_deficit === nothing + ExaModels.add_con(core, + p_pd[_bidx(nBus, it.t, it.b)] + it.gs * v.vm[_bidx(nBus, it.t, it.b)]^2 + for it in kcl_p_init) + else + ExaModels.add_con(core, + p_pd[_bidx(nBus, it.t, it.b)] - active_deficit[_bidx(nBus, it.t, it.b)] + + active_surplus[_bidx(nBus, it.t, it.b)] + + it.gs * v.vm[_bidx(nBus, it.t, it.b)]^2 + for it in kcl_p_init) + end + core, _ = ExaModels.add_con!(core, c_kcl_p, it.brow => -v.pg[it.gcol] for it in gen_bus_items) + core, _ = ExaModels.add_con!(core, c_kcl_p, it.brow => v.p_fr[it.bcol] for it in fr_bus_items) + core, _ = ExaModels.add_con!(core, c_kcl_p, it.brow => v.p_to[it.bcol] for it in to_bus_items) + if nBat > 0 + core, _ = ExaModels.add_con!(core, c_kcl_p, it.brow => -bat.p_dis[it.kcol] for it in bat_bus_items) + core, _ = ExaModels.add_con!(core, c_kcl_p, it.brow => bat.p_ch[it.kcol] for it in bat_bus_items) + end + + # ── Reactive balance (HARD equality; no reactive slack, recourse is active-only) ─ + kcl_q_init = [(t = t, b = b, bs = ft(nd.buses[b].bs)) for t in 1:T for b in 1:nBus] + core, c_kcl_q = ExaModels.add_con(core, + p_qd[_bidx(nBus, it.t, it.b)] - it.bs * v.vm[_bidx(nBus, it.t, it.b)]^2 + for it in kcl_q_init) + core, _ = ExaModels.add_con!(core, c_kcl_q, it.brow => -v.qg[it.gcol] for it in gen_bus_items) + core, _ = ExaModels.add_con!(core, c_kcl_q, it.brow => v.q_fr[it.bcol] for it in fr_bus_items) + core, _ = ExaModels.add_con!(core, c_kcl_q, it.brow => v.q_to[it.bcol] for it in to_bus_items) + return core +end + +""" + add_battery_dynamics!(core, bat, case, T, dt, ft, p_x0) -> core + +Blocks 11–12: the battery initial condition `e[1,b] − e0_b = 0` and the linear +energy balance +`e[t+1,b] − (1−σ_bΔt)e[t,b] − η^ch_bΔt·p_ch + (Δt/η^dis_b)·p_dis = 0`. +""" +function add_battery_dynamics!(core, bat, case::BatteryCase, T::Int, dt::Float64, + ft::Type{<:AbstractFloat}, p_x0) + nBat = length(case.batteries) + nBat == 0 && return core + core, _ = ExaModels.add_con(core, bat.e[_eidx(nBat, 1, k)] - p_x0[k] for k in 1:nBat) + st_items = [(en = _eidx(nBat, t + 1, k), ec = _eidx(nBat, t, k), pc = _kidx(nBat, t, k), + a = ft(1 - b.sigma * dt), bch = ft(b.eta_ch * dt), bdis = ft(dt / b.eta_dis)) + for t in 1:T for (k, b) in enumerate(case.batteries)] + core, _ = ExaModels.add_con(core, + bat.e[it.en] - it.a * bat.e[it.ec] - it.bch * bat.p_ch[it.pc] + it.bdis * bat.p_dis[it.pc] + for it in st_items) + return core +end + +""" + validate_stage_hours(case, dt) + +Shared validation: `Δt` finite and positive, and `σ_b·Δt < 1` for every battery +so the self-discharge retention `(1 − σΔt)` stays positive. +""" +function validate_stage_hours(case::BatteryCase, dt::Real) + (isfinite(dt) && dt > 0) || error("stage_hours (Δt) must be finite and > 0; got $dt") + for b in case.batteries + b.sigma * dt < 1 || + error("battery $(b.id) has σ·Δt = $(b.sigma*dt) ≥ 1; " * + "reduce stage_hours or self_discharge_rate") + end + return nothing +end diff --git a/examples/BatteryStorageOPF/src/battery_data.jl b/examples/BatteryStorageOPF/src/battery_data.jl new file mode 100644 index 0000000..8c3b2a5 --- /dev/null +++ b/examples/BatteryStorageOPF/src/battery_data.jl @@ -0,0 +1,269 @@ +# battery_data.jl +# +# Validated battery structures and the seeded, reproducible case-construction +# API `make_battery_case`. All battery quantities are per unit on the network +# `baseMVA` (power in pu, energy in pu·h), consistent with the one conversion +# layer in network_data.jl. + +using StableRNGs +using Random + +# ── Battery model (see BATTERY_STORAGE_OPF_PLAN.md §4) ───────────────────────── + +""" + BatteryData + +One battery, per unit on the network `baseMVA`. + +* `id` : 1-based battery index; equals placement order (stable). +* `bus_id`,`bus_pos` : host bus original PGLib id and array position. +* `p_charge_max` : max charge power p̄ᶜʰ (pu), applied as `0 ≤ pᶜʰ ≤ p̄ᶜʰ`. +* `p_discharge_max` : max discharge power p̄ᵈⁱˢ (pu), `0 ≤ pᵈⁱˢ ≤ p̄ᵈⁱˢ`. +* `e_min`,`e_max` : energy bounds e (pu·h), `e_min ≤ e ≤ e_max`. +* `e_init` : initial state of charge (pu·h). +* `eta_ch`,`eta_dis` : charge / discharge efficiencies in (0, 1]. +* `sigma` : self-discharge rate σ (per hour); state keeps `(1 − σΔt)`. +* `cycle_cost_per_mwh`: nonnegative throughput/degradation price (\$/MWh). The + objective charges `cycle_cost_per_mwh · baseMVA · Δt · (pᶜʰ + pᵈⁱˢ)` — the + dollar degradation cost of the throughput energy — which keeps the model + continuous and removes any incentive for simultaneous charge/discharge. + +State equation (linear): + `e_{t+1} = (1 − σΔt)·e_t + η_ch·Δt·pᶜʰ_t − (Δt/η_dis)·pᵈⁱˢ_t`, +active injection into the host bus: `p_bat = pᵈⁱˢ − pᶜʰ` (unity power factor). +""" +struct BatteryData + id::Int + bus_id::Int + bus_pos::Int + p_charge_max::Float64 + p_discharge_max::Float64 + e_min::Float64 + e_max::Float64 + e_init::Float64 + eta_ch::Float64 + eta_dis::Float64 + sigma::Float64 + cycle_cost_per_mwh::Float64 +end + +""" + BatteryCase + +A network plus its reproducibly-placed battery fleet and the full record needed +to rebuild it byte-identically (seed, sampling rule, sizing parameters, units). +""" +struct BatteryCase + network::NetworkData + batteries::Vector{BatteryData} + # placement / configuration record (mirrored into the manifest) + case_name::String + seed::Int + number_of_batteries::Int + duration_hours::Float64 + initial_soc::Float64 + charge_efficiency::Float64 + discharge_efficiency::Float64 + fleet_power_fraction::Float64 + self_discharge_rate::Float64 + e_min_fraction::Float64 + cycle_cost_per_mwh::Float64 + eligible_bus_rule::String + eligible_bus_ids::Vector{Int} + selected_bus_ids::Vector{Int} # stable placement order + explicit_buses::Bool + total_load_pu::Float64 + parse_meta::NamedTuple +end + +nbattery(bc::BatteryCase) = length(bc.batteries) + +# ── Eligible-bus rule ───────────────────────────────────────────────────────── + +""" + eligible_load_bus_ids(network) -> Vector{Int} + +Default battery placement pool: original ids of buses hosting at least one +in-service load with strictly positive active demand, sorted ascending. Sorting +by original id makes the sampling reproducible independent of parse order. +""" +function eligible_load_bus_ids(network::NetworkData) + ids = Int[] + for (pos, b) in enumerate(network.buses) + network.bus_pd[pos] > 0 && push!(ids, b.id) + end + return sort!(ids) +end + +const ELIGIBLE_BUS_RULE = "in_service_load_buses_with_positive_active_demand" + +# ── Case construction ───────────────────────────────────────────────────────── + +""" + make_battery_case(case_name; + number_of_batteries = 20, + seed = 20260722, + duration_hours = 4.0, + initial_soc = 0.5, + charge_efficiency = 0.95, + discharge_efficiency= 0.95, + fleet_power_fraction= 0.25, + self_discharge_rate = 0.0, + e_min_fraction = 0.0, + cycle_cost_per_mwh = 2.0, + buses = nothing, + ) -> BatteryCase + +Build a reproducible battery-storage AC-OPF case from a PGLib benchmark. + +Steps: resolve `case_name` in the pinned PGLib artifact; parse it into per-unit +[`NetworkData`]; determine the eligible-bus pool; sample `number_of_batteries` +distinct host buses **uniformly without replacement** with `StableRNG(seed)` +(unless explicit `buses` are given); derive per-battery power/energy ratings +from declared system quantities; validate every bound; and return a typed +[`BatteryCase`]. + +Sizing rule (declared, price-free): the fleet charge/discharge power is +`fleet_power_fraction · Σ load` (pu), split equally across batteries; each +battery's energy capacity is `power · duration_hours` (pu·h). This depends only +on system load and the base, never on generator prices. + +`buses` (optional) is an explicit vector of original PGLib bus ids used in the +given order instead of sampling; its length must equal `number_of_batteries`. + +All numeric inputs are validated; invalid counts, buses, efficiencies, +capacities, SoC, self-discharge, or ambiguous case names raise actionable errors. +""" +function make_battery_case(case_name::AbstractString; + number_of_batteries::Integer = 20, + seed::Integer = 20260722, + duration_hours::Real = 4.0, + initial_soc::Real = 0.5, + charge_efficiency::Real = 0.95, + discharge_efficiency::Real = 0.95, + fleet_power_fraction::Real = 0.25, + self_discharge_rate::Real = 0.0, + e_min_fraction::Real = 0.0, + cycle_cost_per_mwh::Real = 2.0, + buses = nothing) + + # ── Validate scalar inputs up front (actionable messages) ───────────────── + number_of_batteries >= 0 || + error("number_of_batteries must be ≥ 0; got $number_of_batteries") + (isfinite(duration_hours) && duration_hours > 0) || + error("duration_hours must be finite and > 0; got $duration_hours") + (0 <= e_min_fraction < 1) || + error("e_min_fraction must satisfy 0 ≤ e_min_fraction < 1; got $e_min_fraction") + (e_min_fraction <= initial_soc <= 1) || + error("initial_soc must satisfy e_min_fraction ($e_min_fraction) ≤ initial_soc ≤ 1; " * + "got $initial_soc") + (0 < charge_efficiency <= 1) || + error("charge_efficiency must satisfy 0 < η_ch ≤ 1; got $charge_efficiency") + (0 < discharge_efficiency <= 1) || + error("discharge_efficiency must satisfy 0 < η_dis ≤ 1; got $discharge_efficiency") + (isfinite(fleet_power_fraction) && fleet_power_fraction >= 0) || + error("fleet_power_fraction must be finite and ≥ 0; got $fleet_power_fraction") + (isfinite(self_discharge_rate) && 0 <= self_discharge_rate < 1) || + error("self_discharge_rate σ must satisfy 0 ≤ σ < 1 (per hour); got $self_discharge_rate") + (isfinite(cycle_cost_per_mwh) && cycle_cost_per_mwh >= 0) || + error("cycle_cost_per_mwh must be finite and ≥ 0; got $cycle_cost_per_mwh") + + # ── Resolve + parse the PGLib network ───────────────────────────────────── + network, parse_meta = load_pglib_network(case_name) + eligible = eligible_load_bus_ids(network) + total_load_pu = sum(network.bus_pd) + + # ── Determine host buses (explicit or sampled) ──────────────────────────── + explicit = buses !== nothing + local selected::Vector{Int} + if explicit + selected = Int.(collect(buses)) + length(selected) == number_of_batteries || + error("explicit `buses` has $(length(selected)) entries but " * + "number_of_batteries = $number_of_batteries") + length(unique(selected)) == length(selected) || + error("explicit `buses` contains duplicate bus ids") + for bid in selected + haskey(network.bus_id_to_pos, bid) || + error("explicit bus id $bid is not an in-service bus of \"$(network.case_name)\"") + end + else + number_of_batteries <= length(eligible) || + error("cannot place $number_of_batteries batteries: only " * + "$(length(eligible)) eligible load buses in \"$(network.case_name)\". " * + "Reduce number_of_batteries or pass explicit `buses`.") + # Uniform sampling without replacement: a seeded shuffle of the eligible + # pool, first k entries. The shuffle order IS the stable placement order. + rng = StableRNG(UInt64(unsigned(Int64(seed)))) + selected = Random.shuffle(rng, eligible)[1:number_of_batteries] + end + + # ── Derive per-battery ratings from declared system quantities ──────────── + fleet_power = fleet_power_fraction * total_load_pu + per_power = number_of_batteries == 0 ? 0.0 : fleet_power / number_of_batteries + e_max = per_power * duration_hours + e_min = e_min_fraction * e_max + e_init = initial_soc * e_max + + batteries = BatteryData[] + for (i, bid) in enumerate(selected) + push!(batteries, BatteryData(i, bid, network.bus_id_to_pos[bid], + per_power, per_power, + e_min, e_max, e_init, + Float64(charge_efficiency), + Float64(discharge_efficiency), + Float64(self_discharge_rate), + Float64(cycle_cost_per_mwh))) + end + + bc = BatteryCase(network, batteries, network.case_name, Int(seed), + Int(number_of_batteries), Float64(duration_hours), + Float64(initial_soc), Float64(charge_efficiency), + Float64(discharge_efficiency), Float64(fleet_power_fraction), + Float64(self_discharge_rate), Float64(e_min_fraction), + Float64(cycle_cost_per_mwh), + explicit ? "explicit_buses" : ELIGIBLE_BUS_RULE, + eligible, selected, explicit, total_load_pu, parse_meta) + + validate_battery_case(bc) + return bc +end + +""" + validate_battery_case(bc) -> BatteryCase + +Post-construction invariants (defence in depth on top of `make_battery_case`'s +input checks): distinct valid host buses, ordered non-degenerate energy bounds, +initial SoC inside its band, sane efficiencies, nonnegative powers and +throughput cost, and a stable-order selection consistent with the batteries. +Throws on the first violation; returns `bc` when all hold. +""" +function validate_battery_case(bc::BatteryCase) + nd = bc.network + seen = Set{Int}() + for b in bc.batteries + haskey(nd.bus_id_to_pos, b.bus_id) || + error("battery $(b.id) is on unknown bus $(b.bus_id)") + nd.bus_id_to_pos[b.bus_id] == b.bus_pos || + error("battery $(b.id) bus_pos $(b.bus_pos) inconsistent with map") + b.bus_id in seen && error("duplicate battery bus $(b.bus_id)") + push!(seen, b.bus_id) + (b.p_charge_max >= 0 && b.p_discharge_max >= 0) || + error("battery $(b.id) has negative power rating") + (b.e_min <= b.e_max) || + error("battery $(b.id) has e_min ($(b.e_min)) > e_max ($(b.e_max))") + (b.e_min - 1e-12 <= b.e_init <= b.e_max + 1e-12) || + error("battery $(b.id) initial SoC $(b.e_init) outside [$(b.e_min), $(b.e_max)]") + (0 < b.eta_ch <= 1 && 0 < b.eta_dis <= 1) || + error("battery $(b.id) has an efficiency outside (0, 1]") + (0 <= b.sigma < 1) || + error("battery $(b.id) has self-discharge σ = $(b.sigma) outside [0, 1)") + (b.cycle_cost_per_mwh >= 0) || + error("battery $(b.id) has negative cycle cost") + end + length(bc.selected_bus_ids) == length(bc.batteries) || + error("selected_bus_ids length disagrees with batteries") + all(bc.selected_bus_ids[i] == bc.batteries[i].bus_id for i in 1:length(bc.batteries)) || + error("selected_bus_ids order disagrees with battery placement order") + return bc +end diff --git a/examples/BatteryStorageOPF/src/battery_opf_exa.jl b/examples/BatteryStorageOPF/src/battery_opf_exa.jl new file mode 100644 index 0000000..c924827 --- /dev/null +++ b/examples/BatteryStorageOPF/src/battery_opf_exa.jl @@ -0,0 +1,267 @@ +# battery_opf_exa.jl +# +# Full-horizon ExaModels AC-polar deterministic equivalent for the +# battery-storage OPF (BATTERY_STORAGE_OPF_PLAN.md §4). Built directly in +# ExaModels — no JuMP, no MOI. +# +# Model (T stages, stage length Δt hours): +# min Σ_t [ Σ_g (c2_g·pg² + c1_g·pg + c0_g) +# + Σ_b c_cycle·(pᶜʰ_{t,b} + pᵈⁱˢ_{t,b}) ] +# s.t. AC-polar power flow (ref angle, 4 branch-flow eqs, angle limits, +# apparent-power thermal limits at both ends); +# hard active balance incl. battery injection p_bat = pᵈⁱˢ − pᶜʰ; +# hard reactive balance (no reactive slack, unity-PF batteries); +# e_{1,b} = e_init,b; +# e_{t+1,b} = (1−σ_bΔt)·e_{t,b} + η_ch,b·Δt·pᶜʰ − (Δt/η_dis,b)·pᵈⁱˢ; +# 0 ≤ pᶜʰ ≤ p̄ᶜʰ, 0 ≤ pᵈⁱˢ ≤ p̄ᵈⁱˢ, e_min ≤ e ≤ e_max. +# +# There is NO active-recourse variable and NO reactive-slack variable: balance is +# hard on both axes. Batteries support arbitrary (non-consecutive) host-bus ids +# through the network's stable id→position maps. + +using ExaModels +using MadNLP +using NLPModels +using LinearAlgebra + +# The AC-polar equations, index helpers, and branch coefficients live in the +# SHARED acp_core.jl (single source of truth), which this builder and the +# Phase-2 stochastic builder both call. This file keeps the Phase-1 assembly and +# its accepted behavior: no active recourse, no target constraints — the +# hard-balance base-ACP parity artifact. + +""" + BatteryExaProblem + +Holds the ExaModels deterministic equivalent for a [`BatteryCase`]. + +Fields: the `core`/`model`, the tunable parameters (`p_pd`, `p_qd`, `p_e0`), +problem sizes, horizon `T`, stage length `dt` (hours), the per-battery cycle-cost +coefficients (`cycle_coeffs`), and back-references to the case for extraction. +""" +struct BatteryExaProblem + core + model + p_pd # active demand parameter (length T*nBus) + p_qd # reactive demand parameter (length T*nBus) + p_e0 # initial-SoC parameter (length nBat) + nBus::Int + nGen::Int + nBranch::Int + nBat::Int + horizon::Int + dt::Float64 + # Per-battery cycle-cost objective coefficient (length nBat): entry k is + # cycle_cost_per_mwh_k · baseMVA · Δt, the $/pu weight on (pᶜʰ_k + pᵈⁱˢ_k). + cycle_coeffs::Vector{Float64} + case::BatteryCase + float_type::Type +end + +""" + build_battery_de(case, T; backend=nothing, float_type=Float64, + stage_hours=1.0, demand_profile=ones(T)) -> BatteryExaProblem + +Build the `T`-stage ExaModels AC-polar deterministic equivalent for `case`. + +Generator costs are PGLib polynomial USD/HOUR values, so each stage contributes +`Δt·(c2·pg² + c1·pg + c0)` to the objective (Δt = `stage_hours`); the battery +cycle cost `cᵦ·(pᶜʰ + pᵈⁱˢ)` already includes Δt through `cᵦ`. The reported +objective is therefore dollars over the whole horizon. + +* `stage_hours` — Δt, the physical length of one stage in hours (default 1.0). +* `demand_profile` — length-`T` vector of per-stage load multipliers applied to + BOTH active and reactive demand (preserving each bus's power factor); default + all ones (flat base demand). + +`backend=nothing` builds on the CPU. Requires `σ·Δt < 1` for every battery so +the self-discharge factor `(1 − σΔt)` stays positive. +""" +function build_battery_de(case::BatteryCase, T::Int; + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64, + stage_hours::Real = 1.0, + demand_profile::AbstractVector = ones(T)) + T >= 1 || error("horizon T must be ≥ 1; got $T") + length(demand_profile) == T || + error("demand_profile must have length T=$T; got $(length(demand_profile))") + dt = Float64(stage_hours) + (isfinite(dt) && dt > 0) || + error("stage_hours (Δt) must be finite and > 0; got $dt") + for b in case.batteries + b.sigma * dt < 1 || + error("battery $(b.id) has σ·Δt = $(b.sigma*dt) ≥ 1; " * + "reduce stage_hours or self_discharge_rate") + end + nd = case.network + nBus = nbus(nd); nGen = ngen(nd); nBranch = nbranch(nd) + nBat = length(case.batteries) + + # Concrete (immutable) ExaCore + the supported functional builder API + # (`add_var`/`add_par`/`add_obj`/`add_con`/`add_con!`), each returning + # `(new_core, handle)`. Handles carry their own offsets and stay valid as the + # core grows, so we thread `core` through and keep the handles. This is the + # non-deprecated path for ExaModels 0.11.2 (`ExaCore()` without `concrete` + # returns the deprecated LegacyExaCore and warns). + core = ExaModels.ExaCore(float_type; backend = backend, concrete = Val(true)) + + # ── Variables (SHARED blocks, canonical order) ──────────────────────────── + core, v = add_acp_variables!(core, nd, T, float_type) + core, bat = add_battery_variables!(core, case, T, float_type) + + # ── Parameters (per-stage demand + initial SoC) ─────────────────────────── + # Phase 1 bakes the fixed `demand_profile` into the demand parameters. + prof = Float64.(collect(demand_profile)) + init_pd = float_type.([nd.bus_pd[b] * prof[t] for t in 1:T for b in 1:nBus]) + init_qd = float_type.([nd.bus_qd[b] * prof[t] for t in 1:T for b in 1:nBus]) + core, p_pd = ExaModels.add_par(core, init_pd) + core, p_qd = ExaModels.add_par(core, init_qd) + core, p_e0 = ExaModels.add_par(core, float_type.([b.e_init for b in case.batteries])) + + # ── Objective (SHARED): generator cost + battery throughput cost ────────── + core = add_generator_cost!(core, v, nd, T, dt, float_type) + core = add_cycle_cost!(core, bat, case, nd, T, dt, float_type) + + # ── Constraints (SHARED). Phase 1 passes `active_deficit = nothing`: hard + # active and reactive balance with no recourse — the accepted parity artifact. + core = add_acp_network_constraints!(core, v, nd, T, float_type) + core = add_nodal_balance!(core, v, bat, nd, case, T, float_type, p_pd, p_qd; + active_deficit = nothing) + core = add_battery_dynamics!(core, bat, case, T, dt, float_type, p_e0) + + cycle_coeffs = Float64[b.cycle_cost_per_mwh * nd.baseMVA * dt for b in case.batteries] + + model = ExaModels.ExaModel(core) + return BatteryExaProblem(core, model, p_pd, p_qd, p_e0, + nBus, nGen, nBranch, nBat, T, dt, + cycle_coeffs, case, float_type) +end + +# ── Parameter setters ───────────────────────────────────────────────────────── + +""" + set_demand!(prob, pd::AbstractMatrix, qd::AbstractMatrix) + +Overwrite the per-stage active/reactive demand parameters. `pd`, `qd` are +`[T × nBus]` (pu). Both must be supplied so the reactive/active balance stay +consistent. +""" +function set_demand!(prob::BatteryExaProblem, pd::AbstractMatrix, qd::AbstractMatrix) + T, nB = prob.horizon, prob.nBus + size(pd) == (T, nB) || error("pd must be [T=$T × nBus=$nB]") + size(qd) == (T, nB) || error("qd must be [T=$T × nBus=$nB]") + ExaModels.set_parameter!(prob.core, prob.p_pd, + prob.float_type.([pd[t, b] for t in 1:T for b in 1:nB])) + ExaModels.set_parameter!(prob.core, prob.p_qd, + prob.float_type.([qd[t, b] for t in 1:T for b in 1:nB])) + return prob +end + +""" + set_initial_soc!(prob, e0::AbstractVector) + +Overwrite the initial state-of-charge parameter (length `nBat`, pu·h). +""" +function set_initial_soc!(prob::BatteryExaProblem, e0::AbstractVector) + length(e0) == prob.nBat || error("e0 must have length nBat=$(prob.nBat)") + ExaModels.set_parameter!(prob.core, prob.p_e0, prob.float_type.(collect(e0))) + return prob +end + +# ── Solve + structured extraction ───────────────────────────────────────────── + +solve_succeeded(status) = status == MadNLP.SOLVE_SUCCEEDED || + status == MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL + +""" + solve_de!(prob; print_level=MadNLP.ERROR, kwargs...) -> result + +Solve the deterministic equivalent on the CPU with MadNLP. Returns the raw +MadNLP result (with `.status`, `.objective`, `.solution`, `.multipliers`). + +Named `solve_de!` rather than `solve!` to avoid clashing with the `solve!` +that JuMP/MadNLP re-export from CommonSolve. +""" +solve_de!(prob::BatteryExaProblem; print_level = MadNLP.ERROR, kwargs...) = + MadNLP.madnlp(prob.model; print_level = print_level, kwargs...) + +""" + battery_solution(prob, result) -> NamedTuple + +Reshape the flat solution into named `[·×T]` (or `[·×(T+1)]` for SoC) matrices: +`va, vm` (nBus), `pg, qg` (nGen), `p_fr, q_fr, p_to, q_to` (nBranch), +`p_ch, p_dis` (nBat), `soc` (nBat×(T+1)), and derived `p_bat = p_dis − p_ch`. +""" +function battery_solution(prob::BatteryExaProblem, result) + T = prob.horizon; nB = prob.nBus; nG = prob.nGen; nBR = prob.nBranch; nK = prob.nBat + sol = Array(result.solution) + off = 0 + take(n, m) = (v = reshape(sol[off .+ (1:n*m)], n, m); off += n*m; v) + va = take(nB, T) + vm = take(nB, T) + pg = take(nG, T) + qg = take(nG, T) + p_fr = take(nBR, T) + q_fr = take(nBR, T) + p_to = take(nBR, T) + q_to = take(nBR, T) + if nK > 0 + p_ch = take(nK, T) + p_dis = take(nK, T) + soc = take(nK, T + 1) + else + p_ch = zeros(eltype(sol), 0, T) + p_dis = zeros(eltype(sol), 0, T) + soc = zeros(eltype(sol), 0, T + 1) + end + return (va = va, vm = vm, pg = pg, qg = qg, + p_fr = p_fr, q_fr = q_fr, p_to = p_to, q_to = q_to, + p_ch = p_ch, p_dis = p_dis, soc = soc, p_bat = p_dis .- p_ch) +end + +""" + battery_balance_residuals(prob, sol) -> Matrix + +`[nBat × T]` residuals of the state equation +`e[t+1] − (1−σΔt)·e[t] − η_ch·Δt·p_ch + (Δt/η_dis)·p_dis`, which should be ~0. +""" +function battery_balance_residuals(prob::BatteryExaProblem, sol) + T = prob.horizon; nK = prob.nBat + res = zeros(Float64, nK, T) + for (k, bat) in enumerate(prob.case.batteries), t in 1:T + res[k, t] = sol.soc[k, t+1] - (1 - bat.sigma * prob.dt) * sol.soc[k, t] - + bat.eta_ch * prob.dt * sol.p_ch[k, t] + + (prob.dt / bat.eta_dis) * sol.p_dis[k, t] + end + return res +end + +""" + simultaneous_charge_discharge_power(sol) -> Matrix + +`[nBat × T]` **simultaneous charge/discharge power** in pu: the elementwise +smaller of the charge and discharge powers, `min(p_ch, p_dis)` (NOT a product). +It is the amount of power a battery is charging and discharging at the same time; +material positive entries flag unwanted simultaneous operation, which the +nonnegative cycle cost is designed to prevent. +""" +simultaneous_charge_discharge_power(sol) = min.(sol.p_ch, sol.p_dis) + +""" + max_primal_residual(prob, result) -> Float64 + +Largest constraint-bound violation of the returned point, evaluated +independently through the NLPModels interface (max over +`max(lcon − c, c − ucon, 0)`). +""" +function max_primal_residual(prob::BatteryExaProblem, result) + x = Array(result.solution) + c = NLPModels.cons(prob.model, x) + lcon = Array(prob.model.meta.lcon); ucon = Array(prob.model.meta.ucon) + c = Array(c) + viol = 0.0 + @inbounds for i in eachindex(c) + viol = max(viol, lcon[i] - c[i], c[i] - ucon[i], 0.0) + end + return viol +end diff --git a/examples/BatteryStorageOPF/src/battery_policy.jl b/examples/BatteryStorageOPF/src/battery_policy.jl new file mode 100644 index 0000000..c41be0f --- /dev/null +++ b/examples/BatteryStorageOPF/src/battery_policy.jl @@ -0,0 +1,245 @@ +# battery_policy.jl +# +# Battery-SoC reachable target policy (canonical spec: +# docs/src/casestudies/battery_storage_opf.md, "Reachable target policy"). +# +# The policy observes the current battery SoC e_t and the current demand +# features, and outputs the TARGET outgoing energy ê_{t+1}. It never observes +# future atoms; the recurrent encoder carries ξ_{1:t}. +# +# One-stage reachable interval (battery dynamics only — this is NOT a proof of +# ACP network feasibility): +# +# ℓ = max( e_min , (1 − σΔt)·e_prev − (Δt/η_dis)·p̄_dis ) +# u = min( e_max , (1 − σΔt)·e_prev + η_ch·Δt·p̄_ch ) +# +# and the normalized output y is mapped affinely: ê = ℓ + (u − ℓ)·y. +# +# ── Activation (why the safe upper margin matters) ──────────────────────────── +# The canonical default is `stretchedsigmoid`: +# +# y = clamp((sigmoid(z) − 0.03)/0.94, 0, 1 − 1e-3) +# +# A plain sigmoid attains 0/1 only at ±∞ with vanishing gradient, so a squashed +# policy can never learn to sit at the reachable-interval boundary where optimal +# storage decisions frequently live. The gentle stretch attains exactly 0 at +# finite pre-activation while keeping a small margin BELOW the exact upper edge: +# an exact y = 1 (store-max) target under a STRICT equality drives the stage NLP +# onto a measure-zero set that interior-point solvers cannot converge into +# (observed as MAXIMUM_ITERATIONS / spurious INFEASIBLE). The upper normalized +# target must therefore never equal 1. +# +# `hardsigmoidsafe(z) = clamp(0.5 + 0.5z, 0, 1 − 1e-3)` is retained as a +# documented option with the same safe upper margin. +# +# Reachability bounds are physical projection DATA, not learned functions: the +# gradient stops through ℓ and u (ChainRulesCore `@non_differentiable`), and +# flows only through the normalized policy output. Recurrent state is reset at +# every scenario boundary. + +using Flux +using NNlib +using ChainRulesCore +using DecisionRulesExa + +""" + stretchedsigmoid(z) -> y ∈ [0, 1 − 1e-3] + +Canonical boundary-attaining activation +`clamp((sigmoid(z) − 0.03)/0.94, 0, 1 − 1e-3)`. + +Attains exactly `0` for `sigmoid(z) ≤ 0.03` (z ≈ −3.5) and saturates at the +δ-interior upper value `1 − 1e-3` for `sigmoid(z) ≥ 0.969`, so the policy can +learn boundary targets at finite weights while never requesting an exact +store-max target (which is interior-point degenerate under a strict equality). +""" +function stretchedsigmoid(z::Real) + T = float(typeof(z)) + return clamp((NNlib.sigmoid(z) - T(0.03)) / T(0.94), zero(T), one(T) - T(1e-3)) +end + +""" + hardsigmoidsafe(z) -> y ∈ [0, 1 − 1e-3] + +Optional piecewise-linear activation `clamp(0.5 + 0.5z, 0, 1 − 1e-3)`, with the +same safe upper margin as [`stretchedsigmoid`](@ref): it attains exactly `0` at +`z ≤ −1` and `1 − 1e-3` at `z ≥ 1`, with constant interior slope `0.5`. +""" +function hardsigmoidsafe(z::Real) + T = float(typeof(z)) + return clamp(T(0.5) + T(0.5) * z, zero(T), one(T) - T(1e-3)) +end + +# Activations admissible for the target head: range inside [0, 1 − 1e-3] so the +# affine map into the reachable interval stays feasible and never hits the exact +# upper edge. +const BOUNDED_TARGET_ACTIVATIONS = (stretchedsigmoid, hardsigmoidsafe) + +""" + battery_reachable_bounds(e_prev, a, discharge_drop, charge_gain, e_min, e_max) + -> (lower, upper) + +One-stage physical reachability interval, element-wise over batteries, from +precomputed coefficients `a = 1 − σΔt`, `discharge_drop = (Δt/η_dis)·p̄_dis`, +`charge_gain = η_ch·Δt·p̄_ch`, and the energy bounds. + +`upper` is additionally floored at `lower` so a degenerate interval stays +well-defined. This function is physical projection data: it is marked +non-differentiable, so no gradient propagates through the bounds. +""" +function battery_reachable_bounds(e_prev, a, discharge_drop, charge_gain, e_min, e_max) + base = a .* e_prev + lower = max.(e_min, base .- discharge_drop) + upper = min.(e_max, base .+ charge_gain) + upper = max.(upper, lower) + return lower, upper +end + +# Gradients must NOT flow through the reachable bounds (canonical spec: +# "Reachability bounds are physical projection data, not learned functions"). +ChainRulesCore.@non_differentiable battery_reachable_bounds(::Any, ::Any, ::Any, ::Any, ::Any, ::Any) + +""" + BatteryReachablePolicy{E,C,S,V} + +Flux-compatible battery-SoC target policy: a recurrent uncertainty encoder, a +bounded-activation target head, and the affine map into the one-stage reachable +interval. + +# Fields +- `encoder`: recurrent uncertainty encoder (`Chain` of `Flux.LSTM`-style layers); + trainable. +- `combiner`: target head reading `[encoded_uncertainty; e_prev]`, whose output + activation is a bounded target activation; trainable. +- `state`: encoder recurrent state, threaded across stages (not trainable), + restored by `Flux.reset!` at every scenario boundary. +- `a`, `discharge_drop`, `charge_gain`, `e_min`, `e_max`: per-battery + reachability coefficients (physical data; moved across devices with the + policy). +- `n_uncertainty`: number of leading uncertainty features per input. +- `nbat`: number of batteries (state dimension). +""" +mutable struct BatteryReachablePolicy{E,C,S,V} + encoder::E + combiner::C + state::S + a::V + discharge_drop::V + charge_gain::V + e_min::V + e_max::V + n_uncertainty::Int + nbat::Int +end + +Flux.@layer BatteryReachablePolicy trainable=(encoder, combiner) + +""" + (policy::BatteryReachablePolicy)(input) -> target + +One-stage forward pass on `input = vcat(w_t, e_prev)`. Advances the recurrent +encoder by one step, produces the normalized target `y ∈ [0, 1 − 1e-3]`, and +maps it into the one-stage reachable interval derived from `e_prev`. +""" +function (m::BatteryReachablePolicy)(input) + w = input[1:m.n_uncertainty] + e_prev = input[m.n_uncertainty + 1:end] # range index → view, GPU-safe + F = DecisionRulesExa._state_eltype(m.state) + h, new_state = DecisionRulesExa._step_encoder(m.encoder, F.(w), m.state) + m.state = new_state # thread recurrent state + y = m.combiner(vcat(h, e_prev)) # normalized, in [0, 1−1e-3] + lower, upper = battery_reachable_bounds(e_prev, m.a, m.discharge_drop, + m.charge_gain, m.e_min, m.e_max) + return lower .+ (upper .- lower) .* y +end + +""" + Flux.reset!(policy::BatteryReachablePolicy) -> Nothing + +Restore the encoder's recurrent state to `Flux.initialstates`, re-derived from +the (possibly device-moved) encoder weights. Called at every scenario boundary. +""" +function Flux.reset!(m::BatteryReachablePolicy) + m.state = DecisionRulesExa._init_recurrent_state(m.encoder) + return nothing +end + +""" + load_battery_policy!(policy, state) -> policy + +Load a Flux checkpoint state into a [`BatteryReachablePolicy`] and reset the +recurrent state, so the reloaded policy reproduces outputs exactly. +""" +function load_battery_policy!(m::BatteryReachablePolicy, state) + Flux.loadmodel!(m, state) + Flux.reset!(m) + return m +end + +# Cast a Flux model's parameters to the requested precision. Flux builds +# Dense/LSTM in Float32 by default; only Float64 needs an explicit cast. +_cast_model(::Type{Float32}, m) = m +_cast_model(::Type{Float64}, m) = Flux.f64(m) +_cast_model(::Type{<:AbstractFloat}, m) = m + +""" + battery_reachable_policy(case, process; dt=1.0, layers=[32,32], + combiner_layers=Int[], activation=stretchedsigmoid, + encoder_type=Flux.LSTM, float_type=Float32) -> BatteryReachablePolicy + +Construct the reachable target policy for a [`BatteryCase`] and its +[`LoadProcess`]. + +The recurrent encoder reads the `nw = 1 + nregion` per-stage uncertainty +features; the head reads `[encoded_uncertainty; e_prev]` and emits a normalized +target through `activation`, which must be a bounded target activation (range +inside `[0, 1 − 1e-3]`). `dt` MUST match the `stage_hours` of the problem this +policy drives. +""" +function battery_reachable_policy(case::BatteryCase, process::LoadProcess; + dt::Real = 1.0, + layers::AbstractVector{<:Integer} = [32, 32], + combiner_layers::AbstractVector{<:Integer} = Int[], + activation = stretchedsigmoid, + encoder_type = Flux.LSTM, + float_type::Type{<:AbstractFloat} = Float32) + nBat = nbattery(case) + nBat >= 1 || error("battery_reachable_policy needs ≥ 1 battery; got $nBat") + length(layers) >= 1 || error("layers must have ≥ 1 encoder layer") + any(a -> activation === a, BOUNDED_TARGET_ACTIVATIONS) || + throw(ArgumentError("activation must be a bounded target activation with the safe " * + "upper margin (stretchedsigmoid or hardsigmoidsafe)")) + nw = n_uncertainty(process) + Δt = Float64(dt) + validate_stage_hours(case, Δt) + + enc_sizes = vcat(nw, collect(Int, layers)) + enc_layers = [encoder_type(enc_sizes[i] => enc_sizes[i + 1]) for i in 1:length(layers)] + encoder = _cast_model(float_type, Flux.Chain(enc_layers...)) + # Canonical head: the bounded activation is applied at every head layer, + # including the output (DecisionRulesExa._dense_policy_head semantics). + combiner = _cast_model(float_type, + DecisionRulesExa._dense_policy_head(Int(layers[end]) + nBat, nBat, + collect(Int, combiner_layers); + activation = activation)) + + a = float_type.([1 - b.sigma * Δt for b in case.batteries]) + discharge_drop = float_type.([(Δt / b.eta_dis) * b.p_discharge_max for b in case.batteries]) + charge_gain = float_type.([b.eta_ch * Δt * b.p_charge_max for b in case.batteries]) + e_min = float_type.([b.e_min for b in case.batteries]) + e_max = float_type.([b.e_max for b in case.batteries]) + + return BatteryReachablePolicy(encoder, combiner, + DecisionRulesExa._init_recurrent_state(encoder), + a, discharge_drop, charge_gain, e_min, e_max, + nw, nBat) +end + +""" + policy_initial_state(case; float_type=Float32) -> Vector + +Initial battery-SoC state vector `e_init` (length `nBat`) in policy precision — +the `initial_state` argument for `train_tsddr`/`rollout_tsddr`. +""" +policy_initial_state(case::BatteryCase; float_type::Type{<:AbstractFloat} = Float32) = + float_type.([b.e_init for b in case.batteries]) diff --git a/examples/BatteryStorageOPF/src/battery_training.jl b/examples/BatteryStorageOPF/src/battery_training.jl new file mode 100644 index 0000000..333e382 --- /dev/null +++ b/examples/BatteryStorageOPF/src/battery_training.jl @@ -0,0 +1,556 @@ +# battery_training.jl +# +# Beginner-usable TS-DDR training and evaluation entrypoints for the +# battery-storage example (Phase-2 prompt §5), built on the parent +# DecisionRulesExa APIs (`train_tsddr`, `rollout_tsddr`) and the target-constrained +# problem in battery_tsddr.jl. +# +# * `make_replay_sampler` — deterministic scenario replay from a stored +# atom-index matrix (training uses the train +# protocol; evaluation uses the eval protocol). +# * `train_battery_tsddr` — one call that trains a reachable policy on the +# full-horizon target-constrained DE, auditing and +# COUNTING failed solves (never silently dropping). +# * `evaluate_battery_policy` / `evaluate_paired` — honest, non-anticipative +# STAGE-WISE rollout on fixed paired scenarios, with +# the full physical/penalty cost decomposition and a +# compact machine-readable trajectory. +# * checkpoint save/load with exact reload; trajectory serialization. +# +# Physical operating cost and training-only target penalty are kept strictly +# separate everywhere (see `decompose_costs`); improvement is always judged on +# physical operating cost, never on the total objective. + +using DecisionRulesExa +using Flux +using MadNLP +using JSON +using Serialization +using Statistics + +# ── Deterministic scenario replay ───────────────────────────────────────────── + +""" + make_replay_sampler(process, index_matrix) -> (sampler, scenarios) + +Return a zero-argument `sampler()` that yields materialized `w_flat` vectors by +cycling deterministically through the columns of a stage-major `index_matrix` +(one scenario per column). The same matrix always replays the same sequence, so +training is exactly reproducible. `scenarios` is the pre-materialized vector. +""" +function make_replay_sampler(process::LoadProcess, index_matrix::AbstractMatrix{<:Integer}) + scenarios = materialize_all(process, index_matrix) + npaths = length(scenarios) + npaths >= 1 || error("index_matrix must have ≥ 1 path") + i = Ref(0) + sampler = function () + i[] += 1 + return copy(scenarios[((i[] - 1) % npaths) + 1]) + end + return sampler, scenarios +end + +# ── Training ────────────────────────────────────────────────────────────────── + +""" + train_battery_tsddr(policy, de, process, train_index_matrix; kwargs...) -> NamedTuple + +Train a [`BatteryReachablePolicy`] on a full-horizon [`BatteryTSDDRProblem`] `de` +via the parent `train_tsddr`, replaying the training scenarios in +`train_index_matrix` deterministically. + +# Keywords +- `initial_state = policy_initial_state(de.case)`: initial SoC state vector. +- `num_batches::Int = 20`, `num_train_per_batch::Int = 4`. +- `optimizer = Adam(1e-3)` (Flux optimizer/chain). +- `madnlp_kwargs = (print_level=MadNLP.ERROR, tol=1e-6)`. +- `record_loss`: `(iter, model, loss, tag) -> Bool` callback (default prints). +- `warmstart::Bool = true`, `retry_on_failure::Bool = true`. + +Returns `(model, n_ok, n_total, n_failed, failure_counts)` — every failed solve +is counted and reported, never silently discarded. `n_ok`/`n_total` count solves +across all batches. +""" +function train_battery_tsddr(policy, de::BatteryTSDDRProblem, process::LoadProcess, + train_index_matrix::AbstractMatrix{<:Integer}; + initial_state = policy_initial_state(de.case), + num_batches::Int = 20, + num_train_per_batch::Int = 4, + optimizer = Flux.Adam(1f-3), + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-6), + record_loss = (iter, model, loss, tag) -> begin + println(" iter=$iter mean_physical_obj=$(round(loss; digits = 3))") + return false + end, + warmstart::Bool = true, + retry_on_failure::Bool = true) + is_targetless(de) && + error("train_battery_tsddr requires a target mode (:strict or :soft); " * + "a targetless diagnostic problem has no target constraints to differentiate") + size(train_index_matrix, 1) == de.horizon || + error("train_index_matrix has $(size(train_index_matrix,1)) stages but de.horizon=$(de.horizon)") + sampler, _ = make_replay_sampler(process, train_index_matrix) + + failure_counts = Dict{String,Int}() + n_ok = Ref(0); n_total = Ref(0) + diag = (iter, stats) -> begin + n_ok[] += get(stats, "n_ok", 0) + n_total[] += get(stats, "n_total", 0) + for (k, v) in get(stats, "failure_counts", Dict{String,Int}()) + failure_counts[k] = get(failure_counts, k, 0) + v + end + return nothing + end + + train_tsddr(policy, initial_state, de, de.p_x0, de.p_target, de.p_w, sampler; + num_batches = num_batches, num_train_per_batch = num_train_per_batch, + optimizer = optimizer, madnlp_kwargs = madnlp_kwargs, + warmstart = warmstart, retry_on_failure = retry_on_failure, + batch_diagnostics = diag, record_loss = record_loss) + + return (model = policy, n_ok = n_ok[], n_total = n_total[], + n_failed = n_total[] - n_ok[], failure_counts = failure_counts) +end + +# ── Honest stage-wise rollout evaluation + compact trajectory ───────────────── + +""" + evaluate_battery_policy(policy, stage_problem, process, w_flat, atom_row; + reporting_horizon, madnlp_kwargs, warmstart, retry_on_failure) + -> Union{Nothing, NamedTuple} + +Evaluate `policy` on ONE materialized scenario `w_flat` by a non-anticipative +STAGE-WISE rollout of the single-stage `stage_problem` (built with `horizon = 1`). +At each stage the policy sees only the current `w_t` and the realized SoC, sets a +one-stage-reachable target, and the stage AC-OPF is solved; the realized next SoC +feeds the next stage. + +`atom_row` is the scenario's atom-index column (for the trajectory record). +`reporting_horizon` splits reported physical cost from the look-ahead buffer. + +Returns `nothing` if any stage solve fails after retry, otherwise a NamedTuple: +- `reporting_physical_cost`, `lookahead_physical_cost`, `physical_operating_cost`, + `generator_cost`, `battery_throughput_cost`, `target_penalty`; +- `final_soc`; +- `trajectory`: a vector of per-stage records (stage, atom, soc_in, target, + soc_out, p_charge, p_discharge, generator_cost, physical_cost, target_penalty, + status, max_primal_residual, max_balance_residual). + +Uses the parent `rollout_tsddr` for the vetted stage loop (retry, projection), +capturing per-stage physical costs and solutions through its callbacks. +""" +function evaluate_battery_policy(policy, stage_problem::BatteryTSDDRProblem, + process::LoadProcess, w_flat::AbstractVector, + atom_row::AbstractVector{<:Integer}; + reporting_horizon::Int, + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-6), + warmstart::Bool = false, + retry_on_failure::Bool = true) + stage_problem.horizon == 1 || + error("evaluate_battery_policy needs a single-stage stage_problem (horizon=1)") + is_targetless(stage_problem) && + error("evaluate_battery_policy requires a target mode; got a targetless diagnostic") + nw = stage_problem.nw + T = length(w_flat) ÷ nw + length(w_flat) == T * nw || error("w_flat length not a multiple of nw=$nw") + length(atom_row) == T || error("atom_row length $(length(atom_row)) ≠ horizon $T") + 1 <= reporting_horizon <= T || + error("reporting_horizon must satisfy 1 ≤ R ≤ $T; got $reporting_horizon") + nK = stage_problem.nBat + e0 = policy_initial_state(stage_problem.case; float_type = Float32) + + # Per-stage capture buffers (filled inside the rollout callbacks). + stage_ctr = Ref(0) + phys_stage = zeros(Float64, T) + gen_stage = zeros(Float64, T) + cyc_stage = zeros(Float64, T) + rec_cost_stage = zeros(Float64, T) + deficit_mwh_stage = zeros(Float64, T) + surplus_mwh_stage = zeros(Float64, T) + max_deficit_stage = zeros(Float64, T) + max_surplus_stage = zeros(Float64, T) + raw_rec_cost_stage = zeros(Float64, T) + proj_corr_stage = zeros(Float64, T) + lb_viol_stage = zeros(Float64, T) + pen_stage = zeros(Float64, T) + viol_stage = zeros(Float64, T) + status_stage = Vector{Any}(undef, T) + prim_stage = zeros(Float64, T) + bal_stage = zeros(Float64, T) + pch_stage = [zeros(Float64, nK) for _ in 1:T] + pdis_stage = [zeros(Float64, nK) for _ in 1:T] + soc_out_stage = [zeros(Float64, nK) for _ in 1:T] + + set_params! = function (prob, state, w_t, target, stage) + stage_ctr[] = stage + set_tsddr_initial_soc!(prob, state) + set_tsddr_uncertainty!(prob, w_t) # also realizes p_pd/p_qd for this stage + set_tsddr_targets!(prob, target) + return nothing + end + # Physical stage cost (generator + throughput + active-recourse cost; the + # training-only target penalty is stripped). Also records the decomposition, + # keeping the two recourse directions (deficit d⁺ and surplus d⁻) separate. + no_penalty = function (prob, result) + s = stage_ctr[] + d = decompose_costs(prob, result) + gen_stage[s] = d.generator_cost + cyc_stage[s] = d.battery_throughput_cost + rec_cost_stage[s] = d.active_recourse_cost + deficit_mwh_stage[s] = d.active_deficit_energy_mwh + surplus_mwh_stage[s] = d.active_surplus_energy_mwh + max_deficit_stage[s] = d.max_active_deficit_pu + max_surplus_stage[s] = d.max_active_surplus_pu + raw_rec_cost_stage[s] = d.raw_active_recourse_cost + proj_corr_stage[s] = d.active_recourse_projection_correction + lb_viol_stage[s] = d.maximum_active_recourse_lower_bound_violation_pu + phys_stage[s] = d.physical_operating_cost + pen_stage[s] = d.target_penalty + viol_stage[s] = d.target_violation + status_stage[s] = result.status + prim_stage[s] = tsddr_max_primal_residual(prob, result) + solv = tsddr_solution(prob, result) + bal_stage[s] = nK > 0 ? maximum(abs, tsddr_balance_residuals(prob, solv)) : 0.0 + if nK > 0 + pch_stage[s] .= Float64.(solv.p_ch[:, 1]) + pdis_stage[s] .= Float64.(solv.p_dis[:, 1]) + soc_out_stage[s] .= Float64.(solv.soc[:, 2]) + end + return d.physical_operating_cost + end + realized = function (prob, result) + # The realized next SoC (e[:,2]) is the state passed to the next stage. + solv = tsddr_solution(prob, result) + return nK > 0 ? Float32.(solv.soc[:, 2]) : Float32[] + end + + out = rollout_tsddr(policy, e0, stage_problem, Float64.(w_flat); + horizon = T, n_uncertainty = nw, + set_stage_parameters! = set_params!, + realized_state = realized, + objective_no_target_penalty = no_penalty, + madnlp_kwargs = madnlp_kwargs, + warmstart = warmstart, policy_state = :realized, + retry_on_failure = retry_on_failure) + out === nothing && return nothing + + # Assemble the compact trajectory and the reporting/look-ahead split. + traj = Vector{Dict{String,Any}}(undef, T) + for s in 1:T + soc_in = Float64.(out.state_trajectory[s]) + target = Float64.(out.target_trajectory[s]) + traj[s] = Dict{String,Any}( + "stage" => s, + "atom" => Int(atom_row[s]), + "soc_in" => soc_in, + "target" => target, + "soc_out" => soc_out_stage[s], + "p_charge" => pch_stage[s], + "p_discharge" => pdis_stage[s], + "generator_cost" => gen_stage[s], + "battery_throughput_cost" => cyc_stage[s], + "active_recourse_cost" => rec_cost_stage[s], + "active_deficit_energy_mwh" => deficit_mwh_stage[s], + "active_surplus_energy_mwh" => surplus_mwh_stage[s], + "max_active_deficit_pu" => max_deficit_stage[s], + "max_active_surplus_pu" => max_surplus_stage[s], + "raw_active_recourse_cost" => raw_rec_cost_stage[s], + "active_recourse_projection_correction" => proj_corr_stage[s], + "maximum_active_recourse_lower_bound_violation_pu" => lb_viol_stage[s], + "physical_cost" => phys_stage[s], + "target_penalty" => pen_stage[s], + "target_violation" => viol_stage[s], + "status" => string(status_stage[s]), + "max_primal_residual" => prim_stage[s], + "max_balance_residual" => bal_stage[s], + ) + end + reporting = sum(@view phys_stage[1:reporting_horizon]) + lookahead = T > reporting_horizon ? sum(@view phys_stage[reporting_horizon+1:T]) : 0.0 + return (reporting_physical_cost = reporting, + lookahead_physical_cost = lookahead, + physical_operating_cost = sum(phys_stage), + generator_cost = sum(gen_stage), + battery_throughput_cost = sum(cyc_stage), + active_recourse_cost = sum(rec_cost_stage), + active_deficit_energy_mwh = sum(deficit_mwh_stage), + active_surplus_energy_mwh = sum(surplus_mwh_stage), + total_active_recourse_energy_mwh = sum(deficit_mwh_stage) + sum(surplus_mwh_stage), + max_active_deficit_pu = maximum(max_deficit_stage), + max_active_surplus_pu = maximum(max_surplus_stage), + raw_active_recourse_cost = sum(raw_rec_cost_stage), + active_recourse_projection_correction = sum(proj_corr_stage), + maximum_active_recourse_lower_bound_violation_pu = maximum(lb_viol_stage), + target_penalty = sum(pen_stage), + target_violation = sum(viol_stage), + final_soc = Float64.(out.final_state), + trajectory = traj) +end + +""" + evaluate_paired(policy, stage_problem, process, index_matrix; + reporting_horizon, madnlp_kwargs, keep_trajectories=false) -> NamedTuple + +Evaluate `policy` on the fixed paired scenario set defined by a stage-major +`index_matrix` (every method shares the SAME matrix), returning: +- `mean_reporting_physical_cost` (over successful paths), +- `reporting_physical_costs` (per successful path), +- `n_ok`, `n_failed`, `failed_paths`, +- `trajectories` (per path) when `keep_trajectories=true`. + +Improvement is judged on `mean_reporting_physical_cost` — the physical operating +cost over the reporting horizon, with NO target penalty. +""" +function evaluate_paired(policy, stage_problem::BatteryTSDDRProblem, process::LoadProcess, + index_matrix::AbstractMatrix{<:Integer}; + reporting_horizon::Int, + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-6), + keep_trajectories::Bool = false) + horizon, npaths = size(index_matrix) + costs = Float64[] + deficit_mwh = Float64[] + surplus_mwh = Float64[] + max_deficit = Float64[] + max_surplus = Float64[] + raw_rec_cost = Float64[] + proj_corr = Float64[] + lb_viol = Float64[] + failed_paths = Int[] + trajectories = keep_trajectories ? Vector{Any}(undef, npaths) : nothing + for p in 1:npaths + w_flat = materialize_scenario(process, view(index_matrix, :, p); horizon = horizon) + res = evaluate_battery_policy(policy, stage_problem, process, w_flat, + view(index_matrix, :, p); + reporting_horizon = reporting_horizon, + madnlp_kwargs = madnlp_kwargs) + if res === nothing + push!(failed_paths, p) + keep_trajectories && (trajectories[p] = nothing) + else + push!(costs, res.reporting_physical_cost) + push!(deficit_mwh, res.active_deficit_energy_mwh) + push!(surplus_mwh, res.active_surplus_energy_mwh) + push!(max_deficit, res.max_active_deficit_pu) + push!(max_surplus, res.max_active_surplus_pu) + push!(raw_rec_cost, res.raw_active_recourse_cost) + push!(proj_corr, res.active_recourse_projection_correction) + push!(lb_viol, res.maximum_active_recourse_lower_bound_violation_pu) + keep_trajectories && (trajectories[p] = res) + end + end + tot_def = isempty(deficit_mwh) ? 0.0 : sum(deficit_mwh) + tot_sur = isempty(surplus_mwh) ? 0.0 : sum(surplus_mwh) + return (mean_reporting_physical_cost = isempty(costs) ? NaN : mean(costs), + reporting_physical_costs = costs, + n_ok = length(costs), n_failed = length(failed_paths), + failed_paths = failed_paths, + total_active_deficit_energy_mwh = tot_def, + total_active_surplus_energy_mwh = tot_sur, + total_active_recourse_energy_mwh = tot_def + tot_sur, + max_active_deficit_pu = isempty(max_deficit) ? 0.0 : maximum(max_deficit), + max_active_surplus_pu = isempty(max_surplus) ? 0.0 : maximum(max_surplus), + total_raw_active_recourse_cost = isempty(raw_rec_cost) ? 0.0 : sum(raw_rec_cost), + total_active_recourse_projection_correction = isempty(proj_corr) ? 0.0 : sum(proj_corr), + maximum_active_recourse_lower_bound_violation_pu = isempty(lb_viol) ? 0.0 : maximum(lb_viol), + trajectories = trajectories) +end + +# ── Checkpointing (exact reload) ────────────────────────────────────────────── + +""" + battery_checkpoint(policy, de; case, process, extra=Dict()) -> Dict + +Assemble a checkpoint dictionary that identifies everything needed to reload the +policy and reproduce its outputs: the Flux model state, the architecture and +target mode, the reporting/look-ahead horizons, the train/eval seeds, and the +case-manifest / MATPOWER-source / load-process hashes. +""" +function battery_checkpoint(policy::BatteryReachablePolicy, de::BatteryTSDDRProblem; + case::BatteryCase, process::LoadProcess, + extra::AbstractDict = Dict{String,Any}()) + layers = _encoder_layer_sizes(policy) + combiner = _combiner_layer_sizes(policy) + doc = Dict{String,Any}( + "schema" => "battery_tsddr_checkpoint/1", + "flux_state" => Flux.state(policy), + "architecture" => Dict{String,Any}( + "policy" => "BatteryReachablePolicy", + "n_uncertainty" => policy.n_uncertainty, + "nbat" => policy.nbat, + "encoder_layers" => layers, + "combiner_layers" => combiner, + "target_mode" => String(de.mode), + "dt" => de.dt, + # Target-head activation and its safe upper margin: the normalized + # target never reaches exactly 1 (interior-point degeneracy under a + # strict equality). Checkpoints are only comparable across runs with + # the same activation. + "activation" => string(_activation_name(policy)), + "safe_upper_margin" => 1e-3, + ), + "horizons" => Dict{String,Any}( + "reporting_horizon" => de.reporting_horizon, + "lookahead" => de.lookahead, + "horizon" => de.horizon, + ), + "seeds" => Dict{String,Any}( + "train_seed" => process.train_seed, + "eval_seed" => process.eval_seed, + ), + "hashes" => Dict{String,Any}( + "case_manifest_content_hash" => manifest_hash(case), + "matpower_sha256" => case.parse_meta.matpower_sha256, + "load_process_hash" => process_hash(process), + ), + # Training-only target-penalty coefficients (soft mode) and the price used + # by the physical two-sided active-recourse cost. + "target_penalty" => Dict{String,Any}("rho1" => de.rho1, "rho2" => de.rho2), + "active_recourse_cost_per_mwh" => de.active_recourse_cost_per_mwh, + ) + for (k, v) in extra + doc[k] = v + end + return doc +end + +""" + save_checkpoint(path, policy, de; case, process, extra=Dict()) -> String + +Serialize a [`battery_checkpoint`](@ref) to `path` with the `Serialization` +stdlib. Returns `path`. +""" +function save_checkpoint(path::AbstractString, policy::BatteryReachablePolicy, + de::BatteryTSDDRProblem; case::BatteryCase, process::LoadProcess, + extra::AbstractDict = Dict{String,Any}()) + doc = battery_checkpoint(policy, de; case = case, process = process, extra = extra) + open(io -> Serialization.serialize(io, doc), path, "w") + return path +end + +""" + load_checkpoint(path, case, process; dt=nothing, float_type=Float32) + -> (policy, meta) + +Rebuild a [`BatteryReachablePolicy`] from a checkpoint written by +[`save_checkpoint`](@ref): construct a fresh policy with the recorded architecture +(for `case`/`process`), then load the saved Flux state into it exactly. Verifies +that the recorded case-manifest, MATPOWER-source, and load-process hashes match +the supplied `case`/`process`. Reloading reproduces the policy's outputs exactly +on any fixed CPU input. `meta` is the checkpoint dict (without the raw state). +""" +function load_checkpoint(path::AbstractString, case::BatteryCase, process::LoadProcess; + dt::Union{Nothing,Real} = nothing, + float_type::Type{<:AbstractFloat} = Float32) + doc = open(Serialization.deserialize, path) + arch = doc["architecture"] + h = doc["hashes"] + manifest_hash(case) == String(h["case_manifest_content_hash"]) || + error("checkpoint case-manifest hash mismatch") + case.parse_meta.matpower_sha256 == String(h["matpower_sha256"]) || + error("checkpoint MATPOWER-source hash mismatch") + process_hash(process) == String(h["load_process_hash"]) || + error("checkpoint load-process hash mismatch") + + Δt = dt === nothing ? Float64(arch["dt"]) : Float64(dt) + policy = battery_reachable_policy(case, process; + dt = Δt, + layers = Int.(arch["encoder_layers"]), + combiner_layers = Int.(arch["combiner_layers"]), + activation = _activation_from_name(String(arch["activation"])), + float_type = float_type) + load_battery_policy!(policy, doc["flux_state"]) + meta = Dict(k => v for (k, v) in doc if k != "flux_state") + return policy, meta +end + +# Target-head activation recorded in (and restored from) the checkpoint. The +# head applies the bounded activation at every layer, so reading it off the +# output layer identifies the whole head. +function _activation_name(policy::BatteryReachablePolicy) + comb = policy.combiner + layer = comb isa Flux.Dense ? comb : comb.layers[end] + return layer.σ +end + +# Resolve a recorded activation name back to the callable. +function _activation_from_name(name::AbstractString) + name == string(stretchedsigmoid) && return stretchedsigmoid + name == string(hardsigmoidsafe) && return hardsigmoidsafe + error("unknown target activation \"$name\" in checkpoint; expected " * + "$(string(stretchedsigmoid)) or $(string(hardsigmoidsafe))") +end + +# Encoder/combiner size introspection for the checkpoint architecture record. +# The encoder is a Chain of recurrent layers; the combiner is a Dense or a Chain +# of Denses. Sizes are read from the weight matrices so a reloaded policy is +# rebuilt with an identical structure. +function _encoder_layer_sizes(policy::BatteryReachablePolicy) + sizes = Int[] + for layer in policy.encoder.layers + cell = DecisionRulesExa._as_cell(layer) + # LSTM/GRU/RNN cell Wh is (factor·hidden) × hidden; hidden = out features. + push!(sizes, size(cell.Wh, 2)) + end + return sizes +end + +function _combiner_layer_sizes(policy::BatteryReachablePolicy) + comb = policy.combiner + comb isa Flux.Dense && return Int[] # single linear head → no hidden layers + sizes = Int[] + layers = comb.layers + for i in 1:length(layers) - 1 # all but the output layer are "hidden" + push!(sizes, size(layers[i].weight, 1)) + end + return sizes +end + +# ── Trajectory serialization ────────────────────────────────────────────────── + +""" + write_trajectory(path, evaluations; meta=Dict()) -> String + +Write a compact, machine-readable JSON trajectory file. `evaluations` is a vector +whose entries are either `nothing` (a failed path) or the NamedTuple returned by +[`evaluate_battery_policy`](@ref). Each successful path stores its per-stage +records (stage & atom indices, battery SoC in/out, targets, charge/discharge, +generator and physical costs, target penalty, solver status, residuals) plus its +reporting / look-ahead physical costs. Returns `path`. +""" +function write_trajectory(path::AbstractString, evaluations::AbstractVector; meta::AbstractDict = Dict{String,Any}()) + paths = Any[] + for (p, ev) in enumerate(evaluations) + if ev === nothing + push!(paths, Dict{String,Any}("path" => p, "status" => "failed")) + else + push!(paths, Dict{String,Any}( + "path" => p, + "status" => "ok", + "reporting_physical_cost" => ev.reporting_physical_cost, + "lookahead_physical_cost" => ev.lookahead_physical_cost, + "physical_operating_cost" => ev.physical_operating_cost, + "generator_cost" => ev.generator_cost, + "battery_throughput_cost" => ev.battery_throughput_cost, + "active_recourse_cost" => ev.active_recourse_cost, + "active_deficit_energy_mwh" => ev.active_deficit_energy_mwh, + "active_surplus_energy_mwh" => ev.active_surplus_energy_mwh, + "total_active_recourse_energy_mwh" => ev.total_active_recourse_energy_mwh, + "max_active_deficit_pu" => ev.max_active_deficit_pu, + "max_active_surplus_pu" => ev.max_active_surplus_pu, + "raw_active_recourse_cost" => ev.raw_active_recourse_cost, + "active_recourse_projection_correction" => ev.active_recourse_projection_correction, + "maximum_active_recourse_lower_bound_violation_pu" => ev.maximum_active_recourse_lower_bound_violation_pu, + "target_penalty" => ev.target_penalty, + "target_violation" => ev.target_violation, + "final_soc" => ev.final_soc, + "stages" => ev.trajectory, + )) + end + end + doc = Dict{String,Any}("schema" => "battery_tsddr_trajectory/1", "paths" => paths) + for (k, v) in meta + doc[k] = v + end + open(io -> JSON.print(io, doc, 2), path, "w") + return path +end diff --git a/examples/BatteryStorageOPF/src/battery_tsddr.jl b/examples/BatteryStorageOPF/src/battery_tsddr.jl new file mode 100644 index 0000000..fb94c1e --- /dev/null +++ b/examples/BatteryStorageOPF/src/battery_tsddr.jl @@ -0,0 +1,912 @@ +# battery_tsddr.jl +# +# Phase-2 OPERATIONAL / stochastic battery AC-OPF used by TS-DDR training and +# evaluation. Canonical spec: docs/src/casestudies/battery_storage_opf.md. +# +# It EXTENDS the accepted Phase-1 model through the SHARED ACP blocks in +# acp_core.jl (single source of truth) and adds exactly two things: +# +# 1. TWO-SIDED ABSOLUTE ACTIVE RECOURSE — a nonnegative, UNBOUNDED-above pair +# at EVERY bus, +# d⁺[t,i] ≥ 0 (active deficit / injection), +# d⁻[t,i] ≥ 0 (active surplus / absorption), +# entering ONLY the active balance as `− d⁺ + d⁻`. It is an artificial +# active-balance recourse, NOT curtailed customer load: d⁺ may exceed local +# demand and may be positive where p^d = 0. Reactive KCL stays a HARD +# equality: there is NO reactive slack. The pair is priced at VOLL and is +# INCLUDED in physical operating cost. It is an operational safety valve — +# an accepted run leaves BOTH directions at zero within tolerance. +# +# 2. TARGET-STATE PROJECTION on the battery SoC, in one of two modes: +# strict : ê_{t+1,b} − e_{t+1,b} = 0 (no slack, no penalty) +# soft : ê_{t+1,b} − e_{t+1,b} − δ⁺ + δ⁻ = 0, δ± ≥ 0 +# with a documented TRAINING-ONLY penalty ρ1·Σ(δ⁺+δ⁻) + (ρ2/2)·Σ((δ⁺)²+(δ⁻)²) +# that NEVER enters physical operating cost. +# +# Active recourse and target slack are different mechanisms with different names +# and different roles. Target constraints are added LAST so their multipliers +# occupy a contiguous slice (`target_con_range`). +# +# Demand realization follows the canonical pattern: the uncertainty parameter +# `p_w` carries the observed atom features (it appears in no constraint), and +# `set_realized_demand!` / `prepare_solve!` write the realized per-bus demand +# p^d = p^{d,0}·h_t·L·R into the SAME `p_pd`/`p_qd` parameters the shared ACP +# balance uses — so the stochastic and deterministic builders solve identical +# equations for identical realized demand. +# +# Generator prices, pmin/pmax, qmin/qmax, branch limits, voltage limits, +# admittances, and topology are the original PGLib values. Nothing here scales +# generator capacity. + +using ExaModels +using MadNLP +using NLPModels +import DecisionRulesExa: target_multipliers, prepare_solve! + +# w-vector index helpers (stage-major, nw = 1 + nregion): entry 1 of a stage +# block is the realized system multiplier s_t; entry 1+r is region r's multiplier. +@inline _widx_s(nw, t) = (t - 1) * nw + 1 +@inline _widx_r(nw, t, r) = (t - 1) * nw + 1 + r + +""" + DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH + +Default price of the two-sided active nodal recourse, `10_000.0` USD/MWh +(canonical spec — the value of lost load). The stage cost is +`c · S^base · Δt · Σ_i (d⁺_{t,i} + d⁻_{t,i})`. +""" +const DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH = 10_000.0 + +""" + ACTIVE_RECOURSE_LB_TOL_PU + +Declared **per-variable** lower-bound tolerance (pu) for the nominally nonnegative +primal quantities — the two-sided active recourse `d⁺`/`d⁻` and the soft target +slacks `δ⁺`/`δ⁻`. An interior-point solve may return an individual variable a hair +below its zero bound (`~1e-8` pu observed; the loosest solver bound tolerance used +anywhere here is `1e-6`). Reporting VALIDATES that no single value is further below +zero than this tolerance (failing loudly otherwise — that would signal a real +defect, not noise), then PROJECTS the tolerated noise elementwise with `max(·, 0)` +before computing any public physical quantity. This value is `1e-5` pu: 10× the +loosest solver bound tolerance and ≥100× below the smallest material recourse, so +it never trips on noise yet catches a genuinely negative value. +""" +const ACTIVE_RECOURSE_LB_TOL_PU = 1e-5 + +""" + BatteryTSDDRProblem + +Operational target-constrained ExaModels problem for a [`BatteryCase`] and its +[`LoadProcess`]. Usable by `train_tsddr`, `simulate_tsddr`, and (with +`horizon = 1`) `rollout_tsddr`. + +# Key fields +- `p_x0`: initial battery SoC parameter (length `nBat`). +- `p_w`: observed per-stage uncertainty features (length `horizon·nw`); read by + `prepare_solve!`/`set_realized_demand!`, referenced by no constraint. +- `p_pd`, `p_qd`: realized per-bus active/reactive demand actually used by the + shared ACP balance (length `horizon·nBus` each). +- `p_target`: target next-SoC parameter (length `horizon·nBat`). +- `mode`: `:strict` or `:soft`. +- `rho1`, `rho2`: soft target-penalty L1/L2 coefficients (training-only). +- `active_recourse_cost_per_mwh`, `baseMVA`: recourse price and power base. +- `reporting_horizon`, `lookahead`, `dt`: terminal treatment and stage duration. +- `target_con_range`: contiguous multiplier slice of the target constraints. +- `realized_pd`, `realized_qd`: `[T × nBus]` buffers holding the demand most + recently written into `p_pd`/`p_qd` (used by the cost decomposition). +""" +struct BatteryTSDDRProblem + core + model + p_x0 + p_w + p_pd + p_qd + p_target + nBus::Int + nGen::Int + nBranch::Int + nBat::Int + nw::Int + horizon::Int + reporting_horizon::Int + lookahead::Int + dt::Float64 + mode::Symbol + rho1::Float64 + rho2::Float64 + active_recourse_cost_per_mwh::Float64 + baseMVA::Float64 + gen_c2::Vector{Float64} + gen_c1::Vector{Float64} + gen_c0::Vector{Float64} + cycle_coeffs::Vector{Float64} + target_con_range::UnitRange{Int} + realized_pd::Matrix{Float64} + realized_qd::Matrix{Float64} + case::BatteryCase + process::LoadProcess + float_type::Type +end + +nbattery(prob::BatteryTSDDRProblem) = prob.nBat + +""" + build_battery_tsddr_de(case, process; reporting_horizon, lookahead=0, + mode=:soft, rho1=0.0, rho2=:auto, + active_recourse_cost_per_mwh=DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH, + stage_hours=1.0, backend=nothing, float_type=Float64) + -> BatteryTSDDRProblem + +Build the `T = reporting_horizon + lookahead` stage operational problem. + +# Keywords +- `reporting_horizon::Int`: stages whose physical cost is the headline metric. +- `lookahead::Int = 0`: look-ahead buffer appended after the reporting window. +- `mode::Symbol = :strict`: `:strict` (hard `ê=e`, no slack/penalty — the + PRIMARY / default operational and training mode; the two-sided active nodal + recourse gives it complete recourse so it always solves) or `:soft` (split + slacks + training-only penalty — a DIAGNOSTIC/fallback only). +- `rho1`, `rho2`: soft target-penalty L1 and L2 coefficients. `rho2 = :auto` + uses `2·max(c1,c2)` over generators (the project-standard auto scale). +- `active_recourse_cost_per_mwh`: recourse price in USD/MWh (default 10 000). +- `stage_hours::Real = 1.0`: Δt in hours. +- `backend`: `nothing` (CPU) or a CUDA backend (GPU) — same equations either way. + +Generator and network data are never modified. +""" +function _build_battery_problem(case::BatteryCase, process::LoadProcess; + reporting_horizon::Int, + lookahead::Int = 0, + mode::Symbol = :soft, + rho1::Real = 0.0, + rho2::Union{Real,Symbol} = :auto, + active_recourse_cost_per_mwh::Real = DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH, + allow_active_recourse::Bool = true, + stage_hours::Real = 1.0, + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64) + reporting_horizon >= 1 || error("reporting_horizon must be ≥ 1; got $reporting_horizon") + lookahead >= 0 || error("lookahead must be ≥ 0; got $lookahead") + mode in (:strict, :soft, :none) || + error("mode must be :strict, :soft, or :none (internal); got :$mode") + T = reporting_horizon + lookahead + dt = Float64(stage_hours) + validate_stage_hours(case, dt) + (isfinite(active_recourse_cost_per_mwh) && active_recourse_cost_per_mwh >= 0) || + error("active_recourse_cost_per_mwh must be finite and ≥ 0; got $active_recourse_cost_per_mwh") + nd = case.network + length(process.region_of_bus) == nbus(nd) || + error("process.region_of_bus length $(length(process.region_of_bus)) ≠ nbus=$(nbus(nd)); " * + "the LoadProcess must be built for THIS case's network") + + nBus = nbus(nd); nGen = ngen(nd); nBranch = nbranch(nd) + nBat = length(case.batteries) + nw = n_uncertainty(process) + ρ1 = Float64(rho1) + ρ2 = Float64(rho2 === :auto ? + 2.0 * maximum(max(g.cost1, g.cost2) for g in nd.gens; init = 0.0) : rho2) + (isfinite(ρ1) && ρ1 >= 0) || error("rho1 must be finite and ≥ 0; got $ρ1") + (isfinite(ρ2) && ρ2 >= 0) || error("rho2 must be finite and ≥ 0; got $ρ2") + + core = ExaModels.ExaCore(float_type; backend = backend, concrete = Val(true)) + + # ── Variables: shared ACP + battery, then active recourse, then soft slacks ─ + core, v = add_acp_variables!(core, nd, T, float_type) + core, bat = add_battery_variables!(core, case, T, float_type) + + # Two-sided active nodal recourse, both nonnegative and UNBOUNDED above at + # EVERY bus (active-only). They enter the active balance as `− d⁺ + d⁻`: + # active_deficit = d⁺ (injection, covers a local active shortfall — e.g. + # the power to CHARGE a battery at a congested bus) + # active_surplus = d⁻ (absorption, absorbs a local active excess — e.g. + # the power a battery is forced to DISCHARGE into a + # bus whose outgoing branches are saturated) + # They are an artificial active-balance recourse, NOT curtailed customer + # load: d⁺ may exceed local demand and may be positive where p^d = 0. Together + # they give relatively complete recourse, so the stage subproblem is feasible + # for every incoming state and every reachable target. Both are priced at + # VOLL, so an accepted solution leaves both at ~0. `allow_active_recourse = + # false` fixes BOTH blocks to zero (uvar = 0) — used by the Phase-1 parity + # test to recover the hard-balance deterministic model exactly. + rec_ub = allow_active_recourse ? fill(float_type(Inf), T * nBus) : zeros(float_type, T * nBus) + core, active_deficit = ExaModels.add_var(core, T * nBus; lvar = float_type(0), uvar = rec_ub) + core, active_surplus = ExaModels.add_var(core, T * nBus; lvar = float_type(0), uvar = rec_ub) + + slack_pos = nothing; slack_neg = nothing + if mode === :soft && nBat > 0 + core, slack_pos = ExaModels.add_var(core, T * nBat; lvar = float_type(0)) + core, slack_neg = ExaModels.add_var(core, T * nBat; lvar = float_type(0)) + end + + # ── Parameters ──────────────────────────────────────────────────────────── + # Realized demand actually consumed by the shared ACP balance; initialized to + # the deterministic base shape at unit factors and overwritten per solve. + init_pd = Float64[]; init_qd = Float64[] + for t in 1:T + h = process.base_shape[((t - 1) % process.period) + 1] + append!(init_pd, h .* nd.bus_pd) + append!(init_qd, h .* nd.bus_qd) + end + core, p_pd = ExaModels.add_par(core, float_type.(init_pd)) + core, p_qd = ExaModels.add_par(core, float_type.(init_qd)) + # Observed uncertainty features (system + regional multipliers). Referenced + # by NO constraint; consumed by set_realized_demand!/prepare_solve!. + w0 = Float64[] + for t in 1:T + h = process.base_shape[((t - 1) % process.period) + 1] + push!(w0, h) + append!(w0, ones(Float64, process.nregion)) + end + core, p_w = ExaModels.add_par(core, float_type.(w0)) + core, p_x0 = ExaModels.add_par(core, float_type.([b.e_init for b in case.batteries])) + # Target parameter exists ONLY for the target modes. The targetless + # diagnostic (`mode === :none`) carries no target data at all. + p_target = nothing + if mode !== :none + core, p_target = ExaModels.add_par(core, + float_type.([b.e_init for _ in 1:T for b in case.batteries])) + end + + # ── Objective ───────────────────────────────────────────────────────────── + core = add_generator_cost!(core, v, nd, T, dt, float_type) + core = add_cycle_cost!(core, bat, case, nd, T, dt, float_type) + + # Two-sided active-recourse cost: c · S^base · Δt · Σ (d⁺ + d⁻) (physical + # operating cost). Both are recourse POWER in pu, so the cost is linear in + # them directly (no `p^d` factor); pricing both keeps an accepted solution at + # ~0. + rec_coeff = float_type(active_recourse_cost_per_mwh * nd.baseMVA * dt) + if rec_coeff > 0 + rec_items = [(idx = _bidx(nBus, t, b),) for t in 1:T for b in 1:nBus] + core, _ = ExaModels.add_obj(core, + rec_coeff * (active_deficit[it.idx] + active_surplus[it.idx]) for it in rec_items) + end + + # Soft TRAINING-ONLY target penalty ρ1·Σ(δ⁺+δ⁻) + (ρ2/2)·Σ((δ⁺)²+(δ⁻)²). + if mode === :soft && nBat > 0 + pen_items = [(idx = (t - 1) * nBat + k,) for t in 1:T for k in 1:nBat] + if ρ1 > 0 + core, _ = ExaModels.add_obj(core, + float_type(ρ1) * (slack_pos[it.idx] + slack_neg[it.idx]) for it in pen_items) + end + if ρ2 > 0 + core, _ = ExaModels.add_obj(core, + float_type(ρ2 / 2) * (slack_pos[it.idx]^2 + slack_neg[it.idx]^2) + for it in pen_items) + end + end + + # ── Constraints (SHARED blocks; targets appended LAST) ──────────────────── + core = add_acp_network_constraints!(core, v, nd, T, float_type) + core = add_nodal_balance!(core, v, bat, nd, case, T, float_type, p_pd, p_qd; + active_deficit = active_deficit, active_surplus = active_surplus) + core = add_battery_dynamics!(core, bat, case, T, dt, float_type, p_x0) + + n_before = acp_constraint_count(nd, case, T) + if mode !== :none && nBat > 0 + tgt_items = [(pt = (t - 1) * nBat + k, en = _eidx(nBat, t + 1, k), + sl = (t - 1) * nBat + k) for t in 1:T for k in 1:nBat] + if mode === :strict + core, _ = ExaModels.add_con(core, p_target[it.pt] - bat.e[it.en] for it in tgt_items) + else + core, _ = ExaModels.add_con(core, + p_target[it.pt] - bat.e[it.en] - slack_pos[it.sl] + slack_neg[it.sl] + for it in tgt_items) + end + end + # Targetless diagnostic: EMPTY multiplier slice (no target constraints exist). + target_range = mode === :none ? ((n_before + 1):n_before) : + ((n_before + 1):(n_before + T * nBat)) + + gen_c2 = Float64[g.cost2 * dt for g in nd.gens] + gen_c1 = Float64[g.cost1 * dt for g in nd.gens] + gen_c0 = Float64[g.cost0 * dt for g in nd.gens] + cycle_coeffs = Float64[b.cycle_cost_per_mwh * nd.baseMVA * dt for b in case.batteries] + + model = ExaModels.ExaModel(core) + prob = BatteryTSDDRProblem(core, model, p_x0, p_w, p_pd, p_qd, p_target, + nBus, nGen, nBranch, nBat, nw, T, + reporting_horizon, lookahead, dt, mode, ρ1, ρ2, + Float64(active_recourse_cost_per_mwh), Float64(nd.baseMVA), + gen_c2, gen_c1, gen_c0, cycle_coeffs, target_range, + permutedims(reshape(init_pd, nBus, T)), + permutedims(reshape(init_qd, nBus, T)), + case, process, float_type) + return prob +end + +""" + build_battery_stage_problem(case, process; mode=:strict, kwargs...) -> BatteryTSDDRProblem + +Single-stage (`reporting_horizon = 1`, `lookahead = 0`) operational problem for +stage-wise `rollout_tsddr` evaluation. The realized next SoC is `e[:, 2]`. +Defaults to the primary `:strict` mode. +""" +build_battery_stage_problem(case::BatteryCase, process::LoadProcess; mode::Symbol = :strict, kwargs...) = + build_battery_tsddr_de(case, process; reporting_horizon = 1, lookahead = 0, mode = mode, kwargs...) + +""" + build_battery_tsddr_de(case, process; mode = :strict, kwargs...) -> BatteryTSDDRProblem + +PRODUCTION builder. `mode` is `:strict` (default, primary) or `:soft` +(diagnostic/fallback) ONLY; the targetless diagnostic is not a production target +mode and is built by [`build_targetless_diagnostic_de`](@ref). +""" +function build_battery_tsddr_de(case::BatteryCase, process::LoadProcess; + mode::Symbol = :strict, kwargs...) + mode in (:strict, :soft) || + error("mode must be :strict or :soft; got :$mode. The targetless diagnostic " * + "is built with build_targetless_diagnostic_de and is not a production mode.") + return _build_battery_problem(case, process; mode = mode, kwargs...) +end + +""" + is_targetless(prob) -> Bool + +`true` for a targetless PHYSICAL DIAGNOSTIC problem, which carries no target +constraints, no target slack variables, and no target penalty. +""" +is_targetless(prob::BatteryTSDDRProblem) = prob.mode === :none + +""" + build_targetless_diagnostic_de(case, process; reporting_horizon, lookahead=0, + allow_active_recourse=true, kwargs...) -> BatteryTSDDRProblem + +Build the GENUINELY TARGETLESS physical diagnostic. + +The model contains **no target constraints, no target slack variables, and no +target penalty** — not a zero-weight target, but no target data at all +(`p_target === nothing`, `target_con_range` empty). Batteries are freely +dispatchable subject only to their own dynamics and bounds. Everything else — +the shared ACP equations, the realized demand, the two-sided active recourse, +generator and throughput costs, and the battery equations — is identical to the +operational model, and is built from the SAME shared blocks (nothing is copied). + +Its objective is exactly `generator + battery throughput + active-recourse cost`. + +`allow_active_recourse = false` fixes both recourse blocks (d⁺ and d⁻) to zero by +bounds, giving the recourse-forbidden companion that shares every physical +equation and datum. + +Interpretation rules (see the README): an ACCEPTED recourse-forbidden solve +proves a zero-recourse feasible point was FOUND; a FAILED recourse-forbidden +solve is INCONCLUSIVE and never proves infeasibility. Results are local optima of +a nonconvex ACP unless a global certificate exists. + +This is a diagnostic only: `train_tsddr`-style training, target-multiplier +extraction, and target rollout all REJECT it. +""" +build_targetless_diagnostic_de(case::BatteryCase, process::LoadProcess; kwargs...) = + _build_battery_problem(case, process; mode = :none, kwargs...) + +# ── Deterministic, target-consistent primal starts (strict-mode start repair) ── +# +# Only the primal STARTING POINT is affected. No tolerance, iteration budget, +# equation, bound, generator, or network value is touched, and among accepted +# starts the FIRST in the fixed order is taken (never the cheapest). + +""" + variable_offsets(prob) -> NamedTuple + +Zero-based offsets of each variable block in the flat solution/start vector, in +creation order: `va, vm, pg, qg, p_fr, q_fr, p_to, q_to, p_ch, p_dis, e, +active_deficit, active_surplus, [slack_pos, slack_neg]`. +""" +function variable_offsets(prob::BatteryTSDDRProblem) + T = prob.horizon; nB = prob.nBus; nG = prob.nGen; nBR = prob.nBranch; nK = prob.nBat + o = 0 + va = o; o += T*nB + vm = o; o += T*nB + pg = o; o += T*nG + qg = o; o += T*nG + p_fr = o; o += T*nBR + q_fr = o; o += T*nBR + p_to = o; o += T*nBR + q_to = o; o += T*nBR + p_ch = o; o += T*nK + p_dis = o; o += T*nK + e = o; o += (T+1)*nK + active_deficit = o; o += T*nB + active_surplus = o; o += T*nB + slack_pos = o + slack_neg = prob.mode === :soft ? o + T*nK : o + return (va = va, vm = vm, pg = pg, qg = qg, p_fr = p_fr, q_fr = q_fr, + p_to = p_to, q_to = q_to, p_ch = p_ch, p_dis = p_dis, e = e, + active_deficit = active_deficit, active_surplus = active_surplus, + slack_pos = slack_pos, slack_neg = slack_neg) +end + +""" + target_consistent_start!(prob, e_prev, target) -> x0 + +Write a deterministic, TARGET-CONSISTENT battery start into the model's primal +starting point and return it. For each battery, with + +``` +Δe = target − (1 − σ·Δt)·e_prev, +``` + +the start is the exact charge/discharge that realizes `Δe`: + +``` +Δe ≥ 0 : p_charge = Δe/(η_ch·Δt), p_discharge = 0 +Δe < 0 : p_charge = 0, p_discharge = −Δe·η_dis/Δt +``` + +clipped to the battery power bounds (the reachability interval guarantees the +clip is inactive for a reachable target). The SoC trajectory is started at the +implied path (`e_1 = e_prev`, `e_2 = target`). Only the START is set. +""" +function target_consistent_start!(prob::BatteryTSDDRProblem, e_prev::AbstractVector, + target::AbstractVector) + x0 = NLPModels.get_x0(prob.model) + off = variable_offsets(prob) + T = prob.horizon; nK = prob.nBat + nK == 0 && return x0 + length(e_prev) == nK || error("e_prev length must be nBat=$nK") + length(target) == T * nK || error("target length must be horizon·nBat=$(T*nK)") + prev = Float64.(collect(e_prev)) + for t in 1:T + for (k, b) in enumerate(prob.case.batteries) + tgt = Float64(target[(t-1)*nK + k]) + Δe = tgt - (1 - b.sigma * prob.dt) * prev[k] + pch, pdis = if Δe >= 0 + (min(Δe / (b.eta_ch * prob.dt), b.p_charge_max), 0.0) + else + (0.0, min(-Δe * b.eta_dis / prob.dt, b.p_discharge_max)) + end + x0[off.p_ch + (t-1)*nK + k] = pch + x0[off.p_dis + (t-1)*nK + k] = pdis + x0[off.e + t*nK + k] = tgt # e index (t+1) is block t (0-based) + prev[k] = tgt + end + end + for k in 1:nK # e_1 = e_prev + x0[off.e + k] = Float64(e_prev[k]) + end + return x0 +end + +""" + seed_start_from_solution!(prob, x_src) -> x0 + +Copy a previously ACCEPTED primal point into this problem's starting point. +Blocks are copied by the shared variable order for the overlapping length, so a +soft/targetless-diagnostic solution can seed a strict solve of the same size. +""" +function seed_start_from_solution!(prob::BatteryTSDDRProblem, x_src::AbstractVector) + x0 = NLPModels.get_x0(prob.model) + n = min(length(x0), length(x_src)) + @inbounds for i in 1:n + x0[i] = Float64(x_src[i]) + end + return x0 +end + +""" + reset_flat_start!(prob) -> x0 + +Restore the model's default flat start (`vm = 1`, everything else 0, SoC at +`e_init`) — the first entry of the deterministic start sequence. +""" +function reset_flat_start!(prob::BatteryTSDDRProblem) + x0 = NLPModels.get_x0(prob.model) + off = variable_offsets(prob) + T = prob.horizon; nB = prob.nBus; nK = prob.nBat + fill!(x0, 0.0) + @inbounds for i in 1:(T*nB) + x0[off.vm + i] = 1.0 + end + @inbounds for t in 0:T, (k, b) in enumerate(prob.case.batteries) + x0[off.e + t*nK + k] = b.e_init + end + return x0 +end + +# ── Demand realization + parameter setters ──────────────────────────────────── + +""" + set_realized_demand!(prob, w_flat) -> prob + +Compute the realized per-bus demand from the observed uncertainty features and +write it into the `p_pd`/`p_qd` parameters used by the shared ACP balance: + +`p^d_{t,i} = p^{d,0}_i · s_t · R_{r(i),t}` and `q^d_{t,i} = q^{d,0}_i · s_t · R_{r(i),t}`, + +where `s_t` already contains the deterministic daily shape times the atom's +system factor. The SAME multiplier scales active and reactive demand, preserving +every bus's base power factor. Also refreshes the `realized_pd`/`realized_qd` +buffers used by the cost decomposition. +""" +function set_realized_demand!(prob::BatteryTSDDRProblem, w_flat::AbstractVector) + T = prob.horizon; nBus = prob.nBus; nw = prob.nw + w = Float64.(vec(Array(w_flat))) + length(w) == T * nw || + error("w_flat length must be horizon·nw=$(T*nw); got $(length(w))") + nd = prob.case.network + region = prob.process.region_of_bus + @inbounds for t in 1:T + s_t = w[_widx_s(nw, t)] + for b in 1:nBus + f = s_t * w[_widx_r(nw, t, region[b])] + prob.realized_pd[t, b] = nd.bus_pd[b] * f + prob.realized_qd[t, b] = nd.bus_qd[b] * f + end + end + ExaModels.set_parameter!(prob.core, prob.p_pd, + prob.float_type.([prob.realized_pd[t, b] for t in 1:T for b in 1:nBus])) + ExaModels.set_parameter!(prob.core, prob.p_qd, + prob.float_type.([prob.realized_qd[t, b] for t in 1:T for b in 1:nBus])) + return prob +end + +""" + set_tsddr_uncertainty!(prob, w_flat) -> prob + +Set the observed uncertainty parameter AND the realized demand it implies. +""" +function set_tsddr_uncertainty!(prob::BatteryTSDDRProblem, w_flat::AbstractVector) + length(w_flat) == prob.horizon * prob.nw || + error("w_flat length must be horizon·nw=$(prob.horizon*prob.nw); got $(length(w_flat))") + ExaModels.set_parameter!(prob.core, prob.p_w, prob.float_type.(collect(w_flat))) + set_realized_demand!(prob, w_flat) + return prob +end + +""" + set_tsddr_initial_soc!(prob, e0) -> prob + +Set the initial battery-SoC parameter (length `nBat`, pu·h). +""" +function set_tsddr_initial_soc!(prob::BatteryTSDDRProblem, e0::AbstractVector) + length(e0) == prob.nBat || error("e0 length must be nBat=$(prob.nBat); got $(length(e0))") + ExaModels.set_parameter!(prob.core, prob.p_x0, prob.float_type.(collect(e0))) + return prob +end + +""" + set_tsddr_targets!(prob, xhat) -> prob + +Set the target next-SoC parameter (length `horizon·nBat`, stage-major). +""" +function set_tsddr_targets!(prob::BatteryTSDDRProblem, xhat::AbstractVector) + is_targetless(prob) && + error("this is a TARGETLESS diagnostic problem: it has no target parameter, " * + "no target constraints, and no target slacks; targets cannot be set") + length(xhat) == prob.horizon * prob.nBat || + error("xhat length must be horizon·nBat=$(prob.horizon*prob.nBat); got $(length(xhat))") + ExaModels.set_parameter!(prob.core, prob.p_target, prob.float_type.(collect(xhat))) + return prob +end + +""" + prepare_solve!(prob::BatteryTSDDRProblem, init_state, w_flat, xhat_flat) -> Nothing + +Pre-solve hook used by the DecisionRulesExa training loop: after the loop writes +`p_x0`, `p_w`, and `p_target`, this realizes the demand implied by `w_flat` into +`p_pd`/`p_qd` (the canonical pattern — the uncertainty parameter itself appears +in no constraint). +""" +function prepare_solve!(prob::BatteryTSDDRProblem, init_state, w_flat, xhat_flat) + is_targetless(prob) && + error("a TARGETLESS diagnostic problem cannot be used for TS-DDR training/rollout") + set_realized_demand!(prob, w_flat) + return nothing +end + +# ── Envelope multipliers ────────────────────────────────────────────────────── + +""" + target_multipliers(prob::BatteryTSDDRProblem, result) -> λ + +Multipliers of the target constraints, `result.multipliers[target_con_range]`. +With the orientation `ê_{t+1} − e_{t+1} (− δ⁺ + δ⁻) = 0`, these are the envelope +derivatives `∂Q/∂ê` used as the policy-gradient signal, up to the solver's +documented dual convention — the sign AND magnitude are finite-difference tested +rather than assumed. +""" +function target_multipliers(prob::BatteryTSDDRProblem, result) + is_targetless(prob) && + error("target_multipliers is undefined for a TARGETLESS diagnostic problem " * + "(it has no target constraints)") + return result.multipliers[prob.target_con_range] +end + +# ── Solution extraction + cost decomposition ────────────────────────────────── + +""" + tsddr_solution(prob, result) -> NamedTuple + +Reshape the flat solution into named matrices (columns are stages): `va`,`vm` +(nBus), `pg`,`qg` (nGen), `p_fr`,`q_fr`,`p_to`,`q_to` (nBranch), `p_ch`,`p_dis` +(nBat), `soc` (nBat×(T+1)), `active_deficit_pu`/`active_surplus_pu` (nBus×T — the +per-bus two-sided active-recourse power in pu), the soft-mode +`target_slack_pos`/`target_slack_neg` (nBat×T; zeros in strict mode), and derived +`p_bat = p_dis − p_ch`. +""" +function tsddr_solution(prob::BatteryTSDDRProblem, result) + T = prob.horizon; nB = prob.nBus; nG = prob.nGen; nBR = prob.nBranch; nK = prob.nBat + sol = Array(result.solution) + off = 0 + take(n, m) = (v = reshape(sol[off .+ (1:n*m)], n, m); off += n*m; v) + va = take(nB, T) + vm = take(nB, T) + pg = take(nG, T) + qg = take(nG, T) + p_fr = take(nBR, T) + q_fr = take(nBR, T) + p_to = take(nBR, T) + q_to = take(nBR, T) + if nK > 0 + p_ch = take(nK, T) + p_dis = take(nK, T) + soc = take(nK, T + 1) + else + p_ch = zeros(eltype(sol), 0, T) + p_dis = zeros(eltype(sol), 0, T) + soc = zeros(eltype(sol), 0, T + 1) + end + active_deficit_pu = take(nB, T) + active_surplus_pu = take(nB, T) + if prob.mode === :soft && nK > 0 + slack_pos = take(nK, T) + slack_neg = take(nK, T) + else + slack_pos = zeros(eltype(sol), nK, T) + slack_neg = zeros(eltype(sol), nK, T) + end + return (va = va, vm = vm, pg = pg, qg = qg, + p_fr = p_fr, q_fr = q_fr, p_to = p_to, q_to = q_to, + p_ch = p_ch, p_dis = p_dis, soc = soc, + active_deficit_pu = active_deficit_pu, active_surplus_pu = active_surplus_pu, + target_slack_pos = slack_pos, target_slack_neg = slack_neg, + p_bat = p_dis .- p_ch) +end + +""" + decompose_costs(prob, result; sol=tsddr_solution(prob, result), + lb_tol=ACTIVE_RECOURSE_LB_TOL_PU) -> NamedTuple + +Recompute every reported cost independently from the primal solution +(canonical spec, "Stage objective and cost accounting"). + +**Raw solver accounting vs physical reporting are kept separate.** The nominally +nonnegative primal quantities — the two-sided active recourse `d⁺`/`d⁻` and the +soft target slacks `δ⁺`/`δ⁻` — have a hard lower bound of 0, but an interior-point +solve can return an individual value a hair below it. This function therefore: + +1. keeps the RAW solver values for objective reproduction (`raw_*`); +2. VALIDATES that no single value is below `−lb_tol` (per variable), failing loudly + otherwise (a genuinely negative value is a defect, not noise); +3. PROJECTS the tolerated bound noise **elementwise** with `max(·, 0)` — never on + the aggregate, since negative and positive bus values could otherwise cancel; +4. computes every PUBLIC physical quantity from the projected values. + +Public physical fields are consequently always finite and `≥ 0`: +`active_deficit_pu`, `active_surplus_pu`, `active_deficit_energy_mwh`, +`active_surplus_energy_mwh`, `total_active_recourse_energy_mwh`, +`max_active_deficit_pu`, `max_active_surplus_pu`, `active_recourse_cost`, +`target_penalty`, `target_violation`. They satisfy the identities + + total_active_recourse_energy_mwh == active_deficit_energy_mwh + active_surplus_energy_mwh + active_recourse_cost == active_recourse_cost_per_mwh * total_active_recourse_energy_mwh + physical_operating_cost == generator_cost + battery_throughput_cost + active_recourse_cost + total_check == physical_operating_cost + target_penalty (projected) + +Raw diagnostics reproduce the solver objective exactly (the objective uses the raw +primal): `raw_active_deficit_pu`, `raw_active_surplus_pu`, `raw_active_recourse_cost`, +`raw_total_check ≈ result.objective`, `solver_objective_recompute_residual`, plus +`maximum_active_recourse_lower_bound_violation_pu` and +`active_recourse_projection_correction == active_recourse_cost − raw_active_recourse_cost`. + +Recourse is reported by DIRECTION — never merged into one "shed" number and never a +per-load fraction (`d⁺` may exceed local demand or be positive where `p^d = 0`). +""" +function decompose_costs(prob::BatteryTSDDRProblem, result; + sol = tsddr_solution(prob, result), + lb_tol::Real = ACTIVE_RECOURSE_LB_TOL_PU) + T = prob.horizon; nG = prob.nGen; nK = prob.nBat; nB = prob.nBus + Rrep = prob.reporting_horizon + + # ── Validate the per-variable lower-bound noise, then project elementwise ── + # `worst_viol` is how far the MOST-negative single recourse value sits below 0. + worst_viol = 0.0 + if nB > 0 + for t in 1:T, b in 1:nB + worst_viol = max(worst_viol, -Float64(sol.active_deficit_pu[b, t]), + -Float64(sol.active_surplus_pu[b, t])) + end + end + worst_viol <= lb_tol || error( + "active-recourse lower-bound violation $worst_viol pu exceeds the declared " * + "per-variable tolerance $lb_tol pu — a value this far below zero is a defect, " * + "not interior-point noise. Refusing to report a physically impossible " * + "negative recourse quantity.") + + # Generator + throughput are computed from the raw primal (they contribute to the + # solver objective unchanged). Recourse is projected elementwise per bus/stage. + gen_stage = zeros(Float64, T); cyc_stage = zeros(Float64, T); rec_stage = zeros(Float64, T) + def_pu_stage = zeros(Float64, T); sur_pu_stage = zeros(Float64, T) # projected + raw_def_pu_stage = zeros(Float64, T); raw_sur_pu_stage = zeros(Float64, T) # raw + coeff = prob.active_recourse_cost_per_mwh * prob.baseMVA * prob.dt + for t in 1:T + for g in 1:nG + pgv = Float64(sol.pg[g, t]) + gen_stage[t] += prob.gen_c2[g] * pgv^2 + prob.gen_c1[g] * pgv + prob.gen_c0[g] + end + for k in 1:nK + cyc_stage[t] += prob.cycle_coeffs[k] * + (Float64(sol.p_ch[k, t]) + Float64(sol.p_dis[k, t])) + end + for b in 1:nB + rd = Float64(sol.active_deficit_pu[b, t]); rs = Float64(sol.active_surplus_pu[b, t]) + raw_def_pu_stage[t] += rd; raw_sur_pu_stage[t] += rs + def_pu_stage[t] += max(rd, 0.0); sur_pu_stage[t] += max(rs, 0.0) # elementwise + end + rec_stage[t] = coeff * (def_pu_stage[t] + sur_pu_stage[t]) # from PROJECTED values + end + phys_stage = gen_stage .+ cyc_stage .+ rec_stage + generator_cost = sum(gen_stage) + throughput_cost = sum(cyc_stage) + + # Projected (public) aggregates. + deficit_pu = sum(def_pu_stage); surplus_pu = sum(sur_pu_stage) + active_recourse_cost = coeff * (deficit_pu + surplus_pu) + physical = generator_cost + throughput_cost + active_recourse_cost + max_deficit_pu = nB > 0 ? maximum(max(Float64(sol.active_deficit_pu[b, t]), 0.0) + for t in 1:T for b in 1:nB) : 0.0 + max_surplus_pu = nB > 0 ? maximum(max(Float64(sol.active_surplus_pu[b, t]), 0.0) + for t in 1:T for b in 1:nB) : 0.0 + + # Raw (diagnostic) aggregates — reproduce the solver objective. + raw_deficit_pu = sum(raw_def_pu_stage); raw_surplus_pu = sum(raw_sur_pu_stage) + raw_active_recourse_cost = coeff * (raw_deficit_pu + raw_surplus_pu) + projection_correction = active_recourse_cost - raw_active_recourse_cost + + # Soft target penalty — raw for objective reproduction, projected for reporting. + target_penalty = 0.0; target_violation = 0.0 + raw_target_penalty = 0.0 + if prob.mode === :soft && nK > 0 + slack_viol = 0.0 + for x in sol.target_slack_pos + slack_viol = max(slack_viol, -Float64(x)) + end + for x in sol.target_slack_neg + slack_viol = max(slack_viol, -Float64(x)) + end + slack_viol <= lb_tol || error( + "target-slack lower-bound violation $slack_viol pu exceeds the declared " * + "tolerance $lb_tol pu — refusing to report a negative target slack/penalty.") + raw_sp = sum(x -> Float64(x), sol.target_slack_pos) + raw_sn = sum(x -> Float64(x), sol.target_slack_neg) + proj_sp = sum(x -> max(Float64(x), 0.0), sol.target_slack_pos) # elementwise + proj_sn = sum(x -> max(Float64(x), 0.0), sol.target_slack_neg) + raw_target_penalty = prob.rho1 * (raw_sp + raw_sn) + + (prob.rho2 / 2) * (sum(x -> Float64(x)^2, sol.target_slack_pos) + + sum(x -> Float64(x)^2, sol.target_slack_neg)) + target_violation = proj_sp + proj_sn + target_penalty = prob.rho1 * target_violation + + (prob.rho2 / 2) * (sum(x -> max(Float64(x), 0.0)^2, sol.target_slack_pos) + + sum(x -> max(Float64(x), 0.0)^2, sol.target_slack_neg)) + end + + raw_total_check = generator_cost + throughput_cost + raw_active_recourse_cost + raw_target_penalty + solver_objective_recompute_residual = abs(raw_total_check - Float64(result.objective)) + + return (total_solver_objective = Float64(result.objective), + generator_cost = generator_cost, + battery_throughput_cost = throughput_cost, + # ── public (projected, always ≥ 0) ── + active_recourse_cost = active_recourse_cost, + active_deficit_pu = deficit_pu, + active_surplus_pu = surplus_pu, + active_deficit_energy_mwh = prob.baseMVA * prob.dt * deficit_pu, + active_surplus_energy_mwh = prob.baseMVA * prob.dt * surplus_pu, + total_active_recourse_energy_mwh = prob.baseMVA * prob.dt * (deficit_pu + surplus_pu), + max_active_deficit_pu = max_deficit_pu, + max_active_surplus_pu = max_surplus_pu, + physical_operating_cost = physical, + target_penalty = target_penalty, + target_violation = target_violation, + reporting_physical_cost = sum(@view phys_stage[1:Rrep]), + lookahead_physical_cost = T > Rrep ? sum(@view phys_stage[Rrep+1:T]) : 0.0, + total_check = physical + target_penalty, + # ── raw diagnostics (reproduce the solver objective) ── + raw_active_deficit_pu = raw_deficit_pu, + raw_active_surplus_pu = raw_surplus_pu, + raw_active_recourse_cost = raw_active_recourse_cost, + raw_target_penalty = raw_target_penalty, + raw_total_check = raw_total_check, + solver_objective_recompute_residual = solver_objective_recompute_residual, + maximum_active_recourse_lower_bound_violation_pu = worst_viol, + active_recourse_projection_correction = projection_correction, + gen_stage = gen_stage, cyc_stage = cyc_stage, + rec_stage = rec_stage, phys_stage = phys_stage) +end + +""" + solve_stage_with_starts(prob, e_prev, w_t, target; madnlp_kwargs, + soft_seed = nothing, prev_solution = nothing) + -> (result, log) + +Solve one stage trying a FIXED, deterministic sequence of primal STARTS and +accepting the FIRST solver-accepted result (never the cheapest): + +1. flat start (`vm = 1`); +2. target-consistent battery start ([`target_consistent_start!`](@ref)); +3. seed from the corresponding solved soft / targetless-diagnostic ACP point (`soft_seed`); +4. previous-stage accepted solution (`prev_solution`), where applicable. + +Only the starting point varies. Tolerances, iteration limits, equations, bounds, +and all generator/network data are identical across attempts. Every attempt and +its solver status are returned in `log` (a vector of `(start, status, accepted)`). +""" +function solve_stage_with_starts(prob::BatteryTSDDRProblem, e_prev::AbstractVector, + w_t::AbstractVector, target::AbstractVector; + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-6), + soft_seed = nothing, prev_solution = nothing) + is_targetless(prob) && + error("solve_stage_with_starts requires a target mode; got a targetless diagnostic") + set_tsddr_initial_soc!(prob, e_prev) + set_tsddr_uncertainty!(prob, w_t) + set_tsddr_targets!(prob, target) + + attempts = Any[(:flat, () -> reset_flat_start!(prob)), + (:target_consistent, () -> (reset_flat_start!(prob); + target_consistent_start!(prob, e_prev, target)))] + soft_seed === nothing || + push!(attempts, (:soft_seed, () -> seed_start_from_solution!(prob, soft_seed))) + prev_solution === nothing || + push!(attempts, (:previous_stage, () -> seed_start_from_solution!(prob, prev_solution))) + + log = Tuple{Symbol,Any,Bool}[] + local result + for (name, setup) in attempts + setup() + result = MadNLP.madnlp(prob.model; madnlp_kwargs...) + ok = solve_succeeded_result(result) && isfinite(result.objective) + push!(log, (name, result.status, ok)) + ok && return result, log + end + return result, log # all starts failed; caller reports the log +end + +# Accepted-status predicate local to this file (avoids the exported-name clash +# between the Phase-1 status predicate and the DecisionRulesExa result predicate). +solve_succeeded_result(result) = + result.status == MadNLP.SOLVE_SUCCEEDED || + result.status == MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL + +""" + tsddr_balance_residuals(prob, sol) -> Matrix + +`[nBat × T]` residuals of the battery energy balance, which should be ≈ 0. +""" +function tsddr_balance_residuals(prob::BatteryTSDDRProblem, sol) + T = prob.horizon; nK = prob.nBat + res = zeros(Float64, nK, T) + for (k, bat) in enumerate(prob.case.batteries), t in 1:T + res[k, t] = Float64(sol.soc[k, t+1]) - (1 - bat.sigma * prob.dt) * Float64(sol.soc[k, t]) - + bat.eta_ch * prob.dt * Float64(sol.p_ch[k, t]) + + (prob.dt / bat.eta_dis) * Float64(sol.p_dis[k, t]) + end + return res +end + +""" + tsddr_max_primal_residual(prob, result) -> Float64 + +Largest constraint-bound violation of the returned point via the NLPModels +interface (`max(lcon − c, c − ucon, 0)`). +""" +function tsddr_max_primal_residual(prob::BatteryTSDDRProblem, result) + x = Array(result.solution) + c = Array(NLPModels.cons(prob.model, x)) + lcon = Array(prob.model.meta.lcon); ucon = Array(prob.model.meta.ucon) + viol = 0.0 + @inbounds for i in eachindex(c) + viol = max(viol, lcon[i] - c[i], c[i] - ucon[i], 0.0) + end + return viol +end diff --git a/examples/BatteryStorageOPF/src/demand_process.jl b/examples/BatteryStorageOPF/src/demand_process.jl new file mode 100644 index 0000000..a562d0e --- /dev/null +++ b/examples/BatteryStorageOPF/src/demand_process.jl @@ -0,0 +1,735 @@ +# demand_process.jl +# +# Seeded, finite-support, PURE demand-uncertainty process for the battery-storage +# TS-DDR example (BATTERY_STORAGE_OPF_PLAN.md §6). The process is deterministic +# given its parameters and a seed: it turns integer atom indices into realized +# per-stage demand multipliers with NO dependence on generator prices and NO +# renewables. +# +# Structure of the uncertainty at one stage t: +# +# w_t = [ s_t , r_{1,t} , … , r_{R,t} ] (length nw = 1 + nregion) +# +# * s_t = base_shape[t] · L^{a_t} is the realized SYSTEM-WIDE multiplier, +# the deterministic hourly time-of-day shape times the atom's random +# system load factor L; +# * r_{k,t} = R_k^{a_t} is the realized multiplier of REGION k. +# +# Region k of a bus is fixed by a deterministic, price-free, topology-derived +# rule (nearest graph-anchor, see [`assign_regions`](@ref)). The realized demand +# at bus p, stage t is +# +# pd_realized = bus_pd[p] · s_t · r_{region(p),t}, +# qd_realized = bus_qd[p] · s_t · r_{region(p),t}, +# +# so active and reactive demand are scaled by the SAME factor and every bus keeps +# its original power factor. Both the policy (which observes `w_t`) and the +# ExaModels NLP (which multiplies `bus_pd`/`bus_qd` by the parameter entries) +# consume exactly this `w_t`, so there is one canonical uncertainty vector. +# +# The support is a small finite set of joint atoms `(L, R_1, …, R_R)` with +# explicit probabilities, drawn i.i.d. across stages (stagewise independent — the +# form ordinary SDDP backward passes need in a later phase). Atom-index matrices +# are generated with StableRNGs so `(same params, same seed) ⇒ identical indices +# and hashes` across Julia versions and machines. + +using StableRNGs +using Random +using SHA +using JSON + +# ── Atoms and the load process ──────────────────────────────────────────────── + +""" + LoadAtom + +One joint realization of the finite demand support. + +* `system_factor` : the system-wide random load factor `L` (dimensionless). +* `regional_factors` : per-region multipliers `R_k`, length `nregion`. + +The realized system multiplier at a stage is `base_shape[t]·system_factor`; each +region's multiplier is its `regional_factors` entry. Active and reactive demand +are scaled by the same product, preserving every bus's power factor. +""" +struct LoadAtom + system_factor::Float64 + regional_factors::Vector{Float64} +end + +""" + LoadProcess + +A pure, seeded, finite-support demand process for a fixed [`BatteryCase`] network. + +Fields: +* `nregion` : number of demand regions `R` (≥ 1). +* `region_of_bus` : length-`nbus` vector; entry `p` is the region (1..R) of the + bus at array position `p`. Deterministic and topology-derived. +* `anchor_bus_ids` : the `R` original PGLib bus ids used as region anchors, in + region order (region `k`'s anchor is `anchor_bus_ids[k]`). +* `region_rule` : short string naming the assignment rule (provenance). +* `base_shape` : deterministic time-of-day shape, length `period`, values + around 1. Stage `t` uses `base_shape[((t-1) mod period) + 1]`. +* `period` : base-shape period in stages (e.g. 24 for hourly). +* `atoms` : the finite support, a vector of [`LoadAtom`]. +* `probs` : atom probabilities, same length as `atoms`, summing to 1. +* `train_seed` : declared RNG seed for TRAINING scenario generation. +* `eval_seed` : declared RNG seed for EVALUATION scenario generation + (distinct from `train_seed`). + +The per-stage uncertainty dimension is `nw = 1 + nregion`. +""" +struct LoadProcess + nregion::Int + region_of_bus::Vector{Int} + anchor_bus_ids::Vector{Int} + region_rule::String + base_shape::Vector{Float64} + period::Int + atoms::Vector{LoadAtom} + probs::Vector{Float64} + train_seed::Int + eval_seed::Int + preset::Symbol # calibration preset name (:D0/:D1/:D2/:custom) +end + +# ── Calibration ladder (fixed, declared; strongest → weakest) ───────────────── +# +# The demand amplitude is a CASE-DESIGN parameter: too heavy a process makes the +# realized demand unservable and forces nonzero active recourse (deficit d⁺ / +# surplus d⁻), which invalidates a scientific run. These three presets are the +# declared ladder; the first that passes the targetless-diagnostic and soft +# zero-active-recourse gates on case14 AND case300 becomes the public default. +# Rejected (stronger) presets remain available as named experimental presets. No +# interpolation, no unrestricted search. +""" + DEMAND_PRESETS + +Fixed calibration ladder for the default demand process, strongest first: + +| preset | `base_amplitude` | high system factor | scarce regional | off-region | +|:--|:--|:--|:--|:--| +| `:D0` | 0.05 | 1.03 | 1.05 | 0.99 | +| `:D1` | 0.04 | 1.02 | 1.04 | 0.99 | +| `:D2` | 0.03 | 1.01 | 1.03 | 0.995 | + +The finite system/regional atom structure and the calm probability are unchanged +across presets. Passing this ladder is a per-case, per-experiment obligation — +see [`DEFAULT_DEMAND_PRESET`](@ref) for the current (non-passing) status. +""" +const DEMAND_PRESETS = Dict( + :D0 => (base_amplitude = 0.05, high_factor = 1.03, scarce_factor = 1.05, off_scarce_factor = 0.99), + :D1 => (base_amplitude = 0.04, high_factor = 1.02, scarce_factor = 1.04, off_scarce_factor = 0.99), + :D2 => (base_amplitude = 0.03, high_factor = 1.01, scarce_factor = 1.03, off_scarce_factor = 0.995), +) + +""" + DEFAULT_DEMAND_PRESET + +The weakest declared preset (`:D2`), used as the convenient **tutorial / API +default** so the documented one-command examples run. + +!!! warning "Not a scientifically accepted default" + `:D2` is a **conservative tutorial default**, not a gate-passing scientific + benchmark. No member of the declared ladder (`:D0`, `:D1`, `:D2`) passed the + case300 zero-active-recourse acceptance checks. + + * `:D2` passes the current **case14** checks. + * `:D2` does **not** pass the current **case300** zero-active-recourse checks. + * Every scientific experiment must record its preset explicitly and + independently pass the zero-active-recourse gates (both deficit d⁺ and + surplus d⁻ zero within tolerance) for its own case. + * The current **case300 battery placement/configuration is a REJECTED + candidate**, not a frozen benchmark; screening it is Phase-3 work. +""" +const DEFAULT_DEMAND_PRESET = :D2 + +""" + n_uncertainty(process) -> Int + +Per-stage uncertainty dimension `nw = 1 + nregion` (the realized system +multiplier followed by the per-region multipliers). +""" +n_uncertainty(p::LoadProcess) = 1 + p.nregion + +""" + natom(process) -> Int + +Number of atoms in the finite support. +""" +natom(p::LoadProcess) = length(p.atoms) + +# ── Deterministic, price-free region assignment ─────────────────────────────── + +""" + _bus_adjacency(nd) -> Vector{Vector{Int}} + +Undirected adjacency list over bus ARRAY POSITIONS built from the in-service +branches of `nd` (self-loops and duplicates are harmless for BFS). Used only for +the topology-derived region assignment; carries no electrical weighting. +""" +function _bus_adjacency(nd::NetworkData) + adj = [Int[] for _ in 1:nbus(nd)] + for br in nd.branches + push!(adj[br.f_pos], br.t_pos) + push!(adj[br.t_pos], br.f_pos) + end + return adj +end + +""" + _bfs_distances(adj, source) -> Vector{Int} + +Unweighted shortest-path (hop) distances from bus position `source` to every bus +over adjacency `adj`. Unreachable buses get `typemax(Int)`. +""" +function _bfs_distances(adj::Vector{Vector{Int}}, source::Int) + n = length(adj) + dist = fill(typemax(Int), n) + dist[source] = 0 + queue = Int[source] + head = 1 + while head <= length(queue) + u = queue[head]; head += 1 + du = dist[u] + for v in adj[u] + if dist[v] == typemax(Int) + dist[v] = du + 1 + push!(queue, v) + end + end + end + return dist +end + +""" + assign_regions(nd, nregion) -> (region_of_bus, anchor_bus_ids) + +Partition the buses of `nd` into `nregion` connected regions using a +deterministic, physically defensible, price-free topology rule: + +1. anchor 1 is the highest-degree bus (ties broken by smallest original id); +2. each further anchor is the bus whose minimum hop-distance to the already + chosen anchors is largest (farthest-point sampling; ties by smallest id), so + anchors are spread across the network; +3. every bus is assigned to its nearest anchor by hop distance (ties broken by + anchor order); buses unreachable from all anchors fall back to region 1. + +Returns the length-`nbus` region vector (over array positions) and the anchor +original bus ids in region order. + +Depends only on branch topology and bus ids — never on generator prices, so the +region geography is stable under any price change (see the plan's prohibition on +price-derived case design). +""" +function assign_regions(nd::NetworkData, nregion::Int) + nb = nbus(nd) + (1 <= nregion <= nb) || + error("nregion must satisfy 1 ≤ nregion ≤ nbus=$nb; got $nregion") + adj = _bus_adjacency(nd) + deg = [length(a) for a in adj] + + # Anchor 1: highest degree, ties → smallest original id. + order = sortperm(1:nb; by = p -> (-deg[p], nd.buses[p].id)) + anchors = Int[order[1]] + # Farthest-point sampling for the remaining anchors. + mindist = _bfs_distances(adj, anchors[1]) + while length(anchors) < nregion + # Pick the bus maximizing distance to the nearest existing anchor; break + # ties by smallest original id. Cap unreachable (typemax) at nb+1 so an + # isolated component does not silently win every round. + best = 0; best_key = (typemin(Int), 0) + for p in 1:nb + p in anchors && continue + d = mindist[p] == typemax(Int) ? nb + 1 : mindist[p] + key = (d, -nd.buses[p].id) + if key > best_key + best_key = key; best = p + end + end + push!(anchors, best) + dnew = _bfs_distances(adj, best) + @inbounds for p in 1:nb + mindist[p] = min(mindist[p], dnew[p]) + end + end + + # Assign each bus to its nearest anchor (ties → lower region index). + anchor_dists = [_bfs_distances(adj, a) for a in anchors] + region_of_bus = Vector{Int}(undef, nb) + for p in 1:nb + best_r = 1; best_d = anchor_dists[1][p] + for r in 2:nregion + if anchor_dists[r][p] < best_d + best_d = anchor_dists[r][p]; best_r = r + end + end + # Unreachable from every anchor ⇒ region 1 (deterministic fallback). + region_of_bus[p] = best_d == typemax(Int) ? 1 : best_r + end + anchor_bus_ids = [nd.buses[a].id for a in anchors] + return region_of_bus, anchor_bus_ids +end + +# ── Default building blocks ─────────────────────────────────────────────────── + +""" + default_base_shape(period; amplitude=0.15) -> Vector{Float64} + +A smooth deterministic time-of-day shape of length `period`, one sinusoidal peak +per period with mean 1: + +`base_shape[h] = 1 + amplitude·sin(2π(h - 1)/period − π/2)` (min at h=1, peak at +mid-period). `amplitude` (default 0.15) is the peak fractional deviation. Values +stay in `[1 − amplitude, 1 + amplitude] > 0`. +""" +function default_base_shape(period::Int; amplitude::Real = 0.15) + period >= 1 || error("period must be ≥ 1; got $period") + (0 <= amplitude < 1) || error("amplitude must satisfy 0 ≤ amplitude < 1; got $amplitude") + return [1.0 + amplitude * sin(2pi * (h - 1) / period - pi / 2) for h in 1:period] +end + +""" + default_load_atoms(nregion; high_factor=1.05, scarce_factor=1.15, + off_scarce_factor=1.0, calm_prob=0.4) -> (atoms, probs) + +Build the default finite support: `1 + nregion` atoms. + +* atom 1 (`calm`) : `L = 1`, every region multiplier `= 1`; probability + `calm_prob`. +* atom `1+k` (`scarce_k`) : `L = high_factor`, region `k` multiplier + `= scarce_factor`, all other regions `= off_scarce_factor`; probability + `(1 − calm_prob)/nregion` each. + +A small, SDDP-friendly support in which scarcity MOVES between regions across +atoms (the mechanism a later phase exploits) while active/reactive power factor +is preserved. No factor depends on generator prices. +""" +function default_load_atoms(nregion::Int; + high_factor::Real = 1.05, + scarce_factor::Real = 1.15, + off_scarce_factor::Real = 1.0, + calm_prob::Real = 0.4) + nregion >= 1 || error("nregion must be ≥ 1; got $nregion") + (0 < calm_prob < 1) || error("calm_prob must be in (0,1); got $calm_prob") + atoms = LoadAtom[LoadAtom(1.0, ones(Float64, nregion))] + for k in 1:nregion + R = fill(Float64(off_scarce_factor), nregion) + R[k] = Float64(scarce_factor) + push!(atoms, LoadAtom(Float64(high_factor), R)) + end + each = (1.0 - calm_prob) / nregion + probs = vcat(Float64(calm_prob), fill(each, nregion)) + return atoms, probs +end + +# ── Construction ────────────────────────────────────────────────────────────── + +""" + make_load_process(case; kwargs...) -> LoadProcess + +Build the seeded finite-support demand process for a [`BatteryCase`]. + +# Keywords +- `nregion::Int = 3`: number of demand regions. +- `period::Int = 24`: base-shape period in stages. +- `base_amplitude::Real = 0.15`: peak fractional deviation of the base shape. +- `train_seed::Int = 20260722`: declared training seed. +- `eval_seed::Int = 970122`: declared evaluation seed (must differ from + `train_seed`). +- `atoms`, `probs`: optional explicit finite support; by default + [`default_load_atoms`](@ref)`(nregion)` is used. +- `base_shape`: optional explicit base shape (length `period`). + +All parameters are validated: probabilities must be nonnegative, finite, and sum +to 1; atom regional-factor vectors must have length `nregion`; seeds must differ; +factors must be finite and positive. +""" +function make_load_process(case::BatteryCase; + preset::Symbol = DEFAULT_DEMAND_PRESET, + nregion::Int = 3, + period::Int = 24, + base_amplitude::Union{Nothing,Real} = nothing, + calm_prob::Real = 0.4, + train_seed::Int = 20260722, + eval_seed::Int = 970122, + atoms::Union{Nothing,AbstractVector{LoadAtom}} = nothing, + probs::Union{Nothing,AbstractVector{<:Real}} = nothing, + base_shape::Union{Nothing,AbstractVector{<:Real}} = nothing) + nd = case.network + region_of_bus, anchor_bus_ids = assign_regions(nd, nregion) + + # Resolve the calibration preset. Explicit `atoms`/`base_amplitude` override + # it and mark the process `:custom` so the manifest never claims a preset the + # numbers do not match. + haskey(DEMAND_PRESETS, preset) || preset === :custom || + error("unknown demand preset :$preset; known: $(sort(collect(keys(DEMAND_PRESETS))))") + cal = get(DEMAND_PRESETS, preset, DEMAND_PRESETS[DEFAULT_DEMAND_PRESET]) + amp = base_amplitude === nothing ? cal.base_amplitude : Float64(base_amplitude) + resolved_preset = (atoms !== nothing || base_shape !== nothing || + base_amplitude !== nothing) ? :custom : preset + + shape = base_shape === nothing ? default_base_shape(period; amplitude = amp) : + Float64.(collect(base_shape)) + length(shape) == period || + error("base_shape must have length period=$period; got $(length(shape))") + all(x -> isfinite(x) && x > 0, shape) || + error("base_shape entries must be finite and > 0") + + if atoms === nothing + atoms_v, probs_v = default_load_atoms(nregion; + high_factor = cal.high_factor, + scarce_factor = cal.scarce_factor, + off_scarce_factor = cal.off_scarce_factor, + calm_prob = calm_prob) + else + atoms_v = collect(LoadAtom, atoms) + probs_v = probs === nothing ? + error("explicit atoms require explicit probs") : Float64.(collect(probs)) + end + + _validate_support(atoms_v, probs_v, nregion) + train_seed != eval_seed || + error("train_seed and eval_seed must differ (got $train_seed for both)") + + return LoadProcess(nregion, region_of_bus, anchor_bus_ids, + "nearest_graph_anchor_farthest_point", shape, period, + atoms_v, probs_v, Int(train_seed), Int(eval_seed), + resolved_preset) +end + +""" + load_process_from_fields(; nregion, region_of_bus, anchor_bus_ids, region_rule, + base_shape, period, atoms, probs, train_seed, + eval_seed, preset) -> LoadProcess + +Rebuild a [`LoadProcess`](@ref) from EXACT stored field values, with no defaults +and no recomputation (no re-derivation of `base_shape` from an amplitude, no +re-running of the region assignment). This is the reconstruction path used by +the stochastic manifest so a restored process is bit-identical and reproduces its +recorded hash. The support is revalidated before the process is returned. +""" +function load_process_from_fields(; nregion::Integer, + region_of_bus::AbstractVector{<:Integer}, + anchor_bus_ids::AbstractVector{<:Integer}, + region_rule::AbstractString, + base_shape::AbstractVector{<:Real}, + period::Integer, + atoms::AbstractVector{LoadAtom}, + probs::AbstractVector{<:Real}, + train_seed::Integer, eval_seed::Integer, + preset::Union{Symbol,AbstractString}) + atoms_v = collect(LoadAtom, atoms); probs_v = Float64.(collect(probs)) + _validate_support(atoms_v, probs_v, Int(nregion)) + length(base_shape) == period || + error("base_shape length $(length(base_shape)) ≠ period $period") + return LoadProcess(Int(nregion), Int.(collect(region_of_bus)), + Int.(collect(anchor_bus_ids)), String(region_rule), + Float64.(collect(base_shape)), Int(period), + atoms_v, probs_v, Int(train_seed), Int(eval_seed), + Symbol(preset)) +end + +""" + demand_multiplier_summary(process) -> NamedTuple + +Peak demand multipliers implied by the process, reported by the calibration +gates: + +* `max_system_multiplier` — `max_t max_a base_shape[t]·L^a`, the largest + system-wide scaling of nominal demand; +* `max_bus_multiplier` — `max_t max_a max_r base_shape[t]·L^a·R_r^a`, the largest + scaling seen by any individual bus (the quantity that actually decides + servability); +* `min_bus_multiplier` — the corresponding minimum. +""" +function demand_multiplier_summary(process::LoadProcess) + max_sys = -Inf; max_bus = -Inf; min_bus = Inf + for h in process.base_shape, a in process.atoms + s = h * a.system_factor + max_sys = max(max_sys, s) + for r in a.regional_factors + max_bus = max(max_bus, s * r); min_bus = min(min_bus, s * r) + end + end + return (max_system_multiplier = max_sys, + max_bus_multiplier = max_bus, + min_bus_multiplier = min_bus) +end + +""" + _validate_support(atoms, probs, nregion) + +Validate a finite support: matching lengths, correct regional-factor +dimensions, finite positive factors, and probabilities that are nonnegative, +finite, and sum to 1 (within 1e-9). Throws on the first violation. +""" +function _validate_support(atoms::AbstractVector{LoadAtom}, probs::AbstractVector{<:Real}, nregion::Int) + length(atoms) == length(probs) || + error("atoms ($(length(atoms))) and probs ($(length(probs))) length mismatch") + isempty(atoms) && error("finite support must contain at least one atom") + for (a, at) in enumerate(atoms) + length(at.regional_factors) == nregion || + error("atom $a has $(length(at.regional_factors)) regional factors, expected nregion=$nregion") + (isfinite(at.system_factor) && at.system_factor > 0) || + error("atom $a has non-positive/non-finite system_factor $(at.system_factor)") + all(x -> isfinite(x) && x > 0, at.regional_factors) || + error("atom $a has a non-positive/non-finite regional factor") + end + all(p -> isfinite(p) && p >= 0, probs) || + error("atom probabilities must be finite and ≥ 0") + s = sum(probs) + abs(s - 1) <= 1e-9 || error("atom probabilities must sum to 1; got $s") + return nothing +end + +# ── Scenario index generation (stage-major) ─────────────────────────────────── + +""" + scenario_index_matrix(process, horizon, paths; seed) -> Matrix{Int} + +Draw a `horizon × paths` matrix of atom indices (each in `1:natom(process)`), +i.i.d. across stages and paths according to `process.probs`, using +`StableRNG(seed)`. + +Ordering is **stage-major**: entry `[t, p]` is the atom of stage `t` on path `p`, +column `p` being one full scenario. The same `(process, horizon, paths, seed)` +always produces the identical matrix (and hence the identical protocol hash). +""" +function scenario_index_matrix(process::LoadProcess, horizon::Int, paths::Int; seed::Integer) + horizon >= 1 || error("horizon must be ≥ 1; got $horizon") + paths >= 1 || error("paths must be ≥ 1; got $paths") + rng = StableRNG(UInt64(unsigned(Int64(seed)))) + A = natom(process) + # Precompute the cumulative distribution once; sample by inverse-CDF so the + # draw shape (one Float64 per (stage,path)) is fixed and never varies with A. + cdf = cumsum(process.probs) + cdf[end] = 1.0 # guard against fp drift so u ≤ cdf[end] always hits an atom + idx = Matrix{Int}(undef, horizon, paths) + for p in 1:paths, t in 1:horizon # column-major loop matches stage-major storage + u = rand(rng) + a = 1 + @inbounds while a < A && u > cdf[a] + a += 1 + end + idx[t, p] = a + end + return idx +end + +# ── Materialization: atom indices → realized w_flat (stage-major) ───────────── + +""" + materialize_scenario(process, index_row; horizon=length(index_row)) -> Vector{Float64} + +Turn one scenario's atom-index column (`index_row`, length `horizon`) into the +flat per-stage uncertainty vector `w_flat` of length `horizon·nw`, stage-major: +for stage `t` the block is `[ s_t , r_{1,t} , … , r_{R,t} ]` with +`s_t = base_shape[t]·L^{a_t}` and `r_{k,t} = R_k^{a_t}`. + +`index_row` may be a vector or a matrix column view. Every index must be a valid +atom (`1:natom`). +""" +function materialize_scenario(process::LoadProcess, index_row::AbstractVector{<:Integer}; + horizon::Int = length(index_row)) + length(index_row) == horizon || + error("index_row length $(length(index_row)) ≠ horizon $horizon") + nw = n_uncertainty(process) + A = natom(process) + P = process.period + w = Vector{Float64}(undef, horizon * nw) + for t in 1:horizon + a = Int(index_row[t]) + (1 <= a <= A) || error("atom index $a at stage $t out of range 1:$A") + atom = process.atoms[a] + base = process.base_shape[((t - 1) % P) + 1] + off = (t - 1) * nw + w[off + 1] = base * atom.system_factor + @inbounds for k in 1:process.nregion + w[off + 1 + k] = atom.regional_factors[k] + end + end + return w +end + +""" + materialize_all(process, index_matrix) -> Vector{Vector{Float64}} + +Materialize every column of a stage-major `index_matrix` into its `w_flat`, +returning one vector per path (the paired-evaluation scenario set). +""" +function materialize_all(process::LoadProcess, index_matrix::AbstractMatrix{<:Integer}) + horizon = size(index_matrix, 1) + return [materialize_scenario(process, view(index_matrix, :, p); horizon = horizon) + for p in 1:size(index_matrix, 2)] +end + +# ── Serialization: the scenario protocol ────────────────────────────────────── + +""" + process_canonical_content(process) -> String + +Ordered, human-inspectable string capturing every scientific parameter of the +demand process: regions and their anchors, the base shape, the atoms and +probabilities, the per-stage dimension, and the declared seeds. Its SHA-256 is +[`process_hash`](@ref). The process is fully deterministic given this content. +""" +function process_canonical_content(process::LoadProcess) + io = IOBuffer() + println(io, "schema=battery_load_process/2") + println(io, "preset=", process.preset) + println(io, "nregion=", process.nregion) + println(io, "region_rule=", process.region_rule) + println(io, "anchor_bus_ids=", join(process.anchor_bus_ids, ",")) + println(io, "region_of_bus=", join(process.region_of_bus, ",")) + println(io, "period=", process.period) + println(io, "base_shape=", join(_fmt.(process.base_shape), ",")) + println(io, "nw=", n_uncertainty(process)) + println(io, "natom=", natom(process)) + println(io, "train_seed=", process.train_seed) + println(io, "eval_seed=", process.eval_seed) + for (a, atom) in enumerate(process.atoms) + println(io, "atom ", a, + " prob=", _fmt(process.probs[a]), + " L=", _fmt(atom.system_factor), + " R=", join(_fmt.(atom.regional_factors), ",")) + end + return String(take!(io)) +end + +""" + process_hash(process) -> String + +Lowercase hex SHA-256 of [`process_canonical_content`](@ref). Identical process +parameters (regions, base shape, atoms, probabilities, seeds) reproduce the same +hash. +""" +process_hash(process::LoadProcess) = bytes2hex(sha256(process_canonical_content(process))) + +""" + index_matrix_hash(index_matrix) -> String + +Lowercase hex SHA-256 of a stage-major atom-index matrix, computed from its +shape and column-major integer contents so it is stable across runs. +""" +function index_matrix_hash(index_matrix::AbstractMatrix{<:Integer}) + io = IOBuffer() + print(io, size(index_matrix, 1), "x", size(index_matrix, 2), ":") + for v in index_matrix # column-major traversal, deterministic + print(io, Int(v), ",") + end + return bytes2hex(sha256(take!(io))) +end + +""" + write_scenario_protocol(path, process, index_matrix; kind, seed, + extra=Dict()) -> String + +Serialize a paired-scenario protocol to JSON at `path`: all process parameters, +the stage-major atom-index matrix (as nested arrays), the generating `seed` and +`kind` (`"train"`/`"eval"`), the per-stage dimension and units, and the process +and index-matrix hashes. Returns `path`. + +The stored index matrix IS the paired protocol: every method materializes it +with [`materialize_scenario`](@ref) to obtain byte-identical `w_flat`s. +""" +function write_scenario_protocol(path::AbstractString, process::LoadProcess, + index_matrix::AbstractMatrix{<:Integer}; + kind::AbstractString, seed::Integer, + extra::AbstractDict = Dict{String,Any}()) + horizon, paths = size(index_matrix) + doc = Dict{String,Any}( + "schema" => "battery_load_process/1", + "kind" => String(kind), + "seed" => Int(seed), + "horizon" => horizon, + "paths" => paths, + "nw_per_stage" => n_uncertainty(process), + "process_hash_sha256" => process_hash(process), + "index_matrix_hash_sha256" => index_matrix_hash(index_matrix), + "process" => Dict{String,Any}( + "preset" => String(process.preset), + "nregion" => process.nregion, + "region_rule" => process.region_rule, + "anchor_bus_ids" => process.anchor_bus_ids, + "region_of_bus" => process.region_of_bus, + "period" => process.period, + "base_shape" => process.base_shape, + "train_seed" => process.train_seed, + "eval_seed" => process.eval_seed, + "atoms" => [Dict("prob" => process.probs[a], + "system_factor" => process.atoms[a].system_factor, + "regional_factors" => process.atoms[a].regional_factors) + for a in 1:natom(process)], + ), + # Stage-major: outer list = stages, inner list = per-path atom indices. + "index_matrix" => [Int.(index_matrix[t, :]) for t in 1:horizon], + "units" => "atom indices (1-based); w_t = [base_shape·L, R_1..R_nregion]", + ) + for (k, v) in extra + doc[k] = v + end + open(io -> JSON.print(io, doc, 2), path, "w") + return path +end + +""" + reconstruct_scenario_protocol(path) -> (process, index_matrix, meta) + +Rebuild a demand process and its stage-major atom-index matrix from a JSON +protocol written by [`write_scenario_protocol`](@ref), verifying BOTH recorded +hashes: + +1. the reconstructed process reproduces `process_hash_sha256`; +2. the reconstructed index matrix reproduces `index_matrix_hash_sha256`. + +Any mismatch raises an error. `meta` is a NamedTuple with `kind`, `seed`, +`horizon`, `paths`, and `nw_per_stage`. +""" +function reconstruct_scenario_protocol(path::AbstractString) + doc = JSON.parsefile(path) + pr = doc["process"] + nregion = Int(pr["nregion"]) + atoms = LoadAtom[LoadAtom(Float64(a["system_factor"]), + Float64.(a["regional_factors"])) for a in pr["atoms"]] + probs = Float64[Float64(a["prob"]) for a in pr["atoms"]] + # EXACT-field reconstruction: no defaults, no recomputation of base_shape or + # of the region assignment. + process = load_process_from_fields(; + nregion = nregion, + region_of_bus = Int.(pr["region_of_bus"]), + anchor_bus_ids = Int.(pr["anchor_bus_ids"]), + region_rule = String(pr["region_rule"]), + base_shape = Float64.(pr["base_shape"]), + period = Int(pr["period"]), + atoms = atoms, probs = probs, + train_seed = Int(pr["train_seed"]), eval_seed = Int(pr["eval_seed"]), + preset = get(pr, "preset", "custom")) + + got_ph = process_hash(process) + want_ph = String(doc["process_hash_sha256"]) + got_ph == want_ph || + error("process-hash mismatch: protocol recorded $want_ph but rebuilt $got_ph") + + horizon = Int(doc["horizon"]); paths = Int(doc["paths"]) + rows = doc["index_matrix"] + length(rows) == horizon || error("index_matrix has $(length(rows)) stages, expected $horizon") + index_matrix = Matrix{Int}(undef, horizon, paths) + for t in 1:horizon + row = rows[t] + length(row) == paths || error("stage $t has $(length(row)) paths, expected $paths") + @inbounds for p in 1:paths + index_matrix[t, p] = Int(row[p]) + end + end + got_ih = index_matrix_hash(index_matrix) + want_ih = String(doc["index_matrix_hash_sha256"]) + got_ih == want_ih || + error("index-matrix-hash mismatch: protocol recorded $want_ih but rebuilt $got_ih") + + meta = (kind = String(doc["kind"]), seed = Int(doc["seed"]), + horizon = horizon, paths = paths, + nw_per_stage = Int(doc["nw_per_stage"])) + return process, index_matrix, meta +end diff --git a/examples/BatteryStorageOPF/src/manifest.jl b/examples/BatteryStorageOPF/src/manifest.jl new file mode 100644 index 0000000..9db2913 --- /dev/null +++ b/examples/BatteryStorageOPF/src/manifest.jl @@ -0,0 +1,321 @@ +# manifest.jl +# +# Deterministic manifest for a BatteryCase: everything needed to rebuild the +# exact same battery placement and parameters, plus provenance and attribution +# (see BATTERY_STORAGE_OPF_PLAN.md §5). +# +# Two representations: +# * a canonical CONTENT string → SHA-256 `manifest_hash` that depends only on +# the scientific content (case name, MATPOWER source-file SHA-256, seed, +# eligible-bus rule, selected buses, battery params, units, counts) and NOT on +# timestamps, absolute paths, or package versions, so "same seed ⇒ same hash" +# holds across runs of the same PGLib artifact; +# * a JSON manifest file that additionally records provenance, attribution, the +# license (CC BY 4.0), the content hash, and file hashes. + +using SHA +using JSON + +# ── PGLib attribution / upstream version (best effort) ──────────────────────── + +const PGLIB_ATTRIBUTION = string( + "Network data from the Power Grid Library for Benchmarking AC Optimal ", + "Power Flow Algorithms (PGLib-OPF), distributed via PGLib.jl. ", + "Cite: S. Babaeinejadsarookolaee et al., \"The Power Grid Library for ", + "Benchmarking AC Optimal Power Flow Algorithms\", arXiv:1908.02788.", +) + +# PGLib-OPF is released under the Creative Commons Attribution 4.0 license. +const PGLIB_LICENSE_NAME = "Creative Commons Attribution 4.0 International" +const PGLIB_LICENSE_URL = "https://creativecommons.org/licenses/by/4.0/" + +# The PGLib.jl artifact directory is named `pglib-opf-` (e.g. +# `pglib-opf-23.07`), so the upstream benchmark release is read straight from +# the path; fall back to the pinning PGLib.jl package version. +function _pglib_upstream_version() + m = match(r"pglib-opf-([0-9]+\.[0-9]+)", _pglib_case_dir()) + return m === nothing ? "pinned-by-PGLib.jl-" * _pkg_version("PGLib") : m.captures[1] +end + +# ── Canonical content + hash ────────────────────────────────────────────────── + +# Deterministic full-precision rendering of a Float64 (round-trippable shortest +# form is stable for equal values). +_fmt(x::Real) = string(Float64(x)) +_fmt(x::Integer) = string(x) +_fmt(x::AbstractString) = String(x) + +""" + canonical_content(bc) -> String + +Ordered, human-inspectable string capturing exactly the scientific content that +defines the case: the PGLib case name, the SHA-256 of the exact MATPOWER source +bytes, the seed and eligible-bus rule, the ordered battery placement, and every +battery parameter with units. Its SHA-256 is [`manifest_hash`](@ref). +Deliberately excludes timestamps, absolute paths, and package versions so the +hash is reproducible; including the MATPOWER source hash makes the hash sensitive +to the exact network bytes. +""" +function canonical_content(bc::BatteryCase) + io = IOBuffer() + nd = bc.network + println(io, "schema=battery_storage_opf/2") + println(io, "case_name=", bc.case_name) + println(io, "matpower_file=", basename(bc.parse_meta.filepath)) + println(io, "matpower_sha256=", bc.parse_meta.matpower_sha256) + println(io, "baseMVA=", _fmt(nd.baseMVA)) + println(io, "nbus=", nbus(nd), " ngen=", ngen(nd), + " nbranch=", nbranch(nd), " nload=", nload(nd)) + println(io, "total_load_pu=", _fmt(bc.total_load_pu)) + println(io, "seed=", bc.seed) + println(io, "eligible_bus_rule=", bc.eligible_bus_rule) + println(io, "explicit_buses=", bc.explicit_buses) + println(io, "number_of_batteries=", bc.number_of_batteries) + println(io, "duration_hours=", _fmt(bc.duration_hours)) + println(io, "initial_soc=", _fmt(bc.initial_soc)) + println(io, "charge_efficiency=", _fmt(bc.charge_efficiency)) + println(io, "discharge_efficiency=", _fmt(bc.discharge_efficiency)) + println(io, "fleet_power_fraction=", _fmt(bc.fleet_power_fraction)) + println(io, "self_discharge_rate=", _fmt(bc.self_discharge_rate)) + println(io, "e_min_fraction=", _fmt(bc.e_min_fraction)) + println(io, "cycle_cost_per_mwh=", _fmt(bc.cycle_cost_per_mwh)) + println(io, "units=power:pu[baseMVA];energy:pu_hours;time:hours") + println(io, "selected_bus_ids=", join(bc.selected_bus_ids, ",")) + for b in bc.batteries + println(io, "battery ", b.id, + " bus=", b.bus_id, + " pch_max=", _fmt(b.p_charge_max), + " pdis_max=", _fmt(b.p_discharge_max), + " e_min=", _fmt(b.e_min), + " e_max=", _fmt(b.e_max), + " e_init=", _fmt(b.e_init), + " eta_ch=", _fmt(b.eta_ch), + " eta_dis=", _fmt(b.eta_dis), + " sigma=", _fmt(b.sigma), + " cycle_cost_per_mwh=", _fmt(b.cycle_cost_per_mwh)) + end + return String(take!(io)) +end + +""" + manifest_hash(bc) -> String + +Lowercase hex SHA-256 of [`canonical_content`](@ref). Identical inputs (same +PGLib case, seed, and sizing parameters) reproduce the same hash. +""" +manifest_hash(bc::BatteryCase) = bytes2hex(sha256(canonical_content(bc))) + +_sha256_file(path) = bytes2hex(open(sha256, path)) + +# ── JSON manifest ───────────────────────────────────────────────────────────── + +""" + battery_manifest(bc) -> Dict + +Assemble the full machine-readable manifest: scientific content, provenance +(package + Julia + upstream PGLib versions), attribution, units, the horizon / +terminal-treatment note, and the content hash. File hashes are added by +[`write_manifest`](@ref) once files exist on disk. +""" +function battery_manifest(bc::BatteryCase) + nd = bc.network + return Dict( + "schema" => "battery_storage_opf/2", + "content_hash_sha256" => manifest_hash(bc), + "pglib" => Dict( + "case_name" => bc.case_name, + "matpower_file" => basename(bc.parse_meta.filepath), + "matpower_path" => bc.parse_meta.filepath, + "matpower_sha256" => bc.parse_meta.matpower_sha256, + "upstream_release" => _pglib_upstream_version(), + "attribution" => PGLIB_ATTRIBUTION, + "license" => PGLIB_LICENSE_NAME, + "license_url" => PGLIB_LICENSE_URL, + ), + "versions" => Dict( + "julia" => string(VERSION), + "PGLib" => bc.parse_meta.pglib_version, + "PowerModels" => bc.parse_meta.powermodels_version, + ), + "network" => Dict( + "baseMVA" => nd.baseMVA, + "nbus" => nbus(nd), "ngen" => ngen(nd), + "nbranch" => nbranch(nd), "nload" => nload(nd), + "total_load_pu" => bc.total_load_pu, + "per_unit_input" => nd.per_unit_input, + ), + "placement" => Dict( + "seed" => bc.seed, + "eligible_bus_rule" => bc.eligible_bus_rule, + "explicit_buses" => bc.explicit_buses, + "n_eligible_buses" => length(bc.eligible_bus_ids), + "selected_bus_ids" => bc.selected_bus_ids, + ), + "battery_parameters" => Dict( + "number_of_batteries" => bc.number_of_batteries, + "duration_hours" => bc.duration_hours, + "initial_soc" => bc.initial_soc, + "charge_efficiency" => bc.charge_efficiency, + "discharge_efficiency" => bc.discharge_efficiency, + "fleet_power_fraction" => bc.fleet_power_fraction, + "self_discharge_rate" => bc.self_discharge_rate, + "e_min_fraction" => bc.e_min_fraction, + "cycle_cost_per_mwh" => bc.cycle_cost_per_mwh, + "per_battery_p_charge_max_pu" => + isempty(bc.batteries) ? 0.0 : bc.batteries[1].p_charge_max, + "per_battery_p_discharge_max_pu" => + isempty(bc.batteries) ? 0.0 : bc.batteries[1].p_discharge_max, + "per_battery_e_max_pu_h" => + isempty(bc.batteries) ? 0.0 : bc.batteries[1].e_max, + "per_battery_e_init_pu_h" => + isempty(bc.batteries) ? 0.0 : bc.batteries[1].e_init, + ), + "batteries" => [Dict( + "id" => b.id, "bus_id" => b.bus_id, "bus_pos" => b.bus_pos, + "p_charge_max_pu" => b.p_charge_max, + "p_discharge_max_pu" => b.p_discharge_max, + "e_min_pu_h" => b.e_min, "e_max_pu_h" => b.e_max, + "e_init_pu_h" => b.e_init, + "eta_ch" => b.eta_ch, "eta_dis" => b.eta_dis, + "sigma_per_h" => b.sigma, + "cycle_cost_per_mwh" => b.cycle_cost_per_mwh, + ) for b in bc.batteries], + "units" => Dict( + "power" => "per-unit on baseMVA", + "energy" => "per-unit-hours (pu·h)", + "time" => "hours", + "cost" => "USD; generator cost per pu power, battery cost per MWh throughput", + ), + "horizon" => Dict( + "note" => string( + "Phase 1 deterministic foundation: horizon T and stage length ", + "Δt (hours) are chosen at model-build time (build_battery_de). ", + "The stochastic load process, paired protocol, and terminal ", + "treatment are defined in later phases."), + ), + ) +end + +""" + write_manifest(bc, path; extra_files = String[]) -> String + +Write the JSON manifest to `path`, adding SHA-256 hashes of any `extra_files` +(e.g. a human-readable battery file) and of the manifest's own canonical +content. Returns `path`. +""" +function write_manifest(bc::BatteryCase, path::AbstractString; extra_files = String[]) + man = battery_manifest(bc) + man["generated_at_utc"] = _utc_now_string() + fh = Dict{String,String}() + for f in extra_files + isfile(f) && (fh[basename(f)] = _sha256_file(f)) + end + man["file_hashes_sha256"] = fh + open(path, "w") do io + JSON.print(io, man, 2) + end + return path +end + +# ISO-8601 UTC timestamp without pulling in Dates' TimeZones; Libc.strftime is +# stdlib and sufficient for a provenance stamp. +function _utc_now_string() + t = round(Int, time()) + return Libc.strftime("%Y-%m-%dT%H:%M:%SZ", t) +end + +# ── Human-readable battery file ─────────────────────────────────────────────── + +""" + write_battery_file(bc, path) -> String + +Write a human-readable CSV of the battery fleet (one row per battery, original +bus id preserved). Returns `path`. +""" +function write_battery_file(bc::BatteryCase, path::AbstractString) + open(path, "w") do io + println(io, "# BatteryStorageOPF fleet for PGLib case \"", bc.case_name, + "\" (seed=", bc.seed, ", baseMVA=", bc.network.baseMVA, ")") + println(io, "# power/energy are per-unit on baseMVA; energy in pu·h") + println(io, "battery_id,bus_id,p_charge_max_pu,p_discharge_max_pu,", + "e_min_pu_h,e_max_pu_h,e_init_pu_h,eta_ch,eta_dis,sigma_per_h,cycle_cost_per_mwh") + for b in bc.batteries + println(io, b.id, ",", b.bus_id, ",", b.p_charge_max, ",", + b.p_discharge_max, ",", b.e_min, ",", b.e_max, ",", + b.e_init, ",", b.eta_ch, ",", b.eta_dis, ",", b.sigma, + ",", b.cycle_cost_per_mwh) + end + end + return path +end + +# ── Reconstruction ──────────────────────────────────────────────────────────── + +""" + reconstruct_case(manifest_path) -> BatteryCase + +Rebuild the case recorded in a JSON manifest and verify it, returning the case +only when ALL of the following match the manifest: + +1. **Source network bytes** — the SHA-256 of the currently resolved MATPOWER + file equals the manifest's recorded `matpower_sha256`. This is checked FIRST, + directly against the resolved artifact, so a tampered/moved/updated network + file is rejected before anything else. +2. **Ordered placement** — the reconstructed battery buses match the recorded + `selected_bus_ids` in the same order. +3. **Battery parameters** — captured by the content hash: re-running + `make_battery_case` with the recorded case name, seed, and sizing parameters + reproduces the recorded `content_hash_sha256` (which itself includes the + source-file hash). + +Any mismatch raises an error, so a returned case is verified identical in source +network bytes, ordered placement, and battery parameters to the recorded one. +""" +function reconstruct_case(manifest_path::AbstractString) + man = JSON.parsefile(manifest_path) + bp = man["battery_parameters"] + pl = man["placement"] + pg = man["pglib"] + + # 1. Verify the source network bytes FIRST, against the currently resolved + # MATPOWER file, before trusting anything else in the manifest. + recorded_src = String(pg["matpower_sha256"]) + _, filepath = resolve_pglib_case(String(pg["case_name"])) + current_src = _sha256_file(filepath) + current_src == recorded_src || error( + "source-file hash mismatch for \"$(pg["case_name"])\": manifest recorded " * + "$recorded_src but the resolved MATPOWER file hashes to $current_src. " * + "The network data differs from when the manifest was written.") + + explicit = Bool(pl["explicit_buses"]) + kwargs = ( + number_of_batteries = Int(bp["number_of_batteries"]), + seed = Int(pl["seed"]), + duration_hours = Float64(bp["duration_hours"]), + initial_soc = Float64(bp["initial_soc"]), + charge_efficiency = Float64(bp["charge_efficiency"]), + discharge_efficiency = Float64(bp["discharge_efficiency"]), + fleet_power_fraction = Float64(bp["fleet_power_fraction"]), + self_discharge_rate = Float64(bp["self_discharge_rate"]), + e_min_fraction = Float64(bp["e_min_fraction"]), + cycle_cost_per_mwh = Float64(bp["cycle_cost_per_mwh"]), + ) + bc = if explicit + make_battery_case(String(pg["case_name"]); + buses = Int.(pl["selected_bus_ids"]), kwargs...) + else + make_battery_case(String(pg["case_name"]); kwargs...) + end + + # 2. Ordered placement. + Int.(pl["selected_bus_ids"]) == bc.selected_bus_ids || + error("reconstruction produced a different battery placement order") + + # 3. Battery parameters (via the content hash, which embeds the source hash). + got = manifest_hash(bc) + want = String(man["content_hash_sha256"]) + got == want || error( + "reconstruction content-hash mismatch: manifest recorded $want but " * + "rebuilt $got. Battery parameters differ from when the manifest was written.") + return bc +end diff --git a/examples/BatteryStorageOPF/src/network_data.jl b/examples/BatteryStorageOPF/src/network_data.jl new file mode 100644 index 0000000..ae9951e --- /dev/null +++ b/examples/BatteryStorageOPF/src/network_data.jl @@ -0,0 +1,413 @@ +# network_data.jl +# +# Typed parsing of a PGLib / PowerModels AC network into the flat, per-unit +# structures the ExaModels AC-polar builder consumes. +# +# Design rules honored here (see BATTERY_STORAGE_OPF_PLAN.md §5): +# * PGLib component identifiers are NEVER assumed consecutive or equal to +# their array position. Every component keeps its original PGLib `index` +# and we build explicit, stable ID → position maps (`*_id_to_pos`). +# * All unit conversion happens HERE, in one documented data layer, so model +# code sees only per-unit quantities on the `baseMVA` base with angles in +# radians. Nothing downstream re-scales. +# +# We parse the raw MATPOWER data exactly as `PowerModels.parse_file` returns it +# (mixed units, `per_unit = false`: powers in MW/MVAr, voltages in pu, angles +# already converted to radians by the parser). The conversion below is gated on +# the observed `per_unit` flag so a pre-converted dict is also handled safely. + +using PowerModels +using PGLib +using SHA + +# ── Typed, per-unit network components ──────────────────────────────────────── +# Every struct stores the original PGLib id AND the resolved 1-based array +# position of any component it references. + +""" + BusData + +One AC bus, per unit on `baseMVA`. + +* `id` : original PGLib bus index (may be non-consecutive). +* `bus_type` : PowerModels bus type (1 = PQ, 2 = PV, 3 = reference/slack). +* `gs`, `bs` : shunt conductance / susceptance at the bus (pu). +* `vmin`,`vmax`: voltage-magnitude limits (pu). +""" +struct BusData + id::Int + bus_type::Int + gs::Float64 + bs::Float64 + vmin::Float64 + vmax::Float64 +end + +""" + GenData + +One in-service generator, per unit on `baseMVA`. + +* `id` : original PGLib generator index. +* `bus_id`,`bus_pos` : the bus it injects into (original id and array position). +* `pmin`,`pmax` : active-power limits (pu). +* `qmin`,`qmax` : reactive-power limits (pu). +* `cost2`,`cost1`,`cost0`: polynomial cost so that the \$/hour cost of `pg` pu is + `cost2·pg² + cost1·pg + cost0`. Converted once from the MATPOWER \$/MW model. +""" +struct GenData + id::Int + bus_id::Int + bus_pos::Int + pmin::Float64 + pmax::Float64 + qmin::Float64 + qmax::Float64 + cost2::Float64 + cost1::Float64 + cost0::Float64 +end + +""" + BranchData + +One in-service branch, per unit on `baseMVA`, with the raw AC-polar parameters +used by [`_ac_branch_coeffs`](@ref). + +* `id` : original PGLib branch index. +* `f_id`,`f_pos` : from-bus original id and array position. +* `t_id`,`t_pos` : to-bus original id and array position. +* `br_r`,`br_x` : series resistance / reactance (pu). +* `g_fr`,`b_fr`,`g_to`,`b_to`: line-charging shunts (pu). +* `tap`,`shift` : transformer turns ratio and phase shift (rad). +* `rate_a` : apparent-power thermal limit (pu); `Inf` if unlimited. +* `angmin`,`angmax` : angle-difference limits (rad). +""" +struct BranchData + id::Int + f_id::Int + f_pos::Int + t_id::Int + t_pos::Int + br_r::Float64 + br_x::Float64 + g_fr::Float64 + b_fr::Float64 + g_to::Float64 + b_to::Float64 + tap::Float64 + shift::Float64 + rate_a::Float64 + angmin::Float64 + angmax::Float64 +end + +""" + LoadData + +One in-service load, per unit on `baseMVA`. + +* `id` : original PGLib load index. +* `bus_id`,`bus_pos` : the bus it draws from (original id and array position). +* `pd`,`qd` : active / reactive demand (pu). +""" +struct LoadData + id::Int + bus_id::Int + bus_pos::Int + pd::Float64 + qd::Float64 +end + +""" + NetworkData + +Fully parsed, per-unit AC network. Components are stored sorted by original id; +`*_id_to_pos` give the stable inverse maps. `bus_pd`/`bus_qd` aggregate the +loads onto buses (length `nbus`, positional). +""" +struct NetworkData + case_name::String + baseMVA::Float64 + buses::Vector{BusData} + gens::Vector{GenData} + branches::Vector{BranchData} + loads::Vector{LoadData} + ref_bus_positions::Vector{Int} + bus_id_to_pos::Dict{Int,Int} + gen_id_to_pos::Dict{Int,Int} + branch_id_to_pos::Dict{Int,Int} + load_id_to_pos::Dict{Int,Int} + bus_pd::Vector{Float64} # aggregated active demand per bus (pu) + bus_qd::Vector{Float64} # aggregated reactive demand per bus (pu) + per_unit_input::Bool # whether the source dict was already per-unit +end + +nbus(nd::NetworkData) = length(nd.buses) +ngen(nd::NetworkData) = length(nd.gens) +nbranch(nd::NetworkData) = length(nd.branches) +nload(nd::NetworkData) = length(nd.loads) + +# ── PGLib case-name resolution ──────────────────────────────────────────────── + +""" + _pglib_case_dir() -> String + +Absolute directory holding the `pglib_opf_*.m` benchmark files that ship with +the pinned `PGLib.jl` artifact. +""" +_pglib_case_dir() = PGLib.PGLib_opf + +""" + available_pglib_cases() -> Vector{String} + +Sorted canonical short names (the `pglib_opf_` prefix and `.m` suffix stripped) +of every case in the pinned PGLib artifact, e.g. `"case300_ieee"`. +""" +function available_pglib_cases() + dir = _pglib_case_dir() + names = String[] + for f in readdir(dir) + (endswith(f, ".m") && startswith(f, "pglib_opf_")) || continue + push!(names, replace(replace(f, r"\.m$" => ""), r"^pglib_opf_" => "")) + end + return sort!(names) +end + +""" + resolve_pglib_case(name) -> (canonical_name, filepath) + +Resolve a user-supplied PGLib case name to its canonical short name and the +absolute `.m` file path, unambiguously. + +Accepted forms (all mapped to the same case): `"case300_ieee"`, +`"pglib_opf_case300_ieee"`, `"pglib_opf_case300_ieee.m"`. + +Errors: +* a name matching no case lists the closest available names; +* a name matching several cases (only possible via a partial/loose query) is + rejected with the full ambiguous set. Exact canonical matches are never + ambiguous. +""" +function resolve_pglib_case(name::AbstractString) + dir = _pglib_case_dir() + cases = available_pglib_cases() + + # Normalize the query to a canonical short name. + q = String(name) + q = replace(q, r"\.m$" => "") + q = replace(q, r"^pglib_opf_" => "") + + # Exact canonical match wins outright (never ambiguous). + if q in cases + return q, joinpath(dir, "pglib_opf_" * q * ".m") + end + + # Otherwise treat the query as a substring and demand a unique hit. + hits = filter(c -> occursin(q, c), cases) + if isempty(hits) + # Offer nearby suggestions by shared prefix to make the error actionable. + prefix = first(q, min(length(q), 6)) + near = filter(c -> startswith(c, prefix), cases) + suffix = isempty(near) ? "" : "\n Did you mean: " * join(first(near, 8), ", ") + error("PGLib case \"$name\" not found in $(length(cases)) available cases." * + suffix * + "\n Use available_pglib_cases() to list them all.") + elseif length(hits) > 1 + error("PGLib case \"$name\" is ambiguous; it matches $(length(hits)) cases:\n " * + join(hits, ", ") * + "\n Pass an exact canonical name (e.g. one of the above).") + end + c = only(hits) + return c, joinpath(dir, "pglib_opf_" * c * ".m") +end + +# ── Parsing + one documented unit-conversion layer ──────────────────────────── + +# Extract the polynomial generator cost as (cost2, cost1, cost0) with pg in pu. +# MATPOWER model-2 polynomials are in \$ with power in MW; converting to pu +# multiplies the degree-k coefficient by baseMVA^k. Piecewise-linear (model 1) +# cost is not supported by this example and errors loudly. +function _gen_cost_pu(gen, sbase::Float64) + model = Int(get(gen, "model", 2)) + model == 2 || error("generator $(get(gen,"index","?")) uses cost model $model; " * + "only MATPOWER polynomial cost (model 2) is supported") + cost = Float64.(get(gen, "cost", Float64[])) + ncost = Int(get(gen, "ncost", length(cost))) + # Right-align: MATPOWER lists highest-order coefficient first. + c2 = c1 = c0 = 0.0 + if ncost >= 3 + c2, c1, c0 = cost[end-2], cost[end-1], cost[end] + elseif ncost == 2 + c1, c0 = cost[end-1], cost[end] + elseif ncost == 1 + c0 = cost[end] + end + return (c2 * sbase^2, c1 * sbase, c0) +end + +""" + parse_network(data::Dict, case_name) -> NetworkData + +Parse a PowerModels network dict into per-unit [`NetworkData`], keeping only +in-service components (bus_type ≠ 4, and unit/branch/load status = 1) and +building explicit stable id → position maps. + +Power quantities are divided by `baseMVA` iff the dict is not already per-unit +(`data["per_unit"] == false`, the state `PowerModels.parse_file` returns). +""" +function parse_network(data::AbstractDict, case_name::AbstractString) + sbase = Float64(data["baseMVA"]) + per_unit = Bool(get(data, "per_unit", false)) + # Scale factor applied to MW/MVAr quantities (1.0 when already per-unit). + s = per_unit ? 1.0 : 1.0 / sbase + + # ── Shunt admittance (PowerModels keeps shunts as a SEPARATE component, not + # on the bus). Aggregate in-service shunt gs/bs per original bus id so they + # enter the power balance as gs·vm² (active) and −bs·vm² (reactive). Shunt + # gs/bs are already per-unit; the same `s` gate applies for consistency. + shunt_gs = Dict{Int,Float64}(); shunt_bs = Dict{Int,Float64}() + for sh in values(get(data, "shunt", Dict{String,Any}())) + Int(get(sh, "status", 1)) == 1 || continue + bid = Int(sh["shunt_bus"]) + shunt_gs[bid] = get(shunt_gs, bid, 0.0) + Float64(get(sh, "gs", 0.0)) * s + shunt_bs[bid] = get(shunt_bs, bid, 0.0) + Float64(get(sh, "bs", 0.0)) * s + end + + # ── Buses (drop isolated bus_type == 4) ─────────────────────────────────── + bus_raw = collect(values(data["bus"])) + filter!(b -> Int(b["bus_type"]) != 4, bus_raw) + sort!(bus_raw, by = b -> Int(b["index"])) + buses = BusData[] + bus_id_to_pos = Dict{Int,Int}() + for (pos, b) in enumerate(bus_raw) + id = Int(b["index"]) + bus_id_to_pos[id] = pos + # Bus-level gs/bs (rare) plus any separate shunt components at this bus. + gs = Float64(get(b, "gs", 0.0)) * s + get(shunt_gs, id, 0.0) + bs = Float64(get(b, "bs", 0.0)) * s + get(shunt_bs, id, 0.0) + push!(buses, BusData(id, Int(b["bus_type"]), gs, bs, + Float64(get(b, "vmin", 0.9)), + Float64(get(b, "vmax", 1.1)))) + end + ref_bus_positions = [bus_id_to_pos[b.id] for b in buses if b.bus_type == 3] + isempty(ref_bus_positions) && + error("network \"$case_name\" has no reference (type-3) bus") + + # ── Generators (in-service only) ────────────────────────────────────────── + gen_raw = collect(values(data["gen"])) + filter!(g -> Int(get(g, "gen_status", 1)) == 1, gen_raw) + sort!(gen_raw, by = g -> Int(g["index"])) + gens = GenData[] + gen_id_to_pos = Dict{Int,Int}() + for (pos, g) in enumerate(gen_raw) + id = Int(g["index"]) + bus_id = Int(g["gen_bus"]) + haskey(bus_id_to_pos, bus_id) || + error("generator $id references out-of-service/unknown bus $bus_id") + c2, c1, c0 = _gen_cost_pu(g, per_unit ? 1.0 : sbase) + gen_id_to_pos[id] = pos + push!(gens, GenData(id, bus_id, bus_id_to_pos[bus_id], + Float64(get(g, "pmin", 0.0)) * s, + Float64(g["pmax"]) * s, + Float64(get(g, "qmin", -Inf)) * s, + Float64(get(g, "qmax", Inf)) * s, + c2, c1, c0)) + end + + # ── Branches (in-service only) ──────────────────────────────────────────── + br_raw = collect(values(data["branch"])) + filter!(br -> Int(get(br, "br_status", 1)) == 1, br_raw) + sort!(br_raw, by = br -> Int(br["index"])) + branches = BranchData[] + branch_id_to_pos = Dict{Int,Int}() + for (pos, br) in enumerate(br_raw) + id = Int(br["index"]) + f_id = Int(br["f_bus"]); t_id = Int(br["t_bus"]) + (haskey(bus_id_to_pos, f_id) && haskey(bus_id_to_pos, t_id)) || + error("branch $id references an out-of-service/unknown bus " * + "($f_id → $t_id)") + tap = Float64(get(br, "tap", 1.0)); tap = tap ≈ 0 ? 1.0 : tap + rate = Float64(get(br, "rate_a", 0.0)) + rate_pu = rate ≈ 0 ? Inf : rate * s # 0 ⇒ unlimited (PowerModels convention) + branch_id_to_pos[id] = pos + push!(branches, BranchData(id, f_id, bus_id_to_pos[f_id], + t_id, bus_id_to_pos[t_id], + Float64(get(br, "br_r", 0.0)), + Float64(br["br_x"]), + Float64(get(br, "g_fr", 0.0)), + Float64(get(br, "b_fr", 0.0)), + Float64(get(br, "g_to", 0.0)), + Float64(get(br, "b_to", 0.0)), + tap, Float64(get(br, "shift", 0.0)), + rate_pu, + Float64(get(br, "angmin", -pi)), + Float64(get(br, "angmax", pi)))) + end + + # ── Loads (in-service only) ─────────────────────────────────────────────── + load_raw = collect(values(data["load"])) + filter!(l -> Int(get(l, "status", 1)) == 1, load_raw) + sort!(load_raw, by = l -> Int(l["index"])) + loads = LoadData[] + load_id_to_pos = Dict{Int,Int}() + bus_pd = zeros(Float64, length(buses)) + bus_qd = zeros(Float64, length(buses)) + for (pos, l) in enumerate(load_raw) + id = Int(l["index"]) + bus_id = Int(l["load_bus"]) + haskey(bus_id_to_pos, bus_id) || + error("load $id references an out-of-service/unknown bus $bus_id") + bpos = bus_id_to_pos[bus_id] + pd = Float64(get(l, "pd", 0.0)) * s + qd = Float64(get(l, "qd", 0.0)) * s + load_id_to_pos[id] = pos + push!(loads, LoadData(id, bus_id, bpos, pd, qd)) + bus_pd[bpos] += pd + bus_qd[bpos] += qd + end + + return NetworkData(String(case_name), sbase, buses, gens, branches, loads, + ref_bus_positions, bus_id_to_pos, gen_id_to_pos, + branch_id_to_pos, load_id_to_pos, bus_pd, bus_qd, per_unit) +end + +""" + load_pglib_network(name) -> (NetworkData, parse_meta) + +Resolve `name` in the pinned PGLib artifact, parse the `.m` file with +PowerModels, and return the typed [`NetworkData`] plus a small metadata +NamedTuple (`canonical_name`, `filepath`, `pglib_version`, `powermodels_version`) +recorded in the manifest. +""" +function load_pglib_network(name::AbstractString) + canonical, filepath = resolve_pglib_case(name) + data = PowerModels.parse_file(filepath) + nd = parse_network(data, canonical) + meta = (canonical_name = canonical, + filepath = filepath, + # SHA-256 of the exact MATPOWER source bytes (see manifest.jl); this + # pins provenance to the resolved network file. + matpower_sha256 = bytes2hex(open(SHA.sha256, filepath)), + pglib_version = _pkg_version("PGLib"), + powermodels_version = _pkg_version("PowerModels")) + return nd, meta +end + +# Resolve an installed dependency's version string by reading the active +# environment's Manifest.toml with the TOML stdlib. We deliberately avoid +# `using Pkg` at runtime: on this cluster Pkg's native precompile image is +# broken, so the example is kept Pkg-free (Pkg is only used, with +# --pkgimages=no, by the one-shot setup_env.jl). +function _pkg_version(name::AbstractString) + manifest = Base.active_project() === nothing ? nothing : + joinpath(dirname(Base.active_project()), "Manifest.toml") + (manifest === nothing || !isfile(manifest)) && return "unknown" + data = TOML.parsefile(manifest) + deps = get(data, "deps", data) # manifest format 2.0 nests under "deps" + entry = get(deps, name, nothing) + entry === nothing && return "unknown" + # Each package maps to a 1-element array of tables. + rec = entry isa AbstractVector ? first(entry) : entry + return string(get(rec, "version", "unknown")) +end diff --git a/examples/BatteryStorageOPF/src/reference_powermodels.jl b/examples/BatteryStorageOPF/src/reference_powermodels.jl new file mode 100644 index 0000000..ef7ec0c --- /dev/null +++ b/examples/BatteryStorageOPF/src/reference_powermodels.jl @@ -0,0 +1,79 @@ +# reference_powermodels.jl +# +# Independent base-ACP reference. PowerModels builds and solves the standard +# ACPPowerModel from the SAME PGLib file, on an entirely separate modeling stack +# (PowerModels/JuMP/Ipopt). With zero batteries (or zero battery power) our +# ExaModels deterministic equivalent must reproduce this objective within a +# declared tolerance — the model-correctness gate. + +using PowerModels +using JuMP +using Ipopt + +# Silence PowerModels' info logging during solves. +PowerModels.silence() + +""" + reference_ac_opf(case_name; print_level=0) -> NamedTuple + +Solve the base ACP OPF of PGLib `case_name` with PowerModels + Ipopt (no +batteries). Returns `(objective, status, termination, data)` where `objective` +is the operating cost in USD and `data` is the parsed PowerModels dict. +""" +function reference_ac_opf(case_name::AbstractString; print_level::Int = 0) + _, filepath = resolve_pglib_case(case_name) + data = PowerModels.parse_file(filepath) + optimizer = JuMP.optimizer_with_attributes(Ipopt.Optimizer, + "print_level" => print_level, + "tol" => 1e-8) + result = PowerModels.solve_ac_opf(data, optimizer) + return (objective = Float64(result["objective"]), + status = result["termination_status"], + solve_time = get(result, "solve_time", NaN), + data = data) +end + +""" + exa_base_objective(case_name; float_type=Float64, + madnlp_kwargs=(print_level=MadNLP.ERROR, tol=1e-8)) + -> NamedTuple + +Build a single-stage, zero-battery ExaModels deterministic equivalent for +`case_name` and solve it on the CPU, returning `(objective, status, prob, +result)`. This is the ExaModels side of the base-ACP parity check. +""" +function exa_base_objective(case_name::AbstractString; + float_type::Type{<:AbstractFloat} = Float64, + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-8)) + case = make_battery_case(case_name; number_of_batteries = 0) + prob = build_battery_de(case, 1; float_type = float_type) + result = solve_de!(prob; madnlp_kwargs...) + return (objective = Float64(result.objective), + status = result.status, prob = prob, result = result) +end + +""" + check_base_acp_parity(case_name; rtol=1e-3, kwargs...) -> NamedTuple + +Compare the zero-battery ExaModels objective against the PowerModels/Ipopt +reference for `case_name`. Returns +`(reference, exa, abs_diff, rel_diff, within, rtol, ...)`. + +`rtol` default is 1e-3: AC-OPF is nonconvex, so the two independent +solver/model stacks can in principle settle at numerically distinct—but +physically equivalent—KKT points. In practice, with the shunt admittance and +generator costs modeled identically, agreement is far tighter — measured +≈1e-10 on case14_ieee and ≈1e-12 on case300_ieee — so the loose default only +guards against environment-dependent solver noise. +""" +function check_base_acp_parity(case_name::AbstractString; rtol::Real = 1e-3, kwargs...) + ref = reference_ac_opf(case_name) + exa = exa_base_objective(case_name; kwargs...) + absd = abs(exa.objective - ref.objective) + reld = absd / max(abs(ref.objective), eps()) + return (case_name = case_name, + reference = ref.objective, exa = exa.objective, + abs_diff = absd, rel_diff = reld, + within = reld <= rtol, rtol = Float64(rtol), + reference_status = ref.status, exa_status = exa.status) +end diff --git a/examples/BatteryStorageOPF/src/stochastic_manifest.jl b/examples/BatteryStorageOPF/src/stochastic_manifest.jl new file mode 100644 index 0000000..76453c9 --- /dev/null +++ b/examples/BatteryStorageOPF/src/stochastic_manifest.jl @@ -0,0 +1,233 @@ +# stochastic_manifest.jl +# +# Phase-2 experiment manifest: the single artifact from which the whole +# experiment is reconstructed EXACTLY, with no defaults and no environment- +# variable guesses. +# +# It builds on the accepted Phase-1 `battery_manifest` (which records and lets +# `reconstruct_case` verify the MATPOWER source hash) and adds every value needed +# to rebuild the demand process field-for-field: the exact `base_shape` vector, +# `region_of_bus`, region anchors, the complete ORDERED atom definitions and +# probabilities, the period, all seeds, and the selected calibration preset — +# plus the horizons, stage duration, target mode, active-recourse / target-penalty +# configuration, and the policy architecture and activation. +# +# FIVE DISTINCT hashes are stored; none of them is reused for another role: +# 1. load_process_content_hash_sha256 — the process definition +# 2. train_index_matrix_hash_sha256 — training scenario-index matrix +# 3. eval_index_matrix_hash_sha256 — evaluation scenario-index matrix +# 4. train_protocol_file_sha256 — exact training protocol JSON bytes +# 5. eval_protocol_file_sha256 — exact evaluation protocol JSON bytes + +using JSON +using SHA + +""" + stochastic_manifest(case, process; kwargs...) -> Dict + +Assemble the Phase-2 experiment manifest. Required keywords: `reporting_horizon`, +`lookahead`, `mode`. Optional: `stage_hours`, `active_recourse_cost_per_mwh`, +`rho1`, `rho2`, `activation`, `safe_upper_margin`, `policy_layers`, +`policy_combiner_layers`, `policy_seed`, and the four artifact hashes +(`train_index_matrix_hash`, `eval_index_matrix_hash`, +`train_protocol_file_sha256`, `eval_protocol_file_sha256`) with their path counts. +""" +function stochastic_manifest(case::BatteryCase, process::LoadProcess; + reporting_horizon::Int, lookahead::Int, mode::Symbol, + stage_hours::Real = 1.0, + active_recourse_cost_per_mwh::Real = DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH, + rho1::Real = 0.0, rho2::Real = 0.0, + activation::AbstractString = "stretchedsigmoid", + safe_upper_margin::Real = 1e-3, + policy_layers = Int[], policy_combiner_layers = Int[], + policy_seed = nothing, + train_index_matrix_hash = nothing, + eval_index_matrix_hash = nothing, + train_protocol_file_sha256 = nothing, + eval_protocol_file_sha256 = nothing, + train_paths = nothing, eval_paths = nothing) + mult = demand_multiplier_summary(process) + return Dict{String,Any}( + "schema" => "battery_storage_opf_stochastic/3", + "case" => battery_manifest(case), # Phase-1 manifest (with MATPOWER source hash) + + # ── EXACT demand-process definition (field-for-field reconstruction) ── + "load_process" => Dict{String,Any}( + "preset" => String(process.preset), + "nregion" => process.nregion, + "region_rule" => process.region_rule, + "anchor_bus_ids" => process.anchor_bus_ids, + "region_of_bus" => process.region_of_bus, + "period" => process.period, + "base_shape" => process.base_shape, + "atoms" => [Dict("prob" => process.probs[a], + "system_factor" => process.atoms[a].system_factor, + "regional_factors" => process.atoms[a].regional_factors) + for a in 1:natom(process)], + "nw_per_stage" => n_uncertainty(process), + "natom" => natom(process), + "train_seed" => process.train_seed, + "eval_seed" => process.eval_seed, + "max_system_multiplier" => mult.max_system_multiplier, + "max_bus_multiplier" => mult.max_bus_multiplier, + "min_bus_multiplier" => mult.min_bus_multiplier, + ), + + # ── Five distinct hashes (never reused across roles) ───────────────── + "hashes" => Dict{String,Any}( + "load_process_content_hash_sha256" => process_hash(process), + "train_index_matrix_hash_sha256" => train_index_matrix_hash, + "eval_index_matrix_hash_sha256" => eval_index_matrix_hash, + "train_protocol_file_sha256" => train_protocol_file_sha256, + "eval_protocol_file_sha256" => eval_protocol_file_sha256, + ), + "paths" => Dict{String,Any}("train" => train_paths, "eval" => eval_paths), + + "horizon" => Dict{String,Any}( + "reporting_horizon" => reporting_horizon, + "lookahead" => lookahead, + "horizon" => reporting_horizon + lookahead, + "stage_hours" => Float64(stage_hours), + "terminal_treatment" => + "reporting horizon followed by a look-ahead buffer; identical for every method", + ), + "target_mode" => String(mode), + "active_recourse" => Dict{String,Any}( + "formulation" => "two-sided active nodal slack", + "deficit_sign" => "d⁺ ≥ 0 (injection, covers a local active shortfall)", + "surplus_sign" => "d⁻ ≥ 0 (absorption, absorbs a local active excess)", + "unbounded_above" => true, + "active_only" => true, + "active_recourse_cost_per_mwh" => Float64(active_recourse_cost_per_mwh), + "units" => "USD/MWh; stage coefficient = cost · baseMVA · Δt", + "included_in_physical_operating_cost" => true, + "accepted_scientific_requirement" => + "both directions (deficit d⁺ and surplus d⁻) zero within declared tolerance", + ), + "target_penalty" => Dict{String,Any}( + "rho1" => Float64(rho1), "rho2" => Float64(rho2), + "form" => "rho1*sum(slack_pos+slack_neg) + rho2/2*sum(slack_pos^2+slack_neg^2)", + "included_in_physical_operating_cost" => false, + ), + "policy" => Dict{String,Any}( + "architecture" => "BatteryReachablePolicy", + "encoder_layers" => collect(Int, policy_layers), + "combiner_layers" => collect(Int, policy_combiner_layers), + "activation" => String(activation), + "safe_upper_margin" => Float64(safe_upper_margin), + "policy_seed" => policy_seed, + "note" => "normalized target lies in [0, 1 - safe_upper_margin]; never exactly 1", + ), + "cost_definitions" => Dict{String,Any}( + "physical_operating_cost" => + "generator + battery throughput + active-recourse cost (deficit + surplus)", + "training_only" => "target penalty (soft mode)", + "headline_metric" => "reporting-window physical operating cost", + ), + "units" => Dict{String,Any}( + "power" => "per-unit on baseMVA", "energy" => "per-unit-hours (pu·h)", + "time" => "hours", + ), + ) +end + +""" + write_stochastic_manifest(path, case, process; kwargs...) -> String + +Write the [`stochastic_manifest`](@ref) to `path` as JSON (adding a UTC +timestamp). Returns `path`. +""" +function write_stochastic_manifest(path::AbstractString, case::BatteryCase, + process::LoadProcess; kwargs...) + man = stochastic_manifest(case, process; kwargs...) + man["generated_at_utc"] = _utc_now_string() + open(io -> JSON.print(io, man, 2), path, "w") + return path +end + +""" + reconstruct_stochastic_manifest(path) -> (case, process, meta) + +Rebuild the experiment from a Phase-2 manifest and verify it: + +1. the Phase-1 case is reconstructed through the accepted path (the MATPOWER + source hash is re-checked against the resolved artifact and the case content + hash is reproduced); +2. the demand process is rebuilt FIELD-FOR-FIELD from the stored exact values + (base shape, regions, anchors, ordered atoms, probabilities, period, seeds, + preset) with no defaults, and must reproduce + `hashes.load_process_content_hash_sha256`. + +`meta` carries the horizons, stage duration, target mode, VOLL, penalty +coefficients, policy architecture/activation, path counts, and the four artifact +hashes so a caller can verify the protocol files it loads. +""" +function reconstruct_stochastic_manifest(path::AbstractString) + doc = JSON.parsefile(path) + + # 1. Phase-1 case via the accepted verification path. + tmp = tempname() * ".json" + open(io -> JSON.print(io, doc["case"], 2), tmp, "w") + case = try + reconstruct_case(tmp) + finally + isfile(tmp) && rm(tmp; force = true) + end + + # 2. Demand process from EXACT stored fields. + lp = doc["load_process"] + atoms = LoadAtom[LoadAtom(Float64(a["system_factor"]), Float64.(a["regional_factors"])) + for a in lp["atoms"]] + probs = Float64[Float64(a["prob"]) for a in lp["atoms"]] + process = load_process_from_fields(; + nregion = Int(lp["nregion"]), + region_of_bus = Int.(lp["region_of_bus"]), + anchor_bus_ids = Int.(lp["anchor_bus_ids"]), + region_rule = String(lp["region_rule"]), + base_shape = Float64.(lp["base_shape"]), + period = Int(lp["period"]), + atoms = atoms, probs = probs, + train_seed = Int(lp["train_seed"]), eval_seed = Int(lp["eval_seed"]), + preset = String(lp["preset"])) + + h = doc["hashes"] + got = process_hash(process) + want = String(h["load_process_content_hash_sha256"]) + got == want || + error("load-process content-hash mismatch: manifest recorded $want but rebuilt $got") + + hz = doc["horizon"]; pol = doc["policy"]; tp = doc["target_penalty"] + meta = (reporting_horizon = Int(hz["reporting_horizon"]), + lookahead = Int(hz["lookahead"]), + horizon = Int(hz["horizon"]), + stage_hours = Float64(hz["stage_hours"]), + mode = Symbol(doc["target_mode"]), + active_recourse_cost_per_mwh = + Float64(doc["active_recourse"]["active_recourse_cost_per_mwh"]), + rho1 = Float64(tp["rho1"]), rho2 = Float64(tp["rho2"]), + activation = String(pol["activation"]), + encoder_layers = Int.(pol["encoder_layers"]), + combiner_layers = Int.(pol["combiner_layers"]), + train_paths = get(doc["paths"], "train", nothing), + eval_paths = get(doc["paths"], "eval", nothing), + train_index_matrix_hash = get(h, "train_index_matrix_hash_sha256", nothing), + eval_index_matrix_hash = get(h, "eval_index_matrix_hash_sha256", nothing), + train_protocol_file_sha256 = get(h, "train_protocol_file_sha256", nothing), + eval_protocol_file_sha256 = get(h, "eval_protocol_file_sha256", nothing)) + return case, process, meta +end + +""" + verify_protocol_file(path, expected_sha256) -> String + +Verify that a protocol JSON file's EXACT BYTES hash to `expected_sha256` and +return the hash. Raises on mismatch (tamper detection at the file level, distinct +from the index-matrix content hash). +""" +function verify_protocol_file(path::AbstractString, expected_sha256) + got = bytes2hex(open(sha256, path)) + expected_sha256 === nothing && return got + got == String(expected_sha256) || + error("protocol file hash mismatch for $path: expected $expected_sha256, got $got") + return got +end diff --git a/examples/BatteryStorageOPF/test/runtests.jl b/examples/BatteryStorageOPF/test/runtests.jl new file mode 100644 index 0000000..0ee6f07 --- /dev/null +++ b/examples/BatteryStorageOPF/test/runtests.jl @@ -0,0 +1,272 @@ +# runtests.jl — example-local tests for the Phase-1 battery-storage foundation. +# +# Run (from examples/BatteryStorageOPF): +# module load julia +# julia --pkgimages=no --project=. test/runtests.jl +# +# Fast checks (construction / manifest / mapping / validation) use a small PGLib +# case; solve-based checks use case14_ieee, with a construction+residual smoke +# and placement/count assertions on case300_ieee. + +include(joinpath(@__DIR__, "..", "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using Test +using MadNLP +using JSON + +const SMALL = "case14_ieee" # small solve case +const BIG = "case300_ieee" # headline case + +@testset "BatteryStorageOPF Phase 1" begin + + # ── Reproducible placement + manifest hash ─────────────────────────────── + @testset "identical seed ⇒ identical placement and hash" begin + a = make_battery_case(SMALL; number_of_batteries = 3, seed = 123) + b = make_battery_case(SMALL; number_of_batteries = 3, seed = 123) + @test a.selected_bus_ids == b.selected_bus_ids + @test manifest_hash(a) == manifest_hash(b) + # Order is stable and consistent with the battery vector. + @test [bat.bus_id for bat in a.batteries] == a.selected_bus_ids + end + + @testset "different seed ⇒ different placement (when possible)" begin + # 3 batteries drawn from >3 eligible buses: distinct seeds should give a + # different ordered placement (and hence a different hash). + c1 = make_battery_case(SMALL; number_of_batteries = 3, seed = 1) + c2 = make_battery_case(SMALL; number_of_batteries = 3, seed = 2) + @test length(c1.eligible_bus_ids) > 3 + @test c1.selected_bus_ids != c2.selected_bus_ids + @test manifest_hash(c1) != manifest_hash(c2) + end + + # ── Input validation ───────────────────────────────────────────────────── + @testset "invalid inputs fail with clear errors" begin + @test_throws ErrorException make_battery_case("no_such_case_xyz") + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 2, + charge_efficiency = 1.5) + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 2, + discharge_efficiency = 0.0) + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 2, + duration_hours = -1.0) + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 2, + initial_soc = 1.5) + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 2, + e_min_fraction = 0.6, initial_soc = 0.5) + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 2, + self_discharge_rate = 1.2) + # Too many batteries for the eligible pool. + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 10_000) + # Explicit bad bus. + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 1, + buses = [-999]) + # Explicit duplicate buses. + elig = eligible_load_bus_ids(make_battery_case(SMALL; number_of_batteries = 0).network) + @test_throws ErrorException make_battery_case(SMALL; number_of_batteries = 2, + buses = [elig[1], elig[1]]) + end + + # ── Non-consecutive identifier mapping (synthetic, exact) ──────────────── + @testset "arbitrary / non-consecutive id mapping" begin + # Hand-built PowerModels-style dict with deliberately non-consecutive, + # out-of-order bus ids and an isolated (type-4) bus that must be dropped. + data = Dict{String,Any}( + "baseMVA" => 100.0, "per_unit" => false, + "bus" => Dict( + "1" => Dict("index" => 100, "bus_type" => 3, "vmin" => 0.9, "vmax" => 1.1), + "2" => Dict("index" => 5, "bus_type" => 1, "vmin" => 0.9, "vmax" => 1.1), + "3" => Dict("index" => 42, "bus_type" => 1, "vmin" => 0.9, "vmax" => 1.1), + "4" => Dict("index" => 7, "bus_type" => 4, "vmin" => 0.9, "vmax" => 1.1), + ), + "gen" => Dict("1" => Dict("index" => 9, "gen_bus" => 100, "pmax" => 500.0, + "pmin" => 0.0, "model" => 2, "ncost" => 3, + "cost" => [0.1, 20.0, 5.0])), + "branch" => Dict( + "1" => Dict("index" => 3, "f_bus" => 100, "t_bus" => 5, "br_x" => 0.1, + "br_r" => 0.01, "rate_a" => 3.0), + "2" => Dict("index" => 1, "f_bus" => 5, "t_bus" => 42, "br_x" => 0.2, + "br_r" => 0.02, "rate_a" => 0.0)), # 0 ⇒ unlimited + "load" => Dict( + "1" => Dict("index" => 2, "load_bus" => 5, "pd" => 100.0, "qd" => 20.0), + "2" => Dict("index" => 1, "load_bus" => 42, "pd" => 50.0, "qd" => 10.0)), + ) + nd = parse_network(data, "synthetic") + # Buses sorted by original id; type-4 dropped. + @test [b.id for b in nd.buses] == [5, 42, 100] + @test nbus(nd) == 3 + @test nd.bus_id_to_pos[100] == 3 && nd.bus_id_to_pos[5] == 1 + # Gen/branch/load references resolved to positions, ids preserved. + # Components are stored sorted by original id, so look them up by id map + # (never by insertion order). + @test nd.gens[nd.gen_id_to_pos[9]].bus_pos == nd.bus_id_to_pos[100] + br3 = nd.branches[nd.branch_id_to_pos[3]] # branch id 3: bus 100 → bus 5 + @test br3.f_pos == nd.bus_id_to_pos[100] + @test br3.t_pos == nd.bus_id_to_pos[5] + @test isfinite(br3.rate_a) # rate_a = 3 ⇒ limited + br1 = nd.branches[nd.branch_id_to_pos[1]] # branch id 1: rate_a = 0 + @test isinf(br1.rate_a) # 0 rate ⇒ unlimited + # Per-unit conversion (÷ baseMVA) applied to loads and gen limits. + @test nd.bus_pd[nd.bus_id_to_pos[5]] ≈ 1.0 + @test nd.gens[1].pmax ≈ 5.0 + # Cost converted to pu: c1_pu = 20 · 100, c2_pu = 0.1 · 100². + @test nd.gens[1].cost1 ≈ 2000.0 + @test nd.gens[1].cost2 ≈ 1000.0 + @test 100 in nd.ref_bus_positions .|> (p -> nd.buses[p].id) + end + + # ── case300 structure + placement ──────────────────────────────────────── + @testset "case300 has 300 buses and 20 valid battery buses" begin + case = make_battery_case(BIG; number_of_batteries = 20, seed = 20260722) + @test nbus(case.network) == 300 + @test length(case.batteries) == 20 + @test length(unique(case.selected_bus_ids)) == 20 + # Every battery bus is a valid, in-service load bus. + elig = Set(eligible_load_bus_ids(case.network)) + @test all(b -> b.bus_id in elig, case.batteries) + @test all(b -> haskey(case.network.bus_id_to_pos, b.bus_id), case.batteries) + # case300 genuinely has non-consecutive bus ids (position ≠ id somewhere). + @test any(i -> case.network.buses[i].id != i, 1:nbus(case.network)) + validate_battery_case(case) # must not throw + end + + # ── Manifest round-trip / verified reconstruction / provenance ─────────── + @testset "manifest write, verified reconstruction, and tampering" begin + case = make_battery_case(SMALL; number_of_batteries = 3, seed = 77) + dir = mktempdir() + man = joinpath(dir, "manifest.json") + bat = joinpath(dir, "batteries.csv") + write_battery_file(case, bat) + write_manifest(case, man; extra_files = [bat]) + @test isfile(man) && isfile(bat) + + # Provenance fields are present and explicit. + m = JSON.parsefile(man) + @test haskey(m["pglib"], "matpower_file") + @test length(String(m["pglib"]["matpower_sha256"])) == 64 + @test m["pglib"]["license"] == "Creative Commons Attribution 4.0 International" + @test m["pglib"]["license_url"] == "https://creativecommons.org/licenses/by/4.0/" + @test haskey(m["pglib"], "upstream_release") + @test haskey(m["versions"], "PGLib") && haskey(m["versions"], "PowerModels") && + haskey(m["versions"], "julia") + # The MATPOWER source hash participates in the scientific content hash. + @test occursin(String(m["pglib"]["matpower_sha256"]), canonical_content(case)) + + # Normal reconstruction succeeds and verifies source bytes + placement + # + battery parameters. + rc = reconstruct_case(man) + @test rc.selected_bus_ids == case.selected_bus_ids + @test manifest_hash(rc) == manifest_hash(case) + + # Identical inputs ⇒ identical placement and manifest hash. + again = make_battery_case(SMALL; number_of_batteries = 3, seed = 77) + @test again.selected_bus_ids == case.selected_bus_ids + @test manifest_hash(again) == manifest_hash(case) + + # Tampering with the recorded source-file hash must fail reconstruction. + m["pglib"]["matpower_sha256"] = repeat("0", 64) + tampered = joinpath(dir, "tampered.json") + open(io -> JSON.print(io, m, 2), tampered, "w") + @test_throws ErrorException reconstruct_case(tampered) + end + + # ── Solve: status, residuals, no simultaneous charge/discharge ─────────── + @testset "case14 solve: status + residuals + battery operation" begin + case = make_battery_case(SMALL; number_of_batteries = 3, seed = 5, + fleet_power_fraction = 0.3, duration_hours = 4.0) + prob = build_battery_de(case, 3; stage_hours = 1.0, + demand_profile = [0.9, 1.1, 1.0]) + result = solve_de!(prob; print_level = MadNLP.ERROR, tol = 1e-8) + @test solve_succeeded(result.status) + sol = battery_solution(prob, result) + @test max_primal_residual(prob, result) < 1e-5 + @test maximum(abs, battery_balance_residuals(prob, sol)) < 1e-6 + # Cycle cost ⇒ no material simultaneous charge & discharge. + @test maximum(simultaneous_charge_discharge_power(sol)) < 1e-5 + # SoC stays within bounds. + @test all(case.batteries[1].e_min - 1e-6 .<= sol.soc .<= case.batteries[1].e_max + 1e-6) + end + + # ── Base-ACP parity (zero batteries and zero power) ────────────────────── + @testset "zero-battery / zero-power base-ACP parity" begin + # Small case: strict agreement (measured ≈1e-10 relative). + p0 = check_base_acp_parity(SMALL; rtol = 1e-6) + @test p0.within + # Headline case300: same check, conservative tolerance (measured ≈1e-12). + pbig = check_base_acp_parity(BIG; rtol = 1e-4) + @test pbig.within + # Zero battery POWER (fleet_power_fraction = 0) with 3 placed batteries + # must also reproduce the base objective. + ref = reference_ac_opf(SMALL).objective + case0 = make_battery_case(SMALL; number_of_batteries = 3, seed = 9, + fleet_power_fraction = 0.0) + prob0 = build_battery_de(case0, 1) + r0 = solve_de!(prob0; print_level = MadNLP.ERROR, tol = 1e-8) + @test solve_succeeded(r0.status) + @test abs(r0.objective - ref) / abs(ref) < 1e-5 + end + + # ── case300 construction + residual smoke ──────────────────────────────── + @testset "case300 construction + residual smoke" begin + case = make_battery_case(BIG; number_of_batteries = 20, seed = 20260722) + prob = build_battery_de(case, 2; stage_hours = 1.0, + demand_profile = [1.0, 1.02]) + result = solve_de!(prob; print_level = MadNLP.ERROR, tol = 1e-6, max_iter = 1000) + @test solve_succeeded(result.status) + sol = battery_solution(prob, result) + @test max_primal_residual(prob, result) < 1e-4 + @test maximum(abs, battery_balance_residuals(prob, sol)) < 1e-6 + @test maximum(simultaneous_charge_discharge_power(sol)) < 1e-4 + end + + # ── Stage-duration (Δt) generator-cost scaling ─────────────────────────── + # PGLib polynomial costs are USD/hour, so a Δt-hour stage costs Δt× the + # one-hour dispatch. These regressions fail under an implementation that + # omits Δt from the generator cost. + @testset "stage-duration (Δt) cost scaling" begin + ref1h = reference_ac_opf(SMALL).objective # independent 1-hour ACP + case0 = make_battery_case(SMALL; number_of_batteries = 0) + madnlp = (print_level = MadNLP.ERROR, tol = 1e-8) + + # Δt = 1 (retained parity): single stage equals the one-hour reference. + p1 = build_battery_de(case0, 1; stage_hours = 1.0) + r1 = solve_de!(p1; madnlp...) + @test solve_succeeded(r1.status) + @test abs(r1.objective - ref1h) / abs(ref1h) < 1e-5 + + # T=1, Δt=2: objective is ≈ 2× the one-hour reference. + p2 = build_battery_de(case0, 1; stage_hours = 2.0) + r2 = solve_de!(p2; madnlp...) + @test solve_succeeded(r2.status) + @test abs(r2.objective - 2 * ref1h) / abs(2 * ref1h) < 1e-5 + + # T=2, Δt=0.5, flat demand: two half-hour stages sum to the one-hour ref. + p3 = build_battery_de(case0, 2; stage_hours = 0.5, demand_profile = [1.0, 1.0]) + r3 = solve_de!(p3; madnlp...) + @test solve_succeeded(r3.status) + @test abs(r3.objective - ref1h) / abs(ref1h) < 1e-5 + + # Nonpositive / nonfinite stage_hours are rejected. + @test_throws ErrorException build_battery_de(case0, 1; stage_hours = 0.0) + @test_throws ErrorException build_battery_de(case0, 1; stage_hours = -1.0) + @test_throws ErrorException build_battery_de(case0, 1; stage_hours = Inf) + @test_throws ErrorException build_battery_de(case0, 1; stage_hours = NaN) + end + + # ── Per-battery cycle-cost coefficients (not the first battery's) ───────── + @testset "per-battery cycle-cost coefficients" begin + # Two batteries whose cycle price we can distinguish: build with distinct + # per-battery cycle costs by editing the case's battery vector. + case = make_battery_case(SMALL; number_of_batteries = 2, seed = 3) + b1, b2 = case.batteries + case.batteries[1] = BatteryData(b1.id, b1.bus_id, b1.bus_pos, + b1.p_charge_max, b1.p_discharge_max, b1.e_min, b1.e_max, b1.e_init, + b1.eta_ch, b1.eta_dis, b1.sigma, 1.0) # $1/MWh + case.batteries[2] = BatteryData(b2.id, b2.bus_id, b2.bus_pos, + b2.p_charge_max, b2.p_discharge_max, b2.e_min, b2.e_max, b2.e_init, + b2.eta_ch, b2.eta_dis, b2.sigma, 7.0) # $7/MWh + prob = build_battery_de(case, 2; stage_hours = 1.5) + # Coefficient k = cycle_cost_per_mwh_k · baseMVA · Δt, per battery. + @test prob.cycle_coeffs[1] ≈ 1.0 * case.network.baseMVA * 1.5 + @test prob.cycle_coeffs[2] ≈ 7.0 * case.network.baseMVA * 1.5 + @test prob.cycle_coeffs[1] != prob.cycle_coeffs[2] + end +end diff --git a/examples/BatteryStorageOPF/test/runtests_phase2.jl b/examples/BatteryStorageOPF/test/runtests_phase2.jl new file mode 100644 index 0000000..0cf4216 --- /dev/null +++ b/examples/BatteryStorageOPF/test/runtests_phase2.jl @@ -0,0 +1,824 @@ +# runtests_phase2.jl — Phase-2 suite for the REPAIRED implementation. +# +# Covers: shared-ACP parity with the accepted Phase-1 builder; the demand process +# and paired protocol; the reachable policy (stretchedsigmoid / hardsigmoidsafe, +# safe upper margin, non-differentiable bounds); the two-sided active nodal +# recourse (unbounded deficit d⁺ + surplus d⁻ at every bus, hard reactive KCL, +# recourse pricing and inclusion in physical cost) that gives STRICT complete +# recourse; strict absolute recourse across cases and target classes reported by +# DIRECTION; strict and soft target modes with the split-slack soft form; the +# multiplier finite-difference check; checkpointing; and CPU/GPU structural parity. +# +# Run (from examples/BatteryStorageOPF): +# module load julia +# julia --pkgimages=no --project=. test/runtests_phase2.jl +# +# GPU tests are gated on CUDA.functional() and reported as skipped on CPU nodes. + +include(joinpath(@__DIR__, "..", "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using Test +using MadNLP +using JSON +using Flux +using Zygote +using CUDA +using DecisionRulesExa +using LinearAlgebra +using Statistics +using Random + +const SMALL = "case14_ieee" + +# DECLARED numerical tolerances for "zero active recourse". An interior-point +# solution leaves the recourse power at barrier tolerance rather than exactly 0, +# so "zero" is asserted at these declared levels. They are ~4-6 orders of +# magnitude below any material recourse (a genuinely recoursing path carries +# O(1-10) MWh), so they cannot mask a real failure. Deficit d⁺ and surplus d⁻ are +# each held below RECOURSE_ENERGY_TOL_MWH; their sum below 2×. +const RECOURSE_ENERGY_TOL_MWH = 1e-3 +const RECOURSE_PU_TOL = 1e-4 + +# Fresh-solver solve used throughout (avoids CommonSolve.solve! ambiguity). +solve_tsddr(de) = MadNLP.madnlp(de.model; print_level = MadNLP.ERROR, tol = 1e-8) + +# Default (stronger) fixture, used where demand stress is irrelevant. +function fixture(; nbat = 2, nregion = 2, seed = 11) + case = make_battery_case(SMALL; number_of_batteries = nbat, seed = seed, + fleet_power_fraction = 0.3, duration_hours = 4.0) + process = make_load_process(case; nregion = nregion, period = 6, + train_seed = 101, eval_seed = 202) + return case, process +end + +# Mild, demand-feasible fixture (peak ≈1.1×) so an accepted run sheds no load. +function mild_fixture(; nbat = 2, nregion = 2, seed = 11) + case = make_battery_case(SMALL; number_of_batteries = nbat, seed = seed, + fleet_power_fraction = 0.3, duration_hours = 4.0) + atoms = [LoadAtom(1.0, [1.0, 1.0]), LoadAtom(1.03, [1.05, 0.99]), LoadAtom(1.03, [0.99, 1.05])] + process = make_load_process(case; nregion = nregion, period = 6, base_amplitude = 0.04, + atoms = atoms, probs = [0.5, 0.25, 0.25], + train_seed = 101, eval_seed = 202) + return case, process +end + +# Roll the policy forward to build a reachable target trajectory for scenario w. +function policy_targets(pol, process, w, e0, T) + Flux.reset!(pol); prev = e0; xh = Float64[] + nw = n_uncertainty(process) + for t in 1:T + wt = w[(t-1)*nw+1 : t*nw] + xt = pol(vcat(wt, prev)); append!(xh, xt); prev = xt + end + return xh +end + +@testset "BatteryStorageOPF Phase 2 (repaired)" begin + + # ── REPAIR 1: shared ACP source of truth ───────────────────────────────── + @testset "shared-ACP parity: stochastic builder reproduces Phase 1" begin + case, process = mild_fixture() + T = 2 + # Realized demand for a chosen scenario. + w = materialize_scenario(process, [1, 1]; horizon = T) + nw = n_uncertainty(process); nd = case.network + # Phase-1 builder with the SAME realized demand: atom 1 is the calm atom + # (L=1, R=1), so the realized multiplier is exactly the daily shape h_t. + prof = [process.base_shape[((t - 1) % process.period) + 1] for t in 1:T] + p1 = build_battery_de(case, T; stage_hours = 1.0, demand_profile = prof) + r1 = solve_de!(p1; print_level = MadNLP.ERROR, tol = 1e-8) + @test BatteryStorageOPF.solve_succeeded(r1.status) + + # Stochastic builder with targets DISABLED (soft, zero penalty, targets + # pinned to the realized Phase-1 SoC) and active recourse FIXED TO ZERO by + # bounds — not by a zero price, which would make recourse free and is the + # opposite of disabling it. + p2 = build_battery_tsddr_de(case, process; reporting_horizon = T, lookahead = 0, + mode = :soft, rho1 = 0.0, rho2 = 0.0, + allow_active_recourse = false, stage_hours = 1.0) + set_tsddr_initial_soc!(p2, policy_initial_state(case; float_type = Float64)) + set_tsddr_uncertainty!(p2, w) + # Realized demand must match Phase-1's baked profile exactly. + for t in 1:T, b in 1:nbus(nd) + @test isapprox(p2.realized_pd[t, b], nd.bus_pd[b] * prof[t]; rtol = 1e-12) + @test isapprox(p2.realized_qd[t, b], nd.bus_qd[b] * prof[t]; rtol = 1e-12) + end + # Pin the target to Phase-1's realized SoC path and forbid recourse, so + # the two models describe the same optimization problem. + s1 = battery_solution(p1, r1) + set_tsddr_targets!(p2, Float64[s1.soc[k, t+1] for t in 1:T for k in 1:p2.nBat]) + r2 = solve_tsddr(p2) + @test DecisionRulesExa.solve_succeeded(r2) + d2 = decompose_costs(p2, r2) + # Public recourse energy is projected (≥ 0) — assert nonnegativity directly, + # never via abs() on a value that must not be negative in the first place. + @test 0.0 <= d2.active_deficit_energy_mwh < RECOURSE_ENERGY_TOL_MWH # recourse fixed to zero + @test 0.0 <= d2.active_surplus_energy_mwh < RECOURSE_ENERGY_TOL_MWH + # Same physical operating cost as the Phase-1 objective. + @test isapprox(d2.physical_operating_cost, r1.objective; rtol = 1e-5) + end + + # ── Demand process + paired protocol ───────────────────────────────────── + @testset "exact demand replay and distinct train/eval protocols" begin + case, process = fixture() + H, P = 5, 4 + m1 = scenario_index_matrix(process, H, P; seed = process.train_seed) + m2 = scenario_index_matrix(process, H, P; seed = process.train_seed) + @test m1 == m2 + @test index_matrix_hash(m1) == index_matrix_hash(m2) + @test size(m1) == (H, P) + @test all(1 .<= m1 .<= natom(process)) + me = scenario_index_matrix(process, H, P; seed = process.eval_seed) + @test m1 != me + @test process.train_seed != process.eval_seed + end + + @testset "process validation" begin + case, _ = fixture() + @test_throws ErrorException make_load_process(case; nregion = 0) + @test_throws ErrorException make_load_process(case; train_seed = 5, eval_seed = 5) + @test_throws ErrorException make_load_process(case; nregion = 2, + atoms = [LoadAtom(1.0, [1.0, 1.0]), LoadAtom(1.1, [1.2, 0.9])], probs = [0.3, 0.3]) + @test_throws ErrorException make_load_process(case; nregion = 2, + atoms = [LoadAtom(1.0, [1.0])], probs = [1.0]) + end + + @testset "deterministic region assignment" begin + case, _ = fixture(nregion = 3) + r1, a1 = assign_regions(case.network, 3) + r2, a2 = assign_regions(case.network, 3) + @test r1 == r2 && a1 == a2 + @test length(r1) == nbus(case.network) + @test all(1 .<= r1 .<= 3) + @test length(unique(a1)) == 3 + @test length(Set(r1)) == 3 + end + + @testset "power factor preserved under demand scaling" begin + case, process = fixture() + nd = case.network + de = build_battery_tsddr_de(case, process; reporting_horizon = 3, lookahead = 0, + mode = :soft, stage_hours = 1.0) + set_tsddr_uncertainty!(de, materialize_scenario(process, [2, 3, 1]; horizon = 3)) + for t in 1:3, b in 1:nbus(nd) + nd.bus_pd[b] > 0 || continue + @test isapprox(de.realized_qd[t, b] / de.realized_pd[t, b], + nd.bus_qd[b] / nd.bus_pd[b]; rtol = 1e-12) + end + end + + @testset "protocol write, reconstruct, tamper" begin + case, process = fixture() + H, P = 5, 6 + mat = scenario_index_matrix(process, H, P; seed = process.eval_seed) + dir = mktempdir(); path = joinpath(dir, "eval_protocol.json") + write_scenario_protocol(path, process, mat; kind = "eval", seed = process.eval_seed) + pr2, mat2, meta = reconstruct_scenario_protocol(path) + @test mat2 == mat + @test process_hash(pr2) == process_hash(process) + @test meta.kind == "eval" && meta.horizon == H && meta.paths == P + @test materialize_all(pr2, mat2) == materialize_all(process, mat) + doc = JSON.parsefile(path); doc["index_matrix_hash_sha256"] = repeat("0", 64) + bad = joinpath(dir, "bad.json"); open(io -> JSON.print(io, doc, 2), bad, "w") + @test_throws ErrorException reconstruct_scenario_protocol(bad) + end + + @testset "stochastic manifest reconstruct" begin + case, process = fixture() + dir = mktempdir(); path = joinpath(dir, "manifest.json") + write_stochastic_manifest(path, case, process; + reporting_horizon = 4, lookahead = 2, mode = :soft, stage_hours = 1.0, + active_recourse_cost_per_mwh = DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH, rho1 = 0.0, rho2 = 12.5, + activation = string(stretchedsigmoid), policy_layers = [8], + train_index_matrix_hash = repeat("a", 64), + eval_index_matrix_hash = repeat("b", 64), + train_protocol_file_sha256 = repeat("c", 64), + eval_protocol_file_sha256 = repeat("d", 64), + train_paths = 8, eval_paths = 16) + case2, process2, meta = reconstruct_stochastic_manifest(path) + @test manifest_hash(case2) == manifest_hash(case) + @test process_hash(process2) == process_hash(process) + @test meta.reporting_horizon == 4 && meta.lookahead == 2 && meta.horizon == 6 + @test meta.active_recourse_cost_per_mwh == 10_000.0 + doc = JSON.parsefile(path) + # The misleading load_shedding section is gone; active_recourse replaces it. + @test !haskey(doc, "load_shedding") + ar = doc["active_recourse"] + @test ar["active_recourse_cost_per_mwh"] == 10_000.0 + @test ar["formulation"] == "two-sided active nodal slack" + @test ar["unbounded_above"] == true + @test ar["active_only"] == true + @test ar["included_in_physical_operating_cost"] == true + @test occursin("both directions", ar["accepted_scientific_requirement"]) + @test doc["target_penalty"]["included_in_physical_operating_cost"] == false + @test doc["policy"]["safe_upper_margin"] == 1e-3 + end + + # ── REPAIR 2: reachable policy ─────────────────────────────────────────── + @testset "reachability bounds at interior and boundary" begin + case, _ = fixture() + b = case.batteries[1]; dt = 1.0 + a = 1 - b.sigma * dt + dd = (dt / b.eta_dis) * b.p_discharge_max + cg = b.eta_ch * dt * b.p_charge_max + e = (b.e_min + b.e_max) / 2 + lo, up = battery_reachable_bounds([e], [a], [dd], [cg], [b.e_min], [b.e_max]) + @test lo[1] ≈ max(b.e_min, a*e - dd) + @test up[1] ≈ min(b.e_max, a*e + cg) + @test lo[1] <= up[1] + _, up2 = battery_reachable_bounds([b.e_max], [a], [dd], [cg], [b.e_min], [b.e_max]) + @test up2[1] ≈ b.e_max + lo3, _ = battery_reachable_bounds([b.e_min], [a], [dd], [cg], [b.e_min], [b.e_max]) + @test lo3[1] ≈ b.e_min + end + + @testset "activations: safe upper margin, exact lower edge" begin + # stretchedsigmoid: 0 at large negative, 1-1e-3 (NEVER 1) at large positive. + @test stretchedsigmoid(-50.0) == 0.0 + @test stretchedsigmoid(50.0) == 1.0 - 1e-3 + @test stretchedsigmoid(50.0) < 1.0 + @test 0.4 < stretchedsigmoid(0.0) < 0.6 + # hardsigmoidsafe: same safe margin. + @test hardsigmoidsafe(-50.0) == 0.0 + @test hardsigmoidsafe(50.0) == 1.0 - 1e-3 + @test hardsigmoidsafe(50.0) < 1.0 + @test hardsigmoidsafe(0.0) == 0.5 + # Only bounded activations are admissible. + case, process = fixture() + @test_throws ArgumentError battery_reachable_policy(case, process; layers = [4], + activation = Flux.sigmoid) + end + + @testset "policy targets within reachability; default output range" begin + case, process = fixture() + for act in (stretchedsigmoid, hardsigmoidsafe) + pol = battery_reachable_policy(case, process; dt = 1.0, layers = [8], + activation = act, float_type = Float32) + nw = n_uncertainty(process); nB = nbattery(case) + Random.seed!(1) + for _ in 1:15 + Flux.reset!(pol) + e = Float32.(rand(nB) .* [b.e_max for b in case.batteries]) + y = pol(vcat(Float32.(rand(nw)), e)) + lo, up = battery_reachable_bounds(e, pol.a, pol.discharge_drop, + pol.charge_gain, pol.e_min, pol.e_max) + @test all(lo .- 1f-4 .<= y .<= up .+ 1f-4) + # Never at the exact upper edge (safe margin). + @test all(y .<= up .- (up .- lo) .* 1f-4 .+ 1f-5) + end + end + end + + @testset "no gradient through reachable bounds; gradients reach parameters" begin + case, process = fixture() + pol = battery_reachable_policy(case, process; dt = 1.0, layers = [8], + combiner_layers = Int[], float_type = Float64) + nw = n_uncertainty(process); nB = nbattery(case) + w = Float64.(rand(nw)); e0 = policy_initial_state(case; float_type = Float64) + # Gradients reach encoder and combiner parameters. + g = Zygote.gradient(m -> (Flux.reset!(m); sum(m(vcat(w, e0)))), pol)[1] + @test g.combiner !== nothing && all(isfinite, g.combiner.weight) + @test g.encoder !== nothing + # The bound computation itself carries NO gradient (physical data). + gb = Zygote.gradient(e -> sum(sum(battery_reachable_bounds(e, pol.a, pol.discharge_drop, + pol.charge_gain, pol.e_min, pol.e_max))), + e0)[1] + @test gb === nothing || all(iszero, gb) + end + + @testset "recurrent reset and determinism" begin + case, process = fixture() + pol = battery_reachable_policy(case, process; dt = 1.0, layers = [8], float_type = Float32) + nw = n_uncertainty(process); nB = nbattery(case) + inp = vcat(Float32.(rand(nw)), Float32.(fill(0.5, nB))) + Flux.reset!(pol); y1 = copy(pol(inp)) + Flux.reset!(pol); y2 = copy(pol(inp)) + @test y1 ≈ y2 + Flux.reset!(pol); a1 = copy(pol(inp)); b1 = copy(pol(inp)) + Flux.reset!(pol); a2 = copy(pol(inp)); b2 = copy(pol(inp)) + @test a1 ≈ a2 && b1 ≈ b2 + end + + @testset "finite gradients + finite-difference" begin + case, process = fixture() + pol = battery_reachable_policy(case, process; dt = 1.0, layers = [8], + combiner_layers = Int[], float_type = Float64) + nw = n_uncertainty(process); T = 3 + wflat = Float64.(rand(T * nw)); x0 = policy_initial_state(case; float_type = Float64) + + # (a) Multi-stage rollout gradients are finite. NOTE: a finite difference + # of THIS loss will NOT equal the analytic gradient, and that is correct: + # feeding each target forward makes the next stage's reachable bounds + # depend on the parameters, and the gradient through ℓ and u is + # deliberately stopped (physical projection data, canonical spec). + function rollout_loss(m) + Flux.reset!(m); prev = x0; s = 0.0 + for t in 1:T + xt = m(vcat(view(wflat, (t-1)*nw+1 : t*nw), prev)); s += sum(xt); prev = xt + end + return s + end + gr = Zygote.gradient(rollout_loss, pol)[1] + @test all(isfinite, gr.combiner.weight) && all(isfinite, gr.combiner.bias) + @test gr.encoder !== nothing + + # (b) SINGLE-stage loss with a CONSTANT previous state: no parameter path + # runs through the (stopped) bounds, so the analytic gradient and a finite + # difference must agree. + e_fixed = copy(x0) + w1 = Float64.(wflat[1:nw]) + stage_loss(m) = (Flux.reset!(m); sum(m(vcat(w1, e_fixed)))) + g = Zygote.gradient(stage_loss, pol)[1] + @test all(isfinite, g.combiner.weight) && all(isfinite, g.combiner.bias) + W = pol.combiner.weight; i, j = 1, 1; ε = 1e-6 + w0 = W[i, j] + W[i, j] = w0 + ε; Lp = stage_loss(pol) + W[i, j] = w0 - ε; Lm = stage_loss(pol) + W[i, j] = w0 + @test isapprox((Lp - Lm) / (2ε), g.combiner.weight[i, j]; rtol = 1e-4, atol = 1e-8) + end + + # ── REPAIR 3: two-sided active recourse = complete-recourse slack ───────── + # The recourse device is a per-bus two-sided pair d⁺ (deficit / injection) and + # d⁻ (surplus / absorption), each ≥ 0 and UNBOUNDED above, present at EVERY bus + # (an artificial active-balance recourse, NOT a fraction of local demand: d⁺ + # may exceed local demand or be positive where p^d = 0). They enter the active + # balance as `− d⁺ + d⁻`, guaranteeing the stage subproblem is feasible for + # every incoming state and every reachable target; the reactive balance stays a + # HARD equality (recourse is active-only). Priced at VOLL, so accepted runs ~0 + # in BOTH directions. + @testset "two-sided active recourse: unbounded nodal slack, priced, hard reactive KCL" begin + case, process = mild_fixture() + nd = case.network + de = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 0, + mode = :soft, stage_hours = 1.0) + # Recourse bounds: 0 ≤ d⁺,d⁻ < ∞ at EVERY bus (complete recourse everywhere). + lv = Array(de.model.meta.lvar); uv = Array(de.model.meta.uvar) + T = de.horizon; nB = de.nBus; nG = de.nGen; nBR = de.nBranch; nK = de.nBat + off = variable_offsets(de) + for t in 1:T, b in 1:nB + idx = off.active_deficit + (t-1)*nB + b # deficit block + @test lv[idx] == 0.0 + @test uv[idx] == Inf # unbounded at every bus + sidx = off.active_surplus + (t-1)*nB + b # surplus block + @test lv[sidx] == 0.0 + @test uv[sidx] == Inf + end + # allow_active_recourse=false fixes BOTH slack blocks to zero. + dff = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 0, + mode = :soft, stage_hours = 1.0, allow_active_recourse = false) + uvf = Array(dff.model.meta.uvar) + for t in 1:T, b in 1:nB + @test uvf[off.active_deficit + (t-1)*nB + b] == 0.0 + @test uvf[off.active_surplus + (t-1)*nB + b] == 0.0 + end + # A feasible instance uses ~0 recourse (both directions) and reactive is hard. + e0 = policy_initial_state(case; float_type = Float64) + w = materialize_scenario(process, [1, 1]; horizon = T) + set_tsddr_initial_soc!(de, e0); set_tsddr_uncertainty!(de, w) + set_tsddr_targets!(de, Float64[e for _ in 1:T for e in e0]) + r = solve_tsddr(de) + @test DecisionRulesExa.solve_succeeded(r) + d = decompose_costs(de, r) + # Zero-recourse path: public values are projected (≥ 0) and below tolerance — + # asserted WITHOUT abs (a negative value would itself be a defect). + @test 0.0 <= d.active_deficit_energy_mwh < RECOURSE_ENERGY_TOL_MWH # ~zero deficit + @test 0.0 <= d.active_surplus_energy_mwh < RECOURSE_ENERGY_TOL_MWH # ~zero surplus + @test 0.0 <= d.max_active_deficit_pu < RECOURSE_PU_TOL + @test 0.0 <= d.max_active_surplus_pu < RECOURSE_PU_TOL + # Raw per-variable bound noise stayed inside the declared tolerance. + @test d.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + @test tsddr_max_primal_residual(de, r) < 1e-5 # hard reactive KCL satisfied + + # The two-sided slack provides recourse for BATTERY targets that are NOT + # network-deliverable — the whole point of complete recourse. On case300 + # a strict maximal-CHARGE target (upper reachable endpoint) needs power the + # congested network cannot import: the DEFICIT (d⁺) engages and the strict + # stage STILL SOLVES. A strict maximal-DISCHARGE target (lower endpoint) + # needs power the network cannot export: the SURPLUS (d⁻) engages instead. + c300 = make_battery_case("case300_ieee"; number_of_batteries = 20, seed = 20260722) + p300 = make_load_process(c300; preset = DEFAULT_DEMAND_PRESET, nregion = 3, period = 4) + e300 = Float64.(policy_initial_state(c300; float_type = Float64)) + w300 = materialize_scenario(p300, [1]; horizon = 1) + st300 = build_battery_stage_problem(c300, p300; mode = :strict, stage_hours = 1.0) + set_tsddr_initial_soc!(st300, e300); set_tsddr_uncertainty!(st300, w300) + up = [min(b.e_max, (1 - b.sigma)*e300[k] + b.eta_ch*b.p_charge_max) + for (k, b) in enumerate(c300.batteries)] + lo = [max(b.e_min, (1 - b.sigma)*e300[k] - (1/b.eta_dis)*b.p_discharge_max) + for (k, b) in enumerate(c300.batteries)] + # Maximal charge → the DEFICIT direction provides recourse. + set_tsddr_targets!(st300, up) + rc = solve_tsddr(st300) + @test DecisionRulesExa.solve_succeeded(rc) # recourse ⇒ feasible + dc = decompose_costs(st300, rc); solc = tsddr_solution(st300, rc) + @test dc.active_deficit_pu > 1e-3 # projected deficit engaged + # PUBLIC recourse quantities are ALL finite and ≥ 0 (never negative energy/cost). + for x in (dc.active_deficit_pu, dc.active_surplus_pu, dc.active_deficit_energy_mwh, + dc.active_surplus_energy_mwh, dc.total_active_recourse_energy_mwh, + dc.max_active_deficit_pu, dc.max_active_surplus_pu, dc.active_recourse_cost) + @test isfinite(x) && x >= 0.0 + end + # Raw per-variable bound noise stayed inside the declared tolerance … + @test dc.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + # … and projection is ELEMENTWISE: public pu == Σ max(raw, 0) (not max of the sum). + @test isapprox(dc.active_deficit_pu, sum(x -> max(Float64(x), 0.0), solc.active_deficit_pu); rtol = 1e-12) + @test isapprox(dc.active_surplus_pu, sum(x -> max(Float64(x), 0.0), solc.active_surplus_pu); rtol = 1e-12) + # Raw diagnostics keep the un-projected sums (used for objective reproduction). + @test isapprox(dc.raw_active_deficit_pu, sum(Float64, solc.active_deficit_pu); rtol = 1e-12) + @test isapprox(dc.raw_active_surplus_pu, sum(Float64, solc.active_surplus_pu); rtol = 1e-12) + # Public accounting identities. + @test isapprox(dc.total_active_recourse_energy_mwh, + dc.active_deficit_energy_mwh + dc.active_surplus_energy_mwh; rtol = 1e-12) + @test isapprox(dc.active_recourse_cost, + st300.active_recourse_cost_per_mwh * dc.total_active_recourse_energy_mwh; rtol = 1e-12) + @test isapprox(dc.active_deficit_energy_mwh, st300.baseMVA * st300.dt * dc.active_deficit_pu; rtol = 1e-12) + @test isapprox(dc.active_surplus_energy_mwh, st300.baseMVA * st300.dt * dc.active_surplus_pu; rtol = 1e-12) + @test isapprox(dc.physical_operating_cost, + dc.generator_cost + dc.battery_throughput_cost + dc.active_recourse_cost; rtol = 1e-12) + # Projection correction is EXACTLY the projected-minus-raw recourse cost. + @test isapprox(dc.active_recourse_projection_correction, + dc.active_recourse_cost - dc.raw_active_recourse_cost; rtol = 1e-10, atol = 1e-12) + @test dc.active_recourse_projection_correction >= -1e-12 # projection only removes negatives + # RAW diagnostics reproduce the solver objective (which uses raw primal). + @test isapprox(dc.raw_total_check, dc.total_solver_objective; rtol = 1e-6) + @test dc.solver_objective_recompute_residual < 1e-6 + # Maximal discharge → the SURPLUS direction provides recourse. + set_tsddr_targets!(st300, lo) + rd = solve_tsddr(st300) + @test DecisionRulesExa.solve_succeeded(rd) # recourse ⇒ feasible + dd = decompose_costs(st300, rd) + @test dd.active_surplus_pu > 1e-3 # projected surplus engaged + @test dd.active_deficit_pu < dd.active_surplus_pu # discharge target ⇒ surplus dominates + for x in (dd.active_deficit_pu, dd.active_surplus_pu, dd.active_deficit_energy_mwh, + dd.active_surplus_energy_mwh, dd.active_recourse_cost) + @test isfinite(x) && x >= 0.0 + end + @test dd.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + end + + # ── Reporting correctness: raw vs projected, elementwise, fail-loud ─────── + # The reporting layer must never export a negative recourse quantity: it + # validates the raw per-variable bound noise, fails loudly beyond tolerance, + # and projects the tolerated noise ELEMENTWISE (never on the aggregate, where + # a negative bus could cancel a positive one). + @testset "reporting: elementwise projection, fail-loud, raw/projected split" begin + case, process = mild_fixture() + de = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 0, + mode = :soft, rho1 = 1.0, rho2 = 10.0, stage_hours = 1.0) + e0 = policy_initial_state(case; float_type = Float64) + w = materialize_scenario(process, [1, 1]; horizon = 2) + set_tsddr_initial_soc!(de, e0); set_tsddr_uncertainty!(de, w) + set_tsddr_targets!(de, Float64[e for _ in 1:2 for e in e0]) + r = solve_tsddr(de) + @test DecisionRulesExa.solve_succeeded(r) + + # (a) ELEMENTWISE projection ≠ aggregate projection. Inject a deficit bus at + # −0.4 and another at +1.0 (both within the SAME stage). Aggregate = +0.6 → a + # naive max(Σ,0) keeps 0.6; the correct elementwise Σ max(·,0) keeps 1.0. + sol = tsddr_solution(de, r) + sol.active_deficit_pu .= 0.0; sol.active_surplus_pu .= 0.0 + sol.active_deficit_pu[1, 1] = -0.4 + sol.active_deficit_pu[2, 1] = 1.0 + d = decompose_costs(de, r; sol = sol, lb_tol = 1.0) # permit the injected −0.4 + @test isapprox(d.active_deficit_pu, 1.0; rtol = 1e-12) # Σ max(·,0), NOT 0.6 + @test isapprox(d.raw_active_deficit_pu, 0.6; rtol = 1e-12) # raw keeps the sum + @test d.active_deficit_pu >= 0.0 && d.active_surplus_pu >= 0.0 + @test d.active_deficit_energy_mwh >= 0.0 && d.active_surplus_energy_mwh >= 0.0 + @test d.active_recourse_cost >= 0.0 + # projection correction = projected − raw cost (removes the −0.4 bus). + @test isapprox(d.active_recourse_projection_correction, + d.active_recourse_cost - d.raw_active_recourse_cost; rtol = 1e-12) + @test d.active_recourse_projection_correction > 0.0 # a negative was removed + + # (b) FAIL LOUD: a value further below zero than the declared tolerance is a + # defect, not noise — decompose_costs must raise rather than report it. + solbad = tsddr_solution(de, r) + solbad.active_surplus_pu[1, 1] = -10 * ACTIVE_RECOURSE_LB_TOL_PU + @test_throws ErrorException decompose_costs(de, r; sol = solbad) + # A target-slack value below tolerance also fails loudly. + solbad2 = tsddr_solution(de, r) + solbad2.target_slack_pos[1, 1] = -10 * ACTIVE_RECOURSE_LB_TOL_PU + @test_throws ErrorException decompose_costs(de, r; sol = solbad2) + + # (c) On the genuine solve, everything public is ≥ 0, raw reproduces the + # objective, and the declared per-variable tolerance holds. + dg = decompose_costs(de, r) + for x in (dg.active_deficit_pu, dg.active_surplus_pu, dg.active_deficit_energy_mwh, + dg.active_surplus_energy_mwh, dg.total_active_recourse_energy_mwh, + dg.active_recourse_cost, dg.target_penalty, dg.target_violation, + dg.max_active_deficit_pu, dg.max_active_surplus_pu) + @test isfinite(x) && x >= 0.0 + end + @test dg.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + @test isapprox(dg.raw_total_check, r.objective; rtol = 1e-5) + @test dg.solver_objective_recompute_residual < 1e-4 + @test isapprox(dg.active_recourse_cost, + de.active_recourse_cost_per_mwh * dg.total_active_recourse_energy_mwh; rtol = 1e-12) + end + + # ── Targetless physical diagnostic (no targets at all) ─────────────────── + @testset "targetless diagnostic: no target constraints, slacks, or penalty" begin + case, process = mild_fixture() + T = 2 + tl = build_targetless_diagnostic_de(case, process; reporting_horizon = T, + lookahead = 0, stage_hours = 1.0) + soft = build_battery_tsddr_de(case, process; reporting_horizon = T, lookahead = 0, + mode = :soft, rho1 = 0.0, rho2 = 0.0, stage_hours = 1.0) + @test is_targetless(tl) + @test !is_targetless(soft) + + # (a) NO target constraints and NO target parameter. + @test isempty(tl.target_con_range) + @test tl.p_target === nothing + # Soft carries T*nBat target rows AND 2*T*nBat slack variables; the + # targetless model has neither, so it is strictly smaller by exactly that. + @test length(soft.target_con_range) == T * tl.nBat + @test soft.model.meta.ncon - tl.model.meta.ncon == T * tl.nBat + @test soft.model.meta.nvar - tl.model.meta.nvar == 2 * T * tl.nBat + + # (b) Production API refuses to build it, and target ops reject it. + @test_throws ErrorException build_battery_tsddr_de(case, process; + reporting_horizon = T, mode = :none) + @test_throws ErrorException set_tsddr_targets!(tl, zeros(T * tl.nBat)) + @test_throws ErrorException train_battery_tsddr(nothing, tl, process, + scenario_index_matrix(process, T, 1; seed = 1)) + + e0 = policy_initial_state(case; float_type = Float64) + w = materialize_scenario(process, [1, 1]; horizon = T) + set_tsddr_initial_soc!(tl, e0); set_tsddr_uncertainty!(tl, w) + r = solve_tsddr(tl) + @test DecisionRulesExa.solve_succeeded(r) + @test_throws ErrorException target_multipliers(tl, r) + + # (c) PROJECTED physical = generator + throughput + active-recourse cost, and + # every public recourse quantity is finite and ≥ 0. + d = decompose_costs(tl, r) + @test d.target_penalty == 0.0 && d.target_violation == 0.0 + for x in (d.active_deficit_pu, d.active_surplus_pu, d.active_deficit_energy_mwh, + d.active_surplus_energy_mwh, d.total_active_recourse_energy_mwh, + d.active_recourse_cost) + @test isfinite(x) && x >= 0.0 + end + @test isapprox(d.physical_operating_cost, + d.generator_cost + d.battery_throughput_cost + d.active_recourse_cost; + rtol = 1e-12) + # Public reported total: physical + projected penalty (exact, projected side). + @test isapprox(d.total_check, d.physical_operating_cost + d.target_penalty; rtol = 1e-12) + # RAW diagnostics reproduce the solver objective; projection correction explained. + @test isapprox(d.raw_total_check, r.objective; rtol = 1e-6) + @test d.solver_objective_recompute_residual < 1e-6 + @test isapprox(d.active_recourse_projection_correction, + d.active_recourse_cost - d.raw_active_recourse_cost; rtol = 1e-10, atol = 1e-12) + @test d.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + + # (d) Physics residuals within the existing tolerances. + sol = tsddr_solution(tl, r) + @test tsddr_max_primal_residual(tl, r) < 1e-5 + @test maximum(abs, tsddr_balance_residuals(tl, sol)) < 1e-6 + + # (e) Rebuilding with different (unused) target data cannot change it — + # there is no target data to change, so an identical rebuild reproduces + # the same objective bit-for-bit. + tl2 = build_targetless_diagnostic_de(case, process; reporting_horizon = T, + lookahead = 0, stage_hours = 1.0) + set_tsddr_initial_soc!(tl2, e0); set_tsddr_uncertainty!(tl2, w) + r2 = solve_tsddr(tl2) + @test isapprox(r2.objective, r.objective; rtol = 1e-10) + + # (f) recourse-enabled and recourse-forbidden share every physical + # equation/datum: they differ ONLY in the two-sided nodal-slack upper + # bounds (deficit d⁺ and surplus d⁻), with the same var/con counts. + tlf = build_targetless_diagnostic_de(case, process; reporting_horizon = T, + lookahead = 0, stage_hours = 1.0, + allow_active_recourse = false) + @test tlf.model.meta.nvar == tl.model.meta.nvar + @test tlf.model.meta.ncon == tl.model.meta.ncon + @test tlf.active_recourse_cost_per_mwh == tl.active_recourse_cost_per_mwh + off = variable_offsets(tl) + uv, uvf = Array(tl.model.meta.uvar), Array(tlf.model.meta.uvar) + slack_rng = vcat((off.active_deficit + 1):(off.active_deficit + T * tl.nBus), + (off.active_surplus + 1):(off.active_surplus + T * tl.nBus)) + @test all(uvf[slack_rng] .== 0.0) # forbidden: both fixed to 0 + @test all(uv[slack_rng] .== Inf) # enabled: both unbounded + # every OTHER bound is identical + others = setdiff(1:length(uv), slack_rng) + @test uv[others] == uvf[others] + end + + # ── REPAIR 4: strict / soft target modes ───────────────────────────────── + @testset "strict equality, soft split slacks, cost separation" begin + case, process = mild_fixture() + T = 3 + w = materialize_scenario(process, [2, 1, 3]; horizon = T) + e0 = policy_initial_state(case; float_type = Float64) + pol = battery_reachable_policy(case, process; dt = 1.0, layers = [8], float_type = Float64) + + # Strict: hard equality, no slack, no penalty. With the two-sided active + # recourse providing complete recourse, the strict solve ALWAYS succeeds — + # this is asserted unconditionally (a skipped/logged strict failure is a + # test failure). + de_s = build_battery_tsddr_de(case, process; reporting_horizon = T, lookahead = 0, + mode = :strict, stage_hours = 1.0) + xhat = policy_targets(pol, process, w, e0, T) + set_tsddr_initial_soc!(de_s, e0); set_tsddr_uncertainty!(de_s, w); set_tsddr_targets!(de_s, xhat) + rs = solve_tsddr(de_s) + @test DecisionRulesExa.solve_succeeded(rs) # unconditional + sol_s = tsddr_solution(de_s, rs) + @test maximum(abs, vec(sol_s.soc[:, 2:T+1]) .- xhat) < 1e-5 # strict target residual + @test tsddr_max_primal_residual(de_s, rs) < 1e-6 # primal residual + dS = decompose_costs(de_s, rs) + @test dS.target_penalty == 0.0 && dS.target_violation == 0.0 # no penalty/slack + @test all(iszero, sol_s.target_slack_pos) && all(iszero, sol_s.target_slack_neg) + # Public: physical = gen + throughput + projected recourse; total_check is the + # projected reported total; every public recourse value ≥ 0. + @test isapprox(dS.physical_operating_cost, # strict cost decomposition + dS.generator_cost + dS.battery_throughput_cost + dS.active_recourse_cost; + rtol = 1e-12) + @test isapprox(dS.total_check, dS.physical_operating_cost + dS.target_penalty; rtol = 1e-12) + @test dS.active_recourse_cost >= 0.0 && dS.active_deficit_energy_mwh >= 0.0 && + dS.active_surplus_energy_mwh >= 0.0 + # RAW diagnostics reproduce the solver objective (strict ⇒ no penalty). + @test isapprox(dS.raw_total_check, rs.objective; rtol = 1e-6) + @test dS.solver_objective_recompute_residual < 1e-6 + @test dS.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + @test maximum(abs, tsddr_balance_residuals(de_s, sol_s)) < 1e-6 # battery-balance residual + + # Soft: split slacks, L1+L2 training-only penalty excluded from physical. + de_f = build_battery_tsddr_de(case, process; reporting_horizon = T, lookahead = 0, + mode = :soft, rho1 = 1.0, rho2 = 10.0, stage_hours = 1.0) + set_tsddr_initial_soc!(de_f, e0); set_tsddr_uncertainty!(de_f, w) + set_tsddr_targets!(de_f, policy_targets(pol, process, w, e0, T)) + rf = solve_tsddr(de_f) + @test DecisionRulesExa.solve_succeeded(rf) + sol_f = tsddr_solution(de_f, rf) + dF = decompose_costs(de_f, rf) + @test all(sol_f.target_slack_pos .>= -ACTIVE_RECOURSE_LB_TOL_PU) # raw slacks within tol + @test all(sol_f.target_slack_neg .>= -ACTIVE_RECOURSE_LB_TOL_PU) + # Public target penalty/violation are projected (≥ 0), and the reported total + # is physical + projected penalty. + @test dF.target_penalty >= 0.0 && dF.target_violation >= 0.0 + @test isapprox(dF.total_check, dF.physical_operating_cost + dF.target_penalty; rtol = 1e-12) + # Physical cost = generator + throughput + active recourse (NO target penalty). + @test isapprox(dF.physical_operating_cost, + dF.generator_cost + dF.battery_throughput_cost + dF.active_recourse_cost; + rtol = 1e-12) + @test dF.active_recourse_cost >= 0.0 + # RAW diagnostics reproduce the solver objective WITH the raw penalty term. + @test isapprox(dF.raw_total_check, rf.objective; rtol = 1e-5) + @test dF.solver_objective_recompute_residual < 1e-4 + @test isapprox(dF.reporting_physical_cost + dF.lookahead_physical_cost, + dF.physical_operating_cost; rtol = 1e-10) + @test maximum(min.(sol_f.p_ch, sol_f.p_dis)) < 1e-5 # no simultaneous ch/dis + end + + # ── Strict absolute recourse: EVERY reachable target solves ────────────── + # The invariant: with the two-sided active recourse, the strict operational + # model has complete recourse — for every supported PGLib case and every + # dynamically reachable target (INCLUDING the exact reachable endpoints and + # aggressive charging targets that are NOT network-deliverable), the strict + # stage NLP solves and reproduces the target exactly. Network congestion may + # raise the deficit/cost; it can never make the strict solve fail. + @testset "strict absolute recourse across cases and target classes" begin + function reach1(case, e0, dt) + lo = similar(e0); up = similar(e0) + for (k, b) in enumerate(case.batteries) + a = 1 - b.sigma * dt + lo[k] = max(b.e_min, a*e0[k] - (dt/b.eta_dis)*b.p_discharge_max) + up[k] = min(b.e_max, a*e0[k] + b.eta_ch*dt*b.p_charge_max) + end + lo, up + end + for (cn, nb) in (("case14_ieee", 3), ("case118_ieee", 10), ("case300_ieee", 20)) + case = make_battery_case(cn; number_of_batteries = nb, seed = 20260722) + proc = make_load_process(case; preset = DEFAULT_DEMAND_PRESET, nregion = 3, period = 4) + e0 = Float64.(policy_initial_state(case; float_type = Float64)) + lo, up = reach1(case, e0, 1.0) + w = materialize_scenario(proc, [1]; horizon = 1) + st = build_battery_stage_problem(case, proc; mode = :strict, stage_hours = 1.0) + set_tsddr_initial_soc!(st, e0); set_tsddr_uncertainty!(st, w) + classes = Dict( + "hold" => [(1 - case.batteries[k].sigma)*e0[k] for k in 1:nb], + "midpoint" => (lo .+ up) ./ 2, + "lower" => copy(lo), + "upper" => copy(up), + "interior" => lo .+ 0.5 .* (up .- lo), + "charge90" => lo .+ 0.9 .* (up .- lo), + ) + for (nm, tg) in classes + set_tsddr_targets!(st, tg) + r = solve_tsddr(st) + @test DecisionRulesExa.solve_succeeded(r) # ALWAYS feasible + sol = tsddr_solution(st, r) + @test maximum(abs, vec(sol.soc[:, 2]) .- tg) < 1e-5 # target residual ≤ 1e-5 + @test tsddr_max_primal_residual(st, r) < 1e-5 # primal residual ≤ 1e-5 + @test maximum(abs, tsddr_balance_residuals(st, sol)) < 1e-6 + d = decompose_costs(st, r) + @test d.target_penalty == 0.0 # strict: no penalty + # Every PUBLIC recourse quantity is finite and STRICTLY nonnegative — + # the reporting layer projects bound noise elementwise, so a negative + # public value is impossible (and would fail here, not be tolerated). + for x in (d.active_deficit_pu, d.active_surplus_pu, d.active_deficit_energy_mwh, + d.active_surplus_energy_mwh, d.total_active_recourse_energy_mwh, + d.max_active_deficit_pu, d.max_active_surplus_pu, d.active_recourse_cost) + @test isfinite(x) && x >= 0.0 + end + # Raw per-variable noise stayed inside the declared tolerance, and the + # raw diagnostics reproduce the solver objective. + @test d.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + @test isapprox(d.raw_total_check, d.total_solver_objective; rtol = 1e-6) + @test isapprox(d.total_active_recourse_energy_mwh, + d.active_deficit_energy_mwh + d.active_surplus_energy_mwh; rtol = 1e-12) + @test isfinite(d.physical_operating_cost) + end + end + end + + @testset "target multiplier finite-difference (sign and magnitude)" begin + case, process = mild_fixture() + T = 2 + de = build_battery_tsddr_de(case, process; reporting_horizon = T, lookahead = 0, + mode = :soft, rho1 = 0.0, rho2 = 10.0, stage_hours = 1.0) + w = materialize_scenario(process, [1, 2]; horizon = T) + e0 = policy_initial_state(case; float_type = Float64) + pol = battery_reachable_policy(case, process; dt = 1.0, layers = [8], float_type = Float64) + xhat = policy_targets(pol, process, w, e0, T) + set_tsddr_initial_soc!(de, e0); set_tsddr_uncertainty!(de, w); set_tsddr_targets!(de, xhat) + r0 = solve_tsddr(de) + @test DecisionRulesExa.solve_succeeded(r0) + λ = target_multipliers(de, r0) + i = 1; ε = 1e-4 + xp = copy(xhat); xp[i] += ε; set_tsddr_targets!(de, xp); rp = solve_tsddr(de) + xm = copy(xhat); xm[i] -= ε; set_tsddr_targets!(de, xm); rm = solve_tsddr(de) + @test DecisionRulesExa.solve_succeeded(rp) && DecisionRulesExa.solve_succeeded(rm) + fd = (rp.objective - rm.objective) / (2ε) + @test sign(fd) == sign(λ[i]) || abs(λ[i]) < 1e-6 + @test isapprox(fd, λ[i]; rtol = 5e-2, atol = 1e-2) + end + + @testset "reporting/lookahead horizons recorded" begin + case, process = fixture() + de = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 2, + mode = :soft, stage_hours = 1.0) + @test de.horizon == 4 && de.reporting_horizon == 2 && de.lookahead == 2 + end + + # ── REPAIR 5: no unapproved case modification ──────────────────────────── + @testset "no generator/network modification" begin + case, process = fixture() + nd = case.network + de = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 0, + mode = :soft, stage_hours = 1.0) + # gen_capacity_scale must not exist as a keyword any more. + @test_throws MethodError build_battery_tsddr_de(case, process; reporting_horizon = 2, + gen_capacity_scale = 2.0) + # Generator bounds in the model equal the untouched PGLib per-unit values. + lv = Array(de.model.meta.lvar); uv = Array(de.model.meta.uvar) + T = de.horizon; nB = de.nBus; nG = de.nGen + pg_off = 2*T*nB + for t in 1:T, g in 1:nG + idx = pg_off + (t-1)*nG + g + @test lv[idx] == nd.gens[g].pmin + @test uv[idx] == nd.gens[g].pmax + end + end + + # ── Checkpointing ──────────────────────────────────────────────────────── + @testset "checkpoint exact reload" begin + case, process = fixture() + de = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 0, + mode = :soft, stage_hours = 1.0) + pol = battery_reachable_policy(case, process; dt = 1.0, layers = [8], + combiner_layers = [8], float_type = Float32) + nw = n_uncertainty(process); nB = nbattery(case) + inp = vcat(Float32.(rand(nw)), Float32.(fill(0.4, nB))) + Flux.reset!(pol); y_before = copy(pol(inp)) + dir = mktempdir(); ckpt = joinpath(dir, "ckpt.jls") + save_checkpoint(ckpt, pol, de; case = case, process = process) + pol2, meta = load_checkpoint(ckpt, case, process) + Flux.reset!(pol2); y_after = copy(pol2(inp)) + @test y_after == y_before + @test meta["architecture"]["target_mode"] == "soft" + @test meta["architecture"]["activation"] == string(stretchedsigmoid) + @test meta["architecture"]["safe_upper_margin"] == 1e-3 + @test meta["active_recourse_cost_per_mwh"] == DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH + @test meta["hashes"]["load_process_hash"] == process_hash(process) + end + + # ── CPU/GPU structural parity ──────────────────────────────────────────── + @testset "CPU/GPU model structural parity" begin + case, process = fixture() + cpu = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 0, + mode = :soft, stage_hours = 1.0, backend = nothing) + if CUDA.functional() + gpu = build_battery_tsddr_de(case, process; reporting_horizon = 2, lookahead = 0, + mode = :soft, stage_hours = 1.0, backend = CUDABackend()) + @test cpu.model.meta.nvar == gpu.model.meta.nvar + @test cpu.model.meta.ncon == gpu.model.meta.ncon + @test cpu.target_con_range == gpu.target_con_range + @info "GPU structural parity checked" CUDA.name(CUDA.device()) + else + @info "CUDA not functional; GPU structural-parity check SKIPPED (not run)" + @test cpu.model.meta.ncon > 0 + end + end +end diff --git a/examples/BatteryStorageOPF/test/test_artifacts.jl b/examples/BatteryStorageOPF/test/test_artifacts.jl new file mode 100644 index 0000000..7437e34 --- /dev/null +++ b/examples/BatteryStorageOPF/test/test_artifacts.jl @@ -0,0 +1,220 @@ +# test_artifacts.jl — REPAIR B/C: exact experiment reconstruction, five distinct +# hashes, tampering detection, and a FRESH-PROCESS artifact round trip. +# +# Run (from examples/BatteryStorageOPF): +# julia --pkgimages=no --project=. test/test_artifacts.jl +# +# The round-trip step launches a SEPARATE Julia process that is given only the +# saved artifacts (manifest, evaluation protocol, checkpoint) and must reproduce +# the scenario indices, statuses, cost decomposition, active recourse, and aggregate +# physical cost of the in-process evaluation. + +include(joinpath(@__DIR__, "..", "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using Test +using MadNLP +using Flux +using Random +using JSON +using SHA +using Serialization +using DecisionRulesExa + +const CASE = "case14_ieee" +const REPORT, LOOKAH = 2, 1 +const T = REPORT + LOOKAH +const NEVAL = 4 + +@testset "artifacts: exact reconstruction, hashes, tampering, round trip" begin + dir = mktempdir() + case = make_battery_case(CASE; number_of_batteries = 2, seed = 20260722) + process = make_load_process(case; preset = DEFAULT_DEMAND_PRESET, nregion = 2, period = 4) + + train_mat = scenario_index_matrix(process, T, 4; seed = process.train_seed) + eval_mat = scenario_index_matrix(process, T, NEVAL; seed = process.eval_seed) + train_path = joinpath(dir, "train_protocol.json") + eval_path = joinpath(dir, "eval_protocol.json") + write_scenario_protocol(train_path, process, train_mat; kind = "train", seed = process.train_seed) + write_scenario_protocol(eval_path, process, eval_mat; kind = "eval", seed = process.eval_seed) + man_path = joinpath(dir, "manifest.json") + + h_proc = process_hash(process) + h_train = index_matrix_hash(train_mat) + h_eval = index_matrix_hash(eval_mat) + f_train = bytes2hex(open(sha256, train_path)) + f_eval = bytes2hex(open(sha256, eval_path)) + + write_stochastic_manifest(man_path, case, process; + reporting_horizon = REPORT, lookahead = LOOKAH, mode = :soft, stage_hours = 1.0, + active_recourse_cost_per_mwh = DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH, rho1 = 0.0, rho2 = 0.0, + activation = string(stretchedsigmoid), safe_upper_margin = 1e-3, + policy_layers = [8], policy_combiner_layers = Int[], policy_seed = 7, + train_index_matrix_hash = h_train, eval_index_matrix_hash = h_eval, + train_protocol_file_sha256 = f_train, eval_protocol_file_sha256 = f_eval, + train_paths = 4, eval_paths = NEVAL) + + # ── Five DISTINCT hashes, none reused for another role ─────────────────── + @testset "five distinct hashes" begin + doc = JSON.parsefile(man_path); h = doc["hashes"] + vals = [h["load_process_content_hash_sha256"], h["train_index_matrix_hash_sha256"], + h["eval_index_matrix_hash_sha256"], h["train_protocol_file_sha256"], + h["eval_protocol_file_sha256"]] + @test all(x -> x isa String && length(x) == 64, vals) + @test length(unique(vals)) == 5 # all five differ + @test h["load_process_content_hash_sha256"] == h_proc + # The process hash must NOT masquerade as a protocol hash. + @test h["train_protocol_file_sha256"] != h_proc + @test h["eval_protocol_file_sha256"] != h_proc + @test !haskey(doc, "train_protocol_hash") && !haskey(doc, "eval_protocol_hash") + end + + # ── Exact field-for-field reconstruction ───────────────────────────────── + @testset "exact reconstruction from manifest" begin + case2, proc2, meta = reconstruct_stochastic_manifest(man_path) + @test manifest_hash(case2) == manifest_hash(case) + @test process_hash(proc2) == h_proc + @test proc2.base_shape == process.base_shape # exact vector + @test proc2.region_of_bus == process.region_of_bus + @test proc2.anchor_bus_ids == process.anchor_bus_ids + @test proc2.period == process.period + @test proc2.preset == process.preset + @test [a.system_factor for a in proc2.atoms] == [a.system_factor for a in process.atoms] + @test [a.regional_factors for a in proc2.atoms] == [a.regional_factors for a in process.atoms] + @test proc2.probs == process.probs + @test (proc2.train_seed, proc2.eval_seed) == (process.train_seed, process.eval_seed) + @test meta.reporting_horizon == REPORT && meta.lookahead == LOOKAH + @test meta.active_recourse_cost_per_mwh == DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH + @test meta.activation == string(stretchedsigmoid) + # Materialized scenarios are identical after reconstruction. + @test materialize_all(proc2, eval_mat) == materialize_all(process, eval_mat) + end + + # ── Tampering detection ────────────────────────────────────────────────── + @testset "tampering: base_shape, atoms, indices, protocol files" begin + # (a) base_shape tampered in the manifest → process-hash mismatch. + doc = JSON.parsefile(man_path) + doc["load_process"]["base_shape"][1] += 0.01 + bad = joinpath(dir, "bad_shape.json"); open(io -> JSON.print(io, doc, 2), bad, "w") + @test_throws ErrorException reconstruct_stochastic_manifest(bad) + + # (b) atom definition tampered → process-hash mismatch. + doc = JSON.parsefile(man_path) + doc["load_process"]["atoms"][1]["system_factor"] += 0.05 + bad = joinpath(dir, "bad_atom.json"); open(io -> JSON.print(io, doc, 2), bad, "w") + @test_throws ErrorException reconstruct_stochastic_manifest(bad) + + # (c) ORDER of scenario indices tampered → index-matrix hash mismatch. + pdoc = JSON.parsefile(eval_path) + pdoc["index_matrix"][1], pdoc["index_matrix"][2] = + pdoc["index_matrix"][2], pdoc["index_matrix"][1] + badp = joinpath(dir, "bad_order.json"); open(io -> JSON.print(io, pdoc, 2), badp, "w") + @test_throws ErrorException reconstruct_scenario_protocol(badp) + + # (d) protocol FILE bytes tampered → file-hash mismatch (independent of + # the content hash: reformatting alone changes the bytes). + tampered = joinpath(dir, "eval_reformatted.json") + open(io -> JSON.print(io, JSON.parsefile(eval_path)), tampered, "w") # no indent + @test_throws ErrorException verify_protocol_file(tampered, f_eval) + @test verify_protocol_file(eval_path, f_eval) == f_eval # untampered OK + end + + # ── FRESH-PROCESS round trip ───────────────────────────────────────────── + @testset "round trip in a separate Julia process" begin + # Build + save a tiny checkpoint, evaluate in-process, then re-evaluate in + # a NEW Julia process from the artifacts alone and compare. + de = build_battery_tsddr_de(case, process; reporting_horizon = REPORT, + lookahead = LOOKAH, mode = :soft, + rho1 = 0.0, rho2 = 0.0, stage_hours = 1.0) + stage = build_battery_stage_problem(case, process; mode = :soft, + rho1 = 0.0, rho2 = 0.0, stage_hours = 1.0) + Random.seed!(7) + policy = battery_reachable_policy(case, process; dt = 1.0, layers = [8]) + ckpt = joinpath(dir, "ckpt.jls") + save_checkpoint(ckpt, policy, de; case = case, process = process) + + ev = evaluate_paired(policy, stage, process, eval_mat; + reporting_horizon = REPORT, keep_trajectories = true) + @test ev.n_ok == NEVAL + + script = joinpath(dir, "roundtrip.jl") + open(script, "w") do io + println(io, """ + include(raw"$(joinpath(@__DIR__, "..", "src", "BatteryStorageOPF.jl"))") + using .BatteryStorageOPF, MadNLP, Flux, JSON, DecisionRulesExa + case, process, meta = reconstruct_stochastic_manifest(raw"$man_path") + verify_protocol_file(raw"$eval_path", meta.eval_protocol_file_sha256) + _, emat, _ = reconstruct_scenario_protocol(raw"$eval_path") + policy, ck = load_checkpoint(raw"$ckpt", case, process) + stage = build_battery_stage_problem(case, process; mode = meta.mode, + rho1 = meta.rho1, rho2 = meta.rho2, + active_recourse_cost_per_mwh = meta.active_recourse_cost_per_mwh, + stage_hours = meta.stage_hours) + ev = evaluate_paired(policy, stage, process, emat; + reporting_horizon = meta.reporting_horizon, + keep_trajectories = true) + statuses = String[] + for tr in ev.trajectories, st in tr.trajectory + push!(statuses, st["status"]) + end + open(raw"$(joinpath(dir, "roundtrip_out.json"))", "w") do f + JSON.print(f, Dict( + "process_hash" => process_hash(process), + "index_hash" => index_matrix_hash(emat), + "indices" => [Int.(emat[t, :]) for t in 1:size(emat, 1)], + "mean_cost" => ev.mean_reporting_physical_cost, + "costs" => ev.reporting_physical_costs, + "deficit_mwh" => ev.total_active_deficit_energy_mwh, + "surplus_mwh" => ev.total_active_surplus_energy_mwh, + "n_ok" => ev.n_ok, "statuses" => statuses, + "gen" => [tr.generator_cost for tr in ev.trajectories], + "thr" => [tr.battery_throughput_cost for tr in ev.trajectories], + "recourse_cost" => [tr.active_recourse_cost for tr in ev.trajectories], + "raw_recourse_cost" => [tr.raw_active_recourse_cost for tr in ev.trajectories], + "proj_corr" => [tr.active_recourse_projection_correction for tr in ev.trajectories], + "lb_viol" => ev.maximum_active_recourse_lower_bound_violation_pu, + ), 2) + end + """) + end + outfile = joinpath(dir, "roundtrip_out.json") + jl = joinpath(Sys.BINDIR, "julia") + proj = dirname(Base.active_project()) + run(`$jl --pkgimages=no --project=$proj $script`) + @test isfile(outfile) + got = JSON.parsefile(outfile) + + # Identical scenario indices and hashes. + @test got["process_hash"] == h_proc + @test got["index_hash"] == h_eval + @test [Int.(r) for r in got["indices"]] == [Int.(eval_mat[t, :]) for t in 1:T] + # Identical statuses, recourse, decomposition, and aggregate physical cost. + @test got["n_ok"] == NEVAL + @test all(s -> s == "SOLVE_SUCCEEDED" || s == "SOLVED_TO_ACCEPTABLE_LEVEL", + got["statuses"]) + @test isapprox(Float64(got["mean_cost"]), ev.mean_reporting_physical_cost; rtol = 1e-8) + # Both recourse directions round-trip independently and are nonnegative. + @test 0.0 <= Float64(got["deficit_mwh"]) + @test 0.0 <= Float64(got["surplus_mwh"]) + @test isapprox(Float64(got["deficit_mwh"]), ev.total_active_deficit_energy_mwh; + rtol = 1e-6, atol = 1e-9) + @test isapprox(Float64(got["surplus_mwh"]), ev.total_active_surplus_energy_mwh; + rtol = 1e-6, atol = 1e-9) + # The raw lower-bound violation stayed within the declared tolerance. + @test Float64(got["lb_viol"]) <= ACTIVE_RECOURSE_LB_TOL_PU + for (p, c) in enumerate(got["costs"]) + @test isapprox(Float64(c), ev.reporting_physical_costs[p]; rtol = 1e-8) + end + for (p, tr) in enumerate(ev.trajectories) + @test isapprox(Float64(got["gen"][p]), tr.generator_cost; rtol = 1e-8) + @test isapprox(Float64(got["thr"][p]), tr.battery_throughput_cost; rtol = 1e-8) + # Public recourse cost is ≥ 0 and round-trips; raw diagnostic round-trips too. + @test Float64(got["recourse_cost"][p]) >= 0.0 + @test isapprox(Float64(got["recourse_cost"][p]), tr.active_recourse_cost; + rtol = 1e-6, atol = 1e-9) + @test isapprox(Float64(got["raw_recourse_cost"][p]), tr.raw_active_recourse_cost; + rtol = 1e-6, atol = 1e-9) + @test isapprox(Float64(got["proj_corr"][p]), tr.active_recourse_projection_correction; + rtol = 1e-6, atol = 1e-9) + end + end +end diff --git a/examples/BatteryStorageOPF/test/test_e2e_training.jl b/examples/BatteryStorageOPF/test/test_e2e_training.jl new file mode 100644 index 0000000..c1e1b89 --- /dev/null +++ b/examples/BatteryStorageOPF/test/test_e2e_training.jl @@ -0,0 +1,150 @@ +# test_e2e_training.jl — tiny, fixed-seed, end-to-end battery AC-OPF TS-DDR +# training test in the PRIMARY (strict) target mode. +# +# It must: +# * train from a fresh fixed initialization in STRICT mode (hard ê=e, no target +# slack, no target penalty; the strict equality multipliers drive training); +# * complete with accepted solver statuses (failed solves counted, not hidden) — +# the two-sided active nodal recourse gives strict complete recourse so every +# initial, training, and held-out solve succeeds; +# * carry zero active recourse in BOTH directions on the held-out paths +# (feasible demand): deficit d⁺ and surplus d⁻ each below tolerance, and their +# sum below twice that (asserted separately — a single "zero deficit" check is +# insufficient because it inspects only d⁺); +# * reduce the fixed held-out mean PHYSICAL OPERATING COST from initialization +# (physical = generator + battery throughput + two-sided active-recourse cost; +# the target has no penalty in strict mode); +# * save and reload a checkpoint reproducing the policy output AND both recourse +# directions exactly. +# +# Run (from examples/BatteryStorageOPF): +# julia --pkgimages=no --project=. test/test_e2e_training.jl + +include(joinpath(@__DIR__, "..", "src", "BatteryStorageOPF.jl")) +using .BatteryStorageOPF +using Test +using MadNLP +using Flux +using Random +using DecisionRulesExa + +# Declared numerical tolerance for "zero active recourse" (interior-point barrier +# tolerance leaves each recourse power near, not exactly, zero). Deficit d⁺ and +# surplus d⁻ are each held below this; their sum below 2×. +const RECOURSE_ENERGY_TOL_MWH = 1e-3 + +@testset "tiny end-to-end battery TS-DDR training (strict)" begin + case = make_battery_case("case14_ieee"; number_of_batteries = 2, seed = 20260722, + fleet_power_fraction = 0.4, duration_hours = 4.0) + # Mild, demand-feasible process (peak ≈1.1×) so an accepted run needs no recourse. + # (The public default preset is selected by the demand calibration gates.) + mild_atoms = [LoadAtom(1.00, [1.00, 1.00]), + LoadAtom(1.03, [1.05, 0.99]), + LoadAtom(1.03, [0.99, 1.05])] + process = make_load_process(case; nregion = 2, period = 4, base_amplitude = 0.04, + atoms = mild_atoms, probs = [0.5, 0.25, 0.25], + train_seed = 4242, eval_seed = 9999) + REPORT, LOOKAH = 3, 1 + T = REPORT + LOOKAH + MODE = :strict # PRIMARY mode + + train_mat = scenario_index_matrix(process, T, 8; seed = process.train_seed) + eval_mat = scenario_index_matrix(process, T, 8; seed = process.eval_seed) + @test train_mat != eval_mat + + de = build_battery_tsddr_de(case, process; reporting_horizon = REPORT, + lookahead = LOOKAH, mode = MODE, stage_hours = 1.0) + stage = build_battery_stage_problem(case, process; mode = MODE, stage_hours = 1.0) + @test de.mode === :strict && stage.mode === :strict + @test !isempty(de.target_con_range) # strict target rows exist + # Strict carries NO target slack variables (soft would add 2*T*nBat): it has + # the SAME variables as the targetless diagnostic, plus T*nBat target rows. + tldiag = build_targetless_diagnostic_de(case, process; reporting_horizon = REPORT, + lookahead = LOOKAH, stage_hours = 1.0) + @test de.model.meta.nvar == tldiag.model.meta.nvar + @test de.model.meta.ncon == tldiag.model.meta.ncon + de.horizon * de.nBat + + Random.seed!(20260722) # fresh, fixed initialization + policy = battery_reachable_policy(case, process; dt = 1.0, layers = [32, 32], + combiner_layers = [32]) + + fixed_in = vcat(Float32.([0.7, 1.05, 0.95]), Float32.([2.0, 3.0])) + + ev0 = evaluate_paired(policy, stage, process, eval_mat; reporting_horizon = REPORT) + @test ev0.n_ok == size(eval_mat, 2) # all held-out solves accepted + # Zero active recourse at init — deficit d⁺ and surplus d⁻ asserted SEPARATELY, + # and their sum below 2×. A single d⁺-only check would be insufficient. + @test 0.0 <= ev0.total_active_deficit_energy_mwh < RECOURSE_ENERGY_TOL_MWH + @test 0.0 <= ev0.total_active_surplus_energy_mwh < RECOURSE_ENERGY_TOL_MWH + @test 0.0 <= ev0.total_active_recourse_energy_mwh < 2 * RECOURSE_ENERGY_TOL_MWH + @test ev0.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + init_cost = ev0.mean_reporting_physical_cost + @test isfinite(init_cost) + + tr = train_battery_tsddr(policy, de, process, train_mat; + num_batches = 60, num_train_per_batch = 8, + optimizer = Flux.Adam(1f-2), + madnlp_kwargs = (print_level = MadNLP.ERROR, tol = 1e-6), + record_loss = (i, m, l, t) -> false) + @test tr.n_ok == tr.n_total # every training solve accepted + @test tr.n_failed == 0 + + # Strict equality multipliers are what training consumes: on a solved stage + # they must be finite and their count equals the target block. + let w1 = materialize_scenario(process, [eval_mat[1, 1]]; horizon = 1), + e1 = policy_initial_state(case; float_type = Float64) + Flux.reset!(policy) + t1 = Float64.(policy(vcat(Float32.(w1), Float32.(e1)))) + set_tsddr_initial_soc!(stage, e1); set_tsddr_uncertainty!(stage, w1); set_tsddr_targets!(stage, t1) + r1 = MadNLP.madnlp(stage.model; print_level = MadNLP.ERROR, tol = 1e-6) + @test DecisionRulesExa.solve_succeeded(r1) + λ = target_multipliers(stage, r1) + @test length(λ) == stage.nBat && all(isfinite, λ) + # Strict mode carries NO target slack and NO target penalty. Public recourse + # is projected (≥ 0); the RAW diagnostics reproduce the solver objective. + d1 = decompose_costs(stage, r1) + s1 = tsddr_solution(stage, r1) + @test d1.target_penalty == 0.0 && d1.target_violation == 0.0 + @test all(iszero, s1.target_slack_pos) && all(iszero, s1.target_slack_neg) + @test d1.active_recourse_cost >= 0.0 && d1.active_deficit_energy_mwh >= 0.0 && + d1.active_surplus_energy_mwh >= 0.0 + @test d1.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + @test isapprox(d1.raw_total_check, r1.objective; rtol = 1e-6) + @test d1.solver_objective_recompute_residual < 1e-6 + @test isapprox(d1.total_check, d1.physical_operating_cost + d1.target_penalty; rtol = 1e-12) + end + + ev1 = evaluate_paired(policy, stage, process, eval_mat; reporting_horizon = REPORT) + @test ev1.n_ok == size(eval_mat, 2) + # Still zero recourse in BOTH directions after training (asserted separately). + # Public recourse is projected (≥ 0): assert nonnegative-and-below-tolerance + # directly, never abs() on a quantity that must not be negative. + @test 0.0 <= ev1.total_active_deficit_energy_mwh < RECOURSE_ENERGY_TOL_MWH + @test 0.0 <= ev1.total_active_surplus_energy_mwh < RECOURSE_ENERGY_TOL_MWH + @test 0.0 <= ev1.total_active_recourse_energy_mwh < 2 * RECOURSE_ENERGY_TOL_MWH + @test ev1.maximum_active_recourse_lower_bound_violation_pu <= ACTIVE_RECOURSE_LB_TOL_PU + final_cost = ev1.mean_reporting_physical_cost + @info "tiny e2e (strict)" init_cost final_cost improvement = init_cost - final_cost + @info "tiny e2e recourse (MWh)" deficit = ev1.total_active_deficit_energy_mwh surplus = ev1.total_active_surplus_energy_mwh + @test final_cost < init_cost # reduced held-out PHYSICAL cost + + # Checkpoint save + exact reload reproduction on a fixed CPU input. + Flux.reset!(policy); y_before = copy(policy(fixed_in)) + dir = mktempdir(); ckpt = joinpath(dir, "e2e.jls") + save_checkpoint(ckpt, policy, de; case = case, process = process) + pol2, meta = load_checkpoint(ckpt, case, process) + Flux.reset!(pol2); y_after = copy(pol2(fixed_in)) + @test y_after == y_before + @test meta["hashes"]["case_manifest_content_hash"] == manifest_hash(case) + @test meta["architecture"]["activation"] == string(stretchedsigmoid) + @test meta["active_recourse_cost_per_mwh"] == DEFAULT_ACTIVE_RECOURSE_COST_PER_MWH + + # The reloaded policy reproduces BOTH recourse directions exactly on the fixed + # held-out paths (deterministic CPU rollout ⇒ identical deficit and surplus). + ev1r = evaluate_paired(pol2, stage, process, eval_mat; reporting_horizon = REPORT) + @test isapprox(ev1r.total_active_deficit_energy_mwh, ev1.total_active_deficit_energy_mwh; + rtol = 1e-8, atol = 1e-12) + @test isapprox(ev1r.total_active_surplus_energy_mwh, ev1.total_active_surplus_energy_mwh; + rtol = 1e-8, atol = 1e-12) + @test isapprox(ev1r.mean_reporting_physical_cost, final_cost; rtol = 1e-8) +end diff --git a/examples/HydroPowerModels/Project.toml b/examples/HydroPowerModels/Project.toml index 30e1a40..d211cd1 100644 --- a/examples/HydroPowerModels/Project.toml +++ b/examples/HydroPowerModels/Project.toml @@ -12,10 +12,14 @@ JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" MadNLP = "2621e9c9-9eb4-46b1-8089-e8c72242dfb6" MadNLPGPU = "d72a61cc-809d-412f-99be-fd81f4b8a598" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c" Wandb = "ad70616a-06c9-5745-b1f1-6a5f42545108" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" cuDNN = "02a925ec-e4fe-4b08-9a7e-0d78e3d38ccd" -[sources] -DecisionRulesExa = {path = "../.."} +[extras] +CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2" + +[sources.DecisionRulesExa] +path = "../.." diff --git a/examples/HydroPowerModels/README.md b/examples/HydroPowerModels/README.md index 65f91a6..8e4b2b8 100644 --- a/examples/HydroPowerModels/README.md +++ b/examples/HydroPowerModels/README.md @@ -1,111 +1,28 @@ -# HydroPowerModels Example - -Multi-stage hydrothermal scheduling using DecisionRulesExa.jl with DC or AC OPF formulations. - -## Problem description - -A hydro-dominated power system (Bolivia test case) is operated over a planning horizon of up to 96 stages. At each stage, the operator must decide generator dispatch, reservoir outflows, and spillage subject to: - -- **Power flow constraints** (DC linearization or full AC polar OPF) -- **Reservoir dynamics** (water balance with stochastic inflows) -- **Generator and transmission limits** - -The TS-DDR policy (an LSTM network) predicts target reservoir levels at each stage. The deterministic-equivalent NLP projects these targets onto the feasible set via slack-penalized target constraints. Training uses envelope-theorem gradients: dual multipliers on the target constraints give the policy gradient without differentiating through the solver. - -## Formulations - -Set `FORMULATION` in `train_hydro_exa.jl`: - -| Formulation | `FORMULATION` | Variables per stage | Description | -|---|---|---|---| -| DC OPF | `:dc` | ~500 | Linear power flow, fast solves | -| AC Polar OPF | `:ac_polar` | ~2000 | Full nonlinear AC power flow | - -## Data - -The `bolivia/` directory contains: - -- `PowerModels.json` — power system topology (39 buses, 55 branches, 19 generators) -- `hydro.json` — hydro unit parameters (7 reservoirs) -- `inflows.csv` — historical inflow scenarios (144 stages x 200 scenarios x 7 reservoirs) -- `_demand.csv` — per-stage bus demand scaling - -Pre-solved deterministic-equivalent references (MOF format) are provided for validation: -- `DCPPowerModel.mof.json` -- `ACPPowerModel.mof.json` - -## Files - -| File | Description | -|---|---| -| `train_hydro_exa.jl` | Main training script with penalty scheduling, parallel GPU solves, and W&B logging | -| `train_hydro_exa_critic.jl` | Critic/control-variate variant of the main training script; uses normalized hydro features, a replay buffer, and cheap critic rollouts | -| `hydro_power_data.jl` | Data parsing (PowerModels JSON, hydro JSON, inflows CSV) | -| `hydro_power_exa.jl` | ExaModels problem builder for DC and AC OPF formulations | -| `eval_exa_de.jl` | Validation script comparing ExaModels results against JuMP reference | -| `Project.toml` | Example-specific dependencies (W&B, JLD2, CUDA, etc.) | - -## Running - -### GPU training (recommended) - -```julia -# From this directory: -julia --project -t auto train_hydro_exa.jl -``` - -Set `USE_GPU = true` in `train_hydro_exa.jl` (default). Requires a CUDA-capable GPU. - -### GPU training with critic control variate - -```julia -# From this directory: -julia --project -t auto train_hydro_exa_critic.jl -``` - -The critic script keeps the dual-multiplier actor update but adds a damped -control variate (`critic_cv_weight = 0.5`) trained on the stage-wise rollout -objective without target penalty. Its default critic rollout uses -`policy_state = :target`; set `CRITIC_POLICY_STATE = :realized` for closed-loop -critic labels. Deterministic-equivalent critic fitting remains available as an -ablation through `DeterministicEquivalentCriticTarget()`. - -### CPU training - -Set `USE_GPU = false` in `train_hydro_exa.jl`, then run the same command. - -### Configuration - -Key parameters in `train_hydro_exa.jl`: - -| Parameter | Default | Description | -|---|---|---| -| `FORMULATION` | `:ac_polar` | OPF formulation (`:dc` or `:ac_polar`) | -| `NUM_STAGES` | 96 | Planning horizon | -| `NUM_EPOCHS` | 20 | Training epochs | -| `NUM_BATCHES` | 100 | Gradient steps per epoch | -| `NUM_WORKERS` | 4 | Parallel GPU solver instances | -| `LAYERS` | `[128, 128]` | LSTM hidden layer sizes | -| `LR` | 1e-3 | Learning rate | -| `DEFICIT_COST` | 1e5 | Load-shedding penalty ($/pu) | - -### Training features - -- **Penalty scheduling**: target penalty multiplier ramps through phases (0.1 -> 1.0 -> 10.0 -> 30.0) over training -- **Sample scheduling**: `num_train_per_batch` increases from `NUM_WORKERS` to `8 * NUM_WORKERS` -- **Evaluation scheduling**: rollout evaluation starts with 4 scenarios and ramps to 32 at halfway -- **Parallel solves**: independent NLP copies solved concurrently via `Threads.@spawn` worker pool -- **Parallel rollout**: evaluation scenarios distributed across CPU stage-problem copies -- **Critic variant**: optional scalar critic with value and gradient matching, - replay-buffer training, and cheap critic actor samples -- **W&B logging**: training loss, rollout objectives, violation share, penalty multiplier - -## Validation - -Compare the ExaModels formulation against a JuMP/MadNLP reference: - -```julia -julia --project -t auto eval_exa_de.jl -``` - -This loads a pre-solved JuMP reference and solves the same problem in ExaModels, printing a side-by-side comparison of objectives and reservoir trajectories. +# Bolivia hydro example (ExaModels) + +This directory is the ExaModels counterpart of the canonical Bolivia MAIN +hydro example in DecisionRules.jl. + +Canonical invariants: + +- identical MAIN input bytes in both repositories; +- `pd_scale = qd_scale = 0.6`; +- weekly stages (`stage_hours = 168`) and `K = 0.6048`; +- operational active deficit cost `6000 USD/(pu·stage)`; +- strict reachable targets using the stretched sigmoid with a safe `1e-3` + upper margin; +- hard reactive balance and apparent-power thermal limits at both branch ends; +- 96 reporting stages, 30 look-ahead stages, and one saved 126-by-500 + stage-major joint inflow-and-demand protocol. + +The ACP, DC, and SOC-WR MOFs are generated once by the DecisionRules.jl +exporter and copied here without reserialization. Corresponding files must +therefore be byte-identical across the repositories. + +`bolivia/case_manifest.json` records hashes, constants, topology, objective +metadata, and protocol seeds. `bolivia/joint_protocol_500.csv` records the +paired inflow-scenario and demand-atom indices. + +The single retained checkpoint and historical strict result are provenance +references, not final paired evidence. Phase 2A does not recreate missing SDDP +cuts or run production training. diff --git a/examples/HydroPowerModels/compare_paired_evals.jl b/examples/HydroPowerModels/compare_paired_evals.jl new file mode 100644 index 0000000..8456476 --- /dev/null +++ b/examples/HydroPowerModels/compare_paired_evals.jl @@ -0,0 +1,259 @@ +# compare_paired_evals.jl +# +# Verdict report for the cross-package equivalence check between +# DecisionRules.jl (MAIN, JuMP/Ipopt) and DecisionRulesExa.jl (EXA, +# ExaModels/MadNLP) on the 100 paired scenarios. +# +# Loads: +# 1. MAIN ground truth: paired_strict_rollout.jld2 (eval_paired_tsddr.jl) +# 2. MAIN policy reference: paired_policy_reference.jld2 (dump_paired_policy_reference.jl) +# 3. EXA evaluation: paired_exa_strict.jld2 (eval_paired_exa_strict.jl) +# 4. MAIN paired_costs.csv (for the SDDP column) +# +# and prints: +# (a) policy-parity gate results, +# (b) per-scenario cost deltas MAIN-stagewise vs EXA-stagewise, +# (c) EXA-stagewise vs EXA-DE per-scenario deltas, +# (d) trajectory deltas (vol / gen / policy targets), +# (e) mean costs of every leg side by side with SDDP, +# (f) the explicit list of known structural differences with measured +# activity/inertness. +# +# Agreement classification is honest and threshold-based: +# EXACT : relative deviation < 1e-6 +# TIGHT : relative deviation < 1e-3 (solver-tolerance regime) +# DISCREPANT : anything larger +# +# Usage: +# julia --project compare_paired_evals.jl + +using JLD2 +using Statistics +using CSV, Tables + +const SCRIPT_DIR = dirname(@__FILE__) +const MAIN_HPM_DIR = "/storage/scratch1/9/arosemberg3/DecisionRules.jl/examples/HydroPowerModels" + +const MAIN_RESULTS = joinpath(MAIN_HPM_DIR, "bolivia", "ACPPowerModel", "results", + "paired_strict_rollout.jld2") +const MAIN_REFERENCE = joinpath(MAIN_HPM_DIR, "bolivia", "ACPPowerModel", "results", + "paired_policy_reference.jld2") +const MAIN_COSTS_CSV = joinpath(MAIN_HPM_DIR, "bolivia", "ACPPowerModel", "paired_costs.csv") +const EXA_RESULTS = joinpath(SCRIPT_DIR, "bolivia", "ACPPowerModel", "results", + "paired_exa_strict.jld2") + +""" + classify(rel::Real) -> String + +Classify a relative deviation: `"EXACT"` below `1e-6`, `"TIGHT"` below `1e-3` +(solver-tolerance regime), `"DISCREPANT"` otherwise (`"NaN"` if not finite). +""" +function classify(rel::Real) + isfinite(rel) || return "NaN" + rel < 1e-6 && return "EXACT" + rel < 1e-3 && return "TIGHT" + return "DISCREPANT" +end + +""" + report_deltas(label, a, b) -> Float64 + +Print mean/max absolute and relative deviations between two equal-length +vectors (NaN entries dropped pairwise), classify the max relative deviation, +and return it. Relative deviation is `|a - b| / max(|a|, 1)` so near-zero +values do not explode the ratio. +""" +function report_deltas(label, a, b) + mask = .!isnan.(a) .& .!isnan.(b) + n = count(mask) + if n == 0 + println(" $label: no comparable entries (all NaN)") + return NaN + end + av, bv = Float64.(a[mask]), Float64.(b[mask]) + absd = abs.(av .- bv) + reld = absd ./ max.(abs.(av), 1.0) + println(" $label [n=$n]") + println(" mean |Δ| = $(mean(absd)) max |Δ| = $(maximum(absd))") + println(" mean rel = $(mean(reld)) max rel = $(maximum(reld)) → $(classify(maximum(reld)))") + return maximum(reld) +end + +# ── Load everything ──────────────────────────────────────────────────────────── + +main = JLD2.load(MAIN_RESULTS) +ref = JLD2.load(MAIN_REFERENCE) +exa = JLD2.load(EXA_RESULTS) + +main_costs = Float64.(main["costs"]) # [S] +main_vol = Float64.(main["vol_trajectories"]) # [T × S] +main_gen = Float64.(main["gen_trajectories"]) # [T × S] +main_idx = Int.(main["scenario_indices"]) # [T × S] + +exa_costs = Float64.(exa["costs"]) # [S] +exa_vol = Float64.(exa["vol_trajectories"]) # [T × S] +exa_gen = Float64.(exa["gen_trajectories"]) # [T × S] +exa_ok = Bool.(exa["scenario_ok"]) +de_costs = Float64.(exa["de_costs"]) # [N] +de_ok = Bool.(exa["de_ok"]) +n_de = Int(exa["num_de_scenarios"]) + +T, S = size(main_vol) + +# SDDP column from MAIN's paired_costs.csv (quoted headers contain commas, so +# CSV.jl is required for parsing). +sddp_costs = Float64[] +if isfile(MAIN_COSTS_CSV) + cols = Tables.columntable(CSV.File(MAIN_COSTS_CSV)) + sddp_key = findfirst(k -> occursin("SDDP", String(k)), collect(keys(cols))) + if sddp_key !== nothing + global sddp_costs = Float64.(collect(cols[collect(keys(cols))[sddp_key]])) + end +end + +println("=" ^ 72) +println("CROSS-PACKAGE PAIRED EVALUATION — VERDICT REPORT") +println(" MAIN results: $MAIN_RESULTS") +println(" MAIN reference: $MAIN_REFERENCE") +println(" EXA results: $EXA_RESULTS") +println(" T=$T stages, S=$S scenarios, DE scenarios N=$n_de") +println("=" ^ 72) + +# Sanity: both sides must have evaluated the same scenario index matrix. +ref_idx = Int.(ref["scenario_indices"]) +idx_match = main_idx == ref_idx +println("\nScenario-index matrices identical (MAIN results vs reference): $idx_match") +idx_match || println(" !! The two evaluations did NOT use the same scenarios — all cost/trajectory comparisons below are void.") + +# ── (a) Policy-parity gate ───────────────────────────────────────────────────── +println("\n(a) POLICY PARITY GATE (from eval_paired_exa_strict.jl)") +println(" checkpoint loader path: $(exa["gate_loader_mode"])", + exa["gate_loader_mode"] == "stock" ? "" : + " ← EXA's stock loader cannot load MAIN checkpoints (LSTM wrapper vs bare LSTMCell state)") +println(" probe parity (single calls): ", + exa["gate_probe_pass"] ? "PASS" : "FAIL", + " max|Δ| = $(exa["gate_probe_max_dev"])") +println(" open-loop, EXA policy as-is: ", + exa["gate_openloop_asis_pass"] ? "PASS" : "FAIL", + " max|Δ| = $(exa["gate_openloop_asis_max_dev"])") +println(" open-loop, state-threaded: ", + exa["gate_openloop_threaded_pass"] ? "PASS" : "FAIL", + " max|Δ| = $(exa["gate_openloop_threaded_max_dev"])") +if exa["gate_probe_pass"] && !exa["gate_openloop_asis_pass"] && exa["gate_openloop_threaded_pass"] + println(" → INTERPRETATION: weights load correctly; the divergence is the EXA") + println(" policy's MEMORYLESS LSTM encoding (Flux 0.16 restarts recurrent") + println(" state on every call; Flux.reset! is a no-op) vs MAIN's explicit") + println(" state threading across stages. The two packages evaluate") + println(" DIFFERENT functions of the inflow history with the same weights.") +end +println(" inflow indexing/units: ", + exa["gate_inflow_pass"] ? "PASS (exact)" : "FAIL", + " max|Δ| = $(exa["gate_inflow_recon_max_dev"])") +println(" metadata devs: K=$(exa["gate_dev_K"]) min_vol=$(exa["gate_dev_min_vol"]) " * + "max_vol=$(exa["gate_dev_max_vol"]) min_turn=$(exa["gate_dev_min_turn"]) " * + "max_turn=$(exa["gate_dev_max_turn"]) upstream_max=$(exa["gate_dev_upstream_max"])") + +# ── (b) MAIN stage-wise vs EXA stage-wise costs ──────────────────────────────── +println("\n(b) PER-SCENARIO COSTS: MAIN stage-wise (Ipopt) vs EXA stage-wise (MadNLP)") +println(" NOTE: if gate (a) shows the policies compute different targets, part of") +println(" this delta is the policy-function difference, not the subproblem builder.") +report_deltas("total cost per scenario", main_costs, exa_costs) +n_fail = count(.!exa_ok) +n_fail > 0 && println(" EXA stage-wise scenarios failed (excluded): $n_fail") + +# ── (c) EXA stage-wise vs EXA DE ─────────────────────────────────────────────── +println("\n(c) EXA STAGE-WISE vs EXA STRICT FULL-HORIZON DE (first $n_de scenarios)") +println(" Strict-mode claim: identical targets pin the same reservoir path, so the") +println(" DE objective must equal the stage-wise sum for the same scenario.") +report_deltas("total cost per scenario", exa_costs[1:n_de], de_costs) +# Per-stage decomposition of the DE objective vs the stage-wise stage costs. +sc = Float64.(exa["stage_costs"])[:, 1:n_de] +dsc = Float64.(exa["de_stage_costs"]) +report_deltas("per-stage costs (all t, s ≤ N)", vec(sc), vec(dsc)) +n_de_fail = count(.!de_ok) +n_de_fail > 0 && println(" EXA DE scenarios failed (excluded): $n_de_fail") + +# ── (d) Trajectory deltas ────────────────────────────────────────────────────── +println("\n(d) TRAJECTORY DELTAS across all (t, s)") +report_deltas("vol_trajectories (Σ_r volume/0.0036, MW) MAIN vs EXA-stagewise", + vec(main_vol), vec(exa_vol)) +report_deltas("gen_trajectories (Σ_thermal pg·baseMVA) MAIN vs EXA-stagewise", + vec(main_gen), vec(exa_gen)) +# Policy target paths: MAIN open-loop reference vs EXA realized targets. +xhat_ref = Float64.(ref["xhat_trajectories"]) # [T × nHyd × S] +tgt_exa = Float64.(exa["target_trajectories"]) # [nHyd × T × S] +tgt_exa_perm = permutedims(tgt_exa, (2, 1, 3)) # → [T × nHyd × S] +report_deltas("policy target trajectories (all t, r, s) MAIN-ref vs EXA", + vec(xhat_ref), vec(tgt_exa_perm)) +report_deltas("DE vol_trajectories vs EXA-stagewise (s ≤ N)", + vec(Float64.(exa["vol_trajectories"])[:, 1:n_de]), + vec(Float64.(exa["de_vol_trajectories"]))) +report_deltas("DE gen_trajectories vs EXA-stagewise (s ≤ N)", + vec(Float64.(exa["gen_trajectories"])[:, 1:n_de]), + vec(Float64.(exa["de_gen_trajectories"]))) + +# ── (e) Mean costs side by side ──────────────────────────────────────────────── +println("\n(e) MEAN COSTS ($S scenarios; DE over first $n_de)") +_mean_ok(v) = (m = v[.!isnan.(v)]; isempty(m) ? NaN : mean(m)) +println(" MAIN stage-wise strict (Ipopt): $(round(_mean_ok(main_costs); digits=1))") +println(" EXA stage-wise strict (MadNLP): $(round(_mean_ok(exa_costs); digits=1))") +println(" EXA strict DE (MadNLP, N=$n_de): $(round(_mean_ok(de_costs); digits=1))") +println(" MAIN stage-wise mean over s ≤ $n_de: $(round(_mean_ok(main_costs[1:n_de]); digits=1))") +if !isempty(sddp_costs) + println(" SDDP (paired_costs.csv): $(round(_mean_ok(sddp_costs); digits=1))") +else + println(" SDDP column not found in $MAIN_COSTS_CSV") +end +println(" rollout_tsddr cross-check (s=1): $(exa["rollout_tsddr_check_objective"]) vs custom loop $(exa_costs[1])") + +# ── (f) Known structural differences, with measured activity ────────────────── +println("\n(f) KNOWN STRUCTURAL DIFFERENCES (EXA builder vs MAIN mof.json subproblem)") + +max_def_sw = maximum(filter(!isnan, Float64.(exa["max_deficit"])); init = -Inf) +max_def_de = maximum(filter(!isnan, Float64.(exa["de_max_deficit"])); init = -Inf) +max_dq_sw = maximum(filter(!isnan, Float64.(exa["max_abs_deficit_q"])); init = -Inf) +max_dq_de = maximum(filter(!isnan, Float64.(exa["de_max_abs_deficit_q"])); init = -Inf) + +println(""" + 1. DEFICIT COST: MAIN objective uses 6000·Σ deficit (= cost_deficit 60 × baseMVA + 100, scaled at mof export). This evaluation used deficit_cost = + $(exa["deficit_cost_used"]) to match. EXA TRAINING uses 1e5 instead. + Measured max deficit: stage-wise = $max_def_sw pu, DE = $max_def_de pu. + → the 1e5-vs-6000 difference is $(max(max_def_sw, max_def_de) <= 1e-8 ? "INERT (deficit never activates)" : "ACTIVE — deficit occurs, costs are NOT comparable to training runs"). + 2. REACTIVE SLACK: EXA adds a FREE zero-cost deficit_q to every reactive KCL; + MAIN's mof.json has hard reactive balance. Not disableable via kwargs. + Measured max |deficit_q|: stage-wise = $max_dq_sw pu, DE = $max_dq_de pu. + → $(max(max_dq_sw, max_dq_de) <= 1e-6 ? "INERT on these scenarios" : "ACTIVE — the EXA AC feasible set is genuinely relaxed vs MAIN (reactive balance violated at zero cost); expect lower EXA costs"). + 3. THERMAL LIMITS: MAIN enforces quadratic p²+q² ≤ rate_a² per branch end + (62 quadratic constraints in mof.json); EXA box-bounds p_fr/q_fr/p_to/q_to + in [−rate_a, rate_a] — a relaxation in the (|p|,|q|) corners. Not measured + directly here; shows up as cost differences when branch limits bind. + 4. MIN-VIOLATION SLACKS: mof.json has zero-cost min_outflow/min_volume + violation slacks; EXA enforces outflow ≥ min_turn hard. Bolivia has + min_turn ≡ 0 and min_vol ≡ 0 → inert. + 5. GEN COSTS: identical by construction (linear c1 per pu, c2 = 0, from the + same PowerModels.json; verified against mof.json coefficients). + Turbine coupling identical: baseMVA·pg = pf·outflow. + 6. DEMAND: mof.json bakes in 0.6× PowerModels loads (pd and qd); EXA used + load_scaler = $(exa["load_scaler_used"]) → matched. + 7. SPILL UNITS: both sides use spill with coefficient 1 (volume units) in the + water balance and K = 0.0036 on inflow/outflow → identical dynamics. + 8. POLICY RECURRENCE: see gate (a). EXA's encoder is memoryless per stage + (Flux 0.16 LSTM restarts from initialstates each call); MAIN threads LSTM + state across stages. Same weights, different function. This affects EXA + TRAINING and evaluation alike: the EXA pipeline optimizes/evaluates a + policy without inflow memory. + 9. CHECKPOINT LOADER: loader path used = "$(exa["gate_loader_mode"])". If not + "stock", DecisionRulesExa's load_stateconditioned_policy! cannot ingest + MAIN checkpoints (MAIN saves bare LSTMCell states; EXA wraps cells in + Flux.LSTM) — cross-package warmstarts silently depend on a custom loader. + 10. SOLVERS/TOLERANCES: MAIN = Ipopt (mumps, default tol); EXA = MadNLP + (tol = $(exa["solver_tol"])). Agreement at TIGHT (<1e-3 rel) is the + expected ceiling for cost comparisons even with identical models. + 11. INITIAL STATE: EXA training clamps x0 into [min_vol, max_vol]; this + evaluation used MAIN's raw initial_state (denormal ≈ 1e-316 values ≈ 0); + difference ≤ 1e-315 hm³ → inert. +""") +println("=" ^ 72) +println("END OF REPORT") +println("=" ^ 72) diff --git a/examples/HydroPowerModels/eval_paired_exa_strict.jl b/examples/HydroPowerModels/eval_paired_exa_strict.jl new file mode 100644 index 0000000..2ee76ac --- /dev/null +++ b/examples/HydroPowerModels/eval_paired_exa_strict.jl @@ -0,0 +1,1073 @@ +# eval_paired_exa_strict.jl +# +# Cross-package equivalence check against DecisionRules.jl (the JuMP/Ipopt MAIN +# repo). Evaluates the SAME stage-wise-strict TS-DDR checkpoint on the SAME 100 +# paired scenarios that MAIN's eval_paired_tsddr.jl evaluated, using +# +# (a) the stage-wise strict ExaModels rollout (1-stage strict problems, +# closed-loop realized-state feedback), and +# (b) the strict regular full-horizon deterministic equivalent (DE), +# +# and records per-scenario objectives and operative solution values for +# comparison against MAIN's ground truth (paired_strict_rollout.jld2). +# +# DISCREPANCIES ARE THE DELIVERABLE. Nothing here is tuned to force agreement; +# every structural difference is measured and saved. Known structural +# differences handled/documented here (see also compare_paired_evals.jl): +# +# 1. DEFICIT COST. MAIN's ACPPowerModel.mof.json objective is +# Σ_g c1_g·pg_g + 6000·Σ_b deficit_b (linear; c2 = 0 for all g) +# where 6000 = cost_deficit (60 $/MWh) × baseMVA (100): HydroPowerModels +# scales the deficit cost by baseMVA when exporting the subproblem. The +# EXA builder uses `deficit_cost` per pu directly, so this script passes +# deficit_cost = power_data.cost_deficit * power_data.baseMVA = 6000.0. +# NOTE: train_hydro_exa_strict.jl trains with DEFICIT_COST = 1e5 instead — +# inert iff deficit never activates (max deficit is recorded per scenario). +# 2. DEMAND. The mof.json bakes in 0.6 × PowerModels.json loads (see MAIN's +# export_subproblem_mof.jl, which scales pd AND qd by 0.6). The EXA builder +# reproduces this with load_scaler = 0.6 and demand_matrix = nothing. +# 3. REACTIVE SLACK. By default the EXA AC builder adds a FREE, zero-cost +# `deficit_q` variable to every reactive KCL; MAIN's mof.json has no +# reactive slack (hard reactive balance). The builder kwarg +# `reactive_deficit_cost` now controls this; this script reads the +# DR_REACTIVE_DEFICIT env var (default "hard" → Inf, i.e. NO reactive +# slack — MAIN-faithful; "free" → nothing reproduces the historical +# relaxation; a number → linear |deficit_q| penalty at that cost) and +# passes it to BOTH the stage problem and the DE leg. max |deficit_q| is +# still recorded per scenario (exact zeros in hard mode). +# 4. THERMAL LIMITS. mof.json enforces quadratic branch limits +# p² + q² ≤ rate_a² (62 ScalarQuadraticFunction ≤ constraints); the EXA +# builder only box-bounds each of p_fr, q_fr, p_to, q_to in +# [−rate_a, rate_a] — a superset of the disk (relaxation in the corners). +# 5. MIN-VIOLATION SLACKS. mof.json has min_outflow_violation / +# min_volume_violation slack variables with ZERO objective cost (they only +# relax the ≥ min bounds); the EXA builder enforces outflow ≥ min_turn +# hard. For Bolivia min_turn ≡ 0, so both are inert. +# 6. POLICY RECURRENCE (measured by the parity gate below). RESOLVED: the +# EXA HydroReachablePolicy now threads the LSTM state across stages +# explicitly (`DecisionRulesExa._step_encoder`, mirroring MAIN's +# `DecisionRules._step_encoder`), and `Flux.reset!(policy)` is a REAL +# reset back to `Flux.initialstates`. The as-is open-loop gate below is +# therefore expected to PASS with max|Δ| = 0.0 against MAIN. The +# independent manually-threaded diagnostic (gate 4) is retained as a +# cross-check: if as-is and threaded ever disagree, the policy forward +# pass has regressed from MAIN's semantics. +# +# Requires: MAIN's paired_policy_reference.jld2 (produced by +# dump_paired_policy_reference.jl in the MAIN repo) — run that job first. +# +# Environment variables: +# DR_DE_SCENARIOS = "10" (number of scenarios for the full-horizon DE leg) +# DR_MAX_ITER = "9000" +# DR_REACTIVE_DEFICIT = "hard" ("hard" → Inf = no reactive slack (MAIN-faithful); +# "free" → nothing = historical free slack; +# a number → linear |deficit_q| cost) +# DR_CHECKPOINT = (checkpoint to evaluate; default: the +# MAIN reference checkpoint hardcoded below) +# DR_CHECKPOINT_KIND = "main" ("main" → apply ALL MAIN-reference parity gates; +# "exa" → EXA-trained checkpoint: SKIP the +# probe-parity and open-loop-trajectory gates +# (only meaningful for the exact reference +# weights) but KEEP the policy-independent +# inflow-indexing gate; scenario inflow values +# still come from the reference JLD2) +# DR_ENCODER_LAYERS = "128,128" (LSTM encoder widths; must match checkpoint; +# DR_LAYERS is the legacy alias, as in +# train_hydro_exa_strict.jl) +# DR_HEAD_LAYERS = "" (state-conditioned head hidden widths; must +# match checkpoint; "" → linear head) +# DR_CONTEXT = "" (""/"none", "phase", or "phase+progress"; +# defaults to the reference file metadata +# when present) +# DR_CONTEXT_HORIZON = "126" (denominator/horizon used for progress +# context; defaults to reference metadata) +# DR_REFERENCE_FILE = (paired_policy_reference*.jld2 from MAIN) +# DR_OUTPUT_TAG = "" ("" → save to results/paired_exa_strict.jld2; +# "" → results/paired_exa_strict_.jld2 +# with checkpoint path + all knobs recorded) +# +# Usage: +# julia --project -t auto eval_paired_exa_strict.jl + +using DecisionRulesExa +using ExaModels +using MadNLP +using Flux +using Statistics, Random +using JLD2 + +const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) # parse_layers +include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) +include(joinpath(SCRIPT_DIR, "hydro_reachable_policy.jl")) + +# ── Configuration ───────────────────────────────────────────────────────────── + +const CASE_NAME = "bolivia" +const FORMULATION = :ac_polar +const FORM_LABEL = "ACPPowerModel" + +const CASE_DIR = joinpath(SCRIPT_DIR, CASE_NAME) +const PM_FILE = joinpath(CASE_DIR, "PowerModels.json") +const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") +const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") + +# Absolute paths into the MAIN (DecisionRules.jl) repository. +const MAIN_HPM_DIR = "/storage/scratch1/9/arosemberg3/DecisionRules.jl/examples/HydroPowerModels" +const DEFAULT_REFERENCE_FILE = joinpath( + MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "results", "paired_policy_reference.jld2" +) +const REFERENCE_FILE = get(ENV, "DR_REFERENCE_FILE", DEFAULT_REFERENCE_FILE) +# Default checkpoint: the MAIN reference checkpoint against which the parity +# gates below were designed. DR_CHECKPOINT overrides it with any other +# checkpoint (MAIN- or EXA-trained). +const DEFAULT_MODEL_PATH = joinpath( + MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "models", + "bolivia-ACPPowerModel-h126-r96-subproblems-strict-2026-07-01T09:41:53.026.jld2", +) +const MODEL_PATH = get(ENV, "DR_CHECKPOINT", DEFAULT_MODEL_PATH) +# Checkpoint kind: "main" (DecisionRules.jl-trained; MAIN-reference parity +# gates apply) or "exa" (train_hydro_exa_strict.jl-trained; the probe-parity +# and open-loop-trajectory gates are SKIPPED because the reference probe +# outputs / trajectories were produced by the specific reference checkpoint — +# comparing independently trained weights against them is meaningless. The +# inflow-indexing gate is policy-independent and is KEPT, and scenario inflow +# values are still sourced from the reference JLD2 as authoritative data). +const CHECKPOINT_KIND = lowercase(strip(get(ENV, "DR_CHECKPOINT_KIND", "main"))) +CHECKPOINT_KIND in ("main", "exa") || + error("DR_CHECKPOINT_KIND must be \"main\" or \"exa\", got \"$CHECKPOINT_KIND\"") +const MAIN_PARITY_GATES = CHECKPOINT_KIND == "main" + +# Architecture knobs — parsed exactly as train_hydro_exa_strict.jl parses them +# (including the DR_LAYERS legacy alias); they MUST match the checkpoint. +const ENCODER_LAYERS = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +const HEAD_LAYERS = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) + +# Output tag: "" keeps the historical filename results/paired_exa_strict.jld2; +# a non-empty tag saves to results/paired_exa_strict_.jld2 so evaluating a +# new checkpoint never clobbers the reference results. +const OUTPUT_TAG = String(strip(get(ENV, "DR_OUTPUT_TAG", ""))) +const OUT_SUFFIX = isempty(OUTPUT_TAG) ? "" : "_$(OUTPUT_TAG)" +# Record the new provenance knobs in the JLD2 whenever any of them was +# explicitly set; with all of them unset the saved file keeps exactly the +# historical key set (byte-identical default behavior). +const RECORD_KNOBS = any(haskey.(Ref(ENV), + ("DR_CHECKPOINT", "DR_CHECKPOINT_KIND", "DR_ENCODER_LAYERS", "DR_LAYERS", + "DR_HEAD_LAYERS", "DR_CONTEXT", "DR_CONTEXT_HORIZON", "DR_REFERENCE_FILE", + "DR_OUTPUT_TAG", "DR_SNAP_EPS", "DR_ACTIVATION"))) +# mof.json demand = 0.6 × PowerModels.json pd/qd (export_subproblem_mof.jl). +const LOAD_SCALER = 0.6 +const NUM_DE_SCENARIOS = parse(Int, get(ENV, "DR_DE_SCENARIOS", "10")) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +# Reactive-slack control (header item 3): "hard" → Inf (no deficit_q, +# MAIN-faithful), "free" → nothing (historical free slack), number → linear +# |deficit_q| penalty at that cost. Passed to both the stage problem and the +# DE leg builders. +const REACTIVE_DEFICIT_RAW = lowercase(strip(get(ENV, "DR_REACTIVE_DEFICIT", "hard"))) +const REACTIVE_DEFICIT_COST = REACTIVE_DEFICIT_RAW == "hard" ? Inf : + REACTIVE_DEFICIT_RAW == "free" ? nothing : + parse(Float64, REACTIVE_DEFICIT_RAW) + +# Stochastic demand (bolivia/demand_scenarios.csv, single line `s,`): +# i.i.d. per-stage multiplicative demand factor ξ_t ∈ {1−s, 1, 1+s} (P = 1/3), +# independent of the inflow noise — the same model the SDDP baselines register +# via sddp/sddp_demand_noise.jl and train_hydro_exa_strict.jl trains under. +# For the PAIRED protocol, scenario column c uses the SEEDED demand path +# StableRNG(DEMAND_NOISE_SEED + c) (protocol_demand_factors), identical to the +# path the trainer's protocol eval and SDDP.Historical paired evaluator draw +# for that column. Thus the comparison is paired over the complete joint +# inflow-and-demand path, not merely distributional over demand. +# When the file is absent every path below is bit-identical to the historical +# evaluation. +const DEMAND_SPREAD = load_demand_spread(joinpath(CASE_DIR, "demand_scenarios.csv")) +const DEMAND_NOISE = DEMAND_SPREAD !== nothing +# MAIN reference checkpoints have policy input width 2·nHyd (no ξ slot); the +# probe/open-loop parity gates are meaningless and would crash on the wider +# demand-noise policy, so demand noise requires DR_CHECKPOINT_KIND=exa. +DEMAND_NOISE && MAIN_PARITY_GATES && + error("demand_scenarios.csv present: demand-noise evaluation requires DR_CHECKPOINT_KIND=exa " * + "(MAIN reference parity gates only apply to nHyd-input checkpoints)") +DEMAND_NOISE && @info "Stochastic demand ACTIVE (seeded paired demand paths)" DEMAND_SPREAD +# CPU MadNLP: same pattern as train_hydro_exa_strict.jl's SOLVER_KWARGS, but +# this evaluation runs on a CPU node (backend = nothing everywhere). +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) + +# ── Load data ───────────────────────────────────────────────────────────────── + +@info "Loading power system data..." +power_data = load_power_data(PM_FILE) +@info " nBus=$(power_data.nBus) nGen=$(power_data.nGen) baseMVA=$(power_data.baseMVA) cost_deficit=$(power_data.cost_deficit)" + +# Deficit cost matching MAIN's mof.json objective coefficient (see header, item 1): +# 6000 = cost_deficit ($/MWh) × baseMVA, applied per pu of shed load. +const DEFICIT_COST = power_data.cost_deficit * power_data.baseMVA +@info " deficit_cost used for equivalence: $DEFICIT_COST (training used 1e5)" + +@info "Loading hydro data (num_stages = full inflow history)..." +# num_stages = nothing keeps the RAW rows of inflows.csv (47 for Bolivia — +# fewer than the 96 eval stages). MAIN's read_inflow tiles those rows +# vertically for longer horizons (load_hydropowermodels.jl:13-21), so stage t +# maps to raw row mod1(t, nrows); the reconstruction gate below applies the +# same cyclic convention. +hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; num_stages = nothing) +nHyd = hydro_data.nHyd +@info " nHyd=$nHyd nScenarios=$(hydro_data.nScenarios) K=$(hydro_data.K)" + +# ── Load the MAIN policy reference ──────────────────────────────────────────── + +isfile(REFERENCE_FILE) || error( + "Missing reference file $REFERENCE_FILE — run dump_paired_policy_reference.jl " * + "in the MAIN repo first." +) +ref = JLD2.load(REFERENCE_FILE) +const T_EVAL = Int(ref["num_eval_stages"]) +const NUM_SCEN = Int(ref["num_scenarios"]) +inflow_ref = ref["inflow_values"] # [T × nHyd × S] authoritative values +xhat_ref = ref["xhat_trajectories"] # [T × nHyd × S] MAIN open-loop targets +scen_idx = Int.(ref["scenario_indices"]) # [T × S] +x0_ref = Float64.(ref["initial_state"]) +probe_in = Float32.(ref["probe_inputs"]) # [2 nHyd × 3] +probe_out_ref = ref["probe_outputs"] # [nHyd × 3] +@info "Loaded reference: T=$T_EVAL, S=$NUM_SCEN from $REFERENCE_FILE" +@assert size(inflow_ref) == (T_EVAL, nHyd, NUM_SCEN) +@assert length(x0_ref) == nHyd + +# Seeded paired demand paths: column s of the protocol gets the demand factors +# ξ[:, s] = protocol_demand_factors(s) — the identical path the trainer's +# protocol eval uses for that column (column-keyed StableRNG seeding; see +# hydro_power_data.jl). `nothing` when demand is deterministic. +const demand_factors = DEMAND_NOISE ? + reduce(hcat, [protocol_demand_factors(DEMAND_SPREAD, T_EVAL, s) for s in 1:NUM_SCEN]) : + nothing # [T_EVAL × NUM_SCEN] or nothing + +""" + augmented_stage_w(w_t::AbstractVector, t::Int, s::Int) -> Vector{Float64} + +Stage-t uncertainty block for paired scenario `s`: the inflow vector `w_t` +alone (deterministic demand), or `[w_t; ξ_t^{(s)}]` with the seeded paired +demand factor appended (stochastic demand). +""" +augmented_stage_w(w_t::AbstractVector, t::Int, s::Int) = + DEMAND_NOISE ? vcat(Float64.(w_t), demand_factors[t, s]) : Float64.(w_t) + +""" + augmented_flat_w(w_flat::AbstractVector, s::Int) -> Vector{Float64} + +Full-horizon stage-major uncertainty vector for paired scenario `s`: the flat +inflow trajectory unchanged (deterministic demand), or with the seeded paired +demand factor ξ_t^{(s)} interleaved into each stage block (stochastic demand). +""" +augmented_flat_w(w_flat::AbstractVector, s::Int) = + DEMAND_NOISE ? augment_scenario(Float64.(w_flat), demand_factors[:, s]) : Float64.(w_flat) + +# Per-stage uncertainty width used by the rollout machinery below. +const N_UNC = nHyd + (DEMAND_NOISE ? 1 : 0) + +ref_string(key::AbstractString, default::AbstractString) = + haskey(ref, key) ? string(ref[key]) : default +ref_int(key::AbstractString, default::Int) = + haskey(ref, key) ? Int(ref[key]) : default + +const CONTEXT_MODE = canonical_context_mode( + get(ENV, "DR_CONTEXT", ref_string("context_mode", "")) +) +const CONTEXT_PERIOD = ref_int("context_period", countlines(INFLOW_FILE)) +const CONTEXT_HORIZON = parse( + Int, + get(ENV, "DR_CONTEXT_HORIZON", string(ref_int("context_horizon", 126))), +) +CONTEXT_HORIZON >= T_EVAL || + error("DR_CONTEXT_HORIZON=$CONTEXT_HORIZON must cover reference T_EVAL=$T_EVAL") +const STAGE_CONTEXT = build_stage_context(CONTEXT_MODE, CONTEXT_HORIZON, CONTEXT_PERIOD) +const N_CONTEXT = isnothing(STAGE_CONTEXT) ? 0 : size(STAGE_CONTEXT, 1) +@info "Policy context" context_mode=(isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE) CONTEXT_PERIOD CONTEXT_HORIZON N_CONTEXT + +# ── Build the EXA policy and load the requested checkpoint ──────────────────── +# Constructor call mirrors train_hydro_exa_strict.jl exactly (sigmoid activation +# and Flux.LSTM encoder are the constructor defaults); combiner_layers must +# match the checkpoint's head architecture. + +# Target-head activation must match the checkpoint's training activation +# (DR_ACTIVATION, same values as the trainer: sigmoid|hardsigmoid|stretched). +const ACTIVATION = let raw = lowercase(strip(get(ENV, "DR_ACTIVATION", "sigmoid"))) + raw in ("", "sigmoid") ? Flux.NNlib.sigmoid : + raw == "hardsigmoid" ? hardsigmoidsafe : + raw == "stretched" ? stretchedsigmoid : + error("DR_ACTIVATION must be sigmoid, hardsigmoid, or stretched; got $raw") +end +# Diagnostic eval-time snap-to-boundary for sigmoid-trained checkpoints +# (DR_SNAP_EPS > 0): normalized targets within ε of 0/1 become exact boundary +# points — measures the cost of sigmoid's asymptotic boundary gap without any +# retraining. See TARGET_SNAP_EPS in hydro_reachable_policy.jl. +const SNAP_EPS = parse(Float32, get(ENV, "DR_SNAP_EPS", "0")) +0 <= SNAP_EPS < 0.5f0 || error("DR_SNAP_EPS must be in [0, 0.5); got $SNAP_EPS") +TARGET_SNAP_EPS[] = SNAP_EPS +SNAP_EPS > 0 && @info "Eval-time snap-to-boundary ACTIVE" SNAP_EPS + +Random.seed!(42) +base_policy = hydro_reachable_policy( + hydro_data, + ENCODER_LAYERS; + activation = ACTIVATION, + combiner_layers = HEAD_LAYERS, + n_context = N_CONTEXT, + # Demand-noise checkpoints were trained with the encoder observing ξ_t; + # the constructor must match the checkpoint's input width. + n_extra_uncertainty = DEMAND_NOISE ? 1 : 0, +) +policy = isnothing(STAGE_CONTEXT) ? base_policy : ContextualPolicy(base_policy, STAGE_CONTEXT) +policy_core(policy) = policy isa ContextualPolicy ? policy.policy : policy +stage_encoder_input(policy, t::Int, w) = + policy isa ContextualPolicy ? vcat(Float32.(context_at(policy.context, t)), w) : w +@info "Policy built" encoder_layers = ENCODER_LAYERS head_layers = HEAD_LAYERS checkpoint_kind = CHECKPOINT_KIND context_mode=(isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE) +isfile(MODEL_PATH) || error("Checkpoint not found: $MODEL_PATH (set DR_CHECKPOINT)") + +""" + load_main_checkpoint!(policy, model_state) -> String + +Load a DecisionRules.jl (MAIN) checkpoint into the EXA `HydroReachablePolicy`, +returning a string describing which loading path succeeded. + +# Arguments +- `policy`: EXA `HydroReachablePolicy` (encoder is a `Chain` of `Flux.LSTM` + wrapper layers, each holding a `.cell`). +- `model_state`: the checkpoint's `Flux.state`, whose `encoder` field was saved + from MAIN's `Chain` of BARE `LSTMCell`s (MAIN's `_as_cell` strips the `LSTM` + wrapper), i.e. `state.encoder.layers[i] = (Wi = …, Wh = …, bias = …)`. + +# Returns +- `"stock"` if the package loader `load_stateconditioned_policy!` succeeded. +- `"cell-by-cell"` if the stock loader failed and the weights were injected + directly into each `LSTM.cell` plus the combiner. + +# Notes +`load_stateconditioned_policy!` now includes the cell-by-cell MAIN-checkpoint +fallback natively (`DecisionRulesExa._load_encoder_state!`), so the stock path +is expected to SUCCEED and report `"stock"`. EXA-trained checkpoints +(train_hydro_exa_strict.jl saves `Flux.state(cpu(m))` of the same +`HydroReachablePolicy` type, encoder stored as `Flux.LSTM` wrappers) also load +through the stock path. This wrapper's own cell-by-cell +branch is retained as belt-and-braces; reaching it would indicate a loader +regression and is recorded in the output. The cell-by-cell path performs the +mathematically identical weight injection: the wrapper's `cell` has exactly +the fields `(Wi, Wh, bias)` MAIN saved. +""" +function load_main_checkpoint!(policy, model_state) + try + load_stateconditioned_policy!(policy, model_state) + return "stock" + catch err + @warn "Stock EXA loader failed on the MAIN checkpoint (expected: MAIN saves bare LSTMCell states, EXA wraps cells in Flux.LSTM). Falling back to cell-by-cell injection." exception = err + core = policy_core(policy) + inner_state = hasproperty(model_state, :policy) ? getproperty(model_state, :policy) : model_state + enc_state = getproperty(inner_state, :encoder) + layer_states = getproperty(enc_state, :layers) + length(layer_states) == length(core.encoder.layers) || + error("Encoder depth mismatch: checkpoint has $(length(layer_states)) layers, policy has $(length(core.encoder.layers))") + for (layer, lstate) in zip(core.encoder.layers, layer_states) + # MAIN layer state keys (Wi, Wh, bias) match LSTMCell's fields. + Flux.loadmodel!(layer.cell, lstate) + end + Flux.loadmodel!(core.combiner, getproperty(inner_state, :combiner)) + Flux.reset!(policy) + return "cell-by-cell" + end +end + +model_state = JLD2.load(MODEL_PATH, "model_state") +loader_mode = load_main_checkpoint!(policy, model_state) +@info "Checkpoint loaded via: $loader_mode ($MODEL_PATH)" + +# ── Policy-parity gate ──────────────────────────────────────────────────────── + +""" + _max_abs_dev(a, b) -> Float64 + +Maximum absolute element-wise deviation `max_i |a_i − b_i|` between two arrays +of equal size. +""" +_max_abs_dev(a, b) = maximum(abs.(Float64.(a) .- Float64.(b))) + +gate = Dict{String, Any}("loader_mode" => loader_mode) +core_policy = policy_core(policy) + +# (0) Frozen hydro-metadata parity: the bounds the two policies scale into. +gate["dev_K"] = abs(core_policy.K - Float64(ref["policy_K"])) +gate["dev_min_vol"] = _max_abs_dev(core_policy.min_vol, ref["policy_min_vol"]) +gate["dev_max_vol"] = _max_abs_dev(core_policy.max_vol, ref["policy_max_vol"]) +gate["dev_min_turn"] = _max_abs_dev(core_policy.min_turn, ref["policy_min_turn"]) +gate["dev_max_turn"] = _max_abs_dev(core_policy.max_turn, ref["policy_max_turn"]) +gate["dev_upstream_max"] = _max_abs_dev(core_policy.upstream_max_inflow, ref["policy_upstream_max"]) +@info "Metadata deviations" gate["dev_K"] gate["dev_min_vol"] gate["dev_max_vol"] gate["dev_min_turn"] gate["dev_max_turn"] gate["dev_upstream_max"] + +# (1) Probe parity: single policy calls from the reset state. Tests weight +# loading independent of any recurrence-threading semantics. ONLY meaningful +# when evaluating the exact MAIN reference checkpoint the probe outputs were +# generated from — SKIPPED for EXA-trained checkpoints. +probe_out_exa = fill(NaN, nHyd, 3) +if MAIN_PARITY_GATES + for p in 1:3 + Flux.reset!(policy) # REAL reset: each probe starts from initialstates + probe_out_exa[:, p] = Float64.(policy(probe_in[:, p])) + end + gate["probe_outputs_exa"] = probe_out_exa + gate["probe_max_dev"] = _max_abs_dev(probe_out_exa, probe_out_ref) + # Float32 tolerance: relative to the target scale (max_vol up to ~138). + probe_pass = gate["probe_max_dev"] <= 1e-5 * max(1.0, maximum(abs.(probe_out_ref))) + gate["probe_pass"] = probe_pass + println(probe_pass ? "GATE probe parity: PASS" : "GATE probe parity: FAIL", + " (max abs dev = $(gate["probe_max_dev"]))") +else + # The reference probe outputs are the REFERENCE checkpoint's responses; + # an independently trained EXA checkpoint has different weights, so a + # weight-parity comparison would fail by construction and prove nothing. + gate["probe_max_dev"] = NaN + gate["probe_pass"] = "skipped" + println("GATE probe parity: SKIPPED (DR_CHECKPOINT_KIND=exa — reference " * + "probe outputs only characterize the MAIN reference checkpoint)") +end + +# (2) Inflow reconstruction: rebuild w[t, r, s] from the EXA loader's +# scenario_inflows and the reference scenario indices; compare against the +# reference values. This validates scenario indexing and units across loaders +# (the historically buggy spot). +inflow_recon_dev = 0.0 +first_mismatch = nothing +# MAIN's read_inflow tiles the raw inflow rows vertically when num_stages +# exceeds the file length (load_hydropowermodels.jl:13-21), so stage t maps to +# raw row mod1(t, nrows). The EXA loader keeps the raw (untiled) matrix; apply +# the same cyclic convention here so both sides index identical physical data. +n_inflow_rows = size(hydro_data.scenario_inflows[1], 1) +for s in 1:NUM_SCEN, t in 1:T_EVAL, r in 1:nHyd + v_exa = hydro_data.scenario_inflows[r][mod1(t, n_inflow_rows), scen_idx[t, s]] + dev = abs(v_exa - inflow_ref[t, r, s]) + if dev > inflow_recon_dev + global inflow_recon_dev = dev + global first_mismatch = (t = t, r = r, s = s, exa = v_exa, ref = inflow_ref[t, r, s]) + end +end +gate["inflow_recon_max_dev"] = inflow_recon_dev +gate["inflow_pass"] = inflow_recon_dev == 0.0 +if gate["inflow_pass"] + println("GATE inflow indexing: PASS (exact reconstruction)") +else + println("GATE inflow indexing: FAIL (max abs dev = $inflow_recon_dev at $first_mismatch)") + # Diagnose common failure patterns at the first mismatching coordinate. + # All probes use cyclic row indexing so t beyond the raw row count cannot + # itself throw while diagnosing an indexing mismatch. + t, r, s = first_mismatch.t, first_mismatch.r, first_mismatch.s + ω = scen_idx[t, s] + tr = mod1(t, n_inflow_rows) + println(" diagnostics at (t=$t → raw row $tr, r=$r, s=$s, ω=$ω):") + println(" ref value = $(inflow_ref[t, r, s])") + println(" exa [tr, ω] = $(hydro_data.scenario_inflows[r][tr, ω])") + ω <= size(hydro_data.scenario_inflows[r], 1) && tr <= size(hydro_data.scenario_inflows[r], 2) && + println(" exa transposed [ω, tr] = $(hydro_data.scenario_inflows[r][ω, tr])") + println(" exa row-offset [tr+1, ω] = $(hydro_data.scenario_inflows[r][mod1(tr + 1, n_inflow_rows), ω])") + println(" exa unit ratio (exa/ref) = $(hydro_data.scenario_inflows[r][tr, ω] / inflow_ref[t, r, s])") + println(" CONTINUING with the REFERENCE inflow values as authoritative scenario data.") +end + +# (3) Open-loop trajectory with the EXA policy AS-IS (state-threaded forward +# pass — see header item 6). Expected to MATCH MAIN exactly (max|Δ| = 0.0). +""" + open_loop_targets_asis(policy, x0, w_mat) -> Matrix{Float64} + +Open-loop target recursion `x̂_t = π(w_t, x̂_{t-1})`, `x̂_0 = x0`, using the EXA +policy exactly as its own training/eval pipeline calls it. The policy forward +pass now threads the LSTM recurrent state across stages (DecisionRules.jl +semantics), so this trajectory is expected to reproduce MAIN's exactly. + +# Arguments +- `policy`: EXA `HydroReachablePolicy`. +- `x0`: initial reservoir state (length nHyd). +- `w_mat`: `[T × nHyd]` inflow values. + +# Returns +- `[T × nHyd]` matrix of open-loop targets. +""" +function open_loop_targets_asis(policy, x0, w_mat) + Flux.reset!(policy) # REAL reset: start from initialstates + T, nH = size(w_mat) + prev = Float32.(x0) + out = zeros(Float64, T, nH) + for t in 1:T + target = policy(vcat(Float32.(w_mat[t, :]), prev)) + out[t, :] = Float64.(target) + prev = Float32.(target) # open-loop: previous target as next state + end + return out +end + +# (4) Open-loop trajectory with MANUALLY threaded recurrent state, replicating +# MAIN's `_step_encoder` semantics with the SAME loaded weights. Retained as an +# independent cross-check of (3): both are now expected to pass; if (3) fails +# while (4) passes, the policy forward pass has regressed from MAIN's +# threading semantics (weight parity still holds). +""" + open_loop_targets_threaded(policy, x0, w_mat) -> Matrix{Float64} + +Open-loop target recursion with the LSTM state carried across stages, +mirroring DecisionRules.jl's `HydroReachablePolicy` forward pass: + +```math +(h_t^{(l)}, c_t^{(l)}) = \\mathrm{LSTMCell}^{(l)}(h_t^{(l-1)}, (h_{t-1}^{(l)}, c_{t-1}^{(l)})), +\\qquad \\hat{x}_t = \\mathrm{lower} + (\\mathrm{upper} - \\mathrm{lower}) \\cdot \\sigma(\\mathrm{combiner}([h_t; \\hat{x}_{t-1}])) +``` + +followed by the same cascade clamp as the EXA forward pass. Uses +`layer.cell(x, state)` directly (Flux 0.16 `LSTMCell` returns +`(h, (h, c))`). + +# Arguments +- `policy`: EXA `HydroReachablePolicy` with MAIN weights loaded. +- `x0`: initial reservoir state. +- `w_mat`: `[T × nHyd]` inflow values. + +# Returns +- `[T × nHyd]` matrix of open-loop targets under MAIN's recurrence semantics. +""" +function open_loop_targets_threaded(policy, x0, w_mat) + core = policy_core(policy) + cells = [layer.cell for layer in core.encoder.layers] + # Zero initial recurrent state per layer, as Flux.initialstates gives. + states = Any[Flux.initialstates(c) for c in cells] + T, nH = size(w_mat) + prev = Float32.(x0) + out = zeros(Float64, T, nH) + for t in 1:T + w = Float32.(w_mat[t, :]) + # Thread the recurrent state layer by layer across stages. + h = stage_encoder_input(policy, t, w) + for (i, c) in enumerate(cells) + h, states[i] = c(h, states[i]) + end + # Same head + reachable-bounds scaling + cascade clamp as the EXA + # forward pass (these helpers come from hydro_reachable_policy.jl). + y = core.combiner(vcat(h, prev)) + lower, upper = _hydro_reachable_bounds(core, w, prev, y) + raw = lower .+ (upper .- lower) .* y + target = isempty(core.cascade) ? raw : + min.(raw, _cascade_upper_bounds(core, raw, w, prev)) + out[t, :] = Float64.(target) + prev = Float32.(target) + end + return out +end + +if MAIN_PARITY_GATES + w_s1 = inflow_ref[:, :, 1] # scenario 1 inflows [T × nHyd] + xhat_s1_ref = xhat_ref[:, :, 1] # MAIN open-loop trajectory + global xhat_s1_asis = open_loop_targets_asis(policy, x0_ref, w_s1) + global xhat_s1_threaded = open_loop_targets_threaded(policy, x0_ref, w_s1) + gate["openloop_asis_max_dev"] = _max_abs_dev(xhat_s1_asis, xhat_s1_ref) + gate["openloop_threaded_max_dev"] = _max_abs_dev(xhat_s1_threaded, xhat_s1_ref) + tol_traj = 1e-5 * max(1.0, maximum(abs.(xhat_s1_ref))) + gate["openloop_asis_pass"] = gate["openloop_asis_max_dev"] <= tol_traj + gate["openloop_threaded_pass"] = gate["openloop_threaded_max_dev"] <= tol_traj + println("GATE open-loop (EXA policy as-is): ", + gate["openloop_asis_pass"] ? "PASS" : "FAIL", + " (max abs dev = $(gate["openloop_asis_max_dev"]))") + println("GATE open-loop (state-threaded diag): ", + gate["openloop_threaded_pass"] ? "PASS" : "FAIL", + " (max abs dev = $(gate["openloop_threaded_max_dev"]))") +else + # The reference open-loop trajectories (`xhat_trajectories`) were rolled + # out by the MAIN reference checkpoint; comparing an independently trained + # EXA checkpoint's trajectory against them is meaningless (its targets + # SHOULD differ). Its own open-loop targets are still saved via the DE + # leg's de_target_trajectories. + global xhat_s1_asis = zeros(Float64, 0, 0) + global xhat_s1_threaded = zeros(Float64, 0, 0) + gate["openloop_asis_max_dev"] = NaN + gate["openloop_threaded_max_dev"] = NaN + gate["openloop_asis_pass"] = "skipped" + gate["openloop_threaded_pass"] = "skipped" + println("GATE open-loop trajectories: SKIPPED (DR_CHECKPOINT_KIND=exa — " * + "reference trajectories only characterize the MAIN reference checkpoint)") +end + +# ── Stage problem and callbacks (copied from train_hydro_exa_strict.jl) ─────── +# demand_matrix = nothing: the builder bakes in load_scaler × default demand for +# its single stage, matching the mof.json's 0.6-scaled loads at every stage. + +function _build_rollout_de() + build_hydro_de(power_data, hydro_data, 1; + backend = nothing, # CPU node — no GPU backend + float_type = Float64, + formulation = FORMULATION, + target_penalty = :auto, # irrelevant in strict mode + deficit_cost = DEFICIT_COST, # 6000 = MAIN mof.json coefficient + demand_matrix = nothing, + load_scaler = LOAD_SCALER, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, # header item 3 + # Stochastic demand: prepare_solve! multiplies the (constant, 0.6 × + # default) base demand by the ξ_t carried in the stage's wt block. + demand_spread = DEMAND_SPREAD, + ) +end +@info "Building 1-stage strict ExaModels stage problem (CPU, formulation=$FORMULATION)..." +rollout_prob = _build_rollout_de() + +# Same callback as train_hydro_exa_strict.jl's set_hydro_rollout_stage!, minus +# the demand update (demand_mat === nothing here). +function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_inflow, wt) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_target, target) + prepare_solve!(stage_prob, state_in, wt, target) + return stage_prob +end + +# Realized state: in strict mode hydro_solution reads the (cascade-clamped) +# reservoir parameter trajectory, so the realized state equals the clamped +# target — the strict analogue of MAIN reading value(reservoir_out). +hydro_realized_state(stage_prob, result) = + hydro_solution(stage_prob, result).reservoir[:, end] + +# Thermal generator positions, derived with the same hydro.json/PowerModels.json +# logic as MAIN's eval_paired_tsddr.jl lines 74-80: a generator is thermal iff +# its grid index is not any hydro unit's index_grid. +hydro_grid_idx = Set(power_data.gens[h.gen_pos].idx for h in hydro_data.units) +thermal_pos = [pos for (pos, g) in enumerate(power_data.gens) if !(g.idx in hydro_grid_idx)] +@info "Thermal generators: $(length(thermal_pos)) of $(power_data.nGen)" + +""" + volume_to_mw(volume; k = 0.0036) -> Float64 + +Convert a reservoir volume (hm³) to the MW-equivalent used by MAIN's +`eval_paired_tsddr.jl`: `volume / k` with `k = 0.0036`. +""" +volume_to_mw(volume; k = 0.0036) = volume / k + +""" + recompute_stage_costs(sol, power_data, deficit_cost) -> Vector{Float64} + +Recompute per-stage objective values from a `hydro_solution` NamedTuple: + +```math +c_t = \\sum_g (c_{2,g}\\, pg_{g,t}^2 + c_{1,g}\\, pg_{g,t}) + c_d \\sum_b \\mathrm{deficit}_{b,t} +``` + +which is the complete strict-mode EXA objective (no other terms exist). Used +both to split the full-horizon DE objective into stages and as an internal +consistency check on 1-stage solves. + +# Arguments +- `sol`: NamedTuple from [`hydro_solution`](@ref) (`pg` is `[nGen × T]`, + `deficit` is `[nBus × T]`). +- `power_data::PowerData`: generator cost coefficients. +- `deficit_cost::Real`: load-shedding cost per pu (6000 for this check). + +# Returns +- `Vector{Float64}` of length `T` with per-stage objective values. +""" +function recompute_stage_costs(sol, power_data, deficit_cost) + T = size(sol.pg, 2) + costs = zeros(Float64, T) + for t in 1:T + c = 0.0 + for (gpos, g) in enumerate(power_data.gens) + # Quadratic + linear generation cost (c2 = 0 for all Bolivia gens). + c += g.cost2 * sol.pg[gpos, t]^2 + g.cost1 * sol.pg[gpos, t] + end + # Per-bus active load shedding at the deficit cost. + c += deficit_cost * sum(@view sol.deficit[:, t]) + costs[t] = c + end + return costs +end + +# ── Stage-wise leg: closed-loop strict rollout on all paired scenarios ──────── +# Mirrors DecisionRulesExa.rollout_tsddr's fresh-solver path (reuse_solver = +# false, warmstart irrelevant for fresh solvers, retry_on_failure = true) while +# additionally extracting per-stage solution components that rollout_tsddr does +# not expose (pg, deficit, spill, outflow). A rollout_tsddr cross-check on +# scenario 1 verifies the custom loop matches the package machinery. + +@info "Stage-wise leg: $NUM_SCEN scenarios × $T_EVAL stages (cold MadNLP solves)..." + +costs_stagewise = fill(NaN, NUM_SCEN) # per-scenario total objective +stage_costs = fill(NaN, T_EVAL, NUM_SCEN) # per-stage objective +vol_trajectories = fill(NaN, T_EVAL, NUM_SCEN) # Σ_r volume_to_mw(state_r) after stage t +gen_trajectories = fill(NaN, T_EVAL, NUM_SCEN) # Σ_thermal pg × baseMVA at stage t +reservoir_traj = fill(NaN, nHyd, T_EVAL + 1, NUM_SCEN) # realized states incl. x0 +outflow_traj = fill(NaN, nHyd, T_EVAL, NUM_SCEN) +spill_traj = fill(NaN, nHyd, T_EVAL, NUM_SCEN) +target_traj = fill(NaN, nHyd, T_EVAL, NUM_SCEN) # raw policy targets +max_deficit = fill(NaN, NUM_SCEN) # max_b,t deficit (pu) +sum_deficit = fill(NaN, NUM_SCEN) # Σ_b,t deficit (pu) +max_abs_deficit_q = fill(NaN, NUM_SCEN) # max_b,t |deficit_q| (pu) — EXA-only slack +scenario_ok = falses(NUM_SCEN) +n_retries_total = 0 +max_objective_recompute_dev = 0.0 # internal consistency check + +baseMVA = power_data.baseMVA + +for s in 1:NUM_SCEN + Flux.reset!(policy) # REAL reset at the scenario boundary + state = copy(x0_ref) # closed-loop realized state (Float64) + reservoir_traj[:, 1, s] = state + total = 0.0 + failed = false + for t in 1:T_EVAL + # Authoritative inflow values from the MAIN reference (see gate item 2), + # with the seeded paired demand factor ξ_t^{(s)} appended when + # stochastic demand is active ([w_t; ξ_t] block). + w_t = augmented_stage_w(inflow_ref[t, :, s], t, s) + # Policy call with Float32 inputs, exactly as MAIN's closed loop does + # (the policy slices the physical inflow internally). + x_hat = Float64.(policy(vcat(Float32.(w_t), Float32.(state)))) + target_traj[:, t, s] = x_hat + + # Write x_{t-1}, [w_t; ξ_t], x̂_t into the strict stage problem + # (prepare_solve! applies base_demand · ξ_t via set_demand!). + set_hydro_rollout_stage!(rollout_prob, state, w_t, x_hat, t) + + # Cold one-shot solve (rollout_tsddr's fresh-solver path). + result = MadNLP.madnlp(rollout_prob.model; SOLVER_KWARGS...) + if !solve_succeeded(result) || !isfinite(result.objective) + # Single cold retry, mirroring rollout_tsddr's retry_on_failure. + global n_retries_total += 1 + result = MadNLP.madnlp(rollout_prob.model; SOLVER_KWARGS...) + end + if !solve_succeeded(result) || !isfinite(result.objective) + @warn "Stage solve failed after retry" scenario = s stage = t status = result.status + failed = true + break + end + + sol = hydro_solution(rollout_prob, result) + total += result.objective + stage_costs[t, s] = result.objective + # Internal consistency: the strict stage objective must equal the + # recomputed gen+deficit cost exactly (same terms, same data). + recomputed = recompute_stage_costs(sol, power_data, DEFICIT_COST)[1] + global max_objective_recompute_dev = + max(max_objective_recompute_dev, abs(recomputed - result.objective)) + + # Advance the closed loop on the realized state (= clamped target). + state = Float64.(sol.reservoir[:, end]) + reservoir_traj[:, t + 1, s] = state + outflow_traj[:, t, s] = sol.outflow[:, 1] + spill_traj[:, t, s] = sol.spill[:, 1] + + # Operative metrics with the IDENTICAL formulas as MAIN's evaluation. + vol_trajectories[t, s] = sum(volume_to_mw(state[r]) for r in 1:nHyd) + gen_trajectories[t, s] = sum(sol.pg[g, 1] * baseMVA for g in thermal_pos) + + # Deficit activity (quantifies whether cost/slack differences are inert). + def_max = maximum(sol.deficit[:, 1]) + def_sum = sum(sol.deficit[:, 1]) + dq_max = maximum(abs.(sol.deficit_q[:, 1])) + max_deficit[s] = isnan(max_deficit[s]) ? def_max : max(max_deficit[s], def_max) + sum_deficit[s] = isnan(sum_deficit[s]) ? def_sum : sum_deficit[s] + def_sum + max_abs_deficit_q[s] = isnan(max_abs_deficit_q[s]) ? dq_max : max(max_abs_deficit_q[s], dq_max) + end + if !failed + costs_stagewise[s] = total + scenario_ok[s] = true + end + if s % 10 == 0 || s == NUM_SCEN + ok_costs = costs_stagewise[1:s][scenario_ok[1:s]] + running_mean = isempty(ok_costs) ? NaN : mean(ok_costs) + println(" [$s/$NUM_SCEN] cost = $(round(total; digits=1)), running mean = $(round(running_mean; digits=1)), retries = $n_retries_total") + end +end +@info "Stage-wise leg done" n_ok = count(scenario_ok) n_retries_total max_objective_recompute_dev + +# ── rollout_tsddr cross-check on scenario 1 ─────────────────────────────────── +# Runs the actual package machinery (same callbacks) to verify the custom loop +# above reproduces it. Float32 initial state so the policy sees Float32 inputs; +# the solver interface buffers are Float64 regardless. +w_flat_s1 = vec(permutedims(inflow_ref[:, :, 1])) # stage-major [w_1; w_2; …] +rollout_check = rollout_tsddr( + policy, + Float32.(x0_ref), + rollout_prob, + Float32.(w_flat_s1); + horizon = T_EVAL, + n_uncertainty = nHyd, + set_stage_parameters! = set_hydro_rollout_stage!, + realized_state = hydro_realized_state, + madnlp_kwargs = SOLVER_KWARGS, + warmstart = false, + policy_state = :realized, + retry_on_failure = true, +) +rollout_check_obj = rollout_check === nothing ? NaN : rollout_check.objective +@info "rollout_tsddr cross-check (scenario 1)" custom_loop = costs_stagewise[1] rollout_tsddr = rollout_check_obj + +# ── DE leg: strict full-horizon deterministic equivalent ────────────────────── +# Strict-mode claim under test: with the target trajectory generated by the +# open-loop recursion (previous target as next policy state — exactly +# train_hydro_exa_strict.jl's rollout_reachable_targets), the T-stage strict DE +# objective equals the stage-wise sum for the same scenario, because the +# reservoir path is pinned to the same targets and stages decouple. + +@info "DE leg: building $T_EVAL-stage strict DE and solving $NUM_DE_SCENARIOS scenarios..." +de_prob = build_hydro_de(power_data, hydro_data, T_EVAL; + backend = nothing, + float_type = Float64, + formulation = FORMULATION, + target_penalty = :auto, + deficit_cost = DEFICIT_COST, + demand_matrix = nothing, + load_scaler = LOAD_SCALER, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, # header item 3 +) + +""" + rollout_reachable_targets(policy, x0, w_flat, T, nHyd) -> Vector{Float64} + +Roll out the reachable-policy target trajectory for the strict regular DE +(verbatim semantics of train_hydro_exa_strict.jl): start from the feasible +`x0`, feed `[w_t; previous_target]` to the policy, and record each target as +the next previous state, so every target is one-stage reachable from the prior +target by induction. + +# Arguments +- `policy`: reachable hydro policy with input `[inflow; previous_state]`. +- `x0`: initial reservoir state. +- `w_flat`: stage-major flat inflow vector of length `T * nHyd`. +- `T::Int`: number of stages. +- `nHyd::Int`: number of hydro reservoirs. + +# Returns +- `Vector{Float64}`: stage-major target trajectory for + `ExaModels.set_parameter!(prob.core, prob.p_target, targets)`. +""" +function rollout_reachable_targets(policy, x0, w_flat, T, nHyd) + Flux.reset!(policy) + prev = Float32.(x0) + targets = Vector{Vector{Float32}}(undef, T) + for t in 1:T + wt = Float32.(view(w_flat, ((t - 1) * nHyd + 1):(t * nHyd))) + target = policy(vcat(wt, prev)) + targets[t] = Float32.(target) + prev = targets[t] + end + return Float64.(vcat(targets...)) +end + +de_costs = fill(NaN, NUM_DE_SCENARIOS) +de_stage_costs = fill(NaN, T_EVAL, NUM_DE_SCENARIOS) +de_vol = fill(NaN, T_EVAL, NUM_DE_SCENARIOS) +de_gen = fill(NaN, T_EVAL, NUM_DE_SCENARIOS) +de_reservoir = fill(NaN, nHyd, T_EVAL + 1, NUM_DE_SCENARIOS) +de_outflow = fill(NaN, nHyd, T_EVAL, NUM_DE_SCENARIOS) +de_spill = fill(NaN, nHyd, T_EVAL, NUM_DE_SCENARIOS) +de_targets = fill(NaN, nHyd, T_EVAL, NUM_DE_SCENARIOS) +de_max_deficit = fill(NaN, NUM_DE_SCENARIOS) +de_max_abs_deficit_q = fill(NaN, NUM_DE_SCENARIOS) +de_ok = falses(NUM_DE_SCENARIOS) + +for s in 1:NUM_DE_SCENARIOS + # Stage-major flat inflow vector for scenario s from the reference values. + w_flat = vec(permutedims(inflow_ref[:, :, s])) + # Open-loop target trajectory with the EXA policy (previous target as state). + targets = rollout_reachable_targets(policy, x0_ref, w_flat, T_EVAL, nHyd) + de_targets[:, :, s] = reshape(targets, nHyd, T_EVAL) + + # Write parameters and pin the strict reservoir path (prepare_solve! also + # applies the Float64 cascade clamp and enforces p_reservoir[1:nHyd] = x0). + ExaModels.set_parameter!(de_prob.core, de_prob.p_x0, x0_ref) + ExaModels.set_parameter!(de_prob.core, de_prob.p_inflow, w_flat) + ExaModels.set_parameter!(de_prob.core, de_prob.p_target, targets) + prepare_solve!(de_prob, x0_ref, w_flat, targets) + + result = MadNLP.madnlp(de_prob.model; SOLVER_KWARGS...) + if !solve_succeeded(result) || !isfinite(result.objective) + @warn "DE solve failed" scenario = s status = result.status + continue + end + de_ok[s] = true + de_costs[s] = result.objective + + sol = hydro_solution(de_prob, result) + de_stage_costs[:, s] = recompute_stage_costs(sol, power_data, DEFICIT_COST) + de_reservoir[:, :, s] = sol.reservoir + de_outflow[:, :, s] = sol.outflow + de_spill[:, :, s] = sol.spill + for t in 1:T_EVAL + # Operative metrics with the identical MAIN formulas, taken from the + # post-stage reservoir state and per-stage thermal generation. + de_vol[t, s] = sum(volume_to_mw(sol.reservoir[r, t + 1]) for r in 1:nHyd) + de_gen[t, s] = sum(sol.pg[g, t] * baseMVA for g in thermal_pos) + end + de_max_deficit[s] = maximum(sol.deficit) + de_max_abs_deficit_q[s] = maximum(abs.(sol.deficit_q)) + println(" DE [$s/$NUM_DE_SCENARIOS] objective = $(round(result.objective; digits=1)), stagewise total = $(round(costs_stagewise[s]; digits=1))") +end +@info "DE leg done" n_ok = count(de_ok) + +# ── Save everything ──────────────────────────────────────────────────────────── +out_dir = joinpath(SCRIPT_DIR, CASE_NAME, FORM_LABEL, "results") +mkpath(out_dir) +# DR_OUTPUT_TAG suffixes the filename so non-reference evaluations never +# clobber the untagged reference results file. +out_file = joinpath(out_dir, "paired_exa_strict$(OUT_SUFFIX).jld2") +# Provenance knobs, recorded whenever any of the new env vars was set (the +# all-defaults run keeps exactly the historical key set). +knob_extras = RECORD_KNOBS ? ( + checkpoint_path = MODEL_PATH, + checkpoint_kind = CHECKPOINT_KIND, + encoder_layers = ENCODER_LAYERS, + head_layers = HEAD_LAYERS, + context_mode = isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE, + context_period = CONTEXT_PERIOD, + context_horizon = CONTEXT_HORIZON, + n_context = N_CONTEXT, + output_tag = OUTPUT_TAG, + main_parity_gates_applied = MAIN_PARITY_GATES, +) : (;) +jldsave(out_file; + knob_extras..., + # Gate results (policy parity, inflow indexing, metadata) + gate_loader_mode = loader_mode, + gate_probe_max_dev = gate["probe_max_dev"], + gate_probe_pass = gate["probe_pass"], + gate_probe_outputs_exa = probe_out_exa, + gate_inflow_recon_max_dev = gate["inflow_recon_max_dev"], + gate_inflow_pass = gate["inflow_pass"], + gate_openloop_asis_max_dev = gate["openloop_asis_max_dev"], + gate_openloop_asis_pass = gate["openloop_asis_pass"], + gate_openloop_threaded_max_dev = gate["openloop_threaded_max_dev"], + gate_openloop_threaded_pass = gate["openloop_threaded_pass"], + gate_dev_K = gate["dev_K"], + gate_dev_min_vol = gate["dev_min_vol"], + gate_dev_max_vol = gate["dev_max_vol"], + gate_dev_min_turn = gate["dev_min_turn"], + gate_dev_max_turn = gate["dev_max_turn"], + gate_dev_upstream_max = gate["dev_upstream_max"], + xhat_s1_asis = xhat_s1_asis, + xhat_s1_threaded = xhat_s1_threaded, + # Stage-wise leg + costs = costs_stagewise, + stage_costs = stage_costs, + vol_trajectories = vol_trajectories, + gen_trajectories = gen_trajectories, + reservoir_trajectories = reservoir_traj, + outflow_trajectories = outflow_traj, + spill_trajectories = spill_traj, + target_trajectories = target_traj, + max_deficit = max_deficit, + sum_deficit = sum_deficit, + max_abs_deficit_q = max_abs_deficit_q, + scenario_ok = collect(scenario_ok), + n_retries_total = n_retries_total, + max_objective_recompute_dev = max_objective_recompute_dev, + rollout_tsddr_check_objective = rollout_check_obj, + # DE leg + de_costs = de_costs, + de_stage_costs = de_stage_costs, + de_vol_trajectories = de_vol, + de_gen_trajectories = de_gen, + de_reservoir_trajectories = de_reservoir, + de_outflow_trajectories = de_outflow, + de_spill_trajectories = de_spill, + de_target_trajectories = de_targets, + de_max_deficit = de_max_deficit, + de_max_abs_deficit_q = de_max_abs_deficit_q, + de_ok = collect(de_ok), + # Configuration provenance + deficit_cost_used = DEFICIT_COST, + load_scaler_used = LOAD_SCALER, + reactive_deficit_setting = REACTIVE_DEFICIT_RAW, + reactive_deficit_cost_used = + REACTIVE_DEFICIT_COST === nothing ? "free" : Float64(REACTIVE_DEFICIT_COST), + solver_tol = SOLVER_KWARGS.tol, + solver_max_iter = SOLVER_KWARGS.max_iter, + model_path = MODEL_PATH, + reference_file = REFERENCE_FILE, + num_eval_stages = T_EVAL, + num_scenarios = NUM_SCEN, + num_de_scenarios = NUM_DE_SCENARIOS, +) +println("Saved: $out_file") + +# ── Final summary ───────────────────────────────────────────────────────────── + +""" + _split_csv_line(line) -> Vector{String} + +Split one CSV line into fields with minimal double-quote awareness: commas +inside `"…"` do not delimit (needed for the MAIN header column +`"TS-DDR (strict, paired)"`). No escape handling beyond quote toggling. +""" +function _split_csv_line(line::AbstractString) + fields = String[] # accumulated fields + buf = IOBuffer() # current field characters + inq = false # inside a quoted region? + for c in line + if c == '"' + inq = !inq # toggle quoting; quotes are dropped + elseif c == ',' && !inq + push!(fields, String(take!(buf))) # unquoted comma ends the field + else + write(buf, c) + end + end + push!(fields, String(take!(buf))) # trailing field + return fields +end + +""" + _mean_csv_column(path, needle) -> Float64 + +Mean of the numeric column whose header contains `needle` in the CSV at +`path`: + +```math +\\bar{c} = \\frac{1}{n} \\sum_{i=1}^{n} c_i +``` + +Returns `NaN` when the file is missing, the column is not found, or no row +parses — the summary then simply reports `NaN` for that baseline. +""" +function _mean_csv_column(path::AbstractString, needle::AbstractString) + isfile(path) || return NaN # ground truth absent + lines = readlines(path) + length(lines) >= 2 || return NaN # header + ≥1 data row + header = _split_csv_line(lines[1]) # quote-aware header parse + col = findfirst(h -> occursin(needle, h), header) # column by substring + col === nothing && return NaN + vals = Float64[] + for ln in lines[2:end] + fs = split(ln, ',') # data rows are plain numeric + length(fs) == length(header) || continue # skip malformed rows + v = tryparse(Float64, strip(fs[col])) + v === nothing || push!(vals, v) + end + return isempty(vals) ? NaN : mean(vals) +end + +# Successful-scenario cost vectors for the two legs of THIS evaluation. +ok_costs = costs_stagewise[collect(scenario_ok)] +de_ok_costs = de_costs[collect(de_ok)] + +# Untagged ground-truth baselines from the MAIN repo (comparability anchors): +# the .026 reference checkpoint's stage-wise mean (results/paired_strict_rollout.jld2) +# and the SDDP paired mean (SDDP-SOC column of paired_costs.csv). NaN if absent. +main_gt_file = joinpath(MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "results", "paired_strict_rollout.jld2") +main_gt_mean = isfile(main_gt_file) ? mean(Float64.(JLD2.load(main_gt_file, "costs"))) : NaN +sddp_mean = _mean_csv_column( + joinpath(MAIN_HPM_DIR, CASE_NAME, FORM_LABEL, "paired_costs.csv"), "SDDP") + +println("\n" * "=" ^ 64) +println("SUMMARY — paired EXA strict evaluation") +println(" Checkpoint: $MODEL_PATH") +println(" Kind: $CHECKPOINT_KIND (encoder=$(ENCODER_LAYERS), head=$(HEAD_LAYERS))") +println(" Context: $(isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE) (period=$CONTEXT_PERIOD, horizon=$CONTEXT_HORIZON)") +println(" Reference: $REFERENCE_FILE") +println(" Output tag: $(isempty(OUTPUT_TAG) ? "(none)" : OUTPUT_TAG)") +println(" Activation: $(ACTIVATION) snap_eps=$(SNAP_EPS)") +println(" Stage-wise mean: $(round(mean(ok_costs); digits=1)) over $(length(ok_costs))/$NUM_SCEN scenarios") +println(" Stage-wise std: $(round(std(ok_costs); digits=1))") +println(" DE-leg mean: $(round(mean(de_ok_costs); digits=1)) over $(length(de_ok_costs))/$NUM_DE_SCENARIOS scenarios") +println(" MAIN .026 mean (GT): $(round(main_gt_mean; digits=1))") +println(" SDDP mean (GT): $(round(sddp_mean; digits=1))") +println("=" ^ 64) diff --git a/examples/HydroPowerModels/generate_canonical_case_artifacts.jl b/examples/HydroPowerModels/generate_canonical_case_artifacts.jl new file mode 100644 index 0000000..7a6d9b1 --- /dev/null +++ b/examples/HydroPowerModels/generate_canonical_case_artifacts.jl @@ -0,0 +1,277 @@ +#!/usr/bin/env julia + +module HydroCanonicalCase + +using JSON +using SHA + +const ACTIVE_LOAD_FACTOR = 0.6 +const REACTIVE_LOAD_FACTOR = 0.6 +const STAGE_HOURS = 168 +const HYDRO_CONVERSION_K = 0.6048 +const REPORTING_STAGES = 96 +const LOOKAHEAD_STAGES = 30 +const TOTAL_STAGES = REPORTING_STAGES + LOOKAHEAD_STAGES +const ACTIVE_DEFICIT_COST = 6000.0 +const INFLOW_PROTOCOL_SEED = 20260706 +const DEMAND_PROTOCOL_SEED = 20260714 +const PROTOCOL_SCENARIOS = 500 +const DEMAND_ATOMS = (0.9, 1.0, 1.1) +const FORMULATIONS = ("ACPPowerModel", "DCPPowerModel", "SOCWRConicPowerModel") + +const INPUT_HASHES = Dict( + "PowerModels.json" => "1ff598447957f9fc17ca570415bf5b9b5b14e1292ea3bd3163db0ad79911a782", + "inflows.csv" => "5afb275dff3fc879e3e93b6510b81295834faad0bcd2bd1fc070a8e3e6653c77", + "hydro.json" => "1014b90ac06afb36f136a34c66fbd12d45a81c1e700324dcda398a5438d4d98b", + "demand_scenarios.csv" => "abeb664b606ba17b830398e8d891d4d23dd7739d5b221088ca7a30bcfab4e8df", +) + +sha256_file(path::AbstractString) = bytes2hex(open(SHA.sha256, path)) + +function verify_inputs(case_dir::AbstractString) + for (name, expected) in INPUT_HASHES + path = joinpath(case_dir, name) + isfile(path) || error("missing canonical input: $path") + actual = sha256_file(path) + actual == expected || + error("canonical input hash mismatch for $path: expected $expected, got $actual") + end + + hydro = JSON.parsefile(joinpath(case_dir, "hydro.json")) + Int(hydro["stage_hours"]) == STAGE_HOURS || + error("stage_hours must be $STAGE_HOURS") + isapprox(0.0036 * Int(hydro["stage_hours"]), HYDRO_CONVERSION_K; atol=0, rtol=1e-14) || + error("hydro conversion K must be $HYDRO_CONVERSION_K") + + power = JSON.parsefile(joinpath(case_dir, "PowerModels.json")) + Float64(power["baseMVA"]) == 100.0 || error("canonical baseMVA must be 100") + Float64(power["cost_deficit"]) * Float64(power["baseMVA"]) == ACTIVE_DEFICIT_COST || + error("canonical active-deficit coefficient must be $ACTIVE_DEFICIT_COST") + return Dict( + "buses" => length(power["bus"]), + "branches" => length(power["branch"]), + "generators" => length(power["gen"]), + "loads" => length(power["load"]), + "hydro_units" => length(hydro["Hydrogenerators"]), + ) +end + +function scale_main_loads!(alldata) + for stage in alldata + for load in values(stage["powersystem"]["load"]) + load["pd"] = Float64(load["pd"]) * ACTIVE_LOAD_FACTOR + load["qd"] = Float64(load["qd"]) * REACTIVE_LOAD_FACTOR + end + stage["powersystem"]["cost_deficit"] = + ACTIVE_DEFICIT_COST / Float64(stage["powersystem"]["baseMVA"]) + end + return alldata +end + +function base_manifest(case_dir::AbstractString) + counts = verify_inputs(case_dir) + return Dict( + "schema_version" => 1, + "case" => "Bolivia MAIN", + "input_hashes" => INPUT_HASHES, + "initial_volume_repair" => "70% of capacity for corrupted MAIN denormals", + "active_load_factor" => ACTIVE_LOAD_FACTOR, + "reactive_load_factor" => REACTIVE_LOAD_FACTOR, + "stage_hours" => STAGE_HOURS, + "hydro_conversion_K" => HYDRO_CONVERSION_K, + "reporting_stages" => REPORTING_STAGES, + "lookahead_stages" => LOOKAHEAD_STAGES, + "total_stages" => TOTAL_STAGES, + "active_deficit_cost_usd_per_pu_stage" => ACTIVE_DEFICIT_COST, + "active_deficit_cost_derivation" => "60 USD/MWh * 100 MVA", + "strict_targets_primary" => true, + "reactive_balance" => "hard", + "branch_thermal_limits" => "both ends", + "reachable_activation" => "stretchedsigmoid_safe_upper_margin_1e-3", + "demand_atoms" => collect(DEMAND_ATOMS), + "demand_atom_probabilities" => fill(1 / 3, length(DEMAND_ATOMS)), + "topology_counts" => counts, + "mofs" => Dict{String,Any}(), + ) +end + +function write_manifest(case_dir::AbstractString, manifest) + path = joinpath(case_dir, "case_manifest.json") + open(path, "w") do io + JSON.print(io, manifest, 2) + write(io, '\n') + end + return path +end + +function read_manifest(case_dir::AbstractString) + path = joinpath(case_dir, "case_manifest.json") + isfile(path) || error("missing canonical case manifest: $path") + manifest = JSON.parsefile(path) + manifest["active_load_factor"] == ACTIVE_LOAD_FACTOR || + error("manifest active_load_factor must be $ACTIVE_LOAD_FACTOR") + manifest["reactive_load_factor"] == REACTIVE_LOAD_FACTOR || + error("manifest reactive_load_factor must be $REACTIVE_LOAD_FACTOR") + manifest["stage_hours"] == STAGE_HOURS || + error("manifest stage_hours must be $STAGE_HOURS") + manifest["hydro_conversion_K"] == HYDRO_CONVERSION_K || + error("manifest hydro_conversion_K must be $HYDRO_CONVERSION_K") + manifest["active_deficit_cost_usd_per_pu_stage"] == ACTIVE_DEFICIT_COST || + error("manifest active deficit cost must be $ACTIVE_DEFICIT_COST") + return manifest +end + +function objective_terms(mof) + function_object = mof["objective"]["function"] + return get(function_object, "terms", Any[]) +end + +function verify_mof(path::AbstractString, formulation::AbstractString) + mof = JSON.parsefile(path) + mof["objective"]["sense"] == "min" || error("$formulation objective is not Min") + deficit_terms = [ + term for term in objective_terms(mof) + if startswith(get(term, "variable", ""), "deficit[") + ] + length(deficit_terms) == 28 || + error("$formulation must have 28 operational active-deficit objective terms") + all(Float64(term["coefficient"]) == ACTIVE_DEFICIT_COST for term in deficit_terms) || + error("$formulation active-deficit coefficient is not $ACTIVE_DEFICIT_COST") + + hydro_balance = [ + constraint for constraint in mof["constraints"] + if startswith(get(constraint, "name", ""), "hydro_balance[") + ] + length(hydro_balance) == 11 || error("$formulation must have 11 hydro balances") + for constraint in hydro_balance + inflow_terms = [ + term for term in constraint["function"]["terms"] + if startswith(get(term, "variable", ""), "inflow[") + ] + length(inflow_terms) == 1 || error("$formulation hydro balance lacks one inflow") + abs(Float64(only(inflow_terms)["coefficient"])) == HYDRO_CONVERSION_K || + error("$formulation MOF has stale hydro K") + end + + variable_names = String[get(v, "name", "") for v in mof["variables"]] + any(startswith(name, "target_deficit") for name in variable_names) && + error("$formulation MOF mixes operational active deficit with target slack") + + if formulation == "ACPPowerModel" + reactive_balances = [ + constraint for constraint in mof["constraints"] + if any( + startswith(get(term, "variable", ""), "0_q[") + for term in get(constraint["function"], "terms", Any[]) + ) && get(constraint["set"], "type", "") == "EqualTo" + ] + length(reactive_balances) >= 28 || + error("ACP MOF is missing hard reactive balance equations") + quadratic_limits = [ + constraint for constraint in mof["constraints"] + if get(constraint["function"], "type", "") == "ScalarQuadraticFunction" && + get(constraint["set"], "type", "") == "LessThan" + ] + length(quadratic_limits) >= 62 || + error("ACP MOF is missing both-ended apparent-power thermal limits") + end + + return Dict( + "sha256" => sha256_file(path), + "variables" => length(mof["variables"]), + "constraints" => length(mof["constraints"]), + "objective_sense" => mof["objective"]["sense"], + "active_deficit_terms" => length(deficit_terms), + "hydro_balance_terms" => length(hydro_balance), + ) +end + +function finalize_manifest(case_dir::AbstractString) + manifest = read_manifest(case_dir) + manifest["mofs"] = Dict( + formulation => verify_mof( + joinpath(case_dir, formulation * ".mof.json"), + formulation, + ) + for formulation in FORMULATIONS + ) + write_manifest(case_dir, manifest) + return manifest +end + +export ACTIVE_LOAD_FACTOR, REACTIVE_LOAD_FACTOR, STAGE_HOURS, + HYDRO_CONVERSION_K, REPORTING_STAGES, LOOKAHEAD_STAGES, TOTAL_STAGES, + ACTIVE_DEFICIT_COST, INFLOW_PROTOCOL_SEED, DEMAND_PROTOCOL_SEED, + PROTOCOL_SCENARIOS, DEMAND_ATOMS, FORMULATIONS, INPUT_HASHES, + sha256_file, verify_inputs, scale_main_loads!, base_manifest, + write_manifest, read_manifest, verify_mof, finalize_manifest + +end + +using .HydroCanonicalCase +using JSON + +function argument_value(prefix) + arg = findfirst(value -> startswith(value, prefix), ARGS) + return isnothing(arg) ? nothing : split(ARGS[arg], '='; limit=2)[2] +end + +function main() + case_dir = joinpath(@__DIR__, "bolivia") + verify_inputs(case_dir) + + if "--verify-only" in ARGS + manifest = read_manifest(case_dir) + for formulation in FORMULATIONS + verify_mof(joinpath(case_dir, formulation * ".mof.json"), formulation) + end + println(JSON.json(Dict("status" => "verified", "manifest" => manifest))) + return + end + + if !("--finalize-only" in ARGS) + write_manifest(case_dir, base_manifest(case_dir)) + end + "--prepare-only" in ARGS && return + + if !("--finalize-only" in ARGS) + exporter = joinpath(@__DIR__, "export_subproblem_mof.jl") + for formulation in FORMULATIONS + run(`$(Base.julia_cmd()) $exporter bolivia $formulation`) + end + end + + manifest = finalize_manifest(case_dir) + protocol_generator = joinpath(@__DIR__, "generate_joint_protocol.jl") + run(`$(Base.julia_cmd()) $protocol_generator`) + manifest = read_manifest(case_dir) + + exa_root = argument_value("--exa-root=") + if exa_root !== nothing + exa_dir = joinpath(exa_root, "examples", "HydroPowerModels") + exa_case = joinpath(exa_dir, "bolivia") + mkpath(exa_case) + for name in keys(INPUT_HASHES) + cp(joinpath(case_dir, name), joinpath(exa_case, name); force=true) + end + for formulation in FORMULATIONS + name = formulation * ".mof.json" + cp(joinpath(case_dir, name), joinpath(exa_case, name); force=true) + end + for name in ("case_manifest.json", "joint_protocol_500.csv") + cp(joinpath(case_dir, name), joinpath(exa_case, name); force=true) + end + for name in ("generate_canonical_case_artifacts.jl", "generate_joint_protocol.jl") + cp(joinpath(@__DIR__, name), joinpath(exa_dir, name); force=true) + end + end + + println(JSON.json(Dict( + "status" => "generated", + "manifest_sha256" => sha256_file(joinpath(case_dir, "case_manifest.json")), + "mofs" => manifest["mofs"], + "protocol" => manifest["protocol"], + ))) +end + +(abspath(PROGRAM_FILE) == @__FILE__) && main() diff --git a/examples/HydroPowerModels/generate_joint_protocol.jl b/examples/HydroPowerModels/generate_joint_protocol.jl new file mode 100644 index 0000000..4a4383b --- /dev/null +++ b/examples/HydroPowerModels/generate_joint_protocol.jl @@ -0,0 +1,125 @@ +#!/usr/bin/env julia + +using CSV +using DataFrames +using JSON +using SHA +using StableRNGs + +include(joinpath(@__DIR__, "generate_canonical_case_artifacts.jl")) +using .HydroCanonicalCase + +const CASE_DIR = joinpath(@__DIR__, "bolivia") +const OUTPUT = joinpath(CASE_DIR, "joint_protocol_500.csv") + +function inflow_scenario_count(case_dir) + first_row = split(first(eachline(joinpath(case_dir, "inflows.csv"))), ',') + nhyd = length(JSON.parsefile(joinpath(case_dir, "hydro.json"))["Hydrogenerators"]) + length(first_row) % nhyd == 0 || error("inflow columns are not divisible by nHyd") + return length(first_row) ÷ nhyd +end + +function generate_indices(ncen) + inflow = rand( + StableRNG(INFLOW_PROTOCOL_SEED), + 1:ncen, + TOTAL_STAGES, + PROTOCOL_SCENARIOS, + ) + demand = reduce( + hcat, + [ + rand( + StableRNG(DEMAND_PROTOCOL_SEED + scenario), + 1:length(DEMAND_ATOMS), + TOTAL_STAGES, + ) + for scenario in 1:PROTOCOL_SCENARIOS + ], + ) + return inflow, demand +end + +function protocol_table(inflow, demand) + rows = TOTAL_STAGES * PROTOCOL_SCENARIOS + stage = Vector{Int}(undef, rows) + scenario = Vector{Int}(undef, rows) + inflow_index = Vector{Int}(undef, rows) + demand_index = Vector{Int}(undef, rows) + k = 1 + for t in 1:TOTAL_STAGES, s in 1:PROTOCOL_SCENARIOS + stage[k] = t + scenario[k] = s + inflow_index[k] = inflow[t, s] + demand_index[k] = demand[t, s] + k += 1 + end + return DataFrame(; stage, scenario, inflow_index, demand_index) +end + +function verify_protocol(path=OUTPUT) + verify_inputs(CASE_DIR) + ncen = inflow_scenario_count(CASE_DIR) + inflow, demand = generate_indices(ncen) + saved = CSV.read(path, DataFrame) + expected = protocol_table(inflow, demand) + names(saved) == names(expected) || error("joint protocol columns do not match") + saved == expected || error("joint protocol does not reconstruct exactly") + manifest = read_manifest(CASE_DIR) + if haskey(manifest, "protocol") + manifest["protocol"]["csv_sha256"] == sha256_file(path) || + error("joint protocol hash does not match the case manifest") + manifest["protocol"]["source_hashes"] == INPUT_HASHES || + error("joint protocol source hashes do not match the canonical inputs") + end + return (ncen=ncen, protocol_sha256=sha256_file(path)) +end + +if "--verify" in ARGS + result = verify_protocol() + println(JSON.json(Dict( + "status" => "verified", + "inflow_scenarios" => result.ncen, + "protocol_sha256" => result.protocol_sha256, + ))) +else + counts = verify_inputs(CASE_DIR) + ncen = inflow_scenario_count(CASE_DIR) + inflow, demand = generate_indices(ncen) + CSV.write(OUTPUT, protocol_table(inflow, demand)) + protocol_hash = sha256_file(OUTPUT) + metadata = Dict( + "schema_version" => 1, + "layout" => "stage-major rows; stage outer, scenario inner", + "dimensions" => Dict( + "stages" => TOTAL_STAGES, + "scenarios" => PROTOCOL_SCENARIOS, + "inflow_scenarios" => ncen, + "demand_atoms" => length(DEMAND_ATOMS), + ), + "seeds" => Dict( + "inflow_matrix" => INFLOW_PROTOCOL_SEED, + "demand_column_base" => DEMAND_PROTOCOL_SEED, + "demand_column_rule" => "StableRNG(demand_column_base + scenario_id)", + ), + "demand_atom_values" => collect(DEMAND_ATOMS), + "source_hashes" => INPUT_HASHES, + "csv_sha256" => protocol_hash, + "topology_counts" => Dict(string(k) => v for (k, v) in pairs(counts)), + ) + manifest_path = joinpath(CASE_DIR, "case_manifest.json") + manifest = read_manifest(CASE_DIR) + manifest["protocol"] = metadata + open(manifest_path, "w") do io + JSON.print(io, manifest, 2) + write(io, '\n') + end + result = verify_protocol() + result.protocol_sha256 == protocol_hash || error("protocol hash changed during verification") + println(JSON.json(Dict( + "status" => "generated", + "rows" => TOTAL_STAGES * PROTOCOL_SCENARIOS, + "protocol_sha256" => protocol_hash, + "manifest_sha256" => sha256_file(manifest_path), + ))) +end diff --git a/examples/HydroPowerModels/hydro_power_data.jl b/examples/HydroPowerModels/hydro_power_data.jl index bc7f5b4..b79cad7 100644 --- a/examples/HydroPowerModels/hydro_power_data.jl +++ b/examples/HydroPowerModels/hydro_power_data.jl @@ -10,6 +10,7 @@ # (DC-OPF susceptance formula, branch-variable formulation). using JSON, CSV, Tables, Statistics, Random +using StableRNGs # seeded demand-noise draws for the paired eval protocol # ── Power system data structures ───────────────────────────────────────────── @@ -244,7 +245,8 @@ function load_hydro_data(hydro_file::AbstractString, power_data::PowerData; num_stages::Union{Int,Nothing} = nothing) - hydro_json = JSON.parsefile(hydro_file)["Hydrogenerators"] + hydro_root = JSON.parsefile(hydro_file) + hydro_json = hydro_root["Hydrogenerators"] nHyd = length(hydro_json) # Build gen_index → gen_pos map @@ -303,9 +305,18 @@ function load_hydro_data(hydro_file::AbstractString, scenario_inflows[r] = Float64.(allinflows[:, ((r-1)*nScenarios+1):(r*nScenarios)]) end - # Water balance conversion factor K = 0.0036 (standard HydroPowerModels.jl value) - # Converts turbine outflow (m³/s equivalent) to reservoir volume per stage. - K = 0.0036 + # Water-balance conversion factor K = 0.0036 · stage_hours. + # 0.0036 converts a flow of m³/s to reservoir volume (hm³) accumulated over one + # HOUR (3600 s/h · 1e-6 hm³/m³ = 0.0036). Multiplying by the stage duration + # `stage_hours` (hours per stage) gives the per-stage flow→volume factor, so the + # reservoir balance `V_{t+1} = V_t + K·(inflow − outflow) − spill(+upstream)` + # is dimensionally correct for a `stage_hours`-long stage. + # `stage_hours` is read from hydro.json and defaults to 1 (⇒ K = 0.0036) for + # backward compatibility with cases predating the field. This exactly mirrors + # HydroPowerModels.jl `constraint_hydro_balance` (k = 0.0036, coefficient + # k · params["stage_hours"]) so the Exa and JuMP engines share one water balance. + stage_hours = Int(get(hydro_root, "stage_hours", 1)) + K = 0.0036 * stage_hours return HydroData(nHyd, units, upstream_turns, upstream_spills, K, initial_volumes, scenario_inflows, nScenarios, nStagesSample) @@ -390,3 +401,230 @@ function mean_inflow(hydro_data::HydroData, T::Int) end return w end + +# ── Stochastic demand (demand_scenarios.csv) ────────────────────────────────── +# +# Demand model shared with the SDDP engine (see sddp/sddp_demand_noise.jl in +# the MAIN repo): an i.i.d. per-stage MULTIPLICATIVE factor on every bus's +# active demand, +# +# ξ_t ∈ {1 − s, 1, 1 + s}, P = 1/3 each, independent of the inflow noise, +# +# with the spread s read from `/demand_scenarios.csv` (single line +# `s,`). On the ExaModels side ξ_t travels INSIDE the per-stage +# uncertainty vector: an augmented scenario is stage-major +# `[w_t; ξ_t]` (length T·(nHyd+1)), so the policy observes ξ_t exactly like it +# observes the stage inflow, and `prepare_solve!` applies base_demand·ξ_t to +# the p_demand parameter via `set_demand!` before every solve. + +# Seed of the column-keyed demand-noise protocol: eval scenario column c draws +# its demand path from StableRNG(DEMAND_NOISE_SEED + c). Shared by +# train_hydro_exa_strict.jl (protocol eval set) and eval_paired_exa_strict.jl +# (paired-500 protocol), so both see IDENTICAL demand paths per inflow column. +const DEMAND_NOISE_SEED = 20260714 + +""" + load_demand_spread(path::AbstractString) -> Union{Float64, Nothing} + +Read the demand-noise spread `s` from a `demand_scenarios.csv` file. + +The file holds a single data line `s,` (e.g. `s,0.10`) defining the +three-atom multiplicative demand distribution + +```math +\\xi_t \\in \\{1 - s,\\; 1,\\; 1 + s\\}, \\qquad P = \\tfrac{1}{3} \\text{ each}. +``` + +# Arguments +- `path::AbstractString`: path to `demand_scenarios.csv`. + +# Returns +- `Float64` spread `s ∈ [0, 1)` when the file exists. +- `nothing` when the file does not exist (deterministic demand — every code + path then behaves bit-identically to the pre-demand-noise implementation). +""" +function load_demand_spread(path::AbstractString) + # Missing file ⇒ deterministic demand (backwards-compatible default). + isfile(path) || return nothing + # Exactly one non-empty line carries the single `s,` record. + lines = [strip(line) for line in eachline(path) if !isempty(strip(line))] + length(lines) == 1 || error( + "demand_scenarios.csv must contain exactly one non-empty line `s,`; " * + "found $(length(lines))", + ) + line = only(lines) + # Split into the key token and the numeric value. + parts = split(line, ',') + # Enforce the exact two-field `s,` format shared by both engines. + length(parts) == 2 && strip(parts[1]) == "s" || + error("demand_scenarios.csv must contain a single line `s,`; got `$line`") + # Parse the spread value. + s = parse(Float64, strip(parts[2])) + # A spread ≥ 1 would make the low atom non-positive demand; forbid it. + 0.0 <= s < 1.0 || error("demand spread must satisfy 0 ≤ s < 1; got $s") + return s +end + +""" + demand_noise_atoms(spread::Real) -> Vector{Float64} + +Return the three equiprobable multiplicative demand atoms + +```math +\\{1 - s,\\; 1,\\; 1 + s\\} +``` + +for spread `s` (each with probability 1/3, i.i.d. across stages). +""" +demand_noise_atoms(spread::Real) = begin + 0.0 <= spread < 1.0 || error("demand spread must satisfy 0 ≤ s < 1; got $spread") + [1.0 - Float64(spread), 1.0, 1.0 + Float64(spread)] +end + +"""Return the demand-atom IDs for paired-protocol inflow column `column`.""" +function protocol_demand_atom_indices(T::Integer, column::Integer; + seed::Integer=DEMAND_NOISE_SEED) + T >= 0 || throw(ArgumentError("T must be nonnegative; got $T")) + column >= 1 || throw(ArgumentError("column must be positive; got $column")) + return rand(StableRNG(seed + column), 1:3, T) +end + +""" + sample_demand_factors(rng, T::Int, spread::Real) -> Vector{Float64} + +Draw `T` i.i.d. per-stage demand factors `ξ_t` from the three-atom +distribution `{1−s, 1, 1+s}` (probability 1/3 each) using `rng`. + +The draws are made SEQUENTIALLY (one `rand` per stage), so for a fixed seed the +length-`T₁` path is a prefix of the length-`T₂ ≥ T₁` path — training with +`T = 126` and evaluating with `T = 96` therefore share the first 96 factors of +each protocol column. + +# Arguments +- `rng`: any `AbstractRNG` (pass `StableRNG(DEMAND_NOISE_SEED + column)` for + protocol-paired draws). +- `T::Int`: number of stages. +- `spread::Real`: demand spread `s`. + +# Returns +- `Vector{Float64}` of length `T` with entries in `{1−s, 1, 1+s}`. +""" +function sample_demand_factors(rng, T::Int, spread::Real) + # The three equiprobable atoms. + atoms = demand_noise_atoms(spread) + # One sequential draw per stage (prefix property — see docstring). + return [atoms[rand(rng, 1:3)] for _ in 1:T] +end + +""" + protocol_demand_factors(spread::Real, T::Int, column::Int) -> Vector{Float64} + +Seeded demand-factor path for paired-protocol inflow column `column`: + +```math +\\xi^{(c)} = \\mathrm{sample\\_demand\\_factors}(\\mathrm{StableRNG}(\\mathrm{seed} + c),\\; T,\\; s). +``` + +Column-keyed seeding pairs the demand path with the inflow column: any script +evaluating column `c` (training-time protocol eval, paired-500 Exa eval, or +paired SDDP Historical eval) sees the IDENTICAL atom path, making the complete +joint uncertainty trajectory reproducible and paired across methods. + +# Arguments +- `spread::Real`: demand spread `s`. +- `T::Int`: number of stages. +- `column::Int`: paired-protocol scenario column id. + +# Returns +- `Vector{Float64}` of length `T`. +""" +protocol_demand_factors(spread::Real, T::Int, column::Int) = + demand_noise_atoms(spread)[protocol_demand_atom_indices(T, column)] + +""" + augment_scenario(w::AbstractVector, ξ::AbstractVector) -> Vector{Float64} + +Interleave a flat stage-major inflow trajectory `w` (length `T·nHyd`) with +per-stage demand factors `ξ` (length `T`) into the augmented stage-major +uncertainty vector + +```math +[w_1;\\, \\xi_1;\\; w_2;\\, \\xi_2;\\; \\ldots;\\; w_T;\\, \\xi_T] +``` + +of length `T·(nHyd+1)`, i.e. per-stage blocks `[w_t; ξ_t]`. This is the layout +the demand-noise DE builder sizes `p_inflow` for and the layout +`HydroReachablePolicy` slices (`n_uncertainty = nHyd + 1`, physical inflow = +first `nHyd` entries of each block). + +# Arguments +- `w::AbstractVector`: flat inflow trajectory, length divisible by `length(ξ)`. +- `ξ::AbstractVector`: per-stage demand factors, length `T`. + +# Returns +- `Vector{Float64}` of length `T·(nHyd+1)`. +""" +function augment_scenario(w::AbstractVector, ξ::AbstractVector) + # Number of stages comes from the factor vector. + T = length(ξ) + # Per-stage inflow width must divide the flat inflow length exactly. + nHyd, rem = divrem(length(w), T) + rem == 0 || error("length(w)=$(length(w)) is not divisible by T=$T") + # Allocate the augmented stage-major output. + out = Vector{Float64}(undef, T * (nHyd + 1)) + for t in 1:T + # Copy the stage-t inflow block. + out[(t-1)*(nHyd+1)+1 : (t-1)*(nHyd+1)+nHyd] = @view w[(t-1)*nHyd+1 : t*nHyd] + # Append the stage-t demand factor as the block's last entry. + out[t*(nHyd+1)] = ξ[t] + end + return out +end + +""" + sample_scenario(hydro_data, T, demand_spread; rng = Random.default_rng()) + -> Vector{Float64} + +Demand-noise variant of [`sample_scenario`](@ref): draws the joint inflow +trajectory AND i.i.d. per-stage demand factors + +```math +\\xi_t \\sim \\mathrm{Uniform}\\{1-s,\\; 1,\\; 1+s\\} +``` + +(independent of the inflow draw), returning the augmented stage-major vector +`[w_t; ξ_t]` of length `T·(nHyd+1)` (see [`augment_scenario`](@ref)). + +# Arguments +- `hydro_data::HydroData`: inflow scenario data. +- `T::Int`: number of stages. +- `demand_spread::Real`: demand spread `s`. + +# Keywords +- `rng`: random number generator (defaults to the task-global RNG, so + `Random.seed!` seeds it exactly like the 2-arg method). + +# Returns +- `Vector{Float64}` of length `T·(nHyd+1)`. +""" +function sample_scenario(hydro_data::HydroData, T::Int, demand_spread::Real; + rng = Random.default_rng()) + nHyd = hydro_data.nHyd + # The three equiprobable demand atoms. + atoms = demand_noise_atoms(demand_spread) + # Augmented per-stage width: nHyd inflows + 1 demand factor. + nu = nHyd + 1 + w = Vector{Float64}(undef, T * nu) + for t in 1:T + # Cyclic raw-row mapping (same convention as the 2-arg method). + t_row = mod1(t, hydro_data.nStagesSample) + # One joint inflow index per stage — all reservoirs share it. + j = rand(rng, 1:hydro_data.nScenarios) + for r in 1:nHyd + w[(t-1)*nu + r] = hydro_data.scenario_inflows[r][t_row, j] + end + # Independent demand draw for the same stage (separate rand call). + w[t*nu] = atoms[rand(rng, 1:3)] + end + return w +end diff --git a/examples/HydroPowerModels/hydro_power_exa.jl b/examples/HydroPowerModels/hydro_power_exa.jl index 40feee9..9ff05a6 100644 --- a/examples/HydroPowerModels/hydro_power_exa.jl +++ b/examples/HydroPowerModels/hydro_power_exa.jl @@ -15,12 +15,22 @@ # reservoir dynamics (initial condition, water balance, turbine coupling) # p_target[t,r] − reservoir[t+1,r] − δ[t,r] = 0 ← ADDED LAST # -# Target constraints are added last so result.multipliers[target_con_range] -# gives ∇_{x̂} Q directly (envelope theorem for policy gradient). +# Target constraints are added last in non-strict mode so +# target_multipliers(prob, result) gives ∇_{x̂} Q. In strict mode the +# reservoir trajectory is a parameter and target_multipliers transforms +# water-balance duals into target sensitivities. +# +# Strict-mode invariant: with reservoir a parameter, the initial condition +# reservoir[1,r] − p_x0[r] = 0 would be a parameter-only constraint (an +# all-zero Jacobian row), so the builders omit it and the initial condition is +# enforced purely by data: prepare_solve! must keep p_reservoir[1:nHyd] equal +# to the initial state x0 (it writes reservoir_vals = vcat(init, xhat)). Any +# code that updates p_reservoir must preserve p_reservoir[1:nH] == x0. using ExaModels using MadNLP using LinearAlgebra +import DecisionRulesExa: prepare_solve!, target_multipliers # ── Index helpers ───────────────────────────────────────────────────────────── # All arrays are flat, stage-major: index (t, i) → (t-1)*n + i @@ -30,6 +40,15 @@ using LinearAlgebra @inline _bri(nBR, t, br) = (t-1)*nBR + br # branch / pf @inline _ri(nH, t, r) = (t-1)*nH + r # hydro: reservoir / outflow / spill / delta +# ── Cascade link: upstream→downstream water-balance coupling ──────────────── + +struct CascadeLink + downstream::Int # array position of downstream unit + upstream::Int # array position of upstream unit + turn_only::Bool # true if only turbine outflow (not spill) reaches downstream + K_max_turn::Float32 # K × max_turn of the upstream unit +end + # ── Problem struct ──────────────────────────────────────────────────────────── """ @@ -41,8 +60,25 @@ Parameters: - `p_demand` : per-bus per-stage active demand (length T*nBus, stage-major) - `p_reactive_demand` : per-bus per-stage reactive demand (length T*nBus; nothing for DC) - `p_x0` : initial reservoir levels (length nHyd) -- `p_inflow` : inflow trajectory (length T*nHyd, stage-major) +- `p_inflow` : per-stage uncertainty trajectory (length T*n_uncertainty, + stage-major). Without demand noise `n_uncertainty = nHyd` + and this is exactly the inflow trajectory. With demand + noise (`demand_spread` set at build time) + `n_uncertainty = nHyd + 1` and each stage block is + `[w_t; ξ_t]`: the first nHyd entries are inflows, the + last entry is the multiplicative demand factor ξ_t. + The ξ slots appear in NO constraint — `prepare_solve!` + reads them and applies `base_demand[t,:] · ξ_t` to + `p_demand` via `set_demand!` before every solve. - `p_target` : NN-predicted target levels (length T*nHyd, stage-major) + +Demand-noise fields: +- `n_uncertainty::Int` : per-stage uncertainty width (nHyd, or nHyd+1 with noise) +- `base_demand` : `nothing` (deterministic demand — bit-identical legacy + behavior) or the `[T × nBus]` BASE active demand + (already `load_scaler`-scaled) that `prepare_solve!` + multiplies by ξ_t. Its CONTENT may be mutated in place + (the rollout stage problem refreshes row 1 per stage). """ struct HydroExaDEProblem core @@ -65,8 +101,28 @@ struct HydroExaDEProblem horizon::Int # formulation formulation::Symbol # :dc or :ac_polar - # range into result.multipliers for target constraints + # range into result.multipliers for target constraints (or water balance in strict mode) target_con_range::UnitRange{Int} + strict_targets::Bool + # strict mode: reservoir is a parameter, not a variable (length (T+1)*nHyd) + p_reservoir # nothing when !strict_targets + strict_reservoir_values::Vector{Float64} + # cascade data for target clamping (empty if no upstream connections) + cascade::Vector{CascadeLink} + K::Float64 + min_turn::Vector{Float64} + max_vol::Vector{Float64} + # Reactive-slack mode of the AC builder (:none for DC): + # :free — single FREE zero-cost deficit_q variable per bus/stage + # :penalized — deficit_q = δq⁺ − δq⁻ (δq± ≥ 0) with linear cost c·Σ(δq⁺+δq⁻) + # :hard — no reactive slack (hard reactive balance, MAIN-faithful) + reactive_deficit_mode::Symbol + # Per-stage uncertainty width: nHyd (deterministic demand) or nHyd + 1 + # (stochastic demand: each stage block of p_inflow is [w_t; ξ_t]). + n_uncertainty::Int + # nothing (deterministic demand) or the [T × nBus] load_scaler-scaled BASE + # active demand: prepare_solve! sets p_demand[t,:] = base_demand[t,:] · ξ_t. + base_demand::Union{Nothing,Matrix{Float64}} end # ── AC branch coefficient helper ────────────────────────────────────────────── @@ -146,7 +202,10 @@ end backend=nothing, float_type=Float64, formulation=:dc, target_penalty=:auto, target_penalty_l1=:auto, demand_matrix=nothing, - reactive_demand_matrix=nothing, deficit_cost=nothing) + reactive_demand_matrix=nothing, deficit_cost=nothing, + load_scaler=0.6, strict_targets=true, + reactive_deficit_cost=Inf, + demand_spread=nothing) -> HydroExaDEProblem Build the T-stage hydro-power deterministic equivalent. @@ -161,8 +220,8 @@ If nothing, uses `power_data.default_bus_demand` for all stages. `reactive_demand_matrix` is an optional [T × nBus] matrix for reactive demand (AC only). If nothing, uses `power_data.default_bus_reactive_demand`. -`deficit_cost` overrides `power_data.cost_deficit` (the load-shedding penalty per pu). -Pass a large value (e.g. 1e5, >> max thermal cost) to effectively enforce hard KCL. +`deficit_cost` overrides the canonical operational active-deficit coefficient +`power_data.cost_deficit * power_data.baseMVA` (6000 USD/(pu·stage) for MAIN). `target_penalty` sets the L2 coefficient ρ for the `(ρ/2)·δ²` target slack penalty. Pass `:auto` (default) to use `2 × max_gen_cost`, matching JuMP's `penalty_l2 = :auto`. @@ -170,6 +229,37 @@ Pass `:auto` (default) to use `2 × max_gen_cost`, matching JuMP's `penalty_l2 = `target_penalty_l1` sets the L1 coefficient for the `λ·|δ|` target slack penalty. Pass `:auto` (default) to use the same value as L2 ρ. Pass `nothing` to disable L1. The L1 term is reformulated as `λ·(δ⁺ + δ⁻)` with `δ = δ⁺ − δ⁻`, `δ⁺,δ⁻ ≥ 0`. + +`reactive_deficit_cost` (AC only) controls the reactive slack `deficit_q` in the +reactive KCL: +- `nothing` — legacy behavior: one FREE, zero-cost `deficit_q` + variable per bus/stage (relaxes the reactive balance; model is byte-identical + to builds preceding this kwarg). +- finite `c ≥ 0` — `|deficit_q|` is penalized linearly: the slack is split into + `deficit_q = δq⁺ − δq⁻` with `δq⁺, δq⁻ ≥ 0` and objective `c·Σ(δq⁺ + δq⁻)`. + Constraint count and ORDER are unchanged (the split variables enter the same + reactive-KCL rows), so `target_con_range` is unaffected in both strict and + non-strict modes; only variable count changes (+T·nBus vs the default). +- `Inf` — omit `deficit_q` entirely: hard reactive balance, matching MAIN's + ACPPowerModel.mof.json (−T·nBus variables vs the default; constraint + count/order again unchanged). +Passing a non-`nothing` value with `formulation = :dc` throws (DC has no +reactive balance). + +`demand_spread` enables stochastic demand (i.i.d. per-stage multiplicative +factor `ξ_t ∈ {1−s, 1, 1+s}`, probability 1/3 each, independent of the inflow +noise — the same model the SDDP baselines implement via +sddp/sddp_demand_noise.jl): +- `nothing` (default) — deterministic demand; the model is BIT-IDENTICAL to + builds preceding this kwarg. +- spread `s ∈ [0, 1)` — the uncertainty parameter `p_inflow` grows to + `T·(nHyd+1)` with per-stage blocks `[w_t; ξ_t]`; the water-balance inflow + indices stride `nHyd+1`; the ξ slots are referenced by NO constraint and are + consumed by `prepare_solve!`, which applies + `p_demand[t,:] = base_demand[t,:] · ξ_t` via [`set_demand!`](@ref) before + every solve. Only the ACTIVE demand is perturbed (reactive demand stays at + its base value), matching the SDDP override which adds the deviation to the + real-power KCL only. """ function build_hydro_de(power_data::PowerData, hydro_data::HydroData, @@ -182,19 +272,29 @@ function build_hydro_de(power_data::PowerData, demand_matrix = nothing, reactive_demand_matrix = nothing, deficit_cost::Union{Nothing,Real} = nothing, - load_scaler::Real = 1.0) + load_scaler::Real = 0.6, + strict_targets::Bool = true, + reactive_deficit_cost::Union{Nothing,Real} = Inf, + demand_spread::Union{Nothing,Real} = nothing) formulation in (:dc, :ac_polar) || error("formulation must be :dc or :ac_polar, got :$formulation") + # Validate the demand-noise spread once for both formulations. + demand_spread === nothing || 0.0 <= demand_spread < 1.0 || + error("demand_spread must satisfy 0 ≤ s < 1; got $demand_spread") if formulation === :dc + reactive_deficit_cost === nothing || + error("reactive_deficit_cost applies only to formulation = :ac_polar (DC has no reactive balance)") return _build_dc_hydro_de(power_data, hydro_data, T; backend=backend, float_type=float_type, target_penalty=target_penalty, target_penalty_l1=target_penalty_l1, demand_matrix=demand_matrix, deficit_cost=deficit_cost, - load_scaler=load_scaler) + load_scaler=load_scaler, + strict_targets=strict_targets, + demand_spread=demand_spread) else return _build_ac_hydro_de(power_data, hydro_data, T; backend=backend, float_type=float_type, @@ -203,10 +303,33 @@ function build_hydro_de(power_data::PowerData, demand_matrix=demand_matrix, reactive_demand_matrix=reactive_demand_matrix, deficit_cost=deficit_cost, - load_scaler=load_scaler) + load_scaler=load_scaler, + strict_targets=strict_targets, + reactive_deficit_cost=reactive_deficit_cost, + demand_spread=demand_spread) end end +function _build_cascade_links(hydro_data::HydroData) + K = Float64(hydro_data.K) + spill_dests = Dict{Int,Set{Int}}() + for conn in hydro_data.upstream_spills + push!(get!(spill_dests, conn.upstream_pos, Set{Int}()), conn.downstream_pos) + end + cascade = CascadeLink[] + for conn in hydro_data.upstream_turns + d, u = conn.downstream_pos, conn.upstream_pos + has_spill = haskey(spill_dests, u) && d in spill_dests[u] + push!(cascade, CascadeLink(d, u, !has_spill, Float32(K * hydro_data.units[u].max_turn))) + end + for conn in hydro_data.upstream_spills + d, u = conn.downstream_pos, conn.upstream_pos + already = any(c -> c.downstream == d && c.upstream == u, cascade) + already || push!(cascade, CascadeLink(d, u, false, Float32(K * hydro_data.units[u].max_turn))) + end + return cascade +end + # ── DC builder ──────────────────────────────────────────────────────────────── function _build_dc_hydro_de(power_data::PowerData, @@ -218,12 +341,17 @@ function _build_dc_hydro_de(power_data::PowerData, target_penalty_l1::Union{Real,Symbol,Nothing} = :auto, demand_matrix = nothing, deficit_cost::Union{Nothing,Real} = nothing, - load_scaler::Real = 1.0) + load_scaler::Real = 0.6, + strict_targets::Bool = true, + demand_spread::Union{Nothing,Real} = nothing) nBus = power_data.nBus nGen = power_data.nGen nBranch = power_data.nBranch nHyd = hydro_data.nHyd + # Per-stage uncertainty width: +1 slot for the demand factor ξ_t when + # stochastic demand is enabled (see build_hydro_de docstring). + n_unc = nHyd + (demand_spread === nothing ? 0 : 1) K = float_type(hydro_data.K) ρ = float_type(target_penalty === :auto ? auto_target_penalty(power_data, hydro_data) : target_penalty) ρ_l1 = if target_penalty_l1 === :auto @@ -235,7 +363,8 @@ function _build_dc_hydro_de(power_data::PowerData, end use_l1 = ρ_l1 > 0 baseMVA = float_type(power_data.baseMVA) - cd = float_type(deficit_cost !== nothing ? deficit_cost : power_data.cost_deficit) + cd = float_type(deficit_cost !== nothing ? deficit_cost : + power_data.cost_deficit * power_data.baseMVA) core = ExaModels.ExaCore(float_type; backend = backend) @@ -258,9 +387,17 @@ function _build_dc_hydro_de(power_data::PowerData, deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) # Reservoir levels: (T+1)*nHyd - res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) - res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) - reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + # In strict mode, reservoir = [x0; targets] is predetermined → parameter. + # Eliminates (T+1)*nHyd variables and nHyd + T*nHyd constraints. + if strict_targets + p_reservoir = ExaModels.parameter(core, zeros(float_type, (T+1) * nHyd)) + reservoir = p_reservoir + else + p_reservoir = nothing + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + end # Turbine outflow: T*nHyd out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) @@ -271,8 +408,10 @@ function _build_dc_hydro_de(power_data::PowerData, spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) # Target slack: δ = δ⁺ − δ⁻ with δ⁺,δ⁻ ≥ 0 (L1+L2 Lagrangian penalty) - delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) - delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + if !strict_targets + delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + end # ── Parameters ──────────────────────────────────────────────────────────── @@ -282,9 +421,19 @@ function _build_dc_hydro_de(power_data::PowerData, float_type.(load_scaler .* repeat(power_data.default_bus_demand, T)) end + # Uncertainty parameter: length T·n_unc. With demand noise the per-stage + # block is [w_t; ξ_t]; initialize the ξ slots to 1.0 (factor-1 demand) so a + # solve before the first set_parameter! sees the base demand. + init_uncertainty = zeros(float_type, T * n_unc) + if demand_spread !== nothing + for t in 1:T + init_uncertainty[t * n_unc] = one(float_type) # ξ_t slot ← 1.0 + end + end + p_demand = ExaModels.parameter(core, init_demand) p_x0 = ExaModels.parameter(core, zeros(float_type, nHyd)) - p_inflow = ExaModels.parameter(core, zeros(float_type, T * nHyd)) + p_inflow = ExaModels.parameter(core, init_uncertainty) p_target = ExaModels.parameter(core, zeros(float_type, T * nHyd)) p_penalty_half = ExaModels.parameter(core, fill(float_type(ρ / 2), T * nHyd)) p_penalty_l1 = ExaModels.parameter(core, fill(ρ_l1, T * nHyd)) @@ -306,19 +455,22 @@ function _build_dc_hydro_de(power_data::PowerData, for item in def_cost_items ) - # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² - delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] - ExaModels.objective(core, - p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 - for item in delta_items - ) + if !strict_targets + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] - # L1 penalty: λ·(δ⁺ + δ⁻) - if use_l1 + # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² ExaModels.objective(core, - p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 for item in delta_items ) + + # L1 penalty: λ·(δ⁺ + δ⁻) + if use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end end # ── Constraints ─────────────────────────────────────────────────────────── @@ -391,20 +543,29 @@ function _build_dc_hydro_de(power_data::PowerData, ) n_con += T * nBus - # 5. Initial reservoir condition - ic_items = [(r = r,) for r in 1:nHyd] - ExaModels.constraint(core, - reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] - for item in ic_items - ) - n_con += nHyd + # 5. Initial reservoir condition (skip in strict: reservoir is a parameter, + # so reservoir[1,r] − p_x0[r] = 0 would be parameter-only — an all-zero + # Jacobian row. The initial condition is instead maintained by the invariant + # that prepare_solve! writes p_reservoir[1:nHyd] = x0; see file-top comment.) + if !strict_targets + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + end - # 6. Water balance + # 6. Water balance (reservoir is a parameter in strict mode — same expression) + wb_con_start = n_con + 1 wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), out_idx = _ri(nHyd, t, r), spill_idx = _ri(nHyd, t, r), - inflow_p = _ri(nHyd, t, r), + # Inflow slot inside the uncertainty parameter: stride n_unc + # (= nHyd without demand noise — unchanged; = nHyd+1 with it, + # skipping the per-stage ξ_t slot). + inflow_p = _ri(n_unc, t, r), K = K) for t in 1:T for r in 1:nHyd] c_wb = ExaModels.constraint(core, @@ -447,26 +608,45 @@ function _build_dc_hydro_de(power_data::PowerData, n_con += T * nHyd # ── TARGET CONSTRAINTS (ADDED LAST) ─────────────────────────────────────── - # x̂ − x − (δ⁺ − δ⁻) = 0 - target_items = [(param_idx = _ri(nHyd, t, r), - res_idx = _ri(nHyd, t+1, r), - delta_idx = _ri(nHyd, t, r)) - for t in 1:T for r in 1:nHyd] - ExaModels.constraint(core, - p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] - for item in target_items - ) - target_con_range = (n_con + 1):(n_con + T * nHyd) + if strict_targets + # Strict: reservoir is a parameter. target_multipliers transforms these + # water-balance duals into ∇_{x̂} Q by the adjacent-stage chain rule. + target_con_range = wb_con_start:(wb_con_start + T * nHyd - 1) + else + # x̂ − x − (δ⁺ − δ⁻) = 0 + target_items = [(param_idx = _ri(nHyd, t, r), + res_idx = _ri(nHyd, t+1, r), + delta_idx = _ri(nHyd, t, r)) + for t in 1:T for r in 1:nHyd] + ExaModels.constraint(core, + p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] + for item in target_items + ) + target_con_range = (n_con + 1):(n_con + T * nHyd) + end model = ExaModels.ExaModel(core) + cascade = _build_cascade_links(hydro_data) + return HydroExaDEProblem( core, model, p_demand, nothing, p_x0, p_inflow, p_target, p_penalty_half, Float64(ρ / 2), p_penalty_l1, Float64(ρ_l1), nHyd, nBus, nGen, nBranch, T, - :dc, target_con_range, + :dc, target_con_range, strict_targets, + p_reservoir, + strict_targets ? zeros(Float64, (T + 1) * nHyd) : Float64[], + cascade, Float64(K), + Float64.([h.min_turn for h in hydro_data.units]), + Float64.([h.max_vol for h in hydro_data.units]), + :none, # DC has no reactive balance, hence no reactive slack + n_unc, # per-stage uncertainty width (nHyd, or nHyd+1 with demand noise) + # BASE demand [T × nBus] for prepare_solve!'s ξ_t multiplication; + # nothing keeps the deterministic path bit-identical. + demand_spread === nothing ? nothing : + Float64.(permutedims(reshape(init_demand, nBus, T))), ) end @@ -482,12 +662,18 @@ function _build_ac_hydro_de(power_data::PowerData, demand_matrix = nothing, reactive_demand_matrix = nothing, deficit_cost::Union{Nothing,Real} = nothing, - load_scaler::Real = 1.0) + load_scaler::Real = 0.6, + strict_targets::Bool = true, + reactive_deficit_cost::Union{Nothing,Real} = Inf, + demand_spread::Union{Nothing,Real} = nothing) nBus = power_data.nBus nGen = power_data.nGen nBranch = power_data.nBranch nHyd = hydro_data.nHyd + # Per-stage uncertainty width: +1 slot for the demand factor ξ_t when + # stochastic demand is enabled (see build_hydro_de docstring). + n_unc = nHyd + (demand_spread === nothing ? 0 : 1) K = float_type(hydro_data.K) ρ = float_type(target_penalty === :auto ? auto_target_penalty(power_data, hydro_data) : target_penalty) ρ_l1 = if target_penalty_l1 === :auto @@ -499,7 +685,22 @@ function _build_ac_hydro_de(power_data::PowerData, end use_l1 = ρ_l1 > 0 baseMVA = float_type(power_data.baseMVA) - cd = float_type(deficit_cost !== nothing ? deficit_cost : power_data.cost_deficit) + cd = float_type(deficit_cost !== nothing ? deficit_cost : + power_data.cost_deficit * power_data.baseMVA) + + # Reactive-slack mode (see build_hydro_de docstring): + # nothing → :free (default; byte-identical to builds preceding the kwarg), + # finite c ≥ 0 → :penalized (|deficit_q| costed linearly via δq⁺/δq⁻), + # Inf → :hard (no reactive slack — MAIN-faithful hard reactive balance). + rq_mode = if reactive_deficit_cost === nothing + :free + elseif isinf(Float64(reactive_deficit_cost)) + :hard + else + (isnan(Float64(reactive_deficit_cost)) || reactive_deficit_cost < 0) && + error("reactive_deficit_cost must be nothing, a finite cost ≥ 0, or Inf; got $reactive_deficit_cost") + :penalized + end core = ExaModels.ExaCore(float_type; backend = backend) @@ -538,13 +739,28 @@ function _build_ac_hydro_de(power_data::PowerData, # Active deficit (load shedding): T*nBus (non-negative) deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) - # Reactive deficit (free; allows reactive balance at zero cost) - deficit_q = ExaModels.variable(core, T * nBus) + # Reactive deficit — declared in place of the historical free variable so + # the variable ORDER of the :free mode is byte-identical to older builds. + if rq_mode === :free + # Free, zero-cost slack (relaxes the reactive balance). + deficit_q = ExaModels.variable(core, T * nBus) + elseif rq_mode === :penalized + # Split slack deficit_q = δq⁺ − δq⁻ with δq⁺, δq⁻ ≥ 0; the linear cost + # c·Σ(δq⁺ + δq⁻) is added in the objective section below. + deficit_q_pos = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + deficit_q_neg = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + end # :hard → no reactive slack variable at all # Reservoir levels: (T+1)*nHyd - res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) - res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) - reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + if strict_targets + p_reservoir = ExaModels.parameter(core, zeros(float_type, (T+1) * nHyd)) + reservoir = p_reservoir + else + p_reservoir = nothing + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + end # Turbine outflow: T*nHyd out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) @@ -555,8 +771,10 @@ function _build_ac_hydro_de(power_data::PowerData, spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) # Target slack: δ = δ⁺ − δ⁻ with δ⁺,δ⁻ ≥ 0 (L1+L2 Lagrangian penalty) - delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) - delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + if !strict_targets + delta_pos = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + delta_neg = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + end # ── Parameters ──────────────────────────────────────────────────────────── @@ -572,10 +790,20 @@ function _build_ac_hydro_de(power_data::PowerData, float_type.(load_scaler .* repeat(power_data.default_bus_reactive_demand, T)) end + # Uncertainty parameter: length T·n_unc. With demand noise the per-stage + # block is [w_t; ξ_t]; initialize the ξ slots to 1.0 (factor-1 demand) so a + # solve before the first set_parameter! sees the base demand. + init_uncertainty = zeros(float_type, T * n_unc) + if demand_spread !== nothing + for t in 1:T + init_uncertainty[t * n_unc] = one(float_type) # ξ_t slot ← 1.0 + end + end + p_demand = ExaModels.parameter(core, init_demand) p_reactive_demand = ExaModels.parameter(core, init_reactive_demand) p_x0 = ExaModels.parameter(core, zeros(float_type, nHyd)) - p_inflow = ExaModels.parameter(core, zeros(float_type, T * nHyd)) + p_inflow = ExaModels.parameter(core, init_uncertainty) p_target = ExaModels.parameter(core, zeros(float_type, T * nHyd)) p_penalty_half = ExaModels.parameter(core, fill(float_type(ρ / 2), T * nHyd)) p_penalty_l1 = ExaModels.parameter(core, fill(ρ_l1, T * nHyd)) @@ -601,22 +829,43 @@ function _build_ac_hydro_de(power_data::PowerData, for item in def_cost_items ) - # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² - delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] - ExaModels.objective(core, - p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 - for item in delta_items - ) + # Linear reactive-slack penalty c·Σ(δq⁺ + δq⁻) = c·Σ|deficit_q| (only in + # :penalized mode; :free and :hard add no objective term here, keeping the + # default model byte-identical). + if rq_mode === :penalized + cq = float_type(reactive_deficit_cost) + rq_cost_items = [(idx = _bi(nBus, t, b), c = cq) for t in 1:T for b in 1:nBus] + ExaModels.objective(core, + item.c * (deficit_q_pos[item.idx] + deficit_q_neg[item.idx]) + for item in rq_cost_items + ) + end - # L1 penalty: λ·(δ⁺ + δ⁻) - if use_l1 + if !strict_targets + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] + + # L2 penalty: (ρ/2)·(δ⁺ − δ⁻)² ExaModels.objective(core, - p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 for item in delta_items ) + + # L1 penalty: λ·(δ⁺ + δ⁻) + if use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end end # ── Constraints ─────────────────────────────────────────────────────────── + # NOTE (target_con_range audit): the reactive-slack modes differ ONLY in + # variables and objective terms. Constraint COUNT and ORDER are identical + # in all three modes (:penalized adds constraint! terms to EXISTING + # reactive-KCL rows; :hard omits the slack term from those same rows), so + # the n_con accounting below and the resulting target_con_range are + # unaffected in both strict and non-strict modes. n_con = 0 # 1. Reference angle: va[t, ref] = 0 @@ -704,7 +953,29 @@ function _build_ac_hydro_de(power_data::PowerData, ) n_con += T * nBranch - # 7. Active KCL: pd + gs·vm² − deficit − Σpg + Σp_fr + Σp_to = 0 + # 7. Apparent-power thermal limits at BOTH branch ends. PowerModels' + # ACPPowerModel imposes p^2 + q^2 <= rate_a^2 at the from and to ends. + # Component-wise variable bounds alone describe a square and are strictly + # weaker; omitting these circles made the Exa AC feasible set larger than + # the SDDP forward model, especially on deliberately derated corridors. + thermal_ub = float_type.(repeat([br.rate_a^2 for br in power_data.branches], T)) + thermal_items = [(t = t, br = br_pos) + for t in 1:T for br_pos in 1:nBranch] + ExaModels.constraint(core, + p_fr[_bri(nBranch, item.t, item.br)]^2 + + q_fr[_bri(nBranch, item.t, item.br)]^2 + for item in thermal_items; + lcon = fill(float_type(-Inf), T * nBranch), ucon = thermal_ub, + ) + ExaModels.constraint(core, + p_to[_bri(nBranch, item.t, item.br)]^2 + + q_to[_bri(nBranch, item.t, item.br)]^2 + for item in thermal_items; + lcon = fill(float_type(-Inf), T * nBranch), ucon = thermal_ub, + ) + n_con += 2 * T * nBranch + + # 8. Active KCL: pd + gs·vm² − deficit − Σpg + Σp_fr + Σp_to = 0 kcl_p_init = [(t = t, b = b, gs = float_type(power_data.buses[b].gs)) for t in 1:T for b in 1:nBus] c_kcl_p = ExaModels.constraint(core, @@ -738,7 +1009,7 @@ function _build_ac_hydro_de(power_data::PowerData, ) n_con += T * nBus - # 8. Reactive KCL: qd − bs·vm² − deficit_q − Σqg + Σq_fr + Σq_to = 0 + # 9. Reactive KCL: qd − bs·vm² − deficit_q − Σqg + Σq_fr + Σq_to = 0 kcl_q_init = [(t = t, b = b, bs = float_type(power_data.buses[b].bs)) for t in 1:T for b in 1:nBus] c_kcl_q = ExaModels.constraint(core, @@ -764,26 +1035,49 @@ function _build_ac_hydro_de(power_data::PowerData, item.brow => q_to[item.bcol] for item in kcl_qto_items ) - ExaModels.constraint!(core, c_kcl_q, - item.brow => -deficit_q[item.dcol] - for item in kcl_def_items - ) + # Reactive slack term: modifies the EXISTING reactive-KCL rows only — + # constraint count/order identical across all rq_mode values. + if rq_mode === :free + ExaModels.constraint!(core, c_kcl_q, + item.brow => -deficit_q[item.dcol] + for item in kcl_def_items + ) + elseif rq_mode === :penalized + # deficit_q = δq⁺ − δq⁻ enters the KCL as −δq⁺ + δq⁻. + ExaModels.constraint!(core, c_kcl_q, + item.brow => -deficit_q_pos[item.dcol] + for item in kcl_def_items + ) + ExaModels.constraint!(core, c_kcl_q, + item.brow => deficit_q_neg[item.dcol] + for item in kcl_def_items + ) + end # :hard → no slack term: hard reactive balance n_con += T * nBus - # 9. Initial reservoir condition - ic_items = [(r = r,) for r in 1:nHyd] - ExaModels.constraint(core, - reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] - for item in ic_items - ) - n_con += nHyd + # 10. Initial reservoir condition (skip in strict: reservoir is a parameter, + # so reservoir[1,r] − p_x0[r] = 0 would be parameter-only — an all-zero + # Jacobian row. The initial condition is instead maintained by the invariant + # that prepare_solve! writes p_reservoir[1:nHyd] = x0; see file-top comment.) + if !strict_targets + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + end - # 10. Water balance + # 11. Water balance (reservoir is a parameter in strict mode — same expression) + wb_con_start = n_con + 1 wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), out_idx = _ri(nHyd, t, r), spill_idx = _ri(nHyd, t, r), - inflow_p = _ri(nHyd, t, r), + # Inflow slot inside the uncertainty parameter: stride n_unc + # (= nHyd without demand noise — unchanged; = nHyd+1 with it, + # skipping the per-stage ξ_t slot). + inflow_p = _ri(n_unc, t, r), K = K) for t in 1:T for r in 1:nHyd] c_wb = ExaModels.constraint(core, @@ -826,26 +1120,45 @@ function _build_ac_hydro_de(power_data::PowerData, n_con += T * nHyd # ── TARGET CONSTRAINTS (ADDED LAST) ─────────────────────────────────────── - # x̂ − x − (δ⁺ − δ⁻) = 0 - target_items = [(param_idx = _ri(nHyd, t, r), - res_idx = _ri(nHyd, t+1, r), - delta_idx = _ri(nHyd, t, r)) - for t in 1:T for r in 1:nHyd] - ExaModels.constraint(core, - p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] - for item in target_items - ) - target_con_range = (n_con + 1):(n_con + T * nHyd) + if strict_targets + # Strict: reservoir is a parameter. target_multipliers transforms these + # water-balance duals into ∇_{x̂} Q by the adjacent-stage chain rule. + target_con_range = wb_con_start:(wb_con_start + T * nHyd - 1) + else + # x̂ − x − (δ⁺ − δ⁻) = 0 + target_items = [(param_idx = _ri(nHyd, t, r), + res_idx = _ri(nHyd, t+1, r), + delta_idx = _ri(nHyd, t, r)) + for t in 1:T for r in 1:nHyd] + ExaModels.constraint(core, + p_target[item.param_idx] - reservoir[item.res_idx] - delta_pos[item.delta_idx] + delta_neg[item.delta_idx] + for item in target_items + ) + target_con_range = (n_con + 1):(n_con + T * nHyd) + end model = ExaModels.ExaModel(core) + cascade = _build_cascade_links(hydro_data) + return HydroExaDEProblem( core, model, p_demand, p_reactive_demand, p_x0, p_inflow, p_target, p_penalty_half, Float64(ρ / 2), p_penalty_l1, Float64(ρ_l1), nHyd, nBus, nGen, nBranch, T, - :ac_polar, target_con_range, + :ac_polar, target_con_range, strict_targets, + p_reservoir, + strict_targets ? zeros(Float64, (T + 1) * nHyd) : Float64[], + cascade, Float64(K), + Float64.([h.min_turn for h in hydro_data.units]), + Float64.([h.max_vol for h in hydro_data.units]), + rq_mode, + n_unc, # per-stage uncertainty width (nHyd, or nHyd+1 with demand noise) + # BASE demand [T × nBus] for prepare_solve!'s ξ_t multiplication; + # nothing keeps the deterministic path bit-identical. + demand_spread === nothing ? nothing : + Float64.(permutedims(reshape(init_demand, nBus, T))), ) end @@ -868,15 +1181,124 @@ end """ set_inflows!(prob, w) -Set the inflow trajectory. `w` is a flat vector of length `T*nHyd` (stage-major). +Set the uncertainty trajectory. `w` is a flat stage-major vector of length +`T*n_uncertainty`: without demand noise this is exactly the inflow trajectory +(`n_uncertainty = nHyd`); with demand noise each stage block is `[w_t; ξ_t]` +(`n_uncertainty = nHyd + 1`). """ function set_inflows!(prob::HydroExaDEProblem, w::AbstractVector) - expected = prob.horizon * prob.nHyd - length(w) == expected || error("w must have length T*nHyd=$expected") + # Expected flat length follows the problem's per-stage uncertainty width. + expected = prob.horizon * prob.n_uncertainty + length(w) == expected || error("w must have length T*n_uncertainty=$expected") ExaModels.set_parameter!(prob.core, prob.p_inflow, w) return prob end +""" + prepare_solve!(prob::HydroExaDEProblem, init_state, w_flat, xhat_flat) -> Nothing + +Pre-solve hook (called by the training loop, the rollout callback, and the +one-shot scripts immediately before every MadNLP solve). Two jobs: + +1. **Stochastic demand** (only when the problem was built with + `demand_spread`): extract the per-stage demand factors `ξ_t` from the + augmented uncertainty vector (`w_flat[t·n_unc]`, the last entry of each + stage block `[w_t; ξ_t]`) and write the realized demand + + ```math + pd_{t,b} = \\mathrm{base\\_demand}_{t,b} \\cdot \\xi_t + ``` + + into the `p_demand` parameter via [`set_demand!`](@ref). This runs in BOTH + the deterministic-equivalent training path (full-horizon `w`) and the + stage-wise rollout path (1-stage `w`), including inside multi-GPU worker + tasks (each worker's DE carries its own `base_demand`). +2. **Strict-mode reservoir pinning** (only when `strict_targets`): apply the + Float64 cascade clamp to the target trajectory and write + `p_reservoir = [x0; x̂]`, preserving the invariant + `p_reservoir[1:nHyd] == x0` (see file-top comment). + +With `base_demand === nothing` and `p_reservoir === nothing` this is a no-op — +bit-identical to the pre-demand-noise implementation. +""" +function prepare_solve!(prob::HydroExaDEProblem, init_state, w_flat, xhat_flat) + nH = prob.nHyd + # Per-stage uncertainty stride (nHyd, or nHyd+1 with demand noise). + nu = prob.n_uncertainty + + # ── 1. Stochastic demand: p_demand[t,:] = base_demand[t,:] · ξ_t ───────── + if prob.base_demand !== nothing + T = prob.horizon + # Materialize the (possibly GPU) uncertainty vector once on the CPU. + w_cpu = Float64.(vec(Array(w_flat))) + # Guard: the caller must pass the AUGMENTED vector for a noise-enabled DE. + length(w_cpu) == T * nu || + error("demand-noise DE expects w of length T*n_uncertainty=$(T*nu); got $(length(w_cpu))") + # Realized demand matrix: scale each base row by its stage factor. + demand = similar(prob.base_demand) + for t in 1:T + # ξ_t sits in the last slot of stage block t. + ξt = w_cpu[t * nu] + # Row-wise multiplicative demand realization. + demand[t, :] .= prob.base_demand[t, :] .* ξt + end + # Push the realized demand into the ExaModels parameter (device-safe: + # ExaModels.set_parameter! copyto!s a CPU vector into the GPU θ view). + set_demand!(prob, demand) + end + + # ── 2. Strict mode: cascade-clamp targets and pin the reservoir path ───── + if prob.p_reservoir !== nothing + T = prob.horizon + init = Float64.(vec(Array(init_state))) + inflow = Float64.(vec(Array(w_flat))) + xhat = Float64.(vec(Array(xhat_flat))) + if !isempty(prob.cascade) + K = prob.K + for t in 1:T + x_prev = t == 1 ? init : view(xhat, (t-2)*nH+1:(t-1)*nH) + targets_t = view(xhat, (t-1)*nH+1:t*nH) + # PHYSICAL inflow slice: first nH entries of stage block t + # (stride nu skips the ξ_t slot when demand noise is on). + inflow_t = view(inflow, (t-1)*nu+1:(t-1)*nu+nH) + for conn in prob.cascade + u, d = conn.upstream, conn.downstream + R_u = K * inflow_t[u] + x_prev[u] - targets_t[u] + if conn.turn_only + max_contrib = min(Float64(conn.K_max_turn), max(0.0, R_u)) + else + max_contrib = max(0.0, R_u) + end + true_upper = x_prev[d] + K * inflow_t[d] - K * prob.min_turn[d] + max_contrib + true_upper = min(prob.max_vol[d], true_upper) + targets_t[d] = min(targets_t[d], true_upper) + end + end + end + # Strict-mode invariant: p_reservoir[1:nH] must equal x0, because the + # builders omit the (parameter-only) initial-condition constraint in + # strict mode — see file-top comment. vcat(init, xhat) guarantees it. + reservoir_vals = vcat(init, xhat) + copyto!(prob.strict_reservoir_values, reservoir_vals) + ExaModels.set_parameter!(prob.core, prob.p_reservoir, reservoir_vals) + end + return nothing +end + +function target_multipliers(prob::HydroExaDEProblem, result) + λ = result.multipliers[prob.target_con_range] + prob.strict_targets || return λ + + nH = prob.nHyd + T = prob.horizon + raw = vec(Array(λ)) + out = copy(raw) + if T > 1 + out[1:(T - 1) * nH] .-= raw[(nH + 1):(T * nH)] + end + return out +end + # ── Post-processing ─────────────────────────────────────────────────────────── """ @@ -903,12 +1325,20 @@ function hydro_solution(prob::HydroExaDEProblem, result) pg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG pf_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB - res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + if prob.strict_targets + res_sol = reshape(copy(prob.strict_reservoir_values), nH, T+1) + else + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + end out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - delta_sol = dp_sol .- dn_sol + if prob.strict_targets + delta_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + delta_sol = dp_sol .- dn_sol + end return (va=va_sol, pg=pg_sol, pf=pf_sol, deficit=def_sol, reservoir=res_sol, outflow=out_sol, spill=spill_sol, delta=delta_sol) else # :ac_polar @@ -921,13 +1351,33 @@ function hydro_solution(prob::HydroExaDEProblem, result) p_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR q_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB - def_q_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB - res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + # Reactive slack layout depends on the builder's reactive_deficit_mode: + # :free — one free deficit_q block + # :penalized — δq⁺ then δq⁻ blocks; deficit_q = δq⁺ − δq⁻ + # :hard — no slack variable; report exact zeros + if prob.reactive_deficit_mode === :penalized + dqp_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + dqn_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + def_q_sol = dqp_sol .- dqn_sol + elseif prob.reactive_deficit_mode === :hard + def_q_sol = zeros(eltype(sol), nB, T) + else + def_q_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + end + if prob.strict_targets + res_sol = reshape(copy(prob.strict_reservoir_values), nH, T+1) + else + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + end out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH - delta_sol = dp_sol .- dn_sol + if prob.strict_targets + delta_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + delta_sol = dp_sol .- dn_sol + end return (va=va_sol, vm=vm_sol, pg=pg_sol, qg=qg_sol, p_fr=p_fr_sol, q_fr=q_fr_sol, p_to=p_to_sol, q_to=q_to_sol, deficit=def_sol, deficit_q=def_q_sol, diff --git a/examples/HydroPowerModels/hydro_power_exa_embedded.jl b/examples/HydroPowerModels/hydro_power_exa_embedded.jl new file mode 100644 index 0000000..c57cb5b --- /dev/null +++ b/examples/HydroPowerModels/hydro_power_exa_embedded.jl @@ -0,0 +1,1271 @@ +# hydro_power_exa_embedded.jl +# +# Embedded-NN deterministic equivalent for the Hydro power system. +# Target constraints from hydro_power_exa.jl are replaced by a +# VectorNonlinearOracle that evaluates the Flux policy inline. +# +# Slack target constraint (matching regular DE sign convention): +# π_θ(inflow_t, reservoir_t) − reservoir_{t+1,r} − δ⁺_{t,r} + δ⁻_{t,r} = 0 +# +# Strict target constraint: +# π_θ(inflow_t, reservoir_t) − reservoir_{t+1,r} = 0 +# +# Recurrent encoder handling: the per-stage encoder outputs h_t are cached in +# h_cache, computed by threading the recurrent state across stages +# (DecisionRulesExa._step_encoder — DecisionRules.jl semantics; Flux ≥ 0.16 +# cells are stateless so the Chain cannot be called directly per stage). Since +# the encoder reads only inflows, h_t depends on w_{1..t} and never on +# reservoir states; the cache is invalidated when inflows or policy parameters +# change (invalidate_policy_cache!/set_inflows!) and the oracle's analytic +# per-stage Jacobian ∂π_t/∂x_t through the combiner remains the FULL +# derivative — state threading adds no ∂h_t/∂x dependence. +# +# Depends on: hydro_power_data.jl, hydro_power_exa.jl, hydro_reachable_policy.jl + +using Flux +using LinearAlgebra: I +import DecisionRulesExa: set_x0!, set_uncertainty!, set_targets!, invalidate_policy_cache! + +# Keep legacy `include("hydro_power_exa_embedded.jl")` scripts working while +# letting training entrypoints include the policy explicitly. +if !isdefined(@__MODULE__, :HydroReachablePolicy) + include(joinpath(@__DIR__, "hydro_reachable_policy.jl")) +end + +# ── Problem struct ────────────────────────────────────────────────────────────── + +struct EmbeddedHydroExaDEProblem{P, VT <: AbstractVector{Float64}} + core + model + p_demand + p_reactive_demand + p_x0 + p_inflow + p_penalty_half + base_penalty_half::Float64 + p_penalty_l1 + base_penalty_l1::Float64 + policy::P + nHyd::Int + nx::Int + nw::Int + nBus::Int + nGen::Int + nBranch::Int + horizon::Int + formulation::Symbol + target_con_range::UnitRange{Int} + _res_start::Int + _dp_start::Int + _dn_start::Int + _nvar::Int + _inflow_buf::VT + _x0_buf::VT + _h_cache_dirty::Ref{Bool} + strict_targets::Bool +end + +# ── Interface (duck-typing for train_tsddr_embedded) ──────────────────────────── + +""" + set_x0!(prob::EmbeddedHydroExaDEProblem, x0) + +Set the initial reservoir state for the embedded hydro DE. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `x0::AbstractVector`: initial reservoir volumes, one value per hydro unit. + +# Returns +- `prob`. + +# Notes +The oracle closure reads `_x0_buf` directly when evaluating the first-stage +policy input, so this method updates both the ExaModels parameter and the +oracle-side buffer. +""" +function set_x0!(prob::EmbeddedHydroExaDEProblem, x0::AbstractVector) + length(x0) == prob.nHyd || error("x0 length must be nHyd=$(prob.nHyd)") + ExaModels.set_parameter!(prob.core, prob.p_x0, x0) + copyto!(prob._x0_buf, Float64.(x0)) + return prob +end + +""" + set_inflows!(prob::EmbeddedHydroExaDEProblem, w) + +Set the full inflow trajectory parameter. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `w::AbstractVector`: stage-major vector with length `prob.horizon * prob.nHyd`. + +# Returns +- `prob`. + +# Notes +Updating inflows invalidates cached recurrent encoder states because each stage +encoder input has changed. +""" +function set_inflows!(prob::EmbeddedHydroExaDEProblem, w::AbstractVector) + expected = prob.horizon * prob.nHyd + length(w) == expected || error("w must have length T*nHyd=$expected") + ExaModels.set_parameter!(prob.core, prob.p_inflow, w) + copyto!(prob._inflow_buf, Float64.(w)) + prob._h_cache_dirty[] = true + return prob +end + +""" + set_uncertainty!(prob::EmbeddedHydroExaDEProblem, w) + +Set the uncertainty trajectory for generic embedded training code. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `w::AbstractVector`: stage-major inflow trajectory. + +# Returns +- `prob`. + +# Notes +In the hydro example, uncertainty is exactly the inflow trajectory, so this +delegates to [`set_inflows!`](@ref). +""" +function set_uncertainty!(prob::EmbeddedHydroExaDEProblem, w::AbstractVector) + set_inflows!(prob, w) +end + +""" + set_targets!(prob::EmbeddedHydroExaDEProblem, target) + +Ignore external targets for embedded hydro DEs. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: ignored. +- `target::AbstractVector`: ignored. + +# Returns +- `nothing`. + +# Notes +Embedded hydro DEs generate targets inside the NLP through the policy oracle, so +there is no target parameter to update. +""" +function set_targets!(::EmbeddedHydroExaDEProblem, ::AbstractVector) + return nothing +end + +""" + invalidate_policy_cache!(prob::EmbeddedHydroExaDEProblem) + +Mark cached recurrent encoder states dirty after policy parameters change. The +oracle recomputes those states lazily on the next function/Jacobian evaluation. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. + +# Returns +- `prob`. +""" +function invalidate_policy_cache!(prob::EmbeddedHydroExaDEProblem) + prob._h_cache_dirty[] = true + return prob +end + +""" + set_demand!(prob::EmbeddedHydroExaDEProblem, demand_matrix) + +Update active power demand parameters. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `demand_matrix::AbstractMatrix`: `T x nBus` matrix of active demand values. + +# Returns +- `prob`. + +# Notes +The matrix is flattened in stage-major order to match the ExaModels parameter +layout. +""" +function set_demand!(prob::EmbeddedHydroExaDEProblem, demand_matrix::AbstractMatrix) + T, nB = size(demand_matrix) + T == prob.horizon || error("demand_matrix must have T=$(prob.horizon) rows") + nB == prob.nBus || error("demand_matrix must have nBus=$(prob.nBus) cols") + flat = [demand_matrix[t, b] for t in 1:T for b in 1:nB] + ExaModels.set_parameter!(prob.core, prob.p_demand, flat) + return prob +end + +""" + embedded_hydro_realized_states(prob, result) + +Extract realized reservoir states from an embedded hydro solve. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `result`: MadNLP result with a flat primal solution. + +# Returns +- A flat vector containing stages `1:T`, excluding the initial state. +""" +function embedded_hydro_realized_states(prob::EmbeddedHydroExaDEProblem, result) + T = prob.horizon + nH = prob.nHyd + sol = result.solution + return sol[prob._res_start + nH : prob._res_start + (T + 1) * nH - 1] +end + +""" + hydro_solution(prob::EmbeddedHydroExaDEProblem, result) -> NamedTuple + +Reshape the flat NLP solution into named physical blocks. This mirrors the +regular hydro DE post-processing helper so rollout diagnostics can read +reservoirs, generation, deficits, outflows, spills, and target slacks by name. + +# Arguments +- `prob::EmbeddedHydroExaDEProblem`: embedded hydro deterministic equivalent. +- `result`: MadNLP result with a flat primal solution. + +# Returns +- A `NamedTuple` of solution arrays. + +# Notes +For strict targets, slack arrays are returned as zeros because no slack +variables are present in the strict NLP. +""" +function hydro_solution(prob::EmbeddedHydroExaDEProblem, result) + T = prob.horizon + nH = prob.nHyd + nG = prob.nGen + nBR = prob.nBranch + nB = prob.nBus + sol = result.solution + off = 0 + + if prob.formulation === :dc + va_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + pg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG + pf_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + if prob.strict_targets + dp_sol = zeros(eltype(sol), nH, T) + dn_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + end + delta_sol = dp_sol .- dn_sol + return (va=va_sol, pg=pg_sol, pf=pf_sol, deficit=def_sol, + reservoir=res_sol, outflow=out_sol, spill=spill_sol, delta=delta_sol) + else + va_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + vm_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + pg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG + qg_sol = reshape(sol[off .+ (1:T*nG)], nG, T); off += T*nG + p_fr_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + q_fr_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + p_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + q_to_sol = reshape(sol[off .+ (1:T*nBR)], nBR, T); off += T*nBR + def_sol = reshape(sol[off .+ (1:T*nB)], nB, T); off += T*nB + # Canonical ACP keeps reactive nodal balance hard: no reactive slack + # variables are present in the decision vector. + def_q_sol = zeros(eltype(sol), nB, T) + res_sol = reshape(sol[off .+ (1:(T+1)*nH)], nH, T+1); off += (T+1)*nH + out_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + spill_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + if prob.strict_targets + dp_sol = zeros(eltype(sol), nH, T) + dn_sol = zeros(eltype(sol), nH, T) + else + dp_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + dn_sol = reshape(sol[off .+ (1:T*nH)], nH, T); off += T*nH + end + delta_sol = dp_sol .- dn_sol + return (va=va_sol, vm=vm_sol, pg=pg_sol, qg=qg_sol, + p_fr=p_fr_sol, q_fr=q_fr_sol, p_to=p_to_sol, q_to=q_to_sol, + deficit=def_sol, deficit_q=def_q_sol, + reservoir=res_sol, outflow=out_sol, spill=spill_sol, delta=delta_sol) + end +end + +# ── Oracle builder helper ─────────────────────────────────────────────────────── + +function _build_hydro_oracle(policy, T, nHyd, res_start, dp_start, dn_start, + nvar_total, inflow_buf, x0_buf; + strict_targets::Bool = true) + + _res(s) = res_start + (s-1)*nHyd : res_start + s*nHyd - 1 + _dp(t) = dp_start + (t-1)*nHyd : dp_start + t*nHyd - 1 + _dn(t) = dn_start + (t-1)*nHyd : dn_start + t*nHyd - 1 + _inflow(t) = (t-1)*nHyd+1 : t*nHyd + _crow(t) = (t-1)*nHyd+1 : t*nHyd + + encoder = policy.encoder + combiner = policy.combiner + combiner isa Flux.Dense || + throw(ArgumentError("embedded hydro DE currently requires the default single Dense reachable-policy head; use combiner_layers=Int[] or regular strict DE for multilayer state heads")) + n_h = size(combiner.weight, 2) - policy.n_state + reachable_policy = policy isa HydroReachablePolicy + has_spill_cap = reachable_policy && policy.spill_max !== nothing + + jac_r = Int[] + jac_c = Int[] + for t in 1:T, r in 1:nHyd + row = (t-1)*nHyd + r + push!(jac_r, row); push!(jac_c, res_start + t*nHyd + r - 1) + if !strict_targets + push!(jac_r, row); push!(jac_c, dp_start + (t-1)*nHyd + r - 1) + push!(jac_r, row); push!(jac_c, dn_start + (t-1)*nHyd + r - 1) + end + if t > 1 + for j in 1:nHyd + push!(jac_r, row); push!(jac_c, res_start + (t-1)*nHyd + j - 1) + end + end + end + nnzj = length(jac_r) + + const_jac_cpu = zeros(Float64, nnzj) + nn_jac_ranges_flat = Vector{UnitRange{Int}}(undef, T * nHyd) + k = 0 + for t in 1:T, r in 1:nHyd + const_jac_cpu[k+1] = -1.0 + k += 1 + if !strict_targets + const_jac_cpu[k+1] = -1.0 + const_jac_cpu[k+2] = 1.0 + k += 2 + end + if t > 1 + nn_jac_ranges_flat[(t-1)*nHyd + r] = (k+1):(k+nHyd) + k += nHyd + end + end + const_jac_dev = similar(inflow_buf, Float64, nnzj) + copyto!(const_jac_dev, const_jac_cpu) + + W_state = @view combiner.weight[:, n_h+1:end] + act = combiner.σ + has_output_bounds = getfield(policy, :output_lower) !== nothing + output_lower_f32 = similar(x0_buf, Float32, nHyd) + output_scale_f32 = similar(x0_buf, Float32, nHyd) + if has_output_bounds + output_lower_f32 .= getfield(policy, :output_lower) + output_scale_f32 .= getfield(policy, :output_scale) + else + fill!(output_lower_f32, 0f0) + fill!(output_scale_f32, 1f0) + end + min_vol_f32 = similar(x0_buf, Float32, nHyd) + max_vol_f32 = similar(x0_buf, Float32, nHyd) + min_turn_f32 = similar(x0_buf, Float32, nHyd) + max_turn_f32 = similar(x0_buf, Float32, nHyd) + spill_max_f32 = similar(x0_buf, Float32, nHyd) + upstream_max_f32 = similar(x0_buf, Float32, nHyd) + if reachable_policy + min_vol_f32 .= policy.min_vol + max_vol_f32 .= policy.max_vol + min_turn_f32 .= policy.min_turn + max_turn_f32 .= policy.max_turn + upstream_max_f32 .= policy.upstream_max_inflow + if policy.spill_max !== nothing + spill_max_f32 .= policy.spill_max + else + fill!(spill_max_f32, Float32(Inf)) + end + else + fill!(upstream_max_f32, 0f0) + end + + function _act_deriv!(σ_prime, output) + if act === NNlib.sigmoid || act === NNlib.sigmoid_fast + σ_prime .= output .* (one(eltype(output)) .- output) + elseif act === Base.tanh || act === NNlib.tanh_fast + σ_prime .= one(eltype(output)) .- output .* output + elseif act === identity + fill!(σ_prime, one(eltype(σ_prime))) + else + σ_prime .= output .* (one(eltype(output)) .- output) + end + return σ_prime + end + + function _activation!(dest, z) + if act === NNlib.sigmoid || act === NNlib.sigmoid_fast + dest .= inv.(one(eltype(dest)) .+ exp.(-z)) + elseif act === Base.tanh || act === NNlib.tanh_fast + dest .= tanh.(z) + elseif act === identity + dest .= z + else + dest .= act.(z) + end + return dest + end + + # Pre-allocated buffers — all on the same device as x0_buf + x0_f32 = similar(x0_buf, Float32, nHyd) + x_prev_f32 = similar(x0_buf, Float32, nHyd) + infl_f32 = similar(x0_buf, Float32, nHyd) + _comb_in = similar(x0_buf, Float32, n_h + nHyd) + z_buf = similar(x0_buf, Float32, nHyd) + nn_out_f32 = similar(x0_buf, Float32, nHyd) + y_norm_f32 = similar(x0_buf, Float32, nHyd) + lower_f32 = similar(x0_buf, Float32, nHyd) + upper_f32 = similar(x0_buf, Float32, nHyd) + scale_f32 = similar(x0_buf, Float32, nHyd) + dlower_dx_f32 = similar(x0_buf, Float32, nHyd) + dupper_dx_f32 = similar(x0_buf, Float32, nHyd) + dbound_dx_f32 = similar(x0_buf, Float32, nHyd) + σ_prime_buf = similar(x0_buf, Float32, nHyd) + λ_f32_buf = similar(x0_buf, Float32, nHyd) + d_xprev_buf = similar(x0_buf, Float32, nHyd) + J_buf = similar(x0_buf, Float64, nHyd, nHyd) + dbound_dx_f64 = similar(x0_buf, Float64, nHyd) + diag_mask = similar(x0_buf, Float64, nHyd, nHyd) + copyto!(diag_mask, Matrix{Float64}(I, nHyd, nHyd)) + h_cache = similar(x0_buf, Float32, n_h, T) + h_cache_dirty = Ref(true) + + function _reachable_bounds!(inflow, x_prev) + K32 = Float32(policy.K) + upper_f32 .= x_prev .+ K32 .* inflow .- K32 .* min_turn_f32 .+ upstream_max_f32 + dupper_dx_f32 .= ifelse.(upper_f32 .<= max_vol_f32, 1f0, 0f0) + upper_f32 .= min.(max_vol_f32, upper_f32) + + if has_spill_cap + lower_f32 .= x_prev .+ K32 .* inflow .- K32 .* max_turn_f32 .- spill_max_f32 + dlower_dx_f32 .= ifelse.(lower_f32 .>= min_vol_f32, 1f0, 0f0) + lower_f32 .= max.(min_vol_f32, lower_f32) + else + lower_f32 .= min_vol_f32 + fill!(dlower_dx_f32, 0f0) + end + + dupper_dx_f32 .= ifelse.(upper_f32 .< lower_f32, dlower_dx_f32, dupper_dx_f32) + upper_f32 .= max.(upper_f32, lower_f32) + scale_f32 .= upper_f32 .- lower_f32 + return nothing + end + + function _combiner_fwd!(nn_out, h, x_prev) + _comb_in[1:n_h] .= h + _comb_in[n_h+1:end] .= x_prev + mul!(z_buf, combiner.weight, _comb_in) + z_buf .+= combiner.bias + _activation!(y_norm_f32, z_buf) + if reachable_policy + _reachable_bounds!(infl_f32, x_prev) + nn_out .= lower_f32 .+ scale_f32 .* y_norm_f32 + else + nn_out .= output_lower_f32 .+ output_scale_f32 .* y_norm_f32 + end + return nn_out + end + + function _combiner_jac!(J_buf, σ_prime_buf, h, x_prev) + _comb_in[1:n_h] .= h + _comb_in[n_h+1:end] .= x_prev + mul!(z_buf, combiner.weight, _comb_in) + z_buf .+= combiner.bias + _activation!(y_norm_f32, z_buf) + _act_deriv!(σ_prime_buf, y_norm_f32) + if reachable_policy + _reachable_bounds!(infl_f32, x_prev) + σ_prime_buf .*= scale_f32 + else + σ_prime_buf .*= output_scale_f32 + end + J_buf .= reshape(σ_prime_buf, :, 1) .* W_state + if reachable_policy + dbound_dx_f32 .= dlower_dx_f32 .+ y_norm_f32 .* (dupper_dx_f32 .- dlower_dx_f32) + dbound_dx_f64 .= dbound_dx_f32 + J_buf .+= reshape(dbound_dx_f64, :, 1) .* diag_mask + end + return nothing + end + + function _populate_h_cache!() + # Thread the recurrent encoder state across stages explicitly. Flux + # ≥ 0.16 cells are stateless — calling the Chain of LSTM wrappers + # directly (`encoder(x)`) would restart from `initialstates` on every + # stage, making the encoder memoryless. `_step_encoder` advances the + # underlying cells one stage at a time from a fresh initial state, + # exactly like DecisionRules.jl's `_step_encoder` and the threaded + # policy forward pass. h_t therefore depends on the inflow history + # w_{1..t} only (never on reservoir states), which is what makes this + # per-stage cache — and the oracle's per-stage direct Jacobian — + # exact for the embedded policy. + enc_state = DecisionRulesExa._init_recurrent_state(encoder) + for t in 1:T + infl_f32 .= view(inflow_buf, _inflow(t)) + h, enc_state = DecisionRulesExa._step_encoder(encoder, infl_f32, enc_state) + view(h_cache, :, t) .= h + end + h_cache_dirty[] = false + return nothing + end + + function _ensure_h_cache!() + h_cache_dirty[] && _populate_h_cache!() + return nothing + end + + function oracle_f!(c, xv) + _ensure_h_cache!() + x0_f32 .= view(x0_buf, 1:nHyd) + for t in 1:T + infl_f32 .= view(inflow_buf, _inflow(t)) + if t == 1 + x_prev_f32 .= x0_f32 + else + x_prev_f32 .= view(xv, _res(t)) + end + h = view(h_cache, :, t) + _combiner_fwd!(nn_out_f32, h, x_prev_f32) + + c_t = view(c, _crow(t)) + r_v = view(xv, _res(t+1)) + if strict_targets + c_t .= nn_out_f32 .- r_v + else + dp_v = view(xv, _dp(t)) + dn_v = view(xv, _dn(t)) + c_t .= nn_out_f32 .- r_v .- dp_v .+ dn_v + end + end + return nothing + end + + function oracle_jac!(vals, xv) + _ensure_h_cache!() + copyto!(vals, const_jac_dev) + for t in 2:T + infl_f32 .= view(inflow_buf, _inflow(t)) + x_prev_f32 .= view(xv, _res(t)) + h = view(h_cache, :, t) + _combiner_jac!(J_buf, σ_prime_buf, h, x_prev_f32) + for r in 1:nHyd + vals[nn_jac_ranges_flat[(t-1)*nHyd + r]] .= view(J_buf, r, :) + end + end + return nothing + end + + function oracle_vjp!(Jtv, xv, λ) + _ensure_h_cache!() + fill!(Jtv, 0.0) + for t in 1:T + λ_f64 = view(λ, _crow(t)) + view(Jtv, _res(t+1)) .-= λ_f64 + if !strict_targets + view(Jtv, _dp(t)) .-= λ_f64 + view(Jtv, _dn(t)) .+= λ_f64 + end + if t > 1 + h = view(h_cache, :, t) + infl_f32 .= view(inflow_buf, _inflow(t)) + x_prev_f32 .= view(xv, _res(t)) + _comb_in[1:n_h] .= h + _comb_in[n_h+1:end] .= x_prev_f32 + mul!(z_buf, combiner.weight, _comb_in) + z_buf .+= combiner.bias + _activation!(y_norm_f32, z_buf) + _act_deriv!(σ_prime_buf, y_norm_f32) + if reachable_policy + _reachable_bounds!(infl_f32, x_prev_f32) + σ_prime_buf .*= scale_f32 + dbound_dx_f32 .= dlower_dx_f32 .+ + y_norm_f32 .* (dupper_dx_f32 .- dlower_dx_f32) + else + σ_prime_buf .*= output_scale_f32 + fill!(dbound_dx_f32, 0f0) + end + λ_f32_buf .= λ_f64 + σ_prime_buf .*= λ_f32_buf + mul!(d_xprev_buf, W_state', σ_prime_buf) + if reachable_policy + d_xprev_buf .+= dbound_dx_f32 .* λ_f32_buf + end + view(Jtv, _res(t)) .+= d_xprev_buf + end + end + return nothing + end + + oracle = ExaModels.VectorNonlinearOracle( + nvar = nvar_total, + ncon = T * nHyd, + nnzj = nnzj, + jac_rows = jac_r, + jac_cols = jac_c, + lcon = zeros(T * nHyd), + ucon = zeros(T * nHyd), + f! = oracle_f!, + jac! = oracle_jac!, + vjp! = oracle_vjp!, + ) + return oracle, h_cache_dirty +end + +# ── DC builder ────────────────────────────────────────────────────────────────── + +function _build_embedded_dc_hydro_de( + policy, + power_data::PowerData, + hydro_data::HydroData, + T::Int; + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64, + target_penalty::Union{Real,Symbol} = :auto, + target_penalty_l1::Union{Real,Symbol,Nothing} = :auto, + demand_matrix = nothing, + deficit_cost::Union{Nothing,Real} = nothing, + load_scaler::Real = 0.6, + strict_targets::Bool = true, +) + nBus = power_data.nBus + nGen = power_data.nGen + nBranch = power_data.nBranch + nHyd = hydro_data.nHyd + K = float_type(hydro_data.K) + ρ = float_type(target_penalty === :auto ? auto_target_penalty(power_data, hydro_data) : target_penalty) + ρ_l1 = if target_penalty_l1 === :auto + ρ + elseif target_penalty_l1 === nothing + zero(float_type) + else + float_type(target_penalty_l1) + end + use_l1 = ρ_l1 > 0 + baseMVA = float_type(power_data.baseMVA) + cd = float_type(deficit_cost !== nothing ? deficit_cost : + power_data.cost_deficit * power_data.baseMVA) + + core = ExaModels.ExaCore(float_type; backend = backend) + + # ── Variables (track offsets) ───────────────────────────────────────────── + var_offset = 0 + + va = ExaModels.variable(core, T * nBus) + var_offset += T * nBus + + pg_lb = float_type.(repeat([g.pmin for g in power_data.gens], T)) + pg_ub = float_type.(repeat([g.pmax for g in power_data.gens], T)) + pg = ExaModels.variable(core, T * nGen; lvar = pg_lb, uvar = pg_ub) + var_offset += T * nGen + + pf_lb = float_type.(repeat([-b.rate_a for b in power_data.branches], T)) + pf_ub = float_type.(repeat([ b.rate_a for b in power_data.branches], T)) + pf = ExaModels.variable(core, T * nBranch; lvar = pf_lb, uvar = pf_ub) + var_offset += T * nBranch + + deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + var_offset += T * nBus + + res_start = var_offset + 1 + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + var_offset += (T+1) * nHyd + + out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) + out_ub = float_type.(repeat([h.max_turn for h in hydro_data.units], T)) + outflow = ExaModels.variable(core, T * nHyd; lvar = out_lb, uvar = out_ub) + var_offset += T * nHyd + + spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += T * nHyd + + dp_start = strict_targets ? 0 : var_offset + 1 + delta_pos = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + dn_start = strict_targets ? 0 : var_offset + 1 + delta_neg = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + nvar_total = var_offset + + # ── Parameters ──────────────────────────────────────────────────────────── + init_demand = if demand_matrix !== nothing + float_type.(load_scaler .* [demand_matrix[t, b] for t in 1:T for b in 1:nBus]) + else + float_type.(load_scaler .* repeat(power_data.default_bus_demand, T)) + end + + p_demand = ExaModels.parameter(core, init_demand) + p_x0 = ExaModels.parameter(core, zeros(float_type, nHyd)) + p_inflow = ExaModels.parameter(core, zeros(float_type, T * nHyd)) + p_penalty_half = ExaModels.parameter(core, fill(float_type(ρ / 2), T * nHyd)) + p_penalty_l1 = ExaModels.parameter(core, fill(ρ_l1, T * nHyd)) + + # ── Objective ───────────────────────────────────────────────────────────── + gen_cost_items = [(t = t, g = g_pos, + c1 = float_type(g.cost1), c2 = float_type(g.cost2)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.objective(core, + item.c2 * pg[_gi(nGen, item.t, item.g)]^2 + + item.c1 * pg[_gi(nGen, item.t, item.g)] + for item in gen_cost_items + ) + + def_cost_items = [(t = t, b = b, c = cd) for t in 1:T for b in 1:nBus] + ExaModels.objective(core, + item.c * deficit[_bi(nBus, item.t, item.b)] + for item in def_cost_items + ) + + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] + if !strict_targets + ExaModels.objective(core, + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 + for item in delta_items + ) + end + + if !strict_targets && use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end + + # ── Constraints 1–7 ────────────────────────────────────────────────────── + n_con = 0 + + nRef = length(power_data.ref_buses) + ref_items = [(t = t, ref = ref) for t in 1:T for ref in power_data.ref_buses] + ExaModels.constraint(core, + va[_bi(nBus, item.t, item.ref)] + for item in ref_items + ) + n_con += T * nRef + + ohm_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + b = float_type(br.b)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + item.b * (va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + - pf[_bri(nBranch, item.t, item.br)] + for item in ohm_items + ) + n_con += T * nBranch + + ang_lb = float_type.(repeat([br.angmin for br in power_data.branches], T)) + ang_ub = float_type.(repeat([br.angmax for br in power_data.branches], T)) + ang_items = [(t = t, f = br.f_bus, tb = br.t_bus) + for t in 1:T for br in power_data.branches] + ExaModels.constraint(core, + va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)] + for item in ang_items; + lcon = ang_lb, ucon = ang_ub, + ) + n_con += T * nBranch + + kcl_init_items = [(t = t, b = b) for t in 1:T for b in 1:nBus] + c_kcl = ExaModels.constraint(core, + p_demand[_bi(nBus, item.t, item.b)] + for item in kcl_init_items + ) + kcl_gen_items = [(t = t, brow = _bi(nBus, t, g.bus), gcol = _gi(nGen, t, g_pos)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.constraint!(core, c_kcl, + item.brow => -pg[item.gcol] + for item in kcl_gen_items + ) + kcl_fr_items = [(t = t, brow = _bi(nBus, t, br.f_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl, + item.brow => pf[item.bcol] + for item in kcl_fr_items + ) + kcl_to_items = [(t = t, brow = _bi(nBus, t, br.t_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl, + item.brow => -pf[item.bcol] + for item in kcl_to_items + ) + kcl_def_items = [(t = t, brow = _bi(nBus, t, b), dcol = _bi(nBus, t, b)) + for t in 1:T for b in 1:nBus] + ExaModels.constraint!(core, c_kcl, + item.brow => -deficit[item.dcol] + for item in kcl_def_items + ) + n_con += T * nBus + + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + + wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), + out_idx = _ri(nHyd, t, r), spill_idx = _ri(nHyd, t, r), + inflow_p = _ri(nHyd, t, r), K = K) + for t in 1:T for r in 1:nHyd] + c_wb = ExaModels.constraint(core, + reservoir[item.res_next] - reservoir[item.res_curr] + + item.K * outflow[item.out_idx] + + spill[item.spill_idx] + - item.K * p_inflow[item.inflow_p] + for item in wb_items + ) + if !isempty(hydro_data.upstream_turns) + wb_turn_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos), K = K) + for t in 1:T for conn in hydro_data.upstream_turns] + ExaModels.constraint!(core, c_wb, + item.row => -item.K * outflow[item.col] + for item in wb_turn_items + ) + end + if !isempty(hydro_data.upstream_spills) + wb_spill_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos)) + for t in 1:T for conn in hydro_data.upstream_spills] + ExaModels.constraint!(core, c_wb, + item.row => -spill[item.col] + for item in wb_spill_items + ) + end + n_con += T * nHyd + + tc_items = [(gen_col = _gi(nGen, t, h.gen_pos), out_col = _ri(nHyd, t, h.pos), + baseMVA = baseMVA, pf_h = float_type(h.pf)) + for t in 1:T for h in hydro_data.units] + ExaModels.constraint(core, + item.baseMVA * pg[item.gen_col] - item.pf_h * outflow[item.out_col] + for item in tc_items + ) + n_con += T * nHyd + + # ── Oracle (replaces target constraints) ────────────────────────────────── + inflow_buf = backend === nothing ? + zeros(Float64, T * nHyd) : + KernelAbstractions.zeros(backend, Float64, T * nHyd) + x0_buf = backend === nothing ? + zeros(Float64, nHyd) : + KernelAbstractions.zeros(backend, Float64, nHyd) + + oracle, h_cache_dirty = _build_hydro_oracle(policy, T, nHyd, res_start, dp_start, dn_start, + nvar_total, inflow_buf, x0_buf; + strict_targets = strict_targets) + ExaModels.constraint(core, oracle) + target_con_range = (n_con + 1):(n_con + T * nHyd) + + model = ExaModels.ExaModel(core) + + return EmbeddedHydroExaDEProblem( + core, model, + p_demand, nothing, p_x0, p_inflow, + p_penalty_half, Float64(ρ / 2), + p_penalty_l1, Float64(ρ_l1), + policy, nHyd, nHyd, nHyd, + nBus, nGen, nBranch, T, + :dc, target_con_range, + res_start, dp_start, dn_start, nvar_total, + inflow_buf, x0_buf, h_cache_dirty, + strict_targets, + ) +end + +# ── AC polar builder ──────────────────────────────────────────────────────────── + +function _build_embedded_ac_hydro_de( + policy, + power_data::PowerData, + hydro_data::HydroData, + T::Int; + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64, + target_penalty::Union{Real,Symbol} = :auto, + target_penalty_l1::Union{Real,Symbol,Nothing} = :auto, + demand_matrix = nothing, + reactive_demand_matrix = nothing, + deficit_cost::Union{Nothing,Real} = nothing, + load_scaler::Real = 0.6, + strict_targets::Bool = true, +) + nBus = power_data.nBus + nGen = power_data.nGen + nBranch = power_data.nBranch + nHyd = hydro_data.nHyd + K = float_type(hydro_data.K) + ρ = float_type(target_penalty === :auto ? auto_target_penalty(power_data, hydro_data) : target_penalty) + ρ_l1 = if target_penalty_l1 === :auto + ρ + elseif target_penalty_l1 === nothing + zero(float_type) + else + float_type(target_penalty_l1) + end + use_l1 = ρ_l1 > 0 + baseMVA = float_type(power_data.baseMVA) + cd = float_type(deficit_cost !== nothing ? deficit_cost : + power_data.cost_deficit * power_data.baseMVA) + + core = ExaModels.ExaCore(float_type; backend = backend) + + # ── Variables ───────────────────────────────────────────────────────────── + var_offset = 0 + + va = ExaModels.variable(core, T * nBus) + var_offset += T * nBus + + vm_lb = float_type.(repeat([b.vmin for b in power_data.buses], T)) + vm_ub = float_type.(repeat([b.vmax for b in power_data.buses], T)) + vm = ExaModels.variable(core, T * nBus; lvar = vm_lb, uvar = vm_ub, + start = ones(float_type, T * nBus)) + var_offset += T * nBus + + pg_lb = float_type.(repeat([g.pmin for g in power_data.gens], T)) + pg_ub = float_type.(repeat([g.pmax for g in power_data.gens], T)) + pg = ExaModels.variable(core, T * nGen; lvar = pg_lb, uvar = pg_ub) + var_offset += T * nGen + + qg_lb = float_type.(repeat([isfinite(g.qmin) ? g.qmin : -1e4 for g in power_data.gens], T)) + qg_ub = float_type.(repeat([isfinite(g.qmax) ? g.qmax : 1e4 for g in power_data.gens], T)) + qg = ExaModels.variable(core, T * nGen; lvar = qg_lb, uvar = qg_ub) + var_offset += T * nGen + + p_fr_lb = float_type.(repeat([-b.rate_a for b in power_data.branches], T)) + p_fr_ub = float_type.(repeat([ b.rate_a for b in power_data.branches], T)) + p_fr = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + q_fr = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + p_to = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + q_to = ExaModels.variable(core, T * nBranch; lvar = p_fr_lb, uvar = p_fr_ub) + var_offset += 4 * T * nBranch + + deficit = ExaModels.variable(core, T * nBus; lvar = float_type(0)) + var_offset += T * nBus + + # No reactive deficit variable: MAIN-faithful hard reactive nodal balance. + + res_start = var_offset + 1 + res_lb = float_type.(repeat([h.min_vol for h in hydro_data.units], T+1)) + res_ub = float_type.(repeat([h.max_vol for h in hydro_data.units], T+1)) + reservoir = ExaModels.variable(core, (T+1) * nHyd; lvar = res_lb, uvar = res_ub) + var_offset += (T+1) * nHyd + + out_lb = float_type.(repeat([h.min_turn for h in hydro_data.units], T)) + out_ub = float_type.(repeat([h.max_turn for h in hydro_data.units], T)) + outflow = ExaModels.variable(core, T * nHyd; lvar = out_lb, uvar = out_ub) + var_offset += T * nHyd + + spill = ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += T * nHyd + + dp_start = strict_targets ? 0 : var_offset + 1 + delta_pos = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + dn_start = strict_targets ? 0 : var_offset + 1 + delta_neg = strict_targets ? nothing : + ExaModels.variable(core, T * nHyd; lvar = float_type(0)) + var_offset += strict_targets ? 0 : T * nHyd + + nvar_total = var_offset + + # ── Parameters ──────────────────────────────────────────────────────────── + init_demand = if demand_matrix !== nothing + float_type.(load_scaler .* [demand_matrix[t, b] for t in 1:T for b in 1:nBus]) + else + float_type.(load_scaler .* repeat(power_data.default_bus_demand, T)) + end + init_reactive_demand = if reactive_demand_matrix !== nothing + float_type.(load_scaler .* [reactive_demand_matrix[t, b] for t in 1:T for b in 1:nBus]) + else + float_type.(load_scaler .* repeat(power_data.default_bus_reactive_demand, T)) + end + + p_demand = ExaModels.parameter(core, init_demand) + p_reactive_demand = ExaModels.parameter(core, init_reactive_demand) + p_x0 = ExaModels.parameter(core, zeros(float_type, nHyd)) + p_inflow = ExaModels.parameter(core, zeros(float_type, T * nHyd)) + p_penalty_half = ExaModels.parameter(core, fill(float_type(ρ / 2), T * nHyd)) + p_penalty_l1 = ExaModels.parameter(core, fill(ρ_l1, T * nHyd)) + + br_ac = [_ac_branch_coeffs(br, float_type) for br in power_data.branches] + + # ── Objective ───────────────────────────────────────────────────────────── + gen_cost_items = [(t = t, g = g_pos, + c1 = float_type(g.cost1), c2 = float_type(g.cost2)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.objective(core, + item.c2 * pg[_gi(nGen, item.t, item.g)]^2 + + item.c1 * pg[_gi(nGen, item.t, item.g)] + for item in gen_cost_items + ) + + def_cost_items = [(t = t, b = b, c = cd) for t in 1:T for b in 1:nBus] + ExaModels.objective(core, + item.c * deficit[_bi(nBus, item.t, item.b)] + for item in def_cost_items + ) + + delta_items = [(idx = _ri(nHyd, t, r),) for t in 1:T for r in 1:nHyd] + if !strict_targets + ExaModels.objective(core, + p_penalty_half[item.idx] * (delta_pos[item.idx] - delta_neg[item.idx])^2 + for item in delta_items + ) + end + if !strict_targets && use_l1 + ExaModels.objective(core, + p_penalty_l1[item.idx] * (delta_pos[item.idx] + delta_neg[item.idx]) + for item in delta_items + ) + end + + # ── Constraints ─────────────────────────────────────────────────────────── + n_con = 0 + + # 1. Reference angle + nRef = length(power_data.ref_buses) + ref_items = [(t = t, ref = ref) for t in 1:T for ref in power_data.ref_buses] + ExaModels.constraint(core, va[_bi(nBus, item.t, item.ref)] for item in ref_items) + n_con += T * nRef + + # 2. AC from-end active power flow + pfr_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c3 = br_ac[br_pos].c3, c4 = br_ac[br_pos].c4, c5 = br_ac[br_pos].c5) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + p_fr[_bri(nBranch, item.t, item.br)] + - item.c5 * vm[_bi(nBus, item.t, item.f)]^2 + - item.c3 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * cos(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + - item.c4 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * sin(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + for item in pfr_items + ) + n_con += T * nBranch + + # 3. AC from-end reactive power flow + qfr_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c3 = br_ac[br_pos].c3, c4 = br_ac[br_pos].c4, c6 = br_ac[br_pos].c6) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + q_fr[_bri(nBranch, item.t, item.br)] + + item.c6 * vm[_bi(nBus, item.t, item.f)]^2 + + item.c4 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * cos(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + - item.c3 * vm[_bi(nBus, item.t, item.f)] * vm[_bi(nBus, item.t, item.tb)] + * sin(va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)]) + for item in qfr_items + ) + n_con += T * nBranch + + # 4. AC to-end active power flow + pto_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c1 = br_ac[br_pos].c1, c2 = br_ac[br_pos].c2, c7 = br_ac[br_pos].c7) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + p_to[_bri(nBranch, item.t, item.br)] + - item.c7 * vm[_bi(nBus, item.t, item.tb)]^2 + - item.c1 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * cos(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + - item.c2 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * sin(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + for item in pto_items + ) + n_con += T * nBranch + + # 5. AC to-end reactive power flow + qto_items = [(t = t, f = br.f_bus, tb = br.t_bus, br = br_pos, + c1 = br_ac[br_pos].c1, c2 = br_ac[br_pos].c2, c8 = br_ac[br_pos].c8) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint(core, + q_to[_bri(nBranch, item.t, item.br)] + + item.c8 * vm[_bi(nBus, item.t, item.tb)]^2 + + item.c2 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * cos(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + - item.c1 * vm[_bi(nBus, item.t, item.tb)] * vm[_bi(nBus, item.t, item.f)] + * sin(va[_bi(nBus, item.t, item.tb)] - va[_bi(nBus, item.t, item.f)]) + for item in qto_items + ) + n_con += T * nBranch + + # 6. Phase angle difference + ang_lb = float_type.(repeat([br.angmin for br in power_data.branches], T)) + ang_ub = float_type.(repeat([br.angmax for br in power_data.branches], T)) + ang_items = [(t = t, f = br.f_bus, tb = br.t_bus) for t in 1:T for br in power_data.branches] + ExaModels.constraint(core, + va[_bi(nBus, item.t, item.f)] - va[_bi(nBus, item.t, item.tb)] + for item in ang_items; + lcon = ang_lb, ucon = ang_ub, + ) + n_con += T * nBranch + + # Match PowerModels ACPPowerModel: apparent-power limits at both branch + # ends, not merely independent bounds on p and q. + thermal_ub = float_type.(repeat([br.rate_a^2 for br in power_data.branches], T)) + thermal_items = [(t = t, br = br_pos) + for t in 1:T for br_pos in 1:nBranch] + ExaModels.constraint(core, + p_fr[_bri(nBranch, item.t, item.br)]^2 + + q_fr[_bri(nBranch, item.t, item.br)]^2 + for item in thermal_items; + lcon = fill(float_type(-Inf), T * nBranch), ucon = thermal_ub, + ) + ExaModels.constraint(core, + p_to[_bri(nBranch, item.t, item.br)]^2 + + q_to[_bri(nBranch, item.t, item.br)]^2 + for item in thermal_items; + lcon = fill(float_type(-Inf), T * nBranch), ucon = thermal_ub, + ) + n_con += 2 * T * nBranch + + # 7. Active KCL + kcl_p_init = [(t = t, b = b, gs = float_type(power_data.buses[b].gs)) + for t in 1:T for b in 1:nBus] + c_kcl_p = ExaModels.constraint(core, + p_demand[_bi(nBus, item.t, item.b)] + + item.gs * vm[_bi(nBus, item.t, item.b)]^2 + for item in kcl_p_init + ) + kcl_pg_items = [(t = t, brow = _bi(nBus, t, g.bus), gcol = _gi(nGen, t, g_pos)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.constraint!(core, c_kcl_p, + item.brow => -pg[item.gcol] for item in kcl_pg_items) + kcl_pfr_items = [(t = t, brow = _bi(nBus, t, br.f_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_p, + item.brow => p_fr[item.bcol] for item in kcl_pfr_items) + kcl_pto_items = [(t = t, brow = _bi(nBus, t, br.t_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_p, + item.brow => p_to[item.bcol] for item in kcl_pto_items) + kcl_def_items = [(t = t, brow = _bi(nBus, t, b), dcol = _bi(nBus, t, b)) + for t in 1:T for b in 1:nBus] + ExaModels.constraint!(core, c_kcl_p, + item.brow => -deficit[item.dcol] for item in kcl_def_items) + n_con += T * nBus + + # 8. Reactive KCL + kcl_q_init = [(t = t, b = b, bs = float_type(power_data.buses[b].bs)) + for t in 1:T for b in 1:nBus] + c_kcl_q = ExaModels.constraint(core, + p_reactive_demand[_bi(nBus, item.t, item.b)] + - item.bs * vm[_bi(nBus, item.t, item.b)]^2 + for item in kcl_q_init + ) + kcl_qg_items = [(t = t, brow = _bi(nBus, t, g.bus), gcol = _gi(nGen, t, g_pos)) + for t in 1:T for (g_pos, g) in enumerate(power_data.gens)] + ExaModels.constraint!(core, c_kcl_q, + item.brow => -qg[item.gcol] for item in kcl_qg_items) + kcl_qfr_items = [(t = t, brow = _bi(nBus, t, br.f_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_q, + item.brow => q_fr[item.bcol] for item in kcl_qfr_items) + kcl_qto_items = [(t = t, brow = _bi(nBus, t, br.t_bus), bcol = _bri(nBranch, t, br_pos)) + for t in 1:T for (br_pos, br) in enumerate(power_data.branches)] + ExaModels.constraint!(core, c_kcl_q, + item.brow => q_to[item.bcol] for item in kcl_qto_items) + n_con += T * nBus + + # 9. Initial reservoir + ic_items = [(r = r,) for r in 1:nHyd] + ExaModels.constraint(core, + reservoir[_ri(nHyd, 1, item.r)] - p_x0[item.r] + for item in ic_items + ) + n_con += nHyd + + # 10. Water balance + wb_items = [(res_next = _ri(nHyd, t+1, r), res_curr = _ri(nHyd, t, r), + out_idx = _ri(nHyd, t, r), spill_idx = _ri(nHyd, t, r), + inflow_p = _ri(nHyd, t, r), K = K) + for t in 1:T for r in 1:nHyd] + c_wb = ExaModels.constraint(core, + reservoir[item.res_next] - reservoir[item.res_curr] + + item.K * outflow[item.out_idx] + + spill[item.spill_idx] + - item.K * p_inflow[item.inflow_p] + for item in wb_items + ) + if !isempty(hydro_data.upstream_turns) + wb_turn_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos), K = K) + for t in 1:T for conn in hydro_data.upstream_turns] + ExaModels.constraint!(core, c_wb, + item.row => -item.K * outflow[item.col] for item in wb_turn_items) + end + if !isempty(hydro_data.upstream_spills) + wb_spill_items = [(row = _ri(nHyd, t, conn.downstream_pos), + col = _ri(nHyd, t, conn.upstream_pos)) + for t in 1:T for conn in hydro_data.upstream_spills] + ExaModels.constraint!(core, c_wb, + item.row => -spill[item.col] for item in wb_spill_items) + end + n_con += T * nHyd + + # 11. Turbine coupling + tc_items = [(gen_col = _gi(nGen, t, h.gen_pos), out_col = _ri(nHyd, t, h.pos), + baseMVA = baseMVA, pf_h = float_type(h.pf)) + for t in 1:T for h in hydro_data.units] + ExaModels.constraint(core, + item.baseMVA * pg[item.gen_col] - item.pf_h * outflow[item.out_col] + for item in tc_items + ) + n_con += T * nHyd + + # ── Oracle ──────────────────────────────────────────────────────────────── + inflow_buf = backend === nothing ? + zeros(Float64, T * nHyd) : + KernelAbstractions.zeros(backend, Float64, T * nHyd) + x0_buf = backend === nothing ? + zeros(Float64, nHyd) : + KernelAbstractions.zeros(backend, Float64, nHyd) + + oracle, h_cache_dirty = _build_hydro_oracle(policy, T, nHyd, res_start, dp_start, dn_start, + nvar_total, inflow_buf, x0_buf; + strict_targets = strict_targets) + ExaModels.constraint(core, oracle) + target_con_range = (n_con + 1):(n_con + T * nHyd) + + model = ExaModels.ExaModel(core) + + return EmbeddedHydroExaDEProblem( + core, model, + p_demand, p_reactive_demand, p_x0, p_inflow, + p_penalty_half, Float64(ρ / 2), + p_penalty_l1, Float64(ρ_l1), + policy, nHyd, nHyd, nHyd, + nBus, nGen, nBranch, T, + :ac_polar, target_con_range, + res_start, dp_start, dn_start, nvar_total, + inflow_buf, x0_buf, h_cache_dirty, + strict_targets, + ) +end + +# ── Dispatcher ────────────────────────────────────────────────────────────────── + +function build_embedded_hydro_de( + policy, + power_data::PowerData, + hydro_data::HydroData, + T::Int; + formulation::Symbol = :dc, + kwargs..., +) + formulation in (:dc, :ac_polar) || + error("formulation must be :dc or :ac_polar, got :$formulation") + + if formulation === :dc + return _build_embedded_dc_hydro_de(policy, power_data, hydro_data, T; kwargs...) + else + return _build_embedded_ac_hydro_de(policy, power_data, hydro_data, T; kwargs...) + end +end diff --git a/examples/HydroPowerModels/hydro_reachable_policy.jl b/examples/HydroPowerModels/hydro_reachable_policy.jl new file mode 100644 index 0000000..79a1cc6 --- /dev/null +++ b/examples/HydroPowerModels/hydro_reachable_policy.jl @@ -0,0 +1,576 @@ +# hydro_reachable_policy.jl +# +# Hydro-specific feasible target policy for strict regular and embedded DEs. +# This file owns the policy architecture and reachability/cascade bounds; the +# ExaModels problem builders live in hydro_power_exa.jl and +# hydro_power_exa_embedded.jl. + +using Flux +using Zygote +import DecisionRulesExa: load_stateconditioned_policy! + +""" + stretchedsigmoid(x) -> y ∈ [0, 1 − 1e-3] + +Boundary-attaining sigmoid: `clamp((sigmoid(x) − 0.03) / 0.94, 0, 1 − 1e-3)`. + +Plain `sigmoid` reaches 0/1 only at ±∞ with vanishing gradient, so a policy +squashed by it can never exactly attain the reachable-interval boundaries — +where optimal hydro decisions frequently live (SDDP places ~24% of its realized +states within 0.1% of a boundary on the paired protocol). The gentle 6.4% +stretch keeps the interior mapping close to the `sigmoid` shape (warm starts +from sigmoid-trained checkpoints shift by ≤ ~3% absolute) while attaining +exactly 0 for `sigmoid(x) ≤ 0.03` (x ≈ −3.5) and the δ-interior upper value +1 − 1e-3 for `sigmoid(x) ≥ 0.969`. +""" +function stretchedsigmoid(x::Real) + T = float(typeof(x)) + # Gentle stretch: σ ∈ [0.03, 0.97] maps affinely onto the full range, so + # warm starts from sigmoid-trained weights shift interior outputs by ≤ ~3% + # absolute (a 20% stretch amplified near-boundary decisions by 8-10% and + # destroyed warm-started policies), while corners are attained at finite + # pre-activation |x| ≈ 3.5. + # δ-interior upper clamp: exact y = 1 (store-everything when the raw + # reachable upper binds) forces turbine = min_turn AND spill = 0 exactly — + # a measure-zero feasible set that interior-point solvers cannot converge + # into when strict mode imposes the target as an equality (observed as + # MAXIMUM_ITERATIONS / spurious INFEASIBLE). y = 0 keeps a strict interior + # (spill is unbounded above in the stage NLP), so the lower corner is exact. + return clamp((NNlib.sigmoid(x) - T(0.03)) / T(0.94), zero(T), one(T) - T(1e-3)) +end + +""" + hardsigmoidsafe(x) -> y ∈ [0, 1 − 1e-3] + +`hardsigmoid` with the same δ-interior upper clamp as [`stretchedsigmoid`](@ref): +attains exactly 0 on the lower side and 1 − 1e-3 on the upper side, keeping the +strict stage NLP interior-point-solvable at store-max targets. +""" +function hardsigmoidsafe(x::Real) + T = float(typeof(x)) + return min(NNlib.hardsigmoid(x), one(T) - T(1e-3)) +end + +# Activations admissible for the target head: range within [0, 1] so the +# affine map into the reachable interval stays feasible. +const _BOUNDED_ACTIVATIONS = + (sigmoid, NNlib.sigmoid, NNlib.sigmoid_fast, NNlib.hardsigmoid, NNlib.hardσ, + hardsigmoidsafe, stretchedsigmoid) + +# Evaluation-time snap-to-boundary (diagnostic; leave at 0 during training): +# normalized targets y ≤ ε are snapped to exactly 0 and y ≥ 1−ε to exactly 1, +# emulating a boundary-attaining head on a sigmoid-trained checkpoint with no +# retraining. Set via DR_SNAP_EPS in the eval scripts. +const TARGET_SNAP_EPS = Ref(0.0f0) + +""" + hydro_reachable_policy(hydro_data, layers; activation=sigmoid, encoder_type=Flux.LSTM, + spill_max=nothing, combiner_layers=Int[], n_context=0, + n_extra_uncertainty=0) + +Build a state-conditioned hydro policy whose outputs are one-stage reachable +reservoir targets. + +The recurrent encoder reads the per-stage uncertainty block (inflows, plus any +extra uncertainty entries such as the stochastic-demand factor ξ_t). The +nonrecurrent combiner reads `[encoded_uncertainty; reservoir_state]`, emits +normalized targets in `[0, 1]`, and the wrapper maps them into the +reachability interval implied by the current inflow and previous reservoir +state (reachability uses ONLY the physical inflow slice — extra uncertainty +entries never change feasibility logic). + +# Arguments +- `hydro_data::HydroData`: hydro limits, initial volumes, stage duration, and + cascade metadata. +- `layers::AbstractVector{Int}`: recurrent encoder widths over inflows. + +# Keywords +- `activation`: activation used by the target head. It must be sigmoid-style so + normalized targets remain in `[0, 1]`. +- `encoder_type`: Flux recurrent layer constructor, usually `Flux.LSTM`. +- `spill_max`: optional finite spill cap used to tighten lower bounds. +- `combiner_layers`: hidden widths in the nonrecurrent state-conditioned head. +- `n_context`: fixed per-stage context rows PREPENDED before the uncertainty. +- `n_extra_uncertainty`: extra stochastic per-stage entries APPENDED after the + inflow inside the uncertainty block (demand noise ⇒ 1: the block is + `[w_t; ξ_t]`). Widens the encoder input by the same amount; checkpoints are + therefore only compatible across runs with the same value. + +# Returns +- `HydroReachablePolicy`: a Flux-compatible policy with trainable encoder and + combiner parameters plus fixed hydro reachability metadata. + +# Notes +For strict regular deterministic equivalents, rolling this policy from the true +initial state and feeding each previous target into the next policy call gives a +feasible target path by induction. Embedded strict DEs use realized reservoir +states inside the NLP. Cascade links are handled by clamping downstream targets +to the upper bound implied by same-stage upstream targets. + +The recurrent encoder state is threaded across stages explicitly (Flux ≥ 0.16 +cells are stateless, so calling the `LSTM` wrapper directly would restart from +`initialstates` every stage): the forward pass advances `state` by one cell +step per call, mirroring DecisionRules.jl's `HydroReachablePolicy` exactly. +Call `Flux.reset!(policy)` at scenario boundaries to restore the initial state. +""" +mutable struct HydroReachablePolicy{E,C,RS,V,S,I} + encoder::E + combiner::C + state::RS # Encoder recurrent state, threaded across stages + n_context::Int # Number of context dimensions prepended before inflow + n_uncertainty::Int + n_state::Int + min_vol::V + max_vol::V + min_turn::V + max_turn::V + spill_max::S + upstream_max_inflow::V + K::Float64 + output_lower::Nothing + output_scale::Nothing + cascade::Vector{CascadeLink} + cascade_upstream::I + cascade_downstream::I + cascade_turn_only::V + cascade_k_max_turn::V + cascade_reservoir_ids::I +end + +Flux.@layer HydroReachablePolicy trainable=(encoder, combiner) + +""" + _hydro_adapt_bound(x, ref) + +Return `x` as a vector with the same array family and element type as `ref`. + +# Arguments +- `x::AbstractVector`: hydro metadata stored on the policy. +- `ref::AbstractVector`: vector whose device family and element type should be + matched. + +# Returns +- A vector with `length(x)` whose storage is compatible with `ref`. + +# Notes +This keeps metadata such as `min_vol` and `max_vol` on the same device as the +policy forward pass: CPU inputs stay on CPU, GPU inputs stay on GPU. +""" +function _hydro_adapt_bound(x::AbstractVector, ref::AbstractArray) + # `ref` may be a vector (one scenario) or a matrix (batched); either way we + # return a per-reservoir vector on ref's device/eltype. `similar(ref, n)` + # yields a length-n vector even when ref is a matrix, which then broadcasts + # column-wise against batched (nx × N) quantities. + typeof(x) === typeof(ref) && return x + y = similar(ref, length(x)) + copyto!(y, convert.(eltype(ref), x)) + return y +end + +""" + _hydro_adapt_index(x, ref) + +Return integer indices stored on the same device family as `ref`. + +# Arguments +- `x::AbstractVector`: integer indices stored in ordinary Julia memory. +- `ref::AbstractVector`: vector whose device family should be matched. + +# Returns +- An integer vector compatible with `ref`. + +# Notes +This helper is used before GPU gathers such as `inflow[upstream]`; copying to +`similar(ref, Int, ...)` avoids host-side scalar indexing during policy +evaluation. +""" +function _hydro_adapt_index(x::AbstractVector, ref::AbstractVector) + y = similar(ref, Int, length(x)) + copyto!(y, x) + return y +end + +""" + _hydro_reachable_bounds(policy, inflow, x_prev, ref) -> (lower, upper) + +Compute one-stage reservoir target bounds for the hydro water balance. + +# Arguments +- `policy::HydroReachablePolicy`: policy carrying fixed hydro metadata. +- `inflow`: current-stage inflow vector. +- `x_prev`: previous reservoir-state vector. +- `ref`: vector used to choose element type and device family. + +# Returns +- `(lower, upper)`: vectors defining the closed interval into which normalized + policy outputs are mapped. + +# Notes +For each reservoir, the simplified balance is + +```text +x_next = x_prev + K * inflow - K * turbine_out - spill + upstream_contrib. +``` + +The upper bound uses the smallest required turbine outflow (`min_turn`) and zero +spill. The lower bound is the physical minimum volume unless `spill_max` is +finite; with finite spill, the lowest reachable storage uses `max_turn` and +maximum spill. + +Proof sketch: every feasible turbine/spill choice satisfies the same balance +equation and variable bounds. Substituting extremal admissible outflow/spill +values gives an interval containing all one-stage reachable reservoir states. +Mapping a sigmoid output into this interval therefore produces a reachable +target in the relaxed one-stage balance. +""" +function _hydro_reachable_bounds(policy::HydroReachablePolicy, inflow, x_prev, ref) + min_vol = _hydro_adapt_bound(policy.min_vol, ref) + max_vol = _hydro_adapt_bound(policy.max_vol, ref) + min_turn = _hydro_adapt_bound(policy.min_turn, ref) + max_turn = _hydro_adapt_bound(policy.max_turn, ref) + upstream = _hydro_adapt_bound(policy.upstream_max_inflow, ref) + K = convert(eltype(ref), policy.K) + + upper_raw = x_prev .+ K .* inflow .- K .* min_turn .+ upstream + upper = min.(max_vol, upper_raw) + + lower = if policy.spill_max === nothing + min_vol + else + spill_max = _hydro_adapt_bound(policy.spill_max, ref) + lower_raw = x_prev .+ K .* inflow .- K .* max_turn .- spill_max + max.(min_vol, lower_raw) + end + + upper = max.(upper, lower) + return lower, upper +end +Zygote.@nograd _hydro_reachable_bounds + +""" + _cascade_upper_bounds(policy, target, inflow, x_prev) -> upper + +Tighten downstream reservoir upper bounds using same-stage upstream targets. + +# Arguments +- `policy::HydroReachablePolicy`: policy carrying cascade metadata. +- `target`: raw target vector before cascade clamping. +- `inflow`: current-stage inflow vector. +- `x_prev`: previous reservoir-state vector. + +# Returns +- A vector of per-reservoir upper bounds induced by incoming cascade links. + +# Notes +For a cascade link `u -> d`, the upstream target determines the upstream +release implied by the balance: + +```text +release_u = K * inflow_u + x_prev_u - target_u. +``` + +Proof sketch: downstream storage is increasing in upstream contribution. The +largest physically consistent contribution from an upstream target is exactly +the positive release implied by that target, possibly turbine-capped. Therefore +clamping the downstream target to this bound cannot remove any feasible target +that respects the upstream target and cascade balance. + +# Assumptions +- **Single-level cascades.** The implied release + `release_u = K * inflow_u + x_prev_u - target_u` omits the upstream unit's + own incoming cascade contribution: if `u` itself receives water from a unit + further upstream, its true release can be larger than computed here. The + omission is conservative — it can only under-estimate the available + downstream contribution, so the clamp is never infeasible, merely possibly + over-tight for multi-level chains. Bolivia's three links + (COR→SIS turbine-only, ZON→CHU turbine+spill, TAQ1→TAQ2 turbine+spill) are + all single-level, so the bound is exact for that case. +- **No gradient through the clamp.** `Zygote.@nograd` on this function means + that when the cascade clamp binds, the true dependence of the downstream + target on the upstream target carries no gradient — a deliberate + approximation that keeps the policy pullback cheap and well-defined. +- **Physically infeasible edge case.** If the cascade upper bound falls below + the reachable lower bound of `_hydro_reachable_bounds`, the clamped target + can fall below `lower`. There is no policy-level remedy: the stage is + genuinely infeasible, and downstream slack/deficit handling must absorb it. +""" +function _cascade_upper_bounds(policy::HydroReachablePolicy, target, inflow, x_prev) + cascade = policy.cascade + T = eltype(target) + K = T(policy.K) + n = length(target) + isempty(cascade) && return _hydro_adapt_bound(fill(T(Inf), n), target) + + upstream = _hydro_adapt_index(policy.cascade_upstream, target) + downstream = _hydro_adapt_index(policy.cascade_downstream, target) + reservoir_ids = _hydro_adapt_index(policy.cascade_reservoir_ids, target) + turn_only = _hydro_adapt_bound(policy.cascade_turn_only, target) + k_max_turn = _hydro_adapt_bound(policy.cascade_k_max_turn, target) + min_turn = _hydro_adapt_bound(policy.min_turn, target) + max_vol = _hydro_adapt_bound(policy.max_vol, target) + + # Release implied by asking the upstream reservoir to end at `target`. + release = K .* inflow[upstream] .+ x_prev[upstream] .- target[upstream] + positive_release = max.(zero(T), release) + max_contrib = ifelse.(turn_only .> zero(T), min.(k_max_turn, positive_release), positive_release) + + # One upper bound per cascade link, expressed for the downstream reservoir. + link_upper = x_prev[downstream] .+ + K .* inflow[downstream] .- + K .* min_turn[downstream] .+ + max_contrib + link_upper = min.(max_vol[downstream], link_upper) + + # Reduce link-wise bounds to one bound per reservoir. Reservoirs without + # incoming cascade links receive Inf and are left unchanged by `min`. + link_by_reservoir = ifelse.( + reshape(downstream, :, 1) .== reshape(reservoir_ids, 1, :), + reshape(link_upper, :, 1), + T(Inf), + ) + return vec(minimum(link_by_reservoir; dims = 1)) +end +Zygote.@nograd _cascade_upper_bounds + +""" + (policy::HydroReachablePolicy)(input) -> target + +Evaluate the reachable hydro policy. + +# Arguments +- `input`: concatenated vector `[context_t; u_t; x_{t-1}]`, where the + uncertainty block `u_t` is `[inflow_t]` (historical) or `[inflow_t; ξ_t]` + (stochastic demand, `n_uncertainty = n_state + 1`). + +# Returns +- A reservoir target vector in the one-stage reachable set. + +# Notes +The encoder reads `[context_t; u_t]`, threading its recurrent state across +calls (one cell step per stage, stored in `policy.state` — DecisionRules.jl +semantics). Reachability bounds and cascade clamps slice out only the true +inflow entries (the first `n_state` entries of `u_t`), so neither prepended +context nor appended extra uncertainty changes physical feasibility logic. +The combiner reads both encoded inflow/context and previous reservoir state, +emits normalized targets, and those targets are mapped into the reachability +interval before cascade clamping. + +The cascade clamp inherits the assumptions documented on +[`_cascade_upper_bounds`](@ref): + +- Cascades are treated as single-level — the upstream release omits that unit's + own incoming cascade contribution, which is conservative (never infeasible, + possibly over-tight) on multi-level chains. Bolivia's three links + (COR→SIS turn-only, ZON→CHU turn+spill, TAQ1→TAQ2 turn+spill) are all + single-level. +- Both `_hydro_reachable_bounds` and `_cascade_upper_bounds` are + `Zygote.@nograd`, so when the cascade clamp binds, the downstream target's + true dependence on the upstream target carries no gradient (deliberate + approximation). +- In the physically infeasible edge case where the cascade upper bound is + below the reachable lower bound, the returned target can fall below `lower`; + the stage is genuinely infeasible and no policy-level remedy exists. +""" +# Row-slice that works for a vector input (one scenario) OR a matrix input +# (features × batch, one column per scenario). Plain `input[r]` on a matrix does +# LINEAR indexing and silently corrupts batched inputs, so dispatch on ndims. +# The vector method returns exactly `input[r]`, so single-scenario behavior (the +# strict TS-DDR trainer, evaluation) is byte-identical. +_row_slice(input::AbstractVector, r) = input[r] +_row_slice(input::AbstractMatrix, r) = input[r, :] + +function (m::HydroReachablePolicy)(input) + # Split input: optional context first, then the per-stage uncertainty block + # (physical inflow, optionally followed by extra uncertainty entries such + # as the demand factor ξ_t), then the previous state. Physical reachability + # bounds must use ONLY the true inflow slice (first n_state uncertainty + # entries); the encoder reads the FULL uncertainty block. With + # n_uncertainty == n_state (no demand noise) `unc === inflow` and the + # behavior is bit-identical to the historical implementation. + c_end = m.n_context + w_start = c_end + 1 + w_end = c_end + m.n_uncertainty + # Full uncertainty block [w_t; extras] — encoder input. + unc = _row_slice(input, w_start:w_end) + # Physical inflow slice (per-reservoir) — reachability/cascade input. + inflow = m.n_uncertainty == m.n_state ? unc : + _row_slice(input, w_start:(c_end + m.n_state)) + x_prev = _row_slice(input, w_end+1:size(input, 1)) + encoder_input = if c_end == 0 + unc + else + vcat(_row_slice(input, 1:c_end), unc) + end + + # Encode inflow through the recurrent encoder, threading state across calls + # (mirrors DecisionRules.jl: encoded, s_t = _step_encoder(enc, T.(w_t), s_{t-1})). + # Cast to encoder precision for type stability (avoids Zygote codegen bugs). + T = DecisionRulesExa._state_eltype(m.state) + h, new_state = DecisionRulesExa._step_encoder(m.encoder, T.(encoder_input), m.state) + # Thread the recurrent state to the next call. + m.state = new_state + + y = m.combiner(vcat(h, x_prev)) + # Diagnostic eval-time snap (no-op at ε = 0): near-boundary normalized + # targets become exact boundary points, which plain sigmoid cannot attain. + ε = DecisionRulesExa._state_eltype(m.state)(TARGET_SNAP_EPS[]) + if ε > 0 + # Upper snap goes to 1 − 1e-3, not 1: see stretchedsigmoid — exact + # store-max targets are IPM-degenerate under strict equality. + y = ifelse.(y .<= ε, zero(ε), ifelse.(y .>= one(ε) - ε, one(ε) - oftype(ε, 1e-3), y)) + end + lower, upper = _hydro_reachable_bounds(m, inflow, x_prev, y) + # Because `y` is sigmoid-bounded, this affine map stays in [lower, upper]. + raw_target = lower .+ (upper .- lower) .* y + # The single-level cascade clamp uses per-reservoir index gathering (vector + # op) and is `@nograd`. For batched (matrix) input it is skipped: this is + # exact when imitating FEASIBLE targets (SDDP decisions respect the cascade + # balance, so `raw_target <= cascade_upper` and the clamp is a no-op). Single- + # scenario (vector) input keeps the full clamp — strict trainer/eval unchanged. + if !isempty(m.cascade) && ndims(input) == 1 + cascade_upper = _cascade_upper_bounds(m, raw_target, inflow, x_prev) + return min.(raw_target, cascade_upper) + end + return raw_target +end + +""" + Flux.reset!(policy::HydroReachablePolicy) -> Nothing + +Reset the inflow encoder's recurrent state to `Flux.initialstates`, e.g. at +scenario boundaries. + +# Notes +The combiner is feed-forward and does not carry recurrent state. The recurrent +state is re-derived from the (possibly device-moved) encoder weights on every +reset, so the state always matches the encoder's device and element type. +""" +function Flux.reset!(m::HydroReachablePolicy) + # Reinitialize the recurrent state from the encoder's initial states. + m.state = DecisionRulesExa._init_recurrent_state(m.encoder) + return nothing +end + +""" + load_stateconditioned_policy!(policy::HydroReachablePolicy, state) + +Load Flux parameters into a reachable hydro policy. + +# Arguments +- `policy::HydroReachablePolicy`: target policy to update in place. +- `state`: checkpoint object accepted by `Flux.loadmodel!`. + +# Returns +- `policy`. + +# Notes +- If the checkpoint contains only `encoder` and `combiner` fields, those + trainable parts are loaded while hydro reachability metadata from the current + case is preserved. +- Checkpoints trained BEFORE recurrent-state threading (memoryless-encoder era) + have the SAME weight structure and load unchanged — only the runtime + semantics differ (the encoder now carries memory across stages). +- DecisionRules.jl (MAIN) checkpoints save the encoder as a `Chain` of BARE + `LSTMCell`s (layer state `(Wi, Wh, bias)` instead of `(cell = …,)`); those + load through the documented cell-by-cell fallback in + `DecisionRulesExa._load_encoder_state!`. +- The recurrent state is reset after loading so the next rollout starts from + `Flux.initialstates` of the loaded weights. +""" +function load_stateconditioned_policy!(policy::HydroReachablePolicy, state) + try + Flux.loadmodel!(policy, state) + Flux.reset!(policy) + return policy + catch err + hasproperty(state, :encoder) && hasproperty(state, :combiner) || rethrow(err) + @warn "Full HydroReachablePolicy checkpoint load failed; loading encoder/combiner only and keeping hydro reachability bounds" exception=(err, catch_backtrace()) + DecisionRulesExa._load_encoder_state!(policy.encoder, getproperty(state, :encoder)) + Flux.loadmodel!(policy.combiner, getproperty(state, :combiner)) + Flux.reset!(policy) + return policy + end +end + +function hydro_reachable_policy( + hydro_data::HydroData, + layers::AbstractVector{Int}; + activation = sigmoid, + encoder_type = Flux.LSTM, + spill_max = nothing, + combiner_layers = Int[], + n_context::Int = 0, + n_extra_uncertainty::Int = 0, +) + any(a -> activation === a, _BOUNDED_ACTIVATIONS) || + throw(ArgumentError("hydro_reachable_policy requires a [0,1]-bounded activation (sigmoid, hardsigmoid, or stretchedsigmoid) so normalized targets stay in [0, 1]")) + nHyd = hydro_data.nHyd + n_context >= 0 || throw(ArgumentError("n_context must be nonnegative")) + # Extra per-stage uncertainty entries appended after the inflow (e.g. the + # stochastic-demand factor ξ_t: n_extra_uncertainty = 1). They widen the + # encoder input; reachability/cascade bounds keep using only the inflow. + n_extra_uncertainty >= 0 || throw(ArgumentError("n_extra_uncertainty must be nonnegative")) + enc_sizes = vcat(nHyd + n_extra_uncertainty + n_context, layers) + enc_layers = [encoder_type(enc_sizes[i] => enc_sizes[i+1]) + for i in 1:length(layers)] + encoder = Flux.Chain(enc_layers...) + encoder_width = isempty(layers) ? nHyd + n_extra_uncertainty + n_context : layers[end] + combiner = DecisionRulesExa._dense_policy_head( + encoder_width + nHyd, + nHyd, + collect(Int, combiner_layers); + activation = activation, + ) + spill_vec = spill_max === nothing ? nothing : Float32.(collect(spill_max)) + if spill_vec !== nothing && length(spill_vec) != nHyd + throw(ArgumentError("spill_max length must be nHyd=$nHyd")) + end + + K = Float64(hydro_data.K) + upstream_max = zeros(Float32, nHyd) + for conn in hydro_data.upstream_turns + upstream_max[conn.downstream_pos] += Float32(K * hydro_data.units[conn.upstream_pos].max_turn) + end + + spill_dests = Dict{Int,Set{Int}}() + for conn in hydro_data.upstream_spills + push!(get!(spill_dests, conn.upstream_pos, Set{Int}()), conn.downstream_pos) + end + cascade = CascadeLink[] + for conn in hydro_data.upstream_turns + d, u = conn.downstream_pos, conn.upstream_pos + has_spill = haskey(spill_dests, u) && d in spill_dests[u] + push!(cascade, CascadeLink(d, u, !has_spill, Float32(K * hydro_data.units[u].max_turn))) + end + for conn in hydro_data.upstream_spills + d, u = conn.downstream_pos, conn.upstream_pos + already = any(c -> c.downstream == d && c.upstream == u, cascade) + already || push!(cascade, CascadeLink(d, u, false, Float32(K * hydro_data.units[u].max_turn))) + end + + return HydroReachablePolicy( + encoder, combiner, + DecisionRulesExa._init_recurrent_state(encoder), # initial recurrent state + # n_context; n_uncertainty = full per-stage uncertainty width (inflow + + # extras); n_state = nHyd (physical inflow/state width for bounds). + n_context, nHyd + n_extra_uncertainty, nHyd, + Float32.([h.min_vol for h in hydro_data.units]), + Float32.([h.max_vol for h in hydro_data.units]), + Float32.([h.min_turn for h in hydro_data.units]), + Float32.([h.max_turn for h in hydro_data.units]), + spill_vec, + upstream_max, + K, + nothing, nothing, + cascade, + # Typed comprehensions guarantee concrete Vector{Int}/Vector{Float32} + # element types: `getfield.(links, :field)` can infer as Vector{Real} + # (the field symbol does not always constant-propagate through fused + # broadcast), which breaks the struct's I/V type-parameter binding. + Int[c.upstream for c in cascade], + Int[c.downstream for c in cascade], + Float32[c.turn_only for c in cascade], + Float32[c.K_max_turn for c in cascade], + collect(1:nHyd), + ) +end diff --git a/examples/HydroPowerModels/hydro_training_utils.jl b/examples/HydroPowerModels/hydro_training_utils.jl new file mode 100644 index 0000000..b1f8ae6 --- /dev/null +++ b/examples/HydroPowerModels/hydro_training_utils.jl @@ -0,0 +1,50 @@ +# Shared helpers for HydroPowerModels training entrypoints. + +""" + parse_layers(s::AbstractString) -> Vector{Int} + +Parse a comma-separated hidden-layer specification. + +Empty or whitespace-only strings return `Int[]`, which lets environment +variables represent "no hidden layers" without a separate flag. + +# Arguments +- `s::AbstractString`: comma-separated layer widths, with optional whitespace. + +# Returns +- `Vector{Int}`: parsed hidden-layer widths. + +# Examples +```julia +parse_layers("128, 64") == [128, 64] +parse_layers("") == Int[] +parse_layers(" ") == Int[] +``` +""" +function parse_layers(s::AbstractString) + return isempty(strip(s)) ? + Int[] : + [parse(Int, strip(x)) for x in split(s, ",") if !isempty(strip(x))] +end + +function canonical_context_mode(raw_mode::AbstractString) + mode = lowercase(strip(raw_mode)) + mode in ("", "none", "off", "false") && return "" + mode in ("phase", "phase+progress") && return mode + throw(ArgumentError("DR_CONTEXT must be \"\", \"phase\", or \"phase+progress\"; got \"$raw_mode\"")) +end + +function build_stage_context(mode::AbstractString, horizon::Int, period::Int) + isempty(mode) && return nothing + include_progress = mode == "phase+progress" + return DecisionRulesExa.stage_phase_context( + horizon; + period = period, + include_progress = include_progress, + ) +end + +function context_run_tag(mode::AbstractString) + isempty(mode) && return "" + return "-ctx" * replace(mode, "+" => "p") +end diff --git a/examples/HydroPowerModels/train_hydro_exa.jl b/examples/HydroPowerModels/train_hydro_exa.jl index 6ec71e9..4da967a 100644 --- a/examples/HydroPowerModels/train_hydro_exa.jl +++ b/examples/HydroPowerModels/train_hydro_exa.jl @@ -17,6 +17,7 @@ using MadNLPGPU, KernelAbstractions, CUDA using CUDSS, CUDSS_jll, cuDNN const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) @@ -32,33 +33,50 @@ const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") const DEMAND_FILE = joinpath(CASE_DIR, "demand.csv") -const LAYERS = [128, 128] +const LAYERS = parse_layers(get(ENV, "DR_LAYERS", "128,128")) +const HEAD_LAYERS = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) const ACTIVATION = sigmoid const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) const NUM_EPOCHS = parse(Int, get(ENV, "DR_NUM_EPOCHS", "80")) const NUM_BATCHES = 100 const NUM_TRAIN_PER_BATCH = 1 const NUM_EVAL_SCENARIOS = 4 -const EVAL_EVERY = 25 +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) const LR = 1f-3 -const GRAD_CLIP = parse(Float32, get(ENV, "DR_GRAD_CLIP", "10")) +const GRAD_CLIP = parse(Float32, get(ENV, "DR_GRAD_CLIP", "0")) const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = parse(Float64, get(ENV, "DR_TARGET_PENALTY_MULT", "8.0")) const DEFICIT_COST = 1e5 const USE_GPU = true const load_scaler = 0.6 const NUM_WORKERS = 1 -const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "annealed") +const DISCOUNT_GAMMA = parse(Float64, get(ENV, "DR_DISCOUNT_GAMMA", "1.0")) +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) + +const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "const") +const _N_TOTAL = NUM_EPOCHS * NUM_BATCHES +const _ANNEAL_1_END = max(1, div(_N_TOTAL, 100)) +const _ANNEAL_2_END = max(_ANNEAL_1_END + 1, div(_N_TOTAL, 40)) +const _ANNEAL_3_END = max(_ANNEAL_2_END + 1, div(_N_TOTAL, 10)) const PENALTY_SCHEDULE = if _PENALTY_MODE == "annealed" [ - (1, div(NUM_EPOCHS * NUM_BATCHES, 4), 0.1), - (div(NUM_EPOCHS * NUM_BATCHES, 4) + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 2, 1.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 2 + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 3, 10.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 3 + 1, NUM_EPOCHS * NUM_BATCHES, 30.0), + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +elseif _PENALTY_MODE == "annealed_discount" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), ] else - [(1, NUM_EPOCHS * NUM_BATCHES, 1.0)] + [(1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT)] end # Optional: ramp num_train_per_batch and eval scenarios over training. @@ -66,11 +84,21 @@ end const NUM_TRAIN_SCHEDULE = nothing # e.g. [(1,500,1),(501,2000,4),(2001,4000,8)] const EVAL_SCHEDULE = nothing # e.g. [(1,2000,4),(2001,4000,32)] -const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = 9000) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) const _CLIP_TAG = GRAD_CLIP > 0 ? "-clip$(Int(GRAD_CLIP))" : "" -const _SCHED_TAG = _PENALTY_MODE == "annealed" ? "-anneal" : "-const" -const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-deteq-gpu$(_CLIP_TAG)$(_SCHED_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const _DISC_TAG = DISCOUNT_GAMMA < 1.0 ? "-disc$(replace(string(DISCOUNT_GAMMA), "." => ""))" : "" +const _SCHED_TAG = if _PENALTY_MODE == "annealed" + "-anneal" +elseif _PENALTY_MODE == "annealed_discount" + "-anndisc" +else + "-const" +end +const _LAYER_TAG = LAYERS == [128, 128] ? "" : "-L$(join(LAYERS, "_"))" +const _HEAD_TAG = isempty(HEAD_LAYERS) ? "" : "-H$(join(HEAD_LAYERS, "_"))" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-deteq-gpu$(_CLIP_TAG)$(_SCHED_TAG)$(_DISC_TAG)$(_LAYER_TAG)$(_HEAD_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") mkpath(MODEL_DIR) const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") @@ -85,9 +113,10 @@ power_data = load_power_data(PM_FILE) @info "Loading hydro data..." hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; - num_stages = NUM_STAGES * 10) + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) nHyd = hydro_data.nHyd T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES @info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios)" demand_mat = if isfile(DEMAND_FILE) @@ -103,6 +132,7 @@ resolved_pen = TARGET_PEN_ARG === :auto ? auto_target_penalty(power_data, hydro_data) : Float64(TARGET_PEN_ARG) @info "Auto target penalty: ρ=$(round(resolved_pen; digits=2))" +@info "Bolivia hydro target-penalty multiplier default: $(HYDRO_TARGET_PENALTY_MULT)" backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : (@info "Using CPU backend"; nothing) @@ -134,6 +164,8 @@ x0_init = Float32.([clamp(hydro_data.initial_volumes[r], hydro_data.units[r].min_vol, hydro_data.units[r].max_vol) for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) # ── Smoke test ──────────────────────────────────────────────────────────────── @@ -149,15 +181,29 @@ solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding resolved_pen_l1 = prob.base_penalty_l1 +const _discount_weights = Float64[DISCOUNT_GAMMA^(t-1) for t in 1:T for _ in 1:nHyd] +if DISCOUNT_GAMMA < 1.0 + @info "Discount γ=$(DISCOUNT_GAMMA): stage 1 weight=1.0, stage $T weight=$(round(DISCOUNT_GAMMA^(T-1); sigdigits=4))" +end + # ── Policy ──────────────────────────────────────────────────────────────────── -policy = StateConditionedPolicy(nHyd, nHyd, nHyd, LAYERS; - activation = ACTIVATION, - encoder_type = Flux.LSTM) +policy_active_mask = trues(nHyd) +policy = bounded_state_policy(nHyd, target_lower, target_upper, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + active_mask = policy_active_mask, + combiner_layers = HEAD_LAYERS) if !isnothing(PRE_TRAINED) @info "Loading pre-trained model from $(PRE_TRAINED)..." - Flux.loadmodel!(policy, JLD2.load(PRE_TRAINED, "model_state")) + load_stateconditioned_policy!(policy, JLD2.load(PRE_TRAINED, "model_state")) +end + +if USE_GPU + policy = CUDA.cu(policy) + x0_init = CUDA.cu(x0_init) + @info "Policy and x0 moved to GPU" end # ── W&B logging ─────────────────────────────────────────────────────────────── @@ -170,10 +216,13 @@ lg = WandbLogger( "case" => CASE_NAME, "formulation" => FORM_LABEL, "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, "layers" => LAYERS, + "head_layers" => HEAD_LAYERS, "activation" => string(ACTIVATION), "target_penalty" => "auto=$(round(resolved_pen; digits=2))", "target_penalty_l1" => "auto=$(round(resolved_pen_l1; digits=2))", + "hydro_target_penalty_mult" => HYDRO_TARGET_PENALTY_MULT, "deficit_cost" => DEFICIT_COST, "num_epochs" => NUM_EPOCHS, "num_batches" => NUM_BATCHES, @@ -185,6 +234,7 @@ lg = WandbLogger( "backend" => USE_GPU ? "GPU" : "CPU", "load_scaler" => load_scaler, "penalty_schedule" => string(PENALTY_SCHEDULE), + "discount_gamma" => DISCOUNT_GAMMA, "num_train_schedule" => string(something(NUM_TRAIN_SCHEDULE, "fixed")), "eval_schedule" => string(something(EVAL_SCHEDULE, "fixed")), "num_workers" => NUM_WORKERS, @@ -201,10 +251,10 @@ epoch_losses = Float64[] stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] function _build_rollout_de() build_hydro_de(power_data, hydro_data, 1; - backend = nothing, + backend = backend, float_type = Float64, formulation = FORMULATION, - target_penalty = TARGET_PEN_ARG, + target_penalty = resolved_pen * HYDRO_TARGET_PENALTY_MULT, deficit_cost = DEFICIT_COST, demand_matrix = stage_demand, load_scaler = load_scaler, @@ -212,8 +262,8 @@ function _build_rollout_de() end rollout_prob = _build_rollout_de() n_rollout_pool = max(NUM_WORKERS, NUM_EVAL_SCENARIOS) -rollout_pool = [_build_rollout_de() for _ in 1:n_rollout_pool] -@info "Rollout pool ready: $(n_rollout_pool) CPU stage-problem copies" +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:n_rollout_pool] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(n_rollout_pool) stage-problem copies)" : "sequential")" function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) @@ -226,49 +276,43 @@ function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) end hydro_realized_state(stage_prob, result) = - Array(hydro_solution(stage_prob, result).reservoir[:, end]) + hydro_solution(stage_prob, result).reservoir[:, end] + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols + +const _rollout_pen = resolved_pen * HYDRO_TARGET_PENALTY_MULT +const _rollout_pen_l1 = _rollout_pen function hydro_objective_no_target_penalty(stage_prob, result) sol = hydro_solution(stage_prob, result) - delta = Array(sol.delta) - penalty_l2_cost = (resolved_pen / 2) * sum(abs2, delta) - penalty_l1_cost = resolved_pen_l1 * sum(abs, delta) + delta = sol.delta + penalty_l2_cost = (_rollout_pen / 2) * sum(abs2, delta) + penalty_l1_cost = _rollout_pen_l1 * sum(abs, delta) return result.objective - penalty_l2_cost - penalty_l1_cost end Random.seed!(8789) -eval_scenarios = [sample_scenario(hydro_data, T) for _ in 1:NUM_EVAL_SCENARIOS] +eval_scenarios = [sample_scenario(hydro_data, T_ROLLOUT) for _ in 1:NUM_EVAL_SCENARIOS] rollout_evaluation = RolloutEvaluation( rollout_prob, x0_init, eval_scenarios; - horizon = T, + horizon = T_ROLLOUT, n_uncertainty = nHyd, set_stage_parameters! = set_hydro_rollout_stage!, realized_state = hydro_realized_state, objective_no_target_penalty = hydro_objective_no_target_penalty, madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, + warmstart = false, stride = EVAL_EVERY, policy_state = :target, stage_problem_pool = rollout_pool, + retry_on_failure = true, active_scenarios = NUM_EVAL_SCENARIOS, -) -realized_rollout_evaluation = RolloutEvaluation( - rollout_prob, - x0_init, - eval_scenarios; - horizon = T, - n_uncertainty = nHyd, - set_stage_parameters! = set_hydro_rollout_stage!, - realized_state = hydro_realized_state, - objective_no_target_penalty = hydro_objective_no_target_penalty, - madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, - stride = EVAL_EVERY, - policy_state = :realized, - stage_problem_pool = rollout_pool, - active_scenarios = NUM_EVAL_SCENARIOS, + state_bounds = (_min_vols_dev, _max_vols_dev), ) Random.seed!(8788) @@ -281,6 +325,25 @@ function _schedule_value(schedule, iter, default) end current_penalty_mult = Ref(NaN) +last_batch_stats = Ref(Dict{String, Any}()) + +function _merge_batch_stats!(metrics, stats) + isempty(stats) && return metrics + metrics["metrics/train_n_ok"] = get(stats, "n_ok", 0) + metrics["metrics/train_n_total"] = get(stats, "n_total", 0) + metrics["metrics/train_success_share"] = + get(stats, "n_total", 0) == 0 ? NaN : get(stats, "n_ok", 0) / get(stats, "n_total", 0) + for (k, v) in get(stats, "status_counts", Dict{String, Int}()) + metrics["metrics/train_status/$k"] = v + end + for (k, v) in get(stats, "failure_counts", Dict{String, Int}()) + metrics["metrics/train_failure/$k"] = v + end + for (k, v) in get(stats, "retry_counts", Dict{String, Int}()) + metrics["metrics/train_retry/$k"] = v + end + return metrics +end train_tsddr( policy, @@ -300,48 +363,51 @@ train_tsddr( madnlp_kwargs = SOLVER_KWARGS, warmstart = true, problem_pool = problem_pool, + batch_diagnostics = (iter, stats) -> begin + last_batch_stats[] = stats + n_ok = get(stats, "n_ok", 0) + n_total = get(stats, "n_total", 0) + if n_ok < n_total + @warn "Training solve failures at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) failure_counts=get(stats, "failure_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + elseif iter % 10 == 0 + @info "Training solve status at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + end + end, adjust_hyperparameters = (iter, opt_state, n) -> begin mult = _schedule_value(PENALTY_SCHEDULE, iter, last(PENALTY_SCHEDULE)[3]) if mult != current_penalty_mult[] current_penalty_mult[] = mult ρ_half_scaled = prob.base_penalty_half * mult ρ_l1_scaled = prob.base_penalty_l1 * mult - penalty_vals = fill(ρ_half_scaled, T * nHyd) - penalty_l1_vals = fill(ρ_l1_scaled, T * nHyd) + penalty_vals = ρ_half_scaled .* _discount_weights + penalty_l1_vals = ρ_l1_scaled .* _discount_weights for (p, _, _, _) in problem_pool ExaModels.set_parameter!(p.core, p.p_penalty_half, penalty_vals) ExaModels.set_parameter!(p.core, p.p_penalty_l1, penalty_l1_vals) end - @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)))" + @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)), γ=$DISCOUNT_GAMMA)" end if !isnothing(EVAL_SCHEDULE) n_eval = _schedule_value(EVAL_SCHEDULE, iter, NUM_EVAL_SCENARIOS) rollout_evaluation.active_scenarios = n_eval - realized_rollout_evaluation.active_scenarios = n_eval end return isnothing(NUM_TRAIN_SCHEDULE) ? n : _schedule_value(NUM_TRAIN_SCHEDULE, iter, n) end, record_loss = (iter, m, loss, tag) -> begin metrics = Dict{String, Any}(tag => loss, "batch" => iter) + _merge_batch_stats!(metrics, last_batch_stats[]) isfinite(loss) && push!(epoch_losses, loss) if iter % EVAL_EVERY == 0 rollout_evaluation(iter, m) - realized_rollout_evaluation(iter, m) metrics["metrics/rollout_objective_no_target_penalty"] = rollout_evaluation.last_objective_no_target_penalty metrics["metrics/rollout_objective_no_deficit"] = rollout_evaluation.last_objective_no_target_penalty metrics["metrics/rollout_target_violation_share"] = rollout_evaluation.last_violation_share - metrics["metrics/rollout_realized_objective_no_target_penalty"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_objective_no_deficit"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_target_violation_share"] = - realized_rollout_evaluation.last_violation_share metrics["metrics/rollout_n_ok"] = - realized_rollout_evaluation.last_n_ok + rollout_evaluation.last_n_ok end if !isnan(current_penalty_mult[]) diff --git a/examples/HydroPowerModels/train_hydro_exa_critic.jl b/examples/HydroPowerModels/train_hydro_exa_critic.jl index 3d9a5a0..ea93110 100644 --- a/examples/HydroPowerModels/train_hydro_exa_critic.jl +++ b/examples/HydroPowerModels/train_hydro_exa_critic.jl @@ -34,11 +34,12 @@ const DEMAND_FILE = joinpath(CASE_DIR, "demand.csv") const LAYERS = [128, 128] const ACTIVATION = sigmoid -const NUM_STAGES = 96 +const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) const NUM_EPOCHS = 80 const NUM_BATCHES = 100 const MAX_EVAL_SCENARIOS = 32 -const EVAL_EVERY = 25 +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) const EVAL_SCHEDULE = [ (1, div(NUM_EPOCHS * NUM_BATCHES, 2), 4), @@ -56,21 +57,32 @@ const CRITIC_BUFFER_SIZE = 512 const CRITIC_BATCH_SIZE = 32 const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = parse(Float64, get(ENV, "DR_TARGET_PENALTY_MULT", "8.0")) const DEFICIT_COST = 1e5 const USE_GPU = true const load_scaler = 0.6 const NUM_WORKERS = 4 const CRITIC_ROLLOUT_SAMPLES_PER_BATCH = 0 # eval rollouts feed the critic via external_critic_samples const CRITIC_POLICY_STATE = :target # set to :realized for closed-loop critic targets +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) const CRITIC_ROLLOUT_OBJECTIVE = :objective const NUM_CHEAP_CRITIC_SAMPLES_PER_BATCH = 4 * NUM_WORKERS -const PENALTY_SCHEDULE = [ - (1, div(NUM_EPOCHS * NUM_BATCHES, 4), 0.1), - (div(NUM_EPOCHS * NUM_BATCHES, 4) + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 2, 1.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 2 + 1, div(NUM_EPOCHS * NUM_BATCHES, 4) * 3, 10.0), - (div(NUM_EPOCHS * NUM_BATCHES, 4) * 3 + 1, NUM_EPOCHS * NUM_BATCHES, 30.0), -] +const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "const") +const _N_TOTAL = NUM_EPOCHS * NUM_BATCHES +const _ANNEAL_1_END = max(1, div(_N_TOTAL, 100)) +const _ANNEAL_2_END = max(_ANNEAL_1_END + 1, div(_N_TOTAL, 40)) +const _ANNEAL_3_END = max(_ANNEAL_2_END + 1, div(_N_TOTAL, 10)) +const PENALTY_SCHEDULE = if _PENALTY_MODE == "annealed" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +else + [(1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT)] +end const NUM_TRAIN_SCHEDULE = [ (1, div(NUM_EPOCHS * NUM_BATCHES, 5), NUM_WORKERS), @@ -80,9 +92,10 @@ const NUM_TRAIN_SCHEDULE = [ (div(NUM_EPOCHS * NUM_BATCHES, 5) * 4 + 1, NUM_EPOCHS * NUM_BATCHES, 8 * NUM_WORKERS), ] -const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = 9000) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) -const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-deteq-gpu-critic-cv-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-deteq-gpu-critic-cv-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") mkpath(MODEL_DIR) const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") @@ -97,9 +110,10 @@ power_data = load_power_data(PM_FILE) @info "Loading hydro data..." hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; - num_stages = NUM_STAGES * 10) + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) nHyd = hydro_data.nHyd T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES @info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios)" demand_mat = if isfile(DEMAND_FILE) @@ -115,6 +129,7 @@ resolved_pen = TARGET_PEN_ARG === :auto ? auto_target_penalty(power_data, hydro_data) : Float64(TARGET_PEN_ARG) @info "Auto target penalty: ρ=$(round(resolved_pen; digits=2))" +@info "Bolivia hydro target-penalty multiplier default: $(HYDRO_TARGET_PENALTY_MULT)" backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : (@info "Using CPU backend"; nothing) @@ -133,6 +148,7 @@ end @info "Building $(T)-stage ExaModels DE (formulation=$FORMULATION)..." prob = _build_de() +resolved_pen_l1 = prob.base_penalty_l1 @info "Building $(NUM_WORKERS)-worker problem pool..." problem_pool = [(prob, prob.p_x0, prob.p_target, prob.p_inflow)] @@ -146,6 +162,8 @@ x0_init = Float32.([clamp(hydro_data.initial_volumes[r], hydro_data.units[r].min_vol, hydro_data.units[r].max_vol) for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) # ── Critic/control variate ─────────────────────────────────────────────────── @@ -196,13 +214,27 @@ solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding # ── Policy ──────────────────────────────────────────────────────────────────── -policy = StateConditionedPolicy(nHyd, nHyd, nHyd, LAYERS; - activation = ACTIVATION, - encoder_type = Flux.LSTM) +policy_active_mask = trues(nHyd) +policy = bounded_state_policy(nHyd, target_lower, target_upper, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + active_mask = policy_active_mask) if !isnothing(PRE_TRAINED) @info "Loading pre-trained model from $(PRE_TRAINED)..." - Flux.loadmodel!(policy, JLD2.load(PRE_TRAINED, "model_state")) + load_stateconditioned_policy!(policy, JLD2.load(PRE_TRAINED, "model_state")) +end + +if USE_GPU + policy = CUDA.cu(policy) + x0_init = CUDA.cu(x0_init) + control_variate = ScalarCriticControlVariate( + CUDA.cu(control_variate.critic); + featurizer = control_variate.featurizer, + value_loss_weight = control_variate.value_loss_weight, + gradient_loss_weight = control_variate.gradient_loss_weight, + ) + @info "Policy, critic, and x0 moved to GPU" end # ── W&B logging ─────────────────────────────────────────────────────────────── @@ -215,9 +247,12 @@ lg = WandbLogger( "case" => CASE_NAME, "formulation" => FORM_LABEL, "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, "layers" => LAYERS, "activation" => string(ACTIVATION), "target_penalty" => "auto=$(round(resolved_pen; digits=2))", + "target_penalty_l1" => "auto=$(round(resolved_pen_l1; digits=2))", + "hydro_target_penalty_mult" => HYDRO_TARGET_PENALTY_MULT, "deficit_cost" => DEFICIT_COST, "num_epochs" => NUM_EPOCHS, "num_batches" => NUM_BATCHES, @@ -257,7 +292,7 @@ epoch_losses = Float64[] stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] function _build_rollout_de() build_hydro_de(power_data, hydro_data, 1; - backend = nothing, + backend = backend, float_type = Float64, formulation = FORMULATION, target_penalty = TARGET_PEN_ARG, @@ -267,8 +302,8 @@ function _build_rollout_de() ) end rollout_prob = _build_rollout_de() -rollout_pool = [_build_rollout_de() for _ in 1:NUM_WORKERS] -@info "Rollout pool ready: $(NUM_WORKERS) CPU stage-problem copies" +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:NUM_WORKERS] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(NUM_WORKERS) stage-problem copies)" : "sequential")" function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) @@ -281,59 +316,53 @@ function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) end hydro_realized_state(stage_prob, result) = - Array(hydro_solution(stage_prob, result).reservoir[:, end]) + hydro_solution(stage_prob, result).reservoir[:, end] + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols function hydro_objective_no_target_penalty(stage_prob, result) sol = hydro_solution(stage_prob, result) - return result.objective - (resolved_pen / 2) * sum(abs2, Array(sol.delta)) + penalty_l2_cost = (resolved_pen / 2) * sum(abs2, sol.delta) + penalty_l1_cost = resolved_pen_l1 * sum(abs, sol.delta) + return result.objective - penalty_l2_cost - penalty_l1_cost end Random.seed!(8789) -eval_scenarios = [sample_scenario(hydro_data, T) for _ in 1:MAX_EVAL_SCENARIOS] +eval_scenarios = [sample_scenario(hydro_data, T_ROLLOUT) for _ in 1:MAX_EVAL_SCENARIOS] rollout_evaluation = RolloutEvaluation( rollout_prob, x0_init, eval_scenarios; - horizon = T, - n_uncertainty = nHyd, - set_stage_parameters! = set_hydro_rollout_stage!, - realized_state = hydro_realized_state, - objective_no_target_penalty = hydro_objective_no_target_penalty, - madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, - stride = EVAL_EVERY, - policy_state = :target, - stage_problem_pool = rollout_pool, - active_scenarios = 4, -) -realized_rollout_evaluation = RolloutEvaluation( - rollout_prob, - x0_init, - eval_scenarios; - horizon = T, + horizon = T_ROLLOUT, n_uncertainty = nHyd, set_stage_parameters! = set_hydro_rollout_stage!, realized_state = hydro_realized_state, objective_no_target_penalty = hydro_objective_no_target_penalty, madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, + warmstart = false, stride = EVAL_EVERY, policy_state = :realized, stage_problem_pool = rollout_pool, + retry_on_failure = true, active_scenarios = 4, + state_bounds = (_min_vols_dev, _max_vols_dev), ) critic_training_target = RolloutCriticTarget( rollout_prob; - horizon = T, + horizon = T_ROLLOUT, n_uncertainty = nHyd, set_stage_parameters! = set_hydro_rollout_stage!, realized_state = hydro_realized_state, objective_no_target_penalty = hydro_objective_no_target_penalty, madnlp_kwargs = SOLVER_KWARGS, - warmstart = true, + warmstart = false, policy_state = CRITIC_POLICY_STATE, objective_value = CRITIC_ROLLOUT_OBJECTIVE, + state_bounds = (_min_vols_dev, _max_vols_dev), ) Random.seed!(8788) @@ -381,15 +410,17 @@ train_tsddr( if mult != current_penalty_mult[] current_penalty_mult[] = mult ρ_half_scaled = prob.base_penalty_half * mult + ρ_l1_scaled = prob.base_penalty_l1 * mult penalty_vals = fill(ρ_half_scaled, T * nHyd) + penalty_l1_vals = fill(ρ_l1_scaled, T * nHyd) for (p, _, _, _) in problem_pool ExaModels.set_parameter!(p.core, p.p_penalty_half, penalty_vals) + ExaModels.set_parameter!(p.core, p.p_penalty_l1, penalty_l1_vals) end - @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)))" + @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)))" end n_eval = _schedule_value(EVAL_SCHEDULE, iter, MAX_EVAL_SCENARIOS) rollout_evaluation.active_scenarios = n_eval - realized_rollout_evaluation.active_scenarios = n_eval return _schedule_value(NUM_TRAIN_SCHEDULE, iter, n) end, record_loss = (iter, m, loss, tag) -> begin @@ -398,7 +429,6 @@ train_tsddr( if iter % EVAL_EVERY == 0 rollout_evaluation(iter, m) - realized_rollout_evaluation(iter, m) append!(shared_critic_samples, critic_samples_from_evaluation( rollout_evaluation; @@ -410,14 +440,8 @@ train_tsddr( rollout_evaluation.last_objective_no_target_penalty metrics["metrics/rollout_target_violation_share"] = rollout_evaluation.last_violation_share - metrics["metrics/rollout_realized_objective_no_target_penalty"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_objective_no_deficit"] = - realized_rollout_evaluation.last_objective_no_target_penalty - metrics["metrics/rollout_realized_target_violation_share"] = - realized_rollout_evaluation.last_violation_share metrics["metrics/rollout_n_ok"] = - realized_rollout_evaluation.last_n_ok + rollout_evaluation.last_n_ok end if !isnan(current_penalty_mult[]) diff --git a/examples/HydroPowerModels/train_hydro_exa_embedded.jl b/examples/HydroPowerModels/train_hydro_exa_embedded.jl new file mode 100644 index 0000000..5d8c788 --- /dev/null +++ b/examples/HydroPowerModels/train_hydro_exa_embedded.jl @@ -0,0 +1,457 @@ +# train_hydro_exa_embedded.jl +# +# Embedded-NN hydro training with ExaModels + MadNLP (AC or DC OPF). +# Policy is embedded directly in the NLP via VectorNonlinearOracle. +# Gradient: envelope theorem (multiplier-weighted policy Jacobian). +# +# 4 configurations via environment variables: +# DR_PENALTY_SCHEDULE = "const" | "annealed" +# DR_TARGET_PENALTY_MULT = 8.0 (Bolivia default multiplier on :auto penalties) +# DR_PRETRAIN_ITERS = 0 | 500 (regular TSDDR warmup before embedded training) +# DR_PRETRAIN_PENALTY_MULT = 0.1 (penalty multiplier during pretrain, default 0.1) + +using DecisionRulesExa +using ExaModels +using Flux +using Statistics, Random, Dates +using Wandb, Logging +using JLD2 +using MadNLP, MadNLPGPU +using CUDA, CUDSS, KernelAbstractions + +const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) +include(joinpath(SCRIPT_DIR, "hydro_reachable_policy.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa_embedded.jl")) + +# ── Configuration ───────────────────────────────────────────────────────────── + +const CASE_NAME = "bolivia" +const FORMULATION = :ac_polar +const FORM_LABEL = FORMULATION === :ac_polar ? "ACPPowerModel" : "DCPPowerModel" + +const CASE_DIR = joinpath(SCRIPT_DIR, CASE_NAME) +const PM_FILE = joinpath(CASE_DIR, "PowerModels.json") +const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") +const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") +const DEMAND_FILE = joinpath(CASE_DIR, "demand.csv") + +const LAYERS = parse_layers(get(ENV, "DR_LAYERS", "128,128")) +const ACTIVATION = sigmoid +const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) +const NUM_EPOCHS = parse(Int, get(ENV, "DR_NUM_EPOCHS", "80")) +const NUM_BATCHES = 100 +const NUM_TRAIN_PER_BATCH = 1 +const NUM_EVAL_SCENARIOS = 4 +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) +const LR = 1f-3 + +const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = parse(Float64, get(ENV, "DR_TARGET_PENALTY_MULT", "8.0")) +const DEFICIT_COST = 1e5 +const load_scaler = 0.6 + +const PRETRAIN_ITERS = parse(Int, get(ENV, "DR_PRETRAIN_ITERS", "0")) +const PRETRAIN_PENALTY_MULT = parse(Float64, get(ENV, "DR_PRETRAIN_PENALTY_MULT", "0.1")) +const STRICT_EMBEDDED_TARGETS = parse(Bool, get(ENV, "DR_STRICT_EMBEDDED_TARGETS", "false")) + +const DISCOUNT_GAMMA = parse(Float64, get(ENV, "DR_DISCOUNT_GAMMA", "1.0")) +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) + +const _PENALTY_MODE = get(ENV, "DR_PENALTY_SCHEDULE", "const") +const _N_TOTAL = NUM_EPOCHS * NUM_BATCHES +const _ANNEAL_1_END = max(1, div(_N_TOTAL, 100)) +const _ANNEAL_2_END = max(_ANNEAL_1_END + 1, div(_N_TOTAL, 40)) +const _ANNEAL_3_END = max(_ANNEAL_2_END + 1, div(_N_TOTAL, 10)) +const PENALTY_SCHEDULE = if _PENALTY_MODE == "annealed" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +elseif _PENALTY_MODE == "annealed_discount" + [ + (1, _ANNEAL_1_END, 0.1), + (_ANNEAL_1_END + 1, _ANNEAL_2_END, 1.0), + (_ANNEAL_2_END + 1, _ANNEAL_3_END, 4.0), + (_ANNEAL_3_END + 1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT), + ] +else + [(1, _N_TOTAL, HYDRO_TARGET_PENALTY_MULT)] +end + +const USE_GPU = parse(Bool, get(ENV, "DR_USE_GPU", "true")) +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) + +const _DISC_TAG = DISCOUNT_GAMMA < 1.0 ? "-disc$(replace(string(DISCOUNT_GAMMA), "." => ""))" : "" +const _SCHED_TAG = if _PENALTY_MODE == "annealed" + "-anneal" +elseif _PENALTY_MODE == "annealed_discount" + "-anndisc" +else + "-const" +end +const _PRETRAIN_TAG = PRETRAIN_ITERS > 0 ? "-pre$(PRETRAIN_ITERS)" : "" +const _GPU_TAG = USE_GPU ? "-gpu" : "" +const _LAYER_TAG = LAYERS == [128, 128] ? "" : "-L$(join(LAYERS, "_"))" +const _STRICT_TAG = STRICT_EMBEDDED_TARGETS ? "-strict" : "" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-embedded$(_GPU_TAG)$(_SCHED_TAG)$(_DISC_TAG)$(_LAYER_TAG)$(_PRETRAIN_TAG)$(_STRICT_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") +mkpath(MODEL_DIR) +const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") + +# ── Load data ───────────────────────────────────────────────────────────────── + +@info "Loading power system data..." +power_data = load_power_data(PM_FILE) +@info " nBus=$(power_data.nBus) nGen=$(power_data.nGen)" + +@info "Loading hydro data..." +hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) +nHyd = hydro_data.nHyd +T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES +@info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios)" + +demand_mat = if isfile(DEMAND_FILE) + @info "Loading demand from $(DEMAND_FILE)..." + load_demand(DEMAND_FILE, power_data; T = T) +else + nothing +end + +resolved_pen = TARGET_PEN_ARG === :auto ? + auto_target_penalty(power_data, hydro_data) : + Float64(TARGET_PEN_ARG) +@info "Auto target penalty: ρ=$(round(resolved_pen; digits=2))" +@info "Bolivia hydro target-penalty multiplier default: $(HYDRO_TARGET_PENALTY_MULT)" + +backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : + (@info "Using CPU backend"; nothing) + +x0_init = Float32.([clamp(hydro_data.initial_volumes[r], + hydro_data.units[r].min_vol, + hydro_data.units[r].max_vol) + for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) + +const _discount_weights = Float64[DISCOUNT_GAMMA^(t-1) for t in 1:T for _ in 1:nHyd] +if DISCOUNT_GAMMA < 1.0 + @info "Discount γ=$(DISCOUNT_GAMMA): stage 1 weight=1.0, stage $T weight=$(round(DISCOUNT_GAMMA^(T-1); sigdigits=4))" +end + +# ── Policy ──────────────────────────────────────────────────────────────────── + +Random.seed!(42) +policy = if STRICT_EMBEDDED_TARGETS + @info "Using strict embedded hydro policy with one-stage reachable target bounds" + hydro_reachable_policy(hydro_data, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM) +else + bounded_state_policy(nHyd, target_lower, target_upper, LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + active_mask = trues(nHyd)) +end + +# ── Optional pretrain with regular TSDDR ────────────────────────────────────── + +if PRETRAIN_ITERS > 0 + @info "Building regular DE for pretrain ($(PRETRAIN_ITERS) iters)..." + prob_reg = build_hydro_de(power_data, hydro_data, T; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = demand_mat, + load_scaler = load_scaler, + ) + + if PRETRAIN_PENALTY_MULT != 1.0 + ρ_half_pre = prob_reg.base_penalty_half * PRETRAIN_PENALTY_MULT + ρ_l1_pre = prob_reg.base_penalty_l1 * PRETRAIN_PENALTY_MULT + ExaModels.set_parameter!(prob_reg.core, prob_reg.p_penalty_half, + fill(ρ_half_pre, T * nHyd)) + ExaModels.set_parameter!(prob_reg.core, prob_reg.p_penalty_l1, + fill(ρ_l1_pre, T * nHyd)) + @info " Pretrain penalty mult=$(PRETRAIN_PENALTY_MULT) (ρ/2=$(round(ρ_half_pre; digits=2)), l1=$(round(ρ_l1_pre; digits=2)))" + end + + @info "Pretraining policy with regular TSDDR..." + train_tsddr( + policy, x0_init, prob_reg, + prob_reg.p_x0, prob_reg.p_target, prob_reg.p_inflow, + () -> sample_scenario(hydro_data, T); + num_batches = PRETRAIN_ITERS, + num_train_per_batch = 1, + optimizer = Flux.Adam(LR), + madnlp_kwargs = SOLVER_KWARGS, + warmstart = true, + record_loss = (iter, m, loss, tag) -> begin + if iter % 50 == 0 + @info " Pretrain iter $iter/$PRETRAIN_ITERS: loss=$(round(loss; digits=2))" + end + return false + end, + ) + @info "Pretrain done." +end + +# ── Move policy + x0 to GPU (BEFORE embedded DE build so oracle captures GPU policy) ── + +if USE_GPU + policy = CUDA.cu(policy) + x0_init = CUDA.cu(x0_init) + @info "Policy and x0 moved to GPU" +end + +# ── Build embedded DE ───────────────────────────────────────────────────────── + +@info "Building embedded $(T)-stage $(FORMULATION) hydro DE..." +prob_emb = build_embedded_hydro_de(policy, power_data, hydro_data, T; + backend = backend, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = demand_mat, + load_scaler = load_scaler, + strict_targets = STRICT_EMBEDDED_TARGETS, +) +@info " nvar=$(prob_emb._nvar) oracle_cons=$(length(prob_emb.target_con_range)) strict_targets=$(prob_emb.strict_targets)" + +resolved_pen_l1 = prob_emb.base_penalty_l1 + +# ── Smoke test ──────────────────────────────────────────────────────────────── + +w_mean = mean_inflow(hydro_data, T) +set_x0!(prob_emb, x0_init) +set_inflows!(prob_emb, w_mean) +@info "Smoke test: solving embedded DE with mean inflows..." +result0 = MadNLP.madnlp(prob_emb.model; SOLVER_KWARGS..., print_level = MadNLP.WARN) +@info " Status: $(result0.status) Obj: $(round(result0.objective; digits=4))" +solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding anyway" + +# ── W&B logging ─────────────────────────────────────────────────────────────── + +lg = WandbLogger( + project = "RL", + name = RUN_NAME, + save_code = false, + config = Dict( + "case" => CASE_NAME, + "formulation" => FORM_LABEL, + "method" => "embedded_nn", + "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, + "layers" => LAYERS, + "activation" => string(ACTIVATION), + "target_penalty" => "auto=$(round(resolved_pen; digits=2))", + "target_penalty_l1" => "auto=$(round(resolved_pen_l1; digits=2))", + "hydro_target_penalty_mult" => HYDRO_TARGET_PENALTY_MULT, + "deficit_cost" => DEFICIT_COST, + "num_epochs" => NUM_EPOCHS, + "num_batches" => NUM_BATCHES, + "num_train_per_batch" => NUM_TRAIN_PER_BATCH, + "num_eval_scenarios" => NUM_EVAL_SCENARIOS, + "eval_every" => EVAL_EVERY, + "lr" => LR, + "penalty_schedule" => string(PENALTY_SCHEDULE), + "discount_gamma" => DISCOUNT_GAMMA, + "pretrain_iters" => PRETRAIN_ITERS, + "pretrain_penalty_mult" => PRETRAIN_PENALTY_MULT, + "strict_embedded_targets" => STRICT_EMBEDDED_TARGETS, + ), +) + +# ── Rollout evaluation ──────────────────────────────────────────────────────── + +stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] +function _build_rollout_de() + build_hydro_de(power_data, hydro_data, 1; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = stage_demand, + load_scaler = load_scaler, + strict_targets = STRICT_EMBEDDED_TARGETS, + ) +end + +rollout_prob = _build_rollout_de() +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:NUM_EVAL_SCENARIOS] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(NUM_EVAL_SCENARIOS) stage-problem copies)" : "sequential")" + +function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_inflow, wt) + if demand_mat !== nothing + set_demand!(stage_prob, load_scaler .* demand_mat[stage:stage, :]) + end + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_target, target) + prepare_solve!(stage_prob, state_in, wt, target) + return stage_prob +end + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols + +hydro_realized_state(stage_prob, result) = + hydro_solution(stage_prob, result).reservoir[:, end] + +function hydro_objective_no_target_penalty(stage_prob, result) + if STRICT_EMBEDDED_TARGETS + return result.objective + end + sol = hydro_solution(stage_prob, result) + delta = sol.delta + penalty_l2_cost = (resolved_pen / 2) * sum(abs2, delta) + penalty_l1_cost = resolved_pen_l1 * sum(abs, delta) + return result.objective - penalty_l2_cost - penalty_l1_cost +end + +Random.seed!(8789) +eval_scenarios = [sample_scenario(hydro_data, T_ROLLOUT) for _ in 1:NUM_EVAL_SCENARIOS] + +rollout_evaluation = RolloutEvaluation( + rollout_prob, x0_init, eval_scenarios; + horizon = T_ROLLOUT, n_uncertainty = nHyd, + set_stage_parameters! = set_hydro_rollout_stage!, + realized_state = hydro_realized_state, + objective_no_target_penalty = hydro_objective_no_target_penalty, + madnlp_kwargs = SOLVER_KWARGS, + warmstart = false, + stride = EVAL_EVERY, + policy_state = :realized, + stage_problem_pool = rollout_pool, + retry_on_failure = true, + active_scenarios = NUM_EVAL_SCENARIOS, + state_bounds = (_min_vols_dev, _max_vols_dev), +) + +# ── Training ────────────────────────────────────────────────────────────────── + +Random.seed!(8788) + +best_obj = Inf +epoch_losses = Float64[] + +function _schedule_value(schedule, iter, default) + for (lo, hi, val) in schedule + lo <= iter <= hi && return val + end + return default +end + +current_penalty_mult = Ref(NaN) +last_batch_stats = Ref(Dict{String, Any}()) + +function _merge_batch_stats!(metrics, stats) + isempty(stats) && return metrics + metrics["metrics/train_n_ok"] = get(stats, "n_ok", 0) + metrics["metrics/train_n_total"] = get(stats, "n_total", 0) + metrics["metrics/train_success_share"] = + get(stats, "n_total", 0) == 0 ? NaN : get(stats, "n_ok", 0) / get(stats, "n_total", 0) + for (k, v) in get(stats, "status_counts", Dict{String, Int}()) + metrics["metrics/train_status/$k"] = v + end + for (k, v) in get(stats, "failure_counts", Dict{String, Int}()) + metrics["metrics/train_failure/$k"] = v + end + for (k, v) in get(stats, "retry_counts", Dict{String, Int}()) + metrics["metrics/train_retry/$k"] = v + end + return metrics +end + +@info "Starting embedded training: $(NUM_EPOCHS) epochs × $(NUM_BATCHES) batches" + +train_tsddr_embedded( + policy, x0_init, prob_emb, + () -> sample_scenario(hydro_data, T); + num_batches = NUM_EPOCHS * NUM_BATCHES, + num_train_per_batch = NUM_TRAIN_PER_BATCH, + optimizer = Flux.Adam(LR), + madnlp_kwargs = SOLVER_KWARGS, + warmstart = true, + get_realized_states = embedded_hydro_realized_states, + batch_diagnostics = (iter, stats) -> begin + last_batch_stats[] = stats + n_ok = get(stats, "n_ok", 0) + n_total = get(stats, "n_total", 0) + if n_ok < n_total + @warn "Embedded training solve failures at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) failure_counts=get(stats, "failure_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + elseif iter % 10 == 0 + @info "Embedded training solve status at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + end + end, + adjust_hyperparameters = (iter, opt_state, n) -> begin + mult = _schedule_value(PENALTY_SCHEDULE, iter, last(PENALTY_SCHEDULE)[3]) + if mult != current_penalty_mult[] + current_penalty_mult[] = mult + if STRICT_EMBEDDED_TARGETS + @info "Strict target equality active; target slack penalties are not used" + else + ρ_half_scaled = prob_emb.base_penalty_half * mult + ρ_l1_scaled = prob_emb.base_penalty_l1 * mult + penalty_vals = ρ_half_scaled .* _discount_weights + penalty_l1_vals = ρ_l1_scaled .* _discount_weights + ExaModels.set_parameter!(prob_emb.core, prob_emb.p_penalty_half, penalty_vals) + ExaModels.set_parameter!(prob_emb.core, prob_emb.p_penalty_l1, penalty_l1_vals) + @info "Penalty multiplier → $mult (ρ/2 = $(round(ρ_half_scaled; digits=2)), λ_l1 = $(round(ρ_l1_scaled; digits=2)), γ=$DISCOUNT_GAMMA)" + end + end + return n + end, + record_loss = (iter, m, loss, tag) -> begin + metrics = Dict{String, Any}(tag => loss, "batch" => iter) + _merge_batch_stats!(metrics, last_batch_stats[]) + isfinite(loss) && push!(epoch_losses, loss) + + if iter % EVAL_EVERY == 0 + rollout_evaluation(iter, m) + metrics["metrics/rollout_objective_no_target_penalty"] = + rollout_evaluation.last_objective_no_target_penalty + metrics["metrics/rollout_target_violation_share"] = + rollout_evaluation.last_violation_share + metrics["metrics/rollout_n_ok"] = + rollout_evaluation.last_n_ok + end + + if !isnan(current_penalty_mult[]) + metrics["metrics/target_penalty_multiplier"] = current_penalty_mult[] + end + + batch_in_epoch = (iter - 1) % NUM_BATCHES + 1 + if batch_in_epoch == NUM_BATCHES + epoch = (iter - 1) ÷ NUM_BATCHES + 1 + mean_loss = isempty(epoch_losses) ? NaN : mean(epoch_losses) + n_ok = length(epoch_losses) + empty!(epoch_losses) + Wandb.log(lg, Dict("metrics/epoch_objective" => mean_loss, "epoch" => epoch)) + @info "Epoch $epoch/$NUM_EPOCHS mean=$(round(mean_loss; digits=2)) ok=$n_ok/$NUM_BATCHES" + if isfinite(mean_loss) && mean_loss < best_obj + global best_obj = mean_loss + jldsave(MODEL_PATH; model_state = Flux.state(cpu(m))) + @info " → New best: $(round(mean_loss; digits=4)) — saved $MODEL_PATH" + end + end + Wandb.log(lg, metrics) + return false + end, +) + +close(lg) +@info "Done. Best model saved to: $(MODEL_PATH)" diff --git a/examples/HydroPowerModels/train_hydro_exa_strict.jl b/examples/HydroPowerModels/train_hydro_exa_strict.jl new file mode 100644 index 0000000..d35c77a --- /dev/null +++ b/examples/HydroPowerModels/train_hydro_exa_strict.jl @@ -0,0 +1,1042 @@ +# train_hydro_exa_strict.jl +# +# Strict-mode regular DE training with ExaModels + MadNLP (AC polar OPF). +# Uses HydroReachablePolicy (sigmoid-bounded to one-stage reachable set) +# with strict_targets=true (delta variables fixed to zero, no target penalty). +# +# Key insight: ordinary regular DE target generation is open-loop after x0, so +# strict equality is usually unsafe: the optimizer may discover realized states +# different from the target path, and later targets were not computed from those +# realized states. Here the reachable policy is rolled out from the true x0 and +# uses the previous target as the next policy state. Since every target is +# one-stage reachable from the previous target, the full target trajectory is +# feasible by induction and strict regular DE is valid. +# +# Environment variables: +# DR_ENCODER_LAYERS = "128,128" (LSTM encoder layer sizes) +# DR_HEAD_LAYERS = "" (state-conditioned target-head hidden sizes) +# DR_LAYERS = "128,128" (legacy alias for DR_ENCODER_LAYERS) +# DR_NUM_STAGES = "126" +# DR_NUM_ROLLOUT_STAGES = "96" +# DR_NUM_EPOCHS = "80" +# DR_NUM_BATCHES = "100" +# DR_NUM_TRAIN_PER_BATCH = "1" (sampled trajectories per gradient step) +# DR_NUM_TRAIN_SCHEDULE = "" (optional "lo:hi:n,..." schedule) +# DR_CONTEXT = "" (""/"none", "phase", or "phase+progress") +# DR_NUM_EVAL_SCENARIOS = "4" (fixed held-out rollout-selection scenarios) +# DR_EVAL_SCHEDULE = "" (optional "lo:hi:n,..." active-eval schedule) +# DR_EVAL_EVERY = "50" +# DR_SAVE_METRIC = "training" ("training" or "rollout") +# DR_LR = "0.001" +# DR_LR_FINAL = DR_LR (cosine-decay final learning rate) +# DR_LR_WARMUP = "0" (linear warmup iterations from DR_LR/100) +# DR_PRETRAINED_MODEL = "" (optional diagnostic warmstart checkpoint) +# DR_REACTIVE_DEFICIT = "free" ("free", "hard", or a finite penalty) +# DR_GRAD_CLIP = "0" +# DR_MAX_ITER = "9000" +# +# Reproducible recipes: +# +# Historical baseline (all defaults, byte-compatible training semantics): +# julia --project -t auto train_hydro_exa_strict.jl +# +# Fresh fast/fair candidate (headline timing starts at this command): +# DR_REACTIVE_DEFICIT=hard DR_SAVE_METRIC=rollout \ +# DR_NUM_EVAL_SCENARIOS=8 DR_EVAL_EVERY=100 \ +# DR_NUM_TRAIN_PER_BATCH=2 \ +# DR_LR=5e-4 DR_LR_FINAL=5e-5 DR_LR_WARMUP=100 \ +# DR_NUM_EPOCHS=25 DR_NUM_BATCHES=100 \ +# DR_HEAD_LAYERS=128,128 \ +# julia --project -t auto train_hydro_exa_strict.jl +# +# Progressive-sampling Phase B (spend samples only after coarse movement): +# DR_REACTIVE_DEFICIT=hard DR_SAVE_METRIC=rollout \ +# DR_NUM_EVAL_SCENARIOS=12 DR_EVAL_SCHEDULE=1:800:4,801:1400:8,1401:1800:12 \ +# DR_NUM_TRAIN_PER_BATCH=1 DR_NUM_TRAIN_SCHEDULE=1:800:1,801:1400:2,1401:1800:4 \ +# DR_LR=7e-4 DR_LR_FINAL=2e-5 DR_LR_WARMUP=100 \ +# DR_NUM_EPOCHS=18 DR_NUM_BATCHES=100 \ +# DR_HEAD_LAYERS=256,256 \ +# julia --project -t auto train_hydro_exa_strict.jl +# +# Diagnostic warmstart (not a clean headline unless parent time is counted): +# DR_PRETRAINED_MODEL= DR_SAVE_METRIC=rollout \ +# DR_NUM_TRAIN_PER_BATCH=4 DR_LR=1e-4 DR_LR_FINAL=1e-5 DR_LR_WARMUP=50 \ +# DR_NUM_EPOCHS=10 DR_HEAD_LAYERS= \ +# julia --project -t auto train_hydro_exa_strict.jl +# +# Usage: +# julia --project -t auto train_hydro_exa_strict.jl + +using DecisionRulesExa +using StableRNGs +using ExaModels +using Flux +using Statistics, Random, Dates +using Logging # Wandb loaded conditionally below (DR_ENABLE_WANDB) — see note at ENABLE_WANDB +using JLD2 +using MadNLP +using MadNLPGPU, KernelAbstractions, CUDA +using CUDSS, CUDSS_jll, cuDNN + +const SCRIPT_DIR = dirname(@__FILE__) +include(joinpath(SCRIPT_DIR, "hydro_training_utils.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_data.jl")) +include(joinpath(SCRIPT_DIR, "hydro_power_exa.jl")) +include(joinpath(SCRIPT_DIR, "hydro_reachable_policy.jl")) + +# ── Configuration ───────────────────────────────────────────────────────────── + +const CASE_NAME = "bolivia" +const FORMULATION = :ac_polar +const FORM_LABEL = FORMULATION === :ac_polar ? "ACPPowerModel" : "DCPPowerModel" + +const CASE_DIR = joinpath(SCRIPT_DIR, CASE_NAME) +const PM_FILE = joinpath(CASE_DIR, "PowerModels.json") +const HYDRO_FILE = joinpath(CASE_DIR, "hydro.json") +const INFLOW_FILE = joinpath(CASE_DIR, "inflows.csv") + +# Stochastic demand (bolivia/demand_scenarios.csv, single line `s,`): +# i.i.d. per-stage multiplicative factor ξ_t ∈ {1−s, 1, 1+s} (P = 1/3 each) on +# every bus's active demand, independent of the inflow noise — the same model +# the SDDP baselines register via sddp/sddp_demand_noise.jl. When the file is +# ABSENT every code path below is bit-identical to the historical trainer. +# Mechanics: scenarios become augmented stage-major vectors [w_t; ξ_t] +# (sample_scenario 3-arg method / augment_scenario), the DE is built with +# demand_spread (p_inflow sized T·(nHyd+1)), the policy encoder observes ξ_t +# (n_extra_uncertainty = 1), and prepare_solve! applies base_demand·ξ_t via +# set_demand! before every solve (training, rollout, and multi-GPU workers). +const DEMAND_SPREAD = load_demand_spread(joinpath(CASE_DIR, "demand_scenarios.csv")) +const DEMAND_NOISE = DEMAND_SPREAD !== nothing +DEMAND_NOISE && @info "Stochastic demand ACTIVE" DEMAND_SPREAD + +function parse_reactive_deficit_cost(raw::AbstractString) + s = lowercase(strip(raw)) + if s == "free" + return nothing + elseif s == "hard" + return Inf + end + cost = parse(Float64, s) + (isnan(cost) || cost < 0) && + throw(ArgumentError("DR_REACTIVE_DEFICIT must be free, hard, or a finite nonnegative cost; got $raw")) + return cost +end + +function value_tag(x) + return replace(replace(string(x), "." => "p"), "-" => "m") +end + +function parse_int_schedule(raw::AbstractString, name::AbstractString) + s = strip(raw) + if isempty(s) || lowercase(s) in ("fixed", "none", "nothing") + return nothing + end + + schedule = Tuple{Int, Int, Int}[] + for item in split(s, r"[,;]") + token = strip(item) + isempty(token) && continue + + parts = split(token, ":") + local lo::Int + local hi::Int + local value::Int + if length(parts) == 3 + lo = parse(Int, strip(parts[1])) + hi = parse(Int, strip(parts[2])) + value = parse(Int, strip(parts[3])) + elseif length(parts) == 2 && occursin("-", parts[1]) + bounds = split(parts[1], "-") + length(bounds) == 2 || + throw(ArgumentError("$name schedule entry '$token' must be lo:hi:value or lo-hi:value")) + lo = parse(Int, strip(bounds[1])) + hi = parse(Int, strip(bounds[2])) + value = parse(Int, strip(parts[2])) + else + throw(ArgumentError("$name schedule entry '$token' must be lo:hi:value or lo-hi:value")) + end + + lo >= 1 || throw(ArgumentError("$name schedule lower bound must be >= 1 in '$token'")) + hi >= lo || throw(ArgumentError("$name schedule upper bound must be >= lower bound in '$token'")) + value >= 1 || throw(ArgumentError("$name schedule value must be >= 1 in '$token'")) + push!(schedule, (lo, hi, value)) + end + + isempty(schedule) && return nothing + sort!(schedule; by = first) + last_hi = 0 + for (lo, hi, _) in schedule + lo > last_hi || throw(ArgumentError("$name schedule has overlapping entries near iteration $lo")) + last_hi = hi + end + return schedule +end + +function schedule_value(schedule, iter::Int, default::Int) + isnothing(schedule) && return default + for (lo, hi, value) in schedule + lo <= iter <= hi && return value + end + return default +end + +function schedule_tag(schedule, prefix::AbstractString) + isnothing(schedule) && return "" + vals = unique(last.(schedule)) + return "-$(prefix)$(join(vals, "_"))" +end + +const ENCODER_LAYERS = parse_layers(get(ENV, "DR_ENCODER_LAYERS", get(ENV, "DR_LAYERS", "128,128"))) +const HEAD_LAYERS = parse_layers(get(ENV, "DR_HEAD_LAYERS", "")) +# Target-head activation. "sigmoid" is the historical default; it cannot +# exactly attain reachable-interval boundaries (where SDDP places ~24% of its +# realized states), so boundary-attaining alternatives are available: +# DR_ACTIVATION = "sigmoid" | "hardsigmoid" | "stretched" +const ACTIVATION = let raw = lowercase(strip(get(ENV, "DR_ACTIVATION", "stretched"))) + raw in ("", "sigmoid") ? sigmoid : + raw == "hardsigmoid" ? hardsigmoidsafe : + raw == "stretched" ? stretchedsigmoid : + throw(ArgumentError("DR_ACTIVATION must be sigmoid, hardsigmoid, or stretched; got $raw")) +end +const NUM_STAGES = parse(Int, get(ENV, "DR_NUM_STAGES", "126")) +const NUM_ROLLOUT_STAGES = parse(Int, get(ENV, "DR_NUM_ROLLOUT_STAGES", "96")) +const NUM_EPOCHS = parse(Int, get(ENV, "DR_NUM_EPOCHS", "80")) +const NUM_BATCHES = parse(Int, get(ENV, "DR_NUM_BATCHES", "100")) +const NUM_TRAIN_PER_BATCH = parse(Int, get(ENV, "DR_NUM_TRAIN_PER_BATCH", "1")) +const NUM_TRAIN_SCHEDULE = parse_int_schedule(get(ENV, "DR_NUM_TRAIN_SCHEDULE", ""), "DR_NUM_TRAIN_SCHEDULE") +const CONTEXT_MODE = canonical_context_mode(get(ENV, "DR_CONTEXT", "")) +const CONTEXT_PERIOD = countlines(INFLOW_FILE) +const CONTEXT_HORIZON = NUM_STAGES +const _base_context = build_stage_context(CONTEXT_MODE, CONTEXT_HORIZON, CONTEXT_PERIOD) +const STAGE_CONTEXT = _base_context +const N_CONTEXT = isnothing(STAGE_CONTEXT) ? 0 : size(STAGE_CONTEXT, 1) +const NUM_EVAL_SCENARIOS = parse(Int, get(ENV, "DR_NUM_EVAL_SCENARIOS", "4")) +const EVAL_SCHEDULE = parse_int_schedule(get(ENV, "DR_EVAL_SCHEDULE", ""), "DR_EVAL_SCHEDULE") +if !isnothing(EVAL_SCHEDULE) && maximum(last.(EVAL_SCHEDULE)) > NUM_EVAL_SCENARIOS + throw(ArgumentError("DR_EVAL_SCHEDULE cannot exceed DR_NUM_EVAL_SCENARIOS=$(NUM_EVAL_SCENARIOS)")) +end +const EVAL_EVERY = parse(Int, get(ENV, "DR_EVAL_EVERY", "50")) +const SAVE_METRIC = lowercase(strip(get(ENV, "DR_SAVE_METRIC", "training"))) +SAVE_METRIC in ("training", "rollout") || + throw(ArgumentError("DR_SAVE_METRIC must be training or rollout; got $SAVE_METRIC")) +const ENABLE_WANDB = parse(Bool, get(ENV, "DR_ENABLE_WANDB", "true")) +# Load Wandb (PythonCall/CondaPkg) ONLY when enabled. With W&B off this avoids the +# CondaPkg "Downloading artifact: pixi" hang on compute nodes with no/slow internet. +ENABLE_WANDB && @eval using Wandb +const LR = parse(Float32, get(ENV, "DR_LR", "0.001")) +const LR_FINAL = parse(Float32, get(ENV, "DR_LR_FINAL", string(LR))) +const LR_WARMUP = parse(Int, get(ENV, "DR_LR_WARMUP", "0")) +const PRE_TRAINED = strip(get(ENV, "DR_PRETRAINED_MODEL", "")) +const HAS_PRETRAINED = !(isempty(PRE_TRAINED) || lowercase(PRE_TRAINED) == "nothing") +const REACTIVE_DEFICIT_RAW = "hard" +const REACTIVE_DEFICIT_COST = Inf +const GRAD_CLIP = parse(Float32, get(ENV, "DR_GRAD_CLIP", "0")) + +const TARGET_PEN_ARG = :auto +const HYDRO_TARGET_PENALTY_MULT = 8.0 +# MAIN: 60 USD/MWh × 100 MVA = 6000 USD/(pu·stage). +const DEFICIT_COST = 6000.0 +const USE_GPU = true +const load_scaler = 0.6 +const qd_scaler = 0.6 +# Parallel-sample training: solve the `num_train_per_batch` per-gradient DEs +# across worker threads, each with its own MadNLP solver bound to its own CUDA +# stream (see train_tsddr! in src/training.jl). Requires JULIA_NUM_THREADS >= +# DR_NUM_WORKERS. Default 1 = historical sequential path (byte-identical). +const NUM_WORKERS = let n = parse(Int, get(ENV, "DR_NUM_WORKERS", "1")) + n >= 1 || throw(ArgumentError("DR_NUM_WORKERS must be >= 1")) + if n > Threads.nthreads() + @warn "DR_NUM_WORKERS=$n exceeds JULIA_NUM_THREADS=$(Threads.nthreads()); capping" + Threads.nthreads() + else + n + end +end + +const ROLLOUT_PARALLEL = parse(Bool, get(ENV, "DR_ROLLOUT_PARALLEL", "false")) + +const MAX_ITER = parse(Int, get(ENV, "DR_MAX_ITER", "9000")) +const SOLVER_KWARGS = (print_level = MadNLP.ERROR, tol = 1e-6, max_iter = MAX_ITER) + +const _CLIP_TAG = GRAD_CLIP > 0 ? "-clip$(Int(GRAD_CLIP))" : "" +const _ENC_TAG = ENCODER_LAYERS == [128, 128] ? "" : "-E$(join(ENCODER_LAYERS, "_"))" +const _HEAD_TAG = isempty(HEAD_LAYERS) ? "-Hlinear" : "-H$(join(HEAD_LAYERS, "_"))" +const _NT_TAG = isnothing(NUM_TRAIN_SCHEDULE) ? + (NUM_TRAIN_PER_BATCH > 1 ? "-nt$(NUM_TRAIN_PER_BATCH)" : "") : + schedule_tag(NUM_TRAIN_SCHEDULE, "ntsch") +const _CTX_TAG = context_run_tag(CONTEXT_MODE) +const _EV_TAG = schedule_tag(EVAL_SCHEDULE, "evsch") +const _SAVE_TAG = SAVE_METRIC == "rollout" ? "-rollout" : "" +const _WARM_TAG = HAS_PRETRAINED ? "-warm" : "" +const _RQ_TAG = REACTIVE_DEFICIT_COST === nothing ? "" : + isinf(Float64(REACTIVE_DEFICIT_COST)) ? "-rqhard" : + "-rq$(value_tag(REACTIVE_DEFICIT_COST))" +const _ACT_TAG = ACTIVATION === sigmoid ? "" : + ACTIVATION === stretchedsigmoid ? "-actstretch" : "-acthardsig" +# Demand-noise tag: runs with stochastic demand are a different experiment +# family (different SP and different policy input width) — mark them. +const _DN_TAG = DEMAND_NOISE ? "-dnoise$(value_tag(DEMAND_SPREAD))" : "" +const RUN_NAME = "$(CASE_NAME)-$(FORM_LABEL)-h$(NUM_STAGES)-r$(NUM_ROLLOUT_STAGES)-deteq-strict-gpu$(_CLIP_TAG)$(_ENC_TAG)$(_HEAD_TAG)$(_NT_TAG)$(_CTX_TAG)$(_EV_TAG)$(_SAVE_TAG)$(_WARM_TAG)$(_RQ_TAG)$(_ACT_TAG)$(_DN_TAG)-$(Dates.format(now(), "yyyymmdd-HHMMSS"))" +const MODEL_DIR = joinpath(CASE_DIR, FORM_LABEL, "models") +mkpath(MODEL_DIR) +const MODEL_PATH = joinpath(MODEL_DIR, RUN_NAME * ".jld2") +# Crash-safety: independent of SaveBest (which only writes on improvement), the +# LATEST policy is checkpointed every DR_SAVE_LATEST_EVERY gradient steps to a +# separate "_latest.jld2" file (overwritten each time). A crash then loses +# at most that many steps, regardless of whether the run had improved. Set to 0 +# to disable. Default 25 keeps every expensive multi-worker run recoverable. +const SAVE_LATEST_EVERY = parse(Int, get(ENV, "DR_SAVE_LATEST_EVERY", "25")) +const LATEST_PATH = joinpath(MODEL_DIR, RUN_NAME * "_latest.jld2") +const TOTAL_ITERS = NUM_EPOCHS * NUM_BATCHES + +function lr_schedule(iter::Int, total_iters::Int) + if iter <= LR_WARMUP + return LR * (0.01f0 + 0.99f0 * Float32(iter) / Float32(max(LR_WARMUP, 1))) + end + ρ = clamp((iter - LR_WARMUP) / max(total_iters - LR_WARMUP, 1), 0.0, 1.0) + return LR_FINAL + 0.5f0 * (LR - LR_FINAL) * (1f0 + cos(Float32(pi * ρ))) +end + +# ── Advanced scheduler: warm-restart LR per nt-stage + event-based nt stepping + go-back ── +# DR_SCHED_MODE = ""(off, iteration-based default) | "warmrestart" | "event". +# warmrestart: nt walks DR_NT_STAGES on fixed DR_STAGE_ITERS boundaries; the cosine LR +# RESTARTS (warmup→peak→floor) at each nt change — replicates the historical +# warmstart trail (each stage its own LR search) inside one job. +# event: nt advances to the next stage only when the smoothed loss PLATEAUS +# (no improvement for DR_SCHED_PATIENCE iters, min-dwell DR_SCHED_MINSTAGE), +# LR restarts on each advance, and RETREATS (×0.7, "go back") when the smoothed +# loss rises >2% above its best — so it keeps descending, never flat/unlearning. +const SCHED_MODE = lowercase(strip(get(ENV, "DR_SCHED_MODE", ""))) +const NT_STAGES = let s = strip(get(ENV, "DR_NT_STAGES", "")); isempty(s) ? Int[] : parse.(Int, split(s, ',')) end +const STAGE_ITERS = parse(Int, get(ENV, "DR_STAGE_ITERS", "1200")) +const SCHED_PATIENCE = parse(Int, get(ENV, "DR_SCHED_PATIENCE", "150")) +const SCHED_MINSTAGE = parse(Int, get(ENV, "DR_SCHED_MINSTAGE", "150")) + +mutable struct SchedState + ema::Float64; best::Float64; plateau::Int + stage::Int; stage_start::Int; lr_retreat::Float64 +end +const SCHED = SchedState(NaN, Inf, 0, 1, 1, 1.0) + +function sched_observe!(loss::Float64) + isfinite(loss) || return + SCHED.ema = isnan(SCHED.ema) ? loss : 0.9 * SCHED.ema + 0.1 * loss + if SCHED.ema < SCHED.best - 1e-6 * abs(SCHED.best) + SCHED.best = SCHED.ema; SCHED.plateau = 0 + else + SCHED.plateau += 1 + end + if isfinite(SCHED.best) && SCHED.ema > SCHED.best * 1.02 # diverging → retreat LR ("go back") + SCHED.lr_retreat = max(SCHED.lr_retreat * 0.7, 0.05) + end +end + +function sched_nt!(iter::Int) + isempty(NT_STAGES) && return NUM_TRAIN_PER_BATCH + if SCHED_MODE == "event" + if SCHED.stage < length(NT_STAGES) && SCHED.plateau >= SCHED_PATIENCE && + (iter - SCHED.stage_start) >= SCHED_MINSTAGE + SCHED.stage += 1; SCHED.stage_start = iter; SCHED.plateau = 0; SCHED.lr_retreat = 1.0 + @info "event-sched: nt → $(NT_STAGES[SCHED.stage]) (LR restart) at iter $iter" + end + else + ns = min(length(NT_STAGES), 1 + div(iter - 1, max(STAGE_ITERS, 1))) + if ns != SCHED.stage + SCHED.stage = ns; SCHED.stage_start = iter; SCHED.lr_retreat = 1.0 + @info "warmrestart: nt → $(NT_STAGES[SCHED.stage]) (LR restart) at iter $iter" + end + end + return NT_STAGES[SCHED.stage] +end + +function sched_lr(iter::Int) + len = SCHED_MODE == "event" ? max(div(TOTAL_ITERS, max(length(NT_STAGES), 1)), 1) : max(STAGE_ITERS, 1) + k = iter - SCHED.stage_start + lr = if k <= LR_WARMUP + LR * (0.01f0 + 0.99f0 * Float32(k) / Float32(max(LR_WARMUP, 1))) + else + ρ = clamp((k - LR_WARMUP) / max(len - LR_WARMUP, 1), 0.0, 1.0) + LR_FINAL + 0.5f0 * (LR - LR_FINAL) * (1f0 + cos(Float32(pi * ρ))) + end + return Float32(lr * SCHED.lr_retreat) +end + +# ── Load data ───────────────────────────────────────────────────────────────── + +@info "Loading power system data..." +power_data = load_power_data(PM_FILE) +@info " nBus=$(power_data.nBus) nGen=$(power_data.nGen)" + +@info "Loading hydro data..." +hydro_data = load_hydro_data(HYDRO_FILE, INFLOW_FILE, power_data; + num_stages = max(NUM_STAGES, NUM_ROLLOUT_STAGES) * 10) +nHyd = hydro_data.nHyd +T = NUM_STAGES +T_ROLLOUT = NUM_ROLLOUT_STAGES +# Per-stage uncertainty width fed to the policy/DE/rollout machinery: +# nHyd inflows, plus the demand factor ξ_t when stochastic demand is active. +N_UNC = nHyd + (DEMAND_NOISE ? 1 : 0) +@info " nHyd=$(nHyd) nScenarios=$(hydro_data.nScenarios) n_uncertainty=$(N_UNC)" + +demand_mat = nothing + +# Reactive demand at qd_scaler × nominal, independent of load_scaler: the +# builder multiplies any given matrix by load_scaler, so divide it back out. +reactive_mat = qd_scaler == load_scaler ? nothing : + repeat(reshape((qd_scaler / load_scaler) .* power_data.default_bus_reactive_demand, + 1, :), T, 1) + +# ── Build ExaModels DE (strict: delta variables fixed to zero) ──────────────── + +resolved_pen = TARGET_PEN_ARG === :auto ? + auto_target_penalty(power_data, hydro_data) : + Float64(TARGET_PEN_ARG) +@info "Auto target penalty: ρ=$(round(resolved_pen; digits=2)) (not used — strict mode)" + +backend = USE_GPU ? (@info "Using GPU backend"; CUDA.CUDABackend()) : + (@info "Using CPU backend"; nothing) + +function _build_de() + build_hydro_de(power_data, hydro_data, T; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = demand_mat, + reactive_demand_matrix = reactive_mat, + load_scaler = load_scaler, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, + # Stochastic demand: sizes p_inflow for [w_t; ξ_t] blocks and stores + # the base demand for prepare_solve!'s ξ_t multiplication (nothing = + # deterministic, bit-identical legacy model). + demand_spread = DEMAND_SPREAD, + ) +end + +# Multi-GPU device assignment for the worker pool. With USE_GPU, distribute the +# NUM_WORKERS solver DEs round-robin across the visible CUDA devices (SLURM sets +# CUDA_VISIBLE_DEVICES, so ndevices() = the GPUs this job was granted). Worker wi +# runs on device WORKER_DEVICES[wi], and its DE is BUILT on that device below so +# the DE arrays and the worker's solver live together. Single GPU or CPU → all +# zeros / nothing (unchanged behavior). Set DR_NUM_WORKERS = workers_per_gpu × +# n_gpus (e.g. 9 on a 3-GPU job packs 3 workers/GPU). +const N_GPU = USE_GPU ? CUDA.ndevices() : 0 +const WORKER_DEVICES = (USE_GPU && N_GPU > 1) ? + [(i - 1) % N_GPU for i in 1:NUM_WORKERS] : nothing +WORKER_DEVICES === nothing || @info "Multi-GPU worker→device map" N_GPU WORKER_DEVICES + +# Build a DE on a specific CUDA device (arrays allocate on the active device). +# Diagnostics to stderr (flushed) so a hang in the multi-GPU pool build is +# localizable in the SLURM log. +function _build_de_on(dev, i) + if dev === nothing + return _build_de() + end + println(stderr, "[pool $i] building DE on CUDA device $dev ..."); flush(stderr) + CUDA.device!(dev) + de = _build_de() + println(stderr, "[pool $i] DE ready on device $dev (current=$(CUDA.device()))"); flush(stderr) + return de +end + +@info "Building strict $(T)-stage ExaModels DE (formulation=$FORMULATION)..." +prob = _build_de() # metadata DE on the default device (device 0) + +# Multi-GPU: workers build their OWN DE in-task (a main-task-built DE deadlocks +# on the first cross-task solve). `worker_de_builder(wi)` runs inside worker +# `wi`'s task AFTER it has bound its device, so the DE lands on the right GPU. +# Single-GPU / CPU: fall back to the pre-built pool (all on device 0). +worker_de_builder = nothing +problem_pool = nothing +if WORKER_DEVICES !== nothing + worker_de_builder = function (wi) + de = _build_de() # current device is set by the worker's CUDA.device! + (de, de.p_x0, de.p_target, de.p_inflow) + end + @info " Multi-GPU: each of $NUM_WORKERS workers builds its DE in-task" n_gpu=N_GPU +else + problem_pool = [(prob, prob.p_x0, prob.p_target, prob.p_inflow)] + for _ in 2:NUM_WORKERS + p = _build_de() + push!(problem_pool, (p, p.p_x0, p.p_target, p.p_inflow)) + end + @info " Pool ready: $(NUM_WORKERS) DE instances on the default device" +end + +x0_init = Float32.([clamp(hydro_data.initial_volumes[r], + hydro_data.units[r].min_vol, + hydro_data.units[r].max_vol) + for r in 1:nHyd]) +target_lower = Float32.([h.min_vol for h in hydro_data.units]) +target_upper = Float32.([h.max_vol for h in hydro_data.units]) + +# ── Policy (HydroReachablePolicy — one-stage reachable sigmoid bounds) ──────── + +Random.seed!(42) +base_policy = hydro_reachable_policy(hydro_data, ENCODER_LAYERS; + activation = ACTIVATION, + encoder_type = Flux.LSTM, + combiner_layers = HEAD_LAYERS, + n_context = N_CONTEXT, + # Demand noise: the encoder additionally + # observes ξ_t (stage-t revealed demand), + # matching SDDP whose stage subproblem + # sees the realized demand atom. + n_extra_uncertainty = DEMAND_NOISE ? 1 : 0) +policy = isnothing(STAGE_CONTEXT) ? base_policy : ContextualPolicy(base_policy, STAGE_CONTEXT) + +if HAS_PRETRAINED + @info "Loading pre-trained policy checkpoint" PRE_TRAINED + load_stateconditioned_policy!(policy, JLD2.load(PRE_TRAINED, "model_state")) + Flux.reset!(policy) +end + +""" + rollout_reachable_targets(policy, x0, w_flat, T, nHyd) -> Vector{Float64} + +Roll out a [`HydroReachablePolicy`] target trajectory before solving the strict +regular deterministic equivalent. + +The regular DE receives an external target vector, so it cannot query the policy +inside the NLP. This helper constructs that vector in a way that preserves +strict-mode feasibility: it starts from the known feasible initial state `x0`, +feeds `[u_t; previous_target]` to the policy, and stores each reachable target as +the next previous state. By induction, every target in the returned trajectory is +reachable from the prior target under the sampled inflow path. + +# Arguments +- `policy`: reachable hydro policy with input `[uncertainty; previous_state]`. +- `x0`: initial reservoir state. +- `w_flat`: stage-major flat uncertainty vector of length `T * n_uncertainty` + (`n_uncertainty = nHyd` historically, `nHyd + 1` with stochastic demand — + the per-stage stride is derived from `length(w_flat) ÷ T`, so both layouts + work; the policy slices the physical inflow internally). +- `T::Int`: number of stages. +- `nHyd::Int`: number of hydro reservoir state components (kept for call-site + compatibility; the stride no longer depends on it). + +# Returns +- `Vector{Float64}`: stage-major target trajectory suitable for + `ExaModels.set_parameter!(prob.core, prob.p_target, targets)`. + +# Examples +```julia +targets = rollout_reachable_targets(policy, x0_init, mean_inflow(hydro_data, T), T, nHyd) +``` +""" +function rollout_reachable_targets(policy, x0, w_flat, T, nHyd) + Flux.reset!(policy) + prev = x0 + # Per-stage uncertainty stride derived from the vector itself (nHyd, or + # nHyd+1 when the demand factor ξ_t is appended to each stage block). + nu = length(w_flat) ÷ T + targets = Vector{Vector{Float32}}(undef, T) + for t in 1:T + # Full stage-t uncertainty block [w_t] or [w_t; ξ_t]. + wt = Float32.(view(w_flat, ((t - 1) * nu + 1):(t * nu))) + target = policy(vcat(wt, prev)) + targets[t] = Float32.(target) + prev = targets[t] + end + return Float64.(vcat(targets...)) +end + +# ── Smoke test ──────────────────────────────────────────────────────────────── + +# Mean-inflow smoke scenario; with demand noise, append the neutral factor +# ξ_t = 1 to every stage (base demand) so the augmented layout is exercised. +w_mean = DEMAND_NOISE ? augment_scenario(mean_inflow(hydro_data, T), ones(T)) : + mean_inflow(hydro_data, T) +xhat_mean = rollout_reachable_targets(policy, x0_init, w_mean, T, nHyd) +ExaModels.set_parameter!(prob.core, prob.p_x0, x0_init) +ExaModels.set_parameter!(prob.core, prob.p_inflow, w_mean) +ExaModels.set_parameter!(prob.core, prob.p_target, xhat_mean) +prepare_solve!(prob, x0_init, w_mean, xhat_mean) +@info "Smoke test: solving strict DE with mean inflows and reachable policy targets..." +result0 = MadNLP.madnlp(prob.model; SOLVER_KWARGS..., print_level = MadNLP.WARN) +@info " Status: $(result0.status) Objective: $(round(result0.objective; digits=4))" +isfinite(result0.objective) || error("Smoke test returned non-finite objective") +solve_succeeded(result0) || @warn "Smoke test did not fully converge; proceeding anyway" + +Flux.reset!(policy) + +if USE_GPU + policy = policy isa ContextualPolicy ? + ContextualPolicy(CUDA.cu(policy.policy), CUDA.cu(policy.context)) : + CUDA.cu(policy) + Flux.reset!(policy) + x0_init = CUDA.cu(x0_init) + @info "Policy and x0 moved to GPU" +end + +@info "Strict Exa training config" RUN_NAME NUM_STAGES NUM_ROLLOUT_STAGES NUM_EPOCHS NUM_BATCHES NUM_TRAIN_PER_BATCH NUM_TRAIN_SCHEDULE CONTEXT_MODE CONTEXT_PERIOD CONTEXT_HORIZON N_CONTEXT NUM_EVAL_SCENARIOS EVAL_SCHEDULE EVAL_EVERY SAVE_METRIC LR LR_FINAL LR_WARMUP PRE_TRAINED REACTIVE_DEFICIT_RAW REACTIVE_DEFICIT_COST GRAD_CLIP MAX_ITER + +function checkpoint_policy_state(m) + if m isa ContextualPolicy + return Flux.state(ContextualPolicy(cpu(m.policy), Array(m.context))) + end + return Flux.state(cpu(m)) +end + +# ── W&B logging ─────────────────────────────────────────────────────────────── + +lg = ENABLE_WANDB ? WandbLogger( + project = "RL", + name = RUN_NAME, + save_code = false, + config = Dict( + "case" => CASE_NAME, + "formulation" => FORM_LABEL, + "method" => "deteq-strict", + "num_stages" => T, + "num_rollout_stages" => T_ROLLOUT, + "encoder_layers" => ENCODER_LAYERS, + "head_layers" => HEAD_LAYERS, + "activation" => string(ACTIVATION), + "target_penalty" => "strict (disabled)", + "deficit_cost" => DEFICIT_COST, + "reactive_deficit" => REACTIVE_DEFICIT_RAW, + "reactive_deficit_cost" => string(REACTIVE_DEFICIT_COST), + "num_epochs" => NUM_EPOCHS, + "num_batches" => NUM_BATCHES, + "num_train_per_batch" => NUM_TRAIN_PER_BATCH, + "num_train_schedule" => string(something(NUM_TRAIN_SCHEDULE, "fixed")), + "context_mode" => isempty(CONTEXT_MODE) ? "none" : CONTEXT_MODE, + "context_period" => CONTEXT_PERIOD, + "context_horizon" => CONTEXT_HORIZON, + "n_context" => N_CONTEXT, + "num_eval_scenarios" => NUM_EVAL_SCENARIOS, + "eval_schedule" => string(something(EVAL_SCHEDULE, "fixed")), + "eval_every" => EVAL_EVERY, + "save_metric" => SAVE_METRIC, + "lr" => LR, + "lr_final" => LR_FINAL, + "lr_warmup" => LR_WARMUP, + "pre_trained_model" => PRE_TRAINED, + "grad_clip" => GRAD_CLIP, + "backend" => USE_GPU ? "GPU" : "CPU", + "load_scaler" => load_scaler, + # Demand-noise provenance: "none" = deterministic demand. + "demand_spread" => DEMAND_NOISE ? DEMAND_SPREAD : "none", + "strict_targets" => true, + "policy_type" => N_CONTEXT == 0 ? "HydroReachablePolicy" : "ContextualPolicy{HydroReachablePolicy}", + "num_workers" => NUM_WORKERS, + ), +) : nothing +ENABLE_WANDB || @info "W&B disabled (DR_ENABLE_WANDB=false)" + +# ── Training ────────────────────────────────────────────────────────────────── + +Random.seed!(8788) + +best_obj = Inf +epoch_losses = Float64[] + +stage_demand = demand_mat === nothing ? nothing : demand_mat[1:1, :] +function _build_rollout_de() + build_hydro_de(power_data, hydro_data, 1; + backend = backend, + float_type = Float64, + formulation = FORMULATION, + target_penalty = TARGET_PEN_ARG, + deficit_cost = DEFICIT_COST, + demand_matrix = stage_demand, + reactive_demand_matrix = reactive_mat === nothing ? nothing : reactive_mat[1:1, :], + load_scaler = load_scaler, + strict_targets = true, + reactive_deficit_cost = REACTIVE_DEFICIT_COST, + # Stochastic demand: the 1-stage problem also carries a base_demand + # row; set_hydro_rollout_stage! refreshes it per stage and + # prepare_solve! applies the ξ_t carried in the stage's wt block. + demand_spread = DEMAND_SPREAD, + ) +end +rollout_prob = _build_rollout_de() +n_rollout_pool = max(NUM_WORKERS, NUM_EVAL_SCENARIOS) +rollout_pool = ROLLOUT_PARALLEL ? [_build_rollout_de() for _ in 1:n_rollout_pool] : [] +@info "Rollout evaluation: $(ROLLOUT_PARALLEL ? "parallel ($(n_rollout_pool) stage-problem copies)" : "sequential")" + +function set_hydro_rollout_stage!(stage_prob, state_in, wt, target, stage) + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_x0, state_in) + # wt is the full stage uncertainty block ([w_t] or [w_t; ξ_t]); the stage + # problem's p_inflow is sized to match (n_uncertainty entries). + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_inflow, wt) + if demand_mat !== nothing + if DEMAND_NOISE + # Stochastic demand: refresh the 1-stage problem's BASE demand row + # in place; prepare_solve! below multiplies it by the ξ_t carried + # in wt's last entry and writes the product via set_demand!. + stage_prob.base_demand[1, :] .= load_scaler .* @view demand_mat[stage, :] + else + # Deterministic demand: historical direct write (bit-identical). + set_demand!(stage_prob, load_scaler .* demand_mat[stage:stage, :]) + end + end + ExaModels.set_parameter!(stage_prob.core, stage_prob.p_target, target) + prepare_solve!(stage_prob, state_in, wt, target) + return stage_prob +end + +hydro_realized_state(stage_prob, result) = + hydro_solution(stage_prob, result).reservoir[:, end] + +const _min_vols = Float64.([h.min_vol for h in hydro_data.units]) +const _max_vols = Float64.([h.max_vol for h in hydro_data.units]) +const _min_vols_dev = USE_GPU ? CUDA.cu(_min_vols) : _min_vols +const _max_vols_dev = USE_GPU ? CUDA.cu(_max_vols) : _max_vols + +hydro_objective_no_target_penalty(stage_prob, result) = result.objective + +# Held-out evaluation scenarios. Two modes: +# +# 1. DR_EVAL_PROTOCOL_IDS set (comma-separated column ids of the seeded paired +# protocol): the eval set is those exact columns of the 126×500 protocol +# matrix (StableRNG seed 20260706 — identical generation to the paired +# evaluation scripts). With the representative subset found by searching +# 300k candidate subsets against 7 evaluated policies — +# ids 2,39,81,119,130,156,200,206,378,493 (subset seed 77142) — the +# 10-scenario training-time evaluation tracks the full 500-scenario paired +# mean within ~75 cost units and preserves paired differences vs SDDP +# within ~35, so SaveBest selects on (a faithful proxy of) the deployment +# metric. +# 2. Unset (historical): random draws from the inflow process, seed 8789. +# NOTE: arbitrary small draws carry offsets of hundreds-to-thousands of +# cost units vs the paired protocol; never compare their values across +# runs with different eval sets. +const EVAL_PROTOCOL_IDS = let raw = strip(get(ENV, "DR_EVAL_PROTOCOL_IDS", "")) + isempty(raw) ? Int[] : [parse(Int, strip(x)) for x in split(raw, ",")] +end + +""" + protocol_eval_scenario(hydro_data, T, protocol_indices, s) -> Vector{Float64} + +Flat `T × nHyd` inflow vector for paired-protocol scenario column `s`: +stage `t` realizes joint inflow scenario `protocol_indices[t, s]`, with the +cyclic raw-row mapping shared by every paired evaluation script. +""" +function protocol_eval_scenario(hydro_data::HydroData, T::Int, protocol_indices, s::Int) + nHyd = hydro_data.nHyd + w = Vector{Float64}(undef, T * nHyd) + for t in 1:T + t_row = mod1(t, hydro_data.nStagesSample) + j = protocol_indices[t, s] + for r in 1:nHyd + w[(t-1)*nHyd + r] = hydro_data.scenario_inflows[r][t_row, j] + end + end + return w +end + +# Demand-noise augmentation of a protocol column: pair the column's inflows +# with its SEEDED demand path (StableRNG(DEMAND_NOISE_SEED + column) — the +# identical path eval_paired_exa_strict.jl draws for that column). With noise +# off this is the identity, so the historical eval set is unchanged. +_augment_protocol(w, col, T) = DEMAND_NOISE ? + augment_scenario(w, protocol_demand_factors(DEMAND_SPREAD, T, col)) : w + +eval_scenarios = if isempty(EVAL_PROTOCOL_IDS) + Random.seed!(8789) + # Random draws: the 3-arg sampler additionally draws i.i.d. ξ_t from the + # SAME seeded global stream, so the eval set stays reproducible. + [DEMAND_NOISE ? sample_scenario(hydro_data, T_ROLLOUT, DEMAND_SPREAD) : + sample_scenario(hydro_data, T_ROLLOUT) + for _ in 1:NUM_EVAL_SCENARIOS] +else + # Same fixed-shape generation as the paired protocol (126 rows × 500 + # columns; see load_hydropowermodels.jl in the MAIN repo). + protocol_indices = rand(StableRNG(20260706), 1:hydro_data.nScenarios, 126, 500) + @assert T_ROLLOUT <= 126 && all(1 .<= EVAL_PROTOCOL_IDS .<= 500) + @assert length(EVAL_PROTOCOL_IDS) == NUM_EVAL_SCENARIOS "DR_NUM_EVAL_SCENARIOS must match the id count" + @info "Eval set = paired-protocol columns $(EVAL_PROTOCOL_IDS)" + [_augment_protocol(protocol_eval_scenario(hydro_data, T_ROLLOUT, protocol_indices, s), s, T_ROLLOUT) + for s in EVAL_PROTOCOL_IDS] +end +# ── Training sampler ────────────────────────────────────────────────────────── +# Default: fresh random inflow draws (genuine SAA). When DR_TRAIN_PROTOCOL_ALL +# is set, the sampler instead cycles deterministically through the exact 500 +# paired-protocol columns (StableRNG(20260706), full T stages). With +# num_train_per_batch=500 and num_batches a multiple of 1, every gradient step +# is a full-batch step over ALL 500 evaluation scenarios — the "cheating" upper +# bound: can gradient descent directly on the test set beat SDDP? If it can't, +# the policy class is the wall. +const TRAIN_PROTOCOL_ALL = lowercase(strip(get(ENV, "DR_TRAIN_PROTOCOL_ALL", ""))) in ("1", "true", "yes") +train_sampler = if TRAIN_PROTOCOL_ALL + train_protocol_indices = rand(StableRNG(20260706), 1:hydro_data.nScenarios, 126, 500) + @assert T <= 126 + # With demand noise, each protocol column is paired with its seeded demand + # path (prefix-consistent with the T_ROLLOUT-length eval draws — see + # sample_demand_factors' sequential-draw prefix property). + protocol_cols = [_augment_protocol(protocol_eval_scenario(hydro_data, T, train_protocol_indices, s), s, T) + for s in 1:500] + @info "TRAINING on the exact 500 paired-protocol scenarios (cheating upper-bound test)" NUM_TRAIN_PER_BATCH + cyc = Ref(0) + () -> (cyc[] = cyc[] % 500 + 1; protocol_cols[cyc[]]) +elseif DEMAND_NOISE + # Genuine SAA over the PRODUCT distribution: fresh inflow draw + fresh + # i.i.d. per-stage demand factors, returned as augmented [w_t; ξ_t] blocks. + () -> sample_scenario(hydro_data, T, DEMAND_SPREAD) +else + () -> sample_scenario(hydro_data, T) +end + +rollout_evaluation = RolloutEvaluation( + rollout_prob, + x0_init, + eval_scenarios; + horizon = T_ROLLOUT, + # Per-stage uncertainty width: nHyd, or nHyd+1 with demand noise (the + # stage callback then receives the full [w_t; ξ_t] block as wt). + n_uncertainty = N_UNC, + set_stage_parameters! = set_hydro_rollout_stage!, + realized_state = hydro_realized_state, + objective_no_target_penalty = hydro_objective_no_target_penalty, + madnlp_kwargs = SOLVER_KWARGS, + warmstart = false, + stride = EVAL_EVERY, + policy_state = :realized, + stage_problem_pool = rollout_pool, + retry_on_failure = true, + active_scenarios = NUM_EVAL_SCENARIOS, + state_bounds = (_min_vols_dev, _max_vols_dev), +) + +# ── Two-stage SaveBest verification (guards against overfitting the fixed eval +# set). The fixed rep-10 eval is cheap but a policy can overfit to those exact +# 10 scenarios, so SaveBest-on-10 may select an overfit iterate. When +# DR_VERIFY_SCENARIOS > 0, a checkpoint that beats the best on the fixed set is +# only ACCEPTED if it also beats the incumbent best on N FRESHLY-sampled random +# scenarios — re-drawn every trigger, so a policy cannot overfit to them. The +# incumbent is re-evaluated on the SAME fresh draw for a paired comparison. +const VERIFY_SCENARIOS = parse(Int, get(ENV, "DR_VERIFY_SCENARIOS", "0")) +verify_evaluation = if VERIFY_SCENARIOS > 0 && SAVE_METRIC == "rollout" + @info "Two-stage SaveBest: verify accepts on $VERIFY_SCENARIOS fresh random scenarios" + RolloutEvaluation( + _build_rollout_de(), x0_init, + # Fresh verification draws come from the same (augmented, when demand + # noise is on) sampler family as training. + [DEMAND_NOISE ? sample_scenario(hydro_data, T_ROLLOUT, DEMAND_SPREAD) : + sample_scenario(hydro_data, T_ROLLOUT) + for _ in 1:VERIFY_SCENARIOS]; + horizon = T_ROLLOUT, n_uncertainty = N_UNC, + set_stage_parameters! = set_hydro_rollout_stage!, + realized_state = hydro_realized_state, + objective_no_target_penalty = hydro_objective_no_target_penalty, + madnlp_kwargs = SOLVER_KWARGS, warmstart = false, stride = 1, + policy_state = :realized, stage_problem_pool = [], retry_on_failure = true, + active_scenarios = VERIFY_SCENARIOS, + state_bounds = (_min_vols_dev, _max_vols_dev), + ) +else + nothing +end +best_verify_snapshot = Ref{Any}(nothing) # deepcopy of the accepted-best policy + +# Accept `model` as the new best? With verification off, the rep-10 gate that +# already fired is sufficient (return true). With it on, draw fresh scenarios +# and require `model` to beat the incumbent snapshot on them (all must solve). +function verified_improvement(model) + verify_evaluation === nothing && return true + # Fresh draws per trigger; with demand noise the 3-arg sampler pairs each + # fresh inflow path with fresh i.i.d. demand factors ([w_t; ξ_t] blocks). + verify_evaluation.scenarios = + [DEMAND_NOISE ? sample_scenario(hydro_data, T_ROLLOUT, DEMAND_SPREAD) : + sample_scenario(hydro_data, T_ROLLOUT) + for _ in 1:VERIFY_SCENARIOS] + verify_evaluation(1, model) + verify_evaluation.last_n_ok == VERIFY_SCENARIOS || return false + cand = verify_evaluation.last_objective_no_target_penalty + inc = if best_verify_snapshot[] === nothing + Inf + else + verify_evaluation(1, best_verify_snapshot[]) # SAME fresh scenarios + verify_evaluation.last_n_ok == VERIFY_SCENARIOS ? + verify_evaluation.last_objective_no_target_penalty : Inf + end + accept = cand < inc + accept && (best_verify_snapshot[] = deepcopy(model)) + @info " verify on $VERIFY_SCENARIOS fresh: cand=$(round(cand; digits=1)) inc=$(inc == Inf ? "none" : round(inc; digits=1)) => $(accept ? "ACCEPT" : "reject (overfit to fixed set)")" + return accept +end + +current_num_train = Ref(NUM_TRAIN_PER_BATCH) +current_eval_scenarios = Ref(schedule_value(EVAL_SCHEDULE, 1, NUM_EVAL_SCENARIOS)) +rollout_evaluation.active_scenarios = current_eval_scenarios[] + +# The initial rollout eval only sets the checkpoint baseline, and it runs +# SEQUENTIALLY on device 0 (a rollout is 96 dependent stages, one scenario at a +# time) — so it blocks the multi-GPU training start with pure single-GPU work. +# Skip it by default: best_obj stays Inf and the first periodic eval sets the +# baseline, letting the 12 workers engage all GPUs immediately. +const SKIP_INITIAL_EVAL = parse(Bool, get(ENV, "DR_SKIP_INITIAL_EVAL", "true")) +if SAVE_METRIC == "rollout" && !SKIP_INITIAL_EVAL + rollout_evaluation(EVAL_EVERY, policy) + best_obj = rollout_evaluation.last_objective_no_target_penalty + @info "Initial rollout evaluation (checkpoint baseline)" best_obj rollout_evaluation.last_violation_share rollout_evaluation.last_n_ok +elseif SAVE_METRIC == "rollout" + @info "Skipping initial rollout eval (DR_SKIP_INITIAL_EVAL=true) — training starts immediately on all GPUs" +end + +Random.seed!(8788) + +last_batch_stats = Ref(Dict{String, Any}()) + +function _merge_batch_stats!(metrics, stats) + isempty(stats) && return metrics + metrics["metrics/train_n_ok"] = get(stats, "n_ok", 0) + metrics["metrics/train_n_total"] = get(stats, "n_total", 0) + metrics["metrics/train_success_share"] = + get(stats, "n_total", 0) == 0 ? NaN : get(stats, "n_ok", 0) / get(stats, "n_total", 0) + for (k, v) in get(stats, "status_counts", Dict{String, Int}()) + metrics["metrics/train_status/$k"] = v + end + for (k, v) in get(stats, "failure_counts", Dict{String, Int}()) + metrics["metrics/train_failure/$k"] = v + end + for (k, v) in get(stats, "retry_counts", Dict{String, Int}()) + metrics["metrics/train_retry/$k"] = v + end + return metrics +end + +train_tsddr( + policy, + x0_init, + prob, + prob.p_x0, + prob.p_target, + prob.p_inflow, + train_sampler; + num_batches = TOTAL_ITERS, + num_train_per_batch = NUM_TRAIN_PER_BATCH, + optimizer = GRAD_CLIP > 0 ? + Flux.Optimisers.OptimiserChain( + Flux.Optimisers.ClipGrad(GRAD_CLIP), + Flux.Adam(LR), + ) : Flux.Adam(LR), + madnlp_kwargs = SOLVER_KWARGS, + warmstart = true, + problem_pool = problem_pool, + worker_devices = WORKER_DEVICES, + worker_problem_builder = worker_de_builder, + batch_diagnostics = (iter, stats) -> begin + last_batch_stats[] = stats + n_ok = get(stats, "n_ok", 0) + n_total = get(stats, "n_total", 0) + if n_ok < n_total + @warn "Training solve failures at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) failure_counts=get(stats, "failure_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + elseif iter % 10 == 0 + @info "Training solve status at iter $iter" n_ok n_total status_counts=get(stats, "status_counts", nothing) retry_counts=get(stats, "retry_counts", nothing) + end + end, + adjust_hyperparameters = SCHED_MODE != "" ? + ((iter, opt_state, n) -> begin + Flux.Optimisers.adjust!(opt_state, sched_lr(iter)) + n_next = sched_nt!(iter) + if n_next != current_num_train[] + current_num_train[] = n_next + @info "num_train_per_batch → $n_next" + end + n_eval = schedule_value(EVAL_SCHEDULE, iter, NUM_EVAL_SCENARIOS) + if n_eval != current_eval_scenarios[] + current_eval_scenarios[] = n_eval + rollout_evaluation.active_scenarios = n_eval + @info "rollout active scenarios → $n_eval" + end + n_next + end) : + (LR_WARMUP == 0 && LR_FINAL == LR) ? + ((iter, opt_state, n) -> begin + n_next = schedule_value(NUM_TRAIN_SCHEDULE, iter, n) + if n_next != current_num_train[] + current_num_train[] = n_next + @info "num_train_per_batch → $n_next" + end + n_eval = schedule_value(EVAL_SCHEDULE, iter, NUM_EVAL_SCENARIOS) + if n_eval != current_eval_scenarios[] + current_eval_scenarios[] = n_eval + rollout_evaluation.active_scenarios = n_eval + @info "rollout active scenarios → $n_eval" + end + n_next + end) : + ((iter, opt_state, n) -> begin + Flux.Optimisers.adjust!(opt_state, lr_schedule(iter, TOTAL_ITERS)) + n_next = schedule_value(NUM_TRAIN_SCHEDULE, iter, n) + if n_next != current_num_train[] + current_num_train[] = n_next + @info "num_train_per_batch → $n_next" + end + n_eval = schedule_value(EVAL_SCHEDULE, iter, NUM_EVAL_SCENARIOS) + if n_eval != current_eval_scenarios[] + current_eval_scenarios[] = n_eval + rollout_evaluation.active_scenarios = n_eval + @info "rollout active scenarios → $n_eval" + end + n_next + end), + record_loss = (iter, m, loss, tag) -> begin + SCHED_MODE == "" || sched_observe!(Float64(loss)) + metrics = Dict{String, Any}( + tag => loss, + "batch" => iter, + "metrics/lr" => lr_schedule(iter, TOTAL_ITERS), + "metrics/num_train_per_batch" => current_num_train[], + "metrics/active_eval_scenarios" => current_eval_scenarios[], + ) + _merge_batch_stats!(metrics, last_batch_stats[]) + isfinite(loss) && push!(epoch_losses, loss) + + if iter % EVAL_EVERY == 0 + rollout_evaluation(iter, m) + metrics["metrics/rollout_objective_no_target_penalty"] = + rollout_evaluation.last_objective_no_target_penalty + metrics["metrics/rollout_target_violation_share"] = + rollout_evaluation.last_violation_share + metrics["metrics/rollout_n_ok"] = + rollout_evaluation.last_n_ok + if SAVE_METRIC == "rollout" + rollout_score = rollout_evaluation.last_objective_no_target_penalty + # Honest selection: the rollout mean is over the scenarios that + # SOLVED (total / n_ok), so a policy that fails one expensive + # scenario gets a fake bonus of hundreds of cost units. Only + # trust evals where every active scenario succeeded. + if rollout_evaluation.last_n_ok == current_eval_scenarios[] && + isfinite(rollout_score) && rollout_score < best_obj && + verified_improvement(m) # second-stage fresh-scenario gate + global best_obj = rollout_score + jldsave(MODEL_PATH; model_state = checkpoint_policy_state(m)) + @info " -> New best rollout: $(round(rollout_score; digits=4)) -- saved $MODEL_PATH" + end + end + end + + # Crash-safety: overwrite the "_latest" checkpoint every N steps, + # independent of improvement, so a failure loses at most N steps. + if SAVE_LATEST_EVERY > 0 && iter % SAVE_LATEST_EVERY == 0 + jldsave(LATEST_PATH; model_state = checkpoint_policy_state(m)) + @info " latest checkpoint @ iter $iter -> $LATEST_PATH" + end + + batch_in_epoch = (iter - 1) % NUM_BATCHES + 1 + if batch_in_epoch == NUM_BATCHES + epoch = (iter - 1) ÷ NUM_BATCHES + 1 + mean_loss = isempty(epoch_losses) ? NaN : mean(epoch_losses) + n_ok = length(epoch_losses) + empty!(epoch_losses) + lg === nothing || Wandb.log(lg, Dict("metrics/epoch_objective" => mean_loss, "epoch" => epoch)) + @info "Epoch $epoch/$NUM_EPOCHS mean=$(round(mean_loss; digits=2)) ok=$n_ok/$NUM_BATCHES" + if SAVE_METRIC == "training" && isfinite(mean_loss) && mean_loss < best_obj + global best_obj = mean_loss + jldsave(MODEL_PATH; model_state = checkpoint_policy_state(m)) + @info " → New best: $(round(mean_loss; digits=4)) — saved $MODEL_PATH" + end + end + lg === nothing || Wandb.log(lg, metrics) + return false + end, +) + +lg === nothing || close(lg) +@info "Done. Best model saved to: $(MODEL_PATH)" diff --git a/src/DecisionRulesExa.jl b/src/DecisionRulesExa.jl index 4e23112..5c6c17a 100644 --- a/src/DecisionRulesExa.jl +++ b/src/DecisionRulesExa.jl @@ -1,7 +1,39 @@ +""" + DecisionRulesExa + +GPU-accelerated companion to DecisionRules.jl for Two-Stage Deep Decision Rules +(TS-DDR) training with ExaModels and MadNLP. + +DecisionRulesExa implements the same target-projection workflow as +DecisionRules.jl: + +1. a policy predicts target states, +2. an NLP projects those targets onto the feasible set, and +3. target-constraint multipliers provide the policy-gradient signal. + +The package formulates inner optimization problems as `ExaModels.ExaModel` +instances solved by MadNLP, enabling GPU-native solves and warm-started repeated +training solves. + +# Main Types +- [`DeterministicEquivalentProblem`](@ref): deterministic equivalent with + explicit target parameters. +- [`EmbeddedDeterministicEquivalentProblem`](@ref): deterministic equivalent + whose target policy is embedded with `VectorNonlinearOracle`. +- [`StateConditionedPolicy`](@ref): recurrent policy for sequential target + rollout. +- [`MLPPolicy`](@ref): stateless policy for full-horizon target prediction. + +# Main Training APIs +- [`train_tsddr`](@ref): open-loop target-parameter training. +- [`train_tsddr_embedded`](@ref): embedded-policy training. +- [`rollout_tsddr`](@ref): stage-wise deployment-style evaluation. +""" module DecisionRulesExa using ExaModels using MadNLP +using CUDA using NLPModels using LinearAlgebra using Random @@ -12,6 +44,7 @@ using ChainRulesCore include("utils.jl") include("deterministic_equivalent.jl") +include("embedded_deterministic_equivalent.jl") include("policy.jl") include("critic_control_variate.jl") include("training.jl") @@ -39,9 +72,23 @@ export # Policies MLPPolicy, StateConditionedPolicy, + ContextualPolicy, + context_at, + stage_phase_context, + vcat_contexts, + ConstantStatePolicy, + FixedOutputPolicy, + bounded_state_policy, + load_stateconditioned_policy!, + + # Embedded-NN deterministic equivalent + EmbeddedDeterministicEquivalentProblem, + build_embedded_deterministic_equivalent, + invalidate_policy_cache!, # Training solve_succeeded, + prepare_solve!, materialize_tangent, _all_finite_gradient, AbstractCriticControlVariate, @@ -60,6 +107,7 @@ export critic_samples_from_evaluation, simulate_tsddr, train_tsddr, + train_tsddr_embedded, # Stage-wise rollout evaluation rollout_tsddr, diff --git a/src/critic_control_variate.jl b/src/critic_control_variate.jl index a1953d0..0341984 100644 --- a/src/critic_control_variate.jl +++ b/src/critic_control_variate.jl @@ -9,32 +9,82 @@ abstract type AbstractCriticTrainingTarget end """ DeterministicEquivalentCriticTarget() -Train critic value targets from the full deterministic-equivalent objective. -This is useful for ablations and for pure DE control-variate experiments. +Select critic value targets from the full deterministic-equivalent training +objective. + +# Returns +- `DeterministicEquivalentCriticTarget`: a target selector for deterministic- + equivalent critic supervision. + +# Notes +Use this target for ablations or for control-variate experiments tied to the +training surrogate rather than to deployed rollout performance. """ struct DeterministicEquivalentCriticTarget <: AbstractCriticTrainingTarget end """ - RolloutCriticTarget(stage_problem; kwargs...) - -Train critic value targets from stage-wise rollout evaluation. This is the -preferred target when the critic is meant to guide convergence of the deployed -rollout objective rather than the deterministic-equivalent surrogate. - -Required keyword callbacks match `rollout_tsddr`: -- `set_stage_parameters!` -- `realized_state` - -By default `policy_state = :target`, matching the differentiable target -recurrence used by the actor. Set `policy_state = :realized` to train on -closed-loop realized-state rollout targets. - -By default `objective_value = :objective`, so critic value targets include the -same target-penalty contribution that appears in the dual actor signal. Set -`objective_value = :objective_no_target_penalty` to train on the rollout -objective with target-slack penalties removed. + RolloutCriticTarget( + stage_problem; + horizon::Int, + n_uncertainty::Int, + set_stage_parameters!::Function, + realized_state::Function, + objective_no_target_penalty::Function = (prob, result) -> result.objective, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + policy_state::Symbol = :target, + reuse_solver::Bool = false, + objective_value::Symbol = :objective, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, + ) -> RolloutCriticTarget + +Select critic value targets from stage-wise rollout evaluation. + +# Arguments +- `stage_problem`: single-stage optimization problem template used by + [`rollout_tsddr`](@ref). + +# Keywords +- `horizon::Int`: number of rollout stages. +- `n_uncertainty::Int`: dimension of each per-stage uncertainty slice. +- `set_stage_parameters!::Function`: callback that writes state, + uncertainty, target, and stage index into `stage_problem`. +- `realized_state::Function`: callback that extracts the realized next state + from a solved stage. +- `objective_no_target_penalty::Function`: callback returning the stage + objective with target-tracking penalties removed. +- `madnlp_kwargs`: keyword arguments forwarded to the MadNLP solver wrapper. +- `warmstart::Bool`: whether rollout solves should warm-start from previous + stage information. +- `policy_state::Symbol`: `:target` trains against the differentiable target + recurrence; `:realized` trains against closed-loop realized states. +- `reuse_solver::Bool`: whether rollout evaluation may reuse one solver + object across stages. +- `objective_value::Symbol`: `:objective` uses the full rollout objective; + `:objective_no_target_penalty` removes target-slack penalties. +- `state_bounds`: optional `(lower, upper)` projection bounds for realized + states. +- `project_state`: optional custom projection callback for realized states. +- `retry_on_failure::Bool`: whether failed stage solves should be retried with + a cold-start solver. + +# Returns +- `RolloutCriticTarget`: a target selector carrying the rollout callbacks and + options. + +# Throws +- Throws an error if `policy_state` is not `:target` or `:realized`. +- Throws an error if `objective_value` is not `:objective` or + `:objective_no_target_penalty`. + +# Notes +This is the preferred target when the critic is meant to guide convergence of +the deployed rollout objective rather than the deterministic-equivalent +surrogate. """ -struct RolloutCriticTarget{S,R,O,M} <: AbstractCriticTrainingTarget +struct RolloutCriticTarget{S,R,O,M,B,P} <: AbstractCriticTrainingTarget stage_problem horizon::Int n_uncertainty::Int @@ -46,6 +96,9 @@ struct RolloutCriticTarget{S,R,O,M} <: AbstractCriticTrainingTarget policy_state::Symbol reuse_solver::Bool objective_value::Symbol + state_bounds::B + project_state::P + retry_on_failure::Bool end function RolloutCriticTarget( @@ -60,6 +113,9 @@ function RolloutCriticTarget( policy_state::Symbol = :target, reuse_solver::Bool = false, objective_value::Symbol = :objective, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, ) policy_state in (:target, :realized) || error("policy_state must be :target or :realized") @@ -77,32 +133,66 @@ function RolloutCriticTarget( policy_state, reuse_solver, objective_value, + state_bounds, + project_state, + retry_on_failure, ) end """ NoCriticControlVariate() -Default no-op critic configuration. Passing this to `train_tsddr` recovers the -original dual-multiplier actor update. +Construct the no-op critic configuration. + +# Returns +- `NoCriticControlVariate`: a sentinel that disables critic/control-variate + terms. + +# Notes +Passing this value to `train_tsddr` recovers the original dual-multiplier actor +update. """ struct NoCriticControlVariate <: AbstractCriticControlVariate end """ - ScalarCriticControlVariate(critic; featurizer=default_critic_featurizer, - value_loss_weight=0.1, - gradient_loss_weight=1.0) - -Wrap a scalar Flux-compatible critic `C(w, xhat)` for optional TS-DDR -control-variate training. The critic is called as `critic(features)`, where -`features = featurizer(initial_state, uncertainty, xhat)`. - -The critic loss is - - value_loss_weight * mse(C, objective) - + gradient_loss_weight * mse(gradient(xhat -> C, xhat), target_multipliers) - -Either loss weight may be zero. + ScalarCriticControlVariate( + critic; + featurizer = default_critic_featurizer, + value_loss_weight::Real = 0.1, + gradient_loss_weight::Real = 1.0, + ) -> ScalarCriticControlVariate + +Wrap a scalar Flux-compatible critic for optional TS-DDR control-variate +training. + +# Arguments +- `critic`: callable scalar model evaluated as `critic(features)`. + +# Keywords +- `featurizer`: callable + `featurizer(initial_state, uncertainty, xhat) -> features`. +- `value_loss_weight::Real`: nonnegative weight on objective-value matching. +- `gradient_loss_weight::Real`: nonnegative weight on target-gradient + matching. + +# Returns +- `ScalarCriticControlVariate`: critic configuration with loss weights stored + as `Float64`. + +# Throws +- Throws an error if either loss weight is negative. + +# Notes +For each [`CriticSample`](@ref), the critic loss is + +```math +w_v |C(f) - J|^2 ++ w_g \\frac{1}{n}\\|\\nabla_{\\hat{x}} C(f) - \\lambda_{\\hat{x}}\\|_2^2, +``` + +where `f = featurizer(initial_state, uncertainty, xhat)`, `J` is the scalar +objective target, and ``\\lambda_{\\hat{x}}`` is the target multiplier array. +Either weight may be zero. """ struct ScalarCriticControlVariate{C,F} <: AbstractCriticControlVariate critic::C @@ -128,11 +218,33 @@ function ScalarCriticControlVariate( end """ - CriticSample(initial_state, uncertainty, xhat, objective_value, - target_multipliers; metadata=nothing) + CriticSample( + initial_state, + uncertainty, + xhat, + objective_value::Real, + target_multipliers; + metadata = nothing, + ) -> CriticSample + +Store one already-solved TS-DDR scenario as scalar-critic supervision. + +# Arguments +- `initial_state`: initial state used for the scenario. +- `uncertainty`: scenario uncertainty trajectory. +- `xhat`: policy target trajectory. +- `objective_value::Real`: scalar value target for the critic. +- `target_multipliers`: multiplier-like target with the same shape as `xhat`. + +# Keywords +- `metadata`: optional payload retained with the sample. + +# Returns +- `CriticSample`: sample with `objective_value` converted to `Float64`. -Training sample for a scalar critic. Samples are produced from already-solved -TS-DDR scenarios and do not require additional optimization solves. +# Notes +Creating a `CriticSample` does not run any optimization solve; samples are +intended to be built from existing training or rollout results. """ struct CriticSample{I,W,X,L,M} initial_state::I @@ -161,6 +273,21 @@ function CriticSample( ) end +""" + CriticReplayBuffer(max_size::Integer) -> CriticReplayBuffer + +Construct a fixed-capacity FIFO replay buffer for [`CriticSample`](@ref)s. + +# Arguments +- `max_size::Integer`: maximum number of samples retained; negative values are + clamped to zero. + +# Returns +- `CriticReplayBuffer`: empty replay buffer with capacity `max(0, max_size)`. + +# Notes +A capacity of zero disables buffering, so push operations become no-ops. +""" mutable struct CriticReplayBuffer{S} samples::Vector{S} max_size::Int @@ -169,6 +296,22 @@ end CriticReplayBuffer(max_size::Integer) = CriticReplayBuffer{Any}(Any[], max(0, Int(max_size))) +""" + push_critic_sample!(buffer, sample) -> buffer + +Append one [`CriticSample`](@ref) to the buffer, evicting the oldest sample if +the buffer is at capacity. + +# Arguments +- `buffer::CriticReplayBuffer`: replay buffer to mutate. +- `sample::CriticSample`: sample to append. + +# Returns +- `buffer`: the same buffer object, after optional insertion and eviction. + +# Notes +If `buffer.max_size == 0`, the function returns without storing `sample`. +""" function push_critic_sample!(buffer::CriticReplayBuffer, sample::CriticSample) buffer.max_size == 0 && return buffer push!(buffer.samples, sample) @@ -177,6 +320,18 @@ function push_critic_sample!(buffer::CriticReplayBuffer, sample::CriticSample) return buffer end +""" + push_critic_samples!(buffer, samples) -> buffer + +Append multiple [`CriticSample`](@ref)s to the buffer in order. + +# Arguments +- `buffer::CriticReplayBuffer`: replay buffer to mutate. +- `samples`: iterable of [`CriticSample`](@ref) values. + +# Returns +- `buffer`: the same buffer after appending the supplied samples. +""" function push_critic_samples!(buffer::CriticReplayBuffer, samples) for sample in samples push_critic_sample!(buffer, sample) @@ -185,10 +340,17 @@ function push_critic_samples!(buffer::CriticReplayBuffer, samples) end """ - default_critic_featurizer(initial_state, uncertainty, xhat) + default_critic_featurizer(initial_state, uncertainty, xhat) -> AbstractVector + +Concatenate flattened critic inputs into one feature vector. + +# Arguments +- `initial_state`: initial state for the scenario. +- `uncertainty`: uncertainty trajectory for the scenario. +- `xhat`: policy target trajectory for the scenario. -Default critic featurizer: concatenate flattened initial state, uncertainty, and -policy target trajectory. +# Returns +- `AbstractVector`: `vcat(vec(initial_state), vec(uncertainty), vec(xhat))`. """ default_critic_featurizer(initial_state, uncertainty, xhat) = vcat(vec(initial_state), vec(uncertainty), vec(xhat)) @@ -205,9 +367,21 @@ function _critic_value(critic, featurizer, initial_state, uncertainty, xhat) end """ - critic_value(control_variate, initial_state, uncertainty, xhat) + critic_value(control_variate, initial_state, uncertainty, xhat) -> Number Evaluate the scalar critic on one scenario. + +# Arguments +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `initial_state`: initial state for the scenario. +- `uncertainty`: uncertainty trajectory for the scenario. +- `xhat`: policy target trajectory for the scenario. + +# Returns +- `Number`: scalar critic prediction. + +# Throws +- Throws an error if the critic returns a non-scalar array. """ critic_value( cv::ScalarCriticControlVariate, @@ -219,8 +393,23 @@ critic_value( """ critic_xhat_gradient(control_variate, initial_state, uncertainty, xhat) -Return `gradient(xhat -> C(initial_state, uncertainty, xhat), xhat)` and check -that it has the same shape as `xhat`. +Differentiate the scalar critic with respect to the policy target trajectory. + +# Arguments +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `initial_state`: initial state for the scenario. +- `uncertainty`: uncertainty trajectory for the scenario. +- `xhat`: policy target trajectory for the scenario. + +# Returns +- An array with the same shape as `xhat`, equal to + `gradient(x -> critic_value(control_variate, initial_state, uncertainty, x), xhat)`. + +# Throws +- Throws an error if the critic gradient shape differs from `xhat`. + +# Notes +If Zygote reports `nothing`, the gradient is replaced by `zero(xhat)`. """ function critic_xhat_gradient( cv::ScalarCriticControlVariate, @@ -285,9 +474,32 @@ function _critic_loss_with( end """ - critic_loss(control_variate, samples; value_loss_weight, gradient_loss_weight) + critic_loss( + control_variate, + samples; + value_loss_weight = control_variate.value_loss_weight, + gradient_loss_weight = control_variate.gradient_loss_weight, + ) -> Real Compute the scalar critic loss on a collection of `CriticSample`s. + +# Arguments +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `samples`: collection of [`CriticSample`](@ref) values. + +# Keywords +- `value_loss_weight`: nonnegative override for objective-value loss weight. +- `gradient_loss_weight`: nonnegative override for target-gradient loss + weight. + +# Returns +- `Real`: average critic loss over `samples`, or `0.0` for an empty + collection. + +# Throws +- Throws an error if either loss weight is negative. +- Throws an error if a sample's target multipliers or critic gradient do not + match the shape of `xhat`. """ critic_loss(cv::ScalarCriticControlVariate, samples; kwargs...) = _critic_loss_with(cv.critic, cv, samples; kwargs...) @@ -303,10 +515,32 @@ function _critic_minibatch(samples, batch_size) end """ - update_critic!(opt_state, control_variate, samples; batch_size=nothing) - -Run one critic optimizer step and return the numeric loss. Only critic -parameters are updated. + update_critic!( + opt_state, + control_variate, + samples; + batch_size = nothing, + ) -> Float64 + +Run one optimizer step for the scalar critic. + +# Arguments +- `opt_state`: Flux optimizer state for `control_variate.critic`. +- `control_variate::ScalarCriticControlVariate`: critic configuration. +- `samples`: replay samples available for training. + +# Keywords +- `batch_size`: optional minibatch size; `nothing` or a value greater than the + sample count uses all samples. + +# Returns +- `Float64`: critic loss on the selected batch, or `NaN` when the selected + batch is empty. + +# Notes +Only critic parameters are updated. If the materialized gradient is `nothing` +or contains non-finite values, the optimizer update is skipped while the loss +is still reported. """ function update_critic!( opt_state, diff --git a/src/deterministic_equivalent.jl b/src/deterministic_equivalent.jl index bf672b5..6c422da 100644 --- a/src/deterministic_equivalent.jl +++ b/src/deterministic_equivalent.jl @@ -9,15 +9,32 @@ """ DeterministicEquivalentProblem -Holds an ExaModels parametric NLP for the deterministic equivalent subproblem +Container for an ExaModels parametric NLP representing a deterministic +equivalent subproblem. - Q(w, x̂) = min_{x,u,δ} Σ_t stage_cost(t, x_t, u_t, w_t) + (ρ/2)‖δ‖² - s.t. x₁ = x₀ - dynamics(t, x_t, u_t, w_t, x_{t+1}) = 0 - x̂_t - x_t - δ_t = 0 (target constraints, **added last**) +# Fields -We add the target constraints last so their multipliers are a contiguous slice of -`result.multipliers`. +- `core`: ExaModels core used to build variables, parameters, objectives, and + constraints. +- `model`: ExaModels model passed to MadNLP. +- `x`, `u`, `δ`: Flat state, control, and target-slack variables. +- `p_x0`, `p_w`, `p_target`: Initial-state, uncertainty, and target parameters. +- `nx`, `nu`, `nw`, `horizon`: State, control, uncertainty, and time dimensions. +- `target_con_range`: Range of target-constraint multipliers in solver results. + +# Notes + +The subproblem has the form + +```math +Q(w, \\hat{x}) = + \\min_{x,u,\\delta} + \\sum_t c_t(x_t, u_t, w_t) + \\frac{\\rho}{2}\\|\\delta\\|^2 +``` + +subject to the initial condition, dynamics constraints, and target constraints +``\\hat{x}_t - x_t - \\delta_t = 0``. The target constraints are added last so +their multipliers occupy the contiguous slice `target_con_range`. """ struct DeterministicEquivalentProblem core @@ -42,8 +59,17 @@ end """ MadNLPCache -Optional cache to re-use a MadNLP solver instance across repeated solves -(warm-start + reusing symbolic factorizations). +Cache for reusing a MadNLP solver across repeated solves. + +# Fields + +- `solver`: Cached `MadNLP.MadNLPSolver`. +- `last_result`: Most recent solver result, used for optional warm starts. + +# Notes + +Reusing the solver can avoid repeated symbolic setup and can warm-start the +next primal iterate from the previous solution. """ mutable struct MadNLPCache solver @@ -53,27 +79,45 @@ end """ build_deterministic_equivalent(; kwargs...) -> DeterministicEquivalentProblem -Generic builder for a deterministic-equivalent dynamic NLP. - -Keyword arguments: -- `horizon::Int` : number of stages T (states are 1..T, controls are 1..T-1) -- `nx::Int` : state dimension -- `nu::Int` : control dimension (demo assumes `nu == nx` for default dynamics) -- `nw::Int` : disturbance dimension (default: `nx`) -- `backend` : ExaModels backend (e.g., `nothing` for CPU, `CUDABackend()` for GPU) -- `float_type` : numeric type (Float64 recommended for MadNLP) -- `x_bounds` : `(lb, ub)` applied to all state variables -- `u_bounds` : `(lb, ub)` applied to all control variables -- `slack_penalty` : ρ ≥ 0, weight for (ρ/2)‖δ‖² - -- `dynamics_eq` : function `(t, i, x, u, w, nx, nu, nw) -> expr == 0` - returns the scalar equality residual for dimension `i` at stage `t`. -- `stage_cost` : function `(t, i, x, u, w, nx, nu, nw) -> expr` returns a scalar term. - -Notes: -- All variables are stored in flat vectors to keep the interface simple and robust. -- Target constraints are added *last* and written as `x̂ - x - δ = 0` so that their - multipliers are directly the gradient w.r.t. `x̂` (envelope theorem). +Build a deterministic-equivalent dynamic nonlinear program. + +# Keywords + +- `horizon::Int`: Number of stages. States are indexed over `1:horizon`; + controls and uncertainties over `1:(horizon - 1)`. +- `nx::Int`: State dimension. +- `nu::Int = nx`: Control dimension. +- `nw::Int = nx`: Uncertainty dimension. +- `backend = nothing`: ExaModels backend, such as `nothing` for CPU or a CUDA + backend for GPU execution. +- `float_type::Type{<:AbstractFloat} = Float64`: Scalar type used by the model. +- `x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf)`: Lower and upper bounds applied + to every state variable. +- `u_bounds = (-Inf, Inf)`: Control bounds, either a scalar `(lb, ub)` tuple or + a tuple of length-`nu` lower and upper bound vectors. +- `slack_penalty::Real = 1.0`: Nonnegative target-slack penalty weight ``ρ``. +- `dynamics_eq::Function = default_dynamics_eq`: Function + `(t, i, x, u, w, nx, nu, nw) -> residual` defining one scalar dynamics + equality. +- `stage_cost::Function = default_stage_cost`: Function + `(t, i, x, u, w, nx, nu, nw) -> term` defining one scalar objective term. + +# Returns + +- `DeterministicEquivalentProblem`: A mutable problem container with parameters + that can be updated by `set_x0!`, `set_uncertainty!`, and `set_targets!`. + +# Throws + +Throws an error if dimensions are invalid, if vector control bounds have the +wrong length, or if the default dynamics/cost are used with dimensions other +than `nu == nx` and `nw == nx`. + +# Notes + +All variables are stored in flat vectors. Target constraints are added last and +written as ``\\hat{x} - x - \\delta = 0`` so the envelope theorem identifies their +multipliers with gradients with respect to the target trajectory. """ function build_deterministic_equivalent(; horizon::Int, @@ -180,11 +224,28 @@ end """ build_linear_tracking_problem(; kwargs...) -Convenience wrapper that builds a *simple* deterministic equivalent problem with: -- dynamics: x_{t+1} = x_t + u_t + w_t -- stage cost: (1/2) x_t^2 + (1/2) u_t^2 +Build the default linear-quadratic tracking demonstration problem. + +# Keywords + +- `horizon::Int`: Number of stages. +- `nx::Int = 1`: State, control, and uncertainty dimension. +- `backend = nothing`: ExaModels backend. +- `float_type::Type{<:AbstractFloat} = Float64`: Scalar type used by the model. +- `x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf)`: State bounds. +- `u_bounds::Tuple{<:Real,<:Real} = (-1.0, 1.0)`: Control bounds. +- `slack_penalty::Real = 10.0`: Target-slack penalty weight. -This is meant as an end-to-end demo and a template for your real model. +# Returns + +- `DeterministicEquivalentProblem`: Problem with dynamics + ``x_{t+1} = x_t + u_t + w_t`` and stage cost + ``(x_t^2 + u_t^2) / 2``. + +# Notes + +This helper is intended as a small end-to-end example and as a template for +model-specific deterministic-equivalent builders. """ function build_linear_tracking_problem(; horizon::Int, @@ -215,11 +276,28 @@ end # -------------------------- """ -Default per-dimension dynamics residual for the demo: + default_dynamics_eq(t, i, x, u, w, nx::Int, nu::Int, nw::Int) + +Return the default scalar dynamics residual. + +# Arguments - x_{t+1,i} - x_{t,i} - u_{t,i} - w_{t,i} == 0 +- `t`: Stage index. +- `i`: State component index. +- `x`: Flat state variable vector. +- `u`: Flat control variable vector. +- `w`: Flat uncertainty parameter vector. +- `nx::Int`: State dimension. +- `nu::Int`: Control dimension. +- `nw::Int`: Uncertainty dimension. -Assumes `nu == nx` and `nw == nx`. +# Returns + +- Scalar residual ``x_{t+1,i} - x_{t,i} - u_{t,i} - w_{t,i}``. + +# Notes + +The default residual assumes `nu == nx` and `nw == nx`. """ function default_dynamics_eq(t, i, x, u, w, nx::Int, nu::Int, nw::Int) return x[x_index(nx, t + 1, i)] - @@ -229,9 +307,24 @@ function default_dynamics_eq(t, i, x, u, w, nx::Int, nu::Int, nw::Int) end """ -Default per-dimension stage cost term for the demo: + default_stage_cost(t, i, x, u, w, nx::Int, nu::Int, nw::Int) + +Return the default scalar stage-cost term. - (1/2) x_{t,i}^2 + (1/2) u_{t,i}^2 +# Arguments + +- `t`: Stage index. +- `i`: State/control component index. +- `x`: Flat state variable vector. +- `u`: Flat control variable vector. +- `w`: Flat uncertainty parameter vector. +- `nx::Int`: State dimension. +- `nu::Int`: Control dimension. +- `nw::Int`: Uncertainty dimension. + +# Returns + +- Scalar cost ``(x_{t,i}^2 + u_{t,i}^2) / 2``. """ function default_stage_cost(t, i, x, u, w, nx::Int, nu::Int, nw::Int) return (x[x_index(nx, t, i)]^2 + u[u_index(nu, t, i)]^2) / 2 @@ -242,9 +335,22 @@ end # -------------------------- """ - set_x0!(prob, x0) + set_x0!(prob::DeterministicEquivalentProblem, x0::AbstractVector) + +Update the initial-state parameter. + +# Arguments -Update initial state parameter (length nx). +- `prob::DeterministicEquivalentProblem`: Problem to update. +- `x0::AbstractVector`: Initial state with length `prob.nx`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `length(x0) != prob.nx`. """ function set_x0!(prob::DeterministicEquivalentProblem, x0::AbstractVector) length(x0) == prob.nx || error("x0 length must be nx=$(prob.nx), got $(length(x0))") @@ -253,12 +359,29 @@ function set_x0!(prob::DeterministicEquivalentProblem, x0::AbstractVector) end """ - set_uncertainty!(prob, w) + set_uncertainty!(prob::DeterministicEquivalentProblem, w::AbstractVector) + +Update the disturbance-trajectory parameter. -Update disturbance trajectory parameter. -`w` must have length `(T-1)*nw` or `T*nw`; in the latter case the first `(T-1)*nw` -elements are used (the last stage's uncertainty only enters the policy rollout, not -the NLP dynamics). +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem to update. +- `w::AbstractVector`: Disturbance trajectory with length `(T - 1) * nw` or + `T * nw`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `w` has any length other than `(T - 1) * nw` or `T * nw`. + +# Notes + +When `w` has length `T * nw`, only the first `(T - 1) * nw` entries are used in +the NLP dynamics. The final-stage uncertainty may be needed by a policy rollout +but does not enter these dynamics constraints. """ function set_uncertainty!(prob::DeterministicEquivalentProblem, w::AbstractVector) expected = (prob.horizon - 1) * prob.nw @@ -274,9 +397,22 @@ function set_uncertainty!(prob::DeterministicEquivalentProblem, w::AbstractVecto end """ - set_targets!(prob, xhat) + set_targets!(prob::DeterministicEquivalentProblem, xhat::AbstractVector) + +Update the target-trajectory parameter. + +# Arguments -Update target trajectory parameter (length T*nx). +- `prob::DeterministicEquivalentProblem`: Problem to update. +- `xhat::AbstractVector`: Target trajectory with length `prob.horizon * prob.nx`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `length(xhat) != prob.horizon * prob.nx`. """ function set_targets!(prob::DeterministicEquivalentProblem, xhat::AbstractVector) expected = prob.horizon * prob.nx @@ -292,7 +428,20 @@ end """ init_madnlp_cache(prob; solver_kwargs...) -> MadNLPCache -Create and store a `MadNLP.MadNLPSolver` for repeated solves. +Create a cached MadNLP solver for repeated solves. + +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem whose ExaModel will be solved. + +# Keywords + +- `solver_kwargs...`: Keyword arguments forwarded to `MadNLP.MadNLPSolver`. + +# Returns + +- `MadNLPCache`: Cache containing the solver and an initially empty + `last_result`. """ function init_madnlp_cache(prob::DeterministicEquivalentProblem; solver_kwargs...) solver = MadNLP.MadNLPSolver(prob.model; solver_kwargs...) @@ -300,19 +449,62 @@ function init_madnlp_cache(prob::DeterministicEquivalentProblem; solver_kwargs.. end """ - solve!(prob; solver_kwargs...) -> result + solve!(prob::DeterministicEquivalentProblem; solver_kwargs...) -> result -Solve once by instantiating a fresh MadNLP solver (simplest, but allocates). +Solve a deterministic-equivalent problem with a fresh MadNLP solver. + +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem to solve. + +# Keywords + +- `solver_kwargs...`: Keyword arguments forwarded to `MadNLP.madnlp`. + +# Returns + +- `result`: MadNLP result object. + +# Notes + +This path is simple but allocates a new solver for every call. """ function solve!(prob::DeterministicEquivalentProblem; solver_kwargs...) return MadNLP.madnlp(prob.model; solver_kwargs...) end """ - solve!(prob, cache; warmstart=true, solver_kwargs...) -> result + solve!(prob::DeterministicEquivalentProblem, cache::MadNLPCache; warmstart=true, solver_kwargs...) -> result + +Solve a deterministic-equivalent problem with a cached MadNLP solver. -Solve using a cached solver instance. Optionally warm-start the primal -variables from the previous solution. +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem to solve. +- `cache::MadNLPCache`: Cached solver and previous result. + +# Keywords + +- `warmstart::Bool = true`: Whether to copy the previous primal solution into + the model initial point before solving. +- `solver_kwargs...`: Keyword arguments forwarded to `MadNLP.solve!`. + +# Returns + +- `result`: MadNLP result object, also stored in `cache.last_result`. + +# Notes + +Warm starts are used only when `warmstart` is true and `cache.last_result` is +available. + +MadNLP's iteration counter `cnt.k` is cumulative across `solve!` calls on the +same solver instance and is never reset by MadNLP itself. Without resetting it +(together with `cnt.acceptable_cnt` and `cnt.start_time`), repeated solves on a +cached solver eventually exhaust `max_iter` spuriously — the counters are reset +here before each solve so every call gets its intended per-solve iteration and +wall-clock budget. This mirrors the reset in `_solve!` (training.jl) and does +not change any numerics of an individual solve. """ function solve!(prob::DeterministicEquivalentProblem, cache::MadNLPCache; warmstart::Bool = true, @@ -322,6 +514,10 @@ function solve!(prob::DeterministicEquivalentProblem, cache::MadNLPCache; # Warm-start primal from previous solution copyto!(NLPModels.get_x0(prob.model), cache.last_result.solution) end + # Reset per-solve iteration budget (cnt.k is cumulative in MadNLP). + cache.solver.cnt.k = 0 # reset iteration counter + cache.solver.cnt.acceptable_cnt = 0 # reset acceptable-step counter + cache.solver.cnt.start_time = time() # reset wall-clock timer res = MadNLP.solve!(cache.solver; solver_kwargs...) cache.last_result = res return res @@ -332,17 +528,43 @@ end # -------------------------- """ - target_multipliers(prob, result) -> λ + target_multipliers(prob::DeterministicEquivalentProblem, result) -> λ + +Return the dual multipliers associated with target constraints. -Return the dual multipliers associated with the target constraints (∇_{x̂} Q). +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem that defines the multiplier + slice. +- `result`: MadNLP result containing `multipliers`. + +# Returns + +- `λ`: Multipliers in `result.multipliers[prob.target_con_range]`. + +# Notes + +With target constraints written as ``\\hat{x} - x - \\delta = 0``, these +multipliers are the envelope-theorem derivatives with respect to the target +trajectory. """ target_multipliers(prob::DeterministicEquivalentProblem, result) = result.multipliers[prob.target_con_range] """ - solution_components(prob, result) -> (x, u, δ) + solution_components(prob::DeterministicEquivalentProblem, result) -> (x, u, δ) Split the flat solution vector into state, control, and slack components. + +# Arguments + +- `prob::DeterministicEquivalentProblem`: Problem that defines component sizes. +- `result`: MadNLP result containing `solution`. + +# Returns + +- `(x, u, δ)`: Flat slices of the primal solution for states, controls, and + target slacks. """ function solution_components(prob::DeterministicEquivalentProblem, result) n_x = prob.horizon * prob.nx diff --git a/src/embedded_deterministic_equivalent.jl b/src/embedded_deterministic_equivalent.jl new file mode 100644 index 0000000..4975955 --- /dev/null +++ b/src/embedded_deterministic_equivalent.jl @@ -0,0 +1,580 @@ +# embedded_deterministic_equivalent.jl +# +# Embedded-NN deterministic equivalent: the policy π_θ is inside the NLP via +# VectorNonlinearOracle. At convergence the duals λ_t are closed-loop (joint +# NLP) and the gradient ∇_θ Q = Σ_t λ_t · ∇_θ π_θ(w_t, x*_{t-1}) follows +# from the envelope theorem — structurally identical to the open-loop formula +# but with realized states from the coupled solve. +# +# The oracle constraint is: +# π_θ(w_t, x_{t-1}) − x_t − δ_t = 0 ∀t = 1…T +# +# One oracle for all T stages guarantees sequential LSTM evaluation with +# Flux.reset! at the top of each callback invocation. Flux.reset! on the +# threaded policies is a REAL reset (it restores Flux.initialstates), and each +# stage's policy forward advances the recurrent state exactly once, so within +# every callback the policy sees the stage sequence t = 1…T from a fresh +# initial state. Audit of per-stage advancement: oracle_f! calls the policy +# once per stage; oracle_jac!/oracle_vjp! call it inside Zygote.pullback for +# t > 1 (the pullback CONSTRUCTION runs the forward exactly once; calling the +# returned back(·) does not re-run it) and as a bare call for t = 1 — one +# forward, hence one state advance, per stage in all three callbacks. +# +# Jacobian exactness caveat: the oracle callbacks (oracle_jac!, oracle_vjp!) +# compute only the DIRECT partial ∂π_t/∂x_{t-1} via a per-stage Zygote pullback. +# For a recurrent policy whose recurrent layers see x_{t-1}, stage t's output +# also depends on x_{1..t-2} through the hidden state; those cross-stage entries +# are absent from both the sparsity pattern and the pullbacks, so the oracle +# Jacobian is exact for feedforward (stateless-in-x) policies and a structural +# approximation for such recurrent ones. When the recurrent encoder reads only +# the uncertainty w_t (as in StateConditionedPolicy and HydroReachablePolicy, +# where the combiner over [h_t; x_{t-1}] is feedforward in x), the hidden state +# does not depend on x and the direct partial IS the full derivative. This +# exactness claim SURVIVES recurrent-state threading: with the encoder reading +# only w_t, the threaded hidden state depends only on the inflow history +# w_{1..t}, never on x, so threading changes the VALUE of h_t but adds no +# ∂h_t/∂x dependence — ∂π_t/∂x_{t-1} remains the full derivative. + +""" + EmbeddedDeterministicEquivalentProblem + +Container for a deterministic-equivalent NLP whose target constraints are +computed by an embedded Flux policy. + +# Fields + +- `core`: ExaModels core used to build variables, parameters, objectives, and + constraints. +- `model`: ExaModels model passed to MadNLP. +- `x`, `u`, `δ`: Flat state, control, and target-slack variables. +- `p_x0`, `p_w`: Initial-state and uncertainty parameters. +- `policy`: Flux policy evaluated by the nonlinear oracle. +- `nx`, `nu`, `nw`, `horizon`: State, control, uncertainty, and time dimensions. +- `target_con_range`: Range of oracle-constraint multipliers in solver results. +- `_w_buf`, `_x0_buf`: Mutable host buffers captured by oracle callbacks. + +# Notes + +This problem is analogous to `DeterministicEquivalentProblem`, but the explicit +target parameter is replaced by a `VectorNonlinearOracle` enforcing +``\\pi_\\theta(w_t, x_{t-1}) - x_t - \\delta_t = 0``. The oracle closures capture +`policy` by reference, so changing Flux parameters between solves changes the +NLP callbacks without rebuilding the ExaModel. +""" +struct EmbeddedDeterministicEquivalentProblem{P} + core + model + x + u + δ + p_x0 + p_w + policy::P + nx::Int + nu::Int + nw::Int + horizon::Int + target_con_range::UnitRange{Int} + # mutable buffers captured by oracle closures + _w_buf::Vector{Float64} + _x0_buf::Vector{Float64} +end + +""" + set_x0!(prob::EmbeddedDeterministicEquivalentProblem, x0::AbstractVector) + +Update the initial state used by the embedded problem and oracle callbacks. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Embedded problem to update. +- `x0::AbstractVector`: Initial state with length `prob.nx`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `length(x0) != prob.nx`. + +# Notes + +The value is written both to the ExaModels parameter and to the oracle closure +buffer, ensuring the policy callback sees the same initial state as the NLP. +""" +function set_x0!(prob::EmbeddedDeterministicEquivalentProblem, x0::AbstractVector) + length(x0) == prob.nx || error("x0 length must be nx=$(prob.nx), got $(length(x0))") + ExaModels.set_parameter!(prob.core, prob.p_x0, x0) + copyto!(prob._x0_buf, Float64.(x0)) + return prob +end + +""" + set_uncertainty!(prob::EmbeddedDeterministicEquivalentProblem, w::AbstractVector) + +Update the disturbance trajectory used by the embedded problem and oracle +callbacks. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Embedded problem to update. +- `w::AbstractVector`: Disturbance trajectory with length `(T - 1) * nw` or + `T * nw`. + +# Returns + +- `prob`: The updated problem. + +# Throws + +Throws an error if `w` has any length other than `(T - 1) * nw` or `T * nw`. + +# Notes + +The first `(T - 1) * nw` entries are passed to the ExaModels dynamics +parameter. The full length-`T * nw` trajectory is retained in the oracle buffer +when provided, because the policy callback may evaluate a final-stage +disturbance. +""" +function set_uncertainty!(prob::EmbeddedDeterministicEquivalentProblem, w::AbstractVector) + expected = (prob.horizon - 1) * prob.nw + full_len = prob.horizon * prob.nw + n = length(w) + if n == expected + ExaModels.set_parameter!(prob.core, prob.p_w, w) + copyto!(view(prob._w_buf, 1:expected), Float64.(w)) + elseif n == full_len + ExaModels.set_parameter!(prob.core, prob.p_w, view(w, 1:expected)) + copyto!(prob._w_buf, Float64.(w)) + else + error("w length must be (T-1)*nw=$expected or T*nw=$full_len, got $n") + end + return prob +end + +""" + set_targets!(::EmbeddedDeterministicEquivalentProblem, ::AbstractVector) -> nothing + +Ignore explicit targets for embedded deterministic-equivalent problems. + +# Arguments + +- `::EmbeddedDeterministicEquivalentProblem`: Embedded problem whose targets are + generated by the policy oracle. +- `::AbstractVector`: Ignored target vector. + +# Returns + +- `nothing`. + +# Notes + +Embedded problems compute targets inline through the nonlinear oracle instead of +storing them in an NLP parameter. +""" +function set_targets!(::EmbeddedDeterministicEquivalentProblem, ::AbstractVector) + return nothing +end + +""" + invalidate_policy_cache!(embedded_de) + +Invalidate policy-dependent caches for an embedded deterministic-equivalent +problem. + +# Arguments + +- `embedded_de`: Embedded deterministic-equivalent problem. + +# Returns + +- `embedded_de`. + +# Notes + +The generic embedded problem evaluates the policy directly in each oracle +callback and has no cache to invalidate. Specialized embedded problem types may +extend this hook when their oracle stores policy-dependent intermediates across +solver calls. +""" +function invalidate_policy_cache!(embedded_de) + return embedded_de +end + +""" + build_embedded_deterministic_equivalent(policy; kwargs...) + +Build a deterministic-equivalent NLP with a Flux policy embedded as a nonlinear +oracle. + +# Arguments + +- `policy`: Flux model mapping each stage input to an `nx`-vector target. + +# Keywords + +- `horizon::Int`: Number of stages. States are indexed over `1:horizon`; + controls and uncertainties over `1:(horizon - 1)`. +- `nx::Int`: State dimension and policy output dimension. +- `nu::Int = nx`: Control dimension. +- `nw::Int = nx`: Uncertainty dimension and first part of the policy input. +- `backend = nothing`: ExaModels backend, such as `nothing` for CPU or a CUDA + backend for GPU execution. +- `float_type::Type{<:AbstractFloat} = Float64`: Scalar type used by the model. +- `x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf)`: Lower and upper bounds applied + to every state variable. +- `u_bounds = (-Inf, Inf)`: Control bounds, either a scalar `(lb, ub)` tuple or + a tuple of length-`nu` lower and upper bound vectors. +- `slack_penalty::Real = 1.0`: Nonnegative target-slack penalty weight ``ρ``. +- `dynamics_eq::Function = default_dynamics_eq`: Function + `(t, i, x, u, w, nx, nu, nw) -> residual` defining one scalar dynamics + equality. +- `stage_cost::Function = default_stage_cost`: Function + `(t, i, x, u, w, nx, nu, nw) -> term` defining one scalar objective term. + +# Returns + +- `EmbeddedDeterministicEquivalentProblem`: Problem supporting `set_x0!`, + `set_uncertainty!`, `target_multipliers`, and `solution_components`. + +# Throws + +Throws an error if dimensions are invalid, if vector control bounds have the +wrong length, or if the default dynamics/cost are used with dimensions other +than `nu == nx` and `nw == nx`. + +# Notes + +The oracle enforces +``\\pi_\\theta(w_t, x_{t-1}) - x_t - \\delta_t = 0`` and is added last so its +multipliers form a contiguous trailing slice of `result.multipliers`. This +matches the open-loop deterministic-equivalent convention while allowing the +policy to depend on realized previous states. + +The oracle Jacobian and vector-Jacobian callbacks differentiate each stage +independently, providing only the direct partial ``\\partial \\pi_t / +\\partial x_{t-1}``. If the policy's recurrent layers consume ``x_{t-1}``, +stage ``t``'s output also depends on ``x_{1..t-2}`` through the hidden state, +and those cross-stage Jacobian entries are omitted — the reported Jacobian is +then a structural approximation. The Jacobian is exact whenever the recurrent +part of the policy reads only ``w_t`` and the state enters through a +feedforward head (the [`StateConditionedPolicy`](@ref) architecture), because +then the hidden state carries no dependence on ``x``. This holds unchanged +with recurrent-state threading: the threaded hidden state is a function of +``w_{1..t}`` only, so it changes the value of ``h_t`` but introduces no +``\\partial h_t / \\partial x`` term. + +Every callback resets the policy's recurrent state at its top and evaluates +the stages in order, advancing the state exactly once per stage (for +``t > 1`` the forward runs during `Zygote.pullback` construction; the +returned pullback does not re-run it). +""" +function build_embedded_deterministic_equivalent( + policy; + horizon::Int, + nx::Int, + nu::Int = nx, + nw::Int = nx, + backend = nothing, + float_type::Type{<:AbstractFloat} = Float64, + x_bounds::Tuple{<:Real,<:Real} = (-Inf, Inf), + u_bounds = (-Inf, Inf), + slack_penalty::Real = 1.0, + dynamics_eq::Function = default_dynamics_eq, + stage_cost::Function = default_stage_cost, +) + horizon ≥ 2 || error("horizon must be ≥ 2 (got $horizon)") + nx ≥ 1 || error("nx must be ≥ 1 (got $nx)") + nu ≥ 1 || error("nu must be ≥ 1 (got $nu)") + nw ≥ 1 || error("nw must be ≥ 1 (got $nw)") + + if (dynamics_eq === default_dynamics_eq || stage_cost === default_stage_cost) && (nu != nx || nw != nx) + error("Default dynamics/cost assume nu == nx and nw == nx.") + end + + T = horizon + n_x = T * nx + n_u = (T - 1) * nu + + core = ExaModels.ExaCore(float_type; backend = backend) + + function _u_bound(b, side) + v = b[side] + v isa AbstractVector || return float_type(v) + length(v) == nu || error("u_bounds[$side] length must be nu=$nu") + return float_type.(repeat(v, T - 1)) + end + lvar_u = _u_bound(u_bounds, 1) + uvar_u = _u_bound(u_bounds, 2) + + x = ExaModels.variable(core, n_x; + lvar = float_type(x_bounds[1]), + uvar = float_type(x_bounds[2]), + ) + u = ExaModels.variable(core, n_u; + lvar = lvar_u, + uvar = uvar_u, + ) + δ = ExaModels.variable(core, n_x) + + p_x0 = ExaModels.parameter(core, zeros(float_type, nx)) + p_w = ExaModels.parameter(core, zeros(float_type, (T - 1) * nw)) + + ExaModels.objective(core, + stage_cost(t, i, x, u, p_w, nx, nu, nw) + for t in 1:(T - 1), i in 1:nx + ) + ρ = float_type(slack_penalty) + ExaModels.objective(core, + (ρ / 2) * δ[x_index(nx, t, i)]^2 + for t in 1:T, i in 1:nx + ) + + ExaModels.constraint(core, + x[x_index(nx, 1, i)] - p_x0[i] + for i in 1:nx + ) + ExaModels.constraint(core, + dynamics_eq(t, i, x, u, p_w, nx, nu, nw) + for t in 1:(T - 1), i in 1:nx + ) + + n_con_before_oracle = nx + (T - 1) * nx + + # ── Oracle buffers (mutated by set_x0! / set_uncertainty!) ─────────── + w_buf = zeros(Float64, T * nw) + x0_buf = zeros(Float64, nx) + + nvar_total = n_x + n_u + n_x # x, u, δ + x_start = 1 + δ_start = n_x + n_u + 1 + + # ── Pre-allocated oracle buffers ──────────────────────────────────── + _x_prev = zeros(Float32, nx) + _w_t = zeros(Float32, nw) + _input = zeros(Float32, nw + nx) + _J = zeros(Float32, nx, nx) + _e = zeros(Float32, nx) + _λ_t = zeros(Float32, nx) + + function _fill_x_prev!(t, xv) + for i in 1:nx + _x_prev[i] = (t == 1) ? + Float32(x0_buf[i]) : + Float32(xv[x_start + (t-2)*nx + i - 1]) + end + return _x_prev + end + + function _fill_w_t!(t) + for j in 1:nw + _w_t[j] = Float32(w_buf[(t-1)*nw + j]) + end + return _w_t + end + + function _fill_input!(t, xv) + _fill_w_t!(t) + _fill_x_prev!(t, xv) + copyto!(view(_input, 1:nw), _w_t) + copyto!(view(_input, nw+1:nw+nx), _x_prev) + return _input + end + + # ── Oracle callbacks ───────────────────────────────────────────────── + + function oracle_f!(c, xv) + # Real reset: the stage loop below starts from the initial recurrent + # state and each policy call advances it exactly once. + Flux.reset!(policy) + for t in 1:T + _fill_input!(t, xv) + nn_out = policy(_input) + for i in 1:nx + row = (t - 1) * nx + i + xi = x_start + (t-1)*nx + i - 1 + di = δ_start + (t-1)*nx + i - 1 + c[row] = Float64(nn_out[i]) - xv[xi] - xv[di] + end + end + return nothing + end + + # NOTE: this Jacobian holds only the per-stage direct partial ∂π_t/∂x_{t-1}. + # Cross-stage terms through a recurrent hidden state that depends on x are + # not represented (see file-top comment); exact when the recurrent encoder + # reads only w_t, as in StateConditionedPolicy / HydroReachablePolicy. + function oracle_jac!(vals, xv) + # Real reset; each stage advances the recurrent state exactly once: + # for t > 1 the forward runs during Zygote.pullback construction, and + # calling back(·) repeatedly does NOT re-run it; t = 1 is a bare call. + Flux.reset!(policy) + k = 0 + for t in 1:T + _fill_x_prev!(t, xv) + _fill_w_t!(t) + + nn_jac_xprev = if t > 1 + _, back = Zygote.pullback(xp -> policy(vcat(_w_t, xp)), _x_prev) + fill!(_J, 0f0) + for row in 1:nx + fill!(_e, 0f0) + _e[row] = 1.0f0 + col_grad = back(_e)[1] + if col_grad !== nothing + _J[row, :] .= col_grad + end + end + _J + else + # t = 1: no x-Jacobian block, but the forward must still run + # once so the recurrent state advances to stage 2. + policy(vcat(_w_t, _x_prev)) + nothing + end + + for i in 1:nx + k += 1; vals[k] = -1.0 + k += 1; vals[k] = -1.0 + if t > 1 + for j in 1:nx + k += 1; vals[k] = Float64(nn_jac_xprev[i, j]) + end + end + end + end + return nothing + end + + function oracle_vjp!(Jtv, xv, λ) + fill!(Jtv, 0.0) + # Real reset; same one-advance-per-stage discipline as oracle_jac!. + Flux.reset!(policy) + for t in 1:T + _fill_x_prev!(t, xv) + _fill_w_t!(t) + + for i in 1:nx + _λ_t[i] = Float32(λ[(t-1)*nx + i]) + xi = x_start + (t-1)*nx + i - 1 + di = δ_start + (t-1)*nx + i - 1 + Jtv[xi] -= λ[(t-1)*nx + i] + Jtv[di] -= λ[(t-1)*nx + i] + end + + if t > 1 + _, back = Zygote.pullback(xp -> policy(vcat(_w_t, xp)), _x_prev) + dinput = back(_λ_t)[1] + if dinput !== nothing + for j in 1:nx + xj = x_start + (t-2)*nx + j - 1 + Jtv[xj] += Float64(dinput[j]) + end + end + else + # t = 1: no x_prev contribution, but run the forward once so + # the recurrent state advances to stage 2. + policy(vcat(_w_t, _x_prev)) + end + end + return nothing + end + + # ── Sparsity pattern ───────────────────────────────────────────────── + jac_r = Int[] + jac_c = Int[] + for t in 1:T + for i in 1:nx + row = (t - 1) * nx + i + xi = x_start + (t-1)*nx + i - 1 + di = δ_start + (t-1)*nx + i - 1 + push!(jac_r, row); push!(jac_c, xi) # ∂g/∂x_{t,i} + push!(jac_r, row); push!(jac_c, di) # ∂g/∂δ_{t,i} + if t > 1 + for j in 1:nx + xj = x_start + (t-2)*nx + j - 1 + push!(jac_r, row); push!(jac_c, xj) # ∂g/∂x_{t-1,j} + end + end + end + end + + oracle = ExaModels.VectorNonlinearOracle( + nvar = nvar_total, + ncon = n_x, + nnzj = length(jac_r), + jac_rows = jac_r, + jac_cols = jac_c, + lcon = zeros(n_x), + ucon = zeros(n_x), + f! = oracle_f!, + jac! = oracle_jac!, + vjp! = oracle_vjp!, + adapt = Val(true), + ) + ExaModels.constraint(core, oracle) + + model = ExaModels.ExaModel(core) + + target_start = n_con_before_oracle + 1 + target_range = target_start:(target_start + n_x - 1) + + return EmbeddedDeterministicEquivalentProblem( + core, model, x, u, δ, + p_x0, p_w, + policy, + nx, nu, nw, T, + target_range, + w_buf, x0_buf, + ) +end + +""" + target_multipliers(prob::EmbeddedDeterministicEquivalentProblem, result) -> λ + +Return the dual multipliers associated with the embedded oracle constraints. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Problem that defines the + multiplier slice. +- `result`: MadNLP result containing `multipliers`. + +# Returns + +- `λ`: Multipliers in `result.multipliers[prob.target_con_range]`. + +# Notes + +The multipliers correspond to the constraints +``\\pi_\\theta(w_t, x_{t-1}) - x_t - \\delta_t = 0``. +""" +target_multipliers(prob::EmbeddedDeterministicEquivalentProblem, result) = + result.multipliers[prob.target_con_range] + +""" + solution_components(prob::EmbeddedDeterministicEquivalentProblem, result) -> (x, u, δ) + +Split the flat solution vector into state, control, and slack components. + +# Arguments + +- `prob::EmbeddedDeterministicEquivalentProblem`: Problem that defines component + sizes. +- `result`: MadNLP result containing `solution`. + +# Returns + +- `(x, u, δ)`: Flat slices of the primal solution for states, controls, and + target slacks. +""" +function solution_components(prob::EmbeddedDeterministicEquivalentProblem, result) + n_x = prob.horizon * prob.nx + n_u = (prob.horizon - 1) * prob.nu + sol = result.solution + x_sol = sol[1:n_x] + u_sol = sol[(n_x + 1):(n_x + n_u)] + δ_sol = sol[(n_x + n_u + 1):(n_x + n_u + n_x)] + return (x_sol, u_sol, δ_sol) +end diff --git a/src/policy.jl b/src/policy.jl index 8a3a6e8..2ab2db0 100644 --- a/src/policy.jl +++ b/src/policy.jl @@ -11,8 +11,17 @@ """ MLPPolicy(model, output_dim) -Stateless MLP policy: one call with `vcat(x0, w_flat)` returns the full target -trajectory `x̂` as a flat vector of length `T*nx`. +Wrap a stateless Flux model as a full-horizon target policy. + +The wrapped model is called once per scenario with `vcat(x0, w_flat)` and +returns the full target trajectory as a flat vector of length `output_dim`. + +# Arguments +- `model`: Flux-compatible callable. +- `output_dim::Int`: number of target components retained from `vec(model(input))`. + +# Returns +- `MLPPolicy`: a Flux layer wrapper around `model`. """ struct MLPPolicy{M} model::M @@ -21,6 +30,17 @@ end Flux.@layer MLPPolicy +""" + (policy::MLPPolicy)(input) -> AbstractVector + +Evaluate a stateless full-horizon policy. + +# Arguments +- `input`: concatenated initial state and flat uncertainty trajectory. + +# Returns +- The first `policy.output_dim` entries of `vec(policy.model(input))`. +""" function (π::MLPPolicy)(input) y = π.model(input) return vec(y)[1:π.output_dim] @@ -28,6 +48,19 @@ end """ MLPPolicy(input_dim, output_dim; hidden=(64,64), act=tanh) + +Construct a feed-forward [`MLPPolicy`](@ref). + +# Arguments +- `input_dim::Int`: length of the concatenated scenario input. +- `output_dim::Int`: length of the flattened target trajectory. + +# Keywords +- `hidden`: hidden-layer widths. +- `act`: hidden-layer activation. + +# Returns +- `MLPPolicy`: stateless full-horizon policy. """ function MLPPolicy(input_dim::Int, output_dim::Int; hidden = (64, 64), @@ -43,50 +76,531 @@ function MLPPolicy(input_dim::Int, output_dim::Int; return MLPPolicy(Flux.Chain(layers...), output_dim) end -# ── StateConditionedPolicy ──────────────────────────────────────────────────── +# ── ContextualPolicy ────────────────────────────────────────────────────────── + +""" + ContextualPolicy(policy, context) + +Wrap a stage policy so each call receives known exogenous context before the +usual policy input. The wrapped policy is called as +`policy(vcat(context_at(context, t), input))`, with `t` advanced once per call +and reset by `Flux.reset!`. + +This keeps training and rollout loops unchanged: they still pass `[w_t; x_prev]` +to the policy, while the wrapper prepends known stage features such as seasonal +phase or forecast covariates. Only the inner policy is trainable. +""" +mutable struct ContextualPolicy{P,C} + policy::P + context::C + t::Int +end + +ContextualPolicy(policy, context) = ContextualPolicy(policy, context, 0) + +Flux.@layer ContextualPolicy trainable=(policy,) + +""" + context_at(context, t) +Return the context vector for one-based stage `t`. """ - StateConditionedPolicy{E,C} +function context_at(context::AbstractMatrix, t::Integer) + 1 <= t <= size(context, 2) || + throw(BoundsError(context, (:, t))) + return view(context, :, t) +end -Stateful LSTM policy for sequential rollout: +context_at(context::Function, t::Integer) = context(t) - x̂_t = policy(vcat(w_t, x̂_{t-1})) +function (m::ContextualPolicy)(input) + m.t += 1 + return m.policy(vcat(context_at(m.context, m.t), input)) +end -- `encoder`: LSTM chain operating on the uncertainty slice `w_t` -- `combiner`: Dense layer combining encoder output with previous state +function Flux.reset!(m::ContextualPolicy) + m.t = 0 + Flux.reset!(m.policy) + return nothing +end -Call `Flux.reset!(policy)` before each episode. +""" + stage_phase_context(T; period, include_progress=true) -# Flux 0.16 note -LSTM requires ≥2D input. The forward pass reshapes the 1D `w_t` slice to -`(n_uncertainty, 1)` before encoding and squeezes back with `vec`. +Build a `d x T` context matrix with `sin(2*pi*t/period)`, +`cos(2*pi*t/period)`, and optionally normalized horizon progress `t/T`. +The sine/cosine pair preserves cyclic adjacency between the last and first +seasonal positions while using only two bounded input features. """ -struct StateConditionedPolicy{E,C} - encoder::E - combiner::C - n_uncertainty::Int - n_state::Int +function stage_phase_context(T::Integer; period::Integer, include_progress::Bool=true) + T >= 1 || throw(ArgumentError("T must be positive")) + period >= 1 || throw(ArgumentError("period must be positive")) + nrows = include_progress ? 3 : 2 + ctx = Matrix{Float32}(undef, nrows, T) + for t in 1:T + θ = 2f0 * Float32(pi) * Float32(t) / Float32(period) + ctx[1, t] = sin(θ) + ctx[2, t] = cos(θ) + if include_progress + ctx[3, t] = Float32(t) / Float32(T) + end + end + return ctx +end + +""" + vcat_contexts(a, b, ...) + +Vertically concatenate context matrices after checking that they cover the +same number of stages. +""" +function vcat_contexts(contexts::AbstractMatrix...) + isempty(contexts) && return Matrix{Float32}(undef, 0, 0) + T = size(first(contexts), 2) + all(size(c, 2) == T for c in contexts) || + throw(ArgumentError("all contexts must have the same number of columns")) + return vcat(contexts...) +end + +raw""" + _dense_policy_head(input_dim, output_dim, hidden; activation=tanh) + +Build the nonrecurrent target head used by state-conditioned policies. + +The head maps the concatenated features `[encoded_uncertainty; previous_state]` +to the normalized or unbounded target vector. When `hidden` is empty, this is the +historical single `Dense(input_dim => output_dim, activation)` head. When +`hidden = [h₁, …, h_L]`, the returned chain is + +```math +f(z) = +\sigma\left(W_{L+1} + \sigma\left(W_L \cdots \sigma(W_1 z + b_1) \cdots + b_L\right) + + b_{L+1}\right). +``` + +The output layer intentionally uses the same `activation`. This is different +from a generic MLP helper with a linear output: bounded policies need the final +head output to stay in `[0, 1]` when `activation=sigmoid`, before affine scaling +to physical target bounds. + +# Arguments +- `input_dim::Int`: dimension of `[encoded_uncertainty; previous_state]`. +- `output_dim::Int`: number of target components produced by the policy. +- `hidden::AbstractVector{Int}`: hidden widths of the feed-forward target head. +- `activation`: activation applied at every head layer, including the output. + +# Returns +- `Flux.Dense` if `hidden` is empty, otherwise `Flux.Chain` of dense layers. + +# Examples +```julia +head = DecisionRulesExa._dense_policy_head(135, 11, [128, 128]; activation=sigmoid) +``` + +See also: [`StateConditionedPolicy`](@ref), [`bounded_state_policy`](@ref) +""" +function _dense_policy_head( + input_dim::Int, + output_dim::Int, + hidden::AbstractVector{Int}; + activation = tanh, +) + isempty(hidden) && return Flux.Dense(input_dim => output_dim, activation) + layers = Any[Flux.Dense(input_dim => hidden[1], activation)] + for i in 1:(length(hidden) - 1) + push!(layers, Flux.Dense(hidden[i] => hidden[i + 1], activation)) + end + push!(layers, Flux.Dense(hidden[end] => output_dim, activation)) + return Flux.Chain(layers...) +end + +# ── Recurrent-state threading helpers (Flux ≥ 0.16 stateless cells) ────────── +# +# In Flux 0.16 recurrent layers are stateless: `(l::LSTM)(x)` restarts from +# `Flux.initialstates(l)` on EVERY call and `Flux.reset!` on Flux layers is a +# deprecated no-op. A policy that needs cross-stage memory must therefore carry +# the recurrent state itself and advance it with one stateful cell call per +# stage. The helpers below mirror DecisionRules.jl's `_as_cell`, +# `_init_recurrent_state`, `_step_encoder`, and `_state_eltype` +# (src/dense_multilayer_nn.jl) so both packages thread recurrent encoders with +# identical semantics. Unlike DecisionRules.jl, which stores BARE cells, EXA +# encoders keep the `Flux.LSTM` wrapper layers (preserving the weight +# structure of existing checkpoints); `_as_cell` unwraps to the underlying +# cell for the stateful `(x, state) -> (output, new_state)` call. + +""" + _as_cell(layer) + +Return the underlying recurrent cell of `layer`. `Flux.LSTM`/`GRU`/`RNN` wrap a +cell (`LSTMCell`/`GRUCell`/`RNNCell`) in a `.cell` field; if `layer` has no such +field it is already a cell and is returned unchanged. +""" +# Unwrap the .cell field if present (LSTM → LSTMCell); return unchanged otherwise. +_as_cell(layer) = hasfield(typeof(layer), :cell) ? layer.cell : layer + +""" + _init_recurrent_state(encoder) + +Return the initial recurrent state for `encoder`: `Flux.initialstates` of the +underlying cell for a single layer, or a tuple of per-layer initial states for +a `Chain`. + +# Notes +`Flux.initialstates` builds zero states with `zeros_like` on the cell weights, +so the returned state inherits the encoder's device and element type — call +[`Flux.reset!`](@ref) after moving a policy between devices to re-derive the +state on the new device. +""" +# Single layer: initial state of the underlying cell (zeros for LSTM h/c). +_init_recurrent_state(layer) = Flux.initialstates(_as_cell(layer)) +# Chain: one initial state per layer, returned as a tuple. +_init_recurrent_state(chain::Flux.Chain) = map(_init_recurrent_state, chain.layers) + +""" + _step_encoder(encoder, x, state) -> (output, new_state) + +Advance `encoder` by one step on input `x` from recurrent `state`, returning +the output and the updated state. For a `Chain`, each layer's output feeds the +next layer and each layer's state is threaded independently. +""" +# Single layer: one stateful cell call returns (output, new_state). +_step_encoder(layer, x, state) = _as_cell(layer)(x, state) +function _step_encoder(chain::Flux.Chain, x, states::Tuple) + # Delegate to the recursive tuple-based implementation. + return _step_encoder_layers(chain.layers, x, states) +end + +""" + _step_encoder_layers(layers, x, states) -> (output, new_states) + +Recursively advance a tuple of recurrent layers by one time step. + +Each layer receives the output of the previous layer as input and its own +independent recurrent state. The base case (`layers == ()`) returns the input +unchanged with an empty state tuple. + +# Arguments +- `layers::Tuple`: remaining recurrent layers to evaluate. +- `x`: current input (or output of the prior layer). +- `states::Tuple`: per-layer recurrent states, same length as `layers`. + +# Returns +- `output`: output of the last layer in `layers`. +- `new_states::Tuple`: updated recurrent states, one per layer. +""" +_step_encoder_layers(::Tuple{}, x, ::Tuple{}) = x, () +function _step_encoder_layers(layers::Tuple, x, states::Tuple) + # Advance the first layer with its own recurrent state. + out, new_state = _step_encoder(first(layers), x, first(states)) + + # Recurse on remaining layers, feeding this layer's output as input. + rest_out, rest_states = _step_encoder_layers(Base.tail(layers), out, Base.tail(states)) + + # Reassemble the full state tuple: this layer's state followed by the rest. + return rest_out, (new_state, rest_states...) end -Flux.@layer StateConditionedPolicy +""" + _state_eltype(state) -> Type + +Return the scalar element type of a recurrent state. + +For nested tuple states (e.g. LSTM's `(h, c)` or a `Chain`'s tuple of per-layer +states) this recurses into the first element until it reaches an +`AbstractVector`, then returns `eltype(v)`. The result is used to cast inputs +to the encoder's precision before each step. +""" +_state_eltype(state::Tuple) = _state_eltype(first(state)) +_state_eltype(v::AbstractArray) = eltype(v) # AbstractArray (not just Vector) so batched matrix states work + +""" + StateConditionedPolicy{E,C,S} + +Flux-compatible state-conditioned policy for sequential target rollout. + +At each stage the policy is called as + +```julia +xhat_t = policy(vcat(w_t, x_prev)) +``` + +where the recurrent encoder reads only `w_t` and the combiner reads +`[encoded_uncertainty; x_prev]`. +Flux's recurrent cells are stateless (Flux ≥ 0.16): each call returns +`(output, new_state)` instead of mutating internal state, and calling the +`LSTM` wrapper directly would restart from `initialstates` on every call. +`StateConditionedPolicy` therefore carries the encoder's recurrent state itself +in `state`, threading it through one cell call per stage — the same semantics +as DecisionRules.jl's `StateConditionedPolicy`. Call `Flux.reset!(policy)` to +restore it to `Flux.initialstates` at the start of a scenario. + +# Fields +- `encoder`: recurrent uncertainty encoder (`Chain` of `Flux.LSTM`-style layers). +- `combiner`: nonrecurrent target head. +- `state`: current recurrent state ``s_t``, carried across calls (not trainable). +- `n_uncertainty::Int`: number of uncertainty features at each stage. +- `n_state::Int`: number of previous-state features. +- `output_lower`: optional lower bounds for affine output scaling. +- `output_scale`: optional `upper - lower` scale for affine output scaling. + +# Notes +Call `Flux.reset!(policy)` before each scenario — the reset is REAL (it +restores the initial recurrent state), unlike the deprecated Flux-layer +`reset!` no-op. Within a differentiated rollout the threaded state is treated +as data: gradients flow into encoder/combiner parameters through each stage's +forward pass, and the state stored between calls is refreshed by mutation +(mirroring DecisionRules.jl's training semantics). +""" +mutable struct StateConditionedPolicy{E,C,S,L,U} + encoder::E # Recurrent uncertainty encoder (Chain of LSTM-style layers) + combiner::C # Nonrecurrent target head + state::S # Encoder recurrent state, carried across calls + n_uncertainty::Int # Number of uncertainty features per stage + n_state::Int # Number of previous-state features + output_lower::L # Optional affine output lower bounds + output_scale::U # Optional affine output scale (upper - lower) +end + +Flux.@layer StateConditionedPolicy trainable=(encoder, combiner) + +""" + (policy::StateConditionedPolicy)(input) -> AbstractVector + +Evaluate one stage of a state-conditioned policy, threading recurrent state. + +The input is split into the uncertainty portion ``w_t`` and the previous state +``x_{t-1}``, and the forward pass computes + +```math +h_t, s_t = \\text{encoder}(w_t, s_{t-1}), \\qquad +\\hat{x}_t = f_{\\text{combine}}([h_t;\\; x_{t-1}]), +``` + +where ``s_t`` is the updated recurrent state (stored in `policy.state` for the +next call). + +# Arguments +- `input`: concatenated vector `[w_t; x_prev]`. + +# Returns +- Raw combiner output, or affine-scaled output when `output_bounds` were + supplied at construction. +""" function (m::StateConditionedPolicy)(input) - w = reshape(input[1:m.n_uncertainty], :, 1) # (n_unc, 1) for LSTM + # Split the concatenated input into uncertainty w_t and previous state x_{t-1}. + w = input[1:m.n_uncertainty] s = input[m.n_uncertainty+1:end] - h = vec(m.encoder(w)) # (hidden,) - return m.combiner(vcat(h, s)) + + # Cast the uncertainty to the encoder precision (taken from the recurrent + # state), matching DecisionRules.jl's forward pass exactly. + T = _state_eltype(m.state) + + # Advance the recurrent encoder by one step: h_t, s_t = encoder(w_t, s_{t-1}). + h, new_state = _step_encoder(m.encoder, T.(w), m.state) + + # Persist the new recurrent state so the next call starts from s_t. + m.state = new_state + + y = m.combiner(vcat(h, s)) + if m.output_lower === nothing + return y + end + lower = _adapt_policy_bound(m.output_lower, y) + scale = _adapt_policy_bound(m.output_scale, y) + return lower .+ scale .* y end -Flux.reset!(m::StateConditionedPolicy) = Flux.reset!(m.encoder) +""" + Flux.reset!(policy::StateConditionedPolicy) -> Nothing + +Reset the encoder's recurrent state to `Flux.initialstates`, e.g. at the start +of a scenario rollout. +# Notes +The state is re-derived from the (possibly device-moved) encoder weights on +every reset, so calling `Flux.reset!` after `gpu(policy)`/`cpu(policy)` places +the state on the correct device with the correct element type. """ +function Flux.reset!(m::StateConditionedPolicy) + # Reinitialize s_0 to the cell defaults (zeros for LSTM h/c). + m.state = _init_recurrent_state(m.encoder) + return nothing +end + +""" + _adapt_policy_bound(x, ref) + +Move a policy bound vector to the same storage family and element type as +`ref`. + +# Arguments +- `x::AbstractVector`: bound vector stored on the policy. +- `ref::AbstractVector`: output vector whose element type and device should be + matched. + +# Returns +- `x` itself when the concrete vector type already matches `ref`; otherwise a + copied vector compatible with `ref`. +""" +function _adapt_policy_bound(x::AbstractVector, ref::AbstractVector) + typeof(x) === typeof(ref) && return x + y = similar(ref, length(x)) + copyto!(y, convert.(eltype(ref), x)) + return y +end + +""" + _load_encoder_state!(encoder, enc_state) -> encoder + +Load a checkpoint `encoder` state into a recurrent encoder, accepting BOTH +weight structures in use across the two packages: + +1. EXA structure — `Chain` of `Flux.LSTM` wrapper layers, so each layer state + is `(cell = (Wi, Wh, bias),)`. Loaded by the stock `Flux.loadmodel!`. +2. DecisionRules.jl (MAIN) structure — `Chain` of BARE `LSTMCell`s (MAIN's + `_as_cell` strips the wrapper at construction), so each layer state is + `(Wi, Wh, bias)` directly. `Flux.loadmodel!` matches children by key, so + loading a bare-cell layer state into an `LSTM` wrapper (children `(cell,)`) + throws; this helper falls back to cell-by-cell injection, loading each + layer state into `_as_cell(layer)` — the mathematically identical weight + assignment (the wrapper's `cell` has exactly the fields `(Wi, Wh, bias)` + MAIN saved). + +# Arguments +- `encoder`: destination encoder (typically a `Chain` of `Flux.LSTM` layers). +- `enc_state`: the checkpoint's encoder state (from `Flux.state`). + +# Returns +- `encoder`, mutated in place. + +# Throws +- Rethrows the stock `Flux.loadmodel!` error when the fallback does not apply + (non-`Chain` encoder, missing `layers`, or depth mismatch). +""" +function _load_encoder_state!(encoder, enc_state) + try + # Stock path: same weight structure (EXA wrapper layers). + Flux.loadmodel!(encoder, enc_state) + return encoder + catch err + # Fallback applies only to Chain encoders with a per-layer state list. + (encoder isa Flux.Chain && hasproperty(enc_state, :layers)) || rethrow(err) + layer_states = getproperty(enc_state, :layers) + length(layer_states) == length(encoder.layers) || rethrow(err) + for (layer, lstate) in zip(encoder.layers, layer_states) + # Wrapper-style layer state loads into the wrapper; bare-cell + # (MAIN) layer state loads into the unwrapped cell. + dest = hasproperty(lstate, :cell) ? layer : _as_cell(layer) + Flux.loadmodel!(dest, lstate) + end + return encoder + end +end + +""" + load_stateconditioned_policy!(policy, state) + +Load a Flux checkpoint into a [`StateConditionedPolicy`](@ref). + +# Arguments +- `policy::StateConditionedPolicy`: policy to update in place. +- `state`: checkpoint object accepted by `Flux.loadmodel!`. + +# Returns +- `policy`. + +# Notes +- Checkpoints saved before output bounds were added contain only the trainable + encoder and combiner state. In that case, this method restores those + trainable components and keeps the current policy's case-defined output + bounds. +- Checkpoints trained BEFORE recurrent-state threading (memoryless-encoder + era) have the SAME weight structure and load unchanged — only the runtime + semantics differ (the encoder now carries memory across stages). +- DecisionRules.jl (MAIN) checkpoints, whose encoders are `Chain`s of bare + `LSTMCell`s, load through the documented cell-by-cell fallback in + [`_load_encoder_state!`](@ref). +- The loaded policy's recurrent state is reset afterwards so the next rollout + starts from `Flux.initialstates` of the loaded weights. +""" +function load_stateconditioned_policy!(policy::StateConditionedPolicy, state) + try + Flux.loadmodel!(policy, state) + Flux.reset!(policy) + return policy + catch err + hasproperty(state, :encoder) && hasproperty(state, :combiner) || rethrow(err) + @warn "Full StateConditionedPolicy checkpoint load failed; loading encoder/combiner only and keeping current output bounds" exception=(err, catch_backtrace()) + _load_encoder_state!(policy.encoder, getproperty(state, :encoder)) + Flux.loadmodel!(policy.combiner, getproperty(state, :combiner)) + Flux.reset!(policy) + return policy + end +end + +function load_stateconditioned_policy!(policy::ContextualPolicy, state) + inner_state = hasproperty(state, :policy) ? getproperty(state, :policy) : state + load_stateconditioned_policy!(policy.policy, inner_state) + Flux.reset!(policy) + return policy +end + +raw""" StateConditionedPolicy(n_uncertainty, n_state, n_out, layers; - activation=tanh, encoder_type=Flux.LSTM) + activation=tanh, encoder_type=Flux.LSTM, + output_bounds=nothing, combiner_layers=Int[]) + +Construct a state-conditioned sequential target policy. + +The policy separates memory from state conditioning: + +```math +h_t = E_\theta(w_t, h_{t-1}), \qquad +\hat{x}_t = H_\theta([h_t;\, x_{t-1}]). +``` + +The recurrent encoder `Eθ` sees only the stage uncertainty. The previous state +enters only through the feed-forward head `Hθ`, so increasing +`combiner_layers` makes the state-to-target map nonlinear without introducing +recurrence over the state input. -Construct a `StateConditionedPolicy`. +# Arguments +- `n_uncertainty::Int`: number of uncertainty features in each stage input. +- `n_state::Int`: number of previous-state features appended after uncertainty. +- `n_out::Int`: output dimension, usually the state-target dimension. +- `layers::AbstractVector{Int}`: recurrent encoder hidden sizes. -- `layers` : hidden sizes for the LSTM encoder, e.g. `[64, 64]` -- `n_out` : output dimension (= nx = state dimension) +# Keywords +- `activation`: head activation. Use `sigmoid` with `output_bounds` when targets + must remain inside bounds. +- `encoder_type`: recurrent layer constructor, typically `Flux.LSTM`. +- `output_bounds`: optional `(lower, upper)` vectors. If provided, the raw head + output `y` is interpreted as normalized and mapped to + `lower + (upper - lower) .* y`. +- `combiner_layers`: hidden widths for the nonrecurrent target head. `Int[]` + preserves the original single Dense head. + +# Returns +- `StateConditionedPolicy` with trainable `encoder` and `combiner`. + +# Notes +The recurrent encoder receives only uncertainty. The previous state enters +through the feed-forward combiner, so `combiner_layers` increases +state-to-target expressiveness without making the state input recurrent. + +# Examples +```julia +policy = StateConditionedPolicy( + 11, 11, 11, [128, 128]; + activation = sigmoid, + output_bounds = (zeros(Float32, 11), ones(Float32, 11)), + combiner_layers = [128, 128], +) +``` + +See also: [`bounded_state_policy`](@ref) """ function StateConditionedPolicy( n_uncertainty::Int, @@ -95,11 +609,264 @@ function StateConditionedPolicy( layers::AbstractVector{Int}; activation = tanh, encoder_type = Flux.LSTM, + output_bounds = nothing, + combiner_layers = Int[], ) enc_sizes = vcat(n_uncertainty, layers) enc_layers = [encoder_type(enc_sizes[i] => enc_sizes[i+1]) for i in 1:length(layers)] encoder = Flux.Chain(enc_layers...) - combiner = Flux.Dense(layers[end] + n_state => n_out, activation) - return StateConditionedPolicy(encoder, combiner, n_uncertainty, n_state) + combiner = _dense_policy_head( + layers[end] + n_state, + n_out, + collect(Int, combiner_layers); + activation = activation, + ) + if output_bounds === nothing + # Initialize the recurrent state to Flux.initialstates for the encoder. + return StateConditionedPolicy( + encoder, combiner, _init_recurrent_state(encoder), + n_uncertainty, n_state, nothing, nothing, + ) + end + lower, upper = output_bounds + length(lower) == n_out || throw(ArgumentError("output lower bound length must be n_out=$n_out")) + length(upper) == n_out || throw(ArgumentError("output upper bound length must be n_out=$n_out")) + scale = upper .- lower + any(<(zero(eltype(scale))), scale) && + throw(ArgumentError("output upper bounds must be >= lower bounds")) + return StateConditionedPolicy( + encoder, combiner, _init_recurrent_state(encoder), + n_uncertainty, n_state, + collect(lower), collect(scale), + ) +end + +# ── Bounded state-target helpers ────────────────────────────────────────────── + +""" + ConstantStatePolicy(output_template, n_uncertainty, n_state) + +Represent a policy with no trainable target dimensions. + +The policy always returns `output_template`, adapted to the input device. + +# Arguments +- `output_template`: fixed full target vector. +- `n_uncertainty::Int`: number of uncertainty features expected in the input. +- `n_state::Int`: number of state features expected in the input. + +# Returns +- `ConstantStatePolicy`: Flux layer with no trainable parameters. +""" +struct ConstantStatePolicy{O} + output_template::O + n_uncertainty::Int + n_state::Int +end + +Flux.@layer ConstantStatePolicy trainable=() + +""" + (policy::ConstantStatePolicy)(input) -> AbstractVector + +Return the fixed target template, adapted to the input storage family. + +# Arguments +- `input`: vector used only as an adaptation reference. + +# Returns +- The fixed target vector. +""" +(m::ConstantStatePolicy)(input) = _adapt_policy_bound(m.output_template, input) + +""" + Flux.reset!(policy::ConstantStatePolicy) -> Nothing + +No-op reset method for constant policies. +""" +Flux.reset!(::ConstantStatePolicy) = nothing + +""" + FixedOutputPolicy(policy, output_template, output_expansion) + +Expand active target dimensions into a full target vector. + +The wrapped policy predicts only active dimensions. `output_template` stores +constants for inactive dimensions and zeros for active dimensions; +`output_expansion` maps active outputs into the full state vector without +mutation. + +# Arguments +- `policy`: Flux-compatible policy for active dimensions. +- `output_template`: full target vector with inactive constants. +- `output_expansion`: matrix mapping active outputs into full target space. + +# Returns +- `FixedOutputPolicy`: Flux layer wrapper around `policy`. +""" +struct FixedOutputPolicy{P,O,E} + policy::P + output_template::O + output_expansion::E +end + +Flux.@layer FixedOutputPolicy trainable=(policy,) + +""" + (policy::FixedOutputPolicy)(input) -> AbstractVector + +Evaluate the active-dimension policy and expand it into the full target vector. + +# Arguments +- `input`: stage input forwarded to the wrapped policy. + +# Returns +- Full target vector with active outputs inserted and inactive dimensions fixed. +""" +function (m::FixedOutputPolicy)(input) + y = m.policy(input) + return m.output_template .+ m.output_expansion * y +end + +""" + Flux.reset!(policy::FixedOutputPolicy) + +Reset the wrapped active-dimension policy. + +# Returns +- The result of `Flux.reset!(policy.policy)`. +""" +Flux.reset!(m::FixedOutputPolicy) = Flux.reset!(m.policy) + +""" + load_stateconditioned_policy!(policy::FixedOutputPolicy, state) + +Load checkpoint state into the wrapped active-dimension policy. + +# Arguments +- `policy::FixedOutputPolicy`: wrapper whose inner policy is updated. +- `state`: checkpoint object accepted by the inner policy loader. + +# Returns +- The result of `load_stateconditioned_policy!(policy.policy, state)`. +""" +load_stateconditioned_policy!(policy::FixedOutputPolicy, state) = + load_stateconditioned_policy!(policy.policy, state) + +raw""" + bounded_state_policy(n_uncertainty, lower, upper, layers; kwargs...) + +Build a state-conditioned policy whose full output is guaranteed to lie in +`[lower, upper]`, while avoiding trainable outputs for inactive dimensions. + +The returned policy has the same rollout semantics as +[`StateConditionedPolicy`](@ref): + +```math +\hat{x}_t = +\ell + (u - \ell) \odot H_\theta([E_\theta(w_t);\, x_{t-1}]), +``` + +where `Hθ` is a sigmoid head by default. Dimensions with `upper == lower` are +treated as constants unless `active_mask` overrides that choice. + +By default, a dimension is active when `upper > lower`. Fixed dimensions are +returned as constants, so pure pass-through or no-storage state components do +not create meaningless target parameters. Pass `active_mask` to override this +selection for case-specific target relevance. + +Set `combiner_layers` to add hidden layers after `[encoded_uncertainty; state]` +and before the bounded target output. This keeps recurrence confined to the +uncertainty encoder while making the state-to-target map nonlinear. + +# Arguments +- `n_uncertainty::Int`: number of uncertainty features in each stage input. +- `lower::AbstractVector`: lower bound for each full target dimension. +- `upper::AbstractVector`: upper bound for each full target dimension. +- `layers::AbstractVector{Int}`: recurrent uncertainty-encoder hidden sizes. + +# Keywords +- `activation`: head activation, defaulting to `sigmoid`. +- `encoder_type`: recurrent layer constructor. +- `active_mask`: optional Boolean mask selecting trainable output dimensions. +- `fixed_values`: values used for inactive target dimensions. +- `combiner_layers`: hidden widths for the nonrecurrent target head. + +# Returns +- `StateConditionedPolicy` when all dimensions are active. +- `ConstantStatePolicy` when no dimensions are active. +- `FixedOutputPolicy` when only a subset is active. + +# Throws +- `ArgumentError` if bound lengths differ, `active_mask` has the wrong length, + or any upper bound is smaller than its lower bound. + +# Examples +```julia +policy = bounded_state_policy( + 11, + min_volume, + max_volume, + [128, 128]; + combiner_layers = [256, 256], +) +``` +""" +function bounded_state_policy( + n_uncertainty::Int, + lower::AbstractVector, + upper::AbstractVector, + layers::AbstractVector{Int}; + activation = sigmoid, + encoder_type = Flux.LSTM, + active_mask = nothing, + fixed_values = lower, + combiner_layers = Int[], +) + length(lower) == length(upper) || + throw(ArgumentError("lower and upper bound vectors must have the same length")) + n_state = length(lower) + length(fixed_values) == n_state || + throw(ArgumentError("fixed_values length must match state dimension $n_state")) + + scale = upper .- lower + any(<(zero(eltype(scale))), scale) && + throw(ArgumentError("upper bounds must be >= lower bounds")) + + active = if active_mask === nothing + collect(scale .> zero(eltype(scale))) + else + length(active_mask) == n_state || + throw(ArgumentError("active_mask length must match state dimension $n_state")) + collect(Bool.(active_mask)) + end + + if all(active) + return StateConditionedPolicy( + n_uncertainty, n_state, n_state, layers; + activation = activation, + encoder_type = encoder_type, + output_bounds = (lower, upper), + combiner_layers = combiner_layers, + ) + elseif !any(active) + return ConstantStatePolicy(collect(fixed_values), n_uncertainty, n_state) + end + + idx = findall(active) + active_policy = StateConditionedPolicy( + n_uncertainty, n_state, length(idx), layers; + activation = activation, + encoder_type = encoder_type, + output_bounds = (lower[idx], upper[idx]), + combiner_layers = combiner_layers, + ) + template = collect(fixed_values) + template[idx] .= zero(eltype(template)) + expansion = zeros(eltype(template), n_state, length(idx)) + for (j, i) in enumerate(idx) + expansion[i, j] = one(eltype(template)) + end + return FixedOutputPolicy(active_policy, template, expansion) end diff --git a/src/rollout.jl b/src/rollout.jl index f998d5b..3c8dfcd 100644 --- a/src/rollout.jl +++ b/src/rollout.jl @@ -7,35 +7,287 @@ # semantics instead: at each stage, solve one stage, extract the realized next # state, and feed that state into the next policy call. +""" + _target_violation_share(objective::Real, objective_no_target_penalty::Real) -> Float64 + +Compute the fraction of a stage objective attributable to the target-tracking +penalty. + +# Arguments +- `objective::Real`: total stage objective (operational cost + target penalty). +- `objective_no_target_penalty::Real`: stage objective with target penalty + stripped out. + +# Returns +- `Float64`: fraction in ``[0, 1]``, or `NaN` if the ratio is ill-defined. + +# Notes +If the total objective includes both operational cost and a quadratic penalty +``\\lambda \\|x_t - \\hat{x}_t\\|^2``, this function returns + +```math +\\frac{\\text{objective} - \\text{objective\\_no\\_target\\_penalty}}{\\text{objective}}. +``` + +The return value is `NaN` when the objective is non-finite, the penalty is +non-finite, or the objective magnitude is below ``10^{-12}``. +""" function _target_violation_share(objective::Real, objective_no_target_penalty::Real) + # The penalty is the difference between the full and penalty-free objectives. penalty = objective - objective_no_target_penalty + # Guard against non-finite or near-zero denominators that would give NaN/Inf. (isfinite(objective) && isfinite(penalty) && abs(objective) > 1e-12) || return NaN + # Return the penalty's share of the total objective. return penalty / objective end -_cpu_vec(x) = vec(collect(Array(x))) +""" + _to_vec(x) -> AbstractVector + +Flatten `x` into a contiguous one-dimensional vector. +# Arguments +- `x`: any array-like object (matrix, vector, or view). + +# Returns +- `AbstractVector`: a one-dimensional vector with the same elements as `x`. + +# Notes +`SubArray` inputs are materialized with `collect` because downstream solvers +and array operations such as `copyto!` and `vcat` require contiguous storage. +All other array types are reshaped via `vec`. """ - rollout_tsddr(model, initial_state, stage_problem, w_flat; kwargs...) - -Evaluate `model` by solving `stage_problem` sequentially over a materialized -scenario `w_flat`. Unlike deterministic-equivalent evaluation, the optimizer -only receives one uncertainty slice at a time. - -Required callbacks: -- `set_stage_parameters!(stage_problem, state_in, w_t, target, stage)` updates the - one-stage problem before each solve. -- `realized_state(stage_problem, result)` returns the next state realized by the - solved stage problem. - -Optional callbacks: -- `objective_no_target_penalty(stage_problem, result)` returns the stage objective - with target-slack penalty removed. The default is `result.objective`. - -`policy_state = :realized` is the closed-loop deployment semantics. `:target` -keeps the policy recurrence on its own previous target, matching the target -generation used by deterministic-equivalent training while still solving stages -one by one. +_to_vec(x) = vec(x) # reshape in-place for contiguous arrays +_to_vec(x::SubArray) = collect(x) # materialize views to contiguous storage + +""" + _state_bound_vector(bound, ref::AbstractVector) -> Union{Nothing, AbstractVector} + +Convert a user-supplied state bound into a concrete vector whose element type +and device (CPU or GPU) match `ref`. + +# Arguments +- `bound`: the bound specification (`nothing`, a `Real` scalar, or an + `AbstractVector` / `SubArray`). +- `ref::AbstractVector`: reference vector whose length, element type, and + device placement determine the output format. + +# Returns +- `Nothing` if `bound === nothing`. +- `AbstractVector` matching `ref` in length, element type, and device. + +# Throws +- `ArgumentError` if a vector bound has the wrong length or `bound` is an + unsupported type. + +# Notes +Accepted bound forms are `nothing`, a scalar broadcast to all state entries, +or a vector of the same length as `ref`. `SubArray` bounds are materialized via +[`_to_vec`](@ref) before device adaptation so GPU kernels receive contiguous +storage. +""" +function _state_bound_vector(bound, ref::AbstractVector) + # Nothing means "no bound on this side" — propagate the sentinel. + bound === nothing && return nothing + if bound isa AbstractVector || bound isa SubArray + # Vector bounds must be element-wise, so lengths must agree. + length(bound) == length(ref) || + throw(ArgumentError("state bound length must match state length=$(length(ref)), got $(length(bound))")) + # Materialize, cast to ref's element type, and adapt to ref's device. + return _adapt_array(eltype(ref).(_to_vec(bound)), ref) + end + # Scalar bounds are broadcast into a constant vector matching ref's shape. + bound isa Real || + throw(ArgumentError("state bounds must be vectors, scalars, or nothing")) + # Allocate on the same device as ref using `similar`. + out = similar(ref, length(ref)) + # Fill every element with the scalar bound cast to ref's element type. + fill!(out, eltype(ref)(bound)) + return out +end + +""" + _project_state_to_bounds(state::AbstractVector, state_bounds) -> AbstractVector + +Clamp every element of `state` to lie within `[lower, upper]`. + +# Arguments +- `state::AbstractVector`: realized state vector to project. +- `state_bounds`: `nothing` (no projection), or a 2-element tuple/pair + `(lower, upper)` where each side is `nothing`, a scalar, or a vector. + +# Returns +- `AbstractVector`: projected state with the same element type as `state`. + +# Throws +- `ArgumentError` if `state_bounds` is not `nothing` and does not have + exactly two elements. + +# Notes +Given bounds ``l`` and ``u``, the projection is the element-wise clamp + +```math +\\operatorname{proj}(x)_i = \\min\\bigl(\\max(x_i,\\, l_i),\\, u_i\\bigr). +``` + +This repair is intended for solver-tolerance drift at stage interfaces, not as +a substitute for a recourse-feasible model. +""" +function _project_state_to_bounds(state::AbstractVector, state_bounds) + # No bounds means the state passes through unchanged. + state_bounds === nothing && return state + # Bounds must be a (lower, upper) pair. + length(state_bounds) == 2 || + throw(ArgumentError("state_bounds must be a pair/tuple (lower, upper)")) + # Convert each side into a concrete vector (or nothing) matching state. + lower = _state_bound_vector(state_bounds[1], state) + upper = _state_bound_vector(state_bounds[2], state) + # Apply element-wise clamping: first lower, then upper. + projected = state + lower !== nothing && (projected = max.(projected, lower)) # enforce lower bound + upper !== nothing && (projected = min.(projected, upper)) # enforce upper bound + return projected +end + +""" + _project_realized_state(state::AbstractVector, state_bounds, project_state) + -> AbstractVector + +Repair a realized state before it is fed into the next rollout stage. + +# Arguments +- `state::AbstractVector`: raw realized state from the stage solver. +- `state_bounds`: `nothing` or `(lower, upper)` passed to + [`_project_state_to_bounds`](@ref). +- `project_state`: `nothing`, or a callable `f(state) -> projected_state` + that enforces non-box feasibility constraints. + +# Returns +- `AbstractVector`: the doubly-projected state, cast to the element type + of the input `state`. + +# Notes +The function first applies box projection via +[`_project_state_to_bounds`](@ref), then applies `project_state` when a custom +projector is supplied. If `project_state === nothing`, only the box projection +is applied. +""" +function _project_realized_state(state::AbstractVector, state_bounds, project_state) + # First pass: box-clamp the raw realized state. + projected = _project_state_to_bounds(state, state_bounds) + # If no custom projector is provided, the box projection is sufficient. + project_state === nothing && return projected + # Second pass: apply the user's non-box projection and cast back to state's eltype. + return eltype(state).(_to_vec(project_state(projected))) +end + +""" + rollout_tsddr( + model, + initial_state::AbstractVector, + stage_problem, + w_flat::AbstractVector; + horizon::Int, + n_uncertainty::Int, + set_stage_parameters!::Function, + realized_state::Function, + objective_no_target_penalty::Function = (prob, result) -> result.objective, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + policy_state::Symbol = :realized, + solver_state = nothing, + reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, + ) -> Union{Nothing, NamedTuple} + +Evaluate a target-setting decision rule `model` by solving `stage_problem` +sequentially over a materialized uncertainty scenario `w_flat`. + +Unlike the deterministic-equivalent training solve (which sees the full +horizon simultaneously), this rollout mirrors deployment semantics: at each +stage the solver receives only one uncertainty slice, solves, extracts the +realized next state, and feeds that state into the next policy call. + +The stage-wise recursion is + +```math +x_0 = x_{\\text{init}}, \\quad +\\hat{x}_t = \\pi_\\theta(w_t,\\, x_{t-1}), \\quad +x_t = \\operatorname{solve}_t(x_{t-1},\\, w_t,\\, \\hat{x}_t), +``` + +where ``\\pi_\\theta`` is the learned policy (`model`), +``\\operatorname{solve}_t`` solves the single-stage optimization problem, and +``x_t`` is the realized state forwarded to stage ``t+1``. + +# Arguments +- `model`: the target-setting policy ``\\pi_\\theta``. Called as + `model(vcat(w_t, x_{t-1}))` to produce ``\\hat{x}_t``. +- `initial_state::AbstractVector`: initial state ``x_0``. +- `stage_problem`: stage optimization problem (must expose `.model`). +- `w_flat::AbstractVector`: flat uncertainty vector of length + `horizon * n_uncertainty`, sliced into per-stage windows. + +# Keywords +- `horizon::Int`: number of stages ``T``. +- `n_uncertainty::Int`: dimension of each per-stage uncertainty slice. +- `set_stage_parameters!::Function`: callback + `(stage_problem, state, w_t, target, stage) -> nothing` that writes the + current state, uncertainty, and target into the stage problem before each + solve. +- `realized_state::Function`: callback `(stage_problem, result) -> x_t` + that reads the realized next state from the solver result. +- `objective_no_target_penalty::Function`: callback + `(stage_problem, result) -> Float64` returning the stage objective with + the target-tracking penalty removed. Defaults to `result.objective`. +- `madnlp_kwargs`: keyword arguments forwarded to the MadNLP solver + constructor. +- `warmstart::Bool`: whether to warm-start the solver from the previous + stage's dual solution. +- `policy_state::Symbol`: `:realized` feeds the closed-loop realized state + ``x_t`` back to the policy; `:target` feeds the policy's own previous + target ``\\hat{x}_t`` instead, matching deterministic-equivalent training + semantics. +- `solver_state`: optional pre-built solver object to reuse across stages. +- `reuse_solver::Bool`: if `true`, a single solver object is reused (and + warm-started) across all stages. +- `state_bounds`: `nothing` or `(lower, upper)` pair to clamp realized + states via [`_project_state_to_bounds`](@ref). +- `project_state`: `nothing` or a callable `f(state) -> state` for non-box + feasibility repairs via [`_project_realized_state`](@ref). +- `retry_on_failure::Bool`: if `true`, failed or non-finite solves are + retried once with a cold-start solver. + +# Returns +- `nothing` if any stage solve fails after retry. +- A `NamedTuple` with fields: + - `objective::Float64`: cumulative objective over the horizon. + - `objective_no_target_penalty::Float64`: cumulative objective without + target penalties. + - `target_violation_share::Float64`: fraction of objective from target + penalties (see [`_target_violation_share`](@ref)). + - `final_state::AbstractVector`: realized state after the last stage. + - `state_trajectory::Vector`: realized states ``[x_0, x_1, \\ldots, x_T]``. + - `target_trajectory::Vector`: policy targets + ``[\\hat{x}_1, \\ldots, \\hat{x}_T]``. + +# Throws +- `ArgumentError` if `horizon < 1`, `n_uncertainty < 1`, `w_flat` has the + wrong length, or `policy_state` is not `:realized` or `:target`. + +# Examples +```julia +result = rollout_tsddr( + model, x0, stage_problem, w_flat; + horizon = 96, + n_uncertainty = 5, + set_stage_parameters! = my_set_params!, + realized_state = my_realized_state, +) +result !== nothing && @show result.objective +``` """ function rollout_tsddr( model, @@ -52,45 +304,76 @@ function rollout_tsddr( policy_state::Symbol = :realized, solver_state = nothing, reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, ) + # --- Input validation --------------------------------------------------- horizon >= 1 || throw(ArgumentError("horizon must be >= 1")) n_uncertainty >= 1 || throw(ArgumentError("n_uncertainty must be >= 1")) + # w_flat must contain exactly horizon slices of n_uncertainty each. length(w_flat) == horizon * n_uncertainty || throw(ArgumentError("w_flat length must be horizon*n_uncertainty=$(horizon * n_uncertainty), got $(length(w_flat))")) + # Only two feedback modes are supported. policy_state in (:realized, :target) || throw(ArgumentError("policy_state must be :realized or :target, got :$policy_state")) - F = eltype(initial_state) - state = solver_state + # --- Infer numeric types and allocate trajectory storage --------------- + F = eltype(initial_state) # element type (e.g. Float32) + nx = length(initial_state) # state dimension + state = solver_state # optional pre-built solver + w_flat = _adapt_array(F.(w_flat), initial_state) # move w_flat to same device as initial_state + # Reset any recurrent state in the policy network (e.g. LSTM hidden state). Flux.reset!(model) - realized_prev = F.(_cpu_vec(initial_state)) + # x_0: the realized state entering stage 1. + realized_prev = F.(_to_vec(initial_state)) + # Target recurrence state (used when policy_state == :target). target_prev = copy(realized_prev) - state_trajectory = Vector{Vector{F}}(undef, horizon + 1) - target_trajectory = Vector{Vector{F}}(undef, horizon) - state_trajectory[1] = copy(realized_prev) + # Pre-allocate trajectory arrays: states have T+1 entries, targets have T. + state_trajectory = Vector{AbstractVector{F}}(undef, horizon + 1) + target_trajectory = Vector{AbstractVector{F}}(undef, horizon) + state_trajectory[1] = copy(realized_prev) # store x_0 - objective = 0.0 - objective_no_penalty = 0.0 + # Pre-allocate Float64 buffers for the solver interface (solvers use Float64). + state_f64 = similar(initial_state, Float64, nx) # x_{t-1} in Float64 + w_f64 = similar(initial_state, Float64, n_uncertainty) # w_t in Float64 + target_f64 = similar(initial_state, Float64, nx) # xhat_t in Float64 + # Running sums of cumulative cost over the horizon. + objective = 0.0 # total objective (operational + target penalty) + objective_no_penalty = 0.0 # total objective without target penalty + + # --- Stage-wise forward pass ------------------------------------------- for stage in 1:horizon - lo = (stage - 1) * n_uncertainty + 1 - hi = stage * n_uncertainty - wt = F.(_cpu_vec(view(w_flat, lo:hi))) + # Slice out this stage's uncertainty window from the flat vector. + wt = view(w_flat, (stage-1)*n_uncertainty+1 : stage*n_uncertainty) + # Choose the state fed to the policy: realized (closed-loop) or target. policy_input_state = policy_state === :realized ? realized_prev : target_prev + # Evaluate the policy: pi_theta(w_t, x_{t-1}) -> xhat_t. target = model(vcat(wt, policy_input_state)) - target_trajectory[stage] = F.(_cpu_vec(target)) + # Flatten the target to a 1-D vector for downstream use. + target_vec = _to_vec(target) + # Store the target in the trajectory (cast to the state element type). + target_trajectory[stage] = F.(target_vec) + # Copy inputs into the Float64 solver buffers. + copyto!(state_f64, realized_prev) # x_{t-1} + copyto!(w_f64, wt) # w_t + copyto!(target_f64, target_vec) # xhat_t + # Write the current state, uncertainty, and target into the stage problem. set_stage_parameters!( stage_problem, - Float64.(realized_prev), - Float64.(wt), - Float64.(_cpu_vec(target)), + state_f64, + w_f64, + target_f64, stage, ) + # --- Solve the single-stage optimization problem ------------------- if reuse_solver || state !== nothing + # Reuse an existing solver; create one lazily on first call. state === nothing && (state = _make_solver(stage_problem.model, madnlp_kwargs)) result = _solve!( state, @@ -99,6 +382,7 @@ function rollout_tsddr( madnlp_kwargs = madnlp_kwargs, ) else + # Create a fresh solver for this stage (no cross-stage warm-start). stage_state = _make_solver(stage_problem.model, madnlp_kwargs) result = _solve!( stage_state, @@ -108,53 +392,226 @@ function rollout_tsddr( ) end + # --- Retry logic: cold-start fallback on solver failure ------------ + if retry_on_failure && (!solve_succeeded(result) || !isfinite(result.objective)) + # Build a fresh solver and retry without warm-start. + retry_state = _make_solver(stage_problem.model, madnlp_kwargs) + result = _solve!( + retry_state, + stage_problem.model; + warmstart = false, + madnlp_kwargs = madnlp_kwargs, + ) + end + + # Abort the rollout if the solve still failed after retry. solve_succeeded(result) || return nothing + # Abort on non-finite objective (e.g. MadNLP returning 0.0 for infeasible). isfinite(result.objective) || return nothing + # --- Extract penalty-free objective, with retry -------------------- no_penalty = objective_no_target_penalty(stage_problem, result) + if retry_on_failure && !isfinite(no_penalty) + # Non-finite penalty-free cost can indicate a solver glitch; retry. + retry_state = _make_solver(stage_problem.model, madnlp_kwargs) + result = _solve!( + retry_state, + stage_problem.model; + warmstart = false, + madnlp_kwargs = madnlp_kwargs, + ) + solve_succeeded(result) || return nothing + isfinite(result.objective) || return nothing + no_penalty = objective_no_target_penalty(stage_problem, result) + end + # Final guard: abort if the penalty-free cost is still non-finite. isfinite(no_penalty) || return nothing - objective += result.objective - objective_no_penalty += no_penalty - realized_prev = F.(_cpu_vec(realized_state(stage_problem, result))) - target_prev = F.(_cpu_vec(target)) + # --- Accumulate costs and advance the state ------------------------ + objective += result.objective # add stage cost to cumulative total + objective_no_penalty += no_penalty # add penalty-free stage cost + # Read the realized next state x_t from the solver solution. + raw_realized = F.(_to_vec(realized_state(stage_problem, result))) + # Project x_t to feasibility (box bounds + custom projector). + realized_prev = _project_realized_state(raw_realized, state_bounds, project_state) + # Update the target recurrence for :target mode. + target_prev = target_trajectory[stage] + # Record x_t in the state trajectory. state_trajectory[stage + 1] = copy(realized_prev) end + # --- Assemble and return the rollout summary --------------------------- return ( - objective = objective, - objective_no_target_penalty = objective_no_penalty, - target_violation_share = _target_violation_share(objective, objective_no_penalty), - final_state = realized_prev, - state_trajectory = state_trajectory, - target_trajectory = target_trajectory, + objective = objective, # sum of stage objectives + objective_no_target_penalty = objective_no_penalty, # sum minus target penalties + target_violation_share = _target_violation_share(objective, objective_no_penalty), # penalty fraction + final_state = realized_prev, # x_T + state_trajectory = state_trajectory, # [x_0, ..., x_T] + target_trajectory = target_trajectory, # [xhat_1, ..., xhat_T] ) end +""" + RolloutEvaluation + +Store configuration and mutable summaries for periodic rollout evaluation. + +# Fields +- `stage_problem`: the single-stage optimization problem template. +- `initial_state`: initial state ``x_0`` for every rollout. +- `scenarios::Vector`: pre-sampled uncertainty vectors, each of length + `horizon * n_uncertainty`. +- `horizon::Int`: number of stages ``T``. +- `n_uncertainty::Int`: per-stage uncertainty dimension. +- `set_stage_parameters!::Function`: callback to write stage data into the + problem (see [`rollout_tsddr`](@ref)). +- `realized_state::Function`: callback to extract ``x_t`` from a solve + result. +- `objective_no_target_penalty::Function`: callback to extract the + penalty-free stage cost. +- `madnlp_kwargs`: keyword arguments forwarded to MadNLP. +- `warmstart::Bool`: whether to warm-start across stages. +- `stride::Int`: evaluate every `stride` training iterations. +- `policy_state::Symbol`: `:realized` or `:target` (see + [`rollout_tsddr`](@ref)). +- `solver_state`: pre-built solver object (or `nothing`). +- `reuse_solver::Bool`: whether to reuse a single solver across stages. +- `state_bounds`: optional `(lower, upper)` feasibility bounds. +- `project_state`: optional non-box projection callback. +- `retry_on_failure::Bool`: retry failed solves with a cold start. +- `stage_problem_pool::Vector`: pool of stage problems for parallel + evaluation. +- `active_scenarios::Int`: number of scenarios to evaluate (at most + `length(scenarios)`). +- `last_objective::Float64`: mean objective across successful scenarios. +- `last_objective_no_target_penalty::Float64`: mean penalty-free objective. +- `last_violation_share::Float64`: mean target-violation share. +- `last_n_ok::Int`: number of scenarios that solved successfully. +- `last_scenario_data::Vector{Any}`: per-scenario `(index, result)` pairs + from the most recent evaluation. + +# Notes +The struct is callable as `evaluation(iter, model)`. At every `stride`-th +iteration, it runs [`rollout_tsddr`](@ref) on up to `active_scenarios` +scenarios and updates the mutable summary fields. When `stage_problem_pool` +contains more than one entry, scenarios are distributed across the pool with +`Threads.@spawn`. + +Thread-safety: the single `set_stage_parameters!`, `realized_state`, and +`objective_no_target_penalty` callbacks are shared across all spawned tasks +while each task receives its own stage problem from the pool. When +`stage_problem_pool` has more than one entry, these callbacks must therefore be +thread-safe and must write only into the stage problem they are handed — any +shared mutable buffer (e.g. a captured scratch array reused across calls) +races across tasks and silently corrupts results. +""" mutable struct RolloutEvaluation <: Function - stage_problem - initial_state - scenarios::Vector - horizon::Int - n_uncertainty::Int - set_stage_parameters!::Function - realized_state::Function - objective_no_target_penalty::Function - madnlp_kwargs - warmstart::Bool - stride::Int - policy_state::Symbol - solver_state - reuse_solver::Bool - stage_problem_pool::Vector # pool of stage problems for parallel evaluation - active_scenarios::Int # how many scenarios to evaluate (≤ length(scenarios)) - last_objective::Float64 - last_objective_no_target_penalty::Float64 - last_violation_share::Float64 - last_n_ok::Int - last_scenario_data::Vector{Any} + stage_problem # single-stage optimization problem template + initial_state # initial state x_0 for all rollouts + scenarios::Vector # pre-sampled uncertainty vectors + horizon::Int # number of decision stages T + n_uncertainty::Int # per-stage uncertainty dimension + set_stage_parameters!::Function # callback: write stage data into the problem + realized_state::Function # callback: extract realized x_t from result + objective_no_target_penalty::Function # callback: penalty-free stage cost + madnlp_kwargs # solver keyword arguments + warmstart::Bool # warm-start across stages + stride::Int # evaluate every stride-th iteration + policy_state::Symbol # :realized or :target feedback mode + solver_state # pre-built solver (or nothing) + reuse_solver::Bool # reuse single solver across stages + state_bounds # (lower, upper) feasibility bounds or nothing + project_state # non-box projection callback or nothing + retry_on_failure::Bool # retry failed solves with cold start + stage_problem_pool::Vector # pool of stage problems for parallel evaluation + active_scenarios::Int # how many scenarios to evaluate (leq length(scenarios)) + last_objective::Float64 # mean objective from last evaluation + last_objective_no_target_penalty::Float64 # mean penalty-free objective from last evaluation + last_violation_share::Float64 # mean target-violation share from last evaluation + last_n_ok::Int # number of successful scenarios in last evaluation + last_scenario_data::Vector{Any} # per-scenario (index, result) pairs from last eval end +""" + RolloutEvaluation( + stage_problem, + initial_state, + scenarios; + horizon::Int, + n_uncertainty::Int, + set_stage_parameters!::Function, + realized_state::Function, + objective_no_target_penalty::Function = (prob, result) -> result.objective, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + stride::Int = 1, + policy_state::Symbol = :realized, + reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, + stage_problem_pool::Vector = [], + active_scenarios::Int = length(scenarios), + ) -> RolloutEvaluation + +Construct a [`RolloutEvaluation`](@ref) callback for periodic out-of-sample +policy evaluation during training. + +All mutable result fields (`last_objective`, `last_n_ok`, etc.) are +initialized to `NaN` / `0` / empty and are populated on the first call. + +# Arguments +- `stage_problem`: the single-stage optimization problem (must expose + `.model`). +- `initial_state`: initial state ``x_0``. +- `scenarios`: iterable of pre-sampled uncertainty vectors, each of length + `horizon * n_uncertainty`. + +# Keywords +- `horizon::Int`: number of stages ``T``. +- `n_uncertainty::Int`: per-stage uncertainty dimension. +- `set_stage_parameters!::Function`: stage-parameter callback. +- `realized_state::Function`: realized-state extraction callback. +- `objective_no_target_penalty::Function`: penalty-free cost callback. +- `madnlp_kwargs`: solver keyword arguments. +- `warmstart::Bool`: warm-start across stages within each rollout. +- `stride::Int`: evaluate every `stride`-th training iteration. +- `policy_state::Symbol`: `:realized` (closed-loop) or `:target`. +- `reuse_solver::Bool`: reuse a single solver across stages. +- `state_bounds`: `nothing` or `(lower, upper)` for box projection. +- `project_state`: `nothing` or custom non-box projector. +- `retry_on_failure::Bool`: retry failed solves with cold start. +- `stage_problem_pool::Vector`: pool of independent stage problems for + multi-threaded evaluation. An empty pool uses sequential evaluation. + With more than one pool entry, the shared `set_stage_parameters!`, + `realized_state`, and `objective_no_target_penalty` callbacks run + concurrently on different tasks: they must be thread-safe and write only + into the stage problem passed to them (no shared mutable buffers), + otherwise results race. +- `active_scenarios::Int`: cap on how many scenarios to evaluate (defaults + to all). + +# Returns +- `RolloutEvaluation`: callable training callback with empty result summaries. + +# Throws +- `ArgumentError` if `scenarios` is empty, `stride < 1`, or `policy_state` is + not `:realized` or `:target`. + +# Examples +```julia +eval_cb = RolloutEvaluation( + stage_problem, x0, test_scenarios; + horizon = 96, + n_uncertainty = 5, + set_stage_parameters! = my_set_params!, + realized_state = my_realized_state, + stride = 10, +) +# Use as a training callback: +eval_cb(iter, model) +``` +""" function RolloutEvaluation( stage_problem, initial_state, @@ -169,18 +626,24 @@ function RolloutEvaluation( stride::Int = 1, policy_state::Symbol = :realized, reuse_solver::Bool = false, + state_bounds = nothing, + project_state = nothing, + retry_on_failure::Bool = true, stage_problem_pool::Vector = [], active_scenarios::Int = length(scenarios), ) + # At least one scenario is required for meaningful evaluation. isempty(scenarios) && throw(ArgumentError("scenarios must be nonempty")) + # Stride must be positive; stride=1 evaluates every iteration. stride >= 1 || throw(ArgumentError("stride must be >= 1")) + # Validate the feedback mode. policy_state in (:realized, :target) || throw(ArgumentError("policy_state must be :realized or :target, got :$policy_state")) return RolloutEvaluation( stage_problem, initial_state, - collect(scenarios), + collect(scenarios), # materialize to a concrete Vector horizon, n_uncertainty, set_stage_parameters!, @@ -190,18 +653,42 @@ function RolloutEvaluation( warmstart, stride, policy_state, + # Pre-build the solver if reuse is requested; otherwise leave as nothing. reuse_solver ? _make_solver(stage_problem.model, madnlp_kwargs) : nothing, reuse_solver, - collect(stage_problem_pool), + state_bounds, + project_state, + retry_on_failure, + collect(stage_problem_pool), # materialize the pool to a concrete Vector active_scenarios, - NaN, - NaN, - NaN, - 0, - Any[], + NaN, # last_objective: not yet evaluated + NaN, # last_objective_no_target_penalty + NaN, # last_violation_share + 0, # last_n_ok: no successful scenarios yet + Any[], # last_scenario_data: empty until first eval ) end +""" + (evaluation::RolloutEvaluation)(iter, model) -> Nothing + +Evaluate `model` on rollout scenarios when `iter` is aligned with +`evaluation.stride`. + +# Arguments +- `evaluation::RolloutEvaluation`: callback state and rollout configuration. +- `iter`: current training iteration. +- `model`: policy model passed to [`rollout_tsddr`](@ref). + +# Returns +- `nothing`: summary fields on `evaluation` are updated in place. + +# Notes +If `iter % evaluation.stride != 0`, the method only clears stale per-scenario +data and returns. Otherwise it records mean objective, mean penalty-free +objective, mean target-violation share, the number of successful scenarios, and +per-scenario results. +""" function (evaluation::RolloutEvaluation)(iter, model) empty!(evaluation.last_scenario_data) iter % evaluation.stride == 0 || return nothing @@ -232,6 +719,9 @@ function (evaluation::RolloutEvaluation)(iter, model) policy_state = evaluation.policy_state, solver_state = evaluation.solver_state, reuse_solver = evaluation.reuse_solver, + state_bounds = evaluation.state_bounds, + project_state = evaluation.project_state, + retry_on_failure = evaluation.retry_on_failure, ) result === nothing && continue total += result.objective @@ -264,6 +754,9 @@ function (evaluation::RolloutEvaluation)(iter, model) warmstart = evaluation.warmstart, policy_state = evaluation.policy_state, reuse_solver = false, + state_bounds = evaluation.state_bounds, + project_state = evaluation.project_state, + retry_on_failure = evaluation.retry_on_failure, ) push!(tasks, t) end @@ -298,13 +791,30 @@ function (evaluation::RolloutEvaluation)(iter, model) end """ - critic_samples_from_evaluation(eval; objective_key) -> Vector{CriticSample} + critic_samples_from_evaluation( + eval_obj::RolloutEvaluation; + objective_key::Symbol = :objective, + ) -> Vector{CriticSample} Convert the last rollout evaluation results into `CriticSample`s for critic -training. Target multipliers are zero (rollout evaluation does not produce -duals), so these samples contribute only to the value loss term. By default the -critic target uses the full rollout objective; pass -`objective_key = :objective_no_target_penalty` to remove target-slack penalties. +training. + +# Arguments +- `eval_obj::RolloutEvaluation`: evaluation callback containing + `last_scenario_data` from a previous call. + +# Keywords +- `objective_key::Symbol`: field of each rollout result used as the scalar + critic target; commonly `:objective` or `:objective_no_target_penalty`. + +# Returns +- `Vector{CriticSample}`: one sample per successful rollout scenario from the + last evaluation. + +# Notes +Rollout evaluation does not produce dual multipliers, so generated samples use +zero target multipliers and contribute only to the value-loss term unless +combined with other samples. """ function critic_samples_from_evaluation( eval_obj::RolloutEvaluation; diff --git a/src/training.jl b/src/training.jl index 0407ef6..78d5e01 100644 --- a/src/training.jl +++ b/src/training.jl @@ -6,7 +6,7 @@ # 1. uncertainty_sampler() → flat w (length T × nw_per_stage) # 2. Policy rollout: x̂_t = policy(vcat(w_t, x̂_{t-1})) for t = 1..T # 3. ExaModels.set_parameter! for x0, uncertainty, targets → MadNLP.solve! -# 4. λ = result.multipliers[target_con_range] (∇_{x̂} Q, envelope theorem) +# 4. λ = target_multipliers(de, result) (∇_{x̂} Q, envelope theorem) # 5. Zygote: ∇_θ (1/n) Σ_s ⟨λ_s, x̂_s(θ)⟩ → Flux.update! # # The user passes parameter objects (p_x0, p_target, p_uncertainty) exactly as @@ -21,35 +21,265 @@ # ── Gradient materialization ────────────────────────────────────────────────── -_mat(x) = x -_mat(x::Zygote.OneElement) = collect(x) +""" + _mat(x) + _mat(x::Zygote.OneElement) + _mat(x::ChainRulesCore.Tangent) + _mat(x::ChainRulesCore.MutableTangent) + +Recursively materialize lazy Zygote / ChainRules tangent wrappers into plain +Julia values (arrays and named tuples). + +Zygote may return `OneElement` sparse arrays or `Tangent`/`MutableTangent` +wrappers instead of dense arrays or named tuples. `Flux.update!` expects +concrete data, so every tangent node must be materialized before the optimizer +step. + +# Arguments +- `x`: a gradient value — may be a plain array, a `Zygote.OneElement`, a + `ChainRulesCore.Tangent`, or a `ChainRulesCore.MutableTangent`. + +# Returns +- For plain values: returns `x` unchanged. +- For `OneElement`: returns `collect(x)`, a dense array. +- For `Tangent`/`MutableTangent`: returns a `NamedTuple` with recursively + materialized fields. +- For `Base.RefValue` (Zygote's wrapper for tangents of mutated mutable + structs, such as the state-threading policies): unwraps and recurses. +- For plain `NamedTuple`/`Tuple`: recurses so nested wrappers are stripped. +- For `NoTangent`/`ZeroTangent`: returns `nothing`. +""" +_mat(x) = x # plain value — no conversion needed +_mat(x::Zygote.OneElement) = collect(x) # sparse one-hot → dense array function _mat(x::ChainRulesCore.Tangent{<:Any}) - nt = ChainRulesCore.backing(x) - return NamedTuple{keys(nt)}(map(_mat, values(nt))) + nt = ChainRulesCore.backing(x) # extract underlying named tuple + return NamedTuple{keys(nt)}(map(_mat, values(nt))) # recursively materialize each field end function _mat(x::ChainRulesCore.MutableTangent{<:Any}) - nt = ChainRulesCore.backing(x) - return NamedTuple{keys(nt)}(map(_mat, values(nt))) + nt = ChainRulesCore.backing(x) # extract underlying named tuple + return NamedTuple{keys(nt)}(map(_mat, values(nt))) # recursively materialize each field +end +# Structural-zero tangents (non-differentiable fields) map to nothing so +# Flux.update! skips them. +_mat(::ChainRulesCore.NoTangent) = nothing +_mat(::ChainRulesCore.ZeroTangent) = nothing +# Zygote wraps tangents of MUTATED mutable structs (e.g. the state-threading +# policies, whose forward pass stores the new recurrent state via setfield!) in +# Base.RefValue and MutableTangent containers, potentially nested inside plain +# NamedTuples/Tuples. Recurse through those containers so every wrapper is +# stripped before the gradient reaches Flux.update!. +_mat(ref::Base.RefValue) = _mat(ref[]) # unwrap Ref and recurse +_mat(nt::NamedTuple{K}) where {K} = + NamedTuple{K}(map(_mat, values(nt))) # recurse plain named tuples +_mat(t::Tuple) = map(_mat, t) # recurse plain tuples + +""" + materialize_tangent(g) -> Union{Nothing, Any} + +Convert a Zygote gradient `g` into plain Julia arrays and named tuples that +`Flux.update!` can consume. Returns `nothing` when `g` is `nothing` (i.e., no +gradient was produced for that parameter). + +This is the public entry point; internally it delegates to [`_mat`](@ref). + +# Arguments +- `g`: raw gradient returned by `Zygote.gradient`; may be `nothing`. + +# Returns +- `nothing` if `g` is `nothing`. +- A materialized gradient (dense arrays / named tuples) otherwise. + +# Examples +```julia +gs = Zygote.gradient(model) do m + sum(m(x)) end -materialize_tangent(g) = isnothing(g) ? nothing : _mat(g) +grad = materialize_tangent(gs[1]) # NamedTuple or nothing +``` +""" +materialize_tangent(g) = isnothing(g) ? nothing : _mat(g) # guard against nothing gradients -_all_finite_gradient(x::AbstractArray) = all(isfinite, x) -_all_finite_gradient(x::Number) = isfinite(x) -_all_finite_gradient(x::Nothing) = true -_all_finite_gradient(x::NamedTuple) = all(_all_finite_gradient(v) for v in values(x)) -_all_finite_gradient(x::Tuple) = all(_all_finite_gradient(v) for v in x) -_all_finite_gradient(x) = true +""" + _all_finite_gradient(x) -> Bool + +Recursively check that every element of a (possibly nested) gradient structure +is finite (no `NaN` or `±Inf`). + +After BPTT through physics-based dynamics, accumulated Jacobian products can +overflow `Float32` range (> 3.4e38), producing `Inf` or `NaN` values that +would corrupt the Adam optimizer state. This guard prevents +`Flux.update!` from being called with non-finite gradients. + +# Arguments +- `x`: gradient value — may be an `AbstractArray`, `Number`, `Nothing`, + `NamedTuple`, `Tuple`, or any other type. + +# Returns +- `true` if every numeric leaf is finite (or the value is `nothing` / an + unrecognized type that carries no numeric data). +- `false` if any leaf contains `NaN` or `±Inf`. + +# Examples +```julia +_all_finite_gradient([1.0, 2.0]) # true +_all_finite_gradient([1.0, NaN]) # false +_all_finite_gradient((a=[1.0], b=Inf)) # false +_all_finite_gradient(nothing) # true +``` +""" +_all_finite_gradient(x::AbstractArray) = all(isfinite, x) # check every element +_all_finite_gradient(x::Number) = isfinite(x) # scalar check +_all_finite_gradient(x::Nothing) = true # nothing → vacuously finite +_all_finite_gradient(x::NamedTuple) = all(_all_finite_gradient(v) for v in values(x)) # recurse named tuple fields +_all_finite_gradient(x::Tuple) = all(_all_finite_gradient(v) for v in x) # recurse tuple elements +_all_finite_gradient(x) = true # fallback: assume finite for unknown types + +""" + _status_key(status) -> String + +Convert a MadNLP solver status enum value into a safe string key by replacing +non-alphanumeric characters with underscores. Used to build human-readable +diagnostic dictionaries keyed by solver outcome. + +# Arguments +- `status`: a MadNLP status enum (e.g., `MadNLP.SOLVE_SUCCEEDED`). + +# Returns +- A sanitized `String` suitable for use as a dictionary key. + +# Examples +```julia +_status_key(MadNLP.SOLVE_SUCCEEDED) # "SOLVE_SUCCEEDED" +``` +""" +_status_key(status) = replace(string(status), r"[^A-Za-z0-9_]" => "_") # sanitize status to dict-safe key + +""" + _inc_status!(counts::Dict{String, Int}, status) -> Dict{String, Int} + +Increment the counter for the given MadNLP solver `status` in `counts`. +Converts `status` to a string key via [`_status_key`](@ref) before +incrementing. + +# Arguments +- `counts::Dict{String, Int}`: mutable dictionary of status counts. +- `status`: a MadNLP solver status enum. + +# Returns +- The mutated `counts` dictionary. +""" +function _inc_status!(counts::Dict{String, Int}, status) + key = _status_key(status) # convert enum to string key + counts[key] = get(counts, key, 0) + 1 # increment (initialize to 0 if absent) + return counts +end + +""" + _inc_count!(counts::Dict{String, Int}, key::String) -> Dict{String, Int} + +Increment the counter for `key` in the diagnostics dictionary `counts`. +Used to track failure reasons (e.g., `"nonfinite_objective"`) and retry +outcomes during training batches. + +# Arguments +- `counts::Dict{String, Int}`: mutable dictionary of event counts. +- `key::String`: the event identifier to increment. + +# Returns +- The mutated `counts` dictionary. +""" +function _inc_count!(counts::Dict{String, Int}, key::String) + counts[key] = get(counts, key, 0) + 1 # increment (initialize to 0 if absent) + return counts +end + +""" + _adapt_array(x::AbstractVector, ref::AbstractVector) -> AbstractVector + +Move `x` onto the same device (CPU or GPU) as `ref`, allocating a new array +only when the concrete types differ. + +When training runs on GPU, solver outputs (multipliers, states) may be +`CuVector`s while sampled data may arrive as CPU `Vector`s (or vice versa). +This helper ensures type-homogeneous arithmetic by copying `x` into a +`similar` array derived from `ref`. + +# Arguments +- `x::AbstractVector`: the source data to adapt (may be CPU or GPU). +- `ref::AbstractVector`: a reference vector whose concrete type determines the + target device / storage backend. + +# Returns +- `x` itself when `typeof(x) === typeof(ref)` (zero-copy fast path). +- A new array on the same device as `ref`, with the element type of `x` and + the data of `x` copied in. + +# Examples +```julia +using CUDA +cpu_vec = [1.0f0, 2.0f0] +gpu_ref = CUDA.zeros(Float32, 3) +gpu_vec = _adapt_array(cpu_vec, gpu_ref) # CuVector{Float32} +``` +""" +function _adapt_array(x::AbstractVector, ref::AbstractVector) + typeof(x) === typeof(ref) && return x # same type → no copy needed + copyto!(similar(ref, eltype(x), length(x)), x) # allocate on ref's device, copy x into it +end # ── Solve-status check ──────────────────────────────────────────────────────── """ solve_succeeded(result) -> Bool + +Check whether a MadNLP solve result indicates a usable solution. + +MadNLP returns `0.0` objective for failed or infeasible solves, so checking +`isfinite(objective)` alone is not sufficient. This function inspects the +solver status directly. + +# Arguments +- `result`: a MadNLP result object with a `.status` field. + +# Returns +- `true` if `result.status` is `SOLVE_SUCCEEDED` or `SOLVED_TO_ACCEPTABLE_LEVEL`. +- `false` for all other statuses (e.g., `MAXIMUM_ITERATIONS_EXCEEDED`, + `INFEASIBLE_PROBLEM_DETECTED`). + +# Examples +```julia +result = MadNLP.solve!(solver) +if solve_succeeded(result) + @show result.objective +end +``` """ function solve_succeeded(result) - s = result.status - return s == MadNLP.SOLVE_SUCCEEDED || s == MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL + s = result.status # extract MadNLP status enum + return s == MadNLP.SOLVE_SUCCEEDED || s == MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL # accept both convergence levels end +""" + prepare_solve!(de, init_state, w_flat, xhat_flat) + +Hook called after standard parameter updates and before each NLP solve. + +# Arguments +- `de`: deterministic-equivalent problem. +- `init_state`: initial state used for the current sample. +- `w_flat`: flat uncertainty trajectory for the current sample. +- `xhat_flat`: flat target trajectory for the current sample. + +# Returns +- `nothing` by default. + +# Notes +Override this method for problem types that need additional parameter updates, +for example setting a reservoir parameter from `x0` and targets when the +reservoir is not a decision variable. +""" +prepare_solve!(de, init_state, w_flat, xhat_flat) = nothing + # ── Internal: one MadNLP solve with cascade-failure prevention ──────────────── # # After a failed solve the duals (y, zl, zu) are corrupted. Instead of cold- @@ -60,68 +290,256 @@ end # Per-solve iteration budget: MadNLP's cnt.k is CUMULATIVE across calls, so we # reset it before each solve to give each batch a fresh max_iter budget. +""" + _SolverState + +Mutable wrapper around a MadNLP solver that caches the last successful +primal-dual snapshot for warm-start cascade-failure prevention. + +After a failed solve, MadNLP's duals (`y`, `zl`, `zu`) are corrupted. +If the next call uses `reinitialize!()` (warm-start path), it keeps those +corrupted duals and the failure cascades. `_SolverState` stores the +last-good dual values so they can be restored after a failure, breaking +the cascade without paying the cost of a full cold start. + +# Fields +- `solver`: the `MadNLP.MadNLPSolver` instance. +- `last_good_x`: primal snapshot from the last successful solve (CPU or GPU + array), or `nothing` before the first success. +- `last_good_y`: equality dual snapshot, or `nothing`. +- `last_good_zl_vals`: lower-bound dual values snapshot, or `nothing`. +- `last_good_zu_vals`: upper-bound dual values snapshot, or `nothing`. +- `has_fixed_vars::Bool`: `true` when the NLP has fixed variables (via + `MakeParameter`), which requires a fresh solver per solve to avoid stale + KKT factorization state. +""" mutable struct _SolverState - solver - last_good_x # primal snapshot (CPU or GPU array), or nothing - last_good_y # dual snapshot, or nothing - last_good_zl_vals - last_good_zu_vals + solver # MadNLP.MadNLPSolver instance + last_good_x # primal snapshot (CPU or GPU array), or nothing + last_good_y # dual snapshot, or nothing + last_good_zl_vals # lower-bound dual values snapshot, or nothing + last_good_zu_vals # upper-bound dual values snapshot, or nothing + has_fixed_vars::Bool # true when fixed variables exist in the NLP end +""" + _make_solver(nlp, madnlp_kwargs) -> _SolverState + +Construct a [`_SolverState`](@ref) wrapping a fresh `MadNLP.MadNLPSolver`. + +Detects whether the NLP has fixed variables by comparing the solver's internal +variable count against the NLP's variable count. When fixed variables exist +(via `ExaModels.MakeParameter`), the solver cannot be safely reused across +parameter changes, so `has_fixed_vars` is set to `true`. + +# Arguments +- `nlp`: an NLPModels-compatible problem (e.g., `ExaModel`). +- `madnlp_kwargs`: `NamedTuple` of keyword arguments forwarded to + `MadNLP.MadNLPSolver`. + +# Returns +- A fresh [`_SolverState`](@ref) with no cached primal-dual snapshot. + +# Examples +```julia +state = _make_solver(det_equivalent.model, (print_level=0, tol=1e-6)) +``` +""" function _make_solver(nlp, madnlp_kwargs) - solver = MadNLP.MadNLPSolver(nlp; madnlp_kwargs...) - return _SolverState(solver, nothing, nothing, nothing, nothing) + solver = MadNLP.MadNLPSolver(nlp; madnlp_kwargs...) # build MadNLP solver with user options + nvar_solver = length(solver.x.x) # internal (reduced) variable count + nvar_nlp = length(NLPModels.get_x0(nlp)) # NLP-level variable count + has_fixed = nvar_solver != nvar_nlp # mismatch → fixed variables present + return _SolverState(solver, nothing, nothing, nothing, nothing, has_fixed) # no cached duals yet end +""" + _solve!(state::_SolverState, nlp; warmstart::Bool, madnlp_kwargs) + +Solve `nlp` using the MadNLP solver cached in `state`, with cascade-failure +prevention via dual snapshot restore. + +The warm-start logic implements three paths: + +1. **Fixed variables** (`state.has_fixed_vars`): create a fresh solver each + call because MadNLP's KKT factorization becomes stale when `MakeParameter` + changes fixed-variable values between solves. + +2. **Warm start after success**: copy the last-good primal into `x0` and let + `reinitialize!()` keep the cached duals. + +3. **Warm start after failure**: restore the last-good dual snapshot + (`y`, `zl.values`, `zu.values`) and mark the solver as `SOLVE_SUCCEEDED` + so `reinitialize!()` keeps those restored duals rather than the corrupted + ones. If no good snapshot exists, fall back to `INITIAL` (cold start). + +MadNLP's `cnt.k` is cumulative and never reset internally, so we reset it +before each solve to give every call a fresh `max_iter` budget. + +# Arguments +- `state::_SolverState`: solver wrapper with optional cached primal-dual + snapshot. +- `nlp`: the NLPModels-compatible problem to solve. +- `warmstart::Bool`: `true` to reuse duals from prior solves, `false` to + cold-start. +- `madnlp_kwargs`: `NamedTuple` forwarded to `MadNLP.solve!`. + +# Returns +- A MadNLP result object with fields `.status`, `.objective`, + `.multipliers`, `.solution`. + +# Examples +```julia +state = _make_solver(nlp, (print_level=0,)) +result = _solve!(state, nlp; warmstart=true, madnlp_kwargs=(print_level=0,)) +``` +""" function _solve!(state::_SolverState, nlp; warmstart::Bool, madnlp_kwargs) - solver = state.solver + solver = state.solver # cached MadNLP solver instance - # Primal warm-start: seed NLPModel's x0 from last-good primal. + # MadNLP solver reuse with MakeParameter (fixed variables) causes INFEASIBLE + # on subsequent solves even with INITIAL status — stale KKT factorization state. + # Fix: create a fresh solver each time when fixed variables exist. + if state.has_fixed_vars + return MadNLP.madnlp(nlp; madnlp_kwargs...) # one-shot fresh solver + end + + # Normal path (no fixed variables): full warm-start support. if warmstart && state.last_good_x !== nothing - copyto!(NLPModels.get_x0(nlp), state.last_good_x) + copyto!(NLPModels.get_x0(nlp), state.last_good_x) # seed primal with last-good solution end - prev_result = solver.status - prev_failed = (prev_result != MadNLP.INITIAL && - prev_result != MadNLP.SOLVE_SUCCEEDED && + prev_result = solver.status # check previous solve outcome + prev_failed = (prev_result != MadNLP.INITIAL && # any non-success, non-initial status + prev_result != MadNLP.SOLVE_SUCCEEDED && # means the duals may be corrupted prev_result != MadNLP.SOLVED_TO_ACCEPTABLE_LEVEL) if !warmstart - solver.status = MadNLP.INITIAL + solver.status = MadNLP.INITIAL # cold start: reset x, y, zl, zu elseif prev_failed && state.last_good_y !== nothing - solver.y .= state.last_good_y - solver.zl.values .= state.last_good_zl_vals - solver.zu.values .= state.last_good_zu_vals - solver.status = MadNLP.SOLVE_SUCCEEDED + solver.y .= state.last_good_y # restore last-good equality duals + solver.zl.values .= state.last_good_zl_vals # restore last-good lower-bound duals + solver.zu.values .= state.last_good_zu_vals # restore last-good upper-bound duals + solver.status = MadNLP.SOLVE_SUCCEEDED # trick reinitialize!() into warm path elseif prev_failed - solver.status = MadNLP.INITIAL + solver.status = MadNLP.INITIAL # no snapshot available → cold start end - # Reset per-solve iteration budget. - solver.cnt.k = 0 - solver.cnt.acceptable_cnt = 0 - solver.cnt.start_time = time() + # Reset per-solve iteration budget (cnt.k is cumulative in MadNLP). + solver.cnt.k = 0 # reset iteration counter + solver.cnt.acceptable_cnt = 0 # reset acceptable-step counter + solver.cnt.start_time = time() # reset wall-clock timer - res = MadNLP.solve!(solver; madnlp_kwargs...) + res = MadNLP.solve!(solver; madnlp_kwargs...) # run the solver if solve_succeeded(res) - state.last_good_x = copy(solver.x.x) - state.last_good_y = copy(solver.y) - state.last_good_zl_vals = copy(solver.zl.values) - state.last_good_zu_vals = copy(solver.zu.values) + state.last_good_x = copy(solver.x.x) # snapshot primal (GPU-safe copy) + state.last_good_y = copy(solver.y) # snapshot equality duals + state.last_good_zl_vals = copy(solver.zl.values) # snapshot lower-bound duals + state.last_good_zu_vals = copy(solver.zu.values) # snapshot upper-bound duals end return res end +""" + _solve_with_retry!(state::_SolverState, nlp; + warmstart::Bool, madnlp_kwargs, retry_on_failure::Bool) + -> (result, retried::Bool) + +Solve `nlp` via [`_solve!`](@ref), optionally retrying with a fresh cold-start +solver if the first attempt fails or returns a non-finite objective. + +The retry creates a brand-new [`_SolverState`](@ref) (discarding any corrupted +internal state) and solves with `warmstart=false`. This is more expensive than +the dual-restore path in `_solve!` but guarantees a clean factorization. + +# Arguments +- `state::_SolverState`: primary solver state (may have cached duals). +- `nlp`: the NLPModels-compatible problem. +- `warmstart::Bool`: whether the first attempt should warm-start. +- `madnlp_kwargs`: `NamedTuple` forwarded to `MadNLP.solve!`. +- `retry_on_failure::Bool`: if `true`, retry with a fresh solver on failure. + +# Returns +- `result`: the MadNLP result from whichever attempt succeeded (or the retry + result if both failed). +- `retried::Bool`: `true` if the retry path was taken. + +# Examples +```julia +result, retried = _solve_with_retry!( + state, nlp; + warmstart=true, madnlp_kwargs=(print_level=0,), retry_on_failure=true, +) +retried && @warn "solve required retry" +``` +""" +function _solve_with_retry!(state::_SolverState, nlp; warmstart::Bool, madnlp_kwargs, retry_on_failure::Bool) + result = _solve!(state, nlp; warmstart = warmstart, madnlp_kwargs = madnlp_kwargs) # primary attempt + retried = false # track whether retry was needed + if retry_on_failure && (!solve_succeeded(result) || !isfinite(result.objective)) + retry_state = _make_solver(nlp, madnlp_kwargs) # fresh solver, clean factorization + result = _solve!(retry_state, nlp; warmstart = false, madnlp_kwargs = madnlp_kwargs) # cold-start retry + retried = true + if solve_succeeded(result) + state.last_good_x = copy(retry_state.solver.x.x) + state.last_good_y = copy(retry_state.solver.y) + state.last_good_zl_vals = copy(retry_state.solver.zl.values) + state.last_good_zu_vals = copy(retry_state.solver.zu.values) + end + end + return result, retried +end + # ── simulate_tsddr ──────────────────────────────────────────────────────────── """ simulate_tsddr(model, initial_state, det_equivalent, p_x0, p_target, p_uncertainty, - uncertainty_sampler; madnlp_kwargs, warmstart) - -> (objective, lambda) or nothing + uncertainty_sampler; + madnlp_kwargs, warmstart) + -> NamedTuple{(:objective, :lambda)} or nothing + +Perform a single forward pass of the TS-DDR pipeline without a gradient +update: roll out the policy to produce target states, solve the +deterministic-equivalent NLP, and extract the envelope-theorem multipliers. + +The forward pass computes + +```math +\\hat{x}_t = \\pi_\\theta(w_t, \\hat{x}_{t-1}), \\quad t = 1, \\ldots, T, +``` + +then solves ``\\min_z Q(z; \\hat{x}, w, x_0)`` and returns the objective value +and the multipliers ``\\lambda = \\nabla_{\\hat{x}} Q`` from the target +equality constraints. -Single forward pass without a gradient update. +# Arguments +- `model`: Flux policy network (LSTM or MLP). +- `initial_state::AbstractVector`: initial state vector ``x_0``. +- `det_equivalent`: ExaModels NLP with `.core`, `.model`, `.horizon`, + `.target_con_range`. +- `p_x0`: ExaModels parameter handle for the initial state. +- `p_target`: ExaModels parameter handle for policy targets. +- `p_uncertainty`: ExaModels parameter handle for per-stage uncertainty. +- `uncertainty_sampler`: `() -> w_flat` returning a flat vector of length + ``T \\times n_w``. + +# Keywords +- `madnlp_kwargs`: `NamedTuple` forwarded to MadNLP (default `NamedTuple()`). +- `warmstart::Bool`: warm-start MadNLP (default `true`). + +# Returns +- A `NamedTuple` with fields `objective::Float64` and `lambda::Vector{F}`, + or `nothing` if the solve failed or the objective is non-finite. + +# Examples +```julia +result = simulate_tsddr(model, x0, de, p_x0, p_target, p_unc, sampler) +if result !== nothing + @show result.objective +end +``` """ function simulate_tsddr( model, @@ -134,60 +552,131 @@ function simulate_tsddr( madnlp_kwargs = NamedTuple(), warmstart::Bool = true, ) - T = det_equivalent.horizon - F = eltype(initial_state) - nx = length(initial_state) - core = det_equivalent.core - nlp = det_equivalent.model + T = det_equivalent.horizon # number of planning stages + F = eltype(initial_state) # element type (Float32 or Float64) + nx = length(initial_state) # state dimension + core = det_equivalent.core # ExaModels core (for set_parameter!) + nlp = det_equivalent.model # ExaModels NLP model (for MadNLP) - state = _make_solver(nlp, madnlp_kwargs) + state = _make_solver(nlp, madnlp_kwargs) # fresh solver — no warm-start cache - w_flat = uncertainty_sampler() - nw = length(w_flat) ÷ T + # Sample one uncertainty scenario and move to the correct device. + w_flat = uncertainty_sampler() # flat vector of length T * nw_per_stage + nw = length(w_flat) ÷ T # uncertainty dimension per stage + w_dev = _adapt_array(F.(w_flat), initial_state) # move to GPU if initial_state is on GPU - Flux.reset!(model) - xhat_stages = Vector{Vector{F}}(undef, T) - prev = F.(initial_state) + # Roll out the policy to produce target states (outside AD tape). + Flux.reset!(model) # reset LSTM hidden state + xhat_stages = Vector{AbstractVector{F}}(undef, T) # allocate per-stage target storage + prev = initial_state # first policy input is x0 for t in 1:T - wt = F.(w_flat[(t-1)*nw+1 : t*nw]) - xhat_stages[t] = model(vcat(wt, prev)) - prev = xhat_stages[t] + wt = view(w_dev, (t-1)*nw+1 : t*nw) # slice uncertainty for stage t + xhat_stages[t] = model(vcat(wt, prev)) # policy: [w_t; x̂_{t-1}] → x̂_t + prev = xhat_stages[t] # feed x̂_t to next stage end - xhat_flat = vcat(xhat_stages...) + xhat_flat = vcat(xhat_stages...) # flatten targets to a single vector + # Set NLP parameters: initial state, uncertainty, and policy targets. ExaModels.set_parameter!(core, p_x0, initial_state) ExaModels.set_parameter!(core, p_uncertainty, w_flat) - ExaModels.set_parameter!(core, p_target, Float64.(xhat_flat)) + ExaModels.set_parameter!(core, p_target, Float64.(xhat_flat)) # NLP uses Float64 + prepare_solve!(det_equivalent, initial_state, w_flat, xhat_flat) + # Solve the deterministic equivalent (cold start for one-shot simulation). result = _solve!(state, nlp; warmstart = false, madnlp_kwargs = madnlp_kwargs) + # Reject failed or non-finite solves. solve_succeeded(result) || return nothing isfinite(result.objective) || return nothing - λ = result.multipliers[det_equivalent.target_con_range] - return (objective = result.objective, lambda = F.(Array(λ))) + # Extract envelope-theorem multipliers λ = ∇_{x̂} Q. + λ = target_multipliers(det_equivalent, result) + return (objective = result.objective, lambda = F.(λ)) # cast λ to match initial_state eltype end +""" + _rollout_xhat_flat(model, initial_state, w_flat, T::Int, F) -> AbstractVector{F} + +Roll out the policy network over `T` stages and return the concatenated +target trajectory as a single flat vector. + +This is the differentiable inner loop used inside `Zygote.gradient` blocks. +Each stage evaluates + +```math +\\hat{x}_t = \\pi_\\theta([w_t; \\hat{x}_{t-1}]), \\quad t = 1, \\ldots, T, +``` + +and the returned vector is ``[\\hat{x}_1; \\hat{x}_2; \\ldots; \\hat{x}_T]``. + +# Arguments +- `model`: Flux policy (LSTM or MLP). +- `initial_state`: state vector ``x_0`` fed to the first policy call. +- `w_flat`: flat uncertainty vector of length ``T \\times n_w``. +- `T::Int`: number of planning stages. +- `F`: element type (e.g., `Float32`). + +# Returns +- A flat vector of length ``T \\times n_x`` containing all stage targets. +""" function _rollout_xhat_flat(model, initial_state, w_flat, T::Int, F) - nw = length(w_flat) ÷ T - nx = length(initial_state) - Flux.reset!(model) - buf = Zygote.Buffer(zeros(F, nx * T)) - prev = F.(initial_state) + nw = length(w_flat) ÷ T # uncertainty dimension per stage + Flux.reset!(model) # reset LSTM hidden state + prev = F.(initial_state) # cast initial state to element type F + # This runs INSIDE the actor's Zygote closure (critic term), so the flat + # trajectory cannot be built with raw setindex! (Zygote mutation error) nor + # with a shape-growing vcat loop variable (pullback accum mismatch: + # accum(1375, 11) at the first phase-2 gradient). Zygote.Buffer is the + # sanctioned mutation-safe accumulator for exactly this unroll pattern. + nx = length(prev) # state dimension + buf = Zygote.Buffer(prev, nx * T) # AD-safe writable buffer for t in 1:T - wt = F.(w_flat[(t-1)*nw+1 : t*nw]) - xt = model(vcat(wt, prev)) - for i in 1:nx - buf[(t-1)*nx + i] = xt[i] - end - prev = xt + wt = view(w_flat, (t-1)*nw+1 : t*nw) # slice uncertainty for stage t + xt = model(vcat(wt, prev)) # policy forward pass + buf[(t-1)*nx+1 : t*nx] = xt # write stage into buffer + prev = xt # feed target to next stage end - return copy(buf) + return copy(buf) # differentiable flat trajectory end -_has_critic(::NoCriticControlVariate) = false -_has_critic(::AbstractCriticControlVariate) = true +""" + _has_critic(control_variate::AbstractCriticControlVariate) -> Bool + +Return `true` if the control variate wraps an actual critic network, `false` +for the no-op [`NoCriticControlVariate`](@ref). + +# Arguments +- `control_variate`: an [`AbstractCriticControlVariate`](@ref) instance. + +# Returns +- `false` for `NoCriticControlVariate` (recovers the original dual-only update). +- `true` for any concrete critic (e.g., [`ScalarCriticControlVariate`](@ref)). +""" +_has_critic(::NoCriticControlVariate) = false # no-op sentinel → no critic +_has_critic(::AbstractCriticControlVariate) = true # any concrete critic → active + +""" + _validate_critic_training_args(; kwargs...) -> Bool + +Validate critic/control-variate keyword arguments passed to [`train_tsddr`](@ref). + +# Keywords +- `actor_gradient_mode`: must be `:control_variate` or `:surrogate`. +- `critic_cv_weight`: nonnegative control-variate weight. +- `dual_actor_weight`: nonnegative dual-gradient actor weight. +- `critic_actor_weight`: nonnegative critic actor weight. +- `critic_updates_per_batch`: nonnegative number of critic updates. +- `critic_buffer_size`: nonnegative replay-buffer capacity. +- `critic_rollout_samples_per_batch`: nonnegative integer or `nothing`. +- `num_cheap_critic_samples_per_batch`: nonnegative number of extra policy + rollouts. +# Returns +- `true` when all arguments are valid. + +# Throws +- `ErrorException` if any argument is outside its admissible set. +""" function _validate_critic_training_args(; actor_gradient_mode, critic_cv_weight, @@ -214,6 +703,25 @@ function _validate_critic_training_args(; return true end +""" + _resolve_critic_training_target(target, has_critic::Bool) + +Resolve a user-facing critic target configuration to a concrete +`AbstractCriticTrainingTarget`. + +# Arguments +- `target`: critic target object or symbolic alias. +- `has_critic::Bool`: whether critic training is active. + +# Returns +- `DeterministicEquivalentCriticTarget()` when no critic is active or the user + selected deterministic-equivalent critic targets. +- `target` unchanged when it is already an `AbstractCriticTrainingTarget`. + +# Throws +- `ErrorException` when rollout critic training is requested without a + concrete [`RolloutCriticTarget`](@ref) configuration. +""" function _resolve_critic_training_target(target, has_critic::Bool) has_critic || return DeterministicEquivalentCriticTarget() target isa AbstractCriticTrainingTarget && return target @@ -226,6 +734,30 @@ function _resolve_critic_training_target(target, has_critic::Bool) end end +""" + _critic_sample_from_rollout(model, initial_state, target, w_flat, lambda, F, solver_state) + +Build one [`CriticSample`](@ref) by rerunning a solved scenario through +stage-wise rollout. + +# Arguments +- `model`: Flux policy being trained. +- `initial_state`: initial state vector. +- `target::RolloutCriticTarget`: rollout critic-target configuration. +- `w_flat`: uncertainty trajectory from a deterministic-equivalent sample. +- `lambda`: target multipliers from the deterministic-equivalent solve. +- `F`: element type used for critic sample arrays. +- `solver_state`: optional reusable rollout solver state. + +# Returns +- `CriticSample` when rollout succeeds. +- `nothing` when rollout fails. + +# Notes +The rollout objective supplies the critic value target. The deterministic +equivalent multipliers are sliced to the rollout target length and used as the +critic gradient target. +""" function _critic_sample_from_rollout( model, initial_state, @@ -237,11 +769,15 @@ function _critic_sample_from_rollout( ) # Keep both rollout objective variants available; target.objective_value # below selects which one is used as the critic value target. + rollout_len = target.horizon * target.n_uncertainty + length(w_flat) >= rollout_len || + error("rollout critic uncertainty has length $(length(w_flat)); expected at least $rollout_len") + w_rollout = view(w_flat, 1:rollout_len) result = rollout_tsddr( model, initial_state, target.stage_problem, - w_flat; + w_rollout; horizon = target.horizon, n_uncertainty = target.n_uncertainty, set_stage_parameters! = target.set_stage_parameters!, @@ -252,15 +788,42 @@ function _critic_sample_from_rollout( policy_state = target.policy_state, solver_state = solver_state, reuse_solver = target.reuse_solver, + state_bounds = target.state_bounds, + project_state = target.project_state, + retry_on_failure = target.retry_on_failure, ) result === nothing && return nothing objective = target.objective_value === :objective ? result.objective : result.objective_no_target_penalty xhat_flat = F.(vcat(result.target_trajectory...)) - return CriticSample(F.(initial_state), F.(w_flat), xhat_flat, objective, F.(lambda)) + λ_rollout = view(lambda, 1:length(xhat_flat)) + return CriticSample(F.(initial_state), F.(w_rollout), xhat_flat, objective, F.(λ_rollout)) end +""" + _rollout_critic_samples(model, initial_state, target, de_samples, F, max_samples, solver_state) + -> Vector{CriticSample} + +Convert deterministic-equivalent training samples into rollout critic samples. + +# Arguments +- `model`: Flux policy being trained. +- `initial_state`: initial state vector. +- `target::RolloutCriticTarget`: rollout critic-target configuration. +- `de_samples`: deterministic-equivalent samples containing uncertainty and + target multipliers. +- `F`: element type used for critic sample arrays. +- `max_samples`: maximum number of solved scenarios to rerun, or `nothing`. +- `solver_state`: optional reusable rollout solver state. + +# Returns +- `Vector{CriticSample}` containing only successful rollout conversions. + +# Notes +When `max_samples` is smaller than `length(de_samples)`, samples are selected +without replacement using `randperm`. +""" function _rollout_critic_samples( model, initial_state, @@ -298,51 +861,69 @@ end train_tsddr(model, initial_state, det_equivalent, p_x0, p_target, p_uncertainty, uncertainty_sampler; - num_batches, num_train_per_batch, optimizer, - adjust_hyperparameters, record_loss, - madnlp_kwargs, warmstart, - problem_pool) -> model - -TS-DDR policy gradient training. Mirrors `train_multistage` from DecisionRules.jl. - -Arguments: -- `model` : Flux policy (LSTM or MLP) -- `initial_state` : initial state vector -- `det_equivalent` : any ExaModels NLP with fields `.core`, `.model`, - `.horizon`, `.target_con_range` -- `p_x0` : ExaModels parameter for the initial state -- `p_target` : ExaModels parameter for policy targets -- `p_uncertainty` : ExaModels parameter for per-stage uncertainty -- `uncertainty_sampler`: `() -> w_flat` — flat vector of length `T * nw_per_stage`. - For multi-unit problems (e.g., hydro reservoirs) the sampler - should draw one joint scenario index per stage to preserve - spatial correlation; see `sample_scenario` in examples. - -Keyword arguments (mirror `train_multistage`): -- `num_batches` : total gradient steps (default 100) -- `num_train_per_batch` : scenarios averaged per step (default 1) -- `optimizer` : Flux.Optimisers optimizer -- `adjust_hyperparameters` : `(iter, opt_state, n) -> n` -- `record_loss` : `(iter, model, loss, tag) -> Bool`; return `true` to stop -- `madnlp_kwargs` : NamedTuple forwarded to MadNLP -- `warmstart` : warm-start MadNLP between solves (default `true`) -- `problem_pool` : vector of `(de, p_x0, p_target, p_uncertainty)` tuples - for parallel GPU solves; each entry gets its own MadNLP solver - and samples are distributed round-robin across the pool -- `control_variate` : optional `ScalarCriticControlVariate`; default - `NoCriticControlVariate()` recovers the original update -- `critic_training_target` : `RolloutCriticTarget(...)` for rollout-objective - critic fitting, or `DeterministicEquivalentCriticTarget()` - / `:deterministic_equivalent` for DE ablations -- `critic_rollout_samples_per_batch`: number of solved batch scenarios to rerun - through stage-wise rollout for critic targets; - `nothing` uses all successful solved scenarios -- `actor_gradient_mode` : `:control_variate` or `:surrogate` -- `num_cheap_critic_samples_per_batch`: extra policy rollouts used only for - critic actor terms; these do not trigger NLP solves -- `external_critic_samples` : mutable vector; `record_loss` can push - `CriticSample`s (e.g. from `critic_samples_from_evaluation`) - to feed the critic replay buffer without extra solves + num_batches=100, + num_train_per_batch=1, + optimizer, + adjust_hyperparameters, + record_loss, + madnlp_kwargs=NamedTuple(), + warmstart=true, + problem_pool=nothing, + kwargs...) -> model + +Train a TS-DDR policy with open-loop deterministic-equivalent solves. + +The policy rolls out a target trajectory, the ExaModels problem projects that +trajectory onto the feasible set, and target multipliers provide the actor +gradient by the envelope theorem. + +# Arguments +- `model`: Flux policy. +- `initial_state::AbstractVector`: initial state vector. +- `det_equivalent`: ExaModels deterministic-equivalent problem. +- `p_x0`: ExaModels parameter for the initial state. +- `p_target`: ExaModels parameter for policy targets. +- `p_uncertainty`: ExaModels parameter for uncertainty. +- `uncertainty_sampler`: callable returning a flat uncertainty trajectory. + +# Keywords +- `num_batches::Int`: number of gradient steps. +- `num_train_per_batch::Int`: number of scenarios averaged per step. +- `optimizer`: Flux optimizer or optimizer chain. +- `adjust_hyperparameters`: callback `(iter, opt_state, n) -> n`. +- `record_loss`: callback `(iter, model, loss, tag) -> Bool`; return `true` + to stop training. +- `madnlp_kwargs`: keyword arguments forwarded to MadNLP. +- `warmstart::Bool`: warm-start MadNLP between solves. +- `retry_on_failure::Bool`: retry failed solves with a fresh solver state. +- `problem_pool`: optional vector of `(de, p_x0, p_target, p_uncertainty)` + tuples for independent solves. +- `control_variate`: optional critic control variate. +- `actor_gradient_mode::Symbol`: `:control_variate` or `:surrogate`. +- `critic_cv_weight`, `dual_actor_weight`, `critic_actor_weight`: actor loss + weights. +- `critic_updates_per_batch::Int`: critic optimizer steps per batch. +- `critic_buffer_size::Int`: replay-buffer capacity. +- `critic_batch_size`: critic minibatch size, or `nothing` for all samples. +- `critic_training_target`: rollout or deterministic-equivalent critic target. +- `critic_rollout_samples_per_batch`: number of solved samples rerun through + rollout for critic targets; `nothing` uses all successful solved samples. +- `num_cheap_critic_samples_per_batch::Int`: extra policy rollouts used only + for critic actor terms. +- `critic_optimizer`: Flux optimizer for the critic. +- `external_critic_samples`: optional mutable vector of externally produced + `CriticSample`s. +- `batch_diagnostics`: callback `(iter, stats) -> nothing`. +- `reuse_solver::Bool`: force solver reuse when fixed variables have constant + bounds across solves. + +# Returns +- `model`, updated in place. + +# Notes +For multi-unit stochastic processes, `uncertainty_sampler` should preserve +within-stage spatial correlation, for example by drawing one joint scenario +index per stage. """ function train_tsddr( model, @@ -365,6 +946,7 @@ function train_tsddr( end, madnlp_kwargs = NamedTuple(), warmstart::Bool = true, + retry_on_failure::Bool = true, problem_pool = nothing, control_variate::AbstractCriticControlVariate = NoCriticControlVariate(), actor_gradient_mode::Symbol = :control_variate, @@ -379,6 +961,10 @@ function train_tsddr( num_cheap_critic_samples_per_batch::Int = 0, critic_optimizer = Flux.Adam(1f-3), external_critic_samples = nothing, + batch_diagnostics = (iter, stats) -> nothing, + reuse_solver::Bool = false, + worker_devices = nothing, + worker_problem_builder = nothing, ) T = det_equivalent.horizon F = eltype(initial_state) @@ -401,15 +987,38 @@ function train_tsddr( ) # ── Build worker pool ──────────────────────────────────────────────────── - if problem_pool === nothing - _pool = [(det_equivalent, p_x0, p_target, p_uncertainty)] + # Two modes: + # - `worker_problem_builder === nothing` (default): the caller supplies a + # `problem_pool` of already-built DEs (single-device / CPU). + # - `worker_problem_builder = (wi) -> (de, p_x0, p_target, p_uncertainty)`: + # each worker builds its OWN DE INSIDE its task, after binding its GPU. + # Required for multi-GPU: a DE built in the main task and solved from a + # worker task deadlocks on the first cross-task solve (CUDSS/stream/event + # ownership); building in-task keeps DE, solver, stream and events in one + # task/device context. nworkers then comes from `worker_devices`. + _build_in_worker = worker_problem_builder !== nothing + if _build_in_worker + worker_devices !== nothing || + throw(ArgumentError("worker_problem_builder requires worker_devices")) + _pool = nothing + nworkers = length(worker_devices) else - _pool = problem_pool + _pool = problem_pool === nothing ? + [(det_equivalent, p_x0, p_target, p_uncertainty)] : problem_pool + nworkers = length(_pool) + if worker_devices !== nothing + length(worker_devices) == nworkers || + throw(ArgumentError("worker_devices has length $(length(worker_devices)) but there are $nworkers workers")) + end end - nworkers = length(_pool) + _worker_device(wi) = worker_devices === nothing ? nothing : worker_devices[wi] # Single-worker: create solver on main task (no threading needed) - single_state = nworkers == 1 ? _make_solver(_pool[1][1].model, madnlp_kwargs) : nothing + single_state = (nworkers == 1 && !_build_in_worker) ? + _make_solver(_pool[1][1].model, madnlp_kwargs) : nothing + if reuse_solver && single_state !== nothing + single_state.has_fixed_vars = false + end # Multi-worker: persistent worker threads via channels. # Each worker creates its own MadNLP solver on its own thread so that @@ -419,28 +1028,91 @@ function train_tsddr( worker_tasks = Task[] if nworkers > 1 for wi in 1:nworkers - (de, px, pt, pu) = _pool[wi] + _pooled = _build_in_worker ? nothing : _pool[wi] + _builder = worker_problem_builder in_ch = in_channels[wi] out_ch = out_channels[wi] - t = Threads.@spawn begin + _reuse_solver = reuse_solver + _dev = _worker_device(wi) + _wi = wi + t = Threads.@spawn try + # Bind this worker to its GPU (multi-GPU). Must precede DE/solver + # creation so CUDA handles + all CuArray ops on this task target + # `_dev`. Diagnostics go to stderr (flushed) so a hang is + # localizable in the SLURM log even under the WandbLogger. + # The whole body runs under try/catch: a Threads.@spawn task that + # throws dies SILENTLY (exceptions surface only on wait/fetch), so + # without this the main task blocks forever on take!(out_ch) — the + # exact multi-GPU "hang" signature. On error we report loudly and + # close out_ch so the main loop fails fast instead of deadlocking. + println(stderr, "[worker $_wi] task started (dev=$_dev, thread=$(Threads.threadid()))"); flush(stderr) + if _dev !== nothing + CUDA.device!(_dev) + println(stderr, "[worker $_wi] bound to CUDA device $_dev (current=$(CUDA.device()))"); flush(stderr) + end + # Build the DE IN-TASK for multi-GPU (see _build_in_worker note): + # a main-task-built DE deadlocks on the first cross-task solve. + (de, px, pt, pu) = _builder === nothing ? _pooled : _builder(_wi) + _builder === nothing || + (println(stderr, "[worker $_wi] DE built in-task on device $_dev"); flush(stderr)) st = _make_solver(de.model, madnlp_kwargs) + _dev === nothing || (println(stderr, "[worker $_wi] solver ready on device $_dev"); flush(stderr)) + if _reuse_solver + st.has_fixed_vars = false + end while true msg = take!(in_ch) msg === nothing && break (s_idx, init_state, w_flat, xhat_flat) = msg + # Multi-GPU: the main task marshals msg arrays through CPU; + # upload them to THIS worker's device for the solve, but keep + # the CPU w for the reply below — anything sent back must be + # device-neutral (CPU), because the main task consumes it in + # the device-0 gradient (a device-N CuArray there is the + # CUDA-700 illegal access localized by job 10910905). + w_reply = w_flat + if _dev !== nothing + init_state = CUDA.cu(init_state) + w_flat = CUDA.cu(w_flat) + xhat_flat = CUDA.cu(xhat_flat) + end ExaModels.set_parameter!(de.core, px, init_state) ExaModels.set_parameter!(de.core, pu, w_flat) ExaModels.set_parameter!(de.core, pt, Float64.(xhat_flat)) - result = _solve!(st, de.model; warmstart=warmstart, madnlp_kwargs=madnlp_kwargs) - if solve_succeeded(result) && isfinite(result.objective) - λ = result.multipliers[de.target_con_range] + prepare_solve!(de, init_state, w_flat, xhat_flat) + result, retried = _solve_with_retry!( + st, + de.model; + warmstart = warmstart, + madnlp_kwargs = madnlp_kwargs, + retry_on_failure = retry_on_failure, + ) + failure = nothing + if !solve_succeeded(result) + failure = "status_" * _status_key(result.status) + elseif !isfinite(result.objective) + failure = "nonfinite_objective" + else + # Solve succeeded with a finite objective (both negations + # were tested above); only the multipliers remain to check. + λ = target_multipliers(de, result) if all(isfinite, λ) - put!(out_ch, (s_idx, F.(w_flat), F.(Array(λ)), result.objective)) + # Reply with the CPU w (same types the single-GPU + # worker path produces); λ adapts to it → CPU too. + put!(out_ch, (s_idx, F.(w_reply), _adapt_array(F.(λ), w_reply), + result.objective, result.status, nothing, retried)) continue end + failure = "nonfinite_lambda" end - put!(out_ch, (s_idx, nothing, nothing, NaN)) + put!(out_ch, (s_idx, nothing, nothing, NaN, result.status, failure, retried)) end + catch err + # Loud failure + closed channel: the main task's take!(out_ch) + # throws immediately instead of blocking forever on a dead worker. + println(stderr, "[worker $_wi] FATAL: "); showerror(stderr, err, catch_backtrace()); println(stderr); flush(stderr) + close(out_ch) + rethrow() end push!(worker_tasks, t) end @@ -461,24 +1133,36 @@ function train_tsddr( # ── Forward pass: rollout + solve (outside AD tape) ─────────────────── - # Step 1: Roll out policy for all samples (CPU, sequential) - sample_data = Vector{Tuple{Vector{F}, Vector{F}}}(undef, num_train_per_batch) + # Step 1: Roll out policy for all samples + # Precision note: uncertainties are cast to F (typically Float32, the + # policy precision) here, and this F-cast array is later written into + # the Float64 NLP via set_parameter!. The round-trip through Float32 is + # intentional: the NLP must be solved for exactly the targets the policy + # produced from these Float32 inputs, so the resulting λ multipliers + # pair with the same Float32 rollout in the gradient step below + # (train/gradient consistency). Do not "fix" this by keeping w in + # Float64 for the NLP only. + sample_data = Vector{Tuple{AbstractVector{F}, AbstractVector{F}}}(undef, num_train_per_batch) for s in 1:num_train_per_batch w_flat = uncertainty_sampler() nw = length(w_flat) ÷ T + w_dev = _adapt_array(F.(w_flat), initial_state) Flux.reset!(model) - xhat_stages = Vector{Vector{F}}(undef, T) - prev = F.(initial_state) + xhat_stages = Vector{AbstractVector{F}}(undef, T) + prev = initial_state for t in 1:T - wt = F.(w_flat[(t-1)*nw+1 : t*nw]) + wt = view(w_dev, (t-1)*nw+1 : t*nw) xhat_stages[t] = model(vcat(wt, prev)) - prev = xhat_stages[t] + prev = xhat_stages[t] end - sample_data[s] = (F.(w_flat), vcat(xhat_stages...)) + sample_data[s] = (w_dev, vcat(xhat_stages...)) end # Step 2: Solve — parallel across workers if pool provided - solve_ok = Vector{Union{Nothing, Tuple{Vector{F}, Vector{F}, Float64}}}(nothing, num_train_per_batch) + solve_ok = Vector{Union{Nothing, Tuple{AbstractVector{F}, AbstractVector{F}, Float64}}}(nothing, num_train_per_batch) + status_counts = Dict{String, Int}() + failure_counts = Dict{String, Int}() + retry_counts = Dict{String, Int}() if nworkers == 1 (de, px, pt, pu) = _pool[1] @@ -488,12 +1172,30 @@ function train_tsddr( ExaModels.set_parameter!(de.core, px, initial_state) ExaModels.set_parameter!(de.core, pu, w_flat) ExaModels.set_parameter!(de.core, pt, Float64.(xhat_flat)) - result = _solve!(st, de.model; warmstart=warmstart, madnlp_kwargs=madnlp_kwargs) - solve_succeeded(result) || continue - isfinite(result.objective) || continue - λ = result.multipliers[de.target_con_range] - all(isfinite, λ) || continue - solve_ok[s] = (F.(w_flat), F.(Array(λ)), result.objective) + prepare_solve!(de, initial_state, w_flat, xhat_flat) + result, retried = _solve_with_retry!( + st, + de.model; + warmstart = warmstart, + madnlp_kwargs = madnlp_kwargs, + retry_on_failure = retry_on_failure, + ) + retried && _inc_count!(retry_counts, solve_succeeded(result) && isfinite(result.objective) ? "retry_success" : "retry_failure") + _inc_status!(status_counts, result.status) + if !solve_succeeded(result) + _inc_count!(failure_counts, "status_" * _status_key(result.status)) + continue + end + if !isfinite(result.objective) + _inc_count!(failure_counts, "nonfinite_objective") + continue + end + λ = target_multipliers(de, result) + if !all(isfinite, λ) + _inc_count!(failure_counts, "nonfinite_lambda") + continue + end + solve_ok[s] = (F.(w_flat), _adapt_array(F.(λ), initial_state), result.objective) end else for round_start in 1:nworkers:num_train_per_batch @@ -502,19 +1204,47 @@ function train_tsddr( for s in round_start:round_end wi = s - round_start + 1 w_flat, xhat_flat = sample_data[s] - put!(in_channels[wi], (s, initial_state, w_flat, xhat_flat)) + if worker_devices === nothing + put!(in_channels[wi], (s, initial_state, w_flat, xhat_flat)) + else + # Multi-GPU workers run with different current devices. + # Never send a CuArray allocated on device 0 to a worker + # bound to device 1/2; materialize through CPU and let + # the worker copy onto its own device. + put!(in_channels[wi], (s, Array(initial_state), Array(w_flat), Array(xhat_flat))) + end end for wi in 1:round_size - (s_idx, w_out, λ_out, obj_out) = take!(out_channels[wi]) + msg = take!(out_channels[wi]) + if length(msg) == 7 + (s_idx, w_out, λ_out, obj_out, status_out, failure_out, retried_out) = msg + _inc_status!(status_counts, status_out) + failure_out !== nothing && _inc_count!(failure_counts, failure_out) + retried_out && _inc_count!(retry_counts, w_out === nothing ? "retry_failure" : "retry_success") + elseif length(msg) == 6 + (s_idx, w_out, λ_out, obj_out, status_out, failure_out) = msg + _inc_status!(status_counts, status_out) + failure_out !== nothing && _inc_count!(failure_counts, failure_out) + else + (s_idx, w_out, λ_out, obj_out) = msg + end if w_out !== nothing - solve_ok[s_idx] = (w_out, λ_out, obj_out) + # Workers reply device-neutral (CPU) arrays; the actor + # gradient below mixes them with `initial_state`-device + # arrays (vcat/broadcast), so land them on that device + # here — the exact analogue of the single-worker path's + # `_adapt_array(F.(λ), initial_state)`. No-op when the + # types already match (single-GPU pool replies GPU w). + solve_ok[s_idx] = (_adapt_array(w_out, initial_state), + _adapt_array(λ_out, initial_state), + obj_out) end end end end # Step 3: Collect valid results - valid = Vector{Tuple{Vector{F}, Vector{F}}}() + valid = Tuple{AbstractVector{F}, AbstractVector{F}}[] de_samples = CriticSample[] obj_sum = 0.0 for (s, r) in enumerate(solve_ok) @@ -528,6 +1258,13 @@ function train_tsddr( end n_ok = length(valid) mean_obj = n_ok > 0 ? obj_sum / n_ok : NaN + batch_diagnostics(iter, Dict{String, Any}( + "n_ok" => n_ok, + "n_total" => num_train_per_batch, + "status_counts" => copy(status_counts), + "failure_counts" => copy(failure_counts), + "retry_counts" => copy(retry_counts), + )) if has_critic && n_ok > 0 && critic_updates_per_batch > 0 valid_samples = if resolved_critic_training_target isa RolloutCriticTarget @@ -574,18 +1311,18 @@ function train_tsddr( for (w_flat_s, λf) in valid nw = length(w_flat_s) ÷ T Flux.reset!(m) - prev_ad = F.(initial_state) + prev_ad = initial_state for t in 1:T - wt = F.(w_flat_s[(t-1)*nw+1 : t*nw]) + wt = view(w_flat_s, (t-1)*nw+1 : t*nw) xt = m(vcat(wt, prev_ad)) - total = total + sum(λf[(t-1)*nx+1 : t*nx] .* xt) + total = total + sum(view(λf, (t-1)*nx+1 : t*nx) .* xt) prev_ad = xt end end total / F(n_ok) end else - solved_weights = Vector{Tuple{Vector{F}, Vector{F}}}() + solved_weights = Tuple{AbstractVector{F}, AbstractVector{F}}[] for sample in de_samples λf = F.(sample.target_multipliers) if actor_gradient_mode === :control_variate @@ -618,12 +1355,12 @@ function train_tsddr( for (w_flat_s, actor_weight) in solved_weights nw = length(w_flat_s) ÷ T Flux.reset!(m) - prev_ad = F.(initial_state) + prev_ad = initial_state for t in 1:T - wt = F.(w_flat_s[(t-1)*nw+1 : t*nw]) + wt = view(w_flat_s, (t-1)*nw+1 : t*nw) xt = m(vcat(wt, prev_ad)) residual_total = - residual_total + sum(actor_weight[(t-1)*nx+1 : t*nx] .* xt) + residual_total + sum(view(actor_weight, (t-1)*nx+1 : t*nx) .* xt) prev_ad = xt end end @@ -637,7 +1374,7 @@ function train_tsddr( xhat_ad = _rollout_xhat_flat(m, initial_state, w_flat_s, T, F) critic_total = critic_total + critic_value( control_variate, - F.(initial_state), + initial_state, w_flat_s, xhat_ad, ) @@ -659,14 +1396,193 @@ function train_tsddr( end finally - # Shut down worker threads - for ch in in_channels - put!(ch, nothing) + # Shut down worker threads. A worker that already died never drains its + # input channel, so an unconditional put! on a full Channel{Any}(1) + # would block forever; only signal workers that are still running, and + # guard the put! itself against the check-then-put race. + for (i, ch) in enumerate(in_channels) + if !istaskdone(worker_tasks[i]) # skip dead workers (nobody consumes) + try + put!(ch, nothing) # normal shutdown sentinel + catch err + # Worker died between the istaskdone check and the put! + # (or the channel was closed) — nothing left to signal. + @warn "train_tsddr worker $i shutdown signal failed" exception=(err, catch_backtrace()) + end + end end - for t in worker_tasks - wait(t) + for (i, t) in enumerate(worker_tasks) + try + wait(t) # join worker task + catch err + # A failed worker rethrows on wait; log instead of masking the + # original in-flight exception during cleanup. + @warn "train_tsddr worker $i failed" exception=(err, catch_backtrace()) + end end end return model end + +# ── train_tsddr_embedded ───────────────────────────────────────────────────── + +""" + train_tsddr_embedded(model, initial_state, embedded_de, + uncertainty_sampler; kwargs...) -> model + +Train a TS-DDR policy embedded directly in the NLP. + +Unlike [`train_tsddr`](@ref), this function does not roll out targets +externally. The NLP oracle evaluates `model` inline, the solve returns +closed-loop multipliers and realized states, and the actor gradient is computed +from those realized states. + +# Arguments +- `model`: Flux policy captured by the embedded oracle closures. +- `initial_state::AbstractVector`: initial state vector. +- `embedded_de`: embedded deterministic-equivalent problem. +- `uncertainty_sampler`: callable returning a flat uncertainty trajectory. + +# Keywords +- `num_batches::Int`: number of gradient steps. +- `num_train_per_batch::Int`: number of scenarios averaged per step. +- `optimizer`: Flux optimizer or optimizer chain. +- `adjust_hyperparameters`: callback `(iter, opt_state, n) -> n`. +- `record_loss`: callback `(iter, model, loss, tag) -> Bool`; return `true` + to stop training. +- `madnlp_kwargs`: keyword arguments forwarded to MadNLP. +- `warmstart::Bool`: warm-start MadNLP between solves. +- `retry_on_failure::Bool`: retry failed solves with a fresh solver state. +- `get_realized_states`: optional callback `(prob, result) -> x_flat`. +- `batch_diagnostics`: callback `(iter, stats) -> nothing`. + +# Returns +- `model`, updated in place. + +# Notes +The gradient uses +`sum_t dot(lambda_t, pi_theta(w_t, x^*_{t-1}))`, where `x^*` is the realized +state trajectory from the coupled NLP solution. +""" +function train_tsddr_embedded( + model, + initial_state::AbstractVector, + embedded_de, + uncertainty_sampler; + num_batches::Int = 100, + num_train_per_batch::Int = 1, + optimizer = Flux.Optimisers.OptimiserChain( + Flux.Optimisers.ClipGrad(1.0f0), + Flux.Adam(1f-3), + ), + adjust_hyperparameters = (iter, opt_state, n) -> n, + record_loss = (iter, model, loss, tag) -> begin + println("$tag iter=$iter loss=$(round(loss; digits=4))") + return false + end, + madnlp_kwargs = NamedTuple(), + warmstart::Bool = true, + retry_on_failure::Bool = true, + get_realized_states = nothing, + batch_diagnostics = (iter, stats) -> nothing, +) + T = embedded_de.horizon + F = eltype(initial_state) + nx = embedded_de.nx + + _get_states = get_realized_states === nothing ? + (prob, res) -> res.solution[1 : prob.horizon * prob.nx] : + get_realized_states + + state = _make_solver(embedded_de.model, madnlp_kwargs) + opt_state = Flux.setup(optimizer, model) + + for iter in 1:num_batches + num_train_per_batch = adjust_hyperparameters(iter, opt_state, num_train_per_batch) + + valid = Tuple{AbstractVector{F}, AbstractVector{F}, AbstractVector{F}}[] + obj_sum = 0.0 + status_counts = Dict{String, Int}() + failure_counts = Dict{String, Int}() + retry_counts = Dict{String, Int}() + + for s in 1:num_train_per_batch + w_flat = uncertainty_sampler() + + set_x0!(embedded_de, initial_state) + set_uncertainty!(embedded_de, w_flat) + + result, retried = _solve_with_retry!( + state, + embedded_de.model; + warmstart = warmstart, + madnlp_kwargs = madnlp_kwargs, + retry_on_failure = retry_on_failure, + ) + retried && _inc_count!(retry_counts, solve_succeeded(result) && isfinite(result.objective) ? "retry_success" : "retry_failure") + + _inc_status!(status_counts, result.status) + if !solve_succeeded(result) + _inc_count!(failure_counts, "status_" * _status_key(result.status)) + continue + end + if !isfinite(result.objective) + _inc_count!(failure_counts, "nonfinite_objective") + continue + end + + λ = target_multipliers(embedded_de, result) + if !all(isfinite, λ) + _inc_count!(failure_counts, "nonfinite_lambda") + continue + end + + x_sol = _get_states(embedded_de, result) + + λf = _adapt_array(F.(λ), initial_state) + xf = _adapt_array(F.(x_sol), initial_state) + w_dev = _adapt_array(F.(w_flat), initial_state) + push!(valid, (w_dev, λf, xf)) + obj_sum += result.objective + end + + n_ok = length(valid) + mean_obj = n_ok > 0 ? obj_sum / n_ok : NaN + batch_diagnostics(iter, Dict{String, Any}( + "n_ok" => n_ok, + "n_total" => num_train_per_batch, + "status_counts" => copy(status_counts), + "failure_counts" => copy(failure_counts), + "retry_counts" => copy(retry_counts), + )) + + if n_ok > 0 + gs = Zygote.gradient(model) do m + total = zero(F) + for (w_flat_s, λf, x_realized) in valid + nw = length(w_flat_s) ÷ T + Flux.reset!(m) + for t in 1:T + wt = view(w_flat_s, (t-1)*nw+1 : t*nw) + x_prev = (t == 1) ? + initial_state : + view(x_realized, (t-2)*nx+1 : (t-1)*nx) + xt = m(vcat(wt, x_prev)) + total = total + sum(view(λf, (t-1)*nx+1 : t*nx) .* xt) + end + end + total / F(n_ok) + end + + grad = materialize_tangent(gs[1]) + if grad !== nothing && _all_finite_gradient(grad) + Flux.update!(opt_state, model, grad) + end + end + + record_loss(iter, model, mean_obj, "metrics/training_loss") && break + end + + return model +end diff --git a/src/utils.jl b/src/utils.jl index 4ce90ba..208492f 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -2,25 +2,47 @@ # Small helpers shared across the package. """ - x_index(nx, t, i) + x_index(nx, t, i) -> Int -Linear index for state component `i ∈ 1:nx` at stage `t ∈ 1:T` -when state trajectory is stored as a flat vector of length `T*nx`. +Return the flat-vector index for state component `i` at stage `t`. + +# Arguments +- `nx::Int`: number of state components per stage. +- `t`: one-based stage index. +- `i`: one-based state-component index. + +# Returns +- `Int`: index into a stage-major state trajectory of length `T * nx`. """ @inline x_index(nx::Int, t, i) = (t - 1) * nx + i """ - u_index(nu, t, i) + u_index(nu, t, i) -> Int + +Return the flat-vector index for control component `i` at stage `t`. -Linear index for control component `i ∈ 1:nu` at stage `t ∈ 1:(T-1)` -when controls are stored as a flat vector of length `(T-1)*nu`. +# Arguments +- `nu::Int`: number of control components per stage. +- `t`: one-based stage index. +- `i`: one-based control-component index. + +# Returns +- `Int`: index into a stage-major control trajectory of length `(T - 1) * nu`. """ @inline u_index(nu::Int, t, i) = (t - 1) * nu + i """ - w_index(nw, t, i) + w_index(nw, t, i) -> Int + +Return the flat-vector index for uncertainty component `i` at stage `t`. + +# Arguments +- `nw::Int`: number of uncertainty components per stage. +- `t`: one-based stage index. +- `i`: one-based uncertainty-component index. -Linear index for disturbance component `i ∈ 1:nw` at stage `t ∈ 1:(T-1)` -when disturbances are stored as a flat vector of length `(T-1)*nw`. +# Returns +- `Int`: index into a stage-major uncertainty trajectory of length + `(T - 1) * nw`. """ @inline w_index(nw::Int, t, i) = (t - 1) * nw + i diff --git a/test/runtests.jl b/test/runtests.jl index 857e673..407edc7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,6 +1,8 @@ using Test using DecisionRulesExa +using ExaModels using Flux +using MadNLP using Random using Zygote @@ -24,7 +26,7 @@ using Zygote set_uncertainty!(prob, w) set_targets!(prob, xhat) - res = solve!(prob; tol = 1e-6, max_iter = 200) + res = DecisionRulesExa.solve!(prob; tol = 1e-6, max_iter = 200) n_x = T * nx n_u = (T - 1) * nx @@ -110,6 +112,39 @@ end @test isfinite(critic_loss(hybrid, [sample])) end +@testset "Bounded state policy helper" begin + Random.seed!(21) + lower = Float32[0, 1, -2] + upper = Float32[10, 1, 2] + policy = bounded_state_policy(2, lower, upper, [4]; activation = sigmoid) + + y = policy(Float32[0.2, -0.1, 3, 1, -1]) + @test length(y) == 3 + @test lower[1] <= y[1] <= upper[1] + @test y[2] == lower[2] + @test lower[3] <= y[3] <= upper[3] + @test length(policy.policy.combiner.bias) == 2 + + deep_policy = bounded_state_policy( + 2, + lower, + upper, + [4]; + activation = sigmoid, + combiner_layers = [5, 4], + ) + @test deep_policy.policy.combiner isa Flux.Chain + y_deep = deep_policy(Float32[0.2, -0.1, 3, 1, -1]) + @test length(y_deep) == 3 + @test lower[1] <= y_deep[1] <= upper[1] + @test y_deep[2] == lower[2] + @test lower[3] <= y_deep[3] <= upper[3] + + constant_policy = bounded_state_policy(1, Float32[3, -4], Float32[3, -4], [4]) + @test constant_policy(Float32[0, 99, 100]) == Float32[3, -4] + @test isempty(Flux.trainables(constant_policy)) +end + @testset "Critic and actor update separation" begin Random.seed!(11) critic = Chain(Dense(2 => 1, bias = false)) @@ -183,3 +218,388 @@ end @test materialize_tangent(g_cv).layers[1].weight ≈ materialize_tangent(g_dual).layers[1].weight end + +@testset "EmbeddedDeterministicEquivalentProblem (CPU)" begin + T = 5 + nx = 1 + nw = 1 + + Random.seed!(42) + policy = StateConditionedPolicy(nw, nx, nx, [8]; activation = tanh) + + prob = build_embedded_deterministic_equivalent( + policy; + horizon = T, + nx = nx, + nu = nx, + nw = nw, + backend = nothing, + float_type = Float64, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + @test prob isa EmbeddedDeterministicEquivalentProblem + @test prob.horizon == T + @test prob.nx == nx + + x0 = [1.0] + w = randn(T * nw) + + set_x0!(prob, x0) + set_uncertainty!(prob, w) + + n_x = T * nx + n_u = (T - 1) * nx + n_var = n_x + n_u + n_x + n_con_initial = nx + n_con_dynamics = (T - 1) * nx + n_con_oracle = n_x + n_con = n_con_initial + n_con_dynamics + n_con_oracle + + res = MadNLP.madnlp(prob.model; tol = 1e-6, max_iter = 500, print_level = MadNLP.ERROR) + @test DecisionRulesExa.solve_succeeded(res) + @test length(res.solution) == n_var + @test length(res.multipliers) == n_con + + λ = target_multipliers(prob, res) + @test length(λ) == n_x + @test all(isfinite, λ) + + x_sol, u_sol, δ_sol = solution_components(prob, res) + @test length(x_sol) == n_x + @test length(u_sol) == n_u + @test length(δ_sol) == n_x + + # Verify oracle constraint satisfaction: x_t + δ_t ≈ π_θ(w_t, x_{t-1}) + Flux.reset!(policy) + for t in 1:T + x_prev = (t == 1) ? Float32.(x0) : Float32.([x_sol[(t-2)*nx+1:(t-1)*nx]...]) + w_t = Float32.([w[(t-1)*nw+1:t*nw]...]) + nn_out = policy(vcat(w_t, x_prev)) + for i in 1:nx + xi = x_sol[(t-1)*nx+i] + di = δ_sol[(t-1)*nx+i] + @test xi + di ≈ Float64(nn_out[i]) atol = 1e-4 + end + end +end + +@testset "Embedded NN gradient (envelope theorem)" begin + T = 4 + nx = 1 + nw = 1 + + Random.seed!(7) + policy = StateConditionedPolicy(nw, nx, nx, [8]; activation = tanh) + + prob = build_embedded_deterministic_equivalent( + policy; + horizon = T, + nx = nx, + nu = nx, + nw = nw, + backend = nothing, + float_type = Float64, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + x0 = [0.5] + w = randn(T * nw) + + set_x0!(prob, x0) + set_uncertainty!(prob, w) + + res = MadNLP.madnlp(prob.model; tol = 1e-6, max_iter = 500, print_level = MadNLP.ERROR) + @test DecisionRulesExa.solve_succeeded(res) + + λ = target_multipliers(prob, res) + x_sol = res.solution[1 : T * nx] + + # Zygote gradient: ∇_θ Σ_t ⟨λ_t, π_θ(w_t, x*_{t-1})⟩ + gs = Zygote.gradient(policy) do m + total = 0.0f0 + Flux.reset!(m) + for t in 1:T + wt = Float32.([w[(t-1)*nw+1:t*nw]...]) + x_prev = (t == 1) ? + Float32.(x0) : + Float32.([x_sol[(t-2)*nx+1:(t-1)*nx]...]) + xt = m(vcat(wt, x_prev)) + for i in 1:nx + total = total + Float32(λ[(t-1)*nx+i]) * xt[i] + end + end + total + end + + g = materialize_tangent(gs[1]) + @test g !== nothing + @test DecisionRulesExa._all_finite_gradient(g) +end + +@testset "train_tsddr_embedded smoke test" begin + T = 4 + nx = 1 + nw = 1 + + Random.seed!(99) + policy = StateConditionedPolicy(nw, nx, nx, [8]; activation = tanh) + + prob = build_embedded_deterministic_equivalent( + policy; + horizon = T, + nx = nx, + backend = nothing, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + x0 = Float32[1.0] + losses = Float64[] + + train_tsddr_embedded( + policy, x0, prob, + () -> randn(T * nw); + num_batches = 5, + num_train_per_batch = 2, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + warmstart = true, + record_loss = (iter, m, loss, tag) -> begin + push!(losses, loss) + return false + end, + ) + + @test length(losses) == 5 + @test all(isfinite, losses) +end + +@testset "rollout_tsddr (CPU)" begin + horizon = 3 + nx = 1 + + # Stage problem: a horizon-2 linear tracking deterministic equivalent used + # as a one-stage projection problem. The realized next state is x_2. + stage_problem = build_linear_tracking_problem( + horizon = 2, + nx = nx, + backend = nothing, + slack_penalty = 10.0, + u_bounds = (-2.0, 2.0), + ) + + # Callback: write (x_{t-1}, w_t, xhat_t) into the stage problem. The stage-1 + # target equals the incoming state (its constraint is slack-absorbed anyway); + # the stage-2 target is the policy target being projected. + set_stage_params! = (prob, state, w_t, target, stage) -> begin + set_x0!(prob, state) + set_uncertainty!(prob, w_t) + set_targets!(prob, vcat(state, target)) + return nothing + end + # Callback: read the realized next state x_2 from the stage solution. + realized = (prob, result) -> begin + x_sol, _, _ = solution_components(prob, result) + return x_sol[end - prob.nx + 1 : end] + end + + Random.seed!(31) + policy = StateConditionedPolicy(1, 1, 1, [4]; activation = tanh) + x0 = Float32[0.5] + w_flat = Float32.(0.1 .* randn(horizon * 1)) + + result = rollout_tsddr( + policy, x0, stage_problem, w_flat; + horizon = horizon, + n_uncertainty = 1, + set_stage_parameters! = set_stage_params!, + realized_state = realized, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + ) + @test result isa NamedTuple + @test isfinite(result.objective) + @test length(result.state_trajectory) == horizon + 1 + @test length(result.target_trajectory) == horizon + + # The :target feedback mode (policy sees its own previous target) must also run. + result_target = rollout_tsddr( + policy, x0, stage_problem, w_flat; + horizon = horizon, + n_uncertainty = 1, + set_stage_parameters! = set_stage_params!, + realized_state = realized, + policy_state = :target, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + ) + @test result_target isa NamedTuple + @test isfinite(result_target.objective) + + # A wrong-length uncertainty vector must be rejected before any solve. + @test_throws ArgumentError rollout_tsddr( + policy, x0, stage_problem, Float32[0.1]; + horizon = horizon, + n_uncertainty = 1, + set_stage_parameters! = set_stage_params!, + realized_state = realized, + ) +end + +@testset "train_tsddr open-loop smoke test" begin + T = 4 + nx = 1 + + # train_tsddr writes the full length-T*nw uncertainty sample into + # p_uncertainty via ExaModels.set_parameter!, which enforces an exact size + # match. build_linear_tracking_problem's p_w has length (T-1)*nw (dynamics + # stages only), so this test builds the same linear tracking NLP manually + # with a full-length uncertainty parameter, where the final-stage entry + # does not enter the dynamics. + core = ExaModels.ExaCore(Float64) + x = ExaModels.variable(core, T * nx) + u = ExaModels.variable(core, (T - 1) * nx; lvar = -2.0, uvar = 2.0) + δ = ExaModels.variable(core, T * nx) + p_x0 = ExaModels.parameter(core, zeros(nx)) + p_w = ExaModels.parameter(core, zeros(T * nx)) # full length T*nw + p_target = ExaModels.parameter(core, zeros(T * nx)) + # Stage cost (x_t^2 + u_t^2)/2 plus slack penalty (rho/2)*delta^2, rho = 10. + ExaModels.objective(core, + (x[x_index(nx, t, i)]^2 + u[u_index(nx, t, i)]^2) / 2 + for t in 1:(T - 1), i in 1:nx + ) + ExaModels.objective(core, + 5.0 * δ[x_index(nx, t, i)]^2 + for t in 1:T, i in 1:nx + ) + # Initial condition, dynamics x_{t+1} = x_t + u_t + w_t, then targets LAST. + ExaModels.constraint(core, + x[x_index(nx, 1, i)] - p_x0[i] + for i in 1:nx + ) + ExaModels.constraint(core, + x[x_index(nx, t + 1, i)] - x[x_index(nx, t, i)] - + u[u_index(nx, t, i)] - p_w[w_index(nx, t, i)] + for t in 1:(T - 1), i in 1:nx + ) + ExaModels.constraint(core, + p_target[x_index(nx, t, i)] - x[x_index(nx, t, i)] - δ[x_index(nx, t, i)] + for t in 1:T, i in 1:nx + ) + model = ExaModels.ExaModel(core) + target_start = nx + (T - 1) * nx + 1 + prob = DeterministicEquivalentProblem( + core, model, x, u, δ, + p_x0, p_w, p_target, + nx, nx, nx, T, + target_start:(target_start + T * nx - 1), + ) + + Random.seed!(123) + policy = StateConditionedPolicy(nx, nx, nx, [8]; activation = tanh) + x0 = Float32[1.0] + losses = Float64[] + params_before = _state_vector(policy) + + train_tsddr( + policy, x0, prob, + prob.p_x0, prob.p_target, prob.p_w, + () -> randn(T * nx); + num_batches = 2, + num_train_per_batch = 1, + madnlp_kwargs = (tol = 1e-6, max_iter = 300, print_level = MadNLP.ERROR), + warmstart = true, + record_loss = (iter, m, loss, tag) -> begin + push!(losses, loss) + return false + end, + ) + + @test length(losses) == 2 + @test all(isfinite, losses) + @test _state_vector(policy) != params_before +end + +@testset "StateConditionedPolicy recurrent-state threading" begin + Random.seed!(42) + policy = StateConditionedPolicy(2, 1, 1, [4, 3]; activation = tanh) + x = Float32[0.3, -0.2, 0.5] + + # Memory: the same input twice WITHOUT reset must give different outputs + # (the recurrent state advanced between calls). A memoryless (per-call + # restart) encoder would reproduce the first output exactly. + Flux.reset!(policy) + y1 = policy(x) + y2 = policy(x) + @test y1 != y2 + + # Reset restores the exact initial recurrent state: identical output. + Flux.reset!(policy) + @test policy(x) == y1 + + # Manual LSTMCell recursion with the same weights reproduces the policy's + # stage outputs EXACTLY over a 3-stage open-loop sequence. + ws = [Float32[0.3, -0.2], Float32[0.1, 0.4], Float32[-0.5, 0.2]] + Flux.reset!(policy) + prev = Float32[0.5] + outs = Vector{Vector{Float32}}() + for t in 1:3 + y = policy(vcat(ws[t], prev)) + push!(outs, Float32.(y)) + prev = Float32.(y) + end + + cells = [l.cell for l in policy.encoder.layers] + states = Any[Flux.initialstates(c) for c in cells] + prev = Float32[0.5] + for t in 1:3 + h = ws[t] + for (i, c) in enumerate(cells) + h, states[i] = c(h, states[i]) + end + y = policy.combiner(vcat(h, prev)) + @test Float32.(y) == outs[t] + prev = Float32.(y) + end +end + +@testset "Zygote gradient through threaded 3-stage rollout" begin + Random.seed!(43) + policy = StateConditionedPolicy(2, 1, 1, [4]; activation = tanh) + ws = [Float32[0.3, -0.2], Float32[0.1, 0.4], Float32[-0.5, 0.2]] + x0 = Float32[0.5] + + # Same pattern as the training loops: reset inside the gradient block, + # thread the recurrent state through all stages via the policy forward. + gs = Zygote.gradient(policy) do m + Flux.reset!(m) + total = 0.0f0 + prev = x0 + for t in 1:3 + y = m(vcat(ws[t], prev)) + total = total + sum(y) + prev = y + end + total + end + g = materialize_tangent(gs[1]) + @test g !== nothing + @test DecisionRulesExa._all_finite_gradient(g) + + # Encoder gradients must be finite AND nonzero (the LSTM parameters + # participate in every stage of the threaded rollout). + enc_leaves = Float64[] + function _collect_leaves(x) + if x isa AbstractArray && eltype(x) <: Number + append!(enc_leaves, vec(Float64.(x))) + elseif x isa NamedTuple + foreach(_collect_leaves, values(x)) + elseif x isa Tuple + foreach(_collect_leaves, x) + end + return nothing + end + _collect_leaves(g.encoder) + @test !isempty(enc_leaves) + @test any(!=(0.0), enc_leaves) +end