MossPredict

Official repository for Data-Driven Prediction and Interpretation of Mössbauer Parameters in Molecular Iron Complexes.

MossPredict provides molecular data loaders, descriptor extractors, feature-selection methods, machine-learning models, and tools for prediction and interpretation. The repository includes precomputed features, labels, and experiment settings for molecular iron complexes.

Contents

• Pipeline Overview
• Architecture & Directory Layout
• Installation
• Typical Workflow
• Configuration System
• Dataset Loaders and Classes
• Feature Extraction
• ML Model Types
• Deep Learning Trainers
• Feature Selection
• Results and Interpretation
• Extending the Pipeline

Pipeline Overview

Target	Property	Task
IS	Isomer shift	Regression
QS	Quadrupole splitting	Regression
OS	Fe oxidation state	Classification
SPFE	Fe spin state	Classification
SP	Molecular spin state	Classification

IS and QS are expressed in mm/s. Spin labels use multiplicity, 2S+1.

Architecture & Directory Layout

Path	Purpose
moss_predict/__main__.py	Command-line entry point
moss_predict/data.py	Precomputed-data loading and experiment definitions
moss_predict/engine.py	Training, baselines, and metric summaries
moss_predict/external.py	External-case prediction and evaluation
moss_predict/analysis.py	SHAP analysis and figures
src/data/	Dataset loaders and structure augmentation
src/feature/	Molecular, geometric, electronic, and spectroscopic descriptors
src/feature/feature2D/	Spherical projection descriptors
src/feature_select/	Genetic, greedy, and Bayesian search algorithms and clients
src/task/	Task definitions, feature assembly, caching, labels, and splits
src/training/	Tabular, augmentation-aware, CNN, and UniMol2 training
src/infer/	Model fitting and evaluation with specified parameters
src/infer/analysis/	Feature importance and SHAP analysis
src/utils/	Configuration parsing, descriptor initialization, seeds, and plotting
data/features/	Precomputed inputs and experiment settings
models/	Locally generated model weights
results/	Predictions, metrics, and figures

Project scripts

Script	Purpose
scripts/plot_shap_categories.py	Feature-category contribution plots
scripts/check_independence.py	Standalone input and execution checks
scripts/build_reproduction_bundle.py	Project snapshot packaging

Installation

Use Python 3.10 and run commands from the repository root.

python3.10 -m venv .venv
source .venv/bin/activate
python -m pip install -r requirements.txt

The requirements cover the precomputed-feature workflow, SHAP, and plotting. Structure-based modules use user-supplied molecular and electronic-structure data with their chemistry dependencies. Neural-network modules use PyTorch and their corresponding model dependencies; the XGBoost backend uses the xgboost package.

Typical Workflow

1. Select an experiment

python -m moss_predict list

Experiment IDs combine a task prefix with dft_without_mb, dft_with_mb, uma, or uma_aug, for example IS_dft_without_mb.

2. Train models

Train all 20 experiments:

python -m moss_predict train --all --output models/run

Or train one experiment:

python -m moss_predict train \
  --experiment IS_dft_without_mb --output models/is_dft

Training uses the included numerical inputs, specified feature selection, model parameters, random seeds, and fold assignments. Weights are generated locally. Molecular spin models use separate odd and even subsets; a complete run produces 120 fold models.

Add --resume to skip completed experiments in the same run. Use --jobs to set the training job count.

3. View results

python -m moss_predict summarize \
  --runs models/run --output results/run/tables

python -m moss_predict figures \
  --runs models/run --output results/run/figures

python -m moss_predict baseline --output results/run

4. Predict external cases

python -m moss_predict external \
  --runs models/run --output results/run

This evaluates the 31 external UMA cases and the 16 SC-XRD cases. Predictions and metrics are saved under results/run/external_refit/.

Configuration System

Each experiment is defined in data/features/experiments/<experiment>/experiment.json.

Field	Meaning
target	IS, QS, OS, SPFE, or SP
dataset	Loader, sample mode, and structure source
model_type	Model backend
model_seed	Training random seed
best_params	Model hyperparameters
selection	Selected feature extractors and masks
ga_config	Genetic-search settings
metric_aggregation	Aggregation across subsets

Subset directories contain features.npz, feature_manifest.json, labels.json, samples.csv, and splits.json. Augmented inputs also record parent-sample relationships and fold membership.

For structure-based workflows, src/utils/main_helpers.py provides parse_yaml(), including YAML !include support, and init_feature_extractor(). Descriptor definitions use a type, optional params, and a use_cache setting. TaskData manages extraction, caching, labels, and task splits.

Dataset Loaders and Classes

Version	Class	Main characteristics
V1	V1DataLoader	Molecular metadata and XYZ structures
V2	V2DataLoader	Metadata with MOL/XYZ structure support
V2.1	V21DataLoader	Fe-site sample indexing
V2.2	V22DataLoader	Expanded Fe-site metadata and DFT-related data
V3	V3DataLoader	Extended molecular and electronic metadata
V3.1	V31DataLoader	Sample modes and spin-parity filtering
V4	V4DataLoader	Structure-source selection and low-temperature mode

The V3.1 and V4 loaders support all, single, multi, and mol modes. The mol mode selects pure-Fe compounds for molecular spin prediction and supports sp_parity="odd" or "even". V4 also provides all_low for temperatures at or below 100 K.

Core classes

Class	Role
BaseDataLoader	Common dataset interface
BaseAugmentor, MetaFieldAugmentor	Alternative structures with links to parent samples
BaseFeatureExtractor	Per-sample extraction, feature names, and postprocessing
TaskData	Features, labels, caches, augmentation, and cross-validation
TargetType	Prediction-target enumeration
PrecomputedTaskData	Task interface for the included feature matrices

Read precomputed inputs in Python

from pathlib import Path
from moss_predict.data import experiment

settings, subsets = experiment(
    Path("data/features"), "IS_dft_without_mb"
)
task = subsets["all"]

X = task.selected_matrix()
y = task.arrays["y_encoded"]
feature_names = task.manifest["selected_feature_names"]

print(X.shape, y.shape, len(feature_names))
for train_indices, test_indices in task.split_folds.split():
    print(len(train_indices), len(test_indices))

The task.features mapping provides individual descriptor blocks and their names. task.augmentor_features and task.augmentor_labels provide the corresponding augmentation inputs.

Feature Extraction

Category	Modules or configuration types
Geometry and coordination	geometry, coordination_cn, coordination_atoms, coordination_bond_length, coordination_shell_volume
Distortion and symmetry	distortion_cshm, distortion_rmsd, their V2 variants, and csm_c2/c4/c6/ci
Electronic descriptors	Coordination electronegativity, charge statistics, dipoles, quadrupoles, and octupoles
Electrostatic potential	Surface ESP, extrema, coordination-shell ESP, and basin attractors
Molecular descriptors	soap, soapv2, mpi, and complex_charge
Spectroscopic descriptors	mbparam, mbparam_is, and mbparam_qs
State descriptors	spin, spin_per_fe, and oxidation_state
2D projections	Charge, electronegativity, ESP extrema, and basin-attractor spherical maps

Examples of extractor classes include CoordinationNumberFeatureExtractor, CoordinationBondLengthStatsFeatureExtractor, GeometryFeatureExtractor, SOAPFeatureExtractor, and SOAPV2FeatureExtractor.

Extractors follow the BaseFeatureExtractor interface: _extract_single(data_idx) computes a sample descriptor, extract() assembles the matrix, and feature_names identifies its columns. Optional postprocessing is controlled by enable_post_process and _post_process().

ML Model Types

Backend	Classification	Regression
rf	Random forest	Random forest
adaboost	AdaBoost	AdaBoost
xgboost	XGBoost	XGBoost
svm / svr	Support vector classifier (svm)	Support vector regressor (svr)
mlp	Multilayer perceptron	Multilayer perceptron
gb	—	Gradient boosting

src/training/ml_trainer.py implements model training and hyperparameter search. ml_aug_trainer.py handles augmentation-aware training. The corresponding src/infer/ml_infer.py and ml_aug_infer.py functions fit and evaluate models with specified hyperparameters.

Standard experiment commands apply the stored AdaBoost settings. The --model-type and --retune options support additional model experiments.

Deep Learning Trainers

Class	Purpose
BaseDeepLearningTrainer	Shared fold training, logging, and device management
FoldDataLoader	Training and test data for one fold
Feature2DMapBuilder	Assemble spherical maps and optional tabular inputs
Feature2DCNNTrainer	Train a CNN on 2D descriptor maps
UniMol2MoETrainer	UniMol2 training with descriptor tokens, fusion modules, and training phases

Feature2DCNNConfig controls feature-map selection, tabular inputs, channels, hidden layers, dropout, and optimizer settings. UniMol2MoETrainerConfig controls the architecture, pretrained checkpoint, local subgraph, descriptor targets, additional modules, and training phases.

For a TaskData instance containing spherical feature maps:

from src.training.cnn_trainer import Feature2DCNNConfig, Feature2DCNNTrainer

def train_feature_maps(task_data):
    trainer = Feature2DCNNTrainer(
        task_data,
        config=Feature2DCNNConfig(base_channels=32, dropout=0.2),
        epochs=50,
        batch_size=16,
        device="cuda",
        experiment_name="feature_map_cnn",
    )
    return trainer.train()

The trainers accept parallel_folds and parallel_devices for concurrent fold training. UniMol2 training phases specify which modules are frozen and which losses are active.

Feature Selection

Local genetic search

python -m moss_predict select \
  --experiment IS_dft_without_mb --seed 42 \
  --generations 20 --population 24 --output results/selection

The search evaluates candidate feature groups using the experiment's model parameters. Add --retune for per-candidate hyperparameter search. IS_dft_without_mb_selection/selection.json records selected groups, scores, and search history.

Feature masks

A group-level mask selects feature groups in the order defined by feature_manifest.json; each group expands to its constituent extractors. An extractor-level mask selects individual descriptor blocks. For three ordered entries, [1, 0, 1] enables the first and third entries.

Service clients

For a running feature-selection service, the clients in src/feature_select/ support greedy, Bayesian, and genetic searches:

python -m src.feature_select.greedy \
  --host 127.0.0.1 --port 1234 --name greedy_search

python -m src.feature_select.bayesian \
  --host 127.0.0.1 --port 1234 --name bayesian_search --n-trials 100

python -m src.feature_select.ga \
  --host 127.0.0.1 --port 1234 --name genetic_search --random-seed 42

The Bayesian client uses Optuna. The clients communicate through this interface:

Endpoint	Method	Purpose
/handshake	GET	Obtain mask_type, mask_length, mask_labels, and task_target
/train	POST	Evaluate a JSON feature_mask and return metrics

Results and Interpretation

Output	Contents
models/run/<experiment>/result.json	Fold metrics, means, and standard deviations
models/run/<experiment>/predictions.csv	Out-of-fold targets and predictions
models/run/<experiment>/<subset>/fold_*.joblib	Trained models
results/run/tables/paper_comparison.csv	Metrics alongside the values reported in the paper
results/run/figures/	Regression plots and confusion matrices
results/run/baselines_frozen/results.json	Calibration baseline metrics
results/run/external_refit/	External predictions and evaluation summaries

Regression metrics include MAE, RMSE, and R². Classification metrics include accuracy and macro-averaged F1, precision, and recall.

Calculate SHAP values for a trained model:

python -m moss_predict shap \
  --experiment IS_dft_without_mb --subset all \
  --runs models/run --output results/run/shap

Generate the feature-category plot from the included analysis inputs:

python scripts/plot_shap_categories.py --output-dir results/run/feature_categories

src/infer/analysis/ provides feature-importance and SHAP routines. Neural trainers also record fold-level training histories.

Extending the Pipeline

• Implement descriptors as subclasses of BaseFeatureExtractor, set their feature names, and register configuration types in src/utils/main_helpers.py:init_feature_extractor().
• Add dataset formats through BaseDataLoader subclasses in src/data/.
• Add tabular models under src/training/ and their fixed-parameter evaluation routines under src/infer/.
• Build neural trainers on BaseDeepLearningTrainer and use TaskData for sample, feature, and split management.

See the data format and experiment methods for details. Run python -m moss_predict --help to view command options.