Gas transmission networks as differentiable flowsheets (difflow_gas)

This example walks through the difflow_gas plugin on a small meshed network: define the network, let the plugin compute the sequential decomposition from the topology, solve the tears, verify against the full equation set, differentiate through the converged solution, and use the gradients in a small compressor-power optimization.

The physics is the benchmark-standard steady-state isothermal model: squared-pressure Weymouth pipes (\(p_f^2 - p_t^2 = \beta\, q\,|q|\), \(q\) in kg/s, signed), compressor stations as controllable pressure ratios, and a nomination that fixes every boundary flow, with one slack node providing the pressure level.

import jax
import jax.numpy as jnp

import difflow_gas as dg

# a source feeding two sinks through a compressor and a looped grid
net = dg.GasNetwork(
    arcs={
        "p1":  ("src", "a", "pipe"),
        "cs1": ("a", "b", "compressor"),
        "p2":  ("b", "c", "pipe"),
        "p3":  ("b", "d", "pipe"),
        "p4":  ("c", "d", "pipe"),   # closes the loop b-c-d-b
    },
    beta={aid: dg.weymouth_beta(length_m=L, diameter_m=0.6,
                                roughness_m=1e-4)
          for aid, L in [("p1", 20e3), ("p2", 40e3),
                         ("p3", 60e3), ("p4", 80e3)]},
    supply_kg_s={"src": 120.0, "c": -50.0, "d": -70.0},
    pressure_bounds_bar={n: (30.0, 80.0)
                         for n in ["src", "a", "b", "c", "d"]},
)
print(f"{len(net.nodes)} nodes, {len(net.arcs)} arcs, "
      f"cycle rank {net.cycle_rank}")

5 nodes, 5 arcs, cycle rank 1

The computed decomposition

decompose picks a spanning tree (compressors and other non-invertible arcs forced in-tree; the most resistive pipe of each loop becomes the chord/tear), then schedules leaf-to-root mass balances and root-to-leaf pressure propagation. One tear per loop.

dec = dg.decompose(net, root="src")
print("tree arcs :", dec.tree_arc_ids)
print("chords    :", dec.chord_ids)
print("BFS order :", dec.order)

tree arcs : ['cs1', 'p1', 'p2', 'p3']
chords    : ['p4']
BFS order : ['src', 'a', 'b', 'c', 'd']

Build and solve

The builder translates the schedule into difflow units. Two solvers: Anderson-accelerated tears (eager, robust) and a damped fixed-point iteration that is jax.jit- and jax.grad-safe (gas tear maps have negative eigenvalues, so the raw map oscillates; damping \(x \leftarrow x + \alpha(g(x) - x)\) makes it contract while implicit differentiation keeps gradients exact).

fs, dec = dg.build_network_flowsheet(
    net, root="src", p_slack_pa=60e5, ratios={"cs1": 1.3}, dec=dec,
)
streams = fs.solve(tol=1e-8, max_iter=200)   # Anderson, signed flows
print("tear iterations:", fs.last_solve_iterations)

q = dg.verify.arc_flows_kg_s(streams, dec)
p = dg.verify.node_pressures_bar(streams, dec)
print({k: round(v, 3) for k, v in q.items()})
print({k: round(v, 3) for k, v in p.items()})

tear iterations: 46
{'p1': 120.0, 'cs1': 120.0, 'p2': 64.612, 'p3': 55.388, 'p4': 14.612}
{'src': 60.0, 'a': 51.624, 'b': 67.111, 'c': 62.942, 'd': 62.5}

Verification

A sequential solve satisfies most equations by construction; the meaningful check evaluates every equation-oriented residual (all nodal balances, every pipe law) on the solved state.

rep = dg.residual_report(streams, net, dec)
print(f"max node imbalance : {rep.max_node_imbalance_kg_s:.2e} kg/s")
print(f"max pipe residual  : {rep.max_resistance_residual_bar2:.2e} bar^2")
assert rep.ok

max node imbalance : 3.85e-13 kg/s
max pipe residual  : 9.09e-13 bar^2

Exact gradients through the converged tears

obj = fs.make_objective_fn(
    lambda s: dg.total_compressor_power_w(s, dec, net.gas_temp_k))

params = {"cs_cs1.ratio": 1.3, "src_src.P_set": 60e5}
g = jax.grad(obj)(params)
print(f"W                = {float(obj(params))/1e6:.3f} MW")
print(f"dW/dratio        = {float(g['cs_cs1.ratio'])/1e6:.3f} MW per ratio")
print(f"dW/dp_slack      = {float(g['src_src.P_set'])*1e5:.1f} W per bar")

W                = 5.832 MW
dW/dratio        = 17.622 MW per ratio
dW/dp_slack      = -0.0 W per bar

A reduced-space optimization

Minimize shaft power over the compressor ratio subject to a delivery pressure requirement at sink d. Following the GasLib benchmark lesson, the constraint is posed in squared pressure, where the network response is nearly linear (in plain pressure the sqrt makes low-pressure constraints violently nonlinear near their bounds and SQP steps overshoot).

import numpy as np
from scipy.optimize import minimize

P_MIN_D = 55.0  # bar, required delivery pressure at d


def solve_at(ratio):
    return fs._apply_params({"cs_cs1.ratio": ratio}).solve_differentiable()


power_mw = jax.jit(jax.value_and_grad(
    lambda r: dg.total_compressor_power_w(
        solve_at(r), dec, net.gas_temp_k) / 1e6))
margin = jax.jit(lambda r: (solve_at(r)["node_d"]["P"] / 1e5) ** 2
                 - P_MIN_D**2)
margin_grad = jax.jit(jax.grad(margin))


def fun(x):
    v, g = power_mw(x[0])
    return float(v), np.atleast_1d(np.asarray(g))


res = minimize(
    fun,
    x0=[1.05], jac=True, method="SLSQP", bounds=[(1.0, 2.0)],
    constraints=[{
        "type": "ineq",
        "fun": lambda x: np.atleast_1d(np.asarray(margin(x[0]))),
        "jac": lambda x: np.atleast_2d(np.asarray(margin_grad(x[0]))),
    }],
    options={"ftol": 1e-8},
)
r_opt = float(res.x[0])
s_opt = solve_at(r_opt)
print(f"optimal ratio = {r_opt:.6f}, "
      f"W = {float(res.fun)*1e6/1e6:.4f} MW, "
      f"p(d) = {float(s_opt['node_d']['P'])/1e5:.4f} bar")
assert res.success

optimal ratio = 1.165896, W = 3.3690 MW, p(d) = 55.0000 bar

The optimizer pushes the ratio to the point where the delivery constraint is exactly active: the classic shape of gas network power minimization (the same structure the GasLib-40 study found, where one station held a terminal sink at its lower bound and all others idled).