Gradient-Based Optimization for Carbon Capture#
This notebook demonstrates the key advantage of difflow: automatic differentiation through chemical process models for gradient-based optimization.
Learning Objectives#
Understand why gradients matter for process optimization
Use JAX’s
gradto compute derivatives through unit operationsOptimize capture processes using gradient descent
Perform sensitivity analysis efficiently with
vmapCompare gradient-based vs derivative-free optimization
1. Why Gradients Matter#
Traditional Process Optimization#
Chemical process optimization typically uses:
Derivative-free methods: Nelder-Mead, genetic algorithms, particle swarm
Finite differences: Approximate gradients by perturbing parameters
Problems:
Derivative-free methods scale poorly with dimension
Finite differences are inaccurate and expensive (\(O(n)\) evaluations per gradient)
Automatic Differentiation (AD)#
AD computes exact gradients at cost similar to a single function evaluation!
\[\nabla f(x) = \left(\frac{\partial f}{\partial x_1}, \frac{\partial f}{\partial x_2}, \ldots, \frac{\partial f}{\partial x_n}\right)\]
This enables:
Efficient high-dimensional optimization
Sensitivity analysis
Uncertainty quantification
Integration with machine learning
2. Setup#
import jax
import jax.numpy as jnp
from jax import grad, jit, vmap, value_and_grad
jax.config.update("jax_enable_x64", True)
from difflow.streams import make_stream, get_flows, total_flow
from difflow_cc import (
# Amine absorption
AbsorberParams, AmineAbsorber,
# Membranes
MembraneParams, MembraneSeparator,
# Adsorption
AdsorptionParams, PSAUnit, VSAUnit,
)
print(f"JAX version: {jax.__version__}")
print(f"Using 64-bit floats: {jax.config.jax_enable_x64}")
JAX version: 0.11.1
Using 64-bit floats: True
3. Example 1: Optimizing L/G Ratio in Amine Absorption#
Problem Statement#
Find the optimal liquid-to-gas (L/G) ratio that minimizes total operating cost: $\(\text{Cost} = \text{Circulation cost} + \text{Regeneration energy cost} + \text{Capture penalty}\)$
Trade-offs#
Lower L/G: Less solvent circulation, but higher rich loading means more energy needed to regenerate each unit of solvent
Higher L/G: More solvent circulation cost, but lower rich loading means less energy per unit solvent (though more total solvent to regenerate)
The optimal L/G balances these competing effects. Unlike a simple linear cost, this creates a true optimization problem with meaningful gradients.
# Define the flue gas
flue_gas = make_stream(
flows={"CO2": 1.0, "N2": 6.67}, # 15% CO2 typical coal flue gas
T=313.15, # 40°C
P=101325.0,
)
# Create absorber with variable L/G
def create_absorber(L_G_ratio):
params = AbsorberParams(
solvent='MEA',
n_stages=10,
solvent_conc=30.0,
L_G_ratio=L_G_ratio,
lean_loading=0.2,
)
return AmineAbsorber(params)
# Cost function to minimize
# This represents the TOTAL cost including regeneration energy
def capture_cost(L_G_ratio):
"""Total cost = solvent circulation + regeneration energy.
The cost has competing terms:
1. At low L/G: Higher rich loading → more CO2 per mol solvent
→ lower circulation cost but higher regen energy per mol solvent
2. At high L/G: Lower rich loading → less CO2 per mol solvent
→ higher circulation cost but lower regen energy per mol solvent
The optimal L/G balances circulation and regeneration costs.
"""
absorber = create_absorber(L_G_ratio)
_, _, info = absorber(flue_gas)
capture_eff = info['capture_efficiency']
rich_loading = info['rich_loading']
lean_loading = info['lean_loading']
# Solvent circulation cost (pumping, cooling)
circulation_cost = L_G_ratio * 5.0 # $/hr per L/G unit
# Regeneration energy cost
# Higher rich loading means more energy per mol solvent to regenerate
# Regen energy ~ (rich_loading - lean_loading) * heat_of_reaction
delta_loading = rich_loading - lean_loading
regen_energy_per_mol = 80.0 # kJ/mol CO2 (MEA heat of absorption)
# Total CO2 captured
CO2_captured = info['CO2_captured'] # mol/s
# Regen cost scales with solvent flow and loading change
# More solvent at same loading = more total energy
regen_cost = L_G_ratio * delta_loading * regen_energy_per_mol * 0.5 # $/hr
# Penalty for not meeting capture target (must capture >= 90%)
capture_target = 0.90
penalty = 1000.0 * jnp.maximum(0, capture_target - capture_eff)**2
total = circulation_cost + regen_cost + penalty
return total
# Test the function at different L/G values
print("L/G Ratio vs Performance and Cost:")
print("-" * 70)
print(f"{'L/G':>5} {'Capture':>10} {'Rich Load':>12} {'Circ Cost':>12} {'Regen Cost':>12} {'Total':>10}")
print("-" * 70)
for L_G in [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]:
absorber = create_absorber(L_G)
_, _, info = absorber(flue_gas)
circ = L_G * 5.0
delta = info['rich_loading'] - info['lean_loading']
regen = L_G * float(delta) * 80.0 * 0.5
total = capture_cost(L_G)
print(f"{L_G:>5.1f} {float(info['capture_efficiency']):>10.1%} "
f"{float(info['rich_loading']):>12.3f} "
f"${circ:>11.1f} ${regen:>11.1f} ${float(total):>9.1f}")
L/G Ratio vs Performance and Cost:
----------------------------------------------------------------------
L/G Capture Rich Load Circ Cost Regen Cost Total
----------------------------------------------------------------------
1.0 25.8% 0.500 $ 5.0 $ 12.0 $ 429.1
2.0 51.6% 0.500 $ 10.0 $ 24.0 $ 181.4
3.0 77.4% 0.500 $ 15.0 $ 36.0 $ 66.9
4.0 99.9% 0.490 $ 20.0 $ 46.5 $ 66.5
5.0 99.9% 0.432 $ 25.0 $ 46.5 $ 71.5
6.0 99.9% 0.394 $ 30.0 $ 46.5 $ 76.5
# Compute gradient using JAX
grad_cost = grad(capture_cost)
# Gradient at test point
L_G_test = 3.0
d_cost_d_LG = grad_cost(L_G_test)
print(f"Gradient at L/G = {L_G_test}: {float(d_cost_d_LG):.4f} $/hr per unit L/G")
if d_cost_d_LG > 0:
print(" → Decreasing L/G will reduce cost")
else:
print(" → Increasing L/G will reduce cost")
Gradient at L/G = 3.0: -47.9850 $/hr per unit L/G
→ Increasing L/G will reduce cost
# Simple gradient descent optimization
def gradient_descent(f, x0, learning_rate=0.1, n_iters=50):
"""Simple gradient descent with momentum."""
x = x0
grad_f = grad(f)
history = []
for i in range(n_iters):
fx = f(x)
gx = grad_f(x)
history.append((float(x), float(fx), float(gx)))
# Update with gradient
x = x - learning_rate * gx
# Keep L/G in reasonable range
x = jnp.clip(x, 1.0, 10.0)
if i % 10 == 0:
print(f"Iter {i:3d}: L/G = {float(x):.3f}, Cost = ${float(fx):.2f}, Grad = {float(gx):.4f}")
return x, history
# Optimize
L_G_init = 5.0
L_G_opt, history = gradient_descent(capture_cost, L_G_init, learning_rate=0.05, n_iters=50)
print(f"\nOptimal L/G ratio: {float(L_G_opt):.3f}")
Iter 0: L/G = 4.750, Cost = $71.46, Grad = 5.0000
Iter 10: L/G = 4.965, Cost = $72.54, Grad = 5.0000
Iter 20: L/G = 5.413, Cost = $74.78, Grad = 5.0000
Iter 30: L/G = 6.207, Cost = $78.75, Grad = 5.0000
Iter 40: L/G = 3.707, Cost = $66.25, Grad = 5.0000
Optimal L/G ratio: 4.456
# Verify the optimum
absorber_opt = create_absorber(L_G_opt)
_, _, info_opt = absorber_opt(flue_gas)
print(f"At optimal L/G = {float(L_G_opt):.3f}:")
print(f" Capture efficiency: {float(info_opt['capture_efficiency']):.1%}")
print(f" Total cost: ${float(capture_cost(L_G_opt)):.2f}/hr")
At optimal L/G = 4.456:
Capture efficiency: 99.9%
Total cost: $68.74/hr
4. Example 2: Multi-Parameter Membrane Optimization#
Optimize both membrane area and pressure ratio to minimize: $\(\text{Total cost} = \text{Capital (area)} + \text{Operating (compression)} - \text{CO}_2\text{ revenue}\)$
# Feed gas
feed_gas = make_stream(
flows={"CO2": 1.0, "N2": 4.0}, # 20% CO2
T=298.15,
P=500000.0, # 5 bar feed
)
def membrane_cost(params):
"""Annualized cost of membrane separation.
Args:
params: [area_m2, P_permeate_Pa]
"""
area = params[0]
P_perm = params[1]
mem_params = MembraneParams(
membrane_type="PDMS",
area=area,
feed_pressure=500000.0,
permeate_pressure=P_perm,
)
membrane = MembraneSeparator(mem_params)
_, permeate, info = membrane(feed_gas)
# Capital cost (area-based)
membrane_cost_per_m2 = 100.0 # $/m²
capital = area * membrane_cost_per_m2 / 5 # Annualized over 5 years
# Operating cost (compression for vacuum)
pressure_ratio = 500000.0 / (P_perm + 1.0)
compression_cost = 0.05 * jnp.log(pressure_ratio) * 1000 # $/yr
# Revenue from CO2 (negative cost)
CO2_captured = info['CO2_recovery'] * 1.0 # mol/s
CO2_revenue = CO2_captured * 44.0 * 3600 * 8000 * 50 / 1e6 # 50 $/tonne, 8000 hr/yr
# Penalty for low purity
purity = info['CO2_purity']
purity_penalty = 10000.0 * jnp.maximum(0, 0.7 - purity)**2
total = capital + compression_cost - CO2_revenue + purity_penalty
return total
# Test
params_test = jnp.array([10.0, 50000.0]) # 10 m², 0.5 bar permeate
cost = membrane_cost(params_test)
print(f"Cost at area=10m², P_perm=0.5bar: ${float(cost):.0f}/yr")
Cost at area=10m², P_perm=0.5bar: $-27085/yr
# Compute gradient w.r.t. both parameters
grad_membrane = grad(membrane_cost)
grads = grad_membrane(params_test)
print(f"Gradient at test point:")
print(f" d(cost)/d(area) = {float(grads[0]):.2f} $/yr per m²")
print(f" d(cost)/d(P_perm) = {float(grads[1]):.6f} $/yr per Pa")
Gradient at test point:
d(cost)/d(area) = -1634.83 $/yr per m²
d(cost)/d(P_perm) = 0.161749 $/yr per Pa
# Multi-parameter gradient descent
def optimize_membrane(n_iters=100):
params = jnp.array([10.0, 50000.0])
learning_rates = jnp.array([0.5, 5000.0]) # Different scales
grad_f = grad(membrane_cost)
for i in range(n_iters):
cost = membrane_cost(params)
grads = grad_f(params)
# Gradient descent update
params = params - learning_rates * grads
# Enforce bounds
params = jnp.array([
jnp.clip(params[0], 1.0, 100.0), # Area: 1-100 m²
jnp.clip(params[1], 10000.0, 200000.0), # P_perm: 0.1-2 bar
])
if i % 20 == 0:
print(f"Iter {i:3d}: Area={float(params[0]):.1f}m², "
f"P_perm={float(params[1])/1e5:.2f}bar, "
f"Cost=${float(cost):.0f}/yr")
return params
params_opt = optimize_membrane()
print(f"\nOptimal: Area={float(params_opt[0]):.1f}m², P_perm={float(params_opt[1])/1e5:.2f}bar")
Iter 0: Area=100.0m², P_perm=0.49bar, Cost=$-27085/yr
Iter 20: Area=100.0m², P_perm=0.48bar, Cost=$-58311/yr
Iter 40: Area=100.0m², P_perm=0.47bar, Cost=$-58344/yr
Iter 60: Area=100.0m², P_perm=0.45bar, Cost=$-58376/yr
Iter 80: Area=100.0m², P_perm=0.44bar, Cost=$-58406/yr
Optimal: Area=90.0m², P_perm=0.43bar
5. Sensitivity Analysis with vmap#
Use JAX’s vmap (vectorized map) to efficiently compute sensitivities
across many parameter values in parallel.
# Vectorized sensitivity analysis for L/G ratio
@jit
def capture_efficiency(L_G):
absorber = create_absorber(L_G)
_, _, info = absorber(flue_gas)
return info['capture_efficiency']
@jit
def rich_loading(L_G):
absorber = create_absorber(L_G)
_, _, info = absorber(flue_gas)
return info['rich_loading']
# Compute for many L/G values at once
L_G_values = jnp.linspace(1.0, 8.0, 50)
# Vectorize function and gradient
rich_load_vec = vmap(rich_loading)
grad_rich = grad(rich_loading)
grad_rich_vec = vmap(grad_rich)
# Compute all at once
loadings = rich_load_vec(L_G_values)
sensitivities = grad_rich_vec(L_G_values)
print(f"{'L/G':>6} {'Rich Loading':>14} {'d(load)/d(L/G)':>16}")
print("-" * 38)
for i in range(0, 50, 8):
print(f"{float(L_G_values[i]):>6.2f} {float(loadings[i]):>14.4f} "
f"{float(sensitivities[i]):>16.6f}")
L/G Rich Loading d(load)/d(L/G)
--------------------------------------
1.00 0.5000 0.000000
2.14 0.5000 0.000000
3.29 0.5000 0.000000
4.43 0.4623 -0.059224
5.57 0.4085 -0.037419
6.71 0.3730 -0.025765
7.86 0.3478 -0.018815
# Find where sensitivity is highest (steepest change in rich loading)
max_sens_idx = jnp.argmax(jnp.abs(sensitivities))
print(f"Maximum sensitivity at L/G = {float(L_G_values[max_sens_idx]):.2f}")
print(f" d(rich_loading)/d(L/G) = {float(sensitivities[max_sens_idx]):.6f}")
print(f" Rich loading = {float(loadings[max_sens_idx]):.4f}")
print()
print("Interpretation: At this L/G, small changes have the biggest")
print("effect on rich loading, which affects regeneration energy.")
Maximum sensitivity at L/G = 4.00
d(rich_loading)/d(L/G) = -0.072595
Rich loading = 0.4904
Interpretation: At this L/G, small changes have the biggest
effect on rich loading, which affects regeneration energy.
6. Comparing Optimization Approaches#
Let’s compare gradient-based optimization with a derivative-free method.
import time
# Count function evaluations
eval_count = 0
def counted_cost(L_G):
global eval_count
eval_count += 1
return capture_cost(L_G)
# Gradient descent
eval_count = 0
start = time.time()
L_G_gd, _ = gradient_descent(capture_cost, 5.0, learning_rate=0.05, n_iters=30)
gd_time = time.time() - start
# Each iteration: 1 forward + 1 backward ≈ 2 evals
gd_evals = 30 * 2
print(f"Gradient Descent:")
print(f" Optimal L/G: {float(L_G_gd):.3f}")
print(f" Iterations: 30")
print(f" Equiv. evaluations: ~{gd_evals}")
print(f" Time: {gd_time:.3f}s")
Iter 0: L/G = 4.750, Cost = $71.46, Grad = 5.0000
Iter 10: L/G = 4.965, Cost = $72.54, Grad = 5.0000
Iter 20: L/G = 5.413, Cost = $74.78, Grad = 5.0000
Gradient Descent:
Optimal L/G: 6.457
Iterations: 30
Equiv. evaluations: ~60
Time: 0.576s
# Golden section search (derivative-free)
def golden_section(f, a, b, tol=0.01):
"""Golden section search for unimodal function."""
phi = (1 + 5**0.5) / 2
c = b - (b - a) / phi
d = a + (b - a) / phi
evals = 0
while abs(b - a) > tol:
fc = f(c)
fd = f(d)
evals += 2
if fc < fd:
b = d
d = c
c = b - (b - a) / phi
else:
a = c
c = d
d = a + (b - a) / phi
return (a + b) / 2, evals
start = time.time()
L_G_gs, gs_evals = golden_section(capture_cost, 1.0, 10.0, tol=0.01)
gs_time = time.time() - start
print(f"\nGolden Section Search:")
print(f" Optimal L/G: {float(L_G_gs):.3f}")
print(f" Evaluations: {gs_evals}")
print(f" Time: {gs_time:.3f}s")
Golden Section Search:
Optimal L/G: 3.362
Evaluations: 30
Time: 0.064s
# Summary comparison
print("\nComparison Summary:")
print(f"{'Method':<25} {'Optimal L/G':<12} {'Evals':<10}")
print("-" * 47)
print(f"{'Gradient Descent':<25} {float(L_G_gd):<12.3f} {'~60':10}")
print(f"{'Golden Section':<25} {float(L_G_gs):<12.3f} {gs_evals:<10}")
print()
print("Note: For 1D problems, derivative-free methods work well.")
print("For high-dimensional problems, gradients become essential!")
Comparison Summary:
Method Optimal L/G Evals
-----------------------------------------------
Gradient Descent 6.457 ~60
Golden Section 3.362 30
Note: For 1D problems, derivative-free methods work well.
For high-dimensional problems, gradients become essential!
7. Advanced: Jacobian and Hessian#
JAX can compute full Jacobians and Hessians for more sophisticated optimization.
from jax import jacfwd, jacrev, hessian
# Function with vector output
def absorber_outputs(L_G):
absorber = create_absorber(L_G)
_, _, info = absorber(flue_gas)
return jnp.array([
info['capture_efficiency'],
info['rich_loading'],
])
# Jacobian: derivatives of all outputs w.r.t. input
jac = jacfwd(absorber_outputs)
J = jac(3.0)
print("Jacobian at L/G = 3.0:")
print(f" d(capture_eff)/d(L/G) = {float(J[0]):.4f}")
print(f" d(rich_loading)/d(L/G) = {float(J[1]):.4f}")
Jacobian at L/G = 3.0:
d(capture_eff)/d(L/G) = 0.2580
d(rich_loading)/d(L/G) = -0.0000
# Hessian for second-order optimization (Newton's method)
hess_cost = hessian(capture_cost)
# Compute Hessian at different points to see curvature
print("Hessian (curvature) at different L/G values:")
print("-" * 50)
for L_G in [1.5, 2.0, 3.0, 4.0, 5.0]:
H = hess_cost(float(L_G))
cost = capture_cost(float(L_G))
g = grad_cost(float(L_G))
if jnp.abs(H) < 1e-6:
region = "flat"
elif H > 0:
region = "convex ↓"
else:
region = "concave ↑"
print(f"L/G={L_G:.1f}: cost=${float(cost):>6.1f}, grad={float(g):>7.2f}, H={float(H):>8.4f} ({region})")
print()
print("The Hessian shows curvature of the cost function.")
print("Positive H (convex) near a minimum enables Newton-type methods.")
Hessian (curvature) at different L/G values:
--------------------------------------------------
L/G=1.5: cost=$ 288.6, grad=-247.71, H=133.1527 (convex ↓)
L/G=2.0: cost=$ 181.4, grad=-181.14, H=133.1527 (convex ↓)
L/G=3.0: cost=$ 66.9, grad= -47.98, H=133.1527 (convex ↓)
L/G=4.0: cost=$ 66.5, grad= 5.00, H= 0.0000 (flat)
L/G=5.0: cost=$ 71.5, grad= 5.00, H= -0.0000 (flat)
The Hessian shows curvature of the cost function.
Positive H (convex) near a minimum enables Newton-type methods.
# Using scipy.optimize with JAX gradients (recommended for practical use)
from scipy.optimize import minimize
def cost_and_grad(x):
"""Return cost and gradient for scipy.optimize."""
x_jax = jnp.array(x[0])
cost = float(capture_cost(x_jax))
grad_val = float(grad_cost(x_jax))
return cost, jnp.array([grad_val])
# L-BFGS-B with bounds
result = minimize(
lambda x: cost_and_grad(x)[0],
x0=[3.0],
jac=lambda x: cost_and_grad(x)[1],
method='L-BFGS-B',
bounds=[(1.0, 10.0)]
)
print("Scipy L-BFGS-B Optimization:")
print(f" Optimal L/G: {result.x[0]:.3f}")
print(f" Minimum cost: ${result.fun:.2f}/hr")
print(f" Converged: {result.success}")
print(f" Function evals: {result.nfev}")
print()
print("Note: difflow_cc uses optimistix internally for root-finding.")
print("For optimization, scipy.optimize or optax work well with JAX gradients.")
Scipy L-BFGS-B Optimization:
Optimal L/G: 3.360
Minimum cost: $58.21/hr
Converged: True
Function evals: 9
Note: difflow_cc uses optimistix internally for root-finding.
For optimization, scipy.optimize or optax work well with JAX gradients.
8. Key Takeaways#
difflow enables gradient-based optimization of carbon capture processes
JAX’s
gradfunction computes exact derivatives through complex modelsGradient descent efficiently finds optimal operating conditions
vmapenables efficient sensitivity analysis across parameter rangesGradients scale well: Cost is ~2x function evaluation regardless of dimension
Advanced features:
Jacobians for multi-output systems
Hessians for Newton-type optimization
JIT compilation for speed
Next Steps#
Integrate with scipy.optimize for constrained optimization
Use JAX’s optax library for advanced optimizers (Adam, L-BFGS)
Combine with neural networks for hybrid models