MB-AI Repository Overview

This project hosts the training pipeline for Mössbauer AI experiments. Because large datasets, intermediate artifacts, and experiment logs are .gitignore'd, this README explains the code layout and the local directories you must recreate when onboarding the project.

目录

• Pipeline Overview
• Architecture & Directory Layout
• Typical Workflow
• Configuration System
• Dataset Versions
• ML Model Types
• Deep Learning Trainers
• Batch Experiment Scripts
• Extending the Pipeline
• Feature Selection Server Guide

Pipeline Overview

┌──────────────────────────────────────────────────────────────┐
│              Data Formatting & Preprocessing                 │
│  • Structure file conversion (Gaussian → XYZ/MOL)            │
│  • Metadata aggregation (V1/V2/V2.1/V2.2/V3/V3.1 loaders)    │
│  • DFT calculation processing (Multiwfn integration)         │
└──────────────────────┬───────────────────────────────────────┘
                       │
          ┌────────────┴────────────┐
          ▼                         ▼
┌─────────────────────┐   ┌─────────────────────┐
│  Classical ML Path  │   │   Deep Learning Path │
├─────────────────────┤   ├─────────────────────┤
│ Feature Extraction  │   │  UniMol2 Pretrained │
│  • SOAP/SOAPv2      │   │   Representation    │
│  • Geometric        │   │                     │
│  • Electronic       │   │  Trainers:          │
│  • Distortion       │   │  • unimol2          │
│  • DFT-derived      │   │  • unimol2_mtl      │
│         ▼           │   │  • cnn              │
│ Model Training      │   │         ▼           │
│  • RF, AdaBoost     │   │   Model Training    │
│  • XGBoost, SVM     │   │  • Classification   │
│  • MLP, GradBoost   │   │  • Regression       │
│  • Grid search CV   │   │  • K-fold CV        │
└─────────┬───────────┘   └──────────┬──────────┘
          │                          │
          └──────────┬───────────────┘
                     ▼
┌──────────────────────────────────────────────────────────────┐
│            Result Statistics & Visualization                 │
│  • Cross-validation metrics (accuracy, RMSE, R²)             │
│  • Confusion matrices (classification tasks)                 │
│  • Training history plots (loss curves, per-fold analysis)   │
│  • JSON result persistence with timestamps                   │
└──────────────────────────────────────────────────────────────┘

Prediction Targets:

Target	Full Name	Task Type
IS	Isomer Shift	Regression
QS	Quadrupole Splitting	Regression
OS	Oxidation State	Classification
SP	Spin Quantum Number	Classification
SPFE	Spin State per Fe atom	Classification (multi-Fe)

Architecture & Directory Layout

Path	Purpose
main.py	Entry point; parses a YAML config, builds dataset loaders, feature extractors, and dispatches to ML or DL trainers.
server.py	FastAPI server for feature selection optimization via REST API.
config/	Versioned experiment configs organized by {version}_{target}_{method} (e.g., v3_1_OS_ML/).
src/data/	Dataset loaders: V1DataLoader, V2DataLoader, V21DataLoader, V22DataLoader, V3DataLoader, V31DataLoader.
src/feature/	Feature extractors: SOAP, geometric, electronic, distortion, DFT-derived features.
src/feature/feature2D/	2D feature extractors for CNN (spherical projections).
src/task/	TaskData orchestration, TargetType enums, cross-validation logic.
src/training/	ML trainers (ml_trainer.py) and DL trainers (unimol2_trainer.py, unimol2_mtl_trainer.py, cnn_trainer.py).
src/training/dl_base.py	Base class for DL trainers with parallel fold training support.
src/utils/	CLI argument parser, YAML loader, feature initialization (main_helpers.py), result serialization.
script/	Automation scripts (see Script Organization).

Script Organization

Directory	Purpose
script/data_preprocess/	Data preprocessing and format conversion scripts.
script/multiwfn_process/	Batch Multiwfn calculations for quantum chemistry features.
script/process_dft_data/	DFT calculation processing (Gaussian → features).
script/result_process/	Result analysis and cleanup scripts.
script/run_experiment/	Batch experiment scripts for running multiple configs.
script/data_analysis/	Data analysis and statistics scripts.
script/data_check/	Data validation and integrity checks.
script/visualization/	Plotting and visualization tools.
script/structure_gen/	Structure generation pipeline.

Typical Workflow

1. Activate environment

   source ~/PythonVenv/MoessbauerEnv/bin/activate
   # Python interpreter: /data/home/zzb/PythonVenv/MoessbauerEnv/bin/python

2. Prepare data – drop your processed files into dataset/ following the loader's expected structure; run helper scripts under script/data_preprocess/ if needed.

3. Choose a config – copy an existing YAML from config/ and adjust dataset paths, feature lists, targets, or DL parameters.

4. Launch training

   # ML classification
   python main.py --config config/v3_1_OS_ML/all_ALL.yaml --name exp_os --model-type rf
   
   # ML regression
   python main.py --config config/v3_1_IS_ML/all_ALL.yaml --name exp_is --model-type xgboost
   
   # Deep learning
   python main.py --config config/v3_1_OS_DL/all_dl_backbone_freeze.yaml --name exp_os_dl

   The --name flag is required (used for result files). --model-type overrides the model defined in YAML. Results save to result/{config_subdir}/{name}.json.

Configuration System

Config Directory Naming Convention

Configs are organized by: {version}_{target}_{method}/

• Version: v3, v3_1
• Target: IS, QS, OS, SP, SPFE
• Method: ML, DL

Example: config/v3_1_OS_ML/ contains ML configs for Oxidation State prediction using V3.1 dataset.

Config File Naming Convention

• {mode}_ALL.yaml - All features enabled
• {mode}_ALL_no_mbparam.yaml - Exclude MBParam features
• {mode}_ALL_no_ossp.yaml - Exclude OS/SP features (for regression)
• {mode}_dl_backbone_freeze.yaml - DL with frozen backbone
• {mode}_dl_backbone_no_freeze.yaml - DL with trainable backbone

Where mode is: all, single, multi, mol (V3.1 specific).

Key Config Options

dataset:
  loader: v31          # v1, v2, v21, v22, v3, v31
  data_dir: dataset/mb-dataset-latest-20251216
  mode: all            # V3.1 only: all, single, multi, mol
  split_by_sp_parity: false  # V3.1 mol mode: split by odd/even SP

features:
  - type: coordination_bond_length
    use_cache: 1       # 0=force recompute, 1=use cache
  - type: soapv2
    params:
      pca_variance: 0.95

task:
  target: OS           # IS, QS, OS, SP, SPFE
  method_type: ml      # ml or dl
  model_type: rf       # See ML Model Types

dl:
  trainer: unimol2     # unimol2, unimol2_mtl, cnn
  arch_name: 84M       # 84M or 570M
  freeze_backbone: true
  use_local_subgraph: true  # Focus on Fe coordination environment
  expand_hops: 0       # 0 = Fe + coordination shell only
  epochs: 500
  batch_size: 8
  device: cuda

  # Parallel training (train folds simultaneously on multiple GPUs)
  parallel_folds: true
  parallel_devices: [0, 1, 2, 3, 4]
  progress_report_interval: 10  # Report every N epochs

Dataset Versions

Version	Loader	Description
V1	v1	Original dataset, single Fe per structure
V2	v2	Enhanced metadata, better coordination info
V2.1	v21	Multi-Fe support (expanded by Fe atoms)
V2.2	v22	Integrated DFT calculations
V3	v3	Full DFT feature support
V3.1	v31	Mode-based filtering with mode option

V3.1 Modes

Mode	Description
all	All compounds (single + multi-Fe)
single	Single-Fe compounds only
multi	Multi-Fe compounds only
mol	Pure Fe compounds (no other transition metals), SP target only

For mol mode, additional sp_parity filter available:

• all: No filtering (default)
• odd: Only compounds with odd SP values (1, 3, 5, ...)
• even: Only compounds with even SP values (2, 4, 6, ...)

ML Model Types

Classification (OS, SP, SPFE)

Model Type	Algorithm
rf	Random Forest Classifier
adaboost	AdaBoost Classifier
xgboost	XGBoost Classifier
svm	Support Vector Machine
mlp	Multi-Layer Perceptron

Regression (IS, QS)

Model Type	Algorithm
gb	Gradient Boosting Regressor (default)
rf	Random Forest Regressor
adaboost	AdaBoost Regressor
xgboost	XGBoost Regressor
svr	Support Vector Regressor
mlp	Multi-Layer Perceptron Regressor

Deep Learning Trainers

Trainer	Description
unimol2	UniMol2 pretrained molecular representations with custom prediction head
unimol2_mtl	UniMol2 with Multi-Task Learning (auxiliary heads for auxiliary targets)
cnn	Feature2D CNN for 2D spherical projection features

Key DL Options

• freeze_backbone: Whether to freeze UniMol2 backbone (feature extraction only)
• use_local_subgraph: Focus on Fe coordination environment instead of full molecule
• expand_hops: Number of hops to expand from Fe (0 = Fe + coordination shell)
• parallel_folds: Train K-folds in parallel on multiple GPUs
• parallel_devices: List of GPU IDs for parallel training

Batch Experiment Scripts

Scripts in script/run_experiment/ automate running multiple experiments:

# Run experiments with multiple random seeds (parallel, requires parallel_devices)
bash script/run_experiment/run_multi_seed.sh <config_path> <experiment_name> <parallel_devices>
# Example:
bash script/run_experiment/run_multi_seed.sh config/v3_1_OS_DL/all_dl_mtl_all_v2.yaml exp_os_mtl 1,2,3,4,5

# Sequential multi-seed run
bash script/run_experiment/run_multi_seed_seq.sh <config_path> <experiment_name>

# Launch experiments inside a tmux session
bash script/run_experiment/tmux_run.sh

Extending the Pipeline

• Add new feature extractors inside src/feature/ and register them in src/utils/main_helpers.py:init_feature_extractor(). See the Feature Extractor Guide.
• Review existing extractors in the Feature Extractors Guide.
• Extend dataset loaders in src/data/ when the folder schema changes.
• Introduce new trainers under src/training/ but re-use TaskData so configs remain consistent.

Available Feature Extractors

Category	Extractors
Geometric	coordination_cn, coordination_bond_length, coordination_shell_volume, geometry
Electronic	coordination_electronegativity, coordination_charge, surface_esp, surface_esp_coord
Distortion	distortion, distortion_v2, csm
Molecular	soap, soapv2, mbparam, mpi, multipole
ESP	esp_extrema, esp_extrema_coord
Labels	spin_oxidation (OS/SP as features for regression tasks)

Feature Selection Server Guide

The MB-AI feature-selection REST API server lets external optimization algorithms search feature combinations.

Quick Start

# Launch the server (default port: 8000)
python server.py --config config/example_feature_groups.yaml --model-type rf

# Development mode (automatic reload)
python server.py --config config/example_feature_groups.yaml --reload

# Run feature selection clients (default host: 10.137.144.95, port: 1234)
python -m src.feature_select.greedy   --host <host> --port <port> --name exp_greedy
python -m src.feature_select.bayesian --host <host> --port <port> --n-trials 100 --name exp_bayes
python -m src.feature_select.ga       --host <host> --port <port> --name exp_ga

Architecture Overview

┌─────────────────────────┐
│ server.py (FastAPI app) │
└──────────────┬──────────┘
               │ initialize_server()
               ▼
┌──────────────────────────┐
│ parse_yaml(config)       │
└──────────────┬───────────┘
               ▼
┌──────────────────────────┐
│ Data loaders (V1/V2/…)   │
└──────────────┬───────────┘
               ▼
┌──────────────────────────┐
│ init_feature_extractor() │
└──────────────┬───────────┘
               ▼
┌──────────────────────────┐
│ TaskData + cache configs │
└──────────┬─────┬─────────┘
           │     │
   GET /handshake │ exposes mask metadata from TaskData
           │     │
           ▼     ▼
   POST /train → build feature subset → call ML trainers

Key Concepts

Mask Types

Group-level masks (when feature_groups is defined):

feature_groups:
  geometric: [coordination_cn, coordination_bond_length]
  electronic: [coordination_charge_mean]

• Mask length = number of feature groups
• [1, 0] means only the geometric group is active

Extractor-level masks (when feature_groups is not defined):

features:
  - type: coordination_cn
  - type: coordination_bond_length
  - type: soapv2

• Mask length = number of feature extractors
• [1, 0, 1] enables the first and third extractors

API Endpoints

Endpoint	Method	Description
/handshake	GET	Return server configuration and mask information
/train	POST	Submit a feature mask and get training results
/health	GET	Health check
/docs	GET	Interactive API documentation (Swagger UI)

Example Usage

import requests

SERVER = "http://localhost:8000"

# 1. Inspect the mask configuration
info = requests.get(f"{SERVER}/handshake").json()
print(f"Mask type: {info['mask_type']}, length: {info['mask_length']}")

# 2. Train with all features
mask_all = [1] * info['mask_length']
result = requests.post(f"{SERVER}/train", json={"feature_mask": mask_all}).json()
print(f"Accuracy: {result['mean_accuracy']:.4f}")

# Fetch mask configuration
curl http://localhost:8000/handshake | jq

# Submit training job
curl -X POST http://localhost:8000/train \
  -H "Content-Type: application/json" \
  -d '{"feature_mask": [1,1,0,1,0]}' | jq