WaterTAP Engine MCP

⚠️ 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.

Related MCP server: SWMM-MCP

WaterTAP vs QSDsan: Solver Paradigms

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

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

Registry & Discovery Tools

Flowsheet Construction Tools

DOF Management Tools

Scaling Tools

Solver Operations Tools

Zero-Order Tools

Costing Tools

Results Tools

Property Packages

13 supported property packages:

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):

Modified Model Translators (4):

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).