Contributing

We want to make mlipx a truly comprehensive tool for evaluating MLIP models. That means making it easy for you to extend! Whether you’re adding new nodes, recipes, or integrating new MLIP models, this guide will walk you through the process.

For any changes to mlipx, we recommend installing the package from source. You can find detailed instructions in the From Source section.

New Nodes

Nodes are the building blocks of mlipx workflows. They represent individual computational steps in your evaluation.

To create a new node:

  1. Subclass zntrack.Node: Your new node must inherit from the zntrack.Node class and specify all required inputs and outputs. For more information have a look at the ZnTrack documentation.

  2. Implement the run method: This method will contain the core logic of your node.

You can follow the Structure Relaxtion with Custom Nodes for a practical example of writing a new node.

Once you’ve developed your node, here’s how to integrate it:

  1. Create a New Branch:

    (.venv) $ cd mlipx
(.venv) $ git checkout -b <new-branch-name>

  2. Add Dependencies (if any): If your node requires additional Python packages, use uv to add them:

    (.venv) $ uv add <dependency>

    Note

    For less common dependencies, consider adding them as an extra.

  3. Add Your Node File: Create or amend a file in the mlipx/nodes directory with your new node’s implementation.

  4. Make it Importable: Import your new node into mlipx/__init__.pyi and add it to the __all__ list. This makes your node directly importable by users (e.g., from mlipx import MyNewNode).

  5. Commit and Pull Request: Finally, commit your changes and create a pull request on the mlipx GitHub repository.

Tip

We encourage you to provide metadata about the training data used for your model. This enables mlipx to inform users whether models can be compared directly, or if differences in _ab initio_ settings make comparisons unreliable.

To support this, mlipx offers a JSON Schema-based metadata format, defined at mlipx/spec/mlips-schema.json. You can install schema support in VS Code using the CLI command: mlipx install-vscode-schema.

We recommend including an mlips.yaml file in your model package at <your_package>/spec/mlips.yaml. mlipx will automatically attempt to load this file and use it to inform users about the training data behind your model during comparisons.

Training data definitions for MLIPS included in mlipx

These predefined MLIP entries are described in mlipx/spec/mlips.yaml.

mace-mpa-0:
  data:
    type: dataset
    name:
    - MPtrj
    - subsampled_Alexandria

grace-2l-omat:
  data:
    type: dataset
    name:
    - MPtrj
    - subsampled_Alexandria
    - OMat24

chgnet:
  data:
    type: dataset
    name: MPtrj

mattersim:
  data:
    type: dataset
    name: MatterSim

pet-mad:
  data:
    type: dataset
    name: MAD

mace-off:
  data:
    type: dataset
    name: SPICEv1
  metadata:
    doi: https://doi.org/10.1021/jacs.4c07099

We also provide built-in abstractions for several public datasets.

Public datasets supported by mlipx

These datasets are defined in mlipx/spec/datasets.yaml.

MPtrj:
  type: DFT
  code: VASP
  method:
    functional: PBE+U
  dispersion_correction: null
  metadata:
    doi: https://doi.org/10.6084/m9.figshare.23713842.v2


OMat24:
  type: DFT
  code: VASP
  method:
    functional: PBE+U

Alexandria:
  type: DFT
  code: VASP
  method:
    functional: PBE+U

subsampled_Alexandria:
  type: DFT
  code: VASP
  method:
    functional: PBE+U

MatterSim:
  type: DFT
  code: VASP
  method:
    functional: PBE+U
  basis_set:
    type: plane-wave
    plane_wave_cutoff_eV: 520
  metadata:
    doi: https://doi.org/10.48550/arXiv.2405.04967

MAD:
  type: DFT
  code: QuantumEspresso
  method:
    functional: PBEsol
  pseudopotential:
    name: PBEsol

OC20:
  type: DFT
  code: VASP
  method:
    functional: rPBE
  basis_set:
    type: plane-wave
    plane_wave_cutoff_eV: 350

SPICEv1:
  type: DFT
  code: PSI4
  method:
    functional: ωB97M
  dispersion_correction:
    type: DFT-D3(BJ)
  basis_set:
    type: plane-wave
    name: def2-TZVPPD
  metadata:
    doi: https://doi.org/10.1038/s41597-022-01882-6

New Recipes

Recipes are pre-defined workflows that combine multiple nodes to perform a specific evaluation task. They are designed as Jinja2 templates.

To add a new recipe:

  1. Create a Jinja2 Template: In the mlipx/recipes directory, create a new file with the .py.jinja2 extension. Structure your recipe following the examples of existing recipes.

  2. Extend the CLI: Integrate your new recipe into the mlipx command-line interface by modifying the mlipx/recipes/main.py file. We use the Typer library for our CLI, so you can refer to the existing recipes in main.py for guidance on how to add your new command.

New Models

mlipx provides a streamlined way to incorporate new MLIP models for evaluation. All available models are managed in the mlipx/recipes/models.py.jinja2 file.

Models Supported by mlipx.GenericASECalculator

If your model is compatible with the mlipx.GenericASECalculator interface, you can add it directly:

ALL_MODELS["<model-id>"] = mlipx.GenericASECalculator(
    module="<your_module>", # The Python module where your calculator class is located
    class_name="<YourCalculatorClass>", # The name of your calculator class
    device="auto", # Set to "auto" if using PyTorch and your calculator supports a 'device' argument
    kwargs={} # Any additional keyword arguments to pass to your calculator's constructor
)

Replace <model-id>, <your_module>, and <YourCalculatorClass> with your model’s specific details.

Models Not Supported by mlipx.GenericASECalculator

If your model does not fit the mlipx.GenericASECalculator interface, you’ll need to create a custom node. This new node should implement the mlipx.abc.NodeWithCalculator interface and be placed within the mlipx/recipes/models.py.jinja2 file. This ensures mlipx can properly interact with your model for evaluations.