Metadata-Version: 2.4
Name: difflow
Version: 0.1.0
Summary: Differentiable flowsheet framework for chemical processes
Project-URL: Homepage, https://github.com/jkitchin/differentiable-flowsheets
Project-URL: Documentation, https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/
Project-URL: Repository, https://github.com/jkitchin/differentiable-flowsheets
Project-URL: Issues, https://github.com/jkitchin/differentiable-flowsheets/issues
Author-email: "John R. Kitchin" <jkitchin@andrew.cmu.edu>
Maintainer-email: "John R. Kitchin" <jkitchin@andrew.cmu.edu>
License-Expression: MIT
License-File: LICENSE
Keywords: automatic differentiation,chemical engineering,jax,optimization,process simulation,technoeconomic analysis
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Chemistry
Requires-Python: >=3.10
Requires-Dist: blackjax>=1.3
Requires-Dist: diffrax>=0.6.0
Requires-Dist: flax>=0.10.7
Requires-Dist: gpjax>=0.12.2
Requires-Dist: jax>=0.4.0
Requires-Dist: jaxlib>=0.4.0
Requires-Dist: lineax>=0.0.7
Requires-Dist: matplotlib>=3.10.8
Requires-Dist: optimistix>=0.0.6
Requires-Dist: pyyaml>=6.0
Provides-Extra: all
Requires-Dist: ipycytoscape>=1.3.0; extra == 'all'
Requires-Dist: jaxopt>=0.8.0; extra == 'all'
Requires-Dist: jupyter-book<2.0; extra == 'all'
Requires-Dist: jupyter>=1.0.0; extra == 'all'
Requires-Dist: matplotlib>=3.5.0; extra == 'all'
Requires-Dist: networkx>=2.6.0; extra == 'all'
Requires-Dist: optax>=0.1.0; extra == 'all'
Requires-Dist: pytest-cov; extra == 'all'
Requires-Dist: pytest>=7.0; extra == 'all'
Provides-Extra: bio
Provides-Extra: cc
Provides-Extra: cuda11
Requires-Dist: jax[cuda11-local]>=0.4.0; extra == 'cuda11'
Provides-Extra: cuda12
Requires-Dist: jax[cuda12]>=0.4.0; extra == 'cuda12'
Provides-Extra: dev
Requires-Dist: jupyter-book<2.0; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: pytest>=7.0; extra == 'dev'
Provides-Extra: dwsim
Requires-Dist: pythonnet; extra == 'dwsim'
Provides-Extra: examples
Requires-Dist: jaxopt>=0.8.0; extra == 'examples'
Requires-Dist: jupyter>=1.0.0; extra == 'examples'
Requires-Dist: matplotlib>=3.5.0; extra == 'examples'
Requires-Dist: optax>=0.1.0; extra == 'examples'
Provides-Extra: gas
Provides-Extra: pyglenn
Requires-Dist: pyglenn; extra == 'pyglenn'
Provides-Extra: ree
Provides-Extra: visualization
Requires-Dist: ipycytoscape>=1.3.0; extra == 'visualization'
Requires-Dist: networkx>=2.6.0; extra == 'visualization'
Description-Content-Type: text/markdown

# difflow

[![Tests](https://github.com/jkitchin/differentiable-flowsheets/actions/workflows/test.yml/badge.svg)](https://github.com/jkitchin/differentiable-flowsheets/actions/workflows/test.yml)
[![PyPI](https://img.shields.io/pypi/v/difflow.svg)](https://pypi.org/project/difflow/)
[![Python](https://img.shields.io/pypi/pyversions/difflow.svg)](https://pypi.org/project/difflow/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/jkitchin/differentiable-flowsheets/blob/main/LICENSE)

**Differentiable Flowsheet Framework for Chemical Processes**

A JAX-based framework for building and optimizing chemical process flowsheets with automatic differentiation.

- **Source code:** <https://github.com/jkitchin/differentiable-flowsheets>
- **Documentation:** <https://kitchingroup.cheme.cmu.edu/differentiable-flowsheets/>
- **Issue tracker:** <https://github.com/jkitchin/differentiable-flowsheets/issues>

## Features

- **Fully Differentiable**: All unit operations and flowsheet calculations support automatic differentiation via JAX
- **Sensitivity Analysis**: Compute gradients of outputs with respect to any inputs, parameters, or operating conditions
- **Optimization Ready**: Use gradient-based optimization for process design, parameter estimation, and economic optimization
- **Modular Design**: Unit operations can be composed into complex flowsheets with recycle streams
- **Technoeconomic Analysis**: Comprehensive TEA module with equipment costs, operating costs, and profitability metrics (NPV, IRR, MSP)
- **Bio Manufacturing**: Specialized unit operations for biopharmaceutical processes (bioreactors, chromatography, filtration)
- **Gas Networks**: Steady-state gas transmission networks with a topology-computed sequential decomposition and differentiable tear solving

## ⚠️ ALPHA SOFTWARE

This project is under active development and not ready for production use. APIs may change without notice.  This notice will be removed when the project reaches stable release.

It is highly recommended that you confirm the equations and physical properties used in the models you make; these were generated by Claude. We have endeavored to ensure they seem reasonable, but cannot guarantee they are accurate in all cases. 

This package uses jax solvers (e.g. diffrax, optimistix, etc.) and does not rely on IPOPT or pounce, or pyomo / IDAES. It is a pure Python / jax focused package that was developed as a proof of concept. 

## LLM usage

Claude Code is heavily used to generate the code, examples and tests. This has allowed the project to develop faster than it can be used, and to develop more features than are immediately needed. This may mean there are modules that do not match the performance or output of other projects. You should perform your own diligence when using the code to ensure the library does what you expect it to. Ultimately this is a proof of concept in differentiable flowsheets that wouldn't be possible without Claude Code.

We regularly run all of the notebooks to ensure they run without errors, and review them to make sure the results make sense. We are happy to take issues and / or pull requests to fix problems. We also use Claude to review the code to look for issues.

We actually anticipate that Claude Code is used when using this library (See [CLAUDE.md](https://github.com/jkitchin/differentiable-flowsheets/blob/main/CLAUDE.md)). The library is large enough that it would take a long time to learn all the capabilities in addition to learning the nuances of differentiable programming. This repo provides all the information Claude needs to help you translate your flowsheet ideas into differentiable programs.

## Installation

```bash
# From PyPI
pip install difflow

# With examples and tutorials (includes matplotlib, jupyter)
pip install "difflow[examples]"

# Everything
pip install "difflow[all]"
```

For development, install from source:

```bash
git clone https://github.com/jkitchin/differentiable-flowsheets.git
cd differentiable-flowsheets
uv venv
uv pip install -e ".[dev]"

# Install everything
uv pip install -e ".[all]"
```

## Quick Start

```python
import jax.numpy as jnp
import jax

from difflow import (
    make_stream, get_flows,
    IdealThermo, SpeciesData,
    CSTR, CSTRParams,
)

# Define species
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)

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

# Create CSTR
stoich = jnp.array([[-1.0], [+1.0]])  # A → B
cstr_params = CSTRParams(
    V=jnp.array(1.0),
    rate_fn=rate_fn,
    stoich=stoich,
    rate_params={"A": jnp.array(1e6), "Ea": jnp.array(50000.0)},
    species_order=["A", "B"],
)
cstr = CSTR(cstr_params, thermo=thermo, mode="isothermal")

# Run simulation
inlet = make_stream({"A": 10.0, "B": 0.0}, T=300.0, P=101325.0)
outlet, info = cstr(inlet, T_spec=350.0)

print(f"Conversion: {info['conversion']['A']*100:.1f}%")

# Compute gradient of outlet B w.r.t. reactor volume
def outlet_B(V):
    params = CSTRParams(V=V, rate_fn=rate_fn, stoich=stoich,
                        rate_params={"A": jnp.array(1e6), "Ea": jnp.array(50000.0)},
                        species_order=["A", "B"])
    cstr = CSTR(params, thermo=thermo, mode="isothermal")
    outlet, _ = cstr(inlet, T_spec=350.0)
    return outlet["F_B"]

dFB_dV = jax.grad(outlet_B)(jnp.array(1.0))
print(f"dF_B/dV = {dFB_dV:.4f} mol/s per m³")
```

## Unit Operations

### CSTR (Continuous Stirred Tank Reactor)
- Multiple reactions with user-defined kinetics
- Isothermal, adiabatic, or specified heat duty modes
- Automatic material and energy balance solving

### PFR (Plug Flow Reactor)
- ODE-based design equation: dF/dV = stoich @ r
- Isothermal or adiabatic operation
- **GasPFR** variant for gas-phase reactions with:
  - Pressure drop (Ergun equation)
  - Variable volumetric flow from mole change
- RK4 integration via `lax.scan` (fully differentiable)

```python
from difflow import PFR, PFRParams, GasPFR, GasPFRParams

# Liquid-phase PFR
pfr = PFR(PFRParams(V=2.0, rate_fn=rate_fn, stoich=stoich,
                    rate_params=params, species_order=["A", "B"]))
outlet, info = pfr(inlet, T_spec=350.0)

# Gas-phase with pressure drop (A → 2B, mole increase)
gas_pfr = GasPFR(GasPFRParams(V=1.0, rate_fn=rate_fn, stoich=stoich,
                              rate_params=params, species_order=["A", "B"],
                              alpha=50000.0))  # Pressure drop parameter
outlet, info = gas_pfr(inlet, T_spec=500.0)
# info contains: conversion, profiles (V, F, T, P, Q), pressure_drop
```

### Flash Separator
- TP flash (temperature and pressure specified)
- Rachford-Rice equation for VLE
- Ideal thermodynamics (Raoult's law)

### Liquid-Liquid Extraction (LLE)
- **MultistageCascade**: Counter-current or co-current mixer-settler cascade
  - Kremser equation for stage calculations (differentiable in n_stages)
  - Continuous stage relaxation for optimization
- **DifferentialContactor**: Packed column extractor
  - HETP-based equilibrium model
  - Rate-based mass transfer model
- **Equilibrium Models**:
  - Distribution coefficients (K-values) with temperature dependence
  - NRTL activity coefficient model
  - UNIQUAC activity coefficient model

```python
from difflow import (
    MultistageCascade, CascadeParams,
    LLEEquilibrium, DistributionCoeffs,
)

# Define distribution coefficients for rare earth extraction
K_coeffs = DistributionCoeffs(
    species=("La", "Nd", "Dy"),
    K0=(0.5, 2.0, 8.0),  # K at reference temperature
)

equilibrium = LLEEquilibrium(
    solutes=["La", "Nd", "Dy"],
    aqueous_carrier="H2O",
    organic_carrier="Organic",
    K_coeffs=K_coeffs,
)

cascade = MultistageCascade(CascadeParams(
    n_stages=5,
    equilibrium=equilibrium,
    flow_config="counter_current",
))

raffinate, extract, info = cascade(feed_stream, solvent_stream)
```

### Utilities
- **Mixer**: Combine multiple streams
- **Splitter**: Split stream by fraction

### Fed-Batch Reactor
- General-purpose fed-batch (semi-batch) reactor for chemical reactions
- Time-varying feed addition with configurable feed profiles
- RK4 integration for batch dynamics
- Supports multiple reactions with user-defined kinetics

```python
from difflow import FedBatchReactor, FedBatchParams

# Fed-batch reactor with continuous reagent addition
def rate_fn(C, T, params):
    k = params["k0"] * jnp.exp(-params["Ea"] / (8.314 * T))
    return jnp.array([k * C["A"] * C["B"]])

params = FedBatchParams(
    V0=jnp.array(1.0),              # Initial volume (m³)
    rate_fn=rate_fn,
    stoich=jnp.array([[-1.0], [-1.0], [1.0]]),  # A + B → C
    rate_params={"k0": jnp.array(1e6), "Ea": jnp.array(50000.0)},
    species_order=["A", "B", "C"],
    t_final=jnp.array(3600.0),      # Batch time (s)
    n_steps=100,
)
reactor = FedBatchReactor(params)

# Feed profile: constant feed rate
feed = make_stream({"A": 0.0, "B": 1.0, "C": 0.0}, T=300.0, P=101325.0)
def feed_rate(t): return jnp.array(0.001)  # m³/s

final, info = reactor(initial_charge, feed, feed_rate, T_spec=350.0)
# info contains: conversion, profiles (t, V, C, T), yield
```

### Distillation Columns
- **ShortcutColumn**: Fenske-Underwood-Gilliland method for quick design estimates
  - Minimum stages (Fenske equation)
  - Minimum reflux ratio (Underwood equations)
  - Actual stages for given reflux (Gilliland correlation)
- **DistillationColumn**: Rigorous stage-by-stage calculation
  - MESH equations (Material, Equilibrium, Summation, Heat balance)
  - Supports partial/total condenser and reboiler

```python
from difflow import ShortcutColumn, ShortcutColumnParams

params = ShortcutColumnParams(
    species_order=["benzene", "toluene", "xylene"],
    light_key="benzene",
    heavy_key="toluene",
    x_D_LK=0.99,    # 99% benzene recovery in distillate
    x_B_HK=0.99,    # 99% toluene recovery in bottoms
)
column = ShortcutColumn(params, thermo=thermo)

distillate, bottoms, info = column(feed, R_ratio=1.5, q=1.0)
# info contains: N_min, R_min, N_actual, condenser_duty, reboiler_duty
```

### Heat Exchangers
- **Heater/Cooler**: Single-stream with utility (steam, cooling water)
  - Specified duty mode
  - Specified outlet temperature mode
  - Rating mode (given UA and utility temperature)
- **CounterCurrentHX**: Two-stream counter-current (shell-and-tube style)
- **CoCurrentHX**: Two-stream co-current (parallel flow)
- All use effectiveness-NTU method, fully differentiable

```python
from difflow import (
    Heater, HeaterParams,
    CounterCurrentHX, HeatExchangerParams,
    design_heat_exchanger,
)

# Single-stream heater with steam
heater = Heater(HeaterParams(T_out=400.0, Cp=75.0))
heated_feed, info = heater(cold_feed)
# info: Q, T_in, T_out, LMTD (if utility temp specified)

# Two-stream counter-current heat exchanger
hx = CounterCurrentHX(HeatExchangerParams(
    UA=2000.0,       # W/K
    Cp_hot=75.0,     # J/(mol·K)
    Cp_cold=80.0,
))
hot_out, cold_out, info = hx(hot_stream, cold_stream)
# info: Q, effectiveness, NTU, LMTD, approach temperature

# Design: calculate required area
result = design_heat_exchanger(
    Q=jnp.array(100000.0),  # 100 kW
    T_hot_in=jnp.array(450.0), T_hot_out=jnp.array(380.0),
    T_cold_in=jnp.array(300.0), T_cold_out=jnp.array(360.0),
    U=jnp.array(500.0),     # W/(m²·K)
)
print(f"Required area: {result['A']:.1f} m²")
```

## Bio Manufacturing Operations

The `difflow_bio` plugin provides specialized unit operations for biopharmaceutical manufacturing:

### Bioreactors
- **ContinuousBioreactor**: Chemostat with Monod kinetics
- **FedBatchBioreactor**: Fed-batch with substrate feeding strategy

```python
from difflow_bio import (
    ContinuousBioreactor, ContinuousBioreactorParams,
    FedBatchBioreactor, FedBatchBioreactorParams,
    monod_kinetics,
)

# Create a continuous bioreactor (chemostat)
params = ContinuousBioreactorParams(
    V=jnp.array(1000.0),           # Volume (L)
    mu_max=jnp.array(0.3),         # Maximum specific growth rate (1/h)
    Ks=jnp.array(0.5),             # Monod constant (g/L)
    Yxs=jnp.array(0.5),            # Biomass yield
    Yps=jnp.array(0.1),            # Product yield
    D=jnp.array(0.1),              # Dilution rate (1/h)
)
bioreactor = ContinuousBioreactor(params)
outlet = bioreactor(feed_stream)
```

### Downstream Processing
- **DiscStackCentrifuge**: Cell removal with Stokes' law separation
- **Ultrafiltration**: Protein concentration via TFF
- **Diafiltration**: Buffer exchange
- **ProteinAChromatography**: Affinity capture for mAb purification
- **IonExchangeChromatography**: Polishing step (bind-elute or flow-through)
- **SizeExclusionChromatography**: Aggregate removal

```python
from difflow_bio import (
    DiscStackCentrifuge, CentrifugeParams,
    Ultrafiltration, UFParams,
    ProteinAChromatography, ProAParams,
)

# Disc-stack centrifuge for cell removal
centrifuge = DiscStackCentrifuge(CentrifugeParams(
    sigma=jnp.array(5000.0),       # Sigma factor (m²)
    cell_diameter=jnp.array(15e-6), # Cell diameter (m)
))

# Protein A capture
proa = ProteinAChromatography(ProAParams(
    column_volume=jnp.array(10.0),  # CV (L)
    binding_capacity=jnp.array(40.0), # g mAb / L resin
    yield_factor=jnp.array(0.95),
))

# Ultrafiltration for concentration
uf = Ultrafiltration(UFParams(
    membrane_area=jnp.array(1.0),   # m²
    concentration_factor=jnp.array(10.0),
))
```

## Gas Transmission Networks

The `difflow_gas` plugin models steady-state gas transmission networks
as sequential-modular differentiable flowsheets. The sequential
decomposition of a meshed network (spanning tree, tear set, balance
schedule) is computed from the topology, so multi-loop networks need
no hand derivation:

```python
import difflow_gas as dg

net = dg.GasNetwork(
    arcs={
        "p1":  ("src", "a", "pipe"),
        "cs1": ("a", "b", "compressor"),
        "p2":  ("b", "c", "pipe"),
        "p3":  ("b", "d", "pipe"),
        "p4":  ("c", "d", "pipe"),          # closes a loop: the tear
    },
    beta={aid: dg.weymouth_beta(L, 0.6, 1e-4)
          for aid, L in [("p1", 20e3), ("p2", 40e3),
                         ("p3", 60e3), ("p4", 80e3)]},
    supply_kg_s={"src": 120.0, "c": -50.0, "d": -70.0},
)

fs, dec = dg.build_network_flowsheet(net, root="src",
                                     p_slack_pa=60e5,
                                     ratios={"cs1": 1.3})
streams = fs.solve(tol=1e-8)          # signed flows, Anderson tears
assert dg.residual_report(streams, net, dec).ok

# exact gradients through the converged tear iteration
obj = fs.make_objective_fn(
    lambda s: dg.total_compressor_power_w(s, dec, net.gas_temp_k))
dW_dr = jax.grad(obj)({"cs_cs1.ratio": 1.3})
```

Pipes, resistors, compressor stations, open valves, control valves and
short pipes are supported; see `docs/unit-operations-gas.md`.

## Thermodynamics

### Ideal Thermodynamics (for VLE)
- Ideal gas behavior
- Antoine equation for vapor pressures
- Polynomial Cp correlations
- Watson correlation for heat of vaporization

```python
SpeciesData(
    name="species_name",
    MW=100.0,                           # Molecular weight (g/mol)
    Cp_coeffs=(a, b, c, d),            # Cp = a + bT + cT² + dT³
    Hvap_coeffs=(A, n, Tc),            # Hvap = A(1 - T/Tc)^n
    antoine_coeffs=(A, B, C),          # log10(Psat) = A - B/(T+C)
    Hf=0.0,                            # Heat of formation (J/mol)
)
```

### Equations of State (for non-ideal VLE)
- **Peng-Robinson**: Cubic EOS for hydrocarbon and gas systems
- **Soave-Redlich-Kwong (SRK)**: Alternative cubic EOS
- Fugacity coefficients for both vapor and liquid phases
- Flash calculations with non-ideal K-values
- Binary interaction parameters (kij) support

```python
from difflow import PengRobinson, SRK, CriticalProperties, flash_TP_eos

# Define critical properties
props = {
    "methane": CriticalProperties(Tc=190.6, Pc=4.6e6, omega=0.011),
    "ethane": CriticalProperties(Tc=305.3, Pc=4.87e6, omega=0.099),
}

# Create EOS
eos = PengRobinson(props)
# or: eos = SRK(props)

# Compressibility factor
z = eos.compressibility_factor(T=300.0, P=1e6, z=[0.7, 0.3], phase="vapor")

# Fugacity coefficients
phi = eos.fugacity_coefficient(T=300.0, P=1e6, z=[0.7, 0.3], phase="vapor")

# Flash calculation
V_frac, x, y, K = flash_TP_eos(eos, z=[0.5, 0.5], T=250.0, P=2e6)
```

### Property Database
Built-in database with 55+ common species including critical properties and ideal thermo data:

```python
from difflow import (
    get_species_data, get_critical_props, list_species,
    get_alkanes, get_btex, get_common_solvents,
)

# Get species data for ideal thermodynamics
methanol = get_species_data("methanol")
thermo = IdealThermo({"methanol": methanol, "water": get_species_data("water")})

# Get critical properties for EOS
methane = get_critical_props("methane")
eos = PengRobinson({"methane": methane, "ethane": get_critical_props("ethane")})

# Convenience functions
alkanes = get_alkanes()           # methane through n-decane
btex = get_btex()                 # benzene, toluene, ethylbenzene, xylenes
solvents = get_common_solvents()  # water, methanol, ethanol, acetone, etc.

# Alias support: "CO2", "MeOH", "isopropanol", "IPA" all work
co2 = get_critical_props("CO2")

# List all available species
print(list_species())
```

### Activity Coefficient Models (for LLE)
- **NRTL**: Non-Random Two-Liquid model with temperature-dependent parameters
- **UNIQUAC**: Universal Quasi-Chemical model

## Technoeconomic Analysis (TEA)

The `difflow.economics` module provides comprehensive technoeconomic analysis capabilities, all fully differentiable for gradient-based optimization.

### Capital Costs

Equipment cost correlations with CEPCI escalation and installation factors:

```python
import difflow.economics as econ
import jax.numpy as jnp

# Equipment costs (2024 dollars)
reactor_cost = econ.reactor_cost(jnp.array(5.0), "cstr_jacketed")  # 5 m³
hx_cost = econ.heat_exchanger_cost(jnp.array(100.0), "shell_tube_floating")  # 100 m²
pump_cost = econ.pump_cost(jnp.array(10.0), "centrifugal_single")  # 10 kW

# Installed cost with Lang factor
installed = econ.installed_cost(reactor_cost, lang_factor=4.74)

# Total capital investment
tci = econ.total_capital_investment(
    purchased_equipment_cost=reactor_cost + hx_cost + pump_cost,
    lang_factor=4.74,
    working_capital_fraction=0.15,
)
```

Available equipment types:
- **Reactors**: CSTR (jacketed, coil), PFR, batch
- **Vessels**: Pressure vessels, storage tanks, flash drums
- **Heat Exchangers**: Shell-tube, double-pipe, plate-frame, air coolers
- **Columns**: Tray columns, packed columns
- **Pumps**: Centrifugal, reciprocating, gear
- **Compressors**: Centrifugal, reciprocating, screw
- **Separators**: Mixer-settlers, centrifuges, filters, extraction columns

### Utility Costs

```python
# Steam cost from heat duty
heating_cost = econ.steam_cost_from_duty(jnp.array(1e6), "medium_pressure")  # 1 MW

# Cooling water
cooling_cost = econ.cooling_water_cost(jnp.array(500e3))  # 500 kW

# Electricity
electricity_cost = econ.electricity_cost(jnp.array(100.0))  # 100 kW → $/hour
```

### Profitability Metrics

All metrics are JAX-differentiable:

```python
# Net Present Value
cash_flows = jnp.ones(20) * 500000  # $500k/year for 20 years
npv = econ.npv(cash_flows, jnp.array(0.10), jnp.array(2e6))  # 10% discount, $2M investment

# Internal Rate of Return
irr = econ.irr(cash_flows, jnp.array(2e6))

# Minimum Selling Price
msp = econ.minimum_selling_price(
    total_annual_cost=jnp.array(1e6),
    annual_production=jnp.array(50000.0),  # kg/year
)

# Annualized cost for optimization
tac = econ.annualized_cost(
    capital_cost=jnp.array(5e6),
    annual_opex=jnp.array(1e6),
    discount_rate=jnp.array(0.10),
    plant_life=jnp.array(20.0),
)
```

### Gradient-Based Economic Optimization

```python
import jax

def annual_profit(params):
    V, T = params[0], params[1]

    # Simulate process
    outlet, info = simulate_reactor(V, T)

    # Economics
    capex = econ.reactor_cost(V, "cstr_jacketed")
    installed = econ.installed_cost(capex)

    utility_cost = econ.cooling_water_cost(jnp.abs(info["Q"]))
    annual_utility = utility_cost * 8000 * 3600  # $/year

    revenue = outlet["F_product"] * product_price * 8000 * 3600

    crf = econ.capital_recovery_factor(jnp.array(0.10), jnp.array(20.0))
    return revenue - annual_utility - installed * crf

# Optimize design for maximum profit
grad_profit = jax.grad(annual_profit)
# Use gradient for optimization...
```

## Uncertainty Propagation

Leverage JAX's automatic differentiation for uncertainty quantification:

```python
from difflow import linear_propagation, monte_carlo_propagation, sensitivity_analysis

# Define a process model
def reactor_model(params):
    k = params['k0'] * jnp.exp(-params['Ea'] / (8.314 * params['T']))
    conversion = 1 - jnp.exp(-k * params['tau'])
    return conversion

nominal = {'k0': jnp.array(1e6), 'Ea': jnp.array(50000.0),
           'T': jnp.array(350.0), 'tau': jnp.array(100.0)}
uncertainties = {'k0': 1e5, 'Ea': 2000.0, 'T': 5.0, 'tau': 10.0}

# Linear (Jacobian-based) propagation - fast, first-order approximation
mean, std, info = linear_propagation(reactor_model, nominal, uncertainties)
print(f"Conversion: {mean:.3f} ± {std:.3f}")
print(f"Variance contributions: {info['variance_contributions']}")

# Monte Carlo propagation - handles non-linear models
mean_mc, std_mc, info_mc = monte_carlo_propagation(
    reactor_model, nominal, uncertainties, n_samples=10000
)

# Sensitivity analysis with gradient information
sens = sensitivity_analysis(reactor_model, nominal)
# Returns: gradient, elasticity (normalized sensitivity), variance contribution
```

Available functions:
- `linear_propagation()`: First-order Jacobian-based uncertainty propagation
- `monte_carlo_propagation()`: Parallel sampling using JAX vmap
- `sensitivity_analysis()`: Local gradient-based sensitivity with variance contributions
- `sobol_indices()`: Global sensitivity via Sobol sampling
- `propagate_covariance()`: Full covariance matrix propagation for correlated inputs

## Flowsheets with Recycles

```python
from difflow import Flowsheet, make_stream
from difflow.solvers import fixed_point_solve

# Define flowsheet iteration
def flowsheet_step(recycle_arr, args):
    # Unpack recycle, run units, return new recycle
    ...
    return new_recycle_arr

# Solve recycle loop
recycle = fixed_point_solve(
    flowsheet_step,
    initial_guess,
    args,
    max_iter=100,
    damping=0.5,
)
```

## Examples

Jupyter notebooks are in the `examples/` directory:

| Notebook | Description |
|----------|-------------|
| `00_cstr_pfr_basics.ipynb` | CSTR and PFR basics: conventional vs difflow |
| `01_cstr_flash_recycle.ipynb` | Complete flowsheet with CSTR, flash, and recycle |
| `02_cstr_sensitivity.ipynb` | Sensitivity analysis for CSTR parameters |
| `03_optimization.ipynb` | Gradient-based optimization problems |
| `04_rare_earth_extraction.ipynb` | Rare earth recovery using LLE |
| `05_technoeconomic_analysis.ipynb` | Comprehensive TEA with profit optimization |
| `06_uncertainty_propagation.ipynb` | Uncertainty propagation and sensitivity analysis |
| `07_heat_exchangers.ipynb` | Heat exchanger design, rating, and optimization |
| `10_dynamic_modeling.ipynb` | Dynamic simulation, DAE systems, diffrax backend |

```bash
# Launch Jupyter to explore examples
jupyter notebook examples/
```

## Tutorials

The `tutorials/` directory contains comprehensive JAX tutorials for differentiable programming:

| Notebook | Topics |
|----------|--------|
| `01_jax_fundamentals.ipynb` | grad, jit, vmap, pytrees, jacfwd/jacrev, VJP/JVP, HVP |
| `02_inverse_hessian_vector_products.ipynb` | IHVP, conjugate gradient, Newton-CG optimization |
| `02_optimization.ipynb` | Gradient descent, Newton, Adam, constrained optimization |
| `03_differential_equations.ipynb` | ODE solvers, parameter estimation, neural ODEs |
| `04_custom_derivatives.ipynb` | custom_vjp, custom_jvp, stop_gradient |
| `05_machine_learning.ipynb` | Neural networks from scratch, training loops |
| `06_gotchas.ipynb` | Common JAX pitfalls and how to avoid them |

## Key Design Decisions

1. **Streams as Dicts**: Simple `{"F_A": ..., "F_B": ..., "T": ..., "P": ...}` format that's a JAX pytree by default

2. **Property Database Available**: Built-in database with 55+ species, or define custom species data

3. **Function-Based Kinetics**: Maximum flexibility via `rate_fn(C, T, params) → rates`

4. **Unrolled Iteration**: Fixed-point solvers use `lax.scan` for automatic differentiability

5. **Continuous Relaxation**: Discrete parameters (like n_stages) can be relaxed to continuous values for optimization

## Dynamic Modeling

The `difflow.dynamic` module provides a unified framework for transient simulation of process units:

### Basic ODE Integration

```python
from difflow.dynamic import integrate
import jax.numpy as jnp

# Define any ODE system
def harmonic_oscillator(t, y):
    x, v = y[0], y[1]
    return jnp.array([v, -x])  # dx/dt = v, dv/dt = -x

result = integrate(
    harmonic_oscillator,
    y0=jnp.array([1.0, 0.0]),
    t_span=(0.0, 10.0),
    method="RK4",  # or "RK45", "Euler"
)
print(f"Final state: {result.y_final}")
print(f"Trajectory shape: {result.trajectory.y.shape}")
```

### Dynamic Unit Operations

```python
from difflow.dynamic import DynamicCSTR, integrate_unit
from difflow.streams import make_stream

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

# Create dynamic CSTR
cstr = DynamicCSTR(
    volume=1.0,
    rate_fn=rate_fn,
    stoich=jnp.array([[-1], [1]]),  # A -> B
    species_order=["A", "B"],
    rate_params={"k": 1e6, "Ea": 50000.0},
)

# Simulate startup from empty
inlet = make_stream({"A": 1.0, "B": 0.0}, T=350.0, P=101325.0)
result = integrate_unit(
    cstr,
    inputs={"inlet": inlet},
    t_span=(0.0, 1000.0),
    method="RK4",
)
```

### Dynamic Flowsheets

Connect multiple dynamic units for multi-unit transient simulation:

```python
from difflow.dynamic import DynamicFlowsheet, DynamicCSTR, DynamicTank

# Build flowsheet
fs = DynamicFlowsheet(species_order=["A", "B"])
fs.add_feed("feed", inlet_stream)
fs.add_unit(cstr, inlet_names=["feed"], outlet_names=["reactor_out"])
fs.add_unit(tank, inlet_names=["reactor_out"], outlet_names=["product"])

# Simulate entire flowsheet
result = fs.simulate(t_span=(0.0, 1000.0), method="RK4")
```

### DAE (Differential-Algebraic Equations)

For systems with algebraic constraints (e.g., VLE equilibrium):

```python
from difflow.dynamic import DynamicFlashDrum, integrate_dae

# Flash drum with VLE equilibrium constraint
flash = DynamicFlashDrum(
    volume=1.0,
    species_order=["A", "B"],
    K_values={"A": 2.0, "B": 0.5},  # K = y/x
)

result = integrate_dae(
    flash,
    inputs={"inlet": feed},
    t_span=(0.0, 100.0),
    method="RK4",
)
# result.x_final: differential states (moles)
# result.z_final: algebraic states (vapor fraction)
```

### Diffrax Backend (Advanced Solvers)

For stiff systems or when adaptive step control is needed:

```bash
pip install diffrax  # Optional dependency
```

```python
from difflow.dynamic import integrate

# Use diffrax solvers via method string
result = integrate(
    stiff_ode, y0, t_span,
    method="diffrax:kvaerno5",  # Implicit solver for stiff systems
    rtol=1e-6, atol=1e-8,
)

# Available solvers: dopri5, tsit5, dopri8, kvaerno3/4/5, euler, heun
# Default: tsit5 (recommended for most problems)
```

See `docs/dynamic-modeling.md` for complete documentation.

## Limitations

- Rigorous distillation column convergence can be sensitive to initial guesses
- Gradient explosion possible with many iterations (use damping)
- EOS flash limited to two-phase VLE (no three-phase VLLE yet)

## Future Work

- Three-phase (VLLE) flash calculations
- Extended bio operations (viral inactivation, sterile filtration)
- GPU acceleration for large flowsheets
- Integration with experiment databases (e.g., Cantera)

## Citation

If you use difflow in your work, please cite it. Machine-readable metadata is in
[`CITATION.cff`](https://github.com/jkitchin/differentiable-flowsheets/blob/main/CITATION.cff),
and GitHub's "Cite this repository" button will render BibTeX or APA from it.

## License

MIT
