CSTR + Flash Separator with Recycle#
This notebook demonstrates a classic reaction-separation flowsheet that is fully differentiable using JAX.
Flowsheet Overview#
┌─────────┐
Fresh A ──────►│ │ ┌─────────┐
│ CSTR ├─────►│ Flash ├───► Product B (vapor)
Recycle ──────►│ │ │ │
▲ └─────────┘ └────┬────┘
│ │
└───────────────────────────────┘
(liquid recycle)
Process Description:
Fresh feed of reactant A enters the system
CSTR performs the reaction A → B (first-order kinetics)
Flash separator removes product B as vapor (more volatile)
Liquid (unreacted A) is recycled back to the CSTR
Why Differentiability Matters:
Compute sensitivities: How do outputs change with parameters?
Enable gradient-based optimization
Propagate uncertainties through the flowsheet
Setup and Imports#
import jax
import jax.numpy as jnp
from jax import Array
# Enable 64-bit precision for better numerical accuracy
jax.config.update("jax_enable_x64", True)
import optimistix as optx
from difflow.streams import Stream, make_stream, get_flows, combine_streams
from difflow.thermo import IdealThermo, SpeciesData
from difflow.units.cstr import CSTR, CSTRParams
from difflow.units.flash import Flash, FlashParams, Mixer
from difflow.flowsheet import Flowsheet, Unit
print("JAX version:", jax.__version__)
print("64-bit precision enabled:", jax.config.jax_enable_x64)
WARNING:2026-03-01 07:48:32,484:jax._src.xla_bridge:905: Platform 'mps' is experimental and not all JAX functionality may be correctly supported!
WARNING:2026-03-01 07:48:32,513:jax._src.xla_bridge:905: Platform 'METAL' is experimental and not all JAX functionality may be correctly supported!
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1772369312.513448 12631555 mps_client.cc:510] WARNING: JAX Apple GPU support is experimental and not all JAX functionality is correctly supported!
I0000 00:00:1772369312.513655 12631555 service.cc:145] XLA service 0x8bdaa4800 initialized for platform METAL (this does not guarantee that XLA will be used). Devices:
I0000 00:00:1772369312.513661 12631555 service.cc:153] StreamExecutor device (0): Metal, <undefined>
I0000 00:00:1772369312.514232 12631555 mps_client.cc:406] Using Simple allocator.
I0000 00:00:1772369312.514240 12631555 mps_client.cc:384] XLA backend will use up to 55662313472 bytes on device 0 for SimpleAllocator.
Metal device set to: Apple M4 Pro
JAX version: 0.9.0.1
64-bit precision enabled: True
Define Species and Thermodynamic Properties#
We define two species:
Species A (Reactant): Heavy, less volatile (like toluene)
Species B (Product): Light, more volatile (like benzene)
The key thermodynamic property is the K-value (vapor-liquid equilibrium ratio):
K = y/x = Psat/P (for ideal mixtures)
K < 1: Component prefers liquid phase
K > 1: Component prefers vapor phase
We design A to have K < 1 (stays in liquid/recycle) and B to have K > 1 (leaves as product).
# Species A: Reactant (heavy, less volatile - like toluene)
# Species B: Product (light, more volatile - like benzene)
# Antoine coefficients: log10(Psat/Pa) = A - B/(T + C)
species_data = {
"A": SpeciesData(
name="A",
MW=92.0, # g/mol (like toluene)
Cp_coeffs=(75.0, 0.0, 0.0, 0.0), # J/mol/K (constant Cp)
Hvap_coeffs=(35000.0, 0.38, 590.0), # Watson correlation
# At 350K: K_A ≈ 0.035 (stays in liquid)
antoine_coeffs=(10.0, 2000.0, -40.0),
Hf=0.0,
),
"B": SpeciesData(
name="B",
MW=78.0, # g/mol (like benzene)
Cp_coeffs=(50.0, 0.0, 0.0, 0.0), # J/mol/K
Hvap_coeffs=(30000.0, 0.38, 560.0), # Watson correlation
# At 350K: K_B ≈ 1.43 (goes to vapor)
antoine_coeffs=(10.0, 1500.0, -40.0),
Hf=-50000.0, # Exothermic reaction A → B
),
}
thermo = IdealThermo(species_data)
species_order = ["A", "B"]
# Verify K-values at flash conditions
T_flash = 350.0 # K
P_flash = 101325.0 # Pa
print("K-values at flash conditions (T=350K, P=1 atm):")
for species in ["A", "B"]:
Psat = thermo.Psat(species, T_flash)
K = float(Psat / P_flash)
print(f" {species}: Psat = {float(Psat):.0f} Pa, K = {K:.4f}")
print("\n✓ A stays in liquid (K < 1), B goes to vapor (K > 1)")
K-values at flash conditions (T=350K, P=1 atm):
A: Psat = 3535 Pa, K = 0.0349
B: Psat = 144974 Pa, K = 1.4308
✓ A stays in liquid (K < 1), B goes to vapor (K > 1)
Define Reaction Kinetics#
The CSTR performs a first-order reaction: A → B
Rate expression: $\(r = k \cdot C_A\)$
where the rate constant follows Arrhenius kinetics: $\(k = A \cdot \exp\left(-\frac{E_a}{RT}\right)\)$
The stoichiometry matrix defines how each species changes per reaction:
ν_A = -1 (A is consumed)
ν_B = +1 (B is produced)
def rate_function(C: dict[str, Array], T: Array, params: dict) -> Array:
"""First-order reaction: A → B.
Rate = k * C_A where k = A * exp(-Ea / RT)
Args:
C: Concentrations (mol/m³)
T: Temperature (K)
params: {"A": pre-exponential, "Ea": activation energy}
Returns:
Array of reaction rates [r1] (mol/m³/s)
"""
A = params["A"]
Ea = params["Ea"]
R = 8.314 # J/mol/K
k = A * jnp.exp(-Ea / (R * T))
r = k * C["A"]
return jnp.array([r])
# Stoichiometry: A → B means ν_A = -1, ν_B = +1
stoichiometry = jnp.array([
[-1.0], # A
[+1.0], # B
])
print("Reaction: A → B")
print(f"Stoichiometry matrix:\n{stoichiometry}")
Reaction: A → B
Stoichiometry matrix:
[[-1.]
[ 1.]]
The Differentiable Flowsheet Solver#
The key challenge with recycle loops is that the recycle stream depends on itself:
Recycle composition → Reactor inlet → Reactor outlet → Flash → Recycle composition
We solve this using fixed-point iteration:
Guess initial recycle composition
Compute reactor inlet (mix fresh feed + recycle)
Compute reactor outlet
Compute flash products
Liquid from flash = new recycle estimate
Repeat until converged
The solver uses optimistix for fixed-point iteration, which provides implicit differentiation for automatic gradients!
def solve_cstr_flash_recycle(
params: dict[str, Array],
fresh_feed: Stream,
tol: float = 1e-8,
max_iter: int = 100,
) -> dict:
"""Solve CSTR + Flash with recycle as a pure function.
This function is fully differentiable with respect to params.
Args:
params: Dictionary with keys:
- 'V_reactor': Reactor volume (m³)
- 'T_reactor': Reactor temperature (K)
- 'T_flash': Flash temperature (K)
- 'P_flash': Flash pressure (Pa)
- 'k_A': Arrhenius pre-exponential factor (1/s)
- 'k_Ea': Activation energy (J/mol)
fresh_feed: Fresh feed stream
Returns:
Dictionary with all streams and info
"""
# Extract parameters
V_reactor = params["V_reactor"]
T_reactor = params["T_reactor"]
T_flash = params["T_flash"]
P_flash = params["P_flash"]
rate_params = {"A": params["k_A"], "Ea": params["k_Ea"]}
# Create unit operations
cstr_params = CSTRParams(
V=V_reactor,
rate_fn=rate_function,
stoich=stoichiometry,
rate_params=rate_params,
species_order=species_order,
dH_rxn=jnp.array([-50000.0]),
)
cstr = CSTR(cstr_params, thermo=thermo, mode="isothermal")
flash_params = FlashParams(species_order=species_order)
flash = Flash(flash_params, thermo=thermo)
mixer = Mixer(species_order, thermo=thermo)
# Define the fixed-point iteration function
def flowsheet_step(recycle_arr, args):
fresh, T_r, T_f, P_f = args
# Unpack recycle
recycle = make_stream(
{"A": recycle_arr[0], "B": recycle_arr[1]},
T=recycle_arr[2],
P=recycle_arr[3],
)
# Mix fresh feed with recycle
reactor_inlet, _ = mixer(fresh, recycle)
# React in CSTR
reactor_outlet, _ = cstr(reactor_inlet, T_spec=T_r)
# Flash separation
liquid, vapor, _ = flash(reactor_outlet, T=T_f, P=P_f)
# Liquid is new recycle
return jnp.array([
liquid["F_A"],
liquid["F_B"],
liquid["T"],
liquid["P"],
])
# Initial recycle guess
recycle_init = jnp.array([1.0, 0.1, T_flash, P_flash])
args = (fresh_feed, T_reactor, T_flash, P_flash)
# Solve for converged recycle using optimistix
solver = optx.FixedPointIteration(rtol=tol, atol=tol)
solution = optx.fixed_point(
flowsheet_step,
solver,
recycle_init,
args=args,
max_steps=max_iter,
throw=False,
)
recycle_converged = solution.value
# Final evaluation
recycle = make_stream(
{"A": recycle_converged[0], "B": recycle_converged[1]},
T=recycle_converged[2],
P=recycle_converged[3],
)
reactor_inlet, _ = mixer(fresh_feed, recycle)
reactor_outlet, cstr_info = cstr(reactor_inlet, T_spec=T_reactor)
liquid, vapor, flash_info = flash(reactor_outlet, T=T_flash, P=P_flash)
return {
"fresh_feed": fresh_feed,
"recycle": recycle,
"reactor_inlet": reactor_inlet,
"reactor_outlet": reactor_outlet,
"liquid": liquid,
"vapor": vapor,
"cstr_info": cstr_info,
"flash_info": flash_info,
}
print("Flowsheet solver defined ✓")
Flowsheet solver defined ✓
Solve the Flowsheet#
Now let’s solve the flowsheet with specific operating conditions:
Reactor: 2 m³ volume, 350 K
Flash: 350 K, 0.5 atm (vacuum for better vapor-liquid separation)
Kinetics: k = 0.3 s⁻¹ (first-order rate constant)
The parameters are chosen to give meaningful per-pass conversion while maintaining a recycle stream.
# Define parameters
# Note: We balance reactor size and kinetics to get meaningful per-pass conversion
# while still having unreacted A for the recycle loop.
params = {
"V_reactor": jnp.array(2.0), # m³
"T_reactor": jnp.array(350.0), # K
"T_flash": jnp.array(350.0), # K
"P_flash": jnp.array(50000.0), # Pa (vacuum for good separation)
"k_A": jnp.array(0.3), # 1/s (direct rate constant, no Arrhenius)
"k_Ea": jnp.array(0.0), # J/mol (set to 0 to use k directly)
}
# Create fresh feed
fresh_feed = make_stream({"A": 10.0, "B": 0.0}, T=300.0, P=101325.0)
print("Solving flowsheet...")
results = solve_cstr_flash_recycle(params, fresh_feed)
print("Converged! ✓")
Solving flowsheet...
Converged! ✓
View Results#
Let’s examine the stream compositions and unit operation performance.
print("=" * 60)
print("STREAM RESULTS")
print("=" * 60)
for name in ["fresh_feed", "recycle", "reactor_inlet", "reactor_outlet", "liquid", "vapor"]:
stream = results[name]
flows = get_flows(stream)
print(f"\n{name}:")
print(f" F_A = {float(flows['A']):.4f} mol/s")
print(f" F_B = {float(flows['B']):.4f} mol/s")
print(f" T = {float(stream['T']):.2f} K")
print(f" P = {float(stream['P']):.0f} Pa")
============================================================
STREAM RESULTS
============================================================
fresh_feed:
F_A = 10.0000 mol/s
F_B = 0.0000 mol/s
T = 300.00 K
P = 101325 Pa
recycle:
F_A = 0.0000 mol/s
F_B = 0.0000 mol/s
T = 350.00 K
P = 50000 Pa
reactor_inlet:
F_A = 10.0000 mol/s
F_B = 0.0000 mol/s
T = 300.00 K
P = 101325 Pa
reactor_outlet:
F_A = 0.0030 mol/s
F_B = 9.9970 mol/s
T = 350.00 K
P = 101325 Pa
liquid:
F_A = 0.0000 mol/s
F_B = 0.0000 mol/s
T = 350.00 K
P = 50000 Pa
vapor:
F_A = 0.0030 mol/s
F_B = 9.9970 mol/s
T = 350.00 K
P = 50000 Pa
print("=" * 60)
print("UNIT OPERATION PERFORMANCE")
print("=" * 60)
# CSTR Info
cstr_info = results["cstr_info"]
print("\nCSTR:")
print(f" Heat duty Q = {float(cstr_info['Q']):.2f} W")
print(f" Reaction rate = {float(cstr_info['rates'][0]):.4f} mol/m³/s")
print(f" Conversion of A = {float(cstr_info['conversion']['A'])*100:.2f}%")
# Flash Info
flash_info = results["flash_info"]
print("\nFlash Separator:")
print(f" Vapor fraction = {float(flash_info['V_frac'])*100:.2f}%")
print(f" K_A = {float(flash_info['K']['A']):.4f}")
print(f" K_B = {float(flash_info['K']['B']):.4f}")
============================================================
UNIT OPERATION PERFORMANCE
============================================================
CSTR:
Heat duty Q = -475308.50 W
Reaction rate = 4.9985 mol/m³/s
Conversion of A = 99.97%
Flash Separator:
Vapor fraction = 100.00%
K_A = 0.0707
K_B = 2.8995
Sensitivity Analysis with Automatic Differentiation#
This is where the power of differentiable simulation shines!
We can compute exact gradients of any output with respect to any parameter:
How does product B flow change with reactor volume?
How does it change with reactor temperature?
How does it change with flash conditions?
These gradients are computed using automatic differentiation - not finite differences!
def product_B_flow(params: dict) -> Array:
"""Compute product B vapor flow."""
result = solve_cstr_flash_recycle(params, fresh_feed)
return result["vapor"]["F_B"]
print("Computing gradients of product B flow w.r.t. parameters...")
print("(This uses automatic differentiation, not finite differences!)\n")
# Compute gradient of B production with respect to all parameters
grad_fn = jax.grad(lambda p: product_B_flow(p))
grads = grad_fn(params)
print("Sensitivities (∂F_B/∂parameter):")
print(f" ∂F_B/∂V_reactor = {float(grads['V_reactor']):.6f} mol/s per m³")
print(f" ∂F_B/∂k_A = {float(grads['k_A']):.6f} mol/s per (1/s)")
print(f" ∂F_B/∂T_flash = {float(grads['T_flash']):.6f} mol/s per K")
print(f" ∂F_B/∂P_flash = {float(grads['P_flash']):.10f} mol/s per Pa")
Computing gradients of product B flow w.r.t. parameters...
(This uses automatic differentiation, not finite differences!)
Sensitivities (∂F_B/∂parameter):
∂F_B/∂V_reactor = 0.001501 mol/s per m³
∂F_B/∂k_A = 0.010004 mol/s per (1/s)
∂F_B/∂T_flash = -0.000000 mol/s per K
∂F_B/∂P_flash = 0.0000000000 mol/s per Pa
Interpretation of Gradients#
The gradients tell us which parameters have the most influence on product output:
∂F_B/∂V_reactor > 0: Larger reactor → more conversion → more B
∂F_B/∂T_reactor > 0: Higher reactor T → faster kinetics → more B
∂F_B/∂T_flash: Effect depends on volatility
∂F_B/∂P_flash: Higher P → less vaporization
These sensitivities enable:
Optimization: Follow gradients to maximize B production
Uncertainty propagation: How do input uncertainties affect outputs?
Control design: Which parameters to manipulate?
# Practical interpretation
F_B_base = float(product_B_flow(params))
print("Practical Interpretation:")
print(f"\nBase case: F_B = {F_B_base:.4f} mol/s")
print(f"\nIf we increase reactor volume by 0.1 m³:")
print(f" ΔF_B ≈ {float(grads['V_reactor']) * 0.1:.4f} mol/s")
print(f"\nIf we increase rate constant k by 0.1 s⁻¹:")
print(f" ΔF_B ≈ {float(grads['k_A']) * 0.1:.4f} mol/s")
print(f"\nIf we increase flash temperature by 10 K:")
print(f" ΔF_B ≈ {float(grads['T_flash']) * 10:.4f} mol/s")
Practical Interpretation:
Base case: F_B = 9.9970 mol/s
If we increase reactor volume by 0.1 m³:
ΔF_B ≈ 0.0002 mol/s
If we increase rate constant k by 0.1 s⁻¹:
ΔF_B ≈ 0.0010 mol/s
If we increase flash temperature by 10 K:
ΔF_B ≈ -0.0000 mol/s
Summary#
In this notebook, we demonstrated:
Building a recycle flowsheet with CSTR and flash separator
Solving recycle loops using fixed-point iteration
Computing exact gradients using automatic differentiation
Interpreting sensitivities for process insights
The key advantage of differentiable simulation is that gradients are:
Exact (not approximations like finite differences)
Efficient (one backward pass gives all gradients)
Composable (chain rule handles complex flowsheets automatically)