diff --git a/docs/src/dev-docs/utils/export.rst b/docs/src/dev-docs/utils/export.rst deleted file mode 100644 index 24dc2015..00000000 --- a/docs/src/dev-docs/utils/export.rst +++ /dev/null @@ -1,7 +0,0 @@ -Export -###### - -.. automodule:: metatrain.utils.export - :members: - :undoc-members: - :show-inheritance: diff --git a/src/metatrain/experimental/alchemical_model/model.py b/src/metatrain/experimental/alchemical_model/model.py index 8ffbd8dc..e87df9ab 100644 --- a/src/metatrain/experimental/alchemical_model/model.py +++ b/src/metatrain/experimental/alchemical_model/model.py @@ -7,6 +7,7 @@ from metatensor.torch.atomistic import ( MetatensorAtomisticModel, ModelCapabilities, + ModelMetadata, ModelOutput, NeighborListOptions, System, @@ -16,7 +17,6 @@ from ...utils.additive import ZBL from ...utils.data.dataset import DatasetInfo from ...utils.dtype import dtype_to_str -from ...utils.export import export from .utils import systems_to_torch_alchemical_batch @@ -184,7 +184,7 @@ def export(self) -> MetatensorAtomisticModel: dtype=dtype_to_str(dtype), ) - return export(model=self, model_capabilities=capabilities) + return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities) def set_composition_weights( self, diff --git a/src/metatrain/experimental/gap/model.py b/src/metatrain/experimental/gap/model.py index 833325d1..fad045a6 100644 --- a/src/metatrain/experimental/gap/model.py +++ b/src/metatrain/experimental/gap/model.py @@ -14,6 +14,7 @@ from metatensor.torch.atomistic import ( MetatensorAtomisticModel, ModelCapabilities, + ModelMetadata, ModelOutput, System, ) @@ -22,7 +23,6 @@ from metatrain.utils.data.dataset import DatasetInfo from ...utils.additive import ZBL, CompositionModel -from ...utils.export import export class GAP(torch.nn.Module): @@ -253,7 +253,7 @@ def export(self) -> MetatensorAtomisticModel: self._subset_of_regressors.export_torch_script_model() ) - return export(model=self, model_capabilities=capabilities) + return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities) def set_composition_weights( self, diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py index bf567dee..39355b8f 100644 --- a/src/metatrain/experimental/pet/model.py +++ b/src/metatrain/experimental/pet/model.py @@ -8,6 +8,7 @@ from metatensor.torch.atomistic import ( MetatensorAtomisticModel, ModelCapabilities, + ModelMetadata, ModelOutput, NeighborListOptions, System, @@ -20,7 +21,6 @@ from ...utils.additive import ZBL from ...utils.dtype import dtype_to_str -from ...utils.export import export from .utils import systems_to_batch_dict @@ -196,4 +196,4 @@ def export(self) -> MetatensorAtomisticModel: supported_devices=["cpu", "cuda"], # and not __supported_devices__ dtype=dtype_to_str(dtype), ) - return export(model=self, model_capabilities=capabilities) + return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities) diff --git a/src/metatrain/experimental/soap_bpnn/model.py b/src/metatrain/experimental/soap_bpnn/model.py index 556f3ef5..3541141d 100644 --- a/src/metatrain/experimental/soap_bpnn/model.py +++ b/src/metatrain/experimental/soap_bpnn/model.py @@ -8,6 +8,7 @@ from metatensor.torch.atomistic import ( MetatensorAtomisticModel, ModelCapabilities, + ModelMetadata, ModelOutput, System, ) @@ -18,7 +19,6 @@ from ...utils.additive import ZBL, CompositionModel from ...utils.dtype import dtype_to_str -from ...utils.export import export class Identity(torch.nn.Module): @@ -335,7 +335,7 @@ def export(self) -> MetatensorAtomisticModel: dtype=dtype_to_str(dtype), ) - return export(model=self, model_capabilities=capabilities) + return MetatensorAtomisticModel(self.eval(), ModelMetadata(), capabilities) def add_output(self, output_name: str) -> None: """Add a new output to the self.""" diff --git a/src/metatrain/utils/export.py b/src/metatrain/utils/export.py deleted file mode 100644 index db5f5ab7..00000000 --- a/src/metatrain/utils/export.py +++ /dev/null @@ -1,48 +0,0 @@ -import warnings - -import torch -from metatensor.torch.atomistic import ( - MetatensorAtomisticModel, - ModelCapabilities, - ModelMetadata, - is_atomistic_model, -) - - -# TODO: DELETE OR CHANGE THIS FUNCTION. -# EXPORT IS NOW PER-ARCHITECTURE - - -def export( - model: torch.nn.Module, model_capabilities: ModelCapabilities -) -> MetatensorAtomisticModel: - """Export a torch.nn.Module model to a MetatensorAtomisticModel. - - The exoort allows the model to make predictions especially in molecular simulation - engines. Exported models can be be saved to a file with - ``exported_model.save(path)``. - - :param model: model to be exported - :param model_capabilities: capabilities of the model - :returns: exprted model - """ - - if is_atomistic_model(model): - return model - - if model_capabilities.length_unit == "": - warnings.warn( - "No `length_unit` was provided for the model. As a result, lengths " - "and any derived quantities will be passed to MD engines as is.", - stacklevel=1, - ) - - for model_output_name, model_output in model_capabilities.outputs.items(): - if model_output.unit == "": - warnings.warn( - f"No target units were provided for output {model_output_name!r}. " - "As a result, this model output will be passed to MD engines as is.", - stacklevel=1, - ) - - return MetatensorAtomisticModel(model.eval(), ModelMetadata(), model_capabilities) diff --git a/tests/utils/test_export.py b/tests/utils/test_export.py deleted file mode 100644 index 480d411e..00000000 --- a/tests/utils/test_export.py +++ /dev/null @@ -1,107 +0,0 @@ -import os - -import pytest -import torch -from metatensor.torch.atomistic import ModelCapabilities, is_atomistic_model - -from metatrain.experimental.soap_bpnn import __model__ -from metatrain.utils.data import DatasetInfo, TargetInfo -from metatrain.utils.export import export - -from . import MODEL_HYPERS - - -def test_export(tmp_path): - """Tests the export and is_export function""" - os.chdir(tmp_path) - - dataset_info = DatasetInfo( - length_unit="angstrom", - atomic_types={1}, - targets={"energy": TargetInfo(quantity="energy", unit="eV")}, - ) - model = __model__(model_hypers=MODEL_HYPERS, dataset_info=dataset_info) - - capabilities = ModelCapabilities( - length_unit=model.dataset_info.length_unit, - outputs=model.outputs, - atomic_types=list(model.dataset_info.atomic_types), - supported_devices=model.__supported_devices__, - interaction_range=model.hypers["soap"]["cutoff"], - dtype="float32", - ) - - # test `export()` - exported_model = export(model, capabilities) - exported_model.save("model.pt") - - # test `is_export()` - assert not is_atomistic_model(model) - assert is_atomistic_model(exported_model) - assert is_atomistic_model(torch.jit.load("model.pt")) - - -def test_reexport(monkeypatch, tmp_path): - """Test that an already exported model can be loaded and again exported.""" - monkeypatch.chdir(tmp_path) - - dataset_info = DatasetInfo( - length_unit="angstrom", - atomic_types={1}, - targets={"energy": TargetInfo(quantity="energy", unit="eV")}, - ) - model = __model__(model_hypers=MODEL_HYPERS, dataset_info=dataset_info) - - capabilities = ModelCapabilities( - length_unit=model.dataset_info.length_unit, - outputs=model.outputs, - atomic_types=list(model.dataset_info.atomic_types), - supported_devices=model.__supported_devices__, - interaction_range=model.hypers["soap"]["cutoff"], - dtype="float32", - ) - - exported_model = export(model, capabilities) - export(exported_model, capabilities) - - -def test_length_units_warning(): - dataset_info = DatasetInfo( - length_unit="angstrom", - atomic_types={1}, - targets={"energy": TargetInfo(quantity="energy", unit="eV")}, - ) - model = __model__(model_hypers=MODEL_HYPERS, dataset_info=dataset_info) - - capabilities = ModelCapabilities( - outputs=model.outputs, - atomic_types=list(model.dataset_info.atomic_types), - interaction_range=model.hypers["soap"]["cutoff"], - length_unit="", - supported_devices=model.__supported_devices__, - dtype="float32", - ) - - with pytest.warns(match="No `length_unit` was provided for the model."): - export(model, capabilities) - - -def test_units_warning(): - dataset_info = DatasetInfo( - length_unit="angstrom", - atomic_types={1}, - targets={"mtt::output": TargetInfo(quantity="energy")}, - ) - model = __model__(model_hypers=MODEL_HYPERS, dataset_info=dataset_info) - - capabilities = ModelCapabilities( - length_unit=model.dataset_info.length_unit, - outputs=model.outputs, - atomic_types=list(model.dataset_info.atomic_types), - supported_devices=model.__supported_devices__, - interaction_range=model.hypers["soap"]["cutoff"], - dtype="float32", - ) - - with pytest.warns(match="No target units were provided for output 'mtt::output'"): - export(model, capabilities) diff --git a/tox.ini b/tox.ini index f7f69c74..92a12878 100644 --- a/tox.ini +++ b/tox.ini @@ -63,7 +63,7 @@ extras = # architectures used in the package tests soap-bpnn pet allowlist_externals = bash -#commands_pre = bash {toxinidir}/tests/resources/generate-outputs.sh +commands_pre = bash {toxinidir}/tests/resources/generate-outputs.sh commands = pytest \ --numprocesses=auto \