EO vs SM Solver Benchmarks

This notebook compares the equation-oriented (EO) and sequential modular (SM) solvers on various flowsheet problems.

import jax
import jax.numpy as jnp
import time

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

from difflow import (
    CSTR, CSTRParams,
    Mixer, Splitter,
    Flowsheet, Unit,
    EOSolver,
    IdealThermo, SpeciesData,
    make_stream, get_flows,
)
from difflow.benchmarks import compare_solvers

# Thermo setup
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("Setup complete.")

Setup complete.

Benchmark 1: CSTR with Recycle

Compare SM and EO on a CSTR with 30% recycle.

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")
mixer = Mixer(species_order=["A", "B"])
splitter = Splitter(species_order=["A", "B"])

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("mixer", mixer, ["feed", "recycle"], ["mixed"]))
fs.add_unit(Unit("reactor", cstr, ["mixed"], ["reactor_out"],
                 params={"T_spec": 350.0}))
fs.add_unit(Unit("splitter", splitter, ["reactor_out"],
                 ["product", "recycle"], params={"split_frac": 0.7}))
fs.add_recycle("recycle", "recycle")

comparison = compare_solvers(fs, tol=1e-8)

print(f"SM time:  {comparison.sm_time:.4f} s")
print(f"EO time:  {comparison.eo_time:.4f} s")
print(f"EO iterations: {comparison.eo_iterations}")
print(f"Max stream difference: {comparison.max_stream_difference:.2e}")
print(f"EO converged: {comparison.eo_converged}")

SM time:  10.1030 s
EO time:  1.1670 s
EO iterations: 3
Max stream difference: 1.53e-09
EO converged: True

Benchmark 2: EO from SM Init

When initialized from the SM solution, EO should converge in very few iterations.

solver = EOSolver(fs)

# Cold start
result_cold = solver.solve(use_sm_init=False, tol=1e-8)
print(f"Cold start: {result_cold.n_iterations} iterations, "
      f"residual = {result_cold.residual_norm:.2e}, "
      f"time = {result_cold.wall_time:.4f}s")

# Warm start from SM
result_warm = solver.solve(use_sm_init=True, tol=1e-8)
print(f"SM init:    {result_warm.n_iterations} iterations, "
      f"residual = {result_warm.residual_norm:.2e}, "
      f"time = {result_warm.wall_time:.4f}s")

Cold start: 3 iterations, residual = 5.68e-14, time = 0.5897s
SM init:    2 iterations, residual = 4.44e-16, time = 1.5202s

Summary

• SM and EO converge to the same solution (within tolerance)

• EO with SM initialization converges in few Newton iterations

• Both approaches support automatic differentiation through the solution