Equation-Oriented (EO) Solver

Equation-Oriented (EO) Solver#

This notebook demonstrates difflow’s equation-oriented solver, an alternative to the default sequential modular (SM) approach.

SM vs EO: Key Differences#

Feature	Sequential Modular (SM)	Equation-Oriented (EO)
Strategy	Evaluate units one-by-one, iterate on tear streams	Assemble all equations into F(x)=0, solve simultaneously
Solver	Fixed-point iteration (Wegstein/Anderson)	Newton’s method via optimistix
Best for	Simple flowsheets, few recycles	Tightly coupled recycles, sensitivity analysis
Gradients	Implicit differentiation through fixed-point	Implicit differentiation through root_find

Setup#

import jax
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True)

from difflow import (
    CSTR, CSTRParams,
    Flash, FlashParams,
    Mixer, Splitter,
    Heater, HeaterParams,
    Flowsheet, Unit,
    EOSolver, EOSolveResult,
    IdealThermo, SpeciesData,
    make_stream, get_flows,
)

print("JAX version:", jax.__version__)

JAX version: 0.8.0

Define thermodynamics and kinetics#

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)

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

stoich = jnp.array([[-1.0], [+1.0]])
rate_params = {"A": jnp.array(1e6), "Ea": jnp.array(50000.0)}
print("Thermodynamics and kinetics defined.")

Thermodynamics and kinetics defined.

Example 1: Simple CSTR — SM vs EO#

First, let’s solve a simple CSTR (no recycle) with both approaches and verify they agree.

cstr_params = CSTRParams(
    V=jnp.array(1.0),
    rate_fn=rate_fn,
    stoich=stoich,
    rate_params=rate_params,
    species_order=["A", "B"],
)
cstr = CSTR(cstr_params, thermo=thermo, mode="isothermal")

fs = Flowsheet(species_order=["A", "B"])
feed = make_stream({"A": 10.0, "B": 0.0}, T=300.0, P=101325.0)
fs.add_feed("feed", feed)
fs.add_unit(Unit("reactor", cstr, ["feed"], ["reactor_out"],
                 params={"T_spec": 350.0}))

# SM solution
sm_streams = fs.solve()

# EO solution
eo_streams = fs.solve_eo(use_sm_init=False)

print("SM solution:")
sm_flows = get_flows(sm_streams["reactor_out"])
print(f"  F_A = {float(sm_flows['A']):.6f}, F_B = {float(sm_flows['B']):.6f}")

print("\nEO solution:")
eo_flows = get_flows(eo_streams["reactor_out"])
print(f"  F_A = {float(eo_flows['A']):.6f}, F_B = {float(eo_flows['B']):.6f}")

print(f"\nMax difference: {max(abs(float(sm_flows[s] - eo_flows[s])) for s in ['A', 'B']):.2e}")

SM solution:
  F_A = 8.529305, F_B = 1.470695

EO solution:
  F_A = 8.529305, F_B = 1.470695

Max difference: 0.00e+00

Example 2: CSTR + Splitter Recycle#

Now let’s solve a flowsheet with a recycle loop using both solvers.

mixer = Mixer(species_order=["A", "B"])
splitter = Splitter(species_order=["A", "B"])

fs2 = Flowsheet(species_order=["A", "B"])
fs2.add_feed("feed", feed)
fs2.add_unit(Unit("mixer", mixer, ["feed", "recycle"], ["mixed"]))
fs2.add_unit(Unit("reactor", cstr, ["mixed"], ["reactor_out"],
                  params={"T_spec": 350.0}))
fs2.add_unit(Unit("splitter", splitter, ["reactor_out"],
                  ["product", "recycle"],
                  params={"split_frac": 0.7}))
fs2.add_recycle("recycle", "recycle")

# SM solution
sm2 = fs2.solve(tol=1e-8, max_iter=200)

# EO solution (initialized from SM for reliability)
eo2 = fs2.solve_eo(use_sm_init=True, tol=1e-8)

print("Product stream comparison:")
sm_prod = get_flows(sm2["product"])
eo_prod = get_flows(eo2["product"])
print(f"  SM:  F_A = {float(sm_prod['A']):.6f}, F_B = {float(sm_prod['B']):.6f}")
print(f"  EO:  F_A = {float(eo_prod['A']):.6f}, F_B = {float(eo_prod['B']):.6f}")
print(f"  Max diff: {max(abs(float(sm_prod[s] - eo_prod[s])) for s in ['A', 'B']):.2e}")

Product stream comparison:
  SM:  F_A = 8.529305, F_B = 1.470695
  EO:  F_A = 8.529305, F_B = 1.470695
  Max diff: 7.27e-10

Example 3: Differentiation Through the EO Solution#

The EO solver supports automatic differentiation via optimistix’s implicit differentiation. Let’s compute how product B flow changes with reactor volume.

def product_B_vs_volume(V):
    """Compute product B flow as a function of reactor volume."""
    params = CSTRParams(
        V=V,
        rate_fn=rate_fn,
        stoich=stoich,
        rate_params=rate_params,
        species_order=["A", "B"],
    )
    cstr_v = CSTR(params, thermo=thermo, mode="isothermal")

    fs = Flowsheet(species_order=["A", "B"])
    fs.add_feed("feed", feed)
    fs.add_unit(Unit("reactor", cstr_v, ["feed"], ["reactor_out"],
                     params={"T_spec": 350.0}))
    streams = fs.solve_eo(use_sm_init=False)
    return streams["reactor_out"]["F_B"]

V = jnp.array(1.0)
F_B = product_B_vs_volume(V)
dFB_dV = jax.grad(product_B_vs_volume)(V)

print(f"At V = {float(V):.1f} m^3:")
print(f"  F_B = {float(F_B):.6f} mol/s")
print(f"  dF_B/dV = {float(dFB_dV):.6f} mol/s per m^3")
print(f"\nInterpretation: Increasing volume by 0.1 m^3 increases B production by ~{float(dFB_dV)*0.1:.4f} mol/s")

At V = 1.0 m^3:
  F_B = 1.470695 mol/s
  dF_B/dV = 1.254401 mol/s per m^3

Interpretation: Increasing volume by 0.1 m^3 increases B production by ~0.1254 mol/s

Example 4: Using the EOSolver Directly#

For more control, use EOSolver directly to get convergence diagnostics.

fs3 = Flowsheet(species_order=["A", "B"])
fs3.add_feed("feed", feed)
fs3.add_unit(Unit("reactor", cstr, ["feed"], ["reactor_out"],
                  params={"T_spec": 350.0}))

solver = EOSolver(fs3)
result = solver.solve(use_sm_init=False, tol=1e-10)

print(f"Converged: {result.converged}")
print(f"Residual norm: {result.residual_norm:.2e}")
print(f"Newton iterations: {result.n_iterations}")
print(f"Wall time: {result.wall_time:.4f} s")
print(f"\nStreams solved: {list(result.streams.keys())}")

Converged: True
Residual norm: 0.00e+00
Newton iterations: 1
Wall time: 0.4981 s

Streams solved: ['feed', 'reactor_out']

When to Use EO vs SM#

Use SM (default solve()) when:

Flowsheet has no recycles or simple recycles
Units are expensive to evaluate (EO calls them many times per Newton step)
You want simplicity

Use EO (solve_eo()) when:

Flowsheet has tightly coupled recycles that SM struggles with
You need the full system Jacobian
You want implicit differentiation through the entire flowsheet
Newton’s quadratic convergence is beneficial (fewer iterations near solution)