Update to the latest metatensor-torch

drivers/py/pes/ase.py

@@ -18,9 +18,9 @@
 
 ERROR_MSG = """
 This ASE driver requires specification of and ASE calculator
-and an ASE-readable template file that describes the chemical makeup of the structure. 
+and an ASE-readable template file that describes the chemical makeup of the structure.
 
-Example: python driver.py -m ase -u -o template.xyz,model_parameters 
+Example: python driver.py -m ase -u -o template.xyz,model_parameters
 """
 
 

drivers/py/pes/metatensor.py

@@ -1,25 +1,32 @@
-"""Interface with the [metatensor](https://lab-cosmo.github.io/metatensor/latest/atomistic/index.html)
-calculator, that can be used to perform calculations based on different types of machine learning potentials"""
+"""
+Interface with metatensor
+(https://lab-cosmo.github.io/metatensor/latest/atomistic/index.html), that can
+be used to perform calculations based on different types of machine learning
+potentials
+"""
 
 import sys
-from .ase import ASEDriver
 
 from ipi.utils.messages import warning
 
+from .ase import ASEDriver
+
 try:
+    import metatensor.torch
     from metatensor.torch.atomistic.ase_calculator import MetatensorCalculator
-except ImportError:
-    warning("Could not find or import the metatensor module")
+except ImportError as e:
+    warning(f"Could not find or import metatensor.torch: {e}")
     MetatensorCalculator = None
 
 __DRIVER_NAME__ = "metatensor"
 __DRIVER_CLASS__ = "MetatensorDriver"
 
 ERROR_MSG = """
-The metatensor driver requires specification of a .pt torchscript model 
-and an ASE-readable template file that describes the chemical makeup of the structure. 
+The metatensor driver requires specification of a .pt TorchScript model and an
+ASE-readable template file that describes the chemical makeup of the structure.
 
-Example: python driver.py -m metatensor -u -o model.pt,template.xyz
+Example: python driver.py -m metatensor -u -o template.xyz,model.pt,device=cpu,\
+extensions=path/to/extensions,check_consistency=False
 """
 
 
@@ -34,11 +41,57 @@ def check_arguments(self):
         """
 
         if MetatensorCalculator is None:
-            raise ImportError("Couldn't load metatensor bindings")
+            raise ImportError("could not import metatensor.torch, is it installed?")
+
+        metatensor_major, metatensor_minor, *_ = metatensor.torch.__version__.split(".")
+        metatensor_major = int(metatensor_major)
+        metatensor_minor = int(metatensor_minor)
+
+        if metatensor_major != 0 or metatensor_minor != 4:
+            raise ImportError(
+                "this code is only compatible with metatensor-torch v0.4.x, "
+                f"found version v{metatensor.torch.__version__} "
+                f"at '{metatensor.torch.__file__}'"
+            )
+
         super().check_arguments()
 
         if len(self.args) < 2:
             sys.exit(self.error_msg)
         self.model_path = self.args[1]
 
-        self.ase_calculator = MetatensorCalculator(self.model_path)
+        device = None
+        extensions_directory = None
+        check_consistency = False
+        for arg in self.args[2:]:
+            if arg.startswith("device="):
+                device = arg[7:]
+            elif arg.startswith("extensions="):
+                extensions_directory = arg[11:]
+            elif arg.startswith("check_consistency="):
+                arg = arg[18:]
+                if arg == "True":
+                    check_consistency = True
+                elif arg == "False":
+                    check_consistency = False
+                else:
+                    raise ValueError(
+                        "invalid value for check_consistency, expected True or False, "
+                        f"got {arg}"
+                    )
+
+            else:
+                sys.exit(self.error_msg)
+
+        # print("device = ", device)
+        # print(device)
+
+        self.ase_calculator = MetatensorCalculator(
+            self.model_path,
+            device=device,
+            extensions_directory=extensions_directory,
+            check_consistency=check_consistency,
+        )
+
+        # Show the model metadata to the users
+        print(self.ase_calculator.metadata())

drivers/py/pes/pet.py

@@ -26,9 +26,9 @@
 class PET_driver(Dummy_driver):
     def __init__(self, args=None, verbose=False):
         self.error_msg = """
-The PET driver requires specification of a .json model file fitted with 
-    the PET tools, and a template file that describes the chemical makeup of 
-the structure. 
+The PET driver requires specification of a .json model file fitted with
+the PET tools, and a template file that describes the chemical makeup of
+the structure.
 
 Example: python driver.py -m pet -u -o model.json,template.xyz
 """

examples/clients/metatensor/.gitignore

@@ -0,0 +1,7 @@
+collected-extensions/
+nickel-lj-extensions.pt
+
+lj-nickel.out
+lj-nickel.restart
+lj-nickel.xc.xyz
+RESTART

examples/clients/metatensor/README.md

@@ -1,16 +1,44 @@
-metatensor - i-PI example
-========================
-
-Runs a _very_ contrieved example of a metatensor model. 
-The model is just a hard-coded Einstein crystal model for a
-chunk of diamond. No periodic boundary conditions, no ability
-to work with a different number of atoms or just a differently
-oriented sample. Really, this is just to show how to run the
-driver. Given that metatensor model takes a torchsript file
-for a model, the very same machinery can be used for actual
-ML potentials, or torchscript-implemented models. 
+# Metatensor example
+
+Runs an example of a metatensor model, implementing a basic Lennard-Jones model
+with the parameters for Nickel.
+
+This driver can be used with any model following the [metatensor atomistic
+models](https://lab-cosmo.github.io/metatensor/latest/atomistic/index.html)
+interface. This can be used for classical force-fields, but it more intended for
+machine learning potentials.
+
+## Installation
+
+The code is compatible with metatensor-torch v0.4, which you can install with
+
+```bash
+# install all metatensor packages simultaneously
+pip install "metatensor[torch]"
+
+# install packages individually, with explicit control over the installed versions
+pip install "metatensor-torch ==0.4.*" "metatensor-operations ==0.2.*"
+```
+
+## Running the example
 
 ```bash
 i-pi input.xml; sleep 1 &
-i-pi-py_driver -a metatensor -u -m metatensor -o initial.xyz,harmonic-model.pt
+i-pi-py_driver -a metatensor -u -m metatensor -o nickel.xyz,nickel-lj.pt
+
+# with all the optional parameters:
+i-pi-py_driver -a metatensor -u -m metatensor -o nickel.xyz,nickel-lj.pt,device=cpu,extensions=some-extensions-dir/,check_consistency=True
 ```
+
+The options (after `-o`) are as follow:
+
+- the path to a template file, we will use it to get the types for all atoms in
+  the system
+- the path to the model file
+- `device` controls which torch device to use to run the model
+- `extensions` is the path to a directory containing TorchScript extensions. If
+  the model requires such extensions, we will try to load them from this
+  directory first
+- `check_consistency` controls whether we should run some extra internal
+  consistency checks about data given to the model and data returned by the
+  model.

examples/clients/metatensor/create-model.py

@@ -0,0 +1,27 @@
+# Use https://github.com/Luthaf/metatensor-lj-test/ to define a basic LJ model,
+# with and without a custom TorchScript extension
+import metatensor_lj_test
+
+model = metatensor_lj_test.lennard_jones_model(
+    atomic_type=28,
+    cutoff=6.5,
+    sigma=1.5808,
+    epsilon=0.1729,
+    length_unit="Angstrom",
+    energy_unit="eV",
+    with_extension=False,
+)
+
+model.export("nickel-lj.pt")
+
+
+model = metatensor_lj_test.lennard_jones_model(
+    atomic_type=28,
+    cutoff=6.5,
+    sigma=1.5808,
+    epsilon=0.1729,
+    length_unit="Angstrom",
+    energy_unit="eV",
+    with_extension=True,
+)
+model.export("nickel-lj-extensions.pt", collect_extensions="collected-extensions/")

examples/clients/metatensor/input.xml

@@ -1,46 +1,47 @@
 <simulation verbosity="medium">
-<output prefix="harmonic-carbon"> 
-<properties stride="5" filename="out">  [ step, time{picosecond},  conserved{electronvolt}, 
-	temperature{kelvin}, kinetic_md{electronvolt}, potential{electronvolt}, pressure_md{bar}, volume{angstrom3}, 
-	ensemble_temperature{kelvin}, cell_abcABC] </properties>
-	<trajectory filename="xc" stride="5" cell_units="angstrom"> x_centroid{angstrom} </trajectory>
-     	<checkpoint stride="1000"/> 
-   	</output> 
-   	<total_steps> 1000</total_steps> 
-  	<prng><seed>12345</seed></prng>
-	<ffsocket mode='unix' name='driver' pbc="false">
-    		<latency>  1.00000000e-02</latency>
-    		<slots>4</slots>
-    		<port>12345</port>
-    		<timeout>60.000000e+00</timeout>
-   		<address>metatensor</address>
-  	</ffsocket> 
-	<system>
-        	<initialize nbeads="1">
-                	<file mode="ase" units="angstrom"> initial.xyz </file>
-                 	<velocities mode="thermal" units="kelvin"> 250.0 </velocities>
-              	</initialize>
-	     	<forces>
-      			<force forcefield='driver'/>
-   		 </forces>
-		<motion mode="dynamics">
-                	<dynamics mode="nvt">
-                    		<timestep units="femtosecond"> 0.5 </timestep>
-                        	<thermostat mode='gle'>
-                        		<A shape='(5,5)'>
-                        					[	
-				4.498098855452e-3,    6.594810718477e-6,    2.788030342989e-4,   -8.808265165053e-4,    5.605371493938e-3,
-                                -6.726802271646e-6,    2.079069559861e-9,    1.746169548818e-5,   -4.800164465960e-6,    1.025830873432e-5,
-                                -3.586191452340e-4,   -1.746169548818e-5,    3.287481976399e-5,    1.245698716799e-4,   -2.417657162526e-4,
-                                -2.508912543565e-4,    4.800164465960e-6,   -1.245698716799e-4,    6.449207766266e-4,    2.783583234046e-4,
-                                5.273493443008e-3,   -1.025830873432e-5,    2.417657162526e-4,   -2.783583234046e-4,    7.488477456790e-3
-                        					]
-                       		</A>
-                        	</thermostat>
-                 	</dynamics>
-              	</motion>
-	   	<ensemble>
-                	<temperature units="kelvin"> 250.0 </temperature>
-           	</ensemble>
-         </system>
+<output prefix="lj-nickel">
+    <properties stride="5" filename="out">
+		[step, time{picosecond}, conserved{electronvolt}, temperature{kelvin},
+		kinetic_md{electronvolt}, potential{electronvolt}, pressure_md{bar},
+		volume{angstrom3}, ensemble_temperature{kelvin}, cell_abcABC]
+    </properties>
+    <trajectory filename="xc" stride="5" cell_units="angstrom"> x_centroid{angstrom} </trajectory>
+    <checkpoint stride="1000"/>
+</output>
+<total_steps> 1000</total_steps>
+<prng><seed>12345</seed></prng>
+
+<ffsocket mode='unix' name='driver' pbc="false">
+    <port>12345</port>
+    <address>metatensor</address>
+</ffsocket>
+
+<system>
+    <initialize nbeads="1">
+        <file mode="ase" units="angstrom"> nickel.xyz </file>
+        <velocities mode="thermal" units="kelvin"> 250.0 </velocities>
+    </initialize>
+    <forces>
+    	<force forcefield='driver'/>
+    </forces>
+
+    <motion mode="dynamics">
+    	<dynamics mode="nvt">
+        	<timestep units="femtosecond"> 0.5 </timestep>
+            <thermostat mode='gle'>
+            	<A shape='(5,5)'>[
+                	 4.49e-3,   6.59e-6,   2.79e-4,  -8.81e-4,   5.61e-3,
+                    -6.73e-6,   2.08e-9,   1.75e-5,  -4.80e-6,   1.03e-5,
+                    -3.59e-4,  -1.75e-5,   3.29e-5,   1.24e-4,  -2.42e-4,
+                    -2.51e-4,   4.80e-6,  -1.24e-4,   6.45e-4,   2.78e-4,
+                     5.27e-3,  -1.03e-5,   2.42e-4,  -2.78e-4,   7.49e-3
+                ]</A>
+            </thermostat>
+        </dynamics>
+    </motion>
+
+	<ensemble>
+        <temperature units="kelvin"> 250.0 </temperature>
+    </ensemble>
+</system>
 </simulation>