Change atomic_types and gradients from sets to unique lists (#296)

--------- Co-authored-by: frostedoyster <[email protected]>
metatensor · Jul 16, 2024 · 576b2b0 · 576b2b0
1 parent 0024253
commit 576b2b0
Show file tree

Hide file tree

Showing 28 changed files with 273 additions and 134 deletions.
diff --git a/src/metatrain/cli/eval.py b/src/metatrain/cli/eval.py
@@ -269,10 +269,10 @@ def eval_model(
             # TODO: allow the user to specify which outputs to evaluate
             eval_targets = {}
             eval_info_dict = TargetInfoDict()
-            gradients = {"positions"}
+            gradients = ["positions"]
             if all(not torch.all(system.cell == 0) for system in eval_systems):
                 # only add strain if all structures have cells
-                gradients.add("strain")
+                gradients.append("strain")
             for key in model.capabilities().outputs.keys():
                 eval_info_dict[key] = TargetInfo(
                     quantity=model.capabilities().outputs[key].quantity,

diff --git a/src/metatrain/experimental/alchemical_model/model.py b/src/metatrain/experimental/alchemical_model/model.py
@@ -27,7 +27,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
         super().__init__()
         self.hypers = model_hypers
         self.dataset_info = dataset_info
-        self.atomic_types = sorted(dataset_info.atomic_types)
+        self.atomic_types = dataset_info.atomic_types
 
         if len(dataset_info.targets) != 1:
             raise ValueError("The AlchemicalModel only supports a single target")

diff --git a/src/metatrain/experimental/alchemical_model/tests/test_exported.py b/src/metatrain/experimental/alchemical_model/tests/test_exported.py
@@ -18,7 +18,7 @@ def test_to(device, dtype):
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = AlchemicalModel(MODEL_HYPERS, dataset_info).to(dtype=dtype)

diff --git a/src/metatrain/experimental/alchemical_model/tests/test_functionality.py b/src/metatrain/experimental/alchemical_model/tests/test_functionality.py
@@ -14,7 +14,7 @@ def test_prediction_subset_elements():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
 

diff --git a/src/metatrain/experimental/alchemical_model/tests/test_invariance.py b/src/metatrain/experimental/alchemical_model/tests/test_invariance.py
@@ -16,7 +16,7 @@ def test_rotational_invariance():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = AlchemicalModel(MODEL_HYPERS, dataset_info)

diff --git a/src/metatrain/experimental/alchemical_model/tests/test_regression.py b/src/metatrain/experimental/alchemical_model/tests/test_regression.py
@@ -32,7 +32,7 @@ def test_regression_init():
     targets["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=targets
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=targets
     )
     model = AlchemicalModel(MODEL_HYPERS, dataset_info)
 
@@ -97,7 +97,7 @@ def test_regression_train():
     hypers = DEFAULT_HYPERS.copy()
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=target_info_dict
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=target_info_dict
     )
     model = AlchemicalModel(MODEL_HYPERS, dataset_info)
 

diff --git a/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py b/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py
@@ -69,7 +69,7 @@ def test_alchemical_model_inference():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types=set(unique_numbers),
+        atomic_types=unique_numbers,
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
 

diff --git a/src/metatrain/experimental/alchemical_model/tests/test_torchscript.py b/src/metatrain/experimental/alchemical_model/tests/test_torchscript.py
@@ -11,7 +11,7 @@ def test_torchscript():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
 
@@ -24,7 +24,7 @@ def test_torchscript_save_load():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = AlchemicalModel(MODEL_HYPERS, dataset_info)

diff --git a/src/metatrain/experimental/gap/model.py b/src/metatrain/experimental/gap/model.py
@@ -63,7 +63,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
             for key, value in dataset_info.targets.items()
         }
 
-        self.atomic_types = sorted(dataset_info.atomic_types)
+        self.atomic_types = dataset_info.atomic_types
         self.hypers = model_hypers
 
         # creates a composition weight tensor that can be directly indexed by species,

diff --git a/src/metatrain/experimental/gap/tests/test_errors.py b/src/metatrain/experimental/gap/tests/test_errors.py
@@ -64,7 +64,7 @@ def test_ethanol_regression_train_and_invariance():
     )
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=target_info_dict
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=target_info_dict
     )
 
     gap = GAP(hypers["model"], dataset_info)

diff --git a/src/metatrain/experimental/gap/tests/test_regression.py b/src/metatrain/experimental/gap/tests/test_regression.py
@@ -29,7 +29,7 @@ def test_regression_init():
     targets["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=targets
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=targets
     )
     GAP(DEFAULT_HYPERS["model"], dataset_info)
 
@@ -61,7 +61,7 @@ def test_regression_train_and_invariance():
     target_info_dict["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=target_info_dict
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=target_info_dict
     )
 
     gap = GAP(DEFAULT_HYPERS["model"], dataset_info)
@@ -142,7 +142,7 @@ def test_ethanol_regression_train_and_invariance():
     )
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=target_info_dict
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=target_info_dict
     )
 
     gap = GAP(hypers["model"], dataset_info)

diff --git a/src/metatrain/experimental/gap/tests/test_torchscript.py b/src/metatrain/experimental/gap/tests/test_torchscript.py
@@ -17,7 +17,7 @@ def test_torchscript():
     target_info_dict["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=target_info_dict
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=target_info_dict
     )
     conf = {
         "mtt::U0": {
@@ -68,7 +68,7 @@ def test_torchscript_save():
     targets["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=targets
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=targets
     )
     gap = GAP(DEFAULT_HYPERS["model"], dataset_info)
     torch.jit.save(

diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py
@@ -43,7 +43,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
         model_hypers["TARGET_AGGREGATION"] = "sum"
         self.hypers = model_hypers
         self.cutoff = self.hypers["R_CUT"]
-        self.atomic_types: List[int] = sorted(dataset_info.atomic_types)
+        self.atomic_types: List[int] = dataset_info.atomic_types
         self.dataset_info = dataset_info
         self.pet = None
         self.checkpoint_path: Optional[str] = None

diff --git a/src/metatrain/experimental/pet/tests/test_exported.py b/src/metatrain/experimental/pet/tests/test_exported.py
@@ -28,7 +28,7 @@ def test_to(device):
     dtype = torch.float32  # for now
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = WrappedPET(DEFAULT_HYPERS["model"], dataset_info)

diff --git a/src/metatrain/experimental/pet/tests/test_functionality.py b/src/metatrain/experimental/pet/tests/test_functionality.py
@@ -61,7 +61,7 @@ def test_prediction():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = WrappedPET(DEFAULT_HYPERS["model"], dataset_info)
@@ -110,7 +110,7 @@ def test_per_atom_predictions_functionality():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = WrappedPET(DEFAULT_HYPERS["model"], dataset_info)
@@ -160,7 +160,7 @@ def test_selected_atoms_functionality():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = WrappedPET(DEFAULT_HYPERS["model"], dataset_info)

diff --git a/src/metatrain/experimental/pet/tests/test_pet_compatibility.py b/src/metatrain/experimental/pet/tests/test_pet_compatibility.py
@@ -91,16 +91,15 @@ def test_predictions_compatibility(cutoff):
     are consistent with the predictions of the original PET implementation."""
 
     structure = ase.io.read(DATASET_PATH)
-    atomic_types = set(structure.numbers)
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types=atomic_types,
+        atomic_types=structure.numbers,
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     capabilities = ModelCapabilities(
         length_unit="Angstrom",
-        atomic_types=sorted(atomic_types),
+        atomic_types=dataset_info.atomic_types,
         outputs={
             "energy": ModelOutput(
                 quantity="energy",
@@ -116,7 +115,7 @@ def test_predictions_compatibility(cutoff):
     hypers["R_CUT"] = cutoff
     model = WrappedPET(DEFAULT_HYPERS["model"], dataset_info)
     ARCHITECTURAL_HYPERS = Hypers(model.hypers)
-    raw_pet = PET(ARCHITECTURAL_HYPERS, 0.0, len(model.atomic_types))
+    raw_pet = PET(ARCHITECTURAL_HYPERS, 0.0, len(dataset_info.atomic_types))
     model.set_trained_model(raw_pet)
 
     system = systems_to_torch(structure)
@@ -142,7 +141,7 @@ def test_predictions_compatibility(cutoff):
     ARCHITECTURAL_HYPERS = Hypers(DEFAULT_HYPERS["model"])
     batch = get_pyg_graphs(
         [structure],
-        sorted(atomic_types),
+        dataset_info.atomic_types,
         cutoff,
         ARCHITECTURAL_HYPERS.USE_ADDITIONAL_SCALAR_ATTRIBUTES,
         ARCHITECTURAL_HYPERS.USE_LONG_RANGE,

diff --git a/src/metatrain/experimental/pet/tests/test_torchscript.py b/src/metatrain/experimental/pet/tests/test_torchscript.py
@@ -15,7 +15,7 @@ def test_torchscript():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = WrappedPET(DEFAULT_HYPERS["model"], dataset_info)
@@ -30,7 +30,7 @@ def test_torchscript_save_load():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = WrappedPET(DEFAULT_HYPERS["model"], dataset_info)

diff --git a/src/metatrain/experimental/soap_bpnn/model.py b/src/metatrain/experimental/soap_bpnn/model.py
@@ -103,7 +103,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
         self.hypers = model_hypers
         self.dataset_info = dataset_info
         self.new_outputs = list(dataset_info.targets.keys())
-        self.atomic_types = sorted(dataset_info.atomic_types)
+        self.atomic_types = dataset_info.atomic_types
 
         self.soap_calculator = rascaline.torch.SoapPowerSpectrum(
             radial_basis={"Gto": {}}, **self.hypers["soap"]
@@ -198,7 +198,9 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
     def restart(self, dataset_info: DatasetInfo) -> "SoapBpnn":
         # merge old and new dataset info
         merged_info = self.dataset_info.union(dataset_info)
-        new_atomic_types = merged_info.atomic_types - self.dataset_info.atomic_types
+        new_atomic_types = [
+            at for at in merged_info.atomic_types if at not in self.atomic_types
+        ]
         new_targets = merged_info.targets - self.dataset_info.targets
 
         if len(new_atomic_types) > 0:
@@ -212,7 +214,7 @@ def restart(self, dataset_info: DatasetInfo) -> "SoapBpnn":
             self.add_output(output_name)
 
         self.dataset_info = merged_info
-        self.atomic_types = sorted(self.dataset_info.atomic_types)
+        self.atomic_types = sorted(self.atomic_types)
 
         for target_name, target in new_targets.items():
             self.outputs[target_name] = ModelOutput(

diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_continue.py b/src/metatrain/experimental/soap_bpnn/tests/test_continue.py
@@ -26,7 +26,7 @@ def test_continue(monkeypatch, tmp_path):
     target_info_dict["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=target_info_dict
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=target_info_dict
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info)
     output_before = model(systems[:5], {"mtt::U0": model.outputs["mtt::U0"]})

diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py
@@ -18,7 +18,7 @@ def test_to(device, dtype):
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info).to(dtype=dtype)

diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py
@@ -14,7 +14,7 @@ def test_prediction_subset_elements():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
 
@@ -37,7 +37,7 @@ def test_prediction_subset_atoms():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
 
@@ -103,7 +103,7 @@ def test_output_last_layer_features():
     """Tests that the model can output its last layer features."""
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
 
@@ -174,7 +174,7 @@ def test_output_per_atom():
     """Tests that the model can output per-atom quantities."""
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
 

diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_invariance.py b/src/metatrain/experimental/soap_bpnn/tests/test_invariance.py
@@ -15,7 +15,7 @@ def test_rotational_invariance():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info)

diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_regression.py b/src/metatrain/experimental/soap_bpnn/tests/test_regression.py
@@ -25,7 +25,7 @@ def test_regression_init():
     targets["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=targets
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=targets
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info)
 
@@ -76,7 +76,7 @@ def test_regression_train():
     hypers["training"]["num_epochs"] = 2
 
     dataset_info = DatasetInfo(
-        length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=target_info_dict
+        length_unit="Angstrom", atomic_types=[1, 6, 7, 8], targets=target_info_dict
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info)
 

diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py
@@ -14,7 +14,7 @@ def test_torchscript():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info)
@@ -38,7 +38,7 @@ def test_torchscript_with_identity():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     hypers = copy.deepcopy(MODEL_HYPERS)
@@ -64,7 +64,7 @@ def test_torchscript_save_load():
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom",
-        atomic_types={1, 6, 7, 8},
+        atomic_types=[1, 6, 7, 8],
         targets=TargetInfoDict(energy=TargetInfo(quantity="energy", unit="eV")),
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info)