Merge branch 'main' into dataset-repr

docs/src/advanced-concepts/auxiliary-outputs.rst

@@ -1,7 +1,7 @@
 Auxiliary outputs
 =================
 
-These outputs, which are idenfified by the ``mtm::aux::`` prefix,
+These outputs, which are idenfified by the ``mtt::aux::`` prefix,
 represent additional information that the model may provide. They are not
 conventional trainable outputs, and they often correspond to internal
 information that the model is capable of providing, such as its internal
@@ -10,7 +10,7 @@ representation.
 The following auxiliary outputs that are currently supported
 by one or more architectures in the library:
 
-- ``mtm::aux::last_layer_features``: The internal representation
+- ``mtt::aux::last_layer_features``: The internal representation
    of the model at the last layer, before the final linear transformation.
 
 The following table shows the architectures that support each of the
@@ -19,13 +19,13 @@ auxiliary outputs:
 +------------------------------------------+-----------+------------------+-----+
 | Auxiliary output                         | SOAP-BPNN | Alchemical Model | PET |
 +------------------------------------------+-----------+------------------+-----+
-| ``mtm::aux::last_layer_features``        | Yes       |       No         | No  |
+| ``mtt::aux::last_layer_features``        | Yes       |       No         | No  |
 +------------------------------------------+-----------+------------------+-----+
 
 The following tables show the metadata that is expected for each of the
 auxiliary outputs:
 
-mtm::aux::last_layer_features
+mtt::aux::last_layer_features
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 .. list-table:: Metadata for last-layer features

docs/src/advanced-concepts/output-naming.rst

@@ -7,9 +7,9 @@ https://lab-cosmo.github.io/metatensor/latest/atomistic/outputs.html>`_
 package. An immediate example is given by the ``energy`` output.
 
 Any additional outputs present within the library are denoted by the
-``mtm::`` prefix. For example, some models can output their last-layer
-features, which are named as ``mtm::aux::last_layer_features``, where
+``mtt::`` prefix. For example, some models can output their last-layer
+features, which are named as ``mtt::aux::last_layer_features``, where
 ``aux`` denotes an auxiliary output.
 
 Outputs that are specific to a particular model should be named as
-``mtm::<model_name>::<output_name>``.
+``mtt::<model_name>::<output_name>``.

docs/src/getting-started/custom_dataset_conf.rst

@@ -88,11 +88,11 @@ Allows defining multiple target sections, each with a unique name.
 - Commonly, a section named ``energy`` should be defined, which is essential for running
   molecular dynamics simulations. For the ``energy`` section gradients like `forces` and
   `stress` are enabled by default.
-- Other target sections can also be defined, as long as they are prefixed by ``mtm::``.
-  For example, ``mtm::free_energy``. In general, all targets that are not standard
+- Other target sections can also be defined, as long as they are prefixed by ``mtt::``.
+  For example, ``mtt::free_energy``. In general, all targets that are not standard
   outputs of ``metatensor.torch.atomistic`` (see
   https://docs.metatensor.org/latest/atomistic/outputs.html) should be prefixed by
-  ``mtm::``.
+  ``mtt::``.
 
 Target section parameters include:
 

src/metatrain/experimental/alchemical_model/tests/test_regression.py

@@ -29,7 +29,7 @@ def test_regression_init():
     """Perform a regression test on the model at initialization"""
 
     targets = TargetInfoDict()
-    targets["mtm::U0"] = TargetInfo(quantity="energy", unit="eV")
+    targets["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=targets
@@ -64,10 +64,10 @@ def test_regression_init():
 
     # if you need to change the hardcoded values:
     # torch.set_printoptions(precision=12)
-    # print(output["mtm::U0"].block().values)
+    # print(output["mtt::U0"].block().values)
 
     torch.testing.assert_close(
-        output["mtm::U0"].block().values,
+        output["mtt::U0"].block().values,
         expected_output,
     )
 
@@ -79,7 +79,7 @@ def test_regression_train():
     systems = read_systems(DATASET_PATH)
 
     conf = {
-        "mtm::U0": {
+        "mtt::U0": {
             "quantity": "energy",
             "read_from": DATASET_PATH,
             "file_format": ".xyz",
@@ -91,7 +91,7 @@ def test_regression_train():
         }
     }
     targets, target_info_dict = read_targets(OmegaConf.create(conf))
-    dataset = Dataset({"system": systems, "mtm::U0": targets["mtm::U0"]})
+    dataset = Dataset({"system": systems, "mtt::U0": targets["mtt::U0"]})
 
     hypers = DEFAULT_HYPERS.copy()
 
@@ -131,9 +131,9 @@ def test_regression_train():
 
     # if you need to change the hardcoded values:
     # torch.set_printoptions(precision=12)
-    # print(output["mtm::U0"].block().values)
+    # print(output["mtt::U0"].block().values)
 
     torch.testing.assert_close(
-        output["mtm::U0"].block().values,
+        output["mtt::U0"].block().values,
         expected_output,
     )

src/metatrain/experimental/gap/tests/test_regression.py

@@ -26,7 +26,7 @@
 def test_regression_init():
     """Perform a regression test on the model at initialization"""
     targets = TargetInfoDict()
-    targets["mtm::U0"] = TargetInfo(quantity="energy", unit="eV")
+    targets["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=targets
@@ -43,7 +43,7 @@ def test_regression_train_and_invariance():
     systems = read_systems(DATASET_PATH, dtype=torch.float64)
 
     conf = {
-        "mtm::U0": {
+        "mtt::U0": {
             "quantity": "energy",
             "read_from": DATASET_PATH,
             "file_format": ".xyz",
@@ -55,10 +55,10 @@ def test_regression_train_and_invariance():
         }
     }
     targets, _ = read_targets(OmegaConf.create(conf), dtype=torch.float64)
-    dataset = Dataset({"system": systems, "mtm::U0": targets["mtm::U0"]})
+    dataset = Dataset({"system": systems, "mtt::U0": targets["mtt::U0"]})
 
     target_info_dict = TargetInfoDict()
-    target_info_dict["mtm::U0"] = TargetInfo(quantity="energy", unit="eV")
+    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
@@ -69,13 +69,13 @@ def test_regression_train_and_invariance():
     trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".")
 
     # Predict on the first five systems
-    output = gap(systems[:5], {"mtm::U0": gap.outputs["mtm::U0"]})
+    output = gap(systems[:5], {"mtt::U0": gap.outputs["mtt::U0"]})
 
     expected_output = torch.tensor(
         [[-40.5891], [-56.7122], [-76.4146], [-77.3364], [-93.4905]]
     )
 
-    assert torch.allclose(output["mtm::U0"].block().values, expected_output, rtol=0.3)
+    assert torch.allclose(output["mtt::U0"].block().values, expected_output, rtol=0.3)
 
     # Tests that the model is rotationally invariant
     system = ase.io.read(DATASET_PATH)
@@ -86,16 +86,16 @@ def test_regression_train_and_invariance():
 
     original_output = gap(
         [metatensor.torch.atomistic.systems_to_torch(original_system)],
-        {"mtm::U0": gap.outputs["mtm::U0"]},
+        {"mtt::U0": gap.outputs["mtt::U0"]},
     )
     rotated_output = gap(
         [metatensor.torch.atomistic.systems_to_torch(system)],
-        {"mtm::U0": gap.outputs["mtm::U0"]},
+        {"mtt::U0": gap.outputs["mtt::U0"]},
     )
 
     assert torch.allclose(
-        original_output["mtm::U0"].block().values,
-        rotated_output["mtm::U0"].block().values,
+        original_output["mtt::U0"].block().values,
+        rotated_output["mtt::U0"].block().values,
     )
 
 

src/metatrain/experimental/gap/tests/test_torchscript.py

@@ -14,13 +14,13 @@
 def test_torchscript():
     """Tests that the model can be jitted."""
     target_info_dict = TargetInfoDict()
-    target_info_dict["mtm::U0"] = TargetInfo(quantity="energy", unit="eV")
+    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
     )
     conf = {
-        "mtm::U0": {
+        "mtt::U0": {
             "quantity": "energy",
             "read_from": DATASET_PATH,
             "file_format": ".xyz",
@@ -36,29 +36,29 @@ def test_torchscript():
 
     # for system in systems:
     #    system.types = torch.ones(len(system.types), dtype=torch.int32)
-    dataset = Dataset({"system": systems, "mtm::U0": targets["mtm::U0"]})
+    dataset = Dataset({"system": systems, "mtt::U0": targets["mtt::U0"]})
 
     hypers = DEFAULT_HYPERS.copy()
     gap = GAP(DEFAULT_HYPERS["model"], dataset_info)
     trainer = Trainer(hypers["training"])
     trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".")
     scripted_gap = torch.jit.script(gap)
 
-    ref_output = gap.forward(systems[:5], {"mtm::U0": gap.outputs["mtm::U0"]})
+    ref_output = gap.forward(systems[:5], {"mtt::U0": gap.outputs["mtt::U0"]})
     scripted_output = scripted_gap.forward(
-        systems[:5], {"mtm::U0": gap.outputs["mtm::U0"]}
+        systems[:5], {"mtt::U0": gap.outputs["mtt::U0"]}
     )
 
     assert torch.allclose(
-        ref_output["mtm::U0"].block().values,
-        scripted_output["mtm::U0"].block().values,
+        ref_output["mtt::U0"].block().values,
+        scripted_output["mtt::U0"].block().values,
     )
 
 
 def test_torchscript_save():
     """Tests that the model can be jitted and saved."""
     targets = TargetInfoDict()
-    targets["mtm::U0"] = TargetInfo(quantity="energy", unit="eV")
+    targets["mtt::U0"] = TargetInfo(quantity="energy", unit="eV")
 
     dataset_info = DatasetInfo(
         length_unit="Angstrom", atomic_types={1, 6, 7, 8}, targets=targets

src/metatrain/experimental/soap_bpnn/model.py

@@ -120,7 +120,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
         }
 
         # the model is always capable of outputting the last layer features
-        self.outputs["mtm::aux::last_layer_features"] = ModelOutput(per_atom=True)
+        self.outputs["mtt::aux::last_layer_features"] = ModelOutput(per_atom=True)
 
         # creates a composition weight tensor that can be directly indexed by species,
         # this can be left as a tensor of zero or set from the outside using
@@ -189,7 +189,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None:
                     ],
                 )
                 for output_name in self.outputs.keys()
-                if "mtm::aux::" not in output_name
+                if "mtt::aux::" not in output_name
             }
         )
 
@@ -243,14 +243,14 @@ def forward(
         last_layer_features = self.bpnn(soap_features)
 
         # output the hidden features, if requested:
-        if "mtm::aux::last_layer_features" in outputs:
-            last_layer_features_options = outputs["mtm::aux::last_layer_features"]
+        if "mtt::aux::last_layer_features" in outputs:
+            last_layer_features_options = outputs["mtt::aux::last_layer_features"]
             out_features = last_layer_features.keys_to_properties(
                 self.center_type_labels.to(device)
             )
             if not last_layer_features_options.per_atom:
                 out_features = metatensor.torch.sum_over_samples(out_features, ["atom"])
-            return_dict["mtm::aux::last_layer_features"] = (
+            return_dict["mtt::aux::last_layer_features"] = (
                 _remove_center_type_from_properties(out_features)
             )
 

src/metatrain/experimental/soap_bpnn/tests/test_continue.py

@@ -22,16 +22,16 @@ def test_continue(monkeypatch, tmp_path):
     systems = read_systems(DATASET_PATH)
 
     target_info_dict = TargetInfoDict()
-    target_info_dict["mtm::U0"] = TargetInfo(quantity="energy", unit="eV")
+    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
     )
     model = SoapBpnn(MODEL_HYPERS, dataset_info)
-    output_before = model(systems[:5], {"mtm::U0": model.outputs["mtm::U0"]})
+    output_before = model(systems[:5], {"mtt::U0": model.outputs["mtt::U0"]})
 
     conf = {
-        "mtm::U0": {
+        "mtt::U0": {
             "quantity": "energy",
             "read_from": DATASET_PATH,
             "file_format": ".xyz",
@@ -43,7 +43,7 @@ def test_continue(monkeypatch, tmp_path):
         }
     }
     targets, _ = read_targets(OmegaConf.create(conf))
-    dataset = Dataset({"system": systems, "mtm::U0": targets["mtm::U0"]})
+    dataset = Dataset({"system": systems, "mtt::U0": targets["mtt::U0"]})
 
     hypers = DEFAULT_HYPERS.copy()
     hypers["training"]["num_epochs"] = 0
@@ -57,8 +57,8 @@ def test_continue(monkeypatch, tmp_path):
 
     # Predict on the first five systems
     output_before = model_before(
-        systems[:5], {"mtm::U0": model_before.outputs["mtm::U0"]}
+        systems[:5], {"mtt::U0": model_before.outputs["mtt::U0"]}
     )
-    output_after = model_after(systems[:5], {"mtm::U0": model_after.outputs["mtm::U0"]})
+    output_after = model_after(systems[:5], {"mtt::U0": model_after.outputs["mtt::U0"]})
 
-    assert metatensor.torch.allclose(output_before["mtm::U0"], output_after["mtm::U0"])
+    assert metatensor.torch.allclose(output_before["mtt::U0"], output_after["mtt::U0"])

src/metatrain/experimental/soap_bpnn/tests/test_functionality.py

@@ -127,12 +127,12 @@ def test_output_last_layer_features():
         [system],
         {
             "energy": model.outputs["energy"],
-            "mtm::aux::last_layer_features": ll_output_options,
+            "mtt::aux::last_layer_features": ll_output_options,
         },
     )
     assert "energy" in outputs
-    assert "mtm::aux::last_layer_features" in outputs
-    last_layer_features = outputs["mtm::aux::last_layer_features"].block()
+    assert "mtt::aux::last_layer_features" in outputs
+    last_layer_features = outputs["mtt::aux::last_layer_features"].block()
     assert last_layer_features.samples.names == [
         "system",
         "atom",
@@ -155,17 +155,17 @@ def test_output_last_layer_features():
         [system],
         {
             "energy": model.outputs["energy"],
-            "mtm::aux::last_layer_features": ll_output_options,
+            "mtt::aux::last_layer_features": ll_output_options,
         },
     )
     assert "energy" in outputs
-    assert "mtm::aux::last_layer_features" in outputs
-    assert outputs["mtm::aux::last_layer_features"].block().samples.names == ["system"]
-    assert outputs["mtm::aux::last_layer_features"].block().values.shape == (
+    assert "mtt::aux::last_layer_features" in outputs
+    assert outputs["mtt::aux::last_layer_features"].block().samples.names == ["system"]
+    assert outputs["mtt::aux::last_layer_features"].block().values.shape == (
         1,
         128,
     )
-    assert outputs["mtm::aux::last_layer_features"].block().properties.names == [
+    assert outputs["mtt::aux::last_layer_features"].block().properties.names == [
         "properties",
     ]