Merge pull request #30 from lab-cosmo/pre-release-uq

Pre-release the UQ update

Author: abmazitov

.github/workflows/release.yml

@@ -0,0 +1,41 @@
+name: Release
+
+on:
+  push:
+    tags: ["*"]
+
+jobs:
+  build:
+    name: Build distribution
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/project/pet-mad
+    permissions:
+      id-token: write
+      contents: write
+
+    steps:
+    - uses: actions/checkout@v4
+      with:
+        fetch-depth: 0
+    - name: setup Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: "3.13"
+    - run: python -m pip install tox
+    - name: Build package
+      run: tox -e build
+    - name: Publish distribution to PyPI
+      if: startsWith(github.ref, 'refs/tags/v')
+      uses: pypa/gh-action-pypi-publish@release/v1
+    - name: Publish to GitHub release
+      if: startsWith(github.ref, 'refs/tags/v')
+      uses: softprops/action-gh-release@v2
+      with:
+        files: |
+          dist/*.tar.gz
+          dist/*.whl
+        prerelease: ${{ contains(github.ref, '-rc') }}
+      env:
+        GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.gitignore

@@ -7,7 +7,6 @@ __pycache__/
 *.so
 
 # Distribution / packaging
-_version.py
 .Python
 build/
 develop-eggs/

src/pet_mad/__init__.py

@@ -1,3 +1 @@
-from ._models import get_pet_mad, save_pet_mad  # noqa: F401
-
 __version__ = "1.3.0"

src/pet_mad/_models.py

@@ -1,50 +1,40 @@
 import importlib.util
 import logging
 import warnings
+from typing import Optional
 
-from metatomic.torch import AtomisticModel, ModelMetadata
+from metatomic.torch import AtomisticModel
 from metatrain.utils.io import load_model as load_metatrain_model
+from .utils import get_metadata
 
+from packaging.version import Version
 
-METADATA = ModelMetadata(
-    name="PET-MAD",
-    description="A universal interatomic potential for advanced materials modeling",
-    authors=[
-        "Arslan Mazitov (arslan.mazitov@epfl.ch)",
-        "Filippo Bigi",
-        "Matthias Kellner",
-        "Paolo Pegolo",
-        "Davide Tisi",
-        "Guillaume Fraux",
-        "Sergey Pozdnyakov",
-        "Philip Loche",
-        "Michele Ceriotti (michele.ceriotti@epfl.ch)",
-    ],
-    references={
-        "architecture": ["https://arxiv.org/abs/2305.19302v3"],
-        "model": ["http://arxiv.org/abs/2503.14118"],
-    },
-)
-VERSIONS = ("latest", "1.1.0", "1.0.1", "1.0.0")
-BASE_URL = (
-    "https://huggingface.co/lab-cosmo/pet-mad/resolve/{}/models/pet-mad-latest.ckpt"
-)
-
-
-def get_pet_mad(*, version="latest", checkpoint_path=None) -> AtomisticModel:
+from ._version import LATEST_VERSION, AVAILABLE_VERSIONS
+
+BASE_URL = "https://huggingface.co/lab-cosmo/pet-mad/resolve/{tag}/models/pet-mad-{version}.ckpt"
+
+
+def get_pet_mad(
+    *, version: str = "latest", checkpoint_path: Optional[str] = None
+) -> AtomisticModel:
     """Get a metatomic ``AtomisticModel`` for PET-MAD.
 
-    :param version: PET-MAD version to use. Supported versions are "latest", "1.1.0",
-        "1.0.1", "1.0.0". Defaults to "latest".
+    :param version: PET-MAD version to use. Supported versions are
+        "1.1.0", "1.0.1", "1.0.0". Defaults to latest available version.
     :param checkpoint_path: path to a checkpoint file to load the model from. If
         provided, the `version` parameter is ignored.
     """
-    if version not in VERSIONS:
+    if version == "latest":
+        version = Version(LATEST_VERSION)
+    if not isinstance(version, Version):
+        version = Version(version)
+
+    if version not in [Version(v) for v in AVAILABLE_VERSIONS]:
         raise ValueError(
-            f"Version {version} is not supported. Supported versions are {VERSIONS}"
+            f"Version {version} is not supported. Supported versions are {AVAILABLE_VERSIONS}"
         )
 
-    if version == "1.0.0":
+    if version == Version("1.0.0"):
         if not importlib.util.find_spec("pet_neighbors_convert"):
             raise ImportError(
                 f"PET-MAD v{version} is now deprecated. Please consider using the "
@@ -60,9 +50,7 @@ def get_pet_mad(*, version="latest", checkpoint_path=None) -> AtomisticModel:
         path = checkpoint_path
     else:
         logging.info(f"Downloading PET-MAD model version: {version}")
-        path = BASE_URL.format(
-            f"v{version}" if version not in ("latest", "1.1.0") else "main"
-        )
+        path = BASE_URL.format(tag=f"v{version}", version=f"v{version}")
 
     with warnings.catch_warnings():
         warnings.filterwarnings(
@@ -71,25 +59,31 @@ def get_pet_mad(*, version="latest", checkpoint_path=None) -> AtomisticModel:
         )
         model = load_metatrain_model(path)
 
-    return model.export(METADATA)
+    metadata = get_metadata(version)
+    return model.export(metadata)
 
 
-def save_pet_mad(*, version="latest", checkpoint_path=None, output=None):
+def save_pet_mad(*, version: str = "latest", checkpoint_path=None, output=None):
     """
     Save the PET-MAD model to a TorchScript file (``pet-mad-xxx.pt``). These files can
     be used with LAMMPS and other tools to run simulations without Python.
 
-    :param version: PET-MAD version to use. Supported versions are "latest", "1.1.0",
-        "1.0.1", "1.0.0". Defaults to "latest".
+    :param version: PET-MAD version to use. Supported versions are "1.1.0",
+        "1.0.1", "1.0.0". Defaults to the latest version.
     :param checkpoint_path: path to a checkpoint file to load the model from. If
         provided, the `version` parameter is ignored.
     :param output: path to use for the output model, defaults to
         ``pet-mad-{version}.pt`` when using a version, or the checkpoint path when using
         a checkpoint.
     """
+    if version == "latest":
+        version = Version(LATEST_VERSION)
+    if not isinstance(version, Version):
+        version = Version(version)
+
     extensions_directory = None
-    if version == "1.0.0":
-        logging.info("putting TorchScript extensions in `extensions/`")
+    if version == Version("1.0.0"):
+        logging.info("Putting TorchScript extensions in `extensions/`")
         extensions_directory = "extensions"
 
     model = get_pet_mad(version=version, checkpoint_path=checkpoint_path)
@@ -101,4 +95,4 @@ def save_pet_mad(*, version="latest", checkpoint_path=None, output=None):
             raise
 
     model.save(output, collect_extensions=extensions_directory)
-    logging.info(f"saved pet-mad model to {output}")
+    logging.info(f"Saved PET-MAD model to {output}")

src/pet_mad/_version.py

@@ -0,0 +1,2 @@
+LATEST_VERSION = "1.1.0"
+AVAILABLE_VERSIONS = ["1.1.0", "1.0.1", "1.0.0"]

src/pet_mad/calculator.py

@@ -5,7 +5,10 @@
 from metatomic.torch.ase_calculator import MetatomicCalculator
 from platformdirs import user_cache_dir
 
+from packaging.version import Version
+
 from ._models import get_pet_mad
+from ._version import LATEST_VERSION
 
 
 class PETMADCalculator(MetatomicCalculator):
@@ -18,30 +21,37 @@ def __init__(
         version: str = "latest",
         checkpoint_path: Optional[str] = None,
         *,
-        non_conservative=False,
-        check_consistency=False,
-        device=None,
+        check_consistency: bool = False,
+        device: Optional[str] = None,
+        non_conservative: bool = False,
     ):
         """
-        :param version: PET-MAD version to use. Supported versions are "latest", "1.1.0",
-            "1.0.1", "1.0.0". Defaults to "latest".
+        :param version: PET-MAD version to use. Supported versions are
+            "1.1.0", "1.0.1", "1.0.0". Defaults to latest available version.
         :param checkpoint_path: path to a checkpoint file to load the model from. If
             provided, the `version` parameter is ignored.
-        :param non_conservative: whether to use the non-conservative regime of forces and
-            stresses prediction. Defaults to False.
         :param check_consistency: should we check the model for consistency when
             running, defaults to False.
         :param device: torch device to use for the calculation. If `None`, we will try
             the options in the model's `supported_device` in order.
+        :param non_conservative: whether to use the non-conservative regime of forces
+            and stresses prediction. Defaults to False. Only available for PET-MAD
+            version 1.1.0 or higher.
+
         """
 
+        if version == "latest":
+            version = Version(LATEST_VERSION)
+        if not isinstance(version, Version):
+            version = Version(version)
+
         model = get_pet_mad(version=version, checkpoint_path=checkpoint_path)
 
         cache_dir = user_cache_dir("pet-mad", "metatensor")
         os.makedirs(cache_dir, exist_ok=True)
 
         extensions_directory = None
-        if version == "1.0.0":
+        if version == Version("1.0.0"):
             extensions_directory = "extensions"
 
         pt_path = cache_dir + f"/pet-mad-{version}.pt"
@@ -56,9 +66,9 @@ def __init__(
 
         super().__init__(
             pt_path,
+            additional_outputs={},
             extensions_directory=extensions_directory,
             check_consistency=check_consistency,
             device=device,
             non_conservative=non_conservative,
-            additional_outputs={},
         )

src/pet_mad/utils.py

@@ -0,0 +1,23 @@
+from metatomic.torch import ModelMetadata
+
+
+def get_metadata(version: str):
+    return ModelMetadata(
+        name=f"PET-MAD v{version}",
+        description="A universal interatomic potential for advanced materials modeling",
+        authors=[
+            "Arslan Mazitov (arslan.mazitov@epfl.ch)",
+            "Filippo Bigi",
+            "Matthias Kellner",
+            "Paolo Pegolo",
+            "Davide Tisi",
+            "Guillaume Fraux",
+            "Sergey Pozdnyakov",
+            "Philip Loche",
+            "Michele Ceriotti (michele.ceriotti@epfl.ch)",
+        ],
+        references={
+            "architecture": ["https://arxiv.org/abs/2305.19302v3"],
+            "model": ["http://arxiv.org/abs/2503.14118"],
+        },
+    )

tests/test_basic_usage.py

@@ -1,7 +1,28 @@
 from pet_mad.calculator import PETMADCalculator
+from pet_mad._models import get_pet_mad, save_pet_mad
 from ase.build import bulk
+import os
 import pytest
 
+VERSIONS = ("1.1.0", "1.0.1")
+
+@pytest.mark.parametrize(
+    "version",
+    VERSIONS,
+)
+def test_get_pet_mad(version):
+    model = get_pet_mad(version=version)
+    assert model.metadata().name == f"PET-MAD v{version}"
+
+
+@pytest.mark.parametrize(
+    "version",
+    VERSIONS,
+)
+def test_save_pet_mad(version, monkeypatch, tmp_path):
+    monkeypatch.chdir(tmp_path)
+    save_pet_mad(version=version, output=f"pet-mad-{version}.pt")
+    assert os.path.exists(f"pet-mad-{version}.pt")
 
 def test_basic_usage():
     atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
@@ -12,11 +33,7 @@ def test_basic_usage():
 
 @pytest.mark.parametrize(
     "version",
-    [
-        "latest",
-        "1.1.0",
-        "1.0.1",
-    ],
+    VERSIONS,
 )
 def test_version(version):
     atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")

tests/test_metadata.py

@@ -0,0 +1,8 @@
+from pet_mad.utils import get_metadata
+
+def test_get_metadata():
+    metadata = get_metadata("1.1.0")
+    assert metadata.name == "PET-MAD v1.1.0"
+    assert metadata.description == "A universal interatomic potential for advanced materials modeling"
+    assert metadata.authors == ["Arslan Mazitov (arslan.mazitov@epfl.ch)", "Filippo Bigi", "Matthias Kellner", "Paolo Pegolo", "Davide Tisi", "Guillaume Fraux", "Sergey Pozdnyakov", "Philip Loche", "Michele Ceriotti (michele.ceriotti@epfl.ch)"]
+    assert metadata.references == {"architecture": ["https://arxiv.org/abs/2305.19302v3"], "model": ["http://arxiv.org/abs/2503.14118"]}