Developer Guide#

This guide is for contributors who modify Enerzyme core: adding models, extending tasks, changing data pipelines, or maintaining docs and tests. It does not repeat end-user tutorials in Getting Started or configuration reference in User Guide. For class-level API details, see API Documentation.

Development setup#

Editable install#

From the repository root:

pip install -e .

This registers the enerzyme console script via setup.py and exposes the package for local development.

Environment files#

Three dependency contexts matter:

  • Runtime — setup.py install_requires (NumPy, PyTorch, ASE, RDKit, Lightning, etc.)

  • Development — requirements-dev.yaml (conda env for day-to-day coding)

  • Documentation — docs/requirements.yaml (Sphinx, pydata theme, editable install for autodoc)

For a first-time contributor setup, follow Installation, then install in editable mode as above.

Optional dependencies#

Not every contributor needs all optional stacks. Install only what your change touches:

  • PhysNet parity tests — TensorFlow 1.x compatibility mode and the reference PhysNet package

  • NequIP / XPaiNN — nequip, XequiNet and their transitive deps

  • PLUMED workflows — py-plumed and a PLUMED-enabled build

  • QM annotation — TeraChem executable/license, RDKit formal charges

  • Bond assignment — QuantumPDB-style PDB inputs and template SDFs

  • Server mode — Flask/Waitress (already in core install_requires)

Smoke test after setup:

python -c "import enerzyme; print(enerzyme.__file__)"
enerzyme -h
enerzyme train -h
enerzyme predict -h
enerzyme simulate -h

Repository map#

CLI and command wrappers#

enerzyme/cli.py defines argparse subcommands and dispatches to thin wrapper classes:

  • train → enerzyme/train.py (FFTrain)

  • predict → enerzyme/predict.py (FFPredict)

  • simulate → enerzyme/simulate.py (FFSimulate)

  • extract → enerzyme/extract.py (FFExtract, reuses FFPredict)

  • collect → enerzyme/collect.py (FFCollect)

  • annotate → enerzyme/annotate.py (QMAnnotate)

  • bond → enerzyme/bond/bond.py

  • listen / request / kill → enerzyme/listen.py, enerzyme/request.py, HTTP shutdown helper in cli.py

Core packages#

  • enerzyme/data/ — dataset loading, standard fields, transforms, HDF5 preload cache (DataHub, FieldDataset)

  • enerzyme/models/ — architecture cores, layer stack, ModelHub, loss registration

  • enerzyme/tasks/ — training, metrics, simulation, extraction, server, splitting, active-learning picking

  • enerzyme/qm/ — QM driver adapters for annotate

  • enerzyme/bond/ — PDB bond-order assignment

  • enerzyme/utils/ — logging, YAML I/O (YamlHandler)

Reference YAML configs live in enerzyme/config/. After training, the resolved config.yaml in the output directory is the canonical model config for predict and simulate.

Runtime architecture#

Enerzyme is YAML-driven. User-facing behavior is defined by config sections; Python classes implement those sections.

enerzyme <command> -c config.yaml -o out/
    -> YamlHandler.read_yaml()
    -> DataHub / Trainer / task-specific wrapper
    -> ModelHub (for train) or loaded checkpoints (for predict/simulate/extract)
    -> artifacts under out/

Training path (FFTrain):

config.yaml
    -> DataHub(dump_dir=out/, **Datahub)
    -> Trainer(out_dir, Metric, **Trainer)
    -> ModelHub(datahub, trainer, **Modelhub)
    -> FF_single / FF_committee.train() or .active_learn()
    -> out/config.yaml, processed_dataset_<hash>/, FFxx/, logs/

Prediction and simulation load the saved model config (-mc or default model_dir/config.yaml), rebuild active models, and run task code without retraining.

Note

Treat YAML schema and output directory layout as public contracts. External workflows (including Enerzymette) depend on checkpoint names, config.yaml, and per-iteration folder names documented in Active Learning Workflows.

Configuration development rules#

YAML handling#

enerzyme/utils/config_handler.py loads YAML into addict.Dict objects. Training writes the resolved config back to out/config.yaml at startup.

When adding or renaming a config field:

  1. Add a default or example in the relevant file under enerzyme/config/

  2. Wire the field in the consumer (wrapper class or tasks/* module)

  3. Document semantics in User Guide (user-facing) or this guide (developer-facing)

  4. Add a smoke test or minimal pytest case if behavior is non-trivial

Config section ownership#

  • Datahub — consumed by DataHub; may be partially overridden in predict/extract configs

  • Modelhub — consumed by ModelHub (train) or model rebuild (predict/simulate/extract)

  • Trainer / Metric — training loop, early stopping, committee, active learning

  • Simulation / System — enerzyme/tasks/simulator.py

  • Extractor — enerzyme/tasks/extractor.py

  • Supplier / QMDriver — enerzyme/annotate.py

Backward compatibility#

  • Published releases — avoid breaking existing YAML keys or artifact paths without a migration note

  • Unreleased :code:`devel` branch — schema changes may land directly, but update reference YAML and docs in the same PR

Data pipeline extension points#

Standard fields#

Registered in enerzyme/data/datatype.py (N, Ra, Za, E, Fa, Qa, Q, M2, etc.). See Units and Standard Fields.

Custom fields#

Register under Datahub.fields with is_atomic: true/false. Map raw dataset keys via features and targets.

Transforms#

enerzyme/data/transform.py applies dataset-level and global transforms during preload. Common developer touchpoints:

  • negative_gradient — sign flip for QC gradients

  • atomic_energy / total_energy_normalization — energy offsets

Preload cache#

SingleDataHub hashes data path, neighbor-list mode, and transform settings into processed_dataset_<hash>/ with pre_transformed.hdf5 and datahub.yaml. Changing hash inputs invalidates the cache — document this when adding transform keys.

Checklist: new data field or transform#

  1. Register type in datatype.py if it is a new standard name

  2. Ensure loader supports the source format (pickle, npz, or HDF5 cache)

  3. Add transform logic in transform.py if preprocessing is required

  4. Update train.yaml example and /user_guide/data/datahub_reference

  5. Run enerzyme collect -c <yaml> -o <out> to validate mapping without training

Model development#

Layer stack#

enerzyme/models/ff.py builds models from an ordered layers list:

  • get_ff_core(architecture) returns Core, DEFAULT_BUILD_PARAMS, DEFAULT_LAYER_PARAMS

  • build_model() instantiates each layer by name from enerzyme/models/layers/

  • Physics and readout layers (ChargeConservation, ElectrostaticEnergy, EnergyReduce, Force, ShallowEnsembleReduce, etc.) are composed around the architecture Core

ModelHub creates FF_single or FF_committee per active entry in internal_FFs / external_FFs.

Adding an internal architecture#

  1. Create enerzyme/models/<name>/ with core.py, __init__.py, and default params

  2. Register in get_ff_core() inside ff.py

  3. Add a minimal Modelhub block to enerzyme/config/train.yaml (or a dedicated example YAML)

  4. Support required targets (E, Fa, optional Qa, M2) and loss/metric weights

  5. Add tests under test/ (layer parity or forward-pass smoke)

  6. Document in Architecture Catalog

Adding an external wrapper#

External models (NequIP, XPaiNN) live under enerzyme/models/nequip/ and enerzyme/models/xpainn/. Requirements:

  • Lazy or guarded imports so core install still works

  • Add package names to autosummary_mock_imports in docs/conf.py if autodoc cannot import them on RTD

  • Clearly state optional dependencies in User Guide and Getting Started

Checkpoint resolution#

get_pretrain_path() in modelhub.py resolves best vs last, version suffixes, and committee member ranks. Do not rename checkpoint files without updating this logic and downstream docs.

Task and CLI development#

Adding or changing a subcommand#

  1. Parser — add subparser and arguments in cli.py get_parser()

  2. Dispatch — branch in main() and implement <command>(args)

  3. Wrapper — thin class in top-level enerzyme/<task>.py that reads YAML and calls tasks/*

  4. Config — example YAML in enerzyme/config/ if the command is config-driven

  5. Docs — Getting Started tutorial (if user-facing) + User Guide reference + this guide if extension points change

Task modules#

  • trainer.py — training loop, resume, EMA, Lightning multi-GPU, active learning hooks

  • simulator.py — ASE tasks (sp, opt, scan, md, neb, plumed, plumed_scan)

  • extractor.py — uncertainty-based fragment picking

  • server.py — HTTP prediction server used by listen

  • splitter.py — dataset partitions including withheld for internal AL

  • picker.py — active-learning sample selection

Plugin patches#

FFSimulate loads -cp (calculator) and -pp (PLUMED generator) via importlib from user-supplied .py files.

  • Calculator patches — expose a factory function (e.g. get_uma_calculator) referenced by external_calculator.name.

  • PLUMED patches — expose a PlumedConfigGenerator subclass; YAML sets plumed_config_generator.name and method. Enerzyme also accepts legacy callables, but new Enerzymette plugins follow the class-based contract.

See Enhanced Sampling and PLUMED and Enerzymette Integration.

Stable artifact paths#

Avoid renaming without strong reason:

  • out/config.yaml

  • out/processed_dataset_<hash>/

  • out/FF<id>-<arch>[-suffix]/best/ and last/

  • Simulation outputs such as md.traj.xyz, plumed.traj.xyz, neb.xyz

Testing strategy#

Test layout#

Tests live in test/:

  • test_scatter.py — torch_scatter equivalence (CPU/GPU parametrized)

  • test_spookynet.py — SpookyNet layer and forward tests

  • test_physnet.py — PhysNet parity against reference TensorFlow implementation (heavy optional stack)

  • test_scatter_speed.py — performance-oriented scatter checks

Suggested commands#

python -m pytest test/test_scatter.py -q
python -m pytest test/test_spookynet.py -q

PhysNet parity tests require TensorFlow and the external PhysNet reference code — skip unless you are modifying PhysNet layers.

What to add for a PR#

  • Bug fix — regression test when feasible

  • New layer or architecture — at least forward-pass or numerical parity test

  • New CLI flag — argparse help text + docs; optional integration smoke script

  • Config schema change — update reference YAML and run enerzyme collect or a minimal train job locally

Documentation workflow#

Local Sphinx build#

conda env create -f docs/requirements.yaml   # once
conda activate docs_Enerzyme                 # env name from file
sphinx-build -b html docs docs/_build/html

Open docs/_build/html/index.html and confirm the Developer Guide card and toctree link work.

Sphinx configuration#

  • docs/conf.py — extensions, autosummary, mock imports for optional deps

  • .readthedocs.yaml — RTD build uses docs/requirements.yaml

When user-visible behavior changes, update in the same change set:

RST style notes#

  • Prefer bullet lists over wide grid tables (malformed tables break sphinx-build)

  • Section underlines must be at least as long as the title text

  • Cross-link with :doc:`/path` rather than hard-coded HTML paths

Contribution checklist#

New YAML field#

  • [ ] Default in enerzyme/config/*.yaml

  • [ ] Read in wrapper/task code

  • [ ] User Guide section updated

  • [ ] Smoke command or test

New internal model#

  • [ ] get_ff_core() registration

  • [ ] Example train config

  • [ ] Tests in test/

  • [ ] Architecture catalog entry

New task or simulation mode#

  • [ ] simulator.py (or relevant task) implementation

  • [ ] Example YAML

  • [ ] Getting Started + User Guide workflow pages

  • [ ] Artifact names documented

New QM driver or supplier#

  • [ ] enerzyme/qm/ adapter

  • [ ] annotate.yaml example

  • [ ] Dependency and status notes (ORCA/PySCF/Psi4 if partial)

Docs-only change#

  • [ ] sphinx-build passes

  • [ ] No broken :doc: references

Out of scope for V1#

This Developer Guide intentionally does not duplicate:

When to split into sub-pages#

Keep this single page while the contributor surface is still evolving. Split into docs/developer_guide/ when any section grows past ~200 lines or needs its own deep dive (for example, a dedicated “Adding a new architecture” cookbook). The entry docs/developer_guide.rst would then become a short overview plus layered toctree, mirroring User Guide.