# Difflow Documentation

**Difflow** is a JAX-based framework for building and optimizing chemical process flowsheets with automatic differentiation support. This enables gradient-based optimization of process designs, sensitivity analysis, and uncertainty propagation.

## Key Features

- **Automatic Differentiation**: All unit operations and calculations are fully differentiable using JAX
- **Comprehensive Unit Operations**: Reactors, separators, heat exchangers, distillation columns, and more
- **Bio-Manufacturing Support**: Specialized operations for bioreactors, centrifugation, filtration, and chromatography
- **Thermodynamic Models**: Ideal thermodynamics and cubic equations of state (Peng-Robinson, SRK)
- **Technoeconomic Analysis**: Equipment costs, operating costs, and profitability metrics
- **Extensible Architecture**: Plugin system for custom unit operations

## Documentation Contents

### Core Modules

| Document | Description |
|----------|-------------|
| [Getting Started](getting-started.md) | Installation, quick start, and basic examples |
| [Unit Operations - Chemical](unit-operations-chemical.md) | Reactors, separators, heat exchangers, distillation |
| [Unit Operations - Bio](unit-operations-bio.md) | Bioreactors, centrifugation, filtration, chromatography |
| [Unit Operations - REE](unit-operations-ree.md) | Rare earth element extraction, scrubbing, stripping, precipitation |
| [Thermodynamics](thermodynamics.md) | Property calculations, equations of state, databases |
| [Technoeconomics](technoeconomics.md) | Capital costs, operating costs, profitability analysis |
| [Streams and Flowsheets](streams-and-flowsheets.md) | Stream handling, flowsheet solver, recycle calculations |
| [Convergence and Initialization](convergence.md) | Tear-stream guesses, the three acceleration methods, the traced fallback, and what to do when a solve will not converge |
| [Equation-Oriented Solver](eo-solver.md) | Simultaneous solution of the whole flowsheet, SM initialization, adding EO support to a unit |
| [Dynamic Modeling](dynamic-modeling.md) | Transient simulation, ODE/DAE integration, diffrax backend |
| [Moving-Horizon Estimation](moving-horizon-estimation.md) | Constrained MHE over dynamic flowsheets, EKF baseline, delayed and multi-rate data, joint parameter estimation |
| [Operability Screening](operability.md) | Steady-state controllability from AD gains: RGA, singular values, disturbance rejection |
| [Flexibility Analysis](flexibility.md) | Feasibility function, flexibility index, and the feed-vs-parameter uncertainty split |
| [Stochastic Programming](stochastic.md) | Two-stage design under a parameter distribution: scenario sampling, CVaR and chance constraints, VSS and EVPI |
| [Experiment Design and Identifiability](experiment-design.md) | Fisher-information D/A/E-optimal run selection, predicted confidence intervals, structural identifiability |
| [External Solvers](external-solvers.md) | Bridging flowsheets to pounce (NLP, post-optimal sensitivity) and discopt (implicit residual blocks) |
| [Solvers and Utilities](solvers-and-utilities.md) | Numerical methods, uncertainty propagation |
| [Data Reconciliation](data-reconciliation.md) | Constrained least squares on noisy plant data, gross error detection, observability |
| [Delta-Base Planning](planning.md) | AD-generated delta vectors, trust-region LP/MILP planning, sensitivity of the plan |
| [Aspen PIMS Integration](pims-integration.md) | Design proposal: shipping difflow delta vectors into a PIMS planning model (nothing implemented yet) |

## Architecture Overview

```
difflow/
├── streams.py          # Stream data structures
├── thermo.py           # Ideal thermodynamics
├── eos.py              # Cubic equations of state
├── database.py         # Species property database
├── solvers.py          # Numerical solvers
├── flowsheet.py        # Flowsheet management
├── reconciliation/     # Data reconciliation and gross error detection
├── planning/           # Delta-base planning (LP/MILP from flowsheets)
├── uncertainty.py      # Uncertainty propagation
├── cantera_import.py   # Cantera data import
├── pyglenn_import.py   # NASA Glenn (pyglenn) thermo import
├── dwsim_import.py     # DWSIM thermo import (pythonnet; DWSIM 9.0.5)
├── units/              # Unit operations
│   ├── cstr.py         # CSTR reactors
│   ├── pfr.py          # PFR reactors
│   ├── fed_batch.py    # Fed-batch reactors
│   ├── flash.py        # Flash, mixer, splitter
│   ├── distillation.py # Distillation columns
│   ├── heat_exchanger.py # Heat exchangers
│   └── lle.py          # Liquid-liquid extraction
├── dynamic/            # Dynamic (transient) simulation
│   ├── state.py        # State variable specification
│   ├── base.py         # DynamicUnit protocol
│   ├── integrators.py  # ODE integration (RK4, RK45)
│   ├── flowsheet.py    # Multi-unit dynamic flowsheets
│   ├── dae.py          # DAE systems and Newton solver
│   └── diffrax_backend.py # Advanced solvers via diffrax
└── economics/          # Technoeconomic analysis
    ├── capital.py      # Equipment costs
    ├── utilities.py    # Utility costs
    ├── opex.py         # Operating costs
    ├── profitability.py # Financial metrics
    └── indices.py      # Cost indices

difflow_bio/
└── units/              # Bio-manufacturing operations
    ├── bioreactors.py  # Bioreactors
    ├── centrifuge.py   # Centrifugation
    ├── filtration.py   # Membrane filtration
    └── chromatography.py # Chromatography
```

## Quick Example

```python
import jax.numpy as jnp
from difflow.streams import make_stream
from difflow.units.cstr import CSTR, CSTRParams
from difflow.thermo import IdealThermo, SpeciesData

# Define species
species_data = {
    'A': SpeciesData(name='A', MW=50.0, Cp_coeffs=(30.0, 0.01, 0.0, 0.0),
                    Hvap_coeffs=(35000.0, 0.38, 500.0), antoine_coeffs=(10.0, 1500.0, -40.0),
                    Hf=0.0, Tref=298.15),
    'B': SpeciesData(name='B', MW=100.0, Cp_coeffs=(40.0, 0.02, 0.0, 0.0),
                    Hvap_coeffs=(40000.0, 0.38, 550.0), antoine_coeffs=(10.5, 1800.0, -50.0),
                    Hf=-50000.0, Tref=298.15)
}
thermo = IdealThermo(species_data)

# Create inlet stream
inlet = make_stream({'A': 1.0, 'B': 0.0}, T=350.0, P=101325.0)

# Define CSTR parameters
def rate_fn(C, T, rp):
    k = rp['k_ref'] * jnp.exp(-rp['E_a'] / 8.314 * (1.0 / T - 1.0 / rp['T_ref']))
    return jnp.array([k * C['A']])

params = CSTRParams(
    V=1.0,  # m³
    rate_fn=rate_fn,
    stoich=jnp.array([[-1.0], [1.0]]),  # A -> B
    rate_params={'k_ref': 0.1, 'E_a': 50000.0, 'T_ref': 350.0},
    species_order=['A', 'B'],
    dH_rxn=jnp.array([-50000.0]),
    molar_density=10.0,  # mol/m³, sets the concentration basis
)

# Create and run CSTR
cstr = CSTR(params, thermo)
outlet, info = cstr(inlet, T_spec=350.0)

print(f"Conversion: {float(info['conversion']['A']):.2%}")
print(f"Heat duty: {float(info['Q']):.2f} W")
```

## Design Philosophy

### JAX-First Approach

Every calculation in Difflow is designed to be compatible with JAX automatic differentiation:

```python
import jax

# Gradient of product flow with respect to reactor temperature
grad_fn = jax.grad(lambda T_in: cstr(make_stream({'A': 1.0, 'B': 0.0}, T=T_in, P=101325.0), T_spec=T_in)[0]['F_B'])
sensitivity = grad_fn(350.0)
```

### Consistent Interface

All unit operations follow the same calling convention:

```python
outlet, info = cstr(inlet, T_spec=350.0)  # any unit: (inlet, **operating_params)
```

Where:
- `inlet` is a Stream dictionary
- `outlet` is the output Stream
- `info` contains additional results (heat duties, conversions, etc.)

### Pytree Compatibility

Streams are JAX pytrees, allowing seamless use with `jax.vmap`, `jax.jit`, and other transformations.

## Requirements

- Python >= 3.9
- JAX >= 0.4.0
- NumPy
- PyYAML (for Cantera import)
- pyglenn (optional, for NASA Glenn thermo import)

## License

MIT License