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 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.

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 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.

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

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:

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:

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 Marginalizing over random inputs.

Step 5 — Joint design optimization

A single optimize() call minimizes expected outage risk plus reinforcement cost with respect to both the relaxed reinforcement decisions and the substation site:

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

$ python examples/electric_grid/model.py