Example: Distribution-Grid Reinforcement
===========================================

This chapter builds a small electricity distribution network — one
substation feeding several load nodes over a directed graph — where a
fault at the substation cascades downstream along the lines, which lines
get reinforced is a *discrete* design decision, and substation siting is a
*continuous* one. Both live in the same differentiable graph and are
optimized together in one gradient pass. The full runnable source is
``examples/electric_grid/model.py``.

Step 1 — The network
------------------------

Nodes are a :class:`~prophys.structures.PointList` in a shared frame;
edges are a fixed ``(m, 2)`` connectivity array. The substation's position
is a ``Param`` — siting is a design variable — while the load nodes stay
fixed.

.. code-block:: python

   import jax.numpy as jnp
   import prophys as prp

   site = prp.Frame("site", units="km")

   substation = prp.Param("substation_pos", shape=(2,), init=jnp.array([0.0, 0.0]))
   loads = jnp.array([[4.0, 1.0], [7.0, 2.0], [7.0, -3.0], [10.0, 4.0], [10.0, -1.0]])
   nodes = prp.stack([substation, *[loads[i] for i in range(loads.shape[0])]], axis=0)
   edges = jnp.array([[0, 1], [1, 2], [1, 3], [2, 4], [3, 5]])
   grid = prp.Network(nodes, edges, site)

   grid.edge_lengths()     # (5,), differentiable in substation_pos
   grid.total_length()     # scalar, e.g. a cabling-cost proxy

Step 2 — Reinforcement as a differentiable decision
--------------------------------------------------------

Whether each line is reinforced is naturally binary, but it needs a
gradient to be optimized alongside the continuous siting variable. A
:class:`~prophys.symbolic.BinaryState` per edge does exactly that: it
evaluates to a hard 0/1 in ``eval`` mode and a smooth, temperature-relaxed
value in ``opt`` mode, following the model's ``compile(mode=, tau=)``
setting like every other primitive in the library.

.. code-block:: python

   reinforced = prp.stack(
       [prp.BinaryState(f"reinforce_{i}", init=0.3) for i in range(edges.shape[0])], axis=0
   )
   # A reinforced line contains 90% of an incoming fault; an unreinforced
   # line passes 80% of it downstream.
   transmission = 0.8 - 0.7 * reinforced

Step 3 — Cascading a fault along the network
--------------------------------------------------

:meth:`Network.propagate <prophys.structures.Network.propagate>` runs a
K-hop linear cascade of a per-node signal along the directed edges,
scaled by the per-edge transmission weight at every hop — a fault
injected at the substation decays as it crosses reinforced lines and
persists across unreinforced ones:

.. code-block:: python

   fault_severity = prp.RandomVariable("fault_severity", prp.Weibull(scale=1.0, concentration=1.5))

   from prophys.symbolic.expr import op

   def cascade_load_loss(severity):
       n_nodes = grid.n_nodes

       def _source(sev):
           pad = jnp.zeros(jnp.shape(sev) + (n_nodes - 1,))
           return jnp.concatenate([jnp.reshape(sev, jnp.shape(sev) + (1,)), pad], axis=-1)

       source = op("fault_source", _source, severity)
       return grid.propagate(source, steps=3, weights=transmission)

   load_at_risk = cascade_load_loss(fault_severity)  # (6,), stochastic

This is entirely differentiable in both the node coordinates (via
``edge_lengths``/geometry feeding the cost term below) and the
per-edge weights — with ``BinaryState`` weights, the cascade *topology*
itself becomes a design variable.

Step 4 — Outage risk and reinforcement cost
------------------------------------------------

Each load node's outage probability is a smoothed threshold on the
fraction of load at risk; reinforcement has a cost proportional to line
length, charged only for lines actually reinforced:

.. code-block:: python

   outage_prob = prp.clip(prp.sigmoid(10.0 * (load_at_risk[..., 1:] - 0.3)), 1e-4, 1 - 1e-4)
   outage_events = [
       prp.UncertainAttribute(f"outage_node_{i + 1}", prp.Bernoulli(prob=outage_prob[..., i]))
       for i in range(loads.shape[0])
   ]
   model = prp.ProbabilityModel(*outage_events)

   line_cost = grid.edge_lengths() * 50.0
   reinforcement_cost = prp.sum_(line_cost * reinforced)

Note the ``...`` in ``load_at_risk[..., 1:]``: once ``fault_severity`` is
marginalized over Monte-Carlo samples, ``load_at_risk`` carries a leading
sample axis — plain ``[1:]`` would slice that axis instead of the node
axis. This is the same broadcasting discipline covered in
:ref:`marginalization`.

Step 5 — Joint design optimization
----------------------------------------

A single :func:`~prophys.engine.optimize` call minimizes expected outage
risk plus reinforcement cost with respect to *both* the relaxed
reinforcement decisions and the substation site:

.. code-block:: python

   objective = (
       prp.sum_(prp.stack([outage_prob[..., i] for i in range(loads.shape[0])], axis=0))
       + 0.01 * reinforcement_cost
   )
   reinforce_params = [p for p in model.params().values() if p.name.startswith("reinforce_")]
   result = prp.optimize(
       objective,
       wrt=[*reinforce_params, substation],
       n_steps=200,
       learning_rate=0.1,
       extra_env={"fault_severity": 1.0, "__mode__": "opt", "__tau__": 3.0},
   )

The optimizer trades off which lines are worth reinforcing against moving
the substation closer to the load cluster — a joint discrete/continuous
decision that would otherwise need a mixed-integer solver, resolved here
with plain gradient descent.

Running it
-------------

.. code-block:: console

   $ python examples/electric_grid/model.py