diff --git a/drivers/py/pes/ase.py b/drivers/py/pes/ase.py index d17ffe545..7d5597a52 100644 --- a/drivers/py/pes/ase.py +++ b/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 """ diff --git a/drivers/py/pes/metatensor.py b/drivers/py/pes/metatensor.py index 14d4dd527..8c1585e20 100644 --- a/drivers/py/pes/metatensor.py +++ b/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()) diff --git a/drivers/py/pes/pet.py b/drivers/py/pes/pet.py index 4a5234782..1d33f7cf6 100644 --- a/drivers/py/pes/pet.py +++ b/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 """ diff --git a/examples/clients/metatensor/.gitignore b/examples/clients/metatensor/.gitignore new file mode 100644 index 000000000..6052486bc --- /dev/null +++ b/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 diff --git a/examples/clients/metatensor/README.md b/examples/clients/metatensor/README.md index 95409a854..281fbf8f7 100644 --- a/examples/clients/metatensor/README.md +++ b/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. diff --git a/examples/clients/metatensor/create-model.py b/examples/clients/metatensor/create-model.py new file mode 100644 index 000000000..cfd039ad4 --- /dev/null +++ b/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/") diff --git a/examples/clients/metatensor/harmonic-model.pt b/examples/clients/metatensor/harmonic-model.pt deleted file mode 100644 index 809daadc4..000000000 Binary files a/examples/clients/metatensor/harmonic-model.pt and /dev/null differ diff --git a/examples/clients/metatensor/initial.xyz b/examples/clients/metatensor/initial.xyz deleted file mode 100644 index 0f49a2052..000000000 --- a/examples/clients/metatensor/initial.xyz +++ /dev/null @@ -1,56 +0,0 @@ -54 -Lattice="100 0 0 0 100 0 0 0 100" Properties=species:S:1:pos:R:3 pbc="T T T" -C 0.00000000 0.00000000 0.00000000 -C 0.89175000 0.89175000 0.89175000 -C 1.78350000 1.78350000 0.00000000 -C 2.67525000 2.67525000 0.89175000 -C 3.56700000 3.56700000 0.00000000 -C 4.45875000 4.45875000 0.89175000 -C 1.78350000 0.00000000 1.78350000 -C 2.67525000 0.89175000 2.67525000 -C 3.56700000 1.78350000 1.78350000 -C 4.45875000 2.67525000 2.67525000 -C 5.35050000 3.56700000 1.78350000 -C 6.24225000 4.45875000 2.67525000 -C 3.56700000 0.00000000 3.56700000 -C 4.45875000 0.89175000 4.45875000 -C 5.35050000 1.78350000 3.56700000 -C 6.24225000 2.67525000 4.45875000 -C 7.13400000 3.56700000 3.56700000 -C 8.02575000 4.45875000 4.45875000 -C 0.00000000 1.78350000 1.78350000 -C 0.89175000 2.67525000 2.67525000 -C 1.78350000 3.56700000 1.78350000 -C 2.67525000 4.45875000 2.67525000 -C 3.56700000 5.35050000 1.78350000 -C 4.45875000 6.24225000 2.67525000 -C 1.78350000 1.78350000 3.56700000 -C 2.67525000 2.67525000 4.45875000 -C 3.56700000 3.56700000 3.56700000 -C 4.45875000 4.45875000 4.45875000 -C 5.35050000 5.35050000 3.56700000 -C 6.24225000 6.24225000 4.45875000 -C 3.56700000 1.78350000 5.35050000 -C 4.45875000 2.67525000 6.24225000 -C 5.35050000 3.56700000 5.35050000 -C 6.24225000 4.45875000 6.24225000 -C 7.13400000 5.35050000 5.35050000 -C 8.02575000 6.24225000 6.24225000 -C 0.00000000 3.56700000 3.56700000 -C 0.89175000 4.45875000 4.45875000 -C 1.78350000 5.35050000 3.56700000 -C 2.67525000 6.24225000 4.45875000 -C 3.56700000 7.13400000 3.56700000 -C 4.45875000 8.02575000 4.45875000 -C 1.78350000 3.56700000 5.35050000 -C 2.67525000 4.45875000 6.24225000 -C 3.56700000 5.35050000 5.35050000 -C 4.45875000 6.24225000 6.24225000 -C 5.35050000 7.13400000 5.35050000 -C 6.24225000 8.02575000 6.24225000 -C 3.56700000 3.56700000 7.13400000 -C 4.45875000 4.45875000 8.02575000 -C 5.35050000 5.35050000 7.13400000 -C 6.24225000 6.24225000 8.02575000 -C 7.13400000 7.13400000 7.13400000 -C 8.02575000 8.02575000 8.02575000 diff --git a/examples/clients/metatensor/input.xml b/examples/clients/metatensor/input.xml index 8a7dea448..ad58b0a8c 100644 --- a/examples/clients/metatensor/input.xml +++ b/examples/clients/metatensor/input.xml @@ -1,46 +1,47 @@ - - [ step, time{picosecond}, conserved{electronvolt}, - temperature{kelvin}, kinetic_md{electronvolt}, potential{electronvolt}, pressure_md{bar}, volume{angstrom3}, - ensemble_temperature{kelvin}, cell_abcABC] - x_centroid{angstrom} - - - 1000 - 12345 - - 1.00000000e-02 - 4 - 12345 - 60.000000e+00 -
metatensor
-
- - - initial.xyz - 250.0 - - - - - - - 0.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 - ] - - - - - - 250.0 - - + + + [step, time{picosecond}, conserved{electronvolt}, temperature{kelvin}, + kinetic_md{electronvolt}, potential{electronvolt}, pressure_md{bar}, + volume{angstrom3}, ensemble_temperature{kelvin}, cell_abcABC] + + x_centroid{angstrom} + + + 1000 +12345 + + + 12345 +
metatensor
+
+ + + + nickel.xyz + 250.0 + + + + + + + + 0.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 + ] + + + + + + 250.0 + +
diff --git a/examples/clients/metatensor/nickel-lj.pt b/examples/clients/metatensor/nickel-lj.pt new file mode 100644 index 000000000..4f7ca0249 Binary files /dev/null and b/examples/clients/metatensor/nickel-lj.pt differ diff --git a/examples/clients/metatensor/nickel.xyz b/examples/clients/metatensor/nickel.xyz new file mode 100644 index 000000000..ef3986a7e --- /dev/null +++ b/examples/clients/metatensor/nickel.xyz @@ -0,0 +1,34 @@ +32 +Lattice="7.2 0.0 0.0 0.0 7.2 0.0 0.0 0.0 7.2" Properties=species:S:1:pos:R:3 pbc="T T T" +Ni 0.04454187 0.15457874 0.12596108 +Ni 0.04370318 1.95000420 1.91184154 +Ni 1.83573509 0.12884331 1.91958398 +Ni 1.92303186 1.89265686 0.07165795 +Ni 0.19190924 0.04701160 3.61889659 +Ni 0.04116880 1.91517186 5.56317915 +Ni 1.91817295 0.03169102 5.58591245 +Ni 1.90862314 1.82490488 3.69025253 +Ni 0.16624021 3.78437918 0.07411722 +Ni 0.11502732 5.41155667 1.80895247 +Ni 1.88561487 3.68501992 1.91533902 +Ni 1.89092684 5.47334069 0.13616005 +Ni 0.13256096 3.78185764 3.77535577 +Ni 0.01409603 5.46866518 5.46642113 +Ni 1.97184883 3.60053345 5.44955058 +Ni 1.95063962 5.40970450 3.62949639 +Ni 3.61636855 0.00654232 0.19947714 +Ni 3.76729948 1.93207137 1.83696219 +Ni 5.47783206 0.09663835 1.87546754 +Ni 5.51789413 1.88532203 0.15452776 +Ni 3.66412964 0.13315923 3.68647837 +Ni 3.71818983 1.96085601 5.43543213 +Ni 5.54731704 0.16090601 5.59503045 +Ni 5.40527227 1.81865617 3.73525599 +Ni 3.77409855 3.77888268 0.04458552 +Ni 3.77575860 5.56995384 1.83599413 +Ni 5.49352058 3.71215650 1.92995287 +Ni 5.53270737 5.51022677 0.07537065 +Ni 3.60323764 3.73888965 3.76030754 +Ni 3.77150406 5.57225362 5.48239838 +Ni 5.48403719 3.68038819 5.43621244 +Ni 5.51535161 5.56333248 3.64169044