Use Systems in tests when possible (#238)

metatensor · Jun 10, 2024 · f48b865 · f48b865
1 parent b21c058
commit f48b865
Show file tree

Hide file tree

Showing 11 changed files with 115 additions and 75 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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)},
     )