Skip to content

Add MLIPX PaiNN-style MLIP project: data pipeline, model, training, inference, and tests - #3

Open
yeze5421 wants to merge 1 commit into
mainfrom
codex/create-project-for-atomic-potential-calculation-5crq9y
Open

yeze5421 wants to merge 1 commit into
mainfrom
codex/create-project-for-atomic-potential-calculation-5crq9y

Conversation

@yeze5421

Copy link
Copy Markdown
Owner

Motivation

  • Provide a minimal research/engineering codebase for training PaiNN-style atomistic potentials with ASE/extxyz datasets and energy+force supervision.
  • Replace the prior toy example with a structured, testable package that supports training, evaluation, inference, and an ASE Calculator wrapper.
  • Offer reproducible defaults and convenience scripts to create demo data, run training/inference/evaluation, and package the project.

Description

  • Add package layout and metadata including pyproject.toml, README.md, and two config files configs/minimal.yaml and configs/standard.yaml with sensible defaults.
  • Implement configuration utilities in src/mlipx/config.py and data handling with ASEDataset, MLIPDataModule, neighbor graph builder, and GraphBatch types under src/mlipx/data/.
  • Implement model components: RBF and scatter utilities in models/layers.py, a PaiNN-style message-passing network in models/painn.py, and a high-level MLIPEnergyModel with predict_energy_forces in models/model.py.
  • Add training engine mlipx.training.engine.Trainer that handles optimization, LR scheduling, AMP, checkpointing, early stopping, and a CLI entrypoint src/mlipx/train.py plus evaluate.py and infer.py for evaluation and inference workflows.
  • Provide ASE integration with src/mlipx/ase_ext/calculator.py exposing build_calculator, convenience scripts under scripts/ (train.sh, evaluate.sh, infer.sh, make_demo_dataset.py, package.sh), project __init__, and utility modules for metrics and seeding.
  • Remove the old monolithic mlip.py teaching demo and add unit tests in tests/ to cover the data pipeline and a smoke training/model forward pass.

Testing

  • Added unit tests tests/test_data_pipeline.py and tests/test_model_and_train.py that exercise dataset reading, graph collate, model forward and force prediction, and a one-epoch smoke training run using a temporary extxyz dataset.
  • Ran the test suite with pytest -q and the tests completed successfully.

Codex Task

@yeze5421
yeze5421 marked this pull request as draft April 12, 2026 06:26
@yeze5421
yeze5421 marked this pull request as ready for review April 12, 2026 06:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant