EZSolver: Flow Matching-based Enzymatic Reaction Mechanism predictor

EZSolver is a generative model for reaction mechanism prediction that leverages a bidirectional flow matching and bidirectional beam search to predict polar enzymatic reaction mechanisms.

---

🌟 Key Features

• Flow Matching: Models transition pathways of Bond-Electron (BE) matrices from reactant states to product states.
• Bidirectional Beam Search: Simulates pathways from both the reactant front (forward search) and the product front (backward search) simultaneously, running both exo (forward prediction) and endo (inverse prediction) at each step to find intersecting intermediate states and reconstruct the complete mechanism.
• Physical & Chemical Constraints: Integrates rigorous chemistry filters (octet rule checks, biochemical charge sanity, radical suppression) and 2D atom-mapping masks to prune chemically invalid intermediates.
• Interactive Review Interface: Built-in Jupyter Notebook with a custom ipywidgets UI to visualize molecular structures (via RDKit SVG rendering) and label predicted mechanism steps as accurate, plausible, or retrieved.

---

📂 Project Structure

EZSolver/
├── .gitignore             # Git ignore file configuration
├── checkpoints/           # Model checkpoints folder (place FlowER.pt here)
├── EZSolver.py            # Main entry point for predicting reaction mechanisms
├── EZFlow.py              # Evaluation script to compute single-step prediction
├── LICENSE                # MIT License file
├── train.py               # Model training script 
├── settings.py            # Configuration & hyperparameter parser (Args class)
├── requirements.txt       # Python package dependencies list
├── review_mechanisms.ipynb # Jupyter Notebook for visual inspection of results
├── plot_accuracy.py       # Utility for plotting model evaluation accuracies
├── run_train.sh           # Shell script to start the training pipeline
├── run_EZFlow.sh          # Shell script to run model evaluation
├── run_EZSolver.sh        # Shell script to execute mechanism pathfinding
├── data/                  # Reaction dataset and curation folder
│   ├── data_curation/     # Raw dataset processing scripts
│   │   ├── map.py         # Atom mapping & polar/radical classification script
│   │   └── mmff.py        # Force field energy minimization & direction classification script
│   ├── EZFlow/            # Training/evaluation reaction files
│   └── EZSolver/          # Test reaction database JSONs
└── utils/                 # Utilities folder
    ├── predict_utils.py   # Distributed inference helper functions
    ├── data_utils.py      # BE matrices operations and datasets
    ├── physical_filters.py # Octet rule, charge sanity, and chemistry evaluation rules
    ├── train_utils.py     # Noam scheduler and setup loggers
    ├── map_utils.py       # Atom mapping utilities
    └── rounding.py        # Quantization / rounding methods for BE matrices

---

⚙️ Installation & Environment Setup

This project requires Python 3.10.

1. Setup Virtual Environment (Recommended)

You can set up a clean Python virtual environment using venv or Conda:

conda create -n ezsolver python=3.10 -y
conda activate ezsolver

2. Install Dependencies

All dependencies, including PyTorch configured for CUDA 11.8 and chemistry libraries, can be installed directly via pip:

pip install -r requirements.txt

3. Download Model Checkpoints

The pre-trained model checkpoints (such as FlowER.pt) can be downloaded from Hugging Face: 🔗 LunHsin/EZSolver on Hugging Face Repository

(Note: Model descriptions can be viewed directly on the Hugging Face model repository page.)

Please download the checkpoint file and place it inside the existing checkpoints/ directory.

---

📁 Dataset Descriptions

The repository contains several datasets under data/ configured for different training and evaluation purposes:

1. Training & Evaluation Datasets (data/EZFlow/)

• EZFlow/:
  • train.txt / val.txt: Mapped reaction elementary steps datasets without MMFF energy classification.
  • train_exo.txt / val_exo.txt: Mapped reaction elementary steps datasets with MMFF energy classification (ensuring exothermic directionality).
  • test.txt: The test split of 186 elementary steps, which serves as the test file for evaluating single-step prediction accuracy (e.g. via run_EZFlow.sh).
• OOD_Benchmark/:
  • Datasets specifically curated for training comparison against baseline models EZMechanism and Mechfind.

2. Mechanism Search Evaluation Datasets (data/EZSolver/)

The JSON files in this directory are target reactions used to evaluate EZSolver's multi-step mechanism pathfinding:

• EZSolver_test.json: The Out-of-Distribution (OOD) task dataset established in this work to evaluate mechanism prediction.
• OOD_Benchmark.json: The benchmark dataset used to compare EZSolver's mechanism pathfinding performance against EZMechanism and Mechfind.

---

🚀 Execution & Usage

1. Training the Model

To start training the conditional flow matching model on reaction datasets, run:

sh run_train.sh

Configure your data paths (e.g. TRAIN_FILE, VAL_FILE) and hyperparameters (e.g., learning rate LR, batch size, etc.) directly in run_train.sh or through setting environment variables.

2. Evaluating Model Accuracy

To evaluate the top-k prediction performance of the flow matching model for both standard (exo) and inverse (endo) reaction directions:

sh run_EZFlow.sh

Details & Parameters:

• Input test file: Configured via TEST_FILE (default: data/EZFlow/EZFlow/test.txt - a text file containing reaction SMILES).
• Model checkpoint: Loaded via LOAD_FROM (default: checkpoints/EZFlow.pt).
• Key parameters:
  • SAMPLE_SIZE (default 200): Number of ODE integration steps / trajectory samples generated per reaction.
  • SIGMA (default 0.08): Standard deviation for flow matching training/inference.
• Output files:
  • A success log summarizing top-k accuracy statistics of both directions, saved to ${RESULT_PATH}/{SAMPLE_SIZE}_prediction.txt (default: EZFlow_results/{SAMPLE_SIZE}_prediction.txt).
  • Detailed execution log saved in logs/ezflow/.

3. Mechanism Search

To perform a complete bidirectional search on multi-step reactions and reconstruct the pathway of intermediate states:

sh run_EZSolver.sh

Details & Workflow:

• Script executed: EZSolver.py (which runs local evaluation and search algorithm).
• Workflow:
  1. Load Inputs: Reads a JSON database of target chemical reactions from the path in REACTIONS (default: data/EZSolver/EZSolver_test.json).
  2. Bidirectional Search: Simulates reaction pathways from both ends simultaneously. For both the reactant front (forward search) and the product front (backward search), the solver runs both exo (standard flow) and endo (inverse flow) at each step to find intersecting intermediate states.
  3. Pruning & Filtering: Evaluates intermediates and penalizes paths violating chemical constraints (like octet rules, invalid radical formations, strained ring formation) and 2D masks.
  4. Output Results: Reconstructs and sorts the matched paths by length and score, and exports them to a JSON file at MECHANISM_RESULTS (default: EZSolver_results/20.json).
• Configurable variables in run_EZSolver.sh:
  • LOAD_FROM: Path to the model checkpoint .pt file (e.g. checkpoints/EZFlow.pt).
  • BEAM_SIZE (default 10): Beam search width for pathway tracking.
  • MAX_DEPTH (default 4): Maximum search depth (maximum number of steps).
  • SAMPLE_SIZE (default 160): Candidate sample size generated for each intermediate step.

---

📊 Mechanism Review UI & Operation Guidelines

The repository includes review_mechanisms.ipynb to visually inspect and annotate search results using an interactive browser-based UI (ipywidgets + rdkit).

⚙️ How to Start

1. Configure File Path (Cell 2):
   • Modify the JSON_FILE = Path("EZSolver_results/your_result.json") line to target your solved JSON file (e.g., EZSolver_results/EZSolver_20.json).
   • The notebook automatically determines the sample name, sets up column mappings (top1 and top5), and designates the output to your_result_results.csv.
2. Initialize / Load CSV (Cell 3):
   • Run the cell to load any existing annotations. It imports previous labels (fully compatible with legacy column names), pre-fills none for invalid reactions, and populates the time column with the execution times in seconds.
3. Compile Helpers & Launch UI (Cells 4 & 5):
   • Run these cells to render the molecular SVG grid and the control dashboard.

🎮 Control Dashboard Buttons & Actions

• ✔ Accurate / ~ Plausible / ↩ Retrieved: Clicking these buttons labels the current mechanism candidate, immediately auto-saves the CSV file, and jumps to the next unlabeled reaction.
• 〉 Next / ◀ Back: Manually navigate forward or backward to review reactions.
• ↑ Prev Cand / ↓ Next Cand: If the beam search found multiple candidate pathways for a reaction, use these to cycle through and inspect different predicted mechanism candidates.

---

🛠️ Data Curation & Preprocessing

Before training or evaluation, raw reaction datasets can be curated, atom-mapped, and direction-classified using the scripts inside data/data_curation/.

1. Atom Mapping & Classification (map.py)

To map atoms (including hydrogen mapping and proton transfers) and classify reactions by chemistry (Polar vs. Radical) and elements (metal vs. non-metal):

python map.py

• Input: Reads MCSA_reactions.json.
• Process: Uses rxnmapper to generate attention-guided atom mappings, checks mapping consistency, maps hydrogens/protons, and determines if reactions are radical/polar or metal-containing.
• Output: Saves the complete mapped dictionary to mapped_reactions.json and splits them into category JSONs in the mapped_reactions/ directory (e.g., Polar_NM.json, Polar_M.json, Radical_NM.json, Radical_M.json).

2. Thermodynamic Classification (mmff.py)

To classify the reaction directions (Exothermic vs. Endothermic) using molecular force field energy calculations:

python mmff.py

• Input: Reads the JSON files inside the mapped_reactions/ directory.
• Process:
  • Evaluation: Automatically splits disconnected components in reaction mixtures to calculate energies individually and sum them up (avoiding slow 3D conformation searches on disconnected complexes).
  • Unsupported Atom Filtering: Skips conformation embedding for unsupported metal-containing complexes to avoid massive CPU overhead.
  • Direction Classification: Computes MMFF94 force field energy changes ($\Delta E = E_{\text{product}} - E_{\text{reactant}}$) to classify reaction directions.
• Output: Writes training-formatted lists of reaction SMILES (with trailing | format) into the same folder:
  • *_fwd.txt: Reactions that are exothermic ($\Delta E < 0$).
  • *_bwd.txt: Reactions that are endothermic ($\Delta E > 0$).
  • *_combined.txt: Flipped backward reactions combined with forward reactions to keep all pathways exothermic.

---

📄 License

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