MILIA (Machine Intelligent Learning Interface Assistant)

Molecular graph processing and machine learning framework for computational chemistry

MILIA (Machine Intelligent Learning Interface Assistant) is a production-ready, research-oriented Python framework for molecular data processing and graph-based machine learning. It provides a complete, no-code ML/DL workflow — from dataset curation, molecular graph transformation, and molecular descriptor computation through graph neural network training, hyperparameter optimization, and model deployment — requiring only YAML configuration to run. Every level of the pipeline (datasets, transformations, descriptors, models, training, deployment) is fully configurable without writing code, while remaining extensible via plugins for untouched areas of research. Built on PyTorch, PyTorch Geometric, and RDKit, MILIA is designed for researchers, educators, and teams who need a flexible, extensible platform across the full ML/DL stack.

MILIA fills a gap in the molecular ML ecosystem by unifying dataset handling, feature engineering, model training, and deployment into a single, extensible framework with plugin support. Where tools like PyTorch Geometric provide GNN building blocks and DeepChem offers pre-built models, MILIA provides a configuration-driven research workflow — supporting any PyG model architecture through dynamic introspection, any hardware from CPU to TPU, and any dataset through its zero-modification extension architecture. The framework currently ships with 11 dataset implementations spanning the VQM24 family (DFT, DMC, Wavefunction), the QM9 benchmark, the ANI family (ANI-1x, ANI-1ccx, ANI-2x) including coupled-cluster reference data, the rMD17 and xxMD reactive dynamics datasets, the drug-discovery-oriented QDπ dataset, and the QM40 dataset of neutral drug-like ZINC molecules.

🚀 New here? For the shortest reproducible path from zero to a running MILIA install — Docker pull, smoke test, one full walkthrough — see QUICKSTART.md. It is designed to be executable end-to-end in ≤30 minutes on a CPU-only laptop without contacting the author.

Key Features

No-Code ML/DL Workflow

Run the entire machine learning pipeline — dataset curation, molecular graph transformations, molecular descriptor computation, model training, hyperparameter optimization, and prediction — through YAML configuration and CLI commands. No code required at any level. For research requiring capabilities beyond the built-in functionality, the plugin architecture allows code-level extension of datasets, transformations, descriptors, and models without modifying the core framework.

Unlimited Model Flexibility

Access every PyTorch Geometric model (SchNet, DimeNet, GIN, GAT, and all others) simply by naming them in configuration — no model-level code. Define custom architectures from 10 built-in templates, compose multi-model ensembles with parallel, sequential, or hierarchical strategies — all through YAML configuration, no code required. For research beyond built-in capabilities, bring your own models through the model plugin system.

Hardware Agnostic

Train on any device — CPU, CUDA GPU, Apple MPS, or TPU — with automatic device detection or explicit selection through configuration. Scale from single-device to distributed training with 4 strategies (DataParallel, DistributedDataParallel, FSDP, Horovod). Includes memory optimization (AMP, gradient checkpointing) and computation optimization.

Molecular Descriptors

Select from 400+ molecular descriptors across 6 categories — Constitutional (35), Topological (350+), Electronic (8), Geometric (10), Drug-likeness (4), and Fragments (85) — entirely through YAML configuration. MILIA handles the full computation pipeline: multi-format molecular conversion via RDKit, atom-level feature extraction (degree, hybridization, chirality, Mulliken charges), bond-level properties (type, conjugation, stereo, length), automatic conformer generation for 3D descriptors, and result caching. Extend with custom descriptors through the plugin system for advanced research needs.

Advanced Hyperparameter Optimization

Optuna-based optimization with 5 search algorithms (TPE, CMA-ES, Random, Grid, NSGA-II for multi-objective), 5 pruning strategies (Median, Hyperband, Percentile, Patient, Threshold), cross-validation integration, study persistence, and resume support. Includes Neural Architecture Search for GNNs (7 layer types, heterogeneous architectures) and HPO transfer learning with meta-feature extraction and warm-starting across studies.

Production-Ready Deployment

Edge, cloud, and federated deployment strategies with model quantization and pruning. Production monitoring for drift detection and performance tracking. Checkpoint management with training state persistence, model loading, and fine-tuning via transfer learning.

Extensible Graph Transformation System

30+ pre-registered PyG transforms with a 7-layer architecture: dynamic discovery, registry, validation (semantic + dataset-aware), composition with intelligent caching, configuration bridge, error recovery, and production metrics (Prometheus/DataDog export). Three validation levels (Strict/Standard/Lenient) and five validation scopes. Edge-attr-aware parameter injection prevents shape mismatch errors.

Three-Tier Plugin Architecture

Extend descriptors, transformations, and models independently without modifying core code. Plugin discovery with YAML manifests, validation, and security controls. Ships with example plugins and user templates.

Flexible Configuration System

Schema-validated YAML with Pydantic V2 (10 frozen BaseModel containers, 60+ accessor functions). Split-file configs/ directory as the sole configuration source with deep merge, CLI override, and configuration migration. Explicit single-file paths are supported via --config. Each dataset type has colocated configuration files for self-contained setup.

Wavefunction Preprocessing

Modular preprocessing for quantum chemistry data formats (MOLDEN, FCHK) with structural feature filtering per dataset type, wavefunction data extraction, and VQM24 support.

Research-Ready

Experiment configuration with research API for transformation ablation studies, hyperparameter sweeps, and model comparison. Research-grade recommendations per dataset type (DFT precision, DMC uncertainty preservation, wavefunction orbital analysis). Benchmarking and validation reporting in text, JSON, and markdown formats.

Educational Use

The no-code design, comprehensive CLI with interactive mode, and 12+ processing modes make MILIA suitable for teaching molecular ML concepts without requiring students to write model code. Configuration-driven workflows allow focusing on chemistry and ML concepts rather than software engineering.

Installation

MILIA's dependencies (PyTorch, PyTorch Geometric + compiled companions, RDKit, …) are managed with uv from a committed lockfile (uv.lock) for fully reproducible installs. A pre-built Docker image is also published.

Method 1: Docker (Recommended)

The fastest way to get MILIA running. A pre-built image is publicly available on GitHub Container Registry — no authentication required:

# Pull the pre-built CPU image (no authentication required)
docker pull ghcr.io/shahram-boshra/milia:latest

# The image's entry point IS the `milia` CLI — pass CLI args directly:
docker run --rm ghcr.io/shahram-boshra/milia:latest --help

# Process a dataset (runs non-root; mount your data read-write at /data):
docker run --rm --user "$(id -u):$(id -g)" \
  -v ~/Chem_Data/Milia_PyG_Dataset:/data:rw \
  ghcr.io/shahram-boshra/milia:latest --root-dir /data

Or build locally from the Dockerfile:

git clone https://github.com/shahram-boshra/MILIA.git
cd MILIA
docker build -t milia:cpu .                    # CPU (default); GPU: --build-arg ACCEL=cu124
docker run --rm milia:cpu --help

The published image is a lean production runtime (no test tools). The smoke suite runs in CI (and locally via uv run --extra cpu --extra dev pytest -m smoke), not inside the runtime image.

Method 2: uv (Without Docker)

# 1. Install uv (https://docs.astral.sh/uv/)
curl -LsSf https://astral.sh/uv/install.sh | sh

# 2. Clone MILIA
git clone https://github.com/shahram-boshra/MILIA.git
cd MILIA

# 3. Install the full stack from the committed lockfile.
#    Choose exactly ONE accelerator extra: cpu | cu118 | cu121 | cu124.
uv sync --locked --extra cpu

For GPU support, choose a CUDA accelerator extra instead of cpu — e.g. uv sync --locked --extra cu124 (pick the one matching your CUDA runtime; requires an NVIDIA driver + container toolkit at run time).

For development (adds pytest and ruff):

uv sync --locked --extra cpu --extra dev

Quick Start

Command Line

# Process molecular data (auto-detects configs/ directory)
milia --process

# Train a GNN model
milia --train

# Train with hyperparameter optimization (via CLI flag)
milia --train --hpo

# Run predictions on new molecules (supply your own CSV: header smiles,molecule_id)
milia --predict --model-path ./checkpoints/best.pt \
      --test-path ./my_molecules.csv --preds-path ./predictions.csv

# Validate configuration without processing
milia --dry-run

# List available transforms and experimental setups
milia --list-transforms
milia --list-experimental-setups

# Generate statistics from existing data
milia --stats-only

# Explicit config path (equivalent to auto-detection when configs/ exists)
milia --config configs/ --process

Programmatic API

from milia_pipeline import create_cli_manager, setup_logging

# Setup
logger = setup_logging(log_level="INFO")
cli = create_cli_manager(logger=logger)
args = cli.parse_args(['--config', 'configs/', '--process'])

# Load and validate configuration
config = cli.load_and_merge_config(args)
cli.validate_args(args, config)

from milia_pipeline.config import load_config
from milia_pipeline.handlers import create_handler
from milia_pipeline.datasets import miliaDataset

# Load configuration (split-file mode: merges all YAML files in configs/)
config = load_config('configs/')
handler = create_handler(
    dataset_type='DFT',
    config=config,
    logger=logger
)

# Build a PyTorch Geometric dataset
dataset = miliaDataset(
    root='./data',
    handler=handler,
    transform=my_transforms
)

Trying MILIA — Reproducible Walkthrough

This section walks a reviewer through the shortest path from a fresh clone to a trained model and a prediction. Every command below has been validated end-to-end. Paths are relative throughout, so the walkthrough works identically on any machine — Linux, macOS, WSL, or inside the Docker container.

Prerequisites

1. MILIA installed — see Installation above. Method 1 (Docker) is the fastest route for a one-shot evaluation; Method 2 (uv) is recommended if you intend to inspect or modify source.

2. working_root_dir set in your configuration — this is the only path you must configure. It tells MILIA where to download datasets, write processed graphs, and save checkpoints. There is no implicit default — the framework asks you to choose deliberately.

   Open configs/main.yaml and set:

   global_paths:
     working_root_dir: ~/Chem_Data/Milia_PyG_Dataset    # or any directory you prefer
   

   The ~ is expanded to your home directory at runtime, so the same line works for any user. Inside the Docker image the value is preset to /root/Chem_Data/Milia_PyG_Dataset and you can leave it as-is.

3. A laptop-friendly configuration — the shipped configs/models.yaml is preconfigured for low-resource execution (small batch size, few epochs, small ensembles). A reviewer with a CPU-only laptop can run the walkthrough below without a GPU.

Step-by-step

Run each command from the repository root (or, in Docker, from /app/milia). Every command auto-detects the configs/ directory.

# 1. Sanity-check the install (under two minutes on CPU)
pytest -m smoke --tb=short

# 2. Validate the configuration without doing any work
milia --dry-run

# 3. Process the dataset (downloads if needed, writes PyG graphs under working_root_dir)
milia --process

# 4. Inspect the processed dataset
milia --stats-only

# 5. Train a model
#    Best checkpoint written to {working_root_dir}/checkpoints/best.pt
milia --train

# 6. (Optional) Train with hyperparameter optimization instead of step 5
milia --train --hpo

# 7. Run prediction on your own molecules (CSV with header: smiles,molecule_id)
milia --predict \
    --model-path ./checkpoints/best.pt \
    --test-path ./my_molecules.csv \
    --preds-path ./predictions.csv

After step 5 or 6, your {working_root_dir}/checkpoints/ directory contains a best.pt checkpoint. Step 7 reads that checkpoint and writes per-molecule predictions to ./predictions.csv for the molecules in your --test-path input — a CSV with a smiles,molecule_id header and one molecule (SMILES) per row.

Where the outputs live

Every artifact MILIA produces lands under working_root_dir. With the example value ~/Chem_Data/Milia_PyG_Dataset this means:

The CLI resolves checkpoint paths against working_root_dir automatically, so commands that reference ./checkpoints/best.pt work regardless of which directory you run them from.

Introspection commands

These are read-only and useful for orientation:

milia --list-transforms              # 30+ pre-registered PyG transforms + plugin transforms
milia --list-experimental-setups     # Available research/experiment configurations
milia --help                         # Full CLI reference

Prediction input format

milia --predict --test-path <file> expects a CSV with a smiles,molecule_id header and one molecule per row, for example:

smiles,molecule_id
CCO,ethanol
CC(=O)O,acetic_acid
c1ccccc1,benzene

Supply your own file — no sample dataset is bundled with the repository.

Architecture

MILIA is organized into 11 core modules and a split configuration system:

Supporting directories:

Datasets

MILIA ships with 11 production-ready dataset implementations covering the major quantum-chemistry and molecular-dynamics benchmarks used in modern molecular machine learning. Selecting a dataset is a one-line change in configs/main.yaml (dataset_type: <Registry key>) — no code, no rebuild, no glue.

Shipped datasets

The DFT, DMC, and Wavefunction entries are the three deliverables of the Vector-QM24 (VQM24) dataset (Zenodo: 10.5281/zenodo.11164951) — DFT geometries and properties, DMC reference energies, and quantum-mechanical wavefunctions respectively — for which MILIA provides VQM24-aware vibrational refinement and MOLDEN/FCHK readers out of the box.

Per-dataset capability matrix

Each shipped dataset declares an immutable DatasetFeatures record (a Pydantic V2 frozen dataclass) that drives feature-aware code paths in MILIA — descriptor selection, transform compatibility validation, atomization-energy bookkeeping, and orbital-property extraction — without any consumer needing to hard-code dataset names. The five flags below are the most consequential for downstream model targets and are queried at runtime via _get_dataset_feature(dataset_type, feature_name):

These five columns are a subset of the eight flags carried by DatasetFeatures; the other three (rotational_constants, frequency_analysis, mo_energies) gate finer-grained behaviour and are queried programmatically rather than displayed here. The matrix is not a documentation artefact — it is the same registry data that MILIA consults to decide, for example, whether a transform that requires vibrational modes is admissible against a given dataset, or whether atomization-energy targets are derivable. Setting vibrational_analysis: True on a new dataset's DatasetFeatures is therefore sufficient to opt into vibrational-aware processing across the pipeline; no other module needs editing.

Adding a dataset

Adding a new dataset to MILIA is a three-file operation against existing extension points — no core file is modified, and no registration list, switch statement, or import is touched anywhere else in the codebase:

The architecture enforcing this is Protocol + ABC + explicit registry, applied identically to datasets and handlers:

• @register / @register_handler decorators wire the new class into thread-safe (RLock-protected) registries at import time, with no central registration list to update.
• DatasetHandlerProtocol (11 runtime-checkable methods) and the BaseDataset ABC (with __init_subclass__ validation) catch contract violations at import time, not at runtime — a missing required method fails fast before any data is loaded.
• DatasetMetadata, DatasetSchema, DatasetFeatures are immutable Pydantic V2 frozen dataclasses, so the new dataset's metadata is validated structurally before registration succeeds.
• Feature flags (vibrational_analysis, uncertainty_handling, atomization_energy, orbital_analysis, homo_lumo_gap, and three others) declared in the BaseDataset subclass make the new dataset visible to feature-aware code paths in milia_dataset.py and the transform validator without any of those modules learning the new dataset's name.

Once the three files are in place, the new dataset is a first-class citizen of every part of the pipeline — CLI (milia --process, --train, --predict), YAML configuration, hyperparameter optimization, transform validation, descriptors, and prediction — selectable by setting dataset_type: <key> in configs/main.yaml. The 11 shipped datasets follow this exact pattern; the registry has no knowledge that they are "built-in" rather than user-added.

Testing

# Run the full test suite
pytest

# Run specific test categories
pytest -m "not slow"         # Skip slow tests
pytest -m integration        # Integration tests only
pytest -m gpu                # GPU-specific tests only

The test suite includes 127 unit and integration tests covering all core modules.

Contributing

Contributions are welcome. Please see CONTRIBUTING.md for guidelines on how to get started, coding standards, and the pull request process.

Citation

If you use MILIA in your research, please cite it. See CITATION.cff for the preferred citation format.

License

MILIA is released under the MIT License.

Author

• Asadollah (Shahram) Boshra — a.boshra@gmail.com

Links

• Repository: github.com/shahram-boshra/MILIA
• Issues: github.com/shahram-boshra/MILIA/issues
• Changelog: CHANGELOG.md