Difflow Documentation#
Difflow is a JAX-based framework for building and optimizing chemical process flowsheets with automatic differentiation support. This enables gradient-based optimization of process designs, sensitivity analysis, and uncertainty propagation.
Key Features#
Automatic Differentiation: All unit operations and calculations are fully differentiable using JAX
Comprehensive Unit Operations: Reactors, separators, heat exchangers, distillation columns, and more
Bio-Manufacturing Support: Specialized operations for bioreactors, centrifugation, filtration, and chromatography
Thermodynamic Models: Ideal thermodynamics and cubic equations of state (Peng-Robinson, SRK)
Technoeconomic Analysis: Equipment costs, operating costs, and profitability metrics
Extensible Architecture: Plugin system for custom unit operations
Documentation Contents#
Core Modules#
Document |
Description |
|---|---|
Installation, quick start, and basic examples |
|
Reactors, separators, heat exchangers, distillation |
|
Bioreactors, centrifugation, filtration, chromatography |
|
Rare earth element extraction, scrubbing, stripping, precipitation |
|
Property calculations, equations of state, databases |
|
Capital costs, operating costs, profitability analysis |
|
Stream handling, flowsheet solver, recycle calculations |
|
Tear-stream guesses, the three acceleration methods, the traced fallback, and what to do when a solve will not converge |
|
Simultaneous solution of the whole flowsheet, SM initialization, adding EO support to a unit |
|
Transient simulation, ODE/DAE integration, diffrax backend |
|
Constrained MHE over dynamic flowsheets, EKF baseline, delayed and multi-rate data, joint parameter estimation |
|
Steady-state controllability from AD gains: RGA, singular values, disturbance rejection |
|
Feasibility function, flexibility index, and the feed-vs-parameter uncertainty split |
|
Two-stage design under a parameter distribution: scenario sampling, CVaR and chance constraints, VSS and EVPI |
|
Fisher-information D/A/E-optimal run selection, predicted confidence intervals, structural identifiability |
|
Bridging flowsheets to pounce (NLP, post-optimal sensitivity) and discopt (implicit residual blocks) |
|
Numerical methods, uncertainty propagation |
|
Constrained least squares on noisy plant data, gross error detection, observability |
|
AD-generated delta vectors, trust-region LP/MILP planning, sensitivity of the plan |
|
Design proposal: shipping difflow delta vectors into a PIMS planning model (nothing implemented yet) |
Architecture Overview#
difflow/
├── streams.py # Stream data structures
├── thermo.py # Ideal thermodynamics
├── eos.py # Cubic equations of state
├── database.py # Species property database
├── solvers.py # Numerical solvers
├── flowsheet.py # Flowsheet management
├── reconciliation/ # Data reconciliation and gross error detection
├── planning/ # Delta-base planning (LP/MILP from flowsheets)
├── uncertainty.py # Uncertainty propagation
├── cantera_import.py # Cantera data import
├── pyglenn_import.py # NASA Glenn (pyglenn) thermo import
├── dwsim_import.py # DWSIM thermo import (pythonnet; DWSIM 9.0.5)
├── units/ # Unit operations
│ ├── cstr.py # CSTR reactors
│ ├── pfr.py # PFR reactors
│ ├── fed_batch.py # Fed-batch reactors
│ ├── flash.py # Flash, mixer, splitter
│ ├── distillation.py # Distillation columns
│ ├── heat_exchanger.py # Heat exchangers
│ └── lle.py # Liquid-liquid extraction
├── dynamic/ # Dynamic (transient) simulation
│ ├── state.py # State variable specification
│ ├── base.py # DynamicUnit protocol
│ ├── integrators.py # ODE integration (RK4, RK45)
│ ├── flowsheet.py # Multi-unit dynamic flowsheets
│ ├── dae.py # DAE systems and Newton solver
│ └── diffrax_backend.py # Advanced solvers via diffrax
└── economics/ # Technoeconomic analysis
├── capital.py # Equipment costs
├── utilities.py # Utility costs
├── opex.py # Operating costs
├── profitability.py # Financial metrics
└── indices.py # Cost indices
difflow_bio/
└── units/ # Bio-manufacturing operations
├── bioreactors.py # Bioreactors
├── centrifuge.py # Centrifugation
├── filtration.py # Membrane filtration
└── chromatography.py # Chromatography
Quick Example#
import jax.numpy as jnp
from difflow.streams import make_stream
from difflow.units.cstr import CSTR, CSTRParams
from difflow.thermo import IdealThermo, SpeciesData
# Define species
species_data = {
'A': SpeciesData(name='A', MW=50.0, Cp_coeffs=(30.0, 0.01, 0.0, 0.0),
Hvap_coeffs=(35000.0, 0.38, 500.0), antoine_coeffs=(10.0, 1500.0, -40.0),
Hf=0.0, Tref=298.15),
'B': SpeciesData(name='B', MW=100.0, Cp_coeffs=(40.0, 0.02, 0.0, 0.0),
Hvap_coeffs=(40000.0, 0.38, 550.0), antoine_coeffs=(10.5, 1800.0, -50.0),
Hf=-50000.0, Tref=298.15)
}
thermo = IdealThermo(species_data)
# Create inlet stream
inlet = make_stream({'A': 1.0, 'B': 0.0}, T=350.0, P=101325.0)
# Define CSTR parameters
def rate_fn(C, T, rp):
k = rp['k_ref'] * jnp.exp(-rp['E_a'] / 8.314 * (1.0 / T - 1.0 / rp['T_ref']))
return jnp.array([k * C['A']])
params = CSTRParams(
V=1.0, # m³
rate_fn=rate_fn,
stoich=jnp.array([[-1.0], [1.0]]), # A -> B
rate_params={'k_ref': 0.1, 'E_a': 50000.0, 'T_ref': 350.0},
species_order=['A', 'B'],
dH_rxn=jnp.array([-50000.0]),
molar_density=10.0, # mol/m³, sets the concentration basis
)
# Create and run CSTR
cstr = CSTR(params, thermo)
outlet, info = cstr(inlet, T_spec=350.0)
print(f"Conversion: {float(info['conversion']['A']):.2%}")
print(f"Heat duty: {float(info['Q']):.2f} W")
Design Philosophy#
JAX-First Approach#
Every calculation in Difflow is designed to be compatible with JAX automatic differentiation:
import jax
# Gradient of product flow with respect to reactor temperature
grad_fn = jax.grad(lambda T_in: cstr(make_stream({'A': 1.0, 'B': 0.0}, T=T_in, P=101325.0), T_spec=T_in)[0]['F_B'])
sensitivity = grad_fn(350.0)
Consistent Interface#
All unit operations follow the same calling convention:
outlet, info = cstr(inlet, T_spec=350.0) # any unit: (inlet, **operating_params)
Where:
inletis a Stream dictionaryoutletis the output Streaminfocontains additional results (heat duties, conversions, etc.)
Pytree Compatibility#
Streams are JAX pytrees, allowing seamless use with jax.vmap, jax.jit, and other transformations.
Requirements#
Python >= 3.9
JAX >= 0.4.0
NumPy
PyYAML (for Cantera import)
pyglenn (optional, for NASA Glenn thermo import)
License#
MIT License