Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Change mtm:: prefix to mtt:: #254

Merged
merged 1 commit into from
Jun 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions docs/src/advanced-concepts/auxiliary-outputs.rst
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
6 changes: 3 additions & 3 deletions docs/src/advanced-concepts/output-naming.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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>``.
6 changes: 3 additions & 3 deletions docs/src/getting-started/custom_dataset_conf.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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:

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Expand All @@ -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",
Expand All @@ -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()

Expand Down Expand Up @@ -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,
)
20 changes: 10 additions & 10 deletions src/metatrain/experimental/gap/tests/test_regression.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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",
Expand All @@ -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
Expand All @@ -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)
Expand All @@ -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,
)


Expand Down
16 changes: 8 additions & 8 deletions src/metatrain/experimental/gap/tests/test_torchscript.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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
Expand Down
10 changes: 5 additions & 5 deletions src/metatrain/experimental/soap_bpnn/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
)

Expand Down Expand Up @@ -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)
)

Expand Down
14 changes: 7 additions & 7 deletions src/metatrain/experimental/soap_bpnn/tests/test_continue.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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
Expand All @@ -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"])
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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",
]

Expand Down
Loading
Loading