<p align="center"><img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/difflow-logo.svg" alt="difflow -- Differentiable Flowsheets" width="420"></p>
# difflow
[](https://github.com/jkitchin/differentiable-flowsheets/actions/workflows/test.yml)
[](https://pypi.org/project/difflow/)
[](https://pepy.tech/projects/difflow)
[](https://pypi.org/project/difflow/)
[](https://github.com/jkitchin/differentiable-flowsheets/blob/main/LICENSE)
[](https://doi.org/10.5281/zenodo.21881034)
**Differentiable Flowsheet Framework for Chemical Processes**
A JAX-based framework for building and optimizing chemical process flowsheets with automatic differentiation.
- **Source code:** <https://github.com/jkitchin/differentiable-flowsheets>
- **Documentation:** <https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/>
- **Issue tracker:** <https://github.com/jkitchin/differentiable-flowsheets/issues>
## Features
- **Fully Differentiable**: All unit operations and flowsheet calculations support automatic differentiation via JAX
- **Sensitivity Analysis**: Compute gradients of outputs with respect to any inputs, parameters, or operating conditions
- **Optimization Ready**: Use gradient-based optimization for process design, parameter estimation, and economic optimization
- **Modular Design**: Unit operations can be composed into complex flowsheets with recycle streams
- **Technoeconomic Analysis**: Comprehensive TEA module with equipment costs, operating costs, and profitability metrics (NPV, IRR, MSP)
- **Bio Manufacturing**: Specialized unit operations for biopharmaceutical processes (bioreactors, chromatography, filtration)
- **Gas Networks**: Steady-state gas transmission networks with a topology-computed sequential decomposition and differentiable tear solving
- **Flowsheets Without Code**: A machine-readable catalog of every unit, JSON round trip, Python code generation, and parameter sweeps with exact gradients (a browser editor comes with the difflow book)
- **Agent Tools (MCP)**: `difflow mcp` serves difflow to AI agents over the Model Context Protocol: build, solve, diagnose and converge flowsheets, step through a solve with the debugger, and run sensitivity, optimization, uncertainty and cost studies from plain-language requests
- **Flowsheet Debugger**: `difflow debug plant.py` pauses a solve unit by unit and tear iteration by tear iteration, with breakpoints on a unit, a condition or the first NaN; record a solve and replay it with a different tear guess
### Plugins
Six domain plugins ship with difflow:
| | Plugin | Domain |
|---|---|---|
| <img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/plugins/difflow-bio-icon.svg" alt="" width="24"> | [`difflow_bio`](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/unit-operations-bio.html) | Biomanufacturing |
| <img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/plugins/difflow-ree-icon.svg" alt="" width="24"> | [`difflow_ree`](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/unit-operations-ree.html) | Rare earth separations |
| <img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/plugins/difflow-cc-icon.svg" alt="" width="24"> | [`difflow_cc`](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/unit-operations-carbon-capture.html) | Carbon capture |
| <img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/plugins/difflow-gas-icon.svg" alt="" width="24"> | [`difflow_gas`](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/unit-operations-gas.html) | Gas transmission networks |
| <img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/plugins/difflow-power-icon.svg" alt="" width="24"> | [`difflow_power`](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/unit-operations-power.html) | Electrical grids |
| <img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/plugins/difflow-refinery-icon.svg" alt="" width="24"> | [`difflow_refinery`](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/unit-operations-refinery.html) | Petroleum refining |
| <img src="https://raw.githubusercontent.com/jkitchin/differentiable-flowsheets/main/images/plugins/difflow-water-icon.svg" alt="" width="24"> | [`difflow_water`](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/unit-operations-water.html) | Desalination and water treatment |
## ⚠️ ALPHA SOFTWARE
This project is under active development and not ready for production use. APIs may change without notice. This notice will be removed when the project reaches stable release.
It is highly recommended that you confirm the equations and physical properties used in the models you make; these were generated by Claude. We have endeavored to ensure they seem reasonable, but cannot guarantee they are accurate in all cases.
This package uses jax solvers (e.g. diffrax, optimistix, etc.) and does not rely on IPOPT or pounce, or pyomo / IDAES. It is a pure Python / jax focused package that was developed as a proof of concept.
## LLM usage
Claude Code is heavily used to generate the code, examples and tests. This has allowed the project to develop faster than it can be used, and to develop more features than are immediately needed. This may mean there are modules that do not match the performance or output of other projects. You should perform your own diligence when using the code to ensure the library does what you expect it to. Ultimately this is a proof of concept in differentiable flowsheets that wouldn't be possible without Claude Code.
We regularly run all of the notebooks to ensure they run without errors, and review them to make sure the results make sense. We are happy to take issues and / or pull requests to fix problems. We also use Claude to review the code to look for issues.
We actually anticipate that Claude Code is used when using this library (See [CLAUDE.md](https://github.com/jkitchin/differentiable-flowsheets/blob/main/CLAUDE.md)). The library is large enough that it would take a long time to learn all the capabilities in addition to learning the nuances of differentiable programming. This repo provides all the information Claude needs to help you translate your flowsheet ideas into differentiable programs.
## Using difflow with an AI agent (MCP)
difflow includes an [MCP](https://modelcontextprotocol.io) server, so an AI agent such as Claude Code or Claude Desktop can use it directly: find unit operations, build and solve flowsheets, explain why a solve fails and search for settings that converge it, and run sensitivity, optimization, uncertainty and capital-cost studies, all with exact derivatives through the solve.
```bash
pip install "difflow[mcp]" # 0.4.0 or later
claude mcp add difflow -- difflow mcp
```
Then ask in plain language, for example "build a flash drum for an equimolar water and ethanol feed at 362 K and tell me the vapor composition", or "this flowsheet does not converge; find out why". The server discovers operations and plugins from the installed code, so new units are available without changes to it. `difflow mcp --no-exec` leaves out the tools that run Python, for shared or desktop clients. Setup for Claude Desktop and other clients, the full tool list and troubleshooting are in the [agent documentation](https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/docs/agents.html).
### Agent skills, commands and hooks
The repository also ships [Agent Skills](https://agentskills.io) for any
compatible agent (Claude Code, Codex, OpenCode, Gemini CLI, Cursor, Copilot),
plus Claude Code slash commands, a session hook and the MCP server config,
in the layout [knowhow](https://github.com/jkitchin/knowhow) and
[`npx skills`](https://github.com/vercel-labs/skills) install from:
| Item | Name | What it does |
|---|---|---|
| skill | `difflow` | Build, solve, differentiate and debug flowsheets with difflow |
| skill | `difflow-add-unit` | Add a unit operation to the difflow codebase by its conventions |
| skill | `difflow-gradient-check` | Check gradients against finite differences; chase NaNs |
| command | `test`, `test-unit`, `lint`, `example`, `optimize` | Claude Code slash commands for working on difflow |
| hook | `difflow-session-start` | Session-start status: environment, difflow import, JAX backend |
| mcp | `difflow`, `difflow-no-exec` | The `difflow mcp` server, with or without the Python tools |
```bash
uvx knowhow add jkitchin/differentiable-flowsheets --list
uvx knowhow add jkitchin/differentiable-flowsheets -s difflow -g # for every project
uvx knowhow add jkitchin/differentiable-flowsheets --mcp difflow -a claude-code
# or: npx skills add jkitchin/differentiable-flowsheets --skill difflow
```
They live in `skills/`, `commands/`, `hooks/` and `mcp/`; check them with
`uvx knowhow validate .`.
## Installation
```bash
# From PyPI
pip install difflow
# With examples and tutorials (includes matplotlib, jupyter)
pip install "difflow[examples]"
# The MCP server for AI agents (see "Using difflow with an AI agent")
pip install "difflow[mcp]"
# Everything
pip install "difflow[all]"
```
For development, install from source:
```bash
git clone https://github.com/jkitchin/differentiable-flowsheets.git
cd differentiable-flowsheets
uv venv
uv pip install -e ".[dev]"
# Install everything
uv pip install -e ".[all]"
```
## Quick Start
```python
import jax.numpy as jnp
import jax
from difflow import (
make_stream, get_flows,
IdealThermo, SpeciesData,
CSTR, CSTRParams,
)
# Define species
species_data = {
"A": SpeciesData("A", MW=100.0, Cp_coeffs=(75.0, 0.0, 0.0, 0.0),
Hvap_coeffs=(35000.0, 0.38, 500.0),
antoine_coeffs=(10.0, 3000.0, -50.0)),
"B": SpeciesData("B", MW=100.0, Cp_coeffs=(75.0, 0.0, 0.0, 0.0),
Hvap_coeffs=(30000.0, 0.38, 450.0),
antoine_coeffs=(10.0, 2800.0, -40.0)),
}
thermo = IdealThermo(species_data)
# Define reaction kinetics
def rate_fn(C, T, params):
k = params["A"] * jnp.exp(-params["Ea"] / (8.314 * T))
return jnp.array([k * C["A"]])
# Create CSTR
stoich = jnp.array([[-1.0], [+1.0]]) # A → B
cstr_params = CSTRParams(
V=jnp.array(1.0),
rate_fn=rate_fn,
stoich=stoich,
rate_params={"A": jnp.array(1e6), "Ea": jnp.array(50000.0)},
species_order=["A", "B"],
# Concentration basis: tau = V*rho/F. Pass eos=<cubic EOS> +
# reaction_phase instead for the real density at reactor conditions;
# with neither, the CSTR falls back to liquid water and warns.
molar_density=55500.0,
)
cstr = CSTR(cstr_params, thermo=thermo, mode="isothermal")
# Run simulation
inlet = make_stream({"A": 10.0, "B": 0.0}, T=300.0, P=101325.0)
outlet, info = cstr(inlet, T_spec=350.0)
print(f"Conversion: {info['conversion']['A']*100:.1f}%")
# Compute gradient of outlet B w.r.t. reactor volume
def outlet_B(V):
params = CSTRParams(V=V, rate_fn=rate_fn, stoich=stoich,
rate_params={"A": jnp.array(1e6), "Ea": jnp.array(50000.0)},
species_order=["A", "B"])
cstr = CSTR(params, thermo=thermo, mode="isothermal")
outlet, _ = cstr(inlet, T_spec=350.0)
return outlet["F_B"]
dFB_dV = jax.grad(outlet_B)(jnp.array(1.0))
print(f"dF_B/dV = {dFB_dV:.4f} mol/s per m³")
```
## Unit Operations
### CSTR (Continuous Stirred Tank Reactor)
- Multiple reactions with user-defined kinetics
- Isothermal, adiabatic, or specified heat duty modes
- Automatic material and energy balance solving
### PFR (Plug Flow Reactor)
- ODE-based design equation: dF/dV = stoich @ r
- Isothermal or adiabatic operation
- **GasPFR** variant for gas-phase reactions with:
- Pressure drop (Ergun equation)
- Variable volumetric flow from mole change
- RK4 integration via `lax.scan` (fully differentiable)
```python
from difflow import PFR, PFRParams, GasPFR, GasPFRParams
# rate_fn, stoich and inlet are those of the Quick Start
params = {"A": jnp.array(1e6), "Ea": jnp.array(50000.0)} # rate_fn's parameters
# Liquid-phase PFR
pfr = PFR(PFRParams(V=2.0, rate_fn=rate_fn, stoich=stoich,
rate_params=params, species_order=["A", "B"]))
outlet, info = pfr(inlet, volumetric_flow=0.01, T_spec=350.0) # m³/s
# Gas-phase with pressure drop (A → 2B, mole increase)
gas_pfr = GasPFR(GasPFRParams(V=1.0, rate_fn=rate_fn, stoich=stoich,
rate_params=params, species_order=["A", "B"],
alpha=50000.0)) # Pressure drop parameter
outlet, info = gas_pfr(inlet, T_spec=500.0)
# info contains: conversion, profiles (V, F, T, P, Q), pressure_drop
```
### Flash Separator
- TP flash (temperature and pressure specified)
- Rachford-Rice equation for VLE
- Ideal thermodynamics (Raoult's law)
### Liquid-Liquid Extraction (LLE)
- **MultistageCascade**: Counter-current or co-current mixer-settler cascade
- Kremser equation for stage calculations (differentiable in n_stages)
- Continuous stage relaxation for optimization
- **DifferentialContactor**: Packed column extractor
- HETP-based equilibrium model
- Rate-based mass transfer model
- **Equilibrium Models**:
- Distribution coefficients (K-values) with temperature dependence
- NRTL activity coefficient model
- UNIQUAC activity coefficient model
```python
from difflow import (
MultistageCascade, CascadeParams,
LLEEquilibrium, DistributionCoeffs,
)
# Define distribution coefficients for rare earth extraction
K_coeffs = DistributionCoeffs(
species=("La", "Nd", "Dy"),
K0=(0.5, 2.0, 8.0), # K at reference temperature
)
equilibrium = LLEEquilibrium(
solutes=["La", "Nd", "Dy"],
aqueous_carrier="H2O",
organic_carrier="Organic",
K_coeffs=K_coeffs,
)
cascade = MultistageCascade(CascadeParams(
n_stages=5,
equilibrium=equilibrium,
flow_config="counter_current",
))
feed_stream = make_stream({"La": 1.0, "Nd": 1.0, "Dy": 0.5, "H2O": 50.0, "Organic": 0.0},
T=298.15, P=101325.0)
solvent_stream = make_stream({"La": 0.0, "Nd": 0.0, "Dy": 0.0, "H2O": 0.0, "Organic": 50.0},
T=298.15, P=101325.0)
raffinate, extract, info = cascade(feed_stream, solvent_stream)
```
### Utilities
- **Mixer**: Combine multiple streams
- **Splitter**: Split stream by fraction
### Fed-Batch Reactor
- General-purpose fed-batch (semi-batch) reactor for chemical reactions
- Time-varying feed addition with configurable feed profiles
- RK4 integration for batch dynamics
- Supports multiple reactions with user-defined kinetics
```python
from difflow import FedBatchReactor, FedBatchParams
# Fed-batch reactor with continuous reagent addition
def rate_fn(C, T, params):
k = params["k0"] * jnp.exp(-params["Ea"] / (8.314 * T))
return jnp.array([k * C["A"] * C["B"]])
params = FedBatchParams(
V0=jnp.array(1.0), # Initial volume (m³)
rate_fn=rate_fn,
stoich=jnp.array([[-1.0], [-1.0], [1.0]]), # A + B → C
rate_params={"k0": jnp.array(1e6), "Ea": jnp.array(50000.0)},
species_order=["A", "B", "C"],
)
reactor = FedBatchReactor(params)
# Initial charge (mol/m³) and a constant feed of B (mol/m³) at 0.001 m³/s
C0 = {"A": 1000.0, "B": 0.0, "C": 0.0}
def feed_rate(t): return jnp.array(0.001) # m³/s
final, info = reactor(C0, T0=350.0, P=101325.0, t_final=3600.0,
feed_rate_fn=feed_rate,
feed_composition={"A": 0.0, "B": 1000.0, "C": 0.0}, feed_T=350.0)
# info contains the time profiles: t, V, C, T
```
### Distillation Columns
- **ShortcutColumn**: Fenske-Underwood-Gilliland method for quick design estimates
- Minimum stages (Fenske equation)
- Minimum reflux ratio (Underwood equations)
- Actual stages for given reflux (Gilliland correlation)
- **DistillationColumn**: Rigorous stage-by-stage calculation
- MESH equations (Material, Equilibrium, Summation, Heat balance)
- Supports partial/total condenser and reboiler
```python
from difflow import ShortcutColumn, ShortcutColumnParams, IdealThermo
from difflow.database import get_species_data
names = ["benzene", "toluene", "ethylbenzene"]
aromatics = IdealThermo({s: get_species_data(s) for s in names})
params = ShortcutColumnParams(
species_order=names,
light_key="benzene",
heavy_key="toluene",
x_D_LK=0.99, # 99% benzene recovery in distillate
x_B_HK=0.99, # 99% toluene recovery in bottoms
)
column = ShortcutColumn(params, thermo=aromatics)
feed = make_stream({"benzene": 40.0, "toluene": 35.0, "ethylbenzene": 25.0},
T=370.0, P=101325.0)
distillate, bottoms, info = column(feed, R=3.0, q=1.0)
# info contains: N_min, R_min, N, N_feed, Q_condenser, Q_reboiler, ...
```
### Heat Exchangers
- **Heater/Cooler**: Single-stream with utility (steam, cooling water)
- Specified duty mode
- Specified outlet temperature mode
- Rating mode (given UA and utility temperature)
- Constant `Cp`, or a `thermo` for a real enthalpy balance (carries latent heat)
- **CounterCurrentHX**: Two-stream counter-current (shell-and-tube style)
- **CoCurrentHX**: Two-stream co-current (parallel flow)
- **EnthalpyCounterCurrentHX**: Two-stream, closed on EOS enthalpies through phase change
- The constant-Cp units use the effectiveness-NTU method; all are fully differentiable
```python
from difflow import (
Heater, HeaterParams,
CounterCurrentHX, HeatExchangerParams,
design_heat_exchanger,
)
# Streams of the Quick Start's species A and B (mol/s)
cold_feed = make_stream({"A": 10.0, "B": 5.0}, T=300.0, P=101325.0)
hot_stream = make_stream({"A": 10.0, "B": 5.0}, T=450.0, P=101325.0)
cold_stream = make_stream({"A": 10.0, "B": 5.0}, T=300.0, P=101325.0)
# Single-stream heater with steam
heater = Heater(HeaterParams(T_out=400.0, Cp=75.0))
heated_feed, info = heater(cold_feed)
# info: Q, T_in, T_out, LMTD (if utility temp specified)
# Duty from the thermo instead of a constant Cp -- required if the stream
# vaporizes, since a constant Cp carries no latent heat
heater = Heater(HeaterParams(T_out=400.0), thermo=thermo)
heated_feed, info = heater(cold_feed)
# Two-stream counter-current heat exchanger
hx = CounterCurrentHX(HeatExchangerParams(
UA=2000.0, # W/K
Cp_hot=75.0, # J/(mol·K)
Cp_cold=80.0,
))
hot_out, cold_out, info = hx(hot_stream, cold_stream)
# info: Q, effectiveness, NTU, LMTD, approach temperature
# Design: calculate required area
result = design_heat_exchanger(
Q=jnp.array(100000.0), # 100 kW
T_hot_in=jnp.array(450.0), T_hot_out=jnp.array(380.0),
T_cold_in=jnp.array(300.0), T_cold_out=jnp.array(360.0),
U=jnp.array(500.0), # W/(m²·K)
)
print(f"Required area: {result['A']:.1f} m²")
```
## Bio Manufacturing Operations
The `difflow_bio` plugin provides specialized unit operations for biopharmaceutical manufacturing:
### Bioreactors
- **ContinuousBioreactor**: Chemostat with Monod kinetics
- **FedBatchBioreactor**: Fed-batch with substrate feeding strategy
```python
from difflow_bio import (
ContinuousBioreactor, BioreactorParams,
FedBatchBioreactor, FedBatchParams,
monod_kinetics,
)
# Create a continuous bioreactor (chemostat)
params = BioreactorParams(
V=1000.0, # Volume (L)
Y_xs=0.5, # Biomass yield (g/g)
kinetic_fn=monod_kinetics,
kinetic_params={"mu_max": 0.3, "K_s": 0.5}, # 1/h, g/L
alpha=0.1, # Growth-associated product yield
)
bioreactor = ContinuousBioreactor(params)
feed_stream = make_stream({"cells": 0.0, "substrate": 20.0, "product": 0.0},
T=310.0, P=101325.0)
outlet, info = bioreactor(feed_stream, D=0.1) # D: dilution rate (1/h)
```
### Downstream Processing
- **DiscStackCentrifuge**: Cell removal with Stokes' law separation
- **Ultrafiltration**: Protein concentration via TFF
- **Diafiltration**: Buffer exchange
- **ProteinAChromatography**: Affinity capture for mAb purification
- **IonExchangeChromatography**: Polishing step (bind-elute or flow-through)
- **SizeExclusionChromatography**: Aggregate removal
```python
from difflow_bio import (
DiscStackCentrifuge, DiscStackParams,
Ultrafiltration, UltrafiltrationParams,
ProteinAChromatography, ProteinAParams,
)
# Disc-stack centrifuge for cell removal (Sigma from the disc geometry)
centrifuge = DiscStackCentrifuge(DiscStackParams(
n_discs=100, r_outer=0.15, r_inner=0.05, # m
rpm=7000.0,
))
# Protein A capture
proa = ProteinAChromatography(ProteinAParams(
column_volume=10.0, # CV (L)
q_max=40.0, # g mAb / L resin
yield_factor=0.95,
))
# Ultrafiltration for concentration (the concentration factor is a call argument)
uf = Ultrafiltration(UltrafiltrationParams(
membrane_area=1.0, # m²
))
```
## Gas Transmission Networks
The `difflow_gas` plugin models steady-state gas transmission networks
as sequential-modular differentiable flowsheets. The sequential
decomposition of a meshed network (spanning tree, tear set, balance
schedule) is computed from the topology, so multi-loop networks need
no hand derivation:
```python
import difflow_gas as dg
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 a loop: the tear
},
beta={aid: dg.weymouth_beta(L, 0.6, 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},
)
fs, dec = dg.build_network_flowsheet(net, root="src",
p_slack_pa=60e5,
ratios={"cs1": 1.3})
streams = fs.solve(tol=1e-8) # signed flows, Anderson tears
assert dg.residual_report(streams, net, dec).ok
# exact gradients through the converged tear iteration
obj = fs.make_objective_fn(
lambda s: dg.total_compressor_power_w(s, dec, net.gas_temp_k))
dW_dr = jax.grad(obj)({"cs_cs1.ratio": 1.3})
```
Pipes, resistors, compressor stations, open valves, control valves and
short pipes are supported; see `docs/unit-operations-gas.md`.
## Electrical Grids and AC-OPF
The `difflow_power` plugin models steady-state electrical transmission
and distribution networks, and solves the AC optimal power flow — the
nonconvex problem every wholesale market and control centre sits on
top of. Because the model is differentiable, the quantities a grid
study is actually after are derivatives rather than separately-derived
sensitivity factors: locational marginal prices, shift factors,
marginal loss factors, and the value of relaxing any binding limit.
```python
import difflow_power as dp
net = dp.cases.case9() # WSCC 9-bus benchmark
pf = dp.solve_power_flow(net) # Newton-Raphson, implicit-diff gradients
pf.losses_mw # 4.9547 (MATPOWER: 4.9547)
opf = dp.solve_acopf(net) # interior-point AC-OPF, written in JAX
opf.cost # 5296.69 $/h (MATPOWER: 5296.69)
opf.lmp_mw # locational marginal prices, $/MWh
opf.binding() # binding limits and their shadow prices
# the multipliers ARE prices: check them against jax.grad of the optimum
max(opf.check_prices().values()) # ~1e-12 $/MWh
# and everything the classical factor tables give, as derivatives
dp.loss_sensitivity(net) # marginal loss factors
dp.ptdf(net), dp.lodf(net) # shift and outage factors
```
One branch model covers lines, transformers and phase shifters;
generator boxes, voltage limits and thermal ratings are carried as
inequalities by a primal-dual interior-point solver written in JAX
(no IPOPT, so the differentiability survives). Radial feeders also
solve sequentially, by the backward/forward sweep, which agrees with
Newton to 1e-12. Every benchmark result is asserted against MATPOWER's
published answer; see `docs/unit-operations-power.md`.
## Refinery: Crude Distillation
`difflow_refinery.CrudeDistillationUnit` is a crude unit: a TBP assay
characterized into pseudo-components, a fired heater solved together with
an atmospheric column (side strippers, pumparounds, stripping steam), and
products reported as yields, API gravities and TBP ranges. Specs follow a
simulator's degrees of freedom (product rates, pumparound duties,
overflash), and every yield or duty has an implicit-function gradient with
respect to the specs, the feed and the assay.
<!-- doc-test: skip: needs an assay and a spec set; the runnable version is in docs/unit-operations-refinery.md -->
```python
import difflow_refinery as dr
unit = dr.CrudeUnit(assay, column_params) # dr.Assay, dr.column.CrudeColumnParams
res = unit.solve(95_000, T=513.15, P=6e5) # bbl/d at the furnace inlet
print(res.table()) # yields, API, TBP 5/50/95
```
## Refinery: Vacuum Distillation
`difflow_refinery.vacuum` characterizes a crude assay into
pseudocomponents and runs a vacuum distillation unit on the atmospheric
residue. It covers LVGO and HVGO pumparound sections, a wash zone with an
overflash spec, a flash zone fed by the furnace, and a steam-stripped
residue. The MESH equations of every stage are solved simultaneously, and
the result carries exact implicit-function gradients, so the VGO/residue
cut point is a decision variable with a derivative. That includes the
cut's derivative with respect to a single point of the assay's TBP curve.
```python
import difflow_refinery as dr
char = dr.vacuum.characterize(dr.vacuum.heavy_crude()) # 300-800 C cuts + residue lump
feed = dr.vacuum.atmospheric_residue(char, crude_rate_kg_s=100.0)
vdu = dr.VacuumColumn(dr.VacuumColumnParams(components=char.components))
overhead, lvgo, hvgo, slop, residue, info = vdu(feed)
info["properties"]["hvgo"] # rate, SG, S, N, CCR, Ni+V, TBP 5/50/95
```
Any output can be specified in place of the knob that controls it, for
example an HVGO end point instead of the furnace temperature.
The crude unit, the vacuum column and the blend pool can share one
characterization. An `Assay` with a `HeavyEnd` is carried into the vacuum
range, closed by a residue lump, and given sulfur, nitrogen, CCR and metals
per component. The CDU runs on it, and the VDU's components are
`char.pseudo_components()`, so the CDU's `"residue"` outlet feeds the VDU
directly in a `Flowsheet`. The balance closes, and gradients cross the
connection:
<!-- doc-test: skip: needs a HeavyEnd assay and CDU column params; see docs/unit-operations-refinery.md -->
```python
char = dr.characterize(assay) # assay has heavy_end=dr.HeavyEnd()
cdu = dr.CrudeDistillationUnit(dr.CrudeDistillationUnitParams(assay=assay, column=params))
vdu = dr.VacuumColumn(dr.VacuumColumnParams(components=char.pseudo_components()))
grid = dr.BlendCharacterization.from_characterization(char) # products into the pool
```
See `docs/unit-operations-refinery.md` and
`examples/36_crude_to_vacuum.ipynb`.
## Refinery Product Blending
`difflow_refinery.BlendPool` blends component streams into finished
products (gasoline, jet, ULSD, fuel oil) with the nonlinear rules
refiners use: Ethyl RT-70 octane interactions, the RVP^1.25 index (and
Raoult on the pseudocomponents as a check), Hu-Burns flash and cold-flow
indices, Refutas viscosity, and distillation and cetane index computed
from the blend's composition. It is differentiable in the recipe and in
every component property, and it reports signed spec margins with a
smooth-violation option.
```python
from difflow_refinery import BlendComponent, BlendPool
reformate = BlendComponent.from_properties(
"reformate", SG=0.80, RON=98.0, MON=88.0, RVP_psi=3.5, S_ppm=1.0,
olefins_vol=1.0, aromatics_vol=65.0)
alkylate = BlendComponent.from_properties(
"alkylate", SG=0.70, RON=95.0, MON=93.0, RVP_psi=4.5, S_ppm=5.0,
olefins_vol=0.5, aromatics_vol=0.5)
butane = BlendComponent.from_properties(
"butane", SG=0.58, RON=93.0, MON=90.0, RVP_psi=52.0, S_ppm=10.0,
olefins_vol=0.5, aromatics_vol=0.0)
components, recipe = [reformate, alkylate, butane], [0.45, 0.50, 0.05] # volume fractions
pool = BlendPool("gasoline") # RON, MON, RVP, S specs
res = pool(components, recipe) # properties, margins, product stream
pool.linear_blend_error(components, recipe) # what an LP's back-off must cover
pool.as_block(components) # a difflow.planning.Block
```
See `docs/unit-operations-refinery.md` and
`examples/33_refinery_gasoline_blending.ipynb`, which compares the
nonlinear optimum with a linear-by-volume LP plus successive back-off.
Every refinery unit (preheat train, crude and vacuum units, gas plant,
isomerization, hydrotreater, hydrocracker, FCC, reformer, alkylation and
blending) is listed, with its model and how far it has been validated, in
`docs/refinery-summary.md` and `src/difflow_refinery/README.md`.
`examples/40_refinery_flowsheet.ipynb` joins them into a small whole
refinery, from the crude to the product pools and the hydrogen header, and
`difflow_refinery.plant` (`Chain`, `Stage`, `AD_MODES`) composes library units
into one differentiable function across their different AD modes (the
reformer is forward-only, a default hydrotreater reverse-only).
## Data Reconciliation
Plant measurements are noisy and, taken at face value, contradict the
model. `difflow.reconciliation` finds the smallest statistically
weighted adjustment that satisfies the model equations, and returns
estimates sharper than the raw measurements:
```python
from difflow.reconciliation import reconcile, global_test, measurement_test
# a splitter with three metered flows that do not close: F1 = F2 + F3
names = ["F1", "F2", "F3"]
residual_fn = lambda x, params: jnp.array([x[0] - x[1] - x[2]])
y = jnp.array([100.0, 60.0, 38.0]) # measured, kg/s
sigma = jnp.array([1.0, 1.0, 1.0])
res = reconcile(residual_fn, y, sigma, names=names)
print(res.summary()) # measured vs reconciled, with standard errors
global_test(res) # is the data set consistent with the model?
measurement_test(res) # which sensor is lying?
```
An entry of `sigma` set to `inf` marks a variable to be *estimated*
rather than reconciled, so joint parameter estimation and reconciliation
are the same solve. Whether an unknown can be recovered at all is decided
before solving, so an ill-posed problem raises a named error instead of
returning `NaN`. Works with any differentiable residual function; see
`docs/data-reconciliation.md` and
`examples/28_data_reconciliation.ipynb` for the gas-network case, and
`examples/29_model_updating.ipynb` for when to update a model parameter
rather than the data.
## Delta-Base Planning
Refinery and value-chain planning runs on linear programs whose unit
submodels are *base plus delta vectors*, `y ~= y0 + J (u - u0)`. Every
commercial system builds `J` by perturbing a rigorous simulator one
variable at a time, which costs `O(n)` evaluations. A flowsheet is a pure
function with its flash and recycle solves embedded, so `jax.jacobian`
returns the same reduced Jacobian for a cost independent of `n`:
<!-- doc-test: skip: fragment: `ngl` and `power` are Blocks built in docs/planning.md -->
```python
from difflow.planning import Block, Network, DeltaBasePlanner
net = Network([ngl, power], links=[("ngl.residue_F", "power.fuel_F")])
res = DeltaBasePlanner(net, prices={"ngl.NGL_C2": 9.0, "power.Power": 55.0},
specs=[("ngl.T_colfeed", "<=", 236.0)],
radius=0.3).solve()
res.plan # optimal decisions
res.delta_vectors # the J blocks actually used
res.pyomo_model # emitted for the Pyomo/IDAES ecosystem
res.plan_sensitivity(wrt="prices") # d(plan)/d(price), not just the plan
```
Measured on a two-plant chain, the AD gradient costs 1-2 model
evaluations from 5 to 80 decisions while central differences cost `2n`.
Every LP proposal is checked against the *nonlinear* blocks before it is
accepted, violations are charged from the real model rather than from LP
slacks, bang-bang levers are vertex-seeded, and a linearisation that
straddles a phase boundary raises a warning instead of quietly
extrapolating a branch that no longer exists. Pooling/blending
bilinearity and the commercial trappings (assay libraries, blending
correlations, scheduling) are explicitly out of scope. See
`docs/planning.md` and `examples/30_delta_base_planning.ipynb`.
## Thermodynamics
### Ideal Thermodynamics (for VLE)
- Ideal gas behavior
- Antoine equation for vapor pressures
- Polynomial Cp correlations
- Watson correlation for heat of vaporization
<!-- doc-test: skip: signature template with placeholder coefficients -->
```python
SpeciesData(
name="species_name",
MW=100.0, # Molecular weight (g/mol)
Cp_coeffs=(a, b, c, d), # Cp = a + bT + cT² + dT³
Hvap_coeffs=(A, n, Tc), # Hvap = A(1 - T/Tc)^n
antoine_coeffs=(A, B, C), # log10(Psat) = A - B/(T+C)
Hf=0.0, # Heat of formation (J/mol)
)
```
### Equations of State (for non-ideal VLE)
- **Peng-Robinson**: Cubic EOS for hydrocarbon and gas systems
- **Soave-Redlich-Kwong (SRK)**: Alternative cubic EOS
- Fugacity coefficients for both vapor and liquid phases
- Flash calculations with non-ideal K-values
- Binary interaction parameters (kij) support
```python
from difflow import PengRobinson, SRK, CriticalProperties, flash_TP_eos
# Define critical properties
props = {
"methane": CriticalProperties("methane", Tc=190.6, Pc=4.6e6, omega=0.011),
"ethane": CriticalProperties("ethane", Tc=305.3, Pc=4.87e6, omega=0.099),
}
# Create EOS
eos = PengRobinson(props)
# or: eos = SRK(props)
# Compressibility factor
z = eos.solve_Z(T=300.0, P=1e6, y=jnp.array([0.7, 0.3]), phase="vapor")
# Fugacity coefficients
phi = eos.fugacity_coefficient(T=300.0, P=1e6, y=jnp.array([0.7, 0.3]), phase="vapor")
# Flash calculation
V_frac, x, y = flash_TP_eos(eos, z=jnp.array([0.5, 0.5]), T=250.0, P=2e6)
```
### Property Database
Built-in database with 55+ common species including critical properties and ideal thermo data:
```python
from difflow import (
get_species_data, get_critical_props, list_species,
get_alkanes, get_btex, get_common_solvents,
)
# Get species data for ideal thermodynamics
methanol = get_species_data("methanol")
thermo = IdealThermo({"methanol": methanol, "water": get_species_data("water")})
# Get critical properties for EOS
methane = get_critical_props("methane")
eos = PengRobinson({"methane": methane, "ethane": get_critical_props("ethane")})
# Convenience functions
alkanes = get_alkanes() # methane through n-decane
btex = get_btex() # benzene, toluene, ethylbenzene, xylenes
solvents = get_common_solvents() # water, methanol, ethanol, acetone, etc.
# Alias support: "CO2", "MeOH", "isopropanol", "IPA" all work
co2 = get_critical_props("CO2")
# List all available species
print(list_species())
```
### Activity Coefficient Models (for LLE)
- **NRTL**: Non-Random Two-Liquid model with temperature-dependent parameters
- **UNIQUAC**: Universal Quasi-Chemical model
## Technoeconomic Analysis (TEA)
The `difflow.economics` module provides comprehensive technoeconomic analysis capabilities, all fully differentiable for gradient-based optimization.
### Capital Costs
Equipment cost correlations with CEPCI escalation and installation factors:
```python
import difflow.economics as econ
import jax.numpy as jnp
# Equipment costs (2024 dollars)
reactor_cost = econ.reactor_cost(jnp.array(5.0), "cstr_jacketed") # 5 m³
hx_cost = econ.heat_exchanger_cost(jnp.array(100.0), "shell_tube_floating") # 100 m²
pump_cost = econ.pump_cost(jnp.array(10.0), "centrifugal_single") # 10 kW
# Installed cost with Lang factor
installed = econ.installed_cost(reactor_cost, lang_factor=4.74)
# Total capital investment
tci = econ.total_capital_investment(
purchased_equipment_cost=reactor_cost + hx_cost + pump_cost,
lang_factor=4.74,
working_capital_fraction=0.15,
)
```
Available equipment types:
- **Reactors**: CSTR (jacketed, coil), PFR, batch
- **Vessels**: Pressure vessels, storage tanks, flash drums
- **Heat Exchangers**: Shell-tube, double-pipe, plate-frame, air coolers
- **Columns**: Tray columns, packed columns
- **Pumps**: Centrifugal, reciprocating, gear
- **Compressors**: Centrifugal, reciprocating, screw
- **Separators**: Mixer-settlers, centrifuges, filters, extraction columns
### Utility Costs
```python
# Steam cost from heat duty
heating_cost = econ.steam_cost_from_duty(jnp.array(1e6), "medium_pressure") # 1 MW
# Cooling water
cooling_cost = econ.cooling_water_cost(jnp.array(500e3)) # 500 kW
# Electricity
electricity_cost = econ.electricity_cost(jnp.array(100.0)) # 100 kW → $/hour
```
### Profitability Metrics
All metrics are JAX-differentiable:
```python
# Net Present Value
cash_flows = jnp.ones(20) * 500000 # $500k/year for 20 years
npv = econ.npv(cash_flows, jnp.array(0.10), jnp.array(2e6)) # 10% discount, $2M investment
# Internal Rate of Return
irr = econ.irr(cash_flows, jnp.array(2e6))
# Minimum Selling Price
msp = econ.minimum_selling_price(
total_annual_cost=jnp.array(1e6),
annual_production=jnp.array(50000.0), # kg/year
)
# Annualized cost for optimization
tac = econ.annualized_cost(
capital_cost=jnp.array(5e6),
annual_opex=jnp.array(1e6),
discount_rate=jnp.array(0.10),
plant_life=jnp.array(20.0),
)
```
### Gradient-Based Economic Optimization
```python
import jax
def annual_profit(params):
V, T = params[0], params[1]
# Simulate process
outlet, info = simulate_reactor(V, T)
# Economics
capex = econ.reactor_cost(V, "cstr_jacketed")
installed = econ.installed_cost(capex)
utility_cost = econ.cooling_water_cost(jnp.abs(info["Q"]))
annual_utility = utility_cost * 8000 * 3600 # $/year
revenue = outlet["F_product"] * product_price * 8000 * 3600
crf = econ.capital_recovery_factor(jnp.array(0.10), jnp.array(20.0))
return revenue - annual_utility - installed * crf
# Optimize design for maximum profit
grad_profit = jax.grad(annual_profit)
# Use gradient for optimization...
```
## Uncertainty Propagation
Leverage JAX's automatic differentiation for uncertainty quantification:
```python
from difflow import linear_propagation, monte_carlo_propagation, sensitivity_analysis
# Define a process model
def reactor_model(params):
k = params['k0'] * jnp.exp(-params['Ea'] / (8.314 * params['T']))
conversion = 1 - jnp.exp(-k * params['tau'])
return conversion
nominal = {'k0': jnp.array(1e6), 'Ea': jnp.array(50000.0),
'T': jnp.array(350.0), 'tau': jnp.array(100.0)}
uncertainties = {'k0': 1e5, 'Ea': 2000.0, 'T': 5.0, 'tau': 10.0}
# Linear (Jacobian-based) propagation - fast, first-order approximation
mean, std, info = linear_propagation(reactor_model, nominal, uncertainties)
print(f"Conversion: {mean:.3f} ± {std:.3f}")
print(f"Variance by input: {info['variance']}")
# Monte Carlo propagation - handles non-linear models
mean_mc, std_mc, info_mc = monte_carlo_propagation(
reactor_model, nominal, uncertainties, n_samples=10000
)
# Sensitivity analysis with gradient information
sens = sensitivity_analysis(reactor_model, nominal)
# Returns: gradient, elasticity (normalized sensitivity), variance contribution
```
Available functions:
- `linear_propagation()`: First-order Jacobian-based uncertainty propagation
- `monte_carlo_propagation()`: Parallel sampling using JAX vmap
- `sensitivity_analysis()`: Local gradient-based sensitivity with variance contributions
- `sobol_indices()`: Global sensitivity via Sobol sampling
- `propagate_covariance()`: Full covariance matrix propagation for correlated inputs
## Flowsheets with Recycles
```python
from difflow import (Flowsheet, Unit, make_stream, IdealThermo, CSTR, CSTRParams,
Mixer, Splitter, mass_action_kinetics)
from difflow.database import get_species_data
# n-hexane -> 2-methylpentane in a CSTR, 80 % of the effluent recycled
species = ["n_hexane", "2_methylpentane"]
thermo = IdealThermo({s: get_species_data(s) for s in species})
kin = mass_action_kinetics(
[{"reactants": {species[0]: 1.0}, "products": {species[1]: 1.0},
"rate_params": {"A": 0.05}}], species_order=species)
reactor = CSTR(CSTRParams(V=1.0, molar_density=10.0, **kin.params_kwargs()), thermo)
fs = Flowsheet(species_order=species)
fs.add_feed("fresh", make_stream({species[0]: 1.0, species[1]: 0.0}, T=340.0, P=101325.0))
fs.add_unit(Unit("mixer", Mixer(species, thermo), ["fresh", "recycle"], ["reactor_feed"]))
fs.add_unit(Unit("reactor", reactor, ["reactor_feed"], ["effluent"], params={"T_spec": 340.0}))
fs.add_unit(Unit("split", Splitter(species), ["effluent"], ["recycle", "product"],
params={"split_frac": 0.8}))
fs.add_recycle("recycle", "recycle") # tear the splitter outlet
streams = fs.solve(tol=1e-8)
```
### Debugging a recycle solve
When a loop will not close, the flowsheet debugger lets you watch it fail: it
pauses the solve unit by unit and tear iteration by tear iteration, shows every
stream on the way, and stops on breakpoints --- before a unit, at an iteration,
on a condition such as `residual>1e3`, or where the first NaN is born.
```bash
difflow debug plant.py # a difflow-dbg> prompt, before the first pass
difflow debug plant.py --on-error # run; stop only at a failing unit or an unconverged end
difflow debug plant.py --json # the same debugger, as JSON lines for a program
```
<!-- doc-test: skip: interactive prompt -->
```python
from difflow import debug
with debug.attach(): # the same prompt, from a script or notebook
fs.solve(acceleration="wegstein")
```
```text
difflow-dbg> break on nonfinite
difflow-dbg> break if residual>1e3 && iter>5
difflow-dbg> continue
difflow-dbg> p tear # which stream and field the residual is made of
difflow-dbg> set tear recycle.F_n_hexane 0.5 # a different next guess
```
A solve can also be recorded frame by frame and replayed with one tear value
changed, which asks whether a different starting point reaches the same answer
(if not, the loop has more than one steady state):
```python
from difflow.debug.record import compare, summarize
trace = debug.record(fs)
k = summarize(trace)["iterations"][0]["frame"]
replay = debug.record(fs, edits=[{"frame": k, "stream": "recycle",
"field": "F_n_hexane", "value": 0.5}],
baseline=trace["frames"])
compare(trace, replay)["same_answer"]
```
Agents get the same debugger through `difflow mcp` (`debug_start`,
`debug_command`, `debug_record`, `debug_whatif`, ...). See
[docs/debugger.md](docs/debugger.md) and the example notebook
`examples/69_flowsheet_debugger.ipynb`.
## Building Flowsheets Without Code
A flowsheet is a graph with numbers on it. Writing that as Python is the flexible
route, not the only one: difflow can also describe a model as data, write it back
out as a script, and sweep it over a grid of parameters with exact gradients.
### The operation catalog
`difflow.catalog` answers *what you can do with a unit*: how many streams go in
and out, what parameters it takes, which are required, and which hold code rather
than data.
```python
from difflow import catalog, describe_operation
spec = describe_operation("Heater")
spec.ports.inlets # ['inlet']
spec.ports.n_outlets # 1
spec.required_parameters() # [] -- every field has a default
spec.equations # LaTeX governing equations
spec.is_buildable # True: constructible from data alone
spec.to_dict() # JSON-serializable, for a UI or code generator
describe_operation("Flash").is_buildable # False
describe_operation("Flash").constructor_extras # ['thermo'] -- an object, not data
```
All of it is derived by introspection (parameters from `dataclasses.fields`,
ports from the `__call__` signature), so it cannot drift from the code, and
plugin units appear with no extra work.
### JSON round trip and code generation
```python
from difflow import serialize, codegen
serialize.save(fs, "plant.json") # flowsheet -> JSON
fs = serialize.load("plant.json") # JSON -> flowsheet
print(codegen.to_python(fs)) # flowsheet -> runnable script
```
Rate laws are the usual obstacle to writing a reactor down as data, because a
callable is code. `mass_action_kinetics` builds one from plain dictionaries
instead, so a reaction network can be stored and edited as data.
### Parameter sweeps
`sweep` evaluates a flowsheet on a grid of parameter values with `jax.vmap`,
with the exact derivative of every output with respect to every axis:
```python
from difflow import sweep, SweepAxis
result = sweep(
fs,
axes=[SweepAxis("reactor.V", 0.5, 5.0, n=21, label="Reactor volume", units="m³")],
outputs={"isomer_out": lambda streams: streams["product"]["F_2_methylpentane"]},
)
```
### The browser editor
The browser editor (a canvas, an inspector, results and sensitivities, and
the debugger's Debug and Trace panels) is an app that comes with the difflow
book. With it installed, `difflow gui plant.json` opens it. Its engine,
`difflow.workspace.FlowsheetSession`, is part of difflow, and is what the agent
tools (`difflow mcp`) drive.
See [`docs/streams-and-flowsheets.md`](docs/streams-and-flowsheets.md) for the
full documentation.
## Examples
Jupyter notebooks are in the `examples/` directory:
| Notebook | Description |
|----------|-------------|
| `00_cstr_pfr_basics.ipynb` | CSTR and PFR basics: conventional vs difflow |
| `01_cstr_flash_recycle.ipynb` | Complete flowsheet with CSTR, flash, and recycle |
| `02_cstr_sensitivity.ipynb` | Sensitivity analysis for CSTR parameters |
| `03_optimization.ipynb` | Gradient-based optimization problems |
| `04_rare_earth_extraction.ipynb` | Rare earth recovery using LLE |
| `05_technoeconomic_analysis.ipynb` | Comprehensive TEA with profit optimization |
| `06_uncertainty_propagation.ipynb` | Uncertainty propagation and sensitivity analysis |
| `07_heat_exchangers.ipynb` | Heat exchanger design, rating, and optimization |
| `47_evaporators.ipynb` | Single/multiple-effect evaporators, MVR, optimal number of effects |
| `64_serialization.ipynb` | Save/load flowsheets as JSON, `refs=`/`extras=`, export as Python |
| `69_flowsheet_debugger.ipynb` | Step a recycle solve, breakpoints, find where a NaN is born, `set tear`, record and what-if replays |
| `10_dynamic_modeling.ipynb` | Dynamic simulation, DAE systems, diffrax backend |
```bash
# Launch Jupyter to explore examples
jupyter notebook examples/
```
## Tutorials
The `tutorials/` directory contains comprehensive JAX tutorials for differentiable programming:
| Notebook | Topics |
|----------|--------|
| `01_jax_fundamentals.ipynb` | grad, jit, vmap, pytrees, jacfwd/jacrev, VJP/JVP, HVP |
| `02_inverse_hessian_vector_products.ipynb` | IHVP, conjugate gradient, Newton-CG optimization |
| `02_optimization.ipynb` | Gradient descent, Newton, Adam, constrained optimization |
| `03_differential_equations.ipynb` | ODE solvers, parameter estimation, neural ODEs |
| `04_custom_derivatives.ipynb` | custom_vjp, custom_jvp, stop_gradient |
| `05_machine_learning.ipynb` | Neural networks from scratch, training loops |
| `06_gotchas.ipynb` | Common JAX pitfalls and how to avoid them |
## Key Design Decisions
1. **Streams as Dicts**: Simple `{"F_A": ..., "F_B": ..., "T": ..., "P": ...}` format that's a JAX pytree by default
2. **Property Database Available**: Built-in database with 55+ species, or define custom species data
3. **Function-Based Kinetics**: Maximum flexibility via `rate_fn(C, T, params) → rates`
4. **Unrolled Iteration**: Fixed-point solvers use `lax.scan` for automatic differentiability
5. **Continuous Relaxation**: Discrete parameters (like n_stages) can be relaxed to continuous values for optimization
## Dynamic Modeling
The `difflow.dynamic` module provides a unified framework for transient simulation of process units:
### Basic ODE Integration
```python
from difflow.dynamic import integrate
import jax.numpy as jnp
# Define any ODE system
def harmonic_oscillator(t, y):
x, v = y[0], y[1]
return jnp.array([v, -x]) # dx/dt = v, dv/dt = -x
result = integrate(
harmonic_oscillator,
y0=jnp.array([1.0, 0.0]),
t_span=(0.0, 10.0),
method="RK4", # or "RK45", "Euler"
)
print(f"Final state: {result.y_final}")
print(f"Trajectory shape: {result.trajectory.y.shape}")
```
### Dynamic Unit Operations
```python
from difflow.dynamic import DynamicCSTR, integrate_unit
from difflow.streams import make_stream
# Define reaction kinetics
def rate_fn(C, T, params):
k = params["k"] * jnp.exp(-params["Ea"] / (8.314 * T))
return jnp.array([k * C["A"]])
# Create dynamic CSTR
cstr = DynamicCSTR(
volume=1.0,
rate_fn=rate_fn,
stoich=jnp.array([[-1], [1]]), # A -> B
species_order=["A", "B"],
rate_params={"k": 1e6, "Ea": 50000.0},
)
# Simulate startup from empty
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",
)
```
### Dynamic Flowsheets
Connect multiple dynamic units for multi-unit transient simulation:
```python
from difflow.dynamic import DynamicFlowsheet, DynamicCSTR, DynamicTank
def rate_fn(C, T, params): # A -> B, first order
k = params["k0"] * jnp.exp(-params["Ea"] / (8.314 * T))
return jnp.array([k * C["A"]])
stoich = jnp.array([[-1.0], [1.0]])
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")
inlet_stream = make_stream({"A": 1.0, "B": 0.0}, T=350.0, P=101325.0)
# Build flowsheet
fs = DynamicFlowsheet(species_order=["A", "B"])
fs.add_feed("feed", inlet_stream)
fs.add_unit(cstr, inlet_names=["feed"], outlet_names=["reactor_out"])
fs.add_unit(tank, inlet_names=["reactor_out"], outlet_names=["product"])
# Simulate entire flowsheet
result = fs.simulate(t_span=(0.0, 1000.0), method="RK4", n_steps=500)
```
### DAE (Differential-Algebraic Equations)
For systems with algebraic constraints (e.g., VLE equilibrium):
```python
from difflow.dynamic import DynamicFlashDrum, integrate_dae
# Flash drum with VLE equilibrium constraint
flash = DynamicFlashDrum(
volume=1.0,
species_order=["A", "B"],
K_func=lambda T: jnp.array([2.0, 0.5]), # K = y/x for (A, B)
)
feed = make_stream({"A": 0.5, "B": 0.5}, T=350.0, P=101325.0)
result = integrate_dae(
flash,
inputs={"inlet": feed},
t_span=(0.0, 100.0),
method="RK4",
)
# result.x_final: differential states (moles)
# result.z_final: algebraic states (vapor fraction)
```
### Diffrax Backend (Advanced Solvers)
For stiff systems or when adaptive step control is needed:
```bash
pip install diffrax # Optional dependency
```
```python
from difflow.dynamic import integrate
# A stiff linear system, dy/dt = -1000 (y - cos t)
stiff_ode = lambda t, y: -1000.0 * (y - jnp.cos(t))
y0 = jnp.array([0.0])
t_span = (0.0, 1.0)
# Use diffrax solvers via method string
result = integrate(
stiff_ode, y0, t_span,
method="diffrax:kvaerno5", # Implicit solver for stiff systems
rtol=1e-6, atol=1e-8,
)
# Available solvers: dopri5, tsit5, dopri8, kvaerno3/4/5, euler, heun
# Default: tsit5 (recommended for most problems)
```
See `docs/dynamic-modeling.md` for complete documentation.
## Limitations
- Rigorous distillation column convergence can be sensitive to initial guesses
- Gradient explosion possible with many iterations (use damping)
- EOS flash limited to two-phase VLE (no three-phase VLLE yet)
## Future Work
- Three-phase (VLLE) flash calculations
- Extended bio operations (viral inactivation, sterile filtration)
- GPU acceleration for large flowsheets
- Integration with experiment databases (e.g., Cantera)
## Citation
If you use difflow in your work, please cite it. Machine-readable metadata is in
[`CITATION.cff`](https://github.com/jkitchin/differentiable-flowsheets/blob/main/CITATION.cff),
and GitHub's "Cite this repository" button will render BibTeX or APA from it.
## License
MIT