OHMind: Multi-Agent HEM Design System

🎬 OHMind Demo

📺 YouTube | 📺 Bilibili

---

OHMind couples large-language-model (LLM) agents with physics-based simulators to accelerate hydroxide exchange membrane (HEM) discovery. The codebase accompanies the manuscript "Autonomous multi-agent AI accelerates hydroxide exchange membrane discovery through physics-grounded inverse design" and reproduces every experiment described therein.

---

Table of Contents

• Why OHMind?
• System Architecture
• Key Features
• Repository Structure
• Installation
• Quick Start
• OHMind CLI
• Configuration
• Usage Examples
• Multi-Agent Workflow
• MCP Servers
• Core Modules
• Programmatic Access
• Troubleshooting
• Data & Model Availability
• Citation
• Contributing
• Contact & Support

---

Why OHMind?

• 🤖 Agentic orchestration – LangGraph supervisor routes user intents to specialized HEM, chemistry, QM, MD, Multiwfn, RAG, web-search, and validation agents over Model Context Protocol (MCP).
• 🧬 Generative pipeline at scale – a Junction Tree VAE model trained on 223k cations achieves 100% validity and 99.6% uniqueness.
• ⚡ Latent PSO – multi-objective optimization balances conductivity, swelling, stability, and synthetic feasibility.
• 🔬 Physics verification – ORCA DFT tools confirm LUMO/stability trends, while GROMACS workflows quantify morphology and hydroxide diffusion.
• 🌐 Production-ready UI – A modern web interface with multi-workflow support, authentication, and persistent storage.
• 🖥️ Terminal UI & Web Deploy – A modern Textual-based CLI that can run in terminal or be deployed as a web application via textual-serve.
• 🔌 MCP Integration – Five specialized MCP servers (Chemistry, HEM, ORCA, GROMACS, Multiwfn) provide tool access to AI agents.

---

Key Features

🧠 Multi-Agent System

• Supervisor Agent: Intelligent routing to specialized agents
• HEM Agent: Expert in membrane design and PSO optimization
• Chemistry Agent: General molecular operations and cheminformatics
• QM Agent: DFT calculations via ORCA
• MD Agent: Molecular dynamics simulations via GROMACS
• Multiwfn Agent: Wavefunction and orbital analysis
• RAG Agent: Literature retrieval from embedded knowledge base
• Web Search Agent: Real-time scientific information retrieval
• Validation Agent: Human-in-the-loop governance for expensive operations

🔬 Generative AI Pipeline

• Junction Tree VAE: 56-dimensional latent space trained on 223,000 cations
• 100% validity: All generated structures are chemically valid
• 99.6% uniqueness: Minimal duplication in generated candidates
• Vocabulary-based: Molecular fragment vocabulary for robust generation

⚡ Optimization Engine

• Particle Swarm Optimization: Efficient exploration of chemical space
• Multi-objective: Simultaneous optimization of:
  • Ionic conductivity (EC)
  • Water uptake (EWU)
  • Swelling ratio (ESR)
  • Alkaline stability
  • Synthetic accessibility
• Latent space PSO: Direct optimization in VAE latent space
• Desirability functions: Customizable objective balancing

🔌 Model Context Protocol (MCP)

• 5 specialized servers providing 75 tools
• Dual transport: stdio for local clients, HTTP for remote access
• Async execution: Long-running tasks don't block the UI
• Job management: Start, monitor, and terminate computational jobs
• Real-time logging: Monitor optimization and simulation progress

---

Installation

Prerequisites

• Operating System: Linux or macOS (tested on Ubuntu 20.04+)
• Python: 3.10 or 3.11
• Conda: Anaconda or Miniconda
• GPU: CUDA-capable GPU recommended (for VAE training and MD)
• Docker: For PostgreSQL and MinIO services
• Memory: 32GB RAM minimum, 64GB+ recommended
• Disk: 100GB+ free space for models, databases, and results

External Software (Optional but Recommended)

Software	Purpose	Installation
ORCA 6.0+	Quantum chemistry calculations	ORCA Forum
GROMACS 2025+	Molecular dynamics simulations	conda install -c conda-forge gromacs
Multiwfn 3.8+	Wavefunction analysis	Multiwfn
xTB	Semi-empirical QM	conda install -c conda-forge xtb
Qdrant	Vector database for RAG	Docker:docker run -p 6333:6333 qdrant/qdrant

Step 1: Clone the Repository

git clone https://github.com/lunyang/OHMind.git
cd OHMind

Step 2: Create Conda Environment

# Create environment from environment.yml
conda env create -f environment.yml

# Activate the environment
conda activate OHMind

Step 3: Install Additional Dependencies

# Install OHMind as a package (editable mode)
pip install -e .

# Verify installation
python -c "import OHMind; print(OHMind.__version__)"

Step 4: Set Up Environment Variables

Create a .env file in the project root:

# Copy example configuration
cp .env.example .env

# Edit with your settings
nano .env

Minimum required variables:

# LLM Provider (choose one)
OPENAI_API_KEY=your_openai_key_here
# OR
AZURE_OPENAI_API_KEY=your_azure_key_here
AZURE_OPENAI_ENDPOINT=your_azure_endpoint_here
# OR
OPENAI_COMPATIBLE_API_KEY=your_compatible_key_here
OPENAI_COMPATIBLE_BASE_URL=your_compatible_base_url_here

# Workspace (unified directory for all results)
OHMind_workspace=/path/to/your/workspace  # e.g., ./OHMind_workspace

# External Software Paths (if installed)
OHMind_ORCA=/path/to/orca  # e.g., /home/user/orca_6_1_0/orca
MULTIWFN_PATH=/path/to/Multiwfn  # e.g., /home/user/Multiwfn38/Multiwfn

# Optional: Vector database
QDRANT_URL=http://localhost:6333
TAVILY_API_KEY=your_tavily_key_for_web_search  # Optional: for web search agent

Step 5: Start Vector Database Services

   cd OHMind_ui

# Start PostgreSQL and MinIO
   docker compose -f docker-compose-db-only.yml up -d

# Verify services are running
docker compose -f docker-compose-db-only.yml ps

---

Quick Start

Option 1: Use the Startup Script (Recommended)

# From project root
chmod +x start_OHMind.sh
./start_OHMind.sh

This script will:

1. Start the FastAPI backend on port 8005
2. Launch 5 MCP servers (HTTP mode on ports 8101-8105)
3. Start the Chainlit UI on port 8000
4. Set up all environment variables

Access the UI:

http://localhost:8000

Default credentials:

• Username: admin
• Password: admin

The web ui interface:

Option 2: Manual Startup

Terminal 1: Start Backend

cd OHMind
export PYTHONPATH=$(pwd)
python OHMind_backend.py

Terminal 2: Start UI

cd OHMind_ui
chainlit run OHMind_frontend.py --host 0.0.0.0 --port 8000

First Steps in the UI

1. Select a workflow: Choose "HEM Multi-Agent" for full capabilities

2. Configure backend: In chat settings, set backend URL to http://localhost:8005

3. Add MCP servers (sidebar → MCP → Add server):

   • OHMind-Chem: Chemistry operations
   • OHMind-HEMD: HEM optimization
   • OHMind-ORCA: Quantum chemistry (if ORCA installed)
   • OHMind-GROMACS: MD simulations (if GROMACS installed)
   • OHMind-Multiwfn: Wavefunction analysis (if Multiwfn installed)

4. Try a sample query:

   List all available HEM backbones and cation types.
   

5. Run an optimization:

   Design new piperidinium-based cations for PBF_BB_1 backbone 
   optimizing for multi-objective HEM performance. Use 100 particles 
   and 5 steps for a quick test.
   

   

---

OHMind CLI

OHMind CLI (OHMind_cli/) provides a modern terminal user interface (TUI) built with Textual, offering an alternative to the web-based Chainlit UI. It can also be deployed as a web application using textual-serve.

Features

• 🖥️ Modern TUI: Rich terminal interface with syntax highlighting, streaming responses, and keyboard shortcuts
• 🌐 Web Deployment: Deploy as a web application accessible from any browser
• 🎨 Custom Themes: Scientific-inspired color themes (Scientific, Orbital, Synthesis, Quantum)
• 📤 Export: Export chat history to Markdown, HTML, or screenshots
• ⌨️ Keyboard Shortcuts: Efficient navigation with Ctrl+L (clear), Ctrl+T (tools), Ctrl+A (agents), etc.

Running the CLI

Option 1: Using the startup script (recommended)

# Run as terminal UI
./start_OHMind_full.sh

# Run with debug mode
./start_OHMind_full.sh --debug

Option 2: Direct Python execution

# Activate conda environment
conda activate OHMind

# Run the CLI
python -m OHMind_cli

# With debug mode
python -m OHMind_cli --debug

Web Deployment

Deploy OHMind CLI as a web application accessible from any browser:

# Deploy on default port 8000
./start_OHMind_full.sh deploy

# Deploy on custom port
./start_OHMind_full.sh deploy --port 9000

# Deploy with custom host and port
./start_OHMind_full.sh deploy --host 0.0.0.0 --port 8080

# Deploy with custom public URL (for reverse proxy setups)
./start_OHMind_full.sh deploy --public-url https://ohmind.example.com

Direct Python deployment:

python -m OHMind_cli deploy --host 0.0.0.0 --port 8000

Access the web app:

• Local: http://localhost:8000
• Network: http://<your-ip>:8000

Screenshot of the OHMind Cli app:

CLI Commands

Command	Description
/help or /?	Show help screen
/clear	Clear conversation history
/agents	Show available agents
/tools	Show available MCP tools
/export	Export chat history
/export md	Export to Markdown
/export html	Export to HTML
/new	Start new conversation
/debug	Toggle debug mode
/quit or /q	Quit application

---

Usage Examples

Example 1: HEM Optimization (Quick Test)

Natural language query in UI:

Design new piperidinium cations for PBF_BB_1 backbone 
optimizing for multi-objective HEM performance. 
Use 100 particles and 3 steps for a quick test.

What happens:

1. HEM agent validates backbone and cation type
2. Validates optimization configuration
3. Requests user approval (expensive operation)
4. Launches PSO optimization in background
5. Saves results to OHMind_workspace/HEM/
6. Returns top candidates with predicted properties

Example 2: QM Analysis

Query:

For the cation SMILES "C[N+]1(C)CCCCC1", calculate the LUMO energy 
using ORCA and explain what it implies for alkaline stability.

What happens:

1. Chemistry agent validates SMILES
2. QM agent converts SMILES to 3D coordinates
3. Runs ORCA geometry optimization + single-point calculation
4. Extracts LUMO energy and other descriptors
5. Interprets results for alkaline stability

Example 3: MD Simulation

Query:

Build a simple AEM polymer system with 10 chains (DP=25) 
equilibrated at 400 K. Run a GROMACS simulation to estimate 
water uptake and ionic conductivity.

What happens:

1. MD agent prepares polymer structure
2. Parameterizes using appropriate force fields
3. Solvates system with water + counter-ions
4. Runs energy minimization → NVT → NPT → production MD
5. Analyzes trajectories for water uptake and ion diffusion

Example 4: Literature-Guided Design

Query:

Retrieve recent literature on cation designs for high alkaline 
stability in HEMs. Propose 5 new candidate cations that follow 
those design principles.

What happens:

1. RAG agent searches embedded literature database
2. Web search agent finds recent papers (if Tavily key provided)
3. Synthesizes design principles from literature
4. Chemistry/HEM agents propose new structures
5. Validates candidates against design principles

---

Multi-Agent Workflow

Agent Roles and Capabilities

Agent	Role	Key Tools	Typical Tasks
Supervisor	Routes requests to appropriate agents	-	Task decomposition, agent selection, result synthesis
HEM Agent	HEM design and optimization expert	list_hem_backbones, optimize_hem_design, check_optimization_results	PSO optimization, backbone selection, cation design
Chemistry Agent	General chemistry operations	smiles_canonicalization, molecule_similarity, functional_groups	SMILES validation, property calculation, structure manipulation
QM Agent	Quantum chemistry calculations	geometry_optimization, frequency_calculation, polymer_reactivity	DFT calculations, HOMO/LUMO analysis, stability prediction
MD Agent	Molecular dynamics simulations	run_complete_iem_workflow, calculate_msd_tool, analyze_energy_tool	System building, MD simulation, trajectory analysis
Multiwfn Agent	Wavefunction analysis	orbital_analysis, population_analysis, weak_interaction_analysis	Charge distribution, orbital visualization, bonding analysis
RAG Agent	Literature retrieval	search_knowledge_base	Finding relevant papers, extracting design principles
Web Search Agent	Real-time information	web_search	Current protocols, recent publications, property data
Validation Agent	Human-in-the-loop control	request_approval	Expensive operation approval, job termination confirmation

---

MCP Servers

OHMind provides 5 specialized MCP servers with 75 tools total:

Server	Tools	Port (HTTP)	Description
Chem	17	8101	Chemistry operations, SMILES conversion, molecular properties
HEM	7	8102	HEM optimization, PSO configuration, job management
ORCA	10	8103	Quantum chemistry, DFT, HOMO/LUMO, frequencies
GROMACS	25	8105	MD simulations, system building, trajectory analysis
Multiwfn	16	8104	Wavefunction analysis, orbital visualization, charge analysis

MCP Server Details

1. Chemistry Server

17 tools including:

• SMILES validation and canonicalization
• Name ↔ SMILES ↔ IUPAC conversion
• Molecular weight and formula calculation
• Functional group detection
• Structural similarity (Tanimoto)
• 2D structure image generation
• SELFIES conversion
• PubChem integration
• Web search for chemical data

2. HEM Server

7 tools:

1. list_hem_backbones - List available backbones (PBF_BB_1, PP_BB_1, etc.)
2. list_cation_types - List cation families (piperidinium, imidazolium, etc.)
3. validate_optimization_config - Pre-flight validation
4. optimize_hem_design - Launch PSO optimization
5. check_optimization_results - View results
6. show_optimization_logs - Monitor progress
7. kill_optimization_job - Terminate jobs

Supported backbones:

• PBF_BB_1, PBF_BB_2: Polybenzofuran backbones
• PP_BB_1, PP_BB_2: Polyphenylene backbones
• PSU_BB_1: Polysulfone backbone
• PPO_BB_1: Polyphenylene oxide backbone

Supported cations:

• random: General N+ containing structures
• tetraalkylammonium: Quaternary ammonium
• imidazolium: 5-membered ring
• benzimidazolium: Fused ring system
• guanidinium: Guanidinium group
• piperidinium: 6-membered ring
• spiro_undecane: Spirocyclic structure

Optimization properties:

• ec: Effective OH⁻ conductivity (2.4-6.0 range)
• ewu: Effective water uptake (0.1-5.0 range)
• esr: Effective swelling ratio (2.8-6.1, lower better)
• multi: Multi-objective (balances all three)

3. ORCA Server

10 tools including:

• Single-point energy calculation
• Geometry optimization (OPT, TIGHTOPT, VERYTIGHTOPT)
• Frequency calculation (IR, thermochemistry)
• SMILES to XYZ conversion
• Proton affinity and pKa estimation
• Binding energy with BSSE correction
• Ionic solvation energy
• Charge analysis (Mulliken, Löwdin, Hirshfeld)
• Transition state search
• NMR chemical shift prediction
• Polymer reactivity descriptors (HOMO/LUMO)

Supported methods:

• DFT: B3LYP, PBE0, M06-2X, wB97X-D3, etc.
• Wavefunction: HF, MP2, CCSD(T)
• Semi-empirical: HF-3c, PBEh-3c

Dispersion corrections:

• D3BJ (Grimme's D3 with BJ damping)
• D3, D4
• None

4. GROMACS Server

25 tools including:

Building & Parameterization:

• Create polymer from SMILES
• Calculate ions per monomer
• Parameterize with Antechamber
• Generate .itp topology files
• Convert AMBER to GROMACS formats

System Preparation:

• Create system topology
• PACKMOL-based packing
• Add water and ions
• Generate MDP files (EM, NVT, NPT, MD)

Simulation:

• Run energy minimization
• Run NVT/NPT equilibration
• Run production MD
• Complete automated workflows

Analysis:

• MSD and diffusion coefficients
• Energy analysis
• RDF calculation
• Trajectory visualization

Water models:

• SPC/E (recommended for HEM)
• TIP3P, TIP4P, TIP4PEW
• Custom models

5. Multiwfn Server

16 tools including:

• Wavefunction analysis
• Orbital analysis and visualization
• Population analysis (Mulliken, Hirshfeld, RESP, etc.)
• Bond order analysis (Mayer, Wiberg)
• Electron density analysis (AIM, ELF, LOL)
• Weak interaction analysis (RDG, NCI)
• Aromaticity indices (NICS)
• Energy decomposition (LMO-EDA)
• Spectrum simulation (UV-Vis, IR, NMR)
• Cube file generation
• 2D/3D orbital rendering

License

This project is licensed under the MIT License - see the LICENSE file for details.

---

OHMind: Accelerating HEM Discovery with AI

From molecules to membranes, from design to deployment

🔬 🤖 ⚡