Feedback Control: PI(D) and Linear MPC#
difflow.control closes loops around difflow models, at the two levels a plant
runs them:
PIDControlleris a dynamic unit. It sits in aDynamicFlowsheet, measures a stream and drives a port, with output limits and back-calculation anti-windup. The closed loop is one JAX function of the gains, sojax.gradof an integrated squared error tunes it.LinearMPCis a receding-horizon controller on a linear(ized) discrete-time model, adifflow.mhe.StateSpaceModellinearized by autodiff or matrices you give, with bounds on the inputs.run_mpcruns it in closed loop against any JAX plant.
Both are JAX throughout: no Python control flow on traced values, no SciPy.
Table of Contents#
A PI loop on a jacketed reactor#
The controller’s output travels as a signal: a stream with no flows whose
temperature T carries the output (in whatever units it has). A
DynamicCSTR in mode="jacket" takes its coolant temperature from an inlet
port when Tc names one, so the wiring is two add_unit calls:
import jax
import jax.numpy as jnp
from difflow import make_stream
from difflow.dynamic import DynamicCSTR, DynamicFlowsheet, PIDController
jax.config.update("jax_enable_x64", True)
def rate(C, T, p):
return jnp.array([1e-3 * jnp.exp(-4000.0 * (1 / T - 1 / 340.0)) * C["A"]])
fs = DynamicFlowsheet(["A", "B"])
fs.add_feed("feed", make_stream({"A": 0.5, "B": 9.5}, T=330.0, P=2e5))
fs.add_unit(DynamicCSTR(volume=1.0, rate_fn=rate, stoich=jnp.array([[-1.0], [1.0]]),
species_order=["A", "B"], dH_rxn=jnp.array([-60e3]),
mode="jacket", UA=2e3, Tc="coolant", name="reactor"),
inlets={"inlet": "feed", "coolant": "Tc"}, outlets={"outlet": "product"})
fs.add_unit(PIDController(Kc=5.0, tau_I=200.0, setpoint=340.0, bias=330.0,
u_min=300.0, u_max=345.0, name="TC"),
inlets={"measurement": "product"}, outlets={"output": "Tc"})
y0 = fs.initial_state().at[2].set(345.0) # start 5 K hot
res = fs.simulate((0.0, 600.0), y0=y0, method="RK4", n_steps=60)
T = res.trajectory.y[:, 2] # reactor.T
Tc = res.stream_history("Tc")["T"] # the controller output over time
The reactor’s outlet reads its state and its process feed, not the coolant port
(DynamicCSTR declares this through output_dependencies), so the flowsheet
evaluates the reactor outlet, then the controller, then the reactor’s
derivatives, all in one right-hand-side call: there is no algebraic loop to
break. duty="heat" (in "specified_duty" mode) and UA="ua" take a signal
the same way. Anything else reads the signal as inputs[port]["T"] in an
InstantaneousUnit (a valve that turns it into a feed flow, say).
DynamicFlowsheetResult.stream_history(name, params, args) re-evaluates a
stream at every saved time (one vmap); pass the params and args the
simulation used.
PIDController#
The control law is the parallel (“ISA”) form with derivative action on the filtered measurement and back-calculation anti-windup (Åström and Hägglund, Advanced PID Control, ISA 2006, §3.5):
Argument |
Meaning |
|---|---|
|
gain, output units per measured unit |
|
integral time (s); |
|
derivative time (s); |
|
derivative filter time as a fraction of |
|
|
|
|
|
output limits ( |
|
tracking time constant; default |
|
|
|
key of the measured stream ( |
|
initial integral term (output units) |
States: I (the integral term, in output units, so a gain change is bumpless)
and, with derivative action, y_f, which starts at the measurement so there is
no derivative kick at t = 0. The measured stream comes in on port
"measurement" (or the only inlet that is not a setpoint port); the output
leaves on "output".
Every tuning value may be a constant, f(t) or f(t, args), and lives under
its own key in the unit’s params. fs.simulate(params=...) reaches every unit
of the flowsheet, so two controllers in one flowsheet take distinct traced gains
through args (Kc=lambda t, a: a["Kc_level"]), not through params.
Tuning with jax.grad#
Because the gains are read per call from params, the closed-loop trajectory is
a differentiable function of them:
def ise(log_gains):
Kc, tau_I = jnp.exp(log_gains)
r = fs.simulate((0.0, 600.0), y0=y0, method="RK4", n_steps=60,
params={"Kc": Kc, "tau_I": tau_I})
return jnp.sum((r.trajectory.y[:, 2] - 340.0) ** 2)
value, grad = jax.jit(jax.value_and_grad(ise))(jnp.log(jnp.array([5.0, 200.0])))
Hand value_and_grad to scipy.optimize.minimize(jac=True) or any optimizer;
optimize logarithms to keep the gains positive, and add a penalty on output
moves (jnp.diff of stream_history("Tc", params)["T"]) so the optimum is not
a controller that slams the valve.
Limits and anti-windup#
While the output is inside its limits u = v and the tracking term vanishes.
While it is pinned, a plain clipped PI keeps integrating the error: the integral
winds up, and the loop overshoots and takes long to unwind once the error
changes sign. The back-calculation term (u - v)/tau_T pulls I back toward
the value the limit implies, with time constant tau_T. A short tau_T behaves
like conditional integration (stop integrating while saturated); a long one lets
some useful integral build; tau_T=float("inf") is the clipped PI that winds up.
The clip is jnp.minimum/jnp.maximum, so gradients are exact (zero with
respect to the unclipped output while saturated).
LinearMPC#
At every sample the controller predicts horizon steps with a linear model in
deviations from an operating point,
and chooses the input sequence that minimizes
applies the first input and repeats. b, the measured output minus the model’s,
is a constant output disturbance: it makes the loop offset-free on a plant the
linear model does not match. The problem is a box-constrained QP in the
horizon * n_u inputs; box_qp solves it exactly (a primal-dual active-set
method, one dense solve per iteration, with a projected-gradient polish if the
active set has not settled) under lax.fori_loop, so the whole closed loop is a
lax.scan.
from difflow.control import LinearMPC, run_mpc
from difflow.mhe import StateSpaceModel
model = StateSpaceModel.from_ode( # a first-order lag, sampled at 1 s
lambda t, x, u, th: (2.0 * u - x) / 10.0, lambda x, u, th: x,
n_x=1, n_y=1, n_u=1, dt=1.0)
mpc = LinearMPC.from_model(model, x_op=jnp.zeros(1), u_op=jnp.zeros(1), horizon=20,
q=1.0, r_move=0.1, u_min=-1.0, u_max=1.0)
res = run_mpc(mpc, model, x0=jnp.zeros(1), n_steps=60, y_ref=1.5)
res.u.max() # 1.0: the bound binds on the way up
LinearMPC.from_model takes A, B, C as Jacobians of the model’s step and
measurement at (x_op, u_op) (process noise at zero) and c0 = f(x_op, u_op) - x_op, zero at a steady state. Matrices can also be given directly:
LinearMPC(A, B, C, horizon, x_op=..., u_op=..., y_op=...).
run_mpc(mpc, plant, x0, n_steps, y_ref) takes the plant as plant(x, u) -> x_next (any JAX function, e.g. a difflow.dynamic.integrate over one sample of
the nonlinear model) or a StateSpaceModel. measure(x) -> y defaults to the
model’s observe, or y_op + C (x - x_op). y_ref is one value for every
sample or one per sample. It returns MPCResult(x, y, u, kkt): the plant states
and measurements (n_steps + 1 rows), the applied inputs (n_steps rows) and
the QP’s KKT residual at each sample. mpc.solve(xm, u_prev, y_ref, bias) is one
sample, for a loop of your own (with a state estimate from difflow.mhe, say).
From a flowsheet to a state-space model#
StateSpaceModel.from_flowsheet(fs, h, n_y=..., dt=..., n_u=...) makes a
DynamicFlowsheet the discrete-time model both difflow.mhe and LinearMPC
read. Its state is the combined state; the input u reaches the flowsheet as
its args (map it with u_to_args=), and theta as its params:
fs2 = DynamicFlowsheet(["A", "B"])
fs2.add_feed("feed", make_stream({"A": 0.5, "B": 9.5}, T=330.0, P=2e5))
fs2.add_unit(DynamicCSTR(volume=1.0, rate_fn=rate, stoich=jnp.array([[-1.0], [1.0]]),
species_order=["A", "B"], dH_rxn=jnp.array([-60e3]),
mode="jacket", UA=2e3, Tc=lambda t, args: args[0], name="reactor"),
inlets={"inlet": "feed"}, outlets={"outlet": "product"})
plant = StateSpaceModel.from_flowsheet(fs2, lambda x, u, th: x[2:3], n_y=1, n_u=1, dt=10.0)
x_op = fs2.steady_state(args=jnp.array([330.0]))
mpc = LinearMPC.from_model(plant, x_op, jnp.array([330.0]), horizon=20,
r_move=0.05, u_min=300.0, u_max=345.0)
res = run_mpc(mpc, plant, x_op, 60, y_ref=340.0)
Limits of what is here#
Continuous-time PID. The controller integrates with the flowsheet; there is no sample-and-hold or dead time. A digital controller with a long sampling period is a discrete loop of your own.
Input bounds only in the MPC. Output constraints and move-rate limits make the QP general rather than a box and are not offered; penalize them through
qandr_moveinstead.One linearization.
LinearMPCis linear in deviations around one operating point; re-linearize (from_modelat a new point) for a nonlinear plant that moves far.No estimator inside
run_mpc. The model state is propagated open loop from the applied inputs and corrected only through the output bias. With a state estimate, callmpc.solveyourself.The measurement map’s dependence on
uis ignored in the prediction.simulate(params=...)reaches every unit; distinct gains for two controllers go throughargs.
API summary#
from difflow.control import (
PIDController, # also difflow.dynamic.PIDController
LinearMPC, # .from_model, .solve, .prediction_matrices
run_mpc, MPCResult,
box_qp, # min 0.5 u'Hu + g'u, lb <= u <= ub
)
from difflow.mhe import StateSpaceModel # .from_flowsheet
# DynamicFlowsheetResult.stream_history(name, params=None, args=None)
Tests: tests/test_control.py.