Add list_upet() function (#118)

---------

Co-authored-by: Arslan Mazitov <arslan.mazitov@phystech.edu>

Author: ceriottm

README.md

@@ -130,10 +130,10 @@ UPET integrates with the following atomistic simulation engines:
 
 - **Atomic Simulation Environment (ASE)**
 - **LAMMPS** (including KOKKOS support)
+- **GROMACS**
 - **i-PI**
 - **TorchSim**
 - **OpenMM** (coming soon)
-- **GROMACS** (coming soon)
 
 ## Usage
 
@@ -175,6 +175,26 @@ These ASE methods are ideal for single-structure evaluations, but they are
 inefficient for the evaluation on a large number of pre-defined structures. To
 perform efficient batched evaluation in that case, read [here](docs/README_BATCHED.md).
 
+If the `version` argument is not specified, the latest available version of the model
+will be downloaded and used by default.
+
+```python
+from upet.calculator import UPETCalculator
+
+calculator = UPETCalculator(model="pet-mad-s", device="cpu") # uses the latest version of the PET-MAD-S model by default
+```
+
+You can get the list of available versions for a given model using the
+`list_upet` function:
+
+```python
+from upet import list_upet
+
+list_upet(model="pet-mad", size="s") # for PET-MAD model of size S
+list_upet(model="pet-mad") # for all PET-MAD models of all sizes
+list_upet() # for all available UPET models
+```
+
 #### Non-conservative (direct) forces and stresses prediction
 
 UPET models also support the direct prediction of forces and stresses. In that case,

pyproject.toml

@@ -106,8 +106,9 @@ filterwarnings = [
     # metatrain checkpoint upgrade warning
     "ignore: trying to upgrade an old model checkpoint with unknown version, this might fail and require manual modifications",
     # TorchScript deprecation warnings
-    "ignore: `torch.jit.script` is deprecated. Please switch to `torch.compile` or `torch.export`.",
-    "ignore: `torch.jit.load` is deprecated. Please switch to `torch.export`.",
+    "ignore: `torch.jit.script` is deprecated.*:DeprecationWarning",
+    "ignore: `torch.jit.save` is deprecated.*:DeprecationWarning",
+    "ignore: `torch.jit.load` is deprecated.*:DeprecationWarning",
     "ignore: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.",
     # PET-MAD v1.0 deprecation warning
     "ignore:.*is deprecated in favor of the newer PET-MAD-1.5.*:DeprecationWarning",

src/upet/__init__.py

@@ -5,7 +5,7 @@
 import torch
 from warp._src import utils as wp_utils
 
-from ._models import get_upet, save_upet
+from ._models import get_upet, list_upet, save_upet
 
 
 # hides a harmless warning from nvalchemi's neighbor list implmentation
@@ -33,4 +33,4 @@ def _warn_filtered(message, category=None, stacklevel=1):
 # causing "Global alloc not supported yet" errors (cuda 13+) at the time of writing
 torch.jit.set_fusion_strategy([("DYNAMIC", 10)])
 
-__all__ = ["get_upet", "save_upet"]
+__all__ = ["get_upet", "list_upet", "save_upet"]

src/upet/_models.py

@@ -36,6 +36,20 @@ def _get_upet_repo_files() -> List[str]:
     return [f[7:] for f in repo_files if f.startswith("models/")]
 
 
+def get_available_models() -> List[str]:
+    """Get all available base model names from the HuggingFace repository.
+
+    :return: Sorted list of base model names (e.g., ["pet-mad", "pet-omat", ...])
+    """
+    files = _get_upet_repo_files()
+    models = set()
+    for f in files:
+        match = CHECKPOINT_NAME_PATTERN.match(f)
+        if match:
+            models.add(match.group("model"))
+    return sorted(models)
+
+
 def get_sizes_for_model(model: str) -> List[str]:
     """Get all available sizes for a given model from the cached repo files.
 
@@ -278,6 +292,53 @@ def save_upet(
     logging.info(f"Saved UPET model to {output}")
 
 
+def list_upet(
+    *,
+    model: Optional[str] = None,
+    size: Optional[str] = None,
+    print_summary: bool = True,
+) -> List[dict]:
+    """List available UPET models, sizes, and versions.
+
+    When called without arguments, returns all available model/size/version
+    combinations. When ``model`` is given, filters to that model. When both
+    ``model`` and ``size`` are given, filters to that specific combination.
+
+    :param model: Base model name (e.g., "pet-mad", "pet-omat"). If ``None``,
+        lists all available models.
+    :param size: Model size (e.g., "s", "m", "l"). If ``None`` and ``model`` is
+        given, lists all sizes for that model.
+    :param print_summary: Whether to print a human-readable summary to stdout.
+        Defaults to ``True``.
+    :return: A list of dictionaries, each with keys ``"model"``, ``"size"``,
+        and ``"version"``.
+    """
+    if model is None:
+        models = get_available_models()
+    else:
+        models = [model]
+
+    result = []
+    for m in models:
+        if size is None:
+            sizes = get_sizes_for_model(m)
+        else:
+            sizes = [size]
+        for s in sizes:
+            for v in get_versions_for_model(m, s):
+                result.append({"model": m, "size": s, "version": str(v)})
+
+    if print_summary:
+        if not result:
+            print("No UPET models found.")
+        else:
+            print("Available UPET models:")
+            for entry in result:
+                print(f"  - {entry['model']}-{entry['size']} v{entry['version']}")
+
+    return result
+
+
 BASE_URL_PET_MAD_DOS = "https://huggingface.co/lab-cosmo/pet-mad-dos/resolve/{tag}/models/pet-mad-dos-{version}.pt"
 BASE_URL_BANDGAP_MODEL = (
     "https://huggingface.co/lab-cosmo/pet-mad-dos/resolve/{tag}/models/bandgap-model.pt"

tests/upet/test_basic_usage.py

@@ -1,7 +1,15 @@
+import os
+
 import pytest
 from ase.build import bulk, molecule
 
-from upet._models import get_upet, get_versions_for_model, upet_resolve_model
+from upet._models import (
+    get_upet,
+    get_versions_for_model,
+    list_upet,
+    save_upet,
+    upet_resolve_model,
+)
 from upet._version import UPET_AVAILABLE_MODELS
 from upet.calculator import UPETCalculator
 
@@ -50,6 +58,27 @@ def test_get_upet(model_name):
         get_upet(model=model, size=size, version=version)
 
 
+def test_list_models():
+    result = list_upet(print_summary=False)
+    assert len(result) > 0
+    assert all(
+        "model" in entry and "size" in entry and "version" in entry for entry in result
+    )
+    assert any(entry["model"] == "pet-mad" for entry in result)
+
+
+def test_list_sizes_for_model():
+    result = list_upet(model="pet-mad", print_summary=False)
+    assert len(result) > 0
+    assert all(entry["model"] == "pet-mad" for entry in result)
+
+
+def test_list_versions_for_model_and_size():
+    result = list_upet(model="pet-mad", size="s", print_summary=False)
+    assert len(result) > 0
+    assert all(entry["model"] == "pet-mad" and entry["size"] == "s" for entry in result)
+
+
 @pytest.mark.parametrize("model_name", UPET_AVAILABLE_MODELS)
 def test_basic_usage(model_name):
     if "-xl" in model_name or "-l" in model_name:
@@ -72,3 +101,10 @@ def test_basic_usage(model_name):
         assert isinstance(energy, float)
         assert forces.shape == (len(atoms), 3)
         assert virial.shape == (6,)
+
+
+def test_save_upet(tmp_path):
+    output_path = str(tmp_path / "pet-mad-xs.pt")
+    save_upet(model="pet-mad", size="xs", output=output_path)
+    assert os.path.isfile(output_path)
+    assert os.path.getsize(output_path) > 0

tests/upet/test_offline_checkpoint.py

@@ -0,0 +1,17 @@
+from packaging.version import Version
+
+from upet._models import parse_checkpoint_filename
+
+
+def test_parse_standard_checkpoint_name():
+    model, size, version = parse_checkpoint_filename("pet-mad-s-v1.2.3.ckpt")
+    assert model == "pet-mad"
+    assert size == "s"
+    assert version == Version("1.2.3")
+
+
+def test_parse_non_standard_checkpoint_name():
+    model, size, version = parse_checkpoint_filename("custom_name.ckpt")
+    assert model is None
+    assert size is None
+    assert version is None