diff --git a/src/metatrain/experimental/alchemical_model/tests/test_exported.py b/src/metatrain/experimental/alchemical_model/tests/test_exported.py index 9b112798b..434cdac12 100644 --- a/src/metatrain/experimental/alchemical_model/tests/test_exported.py +++ b/src/metatrain/experimental/alchemical_model/tests/test_exported.py @@ -1,7 +1,6 @@ -import ase import pytest import torch -from metatensor.torch.atomistic import ModelEvaluationOptions, systems_to_torch +from metatensor.torch.atomistic import ModelEvaluationOptions, System from metatrain.experimental.alchemical_model import AlchemicalModel from metatrain.utils.data import DatasetInfo, TargetInfo, TargetInfoDict @@ -27,8 +26,11 @@ def test_to(device, dtype): exported.to(device=device) - system = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = systems_to_torch(system, dtype=torch.get_default_dtype()) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) system = get_system_with_neighbor_lists(system, exported.requested_neighbor_lists()) system = system.to(device=device, dtype=dtype) diff --git a/src/metatrain/experimental/alchemical_model/tests/test_functionality.py b/src/metatrain/experimental/alchemical_model/tests/test_functionality.py index a2e4d122a..e5f2f34ee 100644 --- a/src/metatrain/experimental/alchemical_model/tests/test_functionality.py +++ b/src/metatrain/experimental/alchemical_model/tests/test_functionality.py @@ -1,5 +1,5 @@ -import ase -from metatensor.torch.atomistic import ModelEvaluationOptions, systems_to_torch +import torch +from metatensor.torch.atomistic import ModelEvaluationOptions, System from metatrain.experimental.alchemical_model import AlchemicalModel from metatrain.utils.data import DatasetInfo, TargetInfo, TargetInfoDict @@ -20,8 +20,11 @@ def test_prediction_subset_elements(): model = AlchemicalModel(MODEL_HYPERS, dataset_info) - system = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = systems_to_torch(system) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) system = get_system_with_neighbor_lists(system, model.requested_neighbor_lists()) evaluation_options = ModelEvaluationOptions( diff --git a/src/metatrain/experimental/alchemical_model/tests/test_invariance.py b/src/metatrain/experimental/alchemical_model/tests/test_invariance.py index 2ee406d03..178c69327 100644 --- a/src/metatrain/experimental/alchemical_model/tests/test_invariance.py +++ b/src/metatrain/experimental/alchemical_model/tests/test_invariance.py @@ -27,6 +27,8 @@ def test_rotational_invariance(): original_system = get_system_with_neighbor_lists( original_system, model.requested_neighbor_lists() ) + + system.rotate(48, "y") system = systems_to_torch(system) system = get_system_with_neighbor_lists(system, model.requested_neighbor_lists()) diff --git a/src/metatrain/experimental/alchemical_model/tests/test_regression.py b/src/metatrain/experimental/alchemical_model/tests/test_regression.py index cb08f747d..a7cc17ef8 100644 --- a/src/metatrain/experimental/alchemical_model/tests/test_regression.py +++ b/src/metatrain/experimental/alchemical_model/tests/test_regression.py @@ -1,15 +1,19 @@ import random -import ase.io import numpy as np import torch -from metatensor.torch.atomistic import ModelEvaluationOptions, systems_to_torch +from metatensor.torch.atomistic import ModelEvaluationOptions from omegaconf import OmegaConf from metatrain.experimental.alchemical_model import AlchemicalModel, Trainer -from metatrain.utils.data import Dataset, DatasetInfo, TargetInfo +from metatrain.utils.data import ( + Dataset, + DatasetInfo, + TargetInfo, + read_systems, + read_targets, +) from metatrain.utils.data.dataset import TargetInfoDict -from metatrain.utils.data.readers import read_systems, read_targets from metatrain.utils.neighbor_lists import get_system_with_neighbor_lists from . import DATASET_PATH, DEFAULT_HYPERS, MODEL_HYPERS @@ -33,8 +37,7 @@ def test_regression_init(): model = AlchemicalModel(MODEL_HYPERS, dataset_info) # Predict on the first five systems - systems = ase.io.read(DATASET_PATH, ":5") - systems = [systems_to_torch(system) for system in systems] + systems = read_systems(DATASET_PATH)[:5] systems = [ get_system_with_neighbor_lists(system, model.requested_neighbor_lists()) for system in systems diff --git a/src/metatrain/experimental/pet/tests/test_exported.py b/src/metatrain/experimental/pet/tests/test_exported.py index 1e2a5f329..7230f66d5 100644 --- a/src/metatrain/experimental/pet/tests/test_exported.py +++ b/src/metatrain/experimental/pet/tests/test_exported.py @@ -1,11 +1,10 @@ -import ase import pytest import torch from metatensor.torch.atomistic import ( ModelCapabilities, ModelEvaluationOptions, ModelOutput, - systems_to_torch, + System, ) from pet.hypers import Hypers from pet.pet import PET @@ -55,8 +54,11 @@ def test_to(device): exported = export(model, capabilities) exported.to(device=device, dtype=dtype) - system = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = systems_to_torch(system, dtype=torch.get_default_dtype()) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) system = get_system_with_neighbor_lists(system, exported.requested_neighbor_lists()) system = system.to(device=device, dtype=dtype) diff --git a/src/metatrain/experimental/pet/tests/test_functionality.py b/src/metatrain/experimental/pet/tests/test_functionality.py index 91bed9b03..357918e65 100644 --- a/src/metatrain/experimental/pet/tests/test_functionality.py +++ b/src/metatrain/experimental/pet/tests/test_functionality.py @@ -1,4 +1,3 @@ -import ase import torch from metatensor.torch import Labels from metatensor.torch.atomistic import ( @@ -7,7 +6,7 @@ ModelEvaluationOptions, ModelMetadata, ModelOutput, - systems_to_torch, + System, ) from pet.hypers import Hypers from pet.pet import PET @@ -34,8 +33,11 @@ def test_prediction(): raw_pet = PET(ARCHITECTURAL_HYPERS, 0.0, len(model.atomic_types)) model.set_trained_model(raw_pet) - structure = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = systems_to_torch(structure) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) system = get_system_with_neighbor_lists(system, model.requested_neighbor_lists()) evaluation_options = ModelEvaluationOptions( @@ -80,8 +82,11 @@ def test_per_atom_predictions_functionality(): raw_pet = PET(ARCHITECTURAL_HYPERS, 0.0, len(model.atomic_types)) model.set_trained_model(raw_pet) - structure = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = systems_to_torch(structure) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) system = get_system_with_neighbor_lists(system, model.requested_neighbor_lists()) evaluation_options = ModelEvaluationOptions( @@ -127,8 +132,11 @@ def test_selected_atoms_functionality(): raw_pet = PET(ARCHITECTURAL_HYPERS, 0.0, len(model.atomic_types)) model.set_trained_model(raw_pet) - structure = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = systems_to_torch(structure) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) system = get_system_with_neighbor_lists(system, model.requested_neighbor_lists()) evaluation_options = ModelEvaluationOptions( diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py index f0e12e39a..3df1542de 100644 --- a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py +++ b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py @@ -1,7 +1,6 @@ -import ase import pytest import torch -from metatensor.torch.atomistic import ModelEvaluationOptions, systems_to_torch +from metatensor.torch.atomistic import ModelEvaluationOptions, System from metatrain.experimental.soap_bpnn import SoapBpnn from metatrain.utils.data import DatasetInfo, TargetInfo, TargetInfoDict @@ -27,8 +26,11 @@ def test_to(device, dtype): exported.to(device=device) - system = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = systems_to_torch(system, dtype=torch.get_default_dtype()) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) system = get_system_with_neighbor_lists(system, exported.requested_neighbor_lists()) system = system.to(device=device, dtype=dtype) diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py index 75573f7bd..aea3c56d9 100644 --- a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py +++ b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py @@ -1,7 +1,6 @@ -import ase import metatensor.torch import torch -from metatensor.torch.atomistic import ModelOutput, systems_to_torch +from metatensor.torch.atomistic import ModelOutput, System from metatrain.experimental.soap_bpnn import SoapBpnn from metatrain.utils.data import DatasetInfo, TargetInfo, TargetInfoDict @@ -21,9 +20,13 @@ def test_prediction_subset_elements(): model = SoapBpnn(MODEL_HYPERS, dataset_info) - system = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + system = System( + types=torch.tensor([6, 6]), + positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), + cell=torch.zeros(3, 3), + ) model( - [systems_to_torch(system)], + [system], {"energy": model.outputs["energy"]}, ) @@ -43,25 +46,32 @@ def test_prediction_subset_atoms(): # Since we don't yet support atomic predictions, we will test this by # predicting on a system with two monomers at a large distance - system_monomer = ase.Atoms( - "NO2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0]] + system_monomer = System( + types=torch.tensor([7, 8, 8]), + positions=torch.tensor( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0]], + ), + cell=torch.zeros(3, 3), ) energy_monomer = model( - [systems_to_torch(system_monomer)], + [system_monomer], {"energy": ModelOutput(per_atom=False)}, ) - system_far_away_dimer = ase.Atoms( - "N2O4", - positions=[ - [0.0, 0.0, 0.0], - [0.0, 50.0, 0.0], - [0.0, 0.0, 1.0], - [0.0, 0.0, 2.0], - [0.0, 51.0, 0.0], - [0.0, 42.0, 0.0], - ], + system_far_away_dimer = System( + types=torch.tensor([7, 7, 8, 8, 8, 8]), + positions=torch.tensor( + [ + [0.0, 0.0, 0.0], + [0.0, 50.0, 0.0], + [0.0, 0.0, 1.0], + [0.0, 0.0, 2.0], + [0.0, 51.0, 0.0], + [0.0, 42.0, 0.0], + ], + ), + cell=torch.zeros(3, 3), ) selection_labels = metatensor.torch.Labels( @@ -70,12 +80,12 @@ def test_prediction_subset_atoms(): ) energy_dimer = model( - [systems_to_torch(system_far_away_dimer)], + [system_far_away_dimer], {"energy": ModelOutput(per_atom=False)}, ) energy_monomer_in_dimer = model( - [systems_to_torch(system_far_away_dimer)], + [system_far_away_dimer], {"energy": ModelOutput(per_atom=False)}, selected_atoms=selection_labels, ) @@ -99,9 +109,12 @@ def test_output_last_layer_features(): model = SoapBpnn(MODEL_HYPERS, dataset_info) - system = ase.Atoms( - "CHON", - positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], + system = System( + types=torch.tensor([6, 1, 8, 7]), + positions=torch.tensor( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], + ), + cell=torch.zeros(3, 3), ) # last-layer features per atom: @@ -111,7 +124,7 @@ def test_output_last_layer_features(): per_atom=True, ) outputs = model( - [systems_to_torch(system, dtype=torch.get_default_dtype())], + [system], { "energy": model.outputs["energy"], "mtm::aux::last_layer_features": ll_output_options, @@ -139,7 +152,7 @@ def test_output_last_layer_features(): per_atom=False, ) outputs = model( - [systems_to_torch(system, dtype=torch.get_default_dtype())], + [system], { "energy": model.outputs["energy"], "mtm::aux::last_layer_features": ll_output_options, @@ -167,13 +180,16 @@ def test_output_per_atom(): model = SoapBpnn(MODEL_HYPERS, dataset_info) - system = ase.Atoms( - "CHON", - positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], + system = System( + types=torch.tensor([6, 1, 8, 7]), + positions=torch.tensor( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], + ), + cell=torch.zeros(3, 3), ) outputs = model( - [systems_to_torch(system, dtype=torch.get_default_dtype())], + [system], {"energy": model.outputs["energy"]}, ) diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_regression.py b/src/metatrain/experimental/soap_bpnn/tests/test_regression.py index afc323f2a..c76dc6d37 100644 --- a/src/metatrain/experimental/soap_bpnn/tests/test_regression.py +++ b/src/metatrain/experimental/soap_bpnn/tests/test_regression.py @@ -1,9 +1,8 @@ import random -import ase.io import numpy as np import torch -from metatensor.torch.atomistic import ModelOutput, systems_to_torch +from metatensor.torch.atomistic import ModelOutput from omegaconf import OmegaConf from metatrain.experimental.soap_bpnn import SoapBpnn, Trainer @@ -31,10 +30,10 @@ def test_regression_init(): model = SoapBpnn(MODEL_HYPERS, dataset_info) # Predict on the first five systems - systems = ase.io.read(DATASET_PATH, ":5") + systems = read_systems(DATASET_PATH)[:5] output = model( - [systems_to_torch(system) for system in systems], + systems, {"mtm::U0": ModelOutput(quantity="energy", unit="", per_atom=False)}, ) diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py index 6ec3a14fb..73631b1a4 100644 --- a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py +++ b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py @@ -1,8 +1,7 @@ import copy -import ase import torch -from metatensor.torch.atomistic import systems_to_torch +from metatensor.torch.atomistic import System from metatrain.experimental.soap_bpnn import SoapBpnn from metatrain.utils.data import DatasetInfo, TargetInfo, TargetInfoDict @@ -21,12 +20,15 @@ def test_torchscript(): model = SoapBpnn(MODEL_HYPERS, dataset_info) model = torch.jit.script(model) - system = ase.Atoms( - "OHCN", - positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], + system = System( + types=torch.tensor([6, 1, 8, 7]), + positions=torch.tensor( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]] + ), + cell=torch.zeros(3, 3), ) model( - [systems_to_torch(system)], + [system], {"energy": model.outputs["energy"]}, ) @@ -44,12 +46,15 @@ def test_torchscript_with_identity(): model = SoapBpnn(hypers, dataset_info) model = torch.jit.script(model) - system = ase.Atoms( - "OHCN", - positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], + system = System( + types=torch.tensor([6, 1, 8, 7]), + positions=torch.tensor( + [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]] + ), + cell=torch.zeros(3, 3), ) model( - [systems_to_torch(system)], + [system], {"energy": model.outputs["energy"]}, ) diff --git a/tests/utils/test_evaluate_model.py b/tests/utils/test_evaluate_model.py index b829d4428..a7304c6a7 100644 --- a/tests/utils/test_evaluate_model.py +++ b/tests/utils/test_evaluate_model.py @@ -16,9 +16,7 @@ def test_evaluate_model(training, exported): """Test that the evaluate_model function works as intended.""" - systems = read_systems( - RESOURCES_PATH / "alchemical_reduced_10.xyz", dtype=torch.get_default_dtype() - )[:2] + systems = read_systems(RESOURCES_PATH / "alchemical_reduced_10.xyz")[:2] atomic_types = set( torch.unique(torch.concatenate([system.types for system in systems]))