Update to metatrain v2026.3.1 (#153)

* Update pyproject.toml

Try to see if it works with latest metatrain

* Update ASE documentation URL in conf.py

* Add deprecation warning for torch.jit.interface

* Add deprecation warning for torch.jit.set_fusion_strategy

* Update pyproject.toml with new deprecation warnings

Added deprecation warnings for NumPy and model outputs.

* Comment out deprecated warning in pyproject.toml

Comment out deprecated warning for thermalize_momenta

* Uncomment deprecation warning for thermalize_momenta

* Fixed the singular-plural convention in targets

* Added a scm versioning and relaxed the nvidia libraries version requirements

* Cleaned up the __init__.py

---------

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

Author: HowWeiBin

.gitignore

@@ -15,6 +15,7 @@ __pycache__/
 *.ipynb
 
 # Distribution / packaging
+_scm_version.py
 .Python
 build/
 develop-eggs/

docs/src/conf.py

@@ -51,7 +51,7 @@
     "python": ("https://docs.python.org/3", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "torch": ("https://docs.pytorch.org/docs/stable/", None),
-    "ase": ("https://ase-lib.org/", None),
+    "ase": ("https://docs.ase-lib.org/", None),
     "metatensor": ("https://docs.metatensor.org/latest/", None),
     "metatomic": ("https://docs.metatensor.org/metatomic/latest/", None),
     "metatrain": ("https://docs.metatensor.org/metatrain/latest/", None),

pyproject.toml

@@ -23,10 +23,10 @@ maintainers = [
 ]
 
 dependencies = [
-    "metatrain>=2026.2,<2026.3",
+    "metatrain>=2026.3.1, <2026.4",
     "metatomic-ase",
-    "nvalchemi-toolkit-ops==0.3.0",
-    "warp-lang==1.13.0",
+    "nvalchemi-toolkit-ops",
+    "warp-lang",
     "huggingface_hub",
     "hf_xet",
     "packaging",
@@ -60,18 +60,23 @@ repository = "https://github.com/lab-cosmo/upet"
 [build-system]
 requires = [
     "setuptools >= 68",
+    "setuptools_scm>=8",
     "wheel",
 ]
 build-backend = "setuptools.build_meta"
 
 [tool.setuptools.packages.find]
 where = ["src"]
 
-[tool.setuptools.dynamic]
-version = {attr = "upet.__version__"}
+[tool.setuptools_scm]
+version_file = "src/upet/_scm_version.py"
+
+[tool.check-manifest]
+ignore = ["src/upet/_scm_version.py"]
 
 [tool.ruff]
 line-length = 88
+exclude = ["src/upet/_scm_version.py"]
 
 [tool.ruff.lint]
 select = ["E", "F", "B", "I"]
@@ -115,6 +120,8 @@ filterwarnings = [
     "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`.",
+    "ignore: `torch.jit.interface` is deprecated. Please use `torch.compile` instead.",
+    "ignore: `torch.jit.set_fusion_strategy` is deprecated. Please use `torch.compile` instead.",
     # PET-MAD v1.0 deprecation warning
     "ignore:.*is deprecated in favor of the newer PET-MAD-1.5.*:DeprecationWarning",
     # NVAlchemiOps Deprecation warning
@@ -124,4 +131,11 @@ filterwarnings = [
     "ignore:`compute_requested_neighbors_from_options` is deprecated.*:UserWarning",
     # HF Xet deprecation warning
     "ignore:hf_xet.download_files\\(\\) is deprecated. Use XetSession\\(\\):DeprecationWarning",
+    # Numpy warning
+    "ignore: Setting the shape on a NumPy array has been deprecated in NumPy 2.5.",
+    # MD
+    "ignore: Use thermalize_momenta",
+    # New metatomic release warnings
+    "ignore: the 'features' output name is deprecated:UserWarning",
+    "ignore: the 'non_conservative_forces' output name is deprecated:UserWarning",
 ]

src/upet/__init__.py

@@ -1,33 +1,15 @@
-__version__ = "0.2.5"
-
 import warnings
 
 import torch
-from warp._src import utils as wp_utils
 
 from ._models import get_upet, list_upet, save_upet
+from ._scm_version import __version__  # noqa: F401
 
 
-# hides a harmless warning from nvalchemi's neighbor list implmentation
 warnings.filterwarnings(
-    "ignore",
-    category=UserWarning,
-    message="The .grad attribute of a Tensor that is not a leaf Tensor",
+    "ignore", category=DeprecationWarning, message="warp.config.quiet is deprecated"
 )
 
-# we want to suppress a further warning from nvalchemi's usage of warp
-# warp uses an internal warn() helper: we wrap it.
-_orig_warn = wp_utils.warn
-
-
-def _warn_filtered(message, category=None, stacklevel=1):
-    if category is DeprecationWarning and "warp.vec" in str(message):
-        return
-    return _orig_warn(message, category=category, stacklevel=stacklevel)
-
-
-wp_utils.warn = _warn_filtered
-
 # Disable static fusion. Besides the fact that atomistic batches have variable
 # sizes, statically fused CUDA kernels cannot allocate new tensors at runtime,
 # causing "Global alloc not supported yet" errors (cuda 13+) at the time of writing

src/upet/calculator.py

@@ -158,7 +158,7 @@ def __init__(
         if non_conservative:
             selected_variant = None if variants is None else variants.get("energy")
             variant_postfix = f"/{selected_variant}" if selected_variant else ""
-            nc_forces_key = "non_conservative_forces" + variant_postfix
+            nc_forces_key = "non_conservative_force" + variant_postfix
             nc_stress_key = "non_conservative_stress" + variant_postfix
             if nc_forces_key not in model_outputs or nc_stress_key not in model_outputs:
                 raise NotImplementedError(

src/upet/explore/_explorer.py

@@ -23,7 +23,7 @@ class MADExplorer(torch.nn.Module):
     :param output_dim: target low dimensionality for the projected embeddings
     :param device: cpu or cuda
     :param features_output: key to access the PET-MAD feature output.
-        mtt::aux::energy_last_layer_features is used by default
+        mtt::aux::energy_last_layer_feature is used by default
     """
 
     def __init__(
@@ -57,17 +57,17 @@ def forward(
         outputs: Dict[str, mta.ModelOutput],
         selected_atoms: Optional[mts.Labels],
     ) -> Dict[str, mts.TensorMap]:
-        if list(outputs.keys()) != ["features"]:
+        if list(outputs.keys()) != ["feature"]:
             raise ValueError(
                 f"`outputs` keys ({', '.join(outputs.keys())}) contain unsupported "
-                "keys. Only 'features' is supported"
+                "keys. Only 'feature' is supported"
             )
 
         systems = [s.to(self.dtype, self.device) for s in systems]
 
-        per_atom = outputs["features"].per_atom
+        sample_kind = outputs["feature"].sample_kind
         pet_requested_outputs = {
-            self.features_output: mta.ModelOutput(per_atom=per_atom)
+            self.features_output: mta.ModelOutput(sample_kind=sample_kind)
         }
 
         if selected_atoms is not None:
@@ -113,7 +113,7 @@ def forward(
             blocks=[block],
         )
 
-        return {"features": tensor_map}
+        return {"feature": tensor_map}
 
     def _get_features(
         self,
@@ -134,7 +134,7 @@ def _get_features(
         if selected_atoms is not None:
             features = mts.slice(features, "samples", selected_atoms)
 
-        if outputs[self.features_output].per_atom:
+        if outputs[self.features_output].sample_kind == "atom":
             mean = mts.mean_over_samples(features, "atom")
             mean_vals = torch.cat([block.values for block in mean.blocks()], dim=0)
 

src/upet/explore/_featurizer.py

@@ -113,7 +113,7 @@ def __init__(
         explorer = MADExplorer(petmad.module, device=device)
         explorer.load_checkpoint(petmad_explorer_path)
 
-        outputs = {"features": mta.ModelOutput(per_atom=True)}
+        outputs = {"feature": mta.ModelOutput(sample_kind="atom")}
         self.dtype = torch.float64
 
         capabilities = mta.ModelCapabilities(
@@ -155,7 +155,7 @@ def __call__(self, frames, environments):
 
         options = mta.ModelEvaluationOptions(
             length_unit=self.length_unit,
-            outputs={"features": mta.ModelOutput(per_atom=True)},
+            outputs={"feature": mta.ModelOutput(sample_kind="atom")},
             selected_atoms=selected_atoms,
         )
 
@@ -169,7 +169,7 @@ def __call__(self, frames, environments):
                 check_consistency=self.check_consistency,
             )
             outputs.append(
-                batch_outputs["features"].block().values.detach().cpu().numpy()
+                batch_outputs["feature"].block().values.detach().cpu().numpy()
             )
 
         return np.concatenate(outputs)