Merge pull request #1348 from lab-cosmo/metatomic-dev

Author: carlocamilloni

.ci/pytorch/requirements.txt

@@ -1,5 +1,5 @@
 # The extra url makes pip grap the torch fersion we want to use in the regtests
 --extra-index-url=https://download.pytorch.org/whl/cpu
 torch>=2.7
-metatomic-torch>=0.1.3,<0.2
-featomic-torch==0.7.0
+metatomic-torch>=0.1.13,<0.2
+featomic-torch>=0.7.4,<0.8

configure

@@ -10007,8 +10007,8 @@ $as_echo_n "checking libmetatomic without extra libs... " >&6; }
     #include <metatensor.hpp>
     #include <metatensor/torch.hpp>
     #include <metatomic/torch.hpp>
-    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1
-    #error "this code is only compatible with metatomic-torch >=0.1,<0.2"
+    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1 || METATOMIC_TORCH_VERSION_PATCH < 12
+    #error "this code is only compatible with metatomic-torch >=0.1.12,<0.2"
     #endif
     int main() {
       metatomic_torch::version();
@@ -10056,8 +10056,8 @@ $as_echo_n "checking libmetatomic with $all_LIBS... " >&6; }
     #include <metatensor.hpp>
     #include <metatensor/torch.hpp>
     #include <metatomic/torch.hpp>
-    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1
-    #error "this code is only compatible with metatomic-torch >=0.1,<0.2"
+    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1 || METATOMIC_TORCH_VERSION_PATCH < 12
+    #error "this code is only compatible with metatomic-torch >=0.1.12,<0.2"
     #endif
     int main() {
       metatomic_torch::version();
@@ -10105,8 +10105,8 @@ $as_echo_n "checking libmetatomic with -l$testlib... " >&6; }
     #include <metatensor.hpp>
     #include <metatensor/torch.hpp>
     #include <metatomic/torch.hpp>
-    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1
-    #error "this code is only compatible with metatomic-torch >=0.1,<0.2"
+    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1 || METATOMIC_TORCH_VERSION_PATCH < 12
+    #error "this code is only compatible with metatomic-torch >=0.1.12,<0.2"
     #endif
     int main() {
       metatomic_torch::version();

configure.ac

@@ -989,8 +989,8 @@ if test $libmetatomic = true ; then
     #include <metatensor.hpp>
     #include <metatensor/torch.hpp>
     #include <metatomic/torch.hpp>
-    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1
-    #error "this code is only compatible with metatomic-torch >=0.1,<0.2"
+    #if METATOMIC_TORCH_VERSION_MAJOR != 0 || METATOMIC_TORCH_VERSION_MINOR != 1 || METATOMIC_TORCH_VERSION_PATCH < 12
+    #error "this code is only compatible with metatomic-torch >=0.1.12,<0.2"
     #endif
     int main() {
       metatomic_torch::version();

regtest/metatomic/rt-basic/cv.py

@@ -32,30 +32,31 @@ class TestCollectiveVariable(torch.nn.Module):
     CV^2 are returned.
     """
 
-    def __init__(self, cutoff, multiple_properties):
+    def __init__(self, cutoff, multiple_properties, feature_key="feature"):
         super().__init__()
 
         self._nl_request = NeighborListOptions(
             cutoff=cutoff, full_list=True, strict=True
         )
         self._multiple_properties = multiple_properties
+        self._feature_key = feature_key
 
     def forward(
         self,
         systems: List[System],
         outputs: Dict[str, ModelOutput],
         selected_atoms: Optional[Labels],
     ) -> Dict[str, TensorMap]:
-        if "features" not in outputs:
+        if self._feature_key not in outputs:
             return {}
 
         device = torch.device("cpu")
         if len(systems) > 0:
             device = systems[0].positions.device
 
-        output = outputs["features"]
+        output = outputs[self._feature_key]
 
-        if output.per_atom:
+        if output.sample_kind == "atom":
             samples_list: List[List[int]] = []
             for s, system in enumerate(systems):
                 for i in range(len(system)):
@@ -67,6 +68,7 @@ def forward(
                 sample_values.reshape(-1, 2),
             )
         else:
+            assert output.sample_kind == "system"
             samples = Labels(
                 "system", torch.arange(len(systems), device=device).reshape(-1, 1)
             )
@@ -89,14 +91,14 @@ def forward(
             distances = torch.linalg.vector_norm(neighbors.values.reshape(-1, 3), dim=1)
             inv_dist = 1.0 / distances
 
-            if output.per_atom:
+            if output.sample_kind == "atom":
                 sliced = values[system_start:system_stop, 0]
                 sliced += sliced.index_add(0, atom_index, inv_dist)
             else:
                 values[system_i, 0] += inv_dist.sum()
 
             if self._multiple_properties:
-                if output.per_atom:
+                if output.sample_kind == "atom":
                     sliced = values[system_start:system_stop, 1]
                     sliced += sliced.index_add(0, atom_index, inv_dist**2)
                 else:
@@ -116,14 +118,14 @@ def forward(
         )
 
         if selected_atoms is not None:
-            if output.per_atom:
+            if output.sample_kind == "atom":
                 cv = mts.slice(cv, axis="samples", selection=selected_atoms)
             else:
                 raise ValueError(
                     "selected atoms is only supported with per-atom output"
                 )
 
-        return {"features": cv}
+        return {self._feature_key: cv}
 
     def requested_neighbor_lists(self) -> List[NeighborListOptions]:
         return [self._nl_request]
@@ -132,7 +134,7 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
 CUTOFF = 3.5
 
 capabilities = ModelCapabilities(
-    outputs={"features": ModelOutput(per_atom=True)},
+    outputs={"feature": ModelOutput(sample_kind="atom")},
     interaction_range=CUTOFF,
     supported_devices=["cpu", "mps", "cuda"],
     length_unit="A",
@@ -152,7 +154,7 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
 model.save("vector-per-atom.pt")
 
 capabilities = ModelCapabilities(
-    outputs={"features": ModelOutput(per_atom=False)},
+    outputs={"feature": ModelOutput(sample_kind="system")},
     interaction_range=CUTOFF,
     supported_devices=["cpu", "mps", "cuda"],
     length_unit="A",
@@ -169,3 +171,12 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]:
 cv.eval()
 model = AtomisticModel(cv, ModelMetadata(), capabilities)
 model.save("vector-global.pt")
+
+
+cv = TestCollectiveVariable(
+    cutoff=CUTOFF, multiple_properties=False, feature_key="feature/variant"
+)
+capabilities.outputs = {"feature/variant": ModelOutput(sample_kind="system")}
+cv.eval()
+model = AtomisticModel(cv, ModelMetadata(), capabilities)
+model.save("variant-global.pt")

regtest/metatomic/rt-basic/forces.reference

@@ -1,11 +1,11 @@
 9
--4795.88 -5679.30 -5054.72
-X   147.18   -79.30  -178.34
-X  -174.21   -55.39   124.40
-X  -214.90    78.13   151.48
-X  -424.19   203.64   276.74
-X    71.31   -44.07    35.23
-X   230.71  -150.74  -438.92
-X   140.70   -55.57   -67.27
-X   208.55   -18.70   211.20
-X    14.85   122.00  -114.51
+-4822.75 -5711.13 -5083.06
+X   148.00   -79.74  -179.33
+X  -175.19   -55.70   125.09
+X  -216.11    78.57   152.33
+X  -426.55   204.77   278.27
+X    71.71   -44.32    35.43
+X   231.98  -151.58  -441.35
+X   141.49   -55.88   -67.65
+X   209.72   -18.80   212.37
+X    14.94   122.68  -115.16

regtest/metatomic/rt-basic/plumed.dat

@@ -55,15 +55,28 @@ vector_per_atom: METATOMIC ...
 
 PRINT ARG=vector_per_atom FILE=vector_per_atom FMT=%8.2f
 
+variant: METATOMIC ...
+      MODEL=variant-global.pt
+      VARIANT=variant
+      DEVICE=cpu
+
+      SPECIES1=1,4,7
+      SPECIES2=2,3,5,6,8,9
+      SPECIES_TO_TYPES=6,8
+...
+
+PRINT ARG=variant FILE=variant FMT=%8.2f
+
 
 scalar_per_atom_sum: SUM ARG=scalar_per_atom PERIODIC=NO
 vector_global_sum: SUM ARG=vector_global PERIODIC=NO
 vector_per_atom_sum: SUM ARG=vector_per_atom PERIODIC=NO
+variant_sum: SUM ARG=variant PERIODIC=NO
 
 summed: CUSTOM ...
-      ARG=scalar_global,scalar_per_atom_sum,vector_global_sum,vector_per_atom_sum
-      VAR=x,y,z,t
-      FUNC=x+y+z+t
+      ARG=scalar_global,scalar_per_atom_sum,vector_global_sum,vector_per_atom_sum,variant_sum
+      VAR=a,b,c,d,e
+      FUNC=a+b+c+d+e
       PERIODIC=NO
 ...
 

regtest/metatomic/rt-perovskite/BaTiO3-CV.py

@@ -45,7 +45,7 @@ def forward(
         outputs: Dict[str, ModelOutput],
         selected_atoms: Optional[Labels],
     ) -> Dict[str, TensorMap]:
-        if "features" not in outputs:
+        if "feature" not in outputs:
             return {}
 
         spherical_expansion = self.calculator(
@@ -71,14 +71,14 @@ def forward(
             blocks=[block],
         )
 
-        return {"features": cv}
+        return {"feature": cv}
 
 
 cv = BaTiO3_CV()
 cv.eval()
 
 capabilities = ModelCapabilities(
-    outputs={"features": ModelOutput(per_atom=False)},
+    outputs={"feature": ModelOutput(sample_kind="system")},
     interaction_range=3.0,
     supported_devices=["cpu"],
     length_unit="A",

regtest/metatomic/rt-soap/soap_cv.py

@@ -44,11 +44,11 @@ def forward(
         outputs: Dict[str, ModelOutput],
         selected_atoms: Optional[Labels],
     ) -> Dict[str, TensorMap]:
-        if "features" not in outputs:
+        if "feature" not in outputs:
             return {}
 
-        if not outputs["features"].per_atom:
-            raise ValueError("per_atom=False is not supported")
+        if outputs["feature"].sample_kind != "atom":
+            raise ValueError("only per-atom output is supported")
 
         if len(systems[0]) == 0:
             # PLUMED is trying to determine the size of the output
@@ -75,15 +75,15 @@ def forward(
             blocks=[block],
         )
 
-        return {"features": cv}
+        return {"feature": cv}
 
 
 cv = SOAP_CV(species=[1, 6, 7, 8])
 cv.eval()
 
 
 capabilities = ModelCapabilities(
-    outputs={"features": ModelOutput(per_atom=True)},
+    outputs={"feature": ModelOutput(sample_kind="atom")},
     interaction_range=4.0,
     supported_devices=["cpu"],
     length_unit="nm",

src/maketools/codecheck

@@ -75,6 +75,7 @@ fi
       -U__PLUMED_HAS_ASMJIT \
       -D__PLUMED_WRAPPER_CXX_EXPLICIT=explicit \
       --config-exclude=small_vector/ \
+      -i vesin.cpp \
       --template='[{file}:{line}] ({severity}) :{id}: {message}' \
       --enable=all --suppress=missingIncludeSystem --inline-suppr --force \
       $files

src/metatomic/import.sh src/metatomic/import-vesin.shsrc/metatomic/import.sh renamed to src/metatomic/import-vesin.sh

@@ -1,22 +1,19 @@
 #!/usr/bin/env bash
 
-path=$1
+version=$1
 
 if (($# != 1)) ; then
-echo "usage: $0 /path/to/vesin"
-echo
-echo "All the commands to run:"
-echo "    cd /some/dir/"
-echo "    git clone https://github.com/luthaf/vesin"
-echo "    cd /your/plumed2/src/metatomic"
-echo "    ./import.sh /some/dir/vesin"
-exit 0
+echo "usage: $0 <vesin-version>"
+exit 1
 fi
 
-bash -c "$path/create-single-cpp.py > /dev/null"
+rm -f vesin-single-build-v*.tar.gz
+wget https://github.com/Luthaf/vesin/releases/download/v$version/vesin-single-build-v$version.tar.gz
 
-cp $path/vesin/include/vesin.h vesin.h
-mv vesin-single-build.cpp vesin.cpp
+tar xf vesin-single-build-v$version.tar.gz
+
+mv vesin-single-build-nocuda.cpp vesin.cpp
+rm -f vesin-single-build.cpp
 
 # Patch files to follow PLUMED linter
 sed 's|#define VESIN_H|#define VESIN_H\n/*INDENT-OFF*/\n|