Dynamic Modeling

Dynamic Modeling#

The difflow.dynamic module provides a unified framework for transient (time-dependent) simulation of chemical process units. This enables:

Startup/shutdown analysis: Simulate process dynamics from empty to steady state
Disturbance response: Track how processes respond to feed changes
Control system design: Test controllers in simulation before implementation
Batch process modeling: Simulate time-varying batch operations
Parameter estimation: Fit kinetic parameters to dynamic experimental data

All dynamic simulations are fully differentiable via JAX, enabling gradient-based optimization of dynamic systems.

Quick Start#

import jax.numpy as jnp
from difflow.dynamic import integrate, DynamicCSTR, integrate_unit
from difflow.streams import make_stream

# Simple ODE: harmonic oscillator
def oscillator(t, y):
    return jnp.array([y[1], -y[0]])

result = integrate(oscillator, jnp.array([1.0, 0.0]), (0.0, 10.0))
print(f"Final position: {result.y_final[0]:.4f}")

# Dynamic reactor simulation
def rate_fn(C, T, params):
    return jnp.array([params["k"] * C["A"]])

cstr = DynamicCSTR(
    volume=1.0,
    rate_fn=rate_fn,
    stoich=jnp.array([[-1], [1]]),
    species_order=["A", "B"],
    rate_params={"k": 0.1},
)

inlet = make_stream({"A": 1.0, "B": 0.0}, T=350.0, P=101325.0)
result = integrate_unit(cstr, {"inlet": inlet}, (0.0, 100.0))

ODE Integration#

The `integrate()` Function#

The unified interface for ODE integration:

from difflow.dynamic import integrate

def f(t, y):
    """Derivative function: dy/dt = f(t, y)"""
    return -0.5 * y  # Exponential decay

result = integrate(
    f,                          # Derivative function
    y0=jnp.array([1.0]),       # Initial state
    t_span=(0.0, 10.0),        # Time interval
    method="RK4",              # Integration method
    n_steps=100,               # Number of steps (fixed-step methods)
)

# Access results
print(f"Final state: {result.y_final}")
print(f"Success: {result.info.success}")
print(f"Steps taken: {result.info.n_steps}")

# Trajectory (intermediate values)
print(f"Time points: {result.trajectory.t.shape}")
print(f"State history: {result.trajectory.y.shape}")

Available Methods#

Method	Description	Use Case
`"RK4"`	4th order Runge-Kutta	General purpose, fixed step
`"RK45"`	Adaptive RK45 (Dormand-Prince)	When accuracy matters
`"Euler"`	Forward Euler	Simple problems, debugging
`"diffrax"`	Default diffrax solver (Tsit5)	Advanced features
`"diffrax:dopri5"`	Dormand-Prince 5(4)	Good general purpose
`"diffrax:kvaerno5"`	Implicit 5th order	Stiff systems

Integration Options#

y0 = jnp.array([1.0])
t_span = (0.0, 10.0)

# Fixed-step methods (RK4, Euler)
result = integrate(f, y0, t_span, method="RK4", n_steps=1000)

# Adaptive methods (RK45)
result = integrate(f, y0, t_span, method="RK45", rtol=1e-6, atol=1e-8)

# Diffrax methods
result = integrate(
    f, y0, t_span,
    method="diffrax:tsit5",
    rtol=1e-5,
    atol=1e-7,
    max_steps=10000,
)

Dynamic Units#

The DynamicUnit Protocol#

All dynamic units implement the DynamicUnit protocol:

from typing import Protocol

class DynamicUnit(Protocol):
    @property
    def name(self) -> str: ...

    @property
    def state_spec(self) -> StateSpec: ...

    def initial_state(self, inputs: dict) -> Array: ...

    def derivatives(self, t: Array, y: Array, inputs: dict) -> Array: ...

    def outputs(self, t: Array, y: Array, inputs: dict) -> dict[str, Stream]: ...

Refinery units#

difflow_refinery.dynamic builds dynamic versions of the refinery’s hydrotreater and reformer on this protocol – reactor beds as cells with holdups and catalyst heat capacity, furnace lags, separator levels and a recycle-gas loop – from the solved steady units, so that their rest state is the steady answer. See Refinery Unit Operations: Dynamic operation and examples/66_refinery_dynamic_hydrotreater.ipynb.

DynamicCSTR#

Continuous stirred-tank reactor with reaction kinetics:

from difflow.dynamic import DynamicCSTR
import jax.numpy as jnp

# Define rate law: r = k * exp(-Ea/RT) * C_A
def rate_fn(C, T, params):
    k = params["k0"] * jnp.exp(-params["Ea"] / (8.314 * T))
    return jnp.array([k * C["A"]])  # Rate of reaction 1

# A -> B (single reaction)
stoich = jnp.array([
    [-1],  # A consumed
    [+1],  # B produced
])

cstr = DynamicCSTR(
    volume=1.0,                    # m³
    rate_fn=rate_fn,
    stoich=stoich,
    species_order=["A", "B"],
    rate_params={"k0": 1e6, "Ea": 50000.0},
    name="reactor",
)

# State variables: [n_A, n_B] (moles of each species)

mode sets the energy balance: "isothermal" (no temperature state; the reactor runs at the inlet temperature), "adiabatic" (Q = 0), "specified_duty" (Q = duty, W, positive heats the reactor) and "jacket" (Q = UA * (Tc - T), with UA in W/K and coolant temperature Tc in K). duty, UA and Tc are each a constant, a function of time f(t), or f(t, args) reading the args passed to DynamicFlowsheet.simulate. Each may also name an inlet port, whose stream temperature T is then the value: a coolant stream for Tc, or the output signal of a PIDController for any of the three. The outlet does not read those ports, so a controller that measures the outlet and drives the jacket closes the loop in one DynamicFlowsheet without an algebraic loop. They live in the unit’s params under the same keys, so fs.simulate(..., params={"UA": ua}) overrides them for that call and jax.grad goes through them. The heat capacity is fixed at 75 J/(mol K) and the molar density at 55 500 mol/m³ (a water-like liquid).

import jax
import jax.numpy as jnp
from difflow.dynamic import DynamicCSTR, integrate_unit
from difflow.streams import make_stream

def rate_fn(C, T, params):
    return jnp.array([params["k0"] * jnp.exp(-params["Ea"] / (8.314 * T)) * C["A"]])

feed = {"inlet": make_stream({"A": 1.0, "B": 50.0}, T=330.0, P=2e5)}

def T_after_10_min(UA):
    cstr = DynamicCSTR(volume=1.0, rate_fn=rate_fn, stoich=jnp.array([[-1.0], [1.0]]),
                       species_order=["A", "B"], rate_params={"k0": 1e6, "Ea": 50000.0},
                       dH_rxn=jnp.array([-60e3]), mode="jacket",
                       UA=UA, Tc=lambda t: 320.0 + 5.0 * (t > 300.0))  # coolant step at 5 min
    res = integrate_unit(cstr, feed, (0.0, 600.0), method="RK4", n_steps=120)
    return res.y_final[-1]                                    # [n_A, n_B, T]: T last

print(T_after_10_min(2000.0), jax.grad(T_after_10_min)(2000.0))   # K, K/(W/K)

DynamicTank#

Storage tank with variable holdup:

from difflow.dynamic import DynamicTank

tank = DynamicTank(
    max_volume=10.0,              # Maximum volume (m³)
    species_order=["A", "B"],
    name="storage",
)

# State variables: [V, n_A, n_B] (volume + moles; V starts at 1 m³)

Using integrate_unit()#

Convenience wrapper for simulating dynamic units:

from difflow.dynamic import integrate_unit
from difflow.streams import make_stream

inlet = make_stream({"A": 1.0, "B": 0.0}, T=350.0, P=101325.0)

result = integrate_unit(
    cstr,
    inputs={"inlet": inlet},
    t_span=(0.0, 1000.0),
    method="RK4",
    n_steps=500,
)

# Final moles
n_A_final, n_B_final = result.y_final
print(f"Final A: {n_A_final:.4f} mol")
print(f"Final B: {n_B_final:.4f} mol")

State Specification#

StateVar and StateSpec#

Define state variables with metadata:

from difflow.dynamic import StateVar, StateSpec

# Manual specification
spec = StateSpec([
    StateVar("n_A", "moles", "mol", bounds=(0, None)),
    StateVar("n_B", "moles", "mol", bounds=(0, None)),
    StateVar("T", "temperature", "K", bounds=(200, 600)),
])

print(f"State dimension: {spec.n_states}")
print(f"State names: {spec.names}")
print(f"Index of T: {spec.get_index('T')}")

Factory Functions#

Convenience functions for common state patterns:

from difflow.dynamic import (
    molar_states,
    concentration_states,
    thermal_state,
    volume_state,
    reactor_states,
)

# Molar holdup for species
spec = molar_states(["A", "B", "C"])  # n_A, n_B, n_C

# Concentration states
spec = concentration_states(["A", "B"])  # C_A, C_B

# Combine specs
spec = molar_states(["A", "B"]) + thermal_state()  # n_A, n_B, T

# Complete reactor state
spec = reactor_states(["A", "B"])  # n_A, n_B, T (moles + temperature)

StateVector#

Runtime access to state values by name:

from difflow.dynamic import StateVector

spec = molar_states(["A", "B"]) + thermal_state()
y = jnp.array([1.0, 0.5, 350.0])

state = StateVector(y, spec)
print(f"n_A = {state['n_A']}")
print(f"T = {state['T']}")

# Or use spec directly
idx_A = spec.get_index("n_A")
n_A = y[idx_A]

Dynamic Flowsheets#

Building a Flowsheet#

Connect multiple dynamic units:

from difflow.dynamic import DynamicFlowsheet, DynamicCSTR, DynamicTank
from difflow.streams import make_stream

# Create units
cstr = DynamicCSTR(
    volume=1.0,
    rate_fn=rate_fn,
    stoich=stoich,
    species_order=["A", "B"],
    rate_params={"k0": 1e6, "Ea": 50000.0},
    name="reactor",
)

tank = DynamicTank(
    max_volume=10.0,
    species_order=["A", "B"],
    name="storage",
)

# Build flowsheet
fs = DynamicFlowsheet(species_order=["A", "B"])

# Add feed stream
feed = make_stream({"A": 1.0, "B": 0.0}, T=350.0, P=101325.0)
fs.add_feed("feed", feed)

# Add units with connections
fs.add_unit(cstr, inlet_names=["feed"], outlet_names=["reactor_out"])
fs.add_unit(tank, inlet_names=["reactor_out"], outlet_names=["product"])

# Simulate
result = fs.simulate(t_span=(0.0, 1000.0), method="RK4", n_steps=500)

Accessing Results#

# Combined final state
print(f"Final state: {result.y_final}")

# Per-unit states
reactor_state = result.unit_state_at("reactor")
tank_state = result.unit_state_at("storage")

# Trajectory
print(f"Time points: {result.trajectory.t.shape}")

# Output streams at final time
streams = fs.outputs(result.trajectory.t[-1], result.y_final)
print(f"Product stream: {streams['product']}")

Time-Varying Feeds#

# Define feed as function of time
def feed_schedule(t):
    """Feed rate doubles after t=500."""
    base = make_stream({"A": 1.0, "B": 0.0}, T=350.0, P=101325.0)
    scale = jnp.where(t > 500, 2.0, 1.0)  # traced t: use jnp.where, not `if`
    return {k: (v * scale if k.startswith("F_") else v) for k, v in base.items()}

fs.add_feed("feed", feed_schedule)  # Pass function instead of stream

Disturbances as Integrator Args#

A feed written (t, args) reads the args passed to simulate. Those args, and params, go to the integrator as its own args rather than being closed over, so with a diffrax method one compiled right-hand side serves every scenario, and jax.grad with respect to a disturbance does not retrace. Make the values you vary JAX arrays (a Python float is static).

def feed_step(t, args):
    scale = 1.0 + args["step"] * (t > args["t_step"])
    return make_stream({"A": scale, "B": 0.0}, T=350.0, P=101325.0)

fs_args = DynamicFlowsheet(species_order=["A", "B"])
fs_args.add_feed("feed", feed_step)
fs_args.add_unit(cstr, inlets={"inlet": "feed"}, outlets={"outlet": "reactor_out"})
for step in (0.1, 0.2, 0.5):  # compiles once
    result = fs_args.simulate((0.0, 1000.0), method="diffrax:tsit5",
                              args={"step": jnp.array(step), "t_step": jnp.array(500.0)})

integrate(f, y0, t_span, method, args=...) does the same for a bare f(t, y, args).

Ports, Stateless Units and Feed-Effluent Loops#

Units read fixed port names ("inlet" for a CSTR, "hot"/"cold" for DynamicCounterCurrentHX, "feed" for DynamicEOSFlash). Wire them to freely named streams with inlets=/outlets= maps of {port: stream}. A unit with no holdup worth a state (a trim heater or cooler at a setpoint, a mixer) goes in as an InstantaneousUnit, which wraps a steady-state unit operation, or a function fn(inputs, params) -> {port: Stream}, with zero states.

A setpoint that moves with time takes t, chosen by the callable’s arity (as add_feed tells (t) from (t, args)): fn(t, inputs, params), or a call_params callable v(t, params). Add args as a last parameter – fn(t, inputs, params, args), v(t, params, args) – to read the pytree given to simulate(args=...), the same one (t, args) feeds read. Put time-varying disturbances in args (a feed step and a setpoint ramp then share one dict); params is for design parameters.

def ramp(t, params, args):
    return args["T0"] + (args["T1"] - args["T0"]) * jnp.clip((t - args["t0"]) / args["dur"], 0.0, 1.0)

trim = InstantaneousUnit(Cooler(CoolerParams(T_out=320.0), thermo),
                         call_params={"T_out": ramp}, name="trim")

from difflow import Heater, HeaterParams, Cooler, CoolerParams
from difflow.dynamic import (DynamicCounterCurrentHX, DynamicEOSFlash,
                             InstantaneousUnit)

fs = DynamicFlowsheet(species_order=species)
fs.add_feed("naphtha", feed)
fs.add_unit(DynamicCounterCurrentHX(150.0, thermo, name="fehe"),
            inlets={"hot": "effluent", "cold": "naphtha"},
            outlets={"hot_out": "eff_cooled", "cold_out": "preheated"})
fs.add_unit(InstantaneousUnit(Heater(HeaterParams(T_out=460.0), thermo), name="trim"),
            inlets={"inlet": "preheated"}, outlets={"outlet": "rx_in"})
fs.add_unit(reactor, inlets={"inlet": "rx_in"}, outlets={"outlet": "effluent"})
fs.add_unit(InstantaneousUnit(Cooler(CoolerParams(T_out=320.0), thermo), name="cooler"),
            inlets={"inlet": "eff_cooled"}, outlets={"outlet": "cooled"})
fs.add_unit(DynamicEOSFlash(eos, species, P=8e5, name="sep"),
            inlets={"feed": "cooled"}, outlets={"liquid": "liq", "vapor": "vap"})

The exchanger’s cold outlet feeds the heater, the heater the reactor, and the reactor the exchanger’s hot inlet: a cycle through the units’ outputs within one right-hand-side call. It is well posed because each exchanger outlet reads its duty state and its own side’s inlet only. A unit declares that with output_dependencies = {"cold_out": ("cold",), ...} and implements partial_outputs, and the flowsheet orders the evaluation from stream availability (cold side, heater, reactor, then hot side) instead of insertion order. A cycle that no declared dependency breaks raises ValueError (“Algebraic loop …”) naming the units and streams.

Manual Derivative Access#

For custom integration or analysis:

# Get combined initial state
y0 = fs.initial_state()

# Define derivative function
def flowsheet_f(t, y):
    return fs.derivatives(t, y)

# Use with any integrator
result = integrate(flowsheet_f, y0, (0.0, 1000.0), method="RK45")

DAE Systems#

Differential-Algebraic Equations combine ODEs with algebraic constraints:

dx/dt = f(t, x, z)    # Differential equations
0 = g(t, x, z)        # Algebraic constraints

DAE Units#

from difflow.dynamic import DAEUnitBase, AlgebraicSpec, AlgebraicVar

class MyFlashDrum(DAEUnitBase):
    @property
    def algebraic_spec(self) -> AlgebraicSpec:
        return AlgebraicSpec([
            AlgebraicVar("V_frac", "vapor_fraction", "-", bounds=(0, 1)),
        ])

    def algebraic_residual(self, t, x, z, inputs):
        """Return g(t,x,z) - should equal zero at solution."""
        V_frac = z[0]
        # VLE constraint: sum(z_i * (K_i - 1) / (1 + V*(K_i-1))) = 0
        residual = self._rachford_rice(x, V_frac)
        return jnp.array([residual])

    def differential(self, t, x, z, inputs):
        """Return dx/dt given algebraic variables."""
        # Material balances using vapor fraction
        ...

Built-in: DynamicFlashDrum#

from difflow.dynamic import DynamicFlashDrum, integrate_dae

flash = DynamicFlashDrum(
    volume=1.0,
    species_order=["A", "B"],
    K_func=lambda T: jnp.array([2.0, 0.5]),  # K-values (A, B) vs. temperature
    name="flash",
)

inlet = make_stream({"A": 0.5, "B": 0.5}, T=350.0, P=101325.0)

result = integrate_dae(
    flash,
    inputs={"inlet": inlet},
    t_span=(0.0, 100.0),
    method="RK4",
    n_steps=200,
)

print(f"Final moles: {result.x_final}")
print(f"Final vapor fraction: {result.z_final}")

Newton Solver#

The algebraic constraints are solved at each time step:

from difflow.dynamic import newton_solve

def residual(z):
    """System of nonlinear equations."""
    x, y = z[0], z[1]
    return jnp.array([
        x**2 + y**2 - 5.0,  # Circle
        x * y - 2.0,         # Hyperbola
    ])

z0 = jnp.array([2.0, 1.0])
z_solution, info = newton_solve(residual, z0, tol=1e-8, max_iter=50)

print(f"Solution: {z_solution}")
print(f"Converged: {info['converged']}")
print(f"Residual norm: {jnp.linalg.norm(residual(z_solution)):.2e}")

DAE Integration Methods#

# Euler method (simpler, may need smaller steps)
result = integrate_dae(flash, {"inlet": inlet}, (0.0, 100.0), method="Euler", n_steps=1000)

# RK4 method (more accurate)
result = integrate_dae(flash, {"inlet": inlet}, (0.0, 100.0), method="RK4", n_steps=200)

Diffrax Backend#

Diffrax provides advanced ODE/SDE solvers with adaptive step control.

Installation#

pip install diffrax

Basic Usage#

from difflow.dynamic import integrate

# Use diffrax with default solver (Tsit5)
result = integrate(f, y0, t_span, method="diffrax")

# Specify solver
result = integrate(f, y0, t_span, method="diffrax:dopri5")

Available Solvers#

Explicit (non-stiff problems):

dopri5: Dormand-Prince 5(4) - good general purpose
dopri8: Dormand-Prince 8(7) - higher accuracy
tsit5: Tsitouras 5(4) - efficient, recommended default
bosh3: Bogacki-Shampine 3(2)
heun: Heun’s method (2nd order, fixed step)
euler: Forward Euler (1st order, fixed step)

Implicit (stiff problems):

kvaerno3: 3rd order implicit
kvaerno4: 4th order implicit
kvaerno5: 5th order implicit - recommended for stiff
implicit_euler: Backward Euler

Stiff Systems#

Chemical kinetics often involve very different time scales (stiff):

from difflow.dynamic import integrate, integrate_stiff

# Robertson problem - classic stiff test
def robertson(t, y):
    k1, k2, k3 = 0.04, 3e7, 1e4
    A, B, C = y[0], y[1], y[2]
    return jnp.array([
        -k1*A + k3*B*C,
        k1*A - k2*B*B - k3*B*C,
        k2*B*B,
    ])

y0 = jnp.array([1.0, 0.0, 0.0])

# Using implicit solver
result = integrate(
    robertson, y0, (0.0, 1e5),
    method="diffrax:kvaerno5",
    rtol=1e-4, atol=1e-6,
)

# Or use convenience function
result = integrate_stiff(robertson, y0, (0.0, 1e5))

Tolerance Control#

result = integrate(
    f, y0, t_span,
    method="diffrax:tsit5",
    rtol=1e-6,      # Relative tolerance
    atol=1e-8,      # Absolute tolerance
    max_steps=10000, # Maximum integration steps
)

Direct Diffrax API#

For more control:

from difflow.dynamic import integrate_diffrax

result = integrate_diffrax(
    f, y0, (0.0, 100.0),
    solver="tsit5",
    rtol=1e-5,
    atol=1e-7,
    dt0=0.01,              # Initial step size
    saveat=jnp.linspace(0, 100, 101),  # Save at specific times
)

Gradient Computation#

All dynamic simulations are differentiable via JAX.

Gradients Through Integration#

import jax

def loss(y0):
    """Loss based on final state."""
    result = integrate(f, y0, (0.0, 10.0), method="RK4")
    return jnp.sum(result.y_final**2)

# Gradient w.r.t. initial condition
y0 = jnp.array([1.0, 0.0])
grad_y0 = jax.grad(loss)(y0)

Parameter Optimization#

def simulate_with_params(k):
    """Simulate reactor with rate constant k."""
    def rate_fn(C, T, params):
        return jnp.array([k * C["A"]])

    cstr = DynamicCSTR(
        volume=1.0, rate_fn=rate_fn, stoich=stoich,
        species_order=["A", "B"], rate_params={},
    )

    result = integrate_unit(cstr, {"inlet": inlet}, (0.0, 100.0))
    return result.y_final[1]  # Final product amount

# Optimize for maximum product
grad_k = jax.grad(simulate_with_params)(jnp.array(0.1))

Sensitivity Analysis#

from difflow.dynamic import sensitivity_analysis

def f_p(t, y, params):
    """dy/dt with an explicit parameter array (here the decay rate k)."""
    return -params[0] * y

k_nominal = jnp.array([0.5])
result, sens = sensitivity_analysis(
    f_p, jnp.array([1.0]), k_nominal, (0.0, 10.0), method="RK4", n_steps=100,
)
# sens = d y_final / d params (Jacobian, shape (n_states, n_params))

Using integrate_with_grad#

Explicit gradient computation:

from difflow.dynamic import integrate_with_grad

# Returns both result and gradient function
result, grad_fn = integrate_with_grad(f, y0, t_span)

# Compute gradient of final state w.r.t. y0
dy_final_dy0 = grad_fn(jnp.ones_like(result.y_final))

API Reference#

Integration Functions#

Function	Description
`integrate(f, y0, t_span, method, **kwargs)`	Unified ODE integration
`integrate_unit(unit, inputs, t_span, **kwargs)`	Integrate DynamicUnit
`integrate_dae(unit, inputs, t_span, **kwargs)`	Integrate DAE unit
`integrate_rk4(f, y0, t_span, n_steps)`	Fixed-step RK4
`integrate_rk45(f, y0, t_span, rtol, atol)`	Adaptive RK45
`integrate_euler(f, y0, t_span, n_steps)`	Forward Euler

Diffrax Functions#

Function	Description
`integrate_diffrax(f, y0, t_span, solver, **kwargs)`	Direct diffrax integration
`integrate_stiff(f, y0, t_span, **kwargs)`	Stiff system integration
`integrate_dopri5(f, y0, t_span, **kwargs)`	Dormand-Prince 5(4)
`integrate_tsit5(f, y0, t_span, **kwargs)`	Tsitouras 5(4)
`list_diffrax_solvers()`	List available solvers
`check_diffrax_available()`	Check if diffrax installed

Classes#

Class	Description
`DynamicUnit`	Protocol for dynamic units
`DynamicUnitBase`	Base class with utilities
`DynamicCSTR`	Dynamic CSTR reactor
`DynamicTank`	Storage tank with holdup
`PIDController`	PI(D) controller with limits and anti-windup (control)
`DynamicFlowsheet`	Multi-unit flowsheet
`DAEUnit`	Protocol for DAE units
`DAEUnitBase`	Base class for DAE units
`DynamicFlashDrum`	Flash drum with VLE

State Classes#

Class	Description
`StateVar`	Single state variable spec
`StateSpec`	Collection of states
`StateVector`	Runtime state access
`AlgebraicVar`	Algebraic variable spec
`AlgebraicSpec`	Collection of algebraic vars

Result Classes#

Class	Description
`IntegrationResult`	ODE integration result
`Trajectory`	Time series of states
`IntegrationInfo`	Solver statistics
`DAEResult`	DAE integration result
`DynamicFlowsheetResult`	Flowsheet simulation result