WaterTAP Engine

⚠️ DEVELOPMENT STATUS: This project is under active development and is not yet production-ready. APIs, interfaces, and functionality may change without notice. Use at your own risk for evaluation and testing purposes only. Not recommended for production deployments.

A water treatment process simulation engine exposing WaterTAP capabilities through dual adapters for AI agent integration.

Motivation

WaterTAP provides sophisticated equation-oriented models for water treatment processes (RO, NF, crystallizers, evaporators, biological treatment), but the modeling workflow requires substantial domain expertise: DOF management, proper scaling, sequential initialization, and solver diagnostics.

WaterTAP Engine MCP inverts this paradigm by making natural language the primary interface. Instead of manually managing DOF, scaling, and initialization, engineers can work with a companion skill that orchestrates these operations:

"Evaluate hardness removal (IX or lime/soda softening) upstream of RO in this ZLD flowsheet: compare regenerant and chemical costs against natural gas savings from reduced evaporator thermal load, and calculate the IRR on incremental pre-treatment CAPEX."

This enables:

• Collapsed iteration cycles: Build -> solve -> diagnose -> patch -> re-solve without manual DOF tracking
• Explicit hygiene operations: Every fix, scale, and init is visible and controllable
• Structured diagnostics: DOF suggestions, scaling issues, and failure explanations surfaced directly to agents
• Reproducible sessions: Version-controlled flowsheet definitions with deterministic metadata

The goal is not to replace domain expertise, but to remove friction so engineers can focus on design decisions rather than solver mechanics.

WaterTAP vs QSDsan: Solver Paradigms

This server is a sibling to qsdsan-engine-mcp but uses a fundamentally different solver approach:

Aspect	WaterTAP (This Server)	QSDsan
Solver Type	Equation-oriented NLP (IPOPT)	ODE integration (solve_ivp)
Modeling Framework	Pyomo + IDAES	BioSTEAM
Simulation Mode	Steady-state algebraic	Dynamic ODE time-stepping
Simulation Speed	~seconds to minutes per solve	~milliseconds to seconds per run
Optimization	Native gradient-based (pyomo.environ.SolverFactory)	Requires external optimizer
DOF Management	Explicit - requires DOF=0 before solve	Implicit - feed-forward sequential
Scaling	Critical for convergence	Not required
Sensitivities	Automatic via NLP solver	Finite-difference approximation

Steady-State vs Dynamic Simulation

WaterTAP (Equation-Oriented, Steady-State)

• Solves a system of nonlinear algebraic equations: F(x) = 0
• All variables solved simultaneously to find steady-state directly
• Requires DOF = 0 (number of equations = number of unknowns)
• Ideal for design optimization: "What membrane area minimizes LCOW?"
• Supports native sensitivity analysis and gradient-based optimization
• Slower per solve but faster for optimization (exact gradients)

QSDsan (ODE Integration, Dynamic)

• Integrates differential equations forward in time: dx/dt = f(x,t)
• Simulates transient behavior from t=0 to t=T
• Variables update sequentially; no DOF constraint
• Ideal for dynamic behavior: "How does effluent quality change during startup?"
• Supports Monte Carlo sampling and scenario enumeration (fast per run)
• Faster per run but requires many runs for optimization (finite-difference)

When to Use WaterTAP

• Membrane process optimization: RO/NF with detailed concentration polarization, fouling models
• Gradient-based optimization: Minimize LCOW, maximize recovery, optimal pressure
• Rigorous thermodynamics: Multi-component phase equilibria, crystallization
• Design decisions: Optimal membrane area, stage configuration, pressure setpoints
• Sensitivity analysis: How does recovery change with feed salinity?

When to Use QSDsan

• Biological treatment dynamics: Activated sludge transients, digester startup/upset
• Rapid scenario enumeration: Monte Carlo uncertainty quantification, DOE
• Time-series analysis: Diurnal patterns, storm events, process disturbances
• Simpler mass balances: When steady-state is sufficient but dynamics are informative
• Training simulators: Interactive "what-if" exploration with fast feedback

Architecture: Dual Adapters

The engine exposes identical functionality through two adapters:

                    +-------------------------------------+
                    |       WaterTAP Engine Core          |
                    |  (Registries, Pipeline, Sessions)   |
                    +-----------------+-------------------+
                                      |
              +-----------------------+---------------------+
              |                       |                     |
              v                       v                     v
     +----------------+      +----------------+     +----------------+
     |   MCP Adapter  |      |   CLI Adapter  |     |  Python API    |
     |   (server.py)  |      |   (cli.py)     |     |  (direct use)  |
     +----------------+      +----------------+     +----------------+
              |                       |
              v                       v
     +----------------+      +----------------+
     |  MCP Clients   |      |  Agent Skills  |
     |  (Claude, etc) |      |  (Claude Code) |
     +----------------+      +----------------+

MCP Adapter (server.py)

For MCP-compatible clients (Claude Desktop, Cline, etc.):

python server.py

CLI Adapter (cli.py)

For CLI-based agent runtimes and Agent Skills:

python cli.py --help

Tool Surface

Session Management Tools

Tool	MCP	CLI	Description
create_session	create_session	session new	Create new flowsheet session
get_session	get_session	session show	Get session details
list_sessions	list_sessions	session list	List all sessions
delete_session	delete_session	session delete	Remove session

Registry & Discovery Tools

Tool	MCP	CLI	Description
list_units	list_units	units list	List available unit types
list_property_packages	list_property_packages	packages list	List property packages
list_translators	list_translators	translators list	List ASM/ADM translators
get_unit_spec	get_unit_spec	units spec	Get DOF, scaling, init hints

Flowsheet Construction Tools

Tool	MCP	CLI	Description
create_feed	create_feed	flowsheet add-feed	Add feed with state
create_unit	create_unit	flowsheet add-unit	Add unit operation
create_translator	create_translator	flowsheet add-translator	Add ASM/ADM translator
connect_ports	connect_ports	flowsheet connect	Wire units together
update_unit	update_unit	flowsheet update-unit	Modify unit parameters
delete_unit	delete_unit	flowsheet delete-unit	Remove unit
validate_flowsheet	validate_flowsheet	flowsheet validate	Pre-build validation

DOF Management Tools

Tool	MCP	CLI	Description
get_dof_status	get_dof_status	dof status	Get DOF count per unit
fix_variable	fix_variable	dof fix	Fix variable to value
unfix_variable	unfix_variable	dof unfix	Release variable
list_unfixed_vars	list_unfixed_vars	dof unfixed	Show unfixed variables

Scaling Tools

Tool	MCP	CLI	Description
set_scaling_factor	set_scaling_factor	scale set	Set explicit scaling
calculate_scaling_factors	calculate_scaling_factors	scale calculate	Run IDAES scaling
report_scaling_issues	report_scaling_issues	scale report	Find scaling problems

Solver Operations Tools

Tool	MCP	CLI	Description
initialize_unit	initialize_unit	init unit	Initialize single unit
initialize_flowsheet	initialize_flowsheet	init flowsheet	Sequential initialization
propagate_state	propagate_state	init propagate	State between ports
build_and_solve	build_and_solve	solve	Full hygiene pipeline
get_solve_status	get_solve_status	status	Job status
run_diagnostics	run_diagnostics	diagnose	DiagnosticsToolbox

Zero-Order Tools

Tool	MCP	CLI	Description
load_zo_parameters	load_zo_parameters	zo load	Load from database
list_zo_databases	list_zo_databases	zo databases	Available databases
get_zo_unit_parameters	get_zo_unit_parameters	zo params	Unit parameters

Costing Tools

Tool	MCP	CLI	Description
enable_costing	enable_costing	costing enable	Enable costing block
add_unit_costing	add_unit_costing	costing add-unit	Enable costing for unit
disable_unit_costing	disable_unit_costing	costing disable-unit	Disable unit costing
set_costing_parameters	set_costing_parameters	costing set-params	Set cost parameters
list_costed_units	list_costed_units	costing list	Show costing status
compute_costing	compute_costing	costing compute	Calculate LCOW, CapEx, OpEx
get_costing	get_costing	costing results	Get costing results

Results Tools

Tool	MCP	CLI	Description
get_results	get_results	results	Overall solve results
get_stream_results	get_stream_results	results streams	Stream tables
get_unit_results	get_unit_results	results units	Unit performance

Property Packages

13 supported property packages:

Package	Components	Use Case
SEAWATER	H2O + TDS	Seawater RO/NF, mass basis
NACL	H2O + NaCl	Brackish RO, mass basis
NACL_T_DEP	H2O + NaCl	Temperature-dependent thermal
WATER	H2O	Pure water, Liq + Vap
MCAS	Multi-ion	Ion-specific NF, molar basis
ZERO_ORDER	Database-driven	Simple ZO models
ASM1	13	Activated sludge (basic)
ASM2D	19	Activated sludge with bio-P
ASM3	-	Activated sludge extended
MODIFIED_ASM2D	-	Modified ASM2d
ADM1	63	Anaerobic digestion
MODIFIED_ADM1	-	Modified ADM1
ADM1_VAPOR	-	ADM1 vapor phase

Warning - Class-Name Collisions:

• NaClParameterBlock exists in both NaCl_prop_pack.py AND NaCl_T_dep_prop_pack.py
• WaterParameterBlock exists in both water_prop_pack.py AND zero_order_properties.py

Always use the full module path from the registry, not the class name alone.

Translators

Only ASM↔ADM translators exist in WaterTAP! The registry includes 8 translators:

Core Translators (4):

Source	Destination	Translator
ASM1	ADM1	Translator_ASM1_ADM1
ADM1	ASM1	Translator_ADM1_ASM1
ASM2D	ADM1	Translator_ASM2d_ADM1
ADM1	ASM2D	Translator_ADM1_ASM2d

Modified Model Translators (4):

Source	Destination	Translator
ModifiedASM2D	ADM1	Translator_ModifiedASM2d_ADM1
ADM1	ModifiedASM2D	Translator_ADM1_ModifiedASM2d
ASM2D	ModifiedADM1	Translator_ASM2d_ModifiedADM1
ModifiedADM1	ASM2D	Translator_ModifiedADM1_ASM2d

For non-biological flowsheets, use the same property package throughout.

Unit Registry

21+ unit operations available across categories:

• Membrane: ReverseOsmosis0D, ReverseOsmosis1D, Nanofiltration0D, NanofiltrationDSPMDE0D
• Zero-Order: NanofiltrationZO, UltraFiltrationZO, PumpZO, FeedZO
• Thermal: Evaporator, Condenser, Compressor, Crystallization
• Pumps/ERD: Pump, EnergyRecoveryDevice, PressureExchanger
• Utilities: Feed, Product, Mixer, Separator (from IDAES)
• Biological: (Use property package compatibility for ASM/ADM)

# List all units
python cli.py units list --json-out

# Filter by category
python cli.py units list --category membrane

Solver Hygiene Pipeline

The build_and_solve tool runs the full hygiene pipeline:

IDLE -> DOF_CHECK -> SCALING -> INITIALIZATION -> PRE_SOLVE_DIAGNOSTICS -> SOLVING
                                                                              |
COMPLETED <-- POST_SOLVE_DIAGNOSTICS <----------------------------------------+
    ^              | (if failed)                                              |
    +------- RELAXED_SOLVE ---------------------------------------------------+
                   | (if still fails)
                FAILED

Each stage provides structured output for diagnosis:

• DOF_CHECK: Underspecified/overspecified analysis with suggestions
• SCALING: Unscaled variables and constraint warnings
• INITIALIZATION: Per-unit init status and DOF after init
• DIAGNOSTICS: Constraint residuals, bound violations

Quick Start

Using CLI

# Create session
python cli.py session new --property-package SEAWATER --id my_ro

# Add feed and units
python cli.py flowsheet add-feed --session my_ro --flow 100 \
  --tds 35000 --temp 25 --pressure 1.01

python cli.py flowsheet add-unit --session my_ro --type ReverseOsmosis0D \
  --id RO1 --config '{"has_pressure_change": true}'

# Fix DOF
python cli.py dof fix --session my_ro --unit RO1 --var "A_comp[0, H2O]" --value 4.2e-12
python cli.py dof fix --session my_ro --unit RO1 --var "B_comp[0, TDS]" --value 3.5e-8
python cli.py dof fix --session my_ro --unit RO1 --var area --value 50
python cli.py dof fix --session my_ro --unit RO1 --var "permeate.pressure[0]" --value 101325

# Check DOF
python cli.py dof status --session my_ro

# Solve
python cli.py solve --session my_ro --run-full-pipeline

Using MCP

Configure in your MCP client (e.g., Claude Desktop config.json):

{
  "mcpServers": {
    "watertap-engine": {
      "command": "python",
      "args": ["/path/to/watertap-engine-mcp/server.py"]
    }
  }
}

Then use natural language with the companion skill:

"Evaluate hardness removal upstream of RO in this ZLD flowsheet and calculate the IRR on incremental pre-treatment CAPEX based on reduced evaporator thermal load."

Installation

# Clone repository
git clone https://github.com/puran-water/watertap-engine-mcp.git
cd watertap-engine-mcp

# Create virtual environment
python -m venv venv
source venv/bin/activate  # or venv\Scripts\activate on Windows

# Install dependencies
pip install -r requirements.txt

# For full WaterTAP support (optional, for solve operations):
pip install watertap idaes-pse pyomo

Dependencies

• Python 3.10+
• WaterTAP 1.0+ (optional, for solve)
• IDAES-PSE 2.0+ (optional, for solve)
• FastMCP (for MCP adapter)
• Typer (for CLI adapter)

File Structure

watertap-engine-mcp/
├── server.py              # MCP Adapter (FastMCP) - 57 tools
├── cli.py                 # CLI Adapter (typer)
├── worker.py              # Background job worker
├── core/
│   ├── property_registry.py    # 13 property packages
│   ├── translator_registry.py  # ASM/ADM translators
│   ├── unit_registry.py        # Unit specs with DOF, scaling
│   ├── water_state.py          # Feed state abstraction
│   └── session.py              # Session management
├── solver/
│   ├── pipeline.py             # Hygiene pipeline state machine
│   ├── dof_resolver.py         # DOF analysis
│   ├── scaler.py               # Scaling tools
│   ├── initializer.py          # Sequential initialization
│   ├── diagnostics.py          # DiagnosticsToolbox wrapper
│   └── recovery.py             # Failure recovery
├── utils/
│   ├── model_builder.py        # Session -> Pyomo model
│   ├── auto_translator.py      # Translator insertion
│   ├── job_manager.py          # Background jobs
│   ├── state_translator.py     # Feed state conversion
│   └── topo_sort.py            # Topological sort
├── templates/                  # Pre-built flowsheet templates
│   ├── ro_train.py
│   ├── nf_softening.py
│   └── mvc_crystallizer.py
├── jobs/                       # Session/job persistence (runtime)
└── tests/                      # 373 tests

Companion Skill

The watertap-flowsheet-builder skill (at ~/skills/watertap-skill/) provides domain intelligence:

• DOF suggestions with typical values for each unit type
• Property package selection guidance
• Translator selection for biological chains
• Failure diagnosis patterns
• Workflow orchestration for common flowsheets

The skill orchestrates the atomic server tools - server provides explicit operations, skill provides intelligence.

Testing

# Run all tests
pytest tests/ -v

# Skip slow tests
pytest tests/ -v -m "not slow"

License

MIT

Acknowledgments

Built on WaterTAP by the National Alliance for Water Innovation (NAWI).