Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,9 @@ dmypy.json
# Cython debug symbols
cython_debug/

# hatch-vcs generated version file
mlipx/_version.py

# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
Expand Down
1 change: 1 addition & 0 deletions docs/source/concept.rst
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ providing a quick overview of all past evaluations.
concept/data
concept/models
concept/recipes
concept/serve
concept/zntrack
concept/zndraw
concept/metrics
Expand Down
51 changes: 46 additions & 5 deletions docs/source/concept/models.rst
Original file line number Diff line number Diff line change
Expand Up @@ -72,22 +72,63 @@ The :code:`SevenCalc` class, used in the example above, is defined in :code:`src

For more details, refer to the :ref:`custom_nodes` section.

.. _serve-integration:

Serve Integration
-----------------

The dependencies of MLIPs can be incompatible, making it impossible to run multiple MLIPs in the same environment.
To address this, ``mlipx`` can circumvent this using ``uv`` (https://docs.astral.sh/uv/) in combination with ZeroMQ (https://zeromq.org/) to automatically serve models from different environments.

To make a model available for serving, you need to make it installable through an extra for ``mlipx``.
For this, you need to configure a new extra in the ``mlipx`` ``pyproject.toml`` using ``[tool.uv]
conflicts`` to specify dependencies that cannot be installed in the same environment.

Additionally, you need to include this extra in the model definition.

.. code-block:: python

import mlipx

ALL_MODELS["mace-mpa-0"] = mlipx.GenericASECalculator(
module="mace.calculators",
class_name="mace_mp",
device="auto",
extra=["mace"], # mlipx extras for this model
)

**Using served models:**

.. code-block:: python

# Environment variable (global control)
import os
os.environ["MLIPX_USE_SERVE"] = "true"

# Or explicit per-model control
calc = model.get_calculator(use_serve=True)

See :ref:`serve` for details on starting brokers, managing workers, and DVC integration.

.. _update-frames-calc:

Updating Dataset Keys
---------------------

In some cases, models may need to be defined to convert existing dataset keys into the format :code:`mlipx` expects.
For example, you may need to provide isolated atom energies or convert data where energies are stored as :code:`atoms.info['DFT_ENERGY']`
and forces as :code:`atoms.arrays['DFT_FORCES']`.
In some cases, you may need to convert existing dataset keys into the format :code:`mlipx` expects.
For example, your dataset may store energies as :code:`atoms.info['DFT_ENERGY']` and forces as :code:`atoms.arrays['DFT_FORCES']`.

Here’s how to define a model for such a scenario:
Use :code:`mlipx.UpdateFramesCalc` to remap these keys:

.. code-block:: python

import mlipx

REFERENCE = mlipx.UpdateFramesCalc(
# Remap existing dataset keys to mlipx format
data_remapper = mlipx.UpdateFramesCalc(
results_mapping={"energy": "DFT_ENERGY", "forces": "DFT_FORCES"},
info_mapping={mlipx.abc.ASEKeys.isolated_energies.value: "isol_ene"},
)

This is typically used in the :code:`metrics` recipe when comparing MLIP predictions against pre-computed reference data (e.g., DFT results).
The generated :code:`main.py` includes a commented section showing how to configure this.
Loading
Loading