Feedback Control: PI(D) and Linear MPC#

difflow.control closes loops around difflow models, at the two levels a plant runs them:

  • PIDController is a dynamic unit. It sits in a DynamicFlowsheet, 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, so jax.grad of an integrated squared error tunes it.

  • LinearMPC is a receding-horizon controller on a linear(ized) discrete-time model, a difflow.mhe.StateSpaceModel linearized by autodiff or matrices you give, with bounds on the inputs. run_mpc runs it in closed loop against any JAX plant.

Both are JAX throughout: no Python control flow on traced values, no SciPy.

Table of Contents#

  1. A PI loop on a jacketed reactor

  2. PIDController

  3. Tuning with jax.grad

  4. Limits and anti-windup

  5. LinearMPC

  6. From a flowsheet to a state-space model

  7. Limits of what is here

  8. API summary


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):

\[e = s\,(r - y),\qquad v = b + K_c e + I + D,\qquad u = \mathrm{clip}(v, u_\mathrm{min}, u_\mathrm{max})\]
\[\frac{dI}{dt} = \frac{K_c}{\tau_I}\,e + \frac{u - v}{\tau_T},\qquad D = -s K_c \tau_D \frac{dy_f}{dt},\qquad \frac{dy_f}{dt} = \frac{y - y_f}{\alpha\,\tau_D}\]

Argument

Meaning

Kc

gain, output units per measured unit

tau_I

integral time (s); None for no integral state

tau_D

derivative time (s); None for no derivative state

alpha

derivative filter time as a fraction of tau_D (default 0.1)

setpoint

r: a constant, f(t), f(t, args) or an inlet port (cascade)

bias

b, the output at zero error and zero integral

u_min, u_max

output limits (None: unbounded)

tau_T

tracking time constant; default sqrt(tau_I*tau_D) with D, else tau_I; inf disables anti-windup

action

"reverse" (e = r - y, the output falls as the measurement rises) or "direct" (e = y - r)

measure

key of the measured stream ("T", "P", "F_A") or f(stream) -> value

I0

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,

\[\delta x_{k+1} = A\,\delta x_k + B\,\delta u_k + c_0,\qquad y = y_\mathrm{op} + C\,\delta x,\]

and chooses the input sequence that minimizes

\[\sum_{i=1}^{N} \lVert y_i + b - r_i\rVert^2_{Q} + \sum_{k=0}^{N-1} \lVert u_k - u_{k-1}\rVert^2_{R} + \lVert u_k - u_\mathrm{op}\rVert^2_{R_u} \quad\text{s.t.}\quad u_\mathrm{min} \le u_k \le u_\mathrm{max},\]

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 q and r_move instead.

  • One linearization. LinearMPC is linear in deviations around one operating point; re-linearize (from_model at 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, call mpc.solve yourself.

  • The measurement map’s dependence on u is ignored in the prediction.

  • simulate(params=...) reaches every unit; distinct gains for two controllers go through args.

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.