diff --git a/.gitignore b/.gitignore index 11ca4b36fa..33bbd34516 100644 --- a/.gitignore +++ b/.gitignore @@ -188,6 +188,9 @@ docs/src/mtt_README.md # Sphinx architecture docs docs/src/architectures/generated/* docs/src/architectures/default_hypers/* +# Sphinx hooks docs +docs/src/hooks/generated/* +docs/src/hooks/default_hypers/* # JavaScript node_modules/ diff --git a/docs/src/concepts/hooks.rst b/docs/src/concepts/hooks.rst new file mode 100644 index 0000000000..0d0d0caf79 --- /dev/null +++ b/docs/src/concepts/hooks.rst @@ -0,0 +1,6 @@ +.. _hooks: + +Using hooks +=========== + +Hooks are a powerful way to extend the functionality of architectures in ``metatrain``. diff --git a/docs/src/concepts/index.rst b/docs/src/concepts/index.rst index cb983f88db..24a415dcd8 100644 --- a/docs/src/concepts/index.rst +++ b/docs/src/concepts/index.rst @@ -12,6 +12,7 @@ such as output naming, auxiliary outputs, and wrapper models. output-naming fine-tuning loss-functions + hooks scale-targets auxiliary-outputs batch-bounds diff --git a/docs/src/conf.py b/docs/src/conf.py index f719ca87ba..7f17f2715c 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -29,6 +29,7 @@ sys.path.append(os.path.join(ROOT, "docs")) from generate_examples.conf import sphinx_gallery_conf # noqa from src.architectures.generate import setup_architectures_docs # noqa +from src.hooks.generate_hooks_docs import setup_hooks_docs # noqa # -- Project information ----------------------------------------------------- @@ -89,6 +90,7 @@ def generate_examples(): def setup(app): copy_readme() generate_examples() + setup_hooks_docs() setup_architectures_docs() @@ -119,6 +121,7 @@ def setup(app): "sg_execution_times.rst", "architectures/templates/*", "architectures/README.md", + "hooks/templates/*", ] diff --git a/docs/src/hooks/generate_hooks_docs.py b/docs/src/hooks/generate_hooks_docs.py new file mode 100644 index 0000000000..f883c2d9cd --- /dev/null +++ b/docs/src/hooks/generate_hooks_docs.py @@ -0,0 +1,198 @@ +import ast +from pathlib import Path +from typing import TypedDict + +from jinja2 import Environment, FileSystemLoader + +from metatrain.utils import hooks as hooks_module +from metatrain.utils.hooks.helpers import ( + find_all_hooks, + get_hypers_class, + preload_documentation_module, + write_hypers_yaml, +) +from metatrain.utils.hypers import get_hypers_list + + +HOOKS_DIR = Path(__file__).parent +TEMPLATES_DIR = HOOKS_DIR / "templates" +DEFAULT_HYPERS_DIR = HOOKS_DIR / "default_hypers" +GENERATED_DIR = HOOKS_DIR / "generated" + + +JINJA_ENV = Environment( + loader=FileSystemLoader(TEMPLATES_DIR), + trim_blocks=True, + lstrip_blocks=True, +) + + +SECTIONS = [ + "installation", + "hook_hypers", + "references", +] + + +class HookDocVariables(TypedDict): + """Variables to use inside the hook documentation. + + The docstring of the hook will be processed as a + ``jinja`` template. You can find documentation about them + `here `_ , but + the simplest functionality consists of using variables enclosed in + double curly braces ``{{variable_name}}``, which will be replaced by + their corresponding value. + + For example, a file with the following content: + + .. code-block:: rst + + This is the documentation for {{hook}}. + + generates a documentation file that for the hook ``tensor_basis`` would be: + + .. code-block:: rst + + This is the documentation for tensor_basis. + + There are some special variables that start with ``SECTION_``. These contain + the content of different sections of the documentation, and they will be + appended to the docstring if they are not already present. For example, given + the docstring: + + .. code-block:: python + + \""" + My hook + ======= + + This is my hook. + + {{SECTION_DEFAULT_HYPERS}} + + Some important section + ====================== + + Explain something important here. + \""" + + The final documentation will append to the docstring all the sections except + ``SECTION_DEFAULT_HYPERS``, since it is already present. + + Following you can find a description of all the available variables. The + sections are appended in the order documented here. + """ + + SECTION_INSTALLATION: str + """Section containing installation instructions for this hook.""" + SECTION_HOOK_HYPERS: str + """Section containing the description of the hook hyperparameters for + this hook.""" + SECTION_REFERENCES: str + """Section containing references for this hook. It will render the + references that have been used as ``:footcite:p:`` during the hook + documentation.""" + + hook: str + """The name of the hook. + + This excludes any 'experimental.' or 'deprecated.' prefix.""" + default_hypers_path: str + """Path to the yaml file with the default hyperparameters for this + hook. + + This is a path relative to the ``docs/src/hooks/generated`` + directory. + """ + hook_hypers_path: str + """The full python import path to the hook's hypers class of this + hook. + + E.g.: ``"metatrain.utils.hooks..Hypers"`` + """ + hook_hypers: list[str] + """List of hyperparameter names for this hook.""" + + +def setup_hooks_docs(): + """Generate the hook documentation files. + + This function goes through all available hooks, and for each of them + generates a yaml file with the default hyperparameters (so that it can be + easily included in the documentation) and their rst documentation file. + + See :ref:`newarchitecture-documentation-page` for more information. + """ + # If the default_hypers directory does not exist, create it + DEFAULT_HYPERS_DIR.mkdir(exist_ok=True) + # Same for the generated directory + GENERATED_DIR.mkdir(exist_ok=True) + + for name in find_all_hooks(): + # Load documentation module in an isolated way to avoid + # requiring dependencies for every architecture. + preload_documentation_module(name) + + # Write default hypers file + yaml_path = DEFAULT_HYPERS_DIR / f"{name}-default-hypers.yaml" + write_hypers_yaml(name, yaml_path, include_name=True) + + generate_rst(name, yaml_path=yaml_path) + + +def generate_rst( + hook_name: str, + yaml_path: Path, +): + """Generate the rst documentation file for a given hook. + + :param hook_name: The name of the hook to generate the + documentation for. + :param yaml_path: Path to the yaml file with the default hyperparameters + for this architecture. + """ + + # Get the full python import path to the hook + hook_path = f"metatrain.utils.hooks.{hook_name}" + + # Get the docstring from the documentation.py file + doc_file = Path(hooks_module.__file__).parent / hook_name / "documentation.py" + with open(doc_file, "r") as f: + module = ast.parse(f.read(), filename=str(doc_file)) + docstring = ast.get_docstring(module) + if docstring is None: + raise ValueError( + f"The documentation.py file for hook " + f"'{hook_name}' does not have a module docstring." + ) + + hypers_class = get_hypers_class(hook_name) + + # Prepare template variables + template_variables = dict( + hook=hook_name, + default_hypers_path=".." / yaml_path.relative_to(HOOKS_DIR), + hook_hypers_path=f"{hook_path}.documentation.Hypers", + hook_hypers=get_hypers_list(hypers_class), + ) + + # Read section templates and render them + for section in SECTIONS: + template = JINJA_ENV.get_template(f"{section}.rst") + template_variables[f"SECTION_{section.upper()}"] = template.render( + **template_variables + ) + + # Check for missing sections and add them to the end of the docstring + for section in SECTIONS: + section_var = "{{SECTION_" + section.upper() + "}}" + if section_var not in docstring: + docstring += f"\n\n{section_var}" + + # Render docstring template + docstring = JINJA_ENV.from_string(docstring).render(**template_variables) + + # Write to file + with open(GENERATED_DIR / f"{hook_name}.rst", "w") as f: + f.write(docstring + "\n") diff --git a/docs/src/hooks/index.rst b/docs/src/hooks/index.rst new file mode 100644 index 0000000000..e5260ba94c --- /dev/null +++ b/docs/src/hooks/index.rst @@ -0,0 +1,13 @@ +.. _available-hooks: + +Available Hooks +=============== + +This is a list of all hooks available in ``metatrain``. +The concept of hooks is explained in :ref:`hooks`. + +.. toctree:: + :maxdepth: 1 + :glob: + + ./generated/* diff --git a/docs/src/hooks/templates/description.rst b/docs/src/hooks/templates/description.rst new file mode 100644 index 0000000000..049d3c1e9e --- /dev/null +++ b/docs/src/hooks/templates/description.rst @@ -0,0 +1,5 @@ +{{hook}} +============== + +This page gives an overview of the ``{{hook}}`` hook available in +the ``metatrain`` package. diff --git a/docs/src/hooks/templates/hook_hypers.rst b/docs/src/hooks/templates/hook_hypers.rst new file mode 100644 index 0000000000..8f59fe0cc7 --- /dev/null +++ b/docs/src/hooks/templates/hook_hypers.rst @@ -0,0 +1,24 @@ +.. _hook-{{hook}}_hypers: + +Hook hyperparameters +------------------------ + +{% if hook_hypers %} +The default hyperparameters for this hook are: + +.. literalinclude:: {{default_hypers_path}} + :language: yaml + +and here is the documentation for each hyperparameter: + +.. container:: mtt-hypers-remove-classname + + .. + + {% for hyper in hook_hypers %} + .. autoattribute:: {{hook_hypers_path}}.{{hyper}} + + {% endfor %} +{% else %} +This hook has no hyperparameters. +{% endif %} diff --git a/docs/src/hooks/templates/installation.rst b/docs/src/hooks/templates/installation.rst new file mode 100644 index 0000000000..dc0bc7fd2f --- /dev/null +++ b/docs/src/hooks/templates/installation.rst @@ -0,0 +1,13 @@ +.. _hook-{{hook}}_installation: + +Installation +------------ + +To install this hook along with the ``metatrain`` package, run: + +.. code-block:: bash + + pip install metatrain[hook-{{hook}}] + +where the square brackets indicate that you want to install the optional +dependencies required for the ``{{hook}}`` hook. diff --git a/docs/src/hooks/templates/references.rst b/docs/src/hooks/templates/references.rst new file mode 100644 index 0000000000..c03935e38f --- /dev/null +++ b/docs/src/hooks/templates/references.rst @@ -0,0 +1,6 @@ +.. _hook-{{hook}}_references: + +References +---------- + +.. footbibliography:: diff --git a/docs/src/index.rst b/docs/src/index.rst index 118105958d..9273c0cbf2 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -21,6 +21,7 @@ installation getting-started/index architectures/index + hooks/index generated_examples/index concepts/index faq diff --git a/docs/static/refs.bib b/docs/static/refs.bib index 414263aaeb..0ac76c51bb 100644 --- a/docs/static/refs.bib +++ b/docs/static/refs.bib @@ -119,3 +119,20 @@ @misc{dpa3_2025 doi={10.48550/arXiv.2506.01686} } +@article{domina2025representing, + title={Representing spherical tensors with scalar-based machine-learning models}, + author={Domina, Michelangelo and Bigi, Filippo and Pegolo, Paolo and Ceriotti, Michele}, + journal={The Journal of Chemical Physics}, + volume={163}, + number={16}, + year={2025}, + publisher={AIP Publishing} +} + +@article{malosso2026transferable, + title={Transferable machine learning of excited-state dynamics with extremal pooling}, + author={Malosso, Cesare and How, Wei Bin and Mir{\'o}n, Gonzalo D{\'\i}az and Hassanali, Ali and Ceriotti, Michele}, + journal={arXiv preprint arXiv:2606.16859}, + year={2026} +} + diff --git a/src/metatrain/cli/train.py b/src/metatrain/cli/train.py index bc78324d0d..5f02a6f0bf 100644 --- a/src/metatrain/cli/train.py +++ b/src/metatrain/cli/train.py @@ -7,7 +7,7 @@ import re import shutil from pathlib import Path -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Union, cast import numpy as np import torch @@ -578,6 +578,12 @@ def train_model( atomic_types=atomic_types, targets=target_info_dict, extra_data=extra_data_info_dict, + # Resolved to plain containers: ``DatasetInfo`` is an attribute of the + # model, and TorchScript cannot infer the type of an OmegaConf node, + # which is what a hook configured with a nested mapping would give. + hooks=cast(dict, OmegaConf.to_container(options["hooks"], resolve=True)) + if "hooks" in options + else {}, ) ########################### diff --git a/src/metatrain/experimental/mace/model.py b/src/metatrain/experimental/mace/model.py index d73b960349..6a72c751ee 100644 --- a/src/metatrain/experimental/mace/model.py +++ b/src/metatrain/experimental/mace/model.py @@ -29,6 +29,7 @@ sparsify_atomic_basis_target, ) from metatrain.utils.dtype import dtype_to_str +from metatrain.utils.hooks import restart_hooks, setup_hooks from metatrain.utils.metadata import merge_metadata from metatrain.utils.scaler import Scaler from metatrain.utils.sum_over_atoms import sum_over_atoms @@ -263,10 +264,13 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None: # the model during training. train_dataset_info = self._train_dataset_info(dataset_info) + forward_hooks, model_outs = setup_hooks(train_dataset_info) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + # Create heads for each target, store the layout for each of them. self.heads = torch.nn.ModuleDict() self.layouts: Dict[str, TensorMap] = {} - for target_name, target_info in train_dataset_info.targets.items(): + for target_name, target_info in model_outs.items(): self._add_output(target_name, target_info) self.layouts["mtt::aux::mace_features"] = get_e3nn_mts_layout( @@ -279,13 +283,14 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None: ) targets = dataset_info.targets + all_names = set([*train_dataset_info.targets, *model_outs, *self.layouts]) self.outputs = { k: ModelOutput( quantity=targets[k].quantity if k in targets else "", unit=targets[k].unit if k in targets else "", sample_kind="atom", ) - for k in self.layouts + for k in all_names } # --------------------------- @@ -329,9 +334,28 @@ def restart(self, dataset_info: DatasetInfo) -> "MetaMACE": # the model during training. train_dataset_info = self._train_dataset_info(dataset_info) - # Add extra heads for the new targets - for target_name in new_targets: - self._add_output(target_name, train_dataset_info.targets[target_name]) + if dataset_info.targets != self.dataset_info.targets: + # Only re-run the hook setup when the targets changed; ``restart_hooks`` + # does not support rebuilding already-instantiated hooks. + forward_hooks, model_outs = restart_hooks( + list(self.forward_hooks), train_dataset_info + ) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + + # Add extra heads for the new targets. Targets produced by a hook + # are not in ``model_outs``, since the model does not predict them + # directly; they only need to be registered as outputs. + targets = merged_info.targets + for target_name in new_targets: + if target_name in model_outs: + self._add_output(target_name, model_outs[target_name]) + self.outputs[target_name] = ModelOutput( + quantity=targets[target_name].quantity + if target_name in targets + else "", + unit=targets[target_name].unit if target_name in targets else "", + sample_kind="atom", + ) self.dataset_info = merged_info @@ -358,6 +382,14 @@ def forward( selected_atoms: Optional[Labels] = None, ) -> Dict[str, TensorMap]: + # ---------------------------- + # Add outputs needed by hooks + # ---------------------------- + # TODO: In reality, we would have to check if the hook's output is requested + for hook in self.forward_hooks: + requested_inputs = hook.requested_inputs() + outputs.update(requested_inputs) + # -------------------------- # Moving to device and dtype # -------------------------- @@ -460,6 +492,20 @@ def forward( else sum_over_atoms(per_atom_output) ) + # ----------------------------------- + # Apply hooks + # ----------------------------------- + + for hook in self.forward_hooks: + return_dict.update( + hook( + systems, + outputs, + return_dict, + selected_atoms, + ) + ) + # ----------------------------------------- # Undo data preprocessing (eval only) # ----------------------------------------- diff --git a/src/metatrain/experimental/mace/trainer.py b/src/metatrain/experimental/mace/trainer.py index b10211deca..48ba3cfdc8 100644 --- a/src/metatrain/experimental/mace/trainer.py +++ b/src/metatrain/experimental/mace/trainer.py @@ -114,6 +114,10 @@ def get_optimizer_and_scheduler( "name": "scaler", "params": model.scaler.parameters(), }, + { + "name": "forward_hooks", + "params": model.forward_hooks.parameters(), + }, ] ) diff --git a/src/metatrain/experimental/space/model.py b/src/metatrain/experimental/space/model.py index f1271e021c..88abbb305e 100644 --- a/src/metatrain/experimental/space/model.py +++ b/src/metatrain/experimental/space/model.py @@ -34,6 +34,7 @@ ) from metatrain.utils.data.dataset import DatasetInfo, TargetInfo from metatrain.utils.dtype import dtype_to_str +from metatrain.utils.hooks import restart_hooks, setup_hooks from metatrain.utils.metadata import merge_metadata from metatrain.utils.scaler import Scaler @@ -82,7 +83,18 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None: # Two types of model wrapper: one with gradients (training) and one without # (torchscript-based export). - base_model = BaseModel(hypers, train_dataset_info) + # The hooks may replace some targets with the per-atom inputs they + # consume, so they are set up before the heads are built. + forward_hooks, model_outs = setup_hooks(train_dataset_info) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + hooked_dataset_info = DatasetInfo( + length_unit=train_dataset_info.length_unit, + atomic_types=train_dataset_info.atomic_types, + targets=model_outs, + extra_data=train_dataset_info.extra_data, + ) + + base_model = BaseModel(hypers, hooked_dataset_info) self.fake_gradient_model = FakeGradientModel(base_model) self.gradient_model = GradientModel(base_model) self.module = self.fake_gradient_model @@ -123,10 +135,24 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None: self.mlp_head_num_layers = self.hypers["mlp_head_num_layers"] self.target_names: List[str] = [] - for target_name, target_info in train_dataset_info.targets.items(): + for target_name, target_info in model_outs.items(): self.target_names.append(target_name) self._add_output(target_name, target_info) + # Register outputs that are produced only by post-processing hooks (i.e. + # those removed from ``model_outs`` because SPACE does not predict them + # directly). + targets = dataset_info.targets + for target_name in train_dataset_info.targets: + if target_name not in model_outs: + self.outputs[target_name] = ModelOutput( + quantity=targets[target_name].quantity + if target_name in targets + else "", + unit=targets[target_name].unit if target_name in targets else "", + sample_kind="atom", + ) + self.last_layer_feature_size = self.k_max_l[0] # additive models: these are handled by the trainer at training @@ -195,10 +221,19 @@ def restart(self, dataset_info: DatasetInfo) -> "SPACE": # the model during training. train_dataset_info = self._train_dataset_info(dataset_info) - # register new outputs as new last layers - for target_name in new_targets: - self.target_names.append(target_name) - self._add_output(target_name, train_dataset_info.targets[target_name]) + if dataset_info.targets != self.dataset_info.targets: + # Only re-run the hook setup when the targets changed; ``restart_hooks`` + # does not support rebuilding already-instantiated hooks. + forward_hooks, model_outs = restart_hooks( + list(self.forward_hooks), train_dataset_info + ) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + + # register new outputs as new last layers + for target_name in new_targets: + if target_name in model_outs: + self.target_names.append(target_name) + self._add_output(target_name, model_outs[target_name]) self.dataset_info = merged_info @@ -230,6 +265,11 @@ def forward( outputs: Dict[str, ModelOutput], selected_atoms: Optional[Labels] = None, ) -> Dict[str, TensorMap]: + # Add the inputs that the post-processing hooks consume to the outputs + # the model is asked to produce. + for hook in self.forward_hooks: + outputs.update(hook.requested_inputs()) + # transfer labels, if needed device = systems[0].device if self.single_label.values.device != device: @@ -356,7 +396,13 @@ def forward( # remaining outputs (main outputs) for output_name in outputs.keys(): - if output_name == "feature" or output_name.startswith("mtt::aux::"): + if output_name == "feature" or output_name.endswith( + "_last_layer_features" + ): + continue + if output_name not in predictions: + # Produced by a post-processing hook rather than by the model + # itself, so there is nothing to read from the heads here. continue output_as_tensor_dict = predictions[output_name] return_dict[output_name] = TensorMap( @@ -483,6 +529,10 @@ def forward( return_dict[output_name], self.final_scaling ) + # Apply the post-processing hooks + for hook in self.forward_hooks: + return_dict.update(hook(systems, outputs, return_dict, selected_atoms)) + if not self.training: # at evaluation, we also introduce the scaler and additive contributions return_dict = self.scaler( diff --git a/src/metatrain/pet/model.py b/src/metatrain/pet/model.py index 29aacaf270..34332e415e 100644 --- a/src/metatrain/pet/model.py +++ b/src/metatrain/pet/model.py @@ -25,6 +25,7 @@ sparsify_atomic_basis_target, ) from metatrain.utils.dtype import dtype_to_str +from metatrain.utils.hooks import restart_hooks, setup_hooks from metatrain.utils.long_range import DummyLongRangeFeaturizer, LongRangeFeaturizer from metatrain.utils.metadata import merge_metadata from metatrain.utils.scaler import Scaler @@ -130,16 +131,36 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None: # during training. train_dataset_info = self._train_dataset_info(dataset_info) + forward_hooks, model_outs = setup_hooks(train_dataset_info) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + self.output_shapes: Dict[str, Dict[str, List[int]]] = {} self.key_labels: Dict[str, Labels] = {} self.property_labels: Dict[str, List[Labels]] = {} self.component_labels: Dict[str, List[List[Labels]]] = {} self.target_names: List[str] = [] self.last_layer_parameter_names: Dict[str, List[str]] = {} # for LLPR - for target_name, target_info in train_dataset_info.targets.items(): + for target_name, target_info in model_outs.items(): self.target_names.append(target_name) self._add_output(target_name, target_info) + # Register outputs that are produced only by post-processing hooks (i.e. + # those removed from ``model_outs`` because PET itself does not predict + # them directly). + targets = dataset_info.targets + for target_name in train_dataset_info.targets: + if target_name not in model_outs: + self.outputs[target_name] = ModelOutput( + quantity=targets[target_name].quantity + if target_name in targets + else "", + unit=targets[target_name].unit if target_name in targets else "", + sample_kind="atom", + description=targets[target_name].description + if target_name in targets + else "", + ) + # long-range module if self.hypers["long_range"]["enable"]: self.long_range = True @@ -203,20 +224,6 @@ def restart(self, dataset_info: DatasetInfo) -> "PET": new_atomic_types = [ at for at in merged_info.atomic_types if at not in self.atomic_types ] - new_targets = { - key: value - for key, value in merged_info.targets.items() - if key not in self.dataset_info.targets - } - self.has_new_targets = len(new_targets) > 0 - - # Targets that were present before this run but are not part of the current - # run's dataset: with a backbone-altering finetuning method (full/lora), their - # heads are no longer meaningful and are dropped once training starts, by - # ``apply_finetuning_strategy`` (which decides based on the method). - stale_targets = compute_stale_targets( - self.dataset_info.targets, dataset_info.targets - ) if len(new_atomic_types) > 0: raise ValueError( @@ -226,12 +233,56 @@ def restart(self, dataset_info: DatasetInfo) -> "PET": # Modified dataset_info with the targets as they will be seen by PET # during training. - train_dataset_info = self._train_dataset_info(dataset_info) + train_dataset_info = self._train_dataset_info(merged_info) - # register new outputs as new last layers - for target_name in new_targets: - self.target_names.append(target_name) - self._add_output(target_name, train_dataset_info.targets[target_name]) + if dataset_info.targets != self.dataset_info.targets: + forward_hooks, model_outs = restart_hooks( + self.forward_hooks, train_dataset_info + ) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + + new_targets = { + key: value + for key, value in model_outs.items() + if key not in self.target_names + } + self.has_new_targets = len(new_targets) > 0 + + # Targets that were present before this run but are not part of the current + # run's dataset: with a backbone-altering finetuning method (full/lora), + # their heads are no longer meaningful and are dropped once training starts, + # by ``apply_finetuning_strategy`` (which decides based on the method). + stale_targets = compute_stale_targets( + {name: self.dataset_info.targets[name] for name in self.target_names}, + model_outs, + ) + + # register new outputs as new last layers + for target_name in new_targets: + self.target_names.append(target_name) + self._add_output(target_name, model_outs[target_name]) + + # Register outputs that are produced only by post-processing hooks (i.e. + # those removed from ``model_outs`` because PET itself does not predict + # them directly). + targets = dataset_info.targets + for target_name in train_dataset_info.targets: + if target_name not in model_outs: + self.outputs[target_name] = ModelOutput( + quantity=targets[target_name].quantity + if target_name in targets + else "", + unit=targets[target_name].unit + if target_name in targets + else "", + sample_kind="atom", + description=targets[target_name].description + if target_name in targets + else "", + ) + else: + self.has_new_targets = False + stale_targets = [] self.dataset_info = merged_info @@ -384,6 +435,14 @@ def forward( predictions (depending on the ModelOutput configuration) with appropriate metatensor metadata (samples, components, properties). """ + # ---------------------------- + # Add outputs needed by hooks + # ---------------------------- + # TODO: In reality, we would have to check if the hook's output is requested + for hook in self.forward_hooks: + requested_inputs = hook.requested_inputs() + outputs.update(requested_inputs) + device = systems[0].device return_dict: Dict[str, TensorMap] = {} nl_options = self.requested_neighbor_lists()[0] @@ -581,6 +640,19 @@ def forward( h.remove() # ===== END DIAGNOSTIC-RELATED BLOCK + # ----------------------------------- + # Apply hooks + # ----------------------------------- + for hook in self.forward_hooks: + return_dict.update( + hook( + systems, + outputs, + return_dict, + selected_atoms, + ) + ) + # **Post-processing (Evaluation Only)** with torch.profiler.record_function("PET::post-processing"): if not self.training: @@ -598,7 +670,10 @@ def forward( # done before adding the additive contributions, which are also # sparsified (by the additive models themselves, in eval mode). for k in atomic_predictions_dict.keys(): - if self.dataset_info.targets[k].is_atomic_basis: + if ( + k in self.dataset_info.targets + and self.dataset_info.targets[k].is_atomic_basis + ): return_dict[k] = sparsify_atomic_basis_target( systems, return_dict[k], @@ -1034,7 +1109,7 @@ def _add_output(self, target_name: str, target_info: TargetInfo) -> None: :param target_info: TargetInfo object containing details about the target. """ # one output shape for each tensor block, grouped by target (i.e. tensormap) - self.output_shapes[target_name] = {} + self.output_shapes[target_name] = torch.jit.annotate(Dict[str, List[int]], {}) for key, block in target_info.layout.items(): dict_key = target_name for n, k in zip(key.names, key.values, strict=True): @@ -1055,7 +1130,7 @@ def _add_output(self, target_name: str, target_info: TargetInfo) -> None: # Register last-layer parameters, in the same order as they are returned as # last-layer features in the model (the modules live on ``self.backend``). - self.last_layer_parameter_names[target_name] = [] + self.last_layer_parameter_names[target_name] = torch.jit.annotate(List[str], []) for layer_index in range(self.num_readout_layers): for key in self.output_shapes[target_name].keys(): self.last_layer_parameter_names[target_name].append( diff --git a/src/metatrain/share/base_hypers.py b/src/metatrain/share/base_hypers.py index 19803c399c..b6ba679f7e 100644 --- a/src/metatrain/share/base_hypers.py +++ b/src/metatrain/share/base_hypers.py @@ -515,3 +515,8 @@ class BaseHypers(TypedDict): a full dataset specification, or an ``indices`` dict referencing the training source file. """ + + hooks: NotRequired[dict] = {} + """Hooks that the model needs to apply at the end of forward to + produce the final outputs. + """ diff --git a/src/metatrain/soap_bpnn/model.py b/src/metatrain/soap_bpnn/model.py index 548c7ee0ed..f7a2353007 100644 --- a/src/metatrain/soap_bpnn/model.py +++ b/src/metatrain/soap_bpnn/model.py @@ -26,6 +26,7 @@ ) from metatrain.utils.data.dataset import DatasetInfo from metatrain.utils.dtype import dtype_to_str +from metatrain.utils.hooks import restart_hooks, setup_hooks from metatrain.utils.long_range import DummyLongRangeFeaturizer, LongRangeFeaturizer from metatrain.utils.metadata import merge_metadata from metatrain.utils.scaler import Scaler @@ -362,6 +363,9 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None: # the model during training. train_dataset_info = self._train_dataset_info(dataset_info) + forward_hooks, model_outs = setup_hooks(train_dataset_info) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + self.num_properties: Dict[str, Dict[str, int]] = {} # by target and block self.basis_calculators = torch.nn.ModuleDict({}) self.heads = torch.nn.ModuleDict({}) @@ -373,9 +377,23 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None: self.last_layer_parameter_names: Dict[str, List[str]] = {} # for LLPR self.cartesian_rank1_targets: List[str] = [] self.cartesian_rank2_targets: List[str] = [] - for target_name, target in train_dataset_info.targets.items(): + for target_name, target in model_outs.items(): self._add_output(target_name, target) + # Register outputs that are produced only by post-processing hooks (i.e. + # those removed from ``model_outs`` because SOAP-BPNN does not predict + # them directly). + targets = dataset_info.targets + for target_name in train_dataset_info.targets: + if target_name not in model_outs: + self.outputs[target_name] = ModelOutput( + quantity=targets[target_name].quantity + if target_name in targets + else "", + unit=targets[target_name].unit if target_name in targets else "", + sample_kind="atom", + ) + # Pre-compute spherical→Cartesian conversion matrix for rank-2 tensors. # W[i,j,M] maps 9 spherical components (l=0,1,2) to 3×3 Cartesian. # Convention: m=-1→y, m=0→z, m=1→x (same as _to_cartesian_rank_1). @@ -444,9 +462,18 @@ def restart(self, dataset_info: DatasetInfo) -> "SoapBpnn": # the model during training. train_dataset_info = self._train_dataset_info(dataset_info) - # register new outputs as new last layers - for target_name in new_targets: - self._add_output(target_name, train_dataset_info.targets[target_name]) + if dataset_info.targets != self.dataset_info.targets: + # Only re-run the hook setup when the targets changed; ``restart_hooks`` + # does not support rebuilding already-instantiated hooks. + forward_hooks, model_outs = restart_hooks( + list(self.forward_hooks), train_dataset_info + ) + self.forward_hooks = torch.nn.ModuleList(forward_hooks) + + # register new outputs as new last layers + for target_name in new_targets: + if target_name in model_outs: + self._add_output(target_name, model_outs[target_name]) self.dataset_info = merged_info @@ -472,6 +499,11 @@ def forward( outputs: Dict[str, ModelOutput], selected_atoms: Optional[Labels] = None, ) -> Dict[str, TensorMap]: + # Add the inputs that the post-processing hooks consume to the outputs + # the model is asked to produce. + for hook in self.forward_hooks: + outputs.update(hook.requested_inputs()) + device = systems[0].positions.device if self.neighbors_species_labels.device != device: self.neighbors_species_labels = self.neighbors_species_labels.to(device) @@ -813,6 +845,10 @@ def forward( # sum the atomic property to get the total property return_dict[output_name] = sum_over_atoms(atomic_property) + # Apply the post-processing hooks + for hook in self.forward_hooks: + return_dict.update(hook(systems, outputs, return_dict, selected_atoms)) + if not self.training: # at evaluation, we also introduce the scaler and additive contributions return_dict = self.scaler( diff --git a/src/metatrain/utils/data/atomic_basis_helpers.py b/src/metatrain/utils/data/atomic_basis_helpers.py index db574a4235..eb23d6575d 100644 --- a/src/metatrain/utils/data/atomic_basis_helpers.py +++ b/src/metatrain/utils/data/atomic_basis_helpers.py @@ -642,4 +642,5 @@ def densify_atomic_basis_dataset_info(dataset_info: DatasetInfo) -> DatasetInfo: ) for target_name, target_info in dataset_info.targets.items() }, + hooks=dataset_info.hooks, ) diff --git a/src/metatrain/utils/data/dataset.py b/src/metatrain/utils/data/dataset.py index 834aa283f1..8c2aec8a07 100644 --- a/src/metatrain/utils/data/dataset.py +++ b/src/metatrain/utils/data/dataset.py @@ -83,6 +83,7 @@ class DatasetInfo: :param targets: Information about targets in the dataset. :param extra_data: Optional dictionary containing additional data that is not used as a target, but is still relevant to the dataset. + :param hooks: Hooks that models should use to fit the targets in this dataset. """ def __init__( @@ -91,6 +92,7 @@ def __init__( atomic_types: List[int], targets: Dict[str, TargetInfo], extra_data: Optional[Dict[str, TargetInfo]] = None, + hooks: Optional[dict[str, Any]] = None, ): # verify that `length_unit` and `atomic_types` are valid for metatomic _ = ModelCapabilities( @@ -105,6 +107,7 @@ def __init__( self.extra_data: Dict[str, TargetInfo] = ( extra_data if extra_data is not None else {} ) + self.hooks: dict[str, Any] = hooks if hooks is not None else {} @property def atomic_types(self) -> List[int]: @@ -178,6 +181,7 @@ def copy(self) -> "DatasetInfo": atomic_types=self.atomic_types.copy(), targets=self.targets.copy(), extra_data=self.extra_data.copy(), + hooks=self.hooks.copy(), ) @torch.jit.unused @@ -219,6 +223,16 @@ def update(self, other: "DatasetInfo") -> None: ) self.extra_data.update(other.extra_data) + intersecting_hooks_keys = self.hooks.keys() & other.hooks.keys() + for key in intersecting_hooks_keys: + if self.hooks[key] != other.hooks[key]: + raise ValueError( + "Can't update DatasetInfo with different hook information" + f"for key '{key}'." + ) + + self.hooks.update(other.hooks) + def union(self, other: "DatasetInfo") -> "DatasetInfo": """ Return the union of this instance with ``other``. @@ -243,6 +257,7 @@ def __setstate__(self, state: Dict[str, Any]) -> None: self._atomic_types = state["_atomic_types"] self.targets = state["targets"] self.extra_data = state.get("extra_data", {}) + self.hooks = state.get("hooks", {}) def get_stats(dataset: Union[Dataset, Subset], dataset_info: DatasetInfo) -> str: diff --git a/src/metatrain/utils/hooks/__init__.py b/src/metatrain/utils/hooks/__init__.py new file mode 100644 index 0000000000..fcf042a068 --- /dev/null +++ b/src/metatrain/utils/hooks/__init__.py @@ -0,0 +1,6 @@ +from .helpers import restart_hooks, setup_hooks + + +__all__ = ["setup_hooks", "restart_hooks"] + +PostHooksHypers = dict[str, dict | str] diff --git a/src/metatrain/utils/hooks/abc.py b/src/metatrain/utils/hooks/abc.py new file mode 100644 index 0000000000..83f49624bd --- /dev/null +++ b/src/metatrain/utils/hooks/abc.py @@ -0,0 +1,103 @@ +from abc import ABCMeta, abstractmethod +from typing import Generic, Optional, TypeVar + +import torch +from metatensor.torch import Labels, TensorMap +from metatomic.torch import ModelOutput, System + +from metatrain.utils.data import DatasetInfo, TargetInfo + + +HypersType = TypeVar("HypersType") + + +class HookInterface(torch.nn.Module, Generic[HypersType], metaclass=ABCMeta): + """ + Abstract base class for a hook in metatrain. + + All hooks in metatrain must be implemented as sub-class of this class, + and implement the corresponding methods. + + :param hypers: A dictionary with the hook's hyper-parameters. + :param dataset_info: Information containing details about the dataset, such as + target quantities and atomic types. + """ + + __checkpoint_version__: int + """The current version of the model's checkpoint. + + This is used to upgrade checkpoints produced with earlier versions of the code. + See :ref:`ckpt_version` for more information.""" + + def __init__( + self, + hypers: HypersType, + dataset_info: DatasetInfo, + ) -> None: + """""" + super().__init__() + + required_attributes = [ + "__checkpoint_version__", + ] + for attribute in required_attributes: + if not hasattr(self.__class__, attribute): + raise TypeError( + f"missing '{attribute}' class attribute for " + f"'{self.__class__.__module__}.{self.__class__.__name__}'" + ) + + self.hypers = hypers + """The hook hypers passed at initialization""" + + self.dataset_info = dataset_info + """The dataset info passed at initialization""" + + @abstractmethod + def forward( + self, + systems: list[System], + outputs: dict[str, ModelOutput], + inputs: dict[str, TensorMap], + selected_atoms: Optional[Labels] = None, + ) -> dict[str, TensorMap]: + """ + Execute the model for the given ``systems``, computing the requested + ``outputs``, and using the inputs + + :param systems: List of systems to evaluate the model on. + :param outputs: Dictionary of outputs that the model should compute. + :param inputs: Dictionary of input tensors for the model. + :param selected_atoms: Optional ``Labels`` specifying a subset of atoms to + compute the outputs for. If ``None``, the outputs are computed for all + atoms in each system. + + :return: A dictionary mapping each requested output name to the corresponding + ``TensorMap`` containing the computed values. + """ + + @abstractmethod + def requested_target_infos(self) -> dict[str, TargetInfo]: + """ + Returns the inputs that the hook requires, with their layout. + + :return: The requested inputs. Each key is the name of the requested input, + and the value is a :py:class:`TargetInfo` object describing the layout of + the requested input. + """ + + @abstractmethod + def requested_inputs(self) -> dict[str, ModelOutput]: + """ + Returns the list of requested inputs for the hook. + + :return: A list of requested input names. + """ + + @abstractmethod + def supported_outputs(self) -> dict[str, ModelOutput]: + """ + Get the outputs currently supported by this hook. + + :return: A dictionary of the supported outputs by this hook. + """ diff --git a/src/metatrain/utils/hooks/global_multipoles/__init__.py b/src/metatrain/utils/hooks/global_multipoles/__init__.py new file mode 100644 index 0000000000..1780991309 --- /dev/null +++ b/src/metatrain/utils/hooks/global_multipoles/__init__.py @@ -0,0 +1,4 @@ +from .hook import GlobalMultipole + + +__hook__ = GlobalMultipole diff --git a/src/metatrain/utils/hooks/global_multipoles/documentation.py b/src/metatrain/utils/hooks/global_multipoles/documentation.py new file mode 100644 index 0000000000..c7be0fcbc2 --- /dev/null +++ b/src/metatrain/utils/hooks/global_multipoles/documentation.py @@ -0,0 +1,102 @@ +r""" +Global multipoles +================= + +Computes a global multipole from local predictions. + +Predicting the multipole of a system by simply summing local contributions misses +important contributions from the relative positions of the atoms in the system. As an +example, one can think of predicting the dipole of a system with two atoms of opposite +charge. If the two atoms don't see each other given the cutoff of the model, each atom +will predict a local dipole of zero, and the sum of the two local dipoles will also be +zero. However, it is obvious that this system has a non-zero dipole. + +In general, predicting a global multipole of order :math:`\ell` from a simple sum of +local contributions of that :math:`\ell` is not valid whenever the local regions have a +non-zero multipole of order :math:`< \ell`. For example, in the case of the dipole, +whenever the local regions have a net charge (non-zero monopole), the global dipole +can't be computed as the sum of local dipoles. Instead, one needs to add the +contributions from all the lower order multipoles, taking into account their origin +(i.e. the position of the atom). For the case of the dipole, the global dipole can be +computed as: + +.. math:: + \mathbf{p} = \sum_i \mathbf{p}_i + \sum_i q_i \mathbf{r}_i + +This hook implements this concept by taking care of an output that is global and of +order :math:`\ell`, asking for local contributions of all orders :math:`\leq \ell` and +computing the global multipole accounting for the positions of the atoms. + +Origin of the positions +----------------------- + +The expression above is only independent of the choice of origin when all multipoles of +order lower than :math:`\ell` vanish. For the dipole that means :math:`\sum_i q_i = 0`, +which is not guaranteed when the charges are learned: translating a system by +:math:`\mathbf{a}` changes the prediction by :math:`\left(\sum_i q_i\right)\mathbf{a}`. + +To make the prediction invariant under a rigid translation, the positions of each system +are by default referred to its centre of nuclear charge, + +.. math:: + \mathbf{R} = \frac{\sum_i Z_i \mathbf{r}_i}{\sum_i Z_i} + +before the multipole is assembled. Because the origin is defined by the system itself it +translates with it, so the dependence cancels exactly, whatever the value of +:math:`\sum_i q_i`. The centre of nuclear charge is a convention used by some electronic +structure codes, but not all. + +This is controlled by the ``origin`` hyperparameter: set it to ``"absolute"`` to keep +the positions as they are stored, and recover the origin-dependent behaviour described +above. + +.. note:: + + This hook only makes sense for non-periodic systems: :math:`\sum_i q_i \mathbf{r}_i` + depends on how the atoms are wrapped into the cell, and under periodic boundary + conditions the dipole is only defined modulo a quantum. This is not handled in the + current implementation. +""" + +from typing import Literal, Optional, TypedDict + + +class Hypers(TypedDict): + """ + Hyperparameters for the global multipole hook. + """ + + inputs: Optional[str | list[str]] = None + """ + Name or names for the inputs that this hook will request. + + If ``None``, the hook will request an input named + ``mtt::aux::local_multipoles::{output_name.replace('mtt::', '')}`` + for each output name. + """ + + outputs: Optional[str | list[str]] = None + """ + Name or names for the outputs that this hook must produce. + + These targets must be spherical and global, i.e. with sample + kind ``"system"``. + """ + + origin: Literal["center_of_charge", "absolute"] = "center_of_charge" + """ + Origin (per-system) that the atomic positions are referred to when assembling the + multipole. + + - ``"center_of_charge"`` (default): subtract each system's centre of nuclear charge, + :math:`\\mathbf{R} = \\sum_i Z_i \\mathbf{r}_i / \\sum_i Z_i`, from its positions. + Because the origin is defined by the system itself it translates with it, so the + predicted multipole is invariant under a rigid translation whatever the value of + :math:`\\sum_i q_i`. This is a convention used by some electronic structure codes, + but not all. + - ``"absolute"``: use the positions as they are stored, without shifting the origin. + The predicted multipole is then **not** translationally invariant unless the local + charges happen to sum to zero: translating a system by :math:`\\mathbf{a}` changes + it by :math:`\\left(\\sum_i q_i\\right)\\mathbf{a}`. Use this only when the + absolute frame is meaningful, or to reproduce earlier behaviour. + """ diff --git a/src/metatrain/utils/hooks/global_multipoles/hook.py b/src/metatrain/utils/hooks/global_multipoles/hook.py new file mode 100644 index 0000000000..2d0887879c --- /dev/null +++ b/src/metatrain/utils/hooks/global_multipoles/hook.py @@ -0,0 +1,252 @@ +from typing import Optional + +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap +from metatomic.torch import ModelOutput, System + +from metatrain.utils.data import DatasetInfo, TargetInfo +from metatrain.utils.data.target_info import get_generic_target_info +from metatrain.utils.sum_over_atoms import sum_over_atoms + +from ..abc import HookInterface +from .documentation import Hypers + + +class GlobalMultipole(HookInterface[Hypers]): + """ + Computes a global multipole from local predictions. + + :param hypers: A dictionary with the hook's hyper-parameters. + :param dataset_info: Information containing details about the dataset, such as + target quantities and atomic types. + """ + + __checkpoint_version__ = 1 + + def __init__(self, hypers: Hypers, dataset_info: DatasetInfo): + super().__init__(hypers, dataset_info) + + self.hypers = hypers + + # Origin the positions are referred to. Stored as a plain string so + # that the forward stays TorchScript-friendly. + origin = hypers.get("origin", "center_of_charge") + if origin not in ("center_of_charge", "absolute"): + raise ValueError( + f"Invalid 'origin' for the global multipoles hook: {origin!r}. " + f"Expected either 'center_of_charge' or 'absolute'." + ) + self.origin: str = origin + + # Get the information about the output targets from the dataset info + out_names = hypers["outputs"] + if isinstance(out_names, str): + out_names = [out_names] + if out_names is None: + raise ValueError("GlobalMultipole hook requires at least one output target") + self.out_targets = {k: dataset_info.targets[k] for k in out_names} + + # Get the information for the input targets. + input_names = hypers.get("inputs") + if isinstance(input_names, str): + input_names = [input_names] + elif input_names is None: + input_names = [ + f"mtt::aux::local_multipoles::{out_name.replace('mtt::', '')}" + for out_name in out_names + ] + + if len(input_names) != len(out_names): + raise ValueError( + f"Global multipoles hook expects the same number of input and output " + f"targets, but got {len(input_names)} inputs and " + f"{len(out_names)} outputs" + ) + + # Build the target infos that the hook will request. + self._input_target_infos = {} + for in_name, (out_name, out_target) in zip( + input_names, self.out_targets.items(), strict=True + ): + if out_target.sample_kind != "system": + raise ValueError( + f"Global multipoles hook only supports system-level outputs, " + f"but {out_name} has sample kind " + f"{out_target.sample_kind}" + ) + if not out_target.is_spherical: + raise ValueError( + f"Global multipoles hook only supports spherical outputs, " + f"but {out_name} has components " + f"{out_target.layout.components}" + ) + + max_degree = out_target.layout.keys["o3_lambda"].max().item() + if max_degree > 1: + raise ValueError( + f"Global multipoles hook only supports multipoles up " + f"to l=1 for now, but {out_name} has max degree " + f"{max_degree}" + ) + + # Get the input that we will request from the model, which are + # the local multipoles up to the degree of the output multipole. + self._input_target_infos[in_name] = get_generic_target_info( + in_name, + { + "quantity": "", + "unit": "", + "type": { + "spherical": { + "irreps": [ + {"o3_lambda": i, "o3_sigma": 1} + for i in range(max_degree + 1) + ] + } + }, + "num_subtargets": 1, + "sample_kind": "atom", + }, + ) + out_target.layout.set_info("max_degree", str(max_degree)) + out_target.layout.set_info("input_name", in_name) + + def requested_target_infos(self) -> dict[str, TargetInfo]: + """ + Returns the list of requested target infos for the hook. + + :return: A list of requested target names. + """ + return self._input_target_infos + + def requested_inputs(self) -> dict[str, ModelOutput]: + """ + Returns the list of requested inputs for the hook. + + :return: A list of requested input names. + """ + return { + in_name: ModelOutput( + quantity=target_info.quantity, + unit=target_info.unit, + sample_kind="atom", + ) + for in_name, target_info in self._input_target_infos.items() + } + + def supported_outputs(self) -> dict[str, ModelOutput]: + """ + Returns the supported outputs for the hook. + + :return: A list of supported output names. + """ + return { + out_name: ModelOutput( + quantity=target.quantity, + unit=target.unit, + sample_kind="system", + ) + for out_name, target in self.out_targets.items() + } + + def forward( + self, + systems: list[System], + outputs: dict[str, ModelOutput], + inputs: dict[str, TensorMap], + selected_atoms: Optional[Labels] = None, + ) -> dict[str, TensorMap]: + requested_outs: list[str] = [] + for out_name in self.out_targets: + if out_name in outputs: + requested_outs.append(out_name) + if not requested_outs: + # Quick exit if no output of this hook is requested + return {} + + device = systems[0].positions.device + + positions = torch.cat([system.positions for system in systems], dim=0) + + if self.origin == "center_of_charge": + # Enforce origin independence by referring the positions of each + # system to its own centre of nuclear charge. The per-system sums + # are done with a scatter over the concatenated batch; note that + # the origin is computed per system and never over the whole batch, + # which would make the prediction depend on how systems are + # batched together. + nuclear_charges = torch.cat( + [system.types for system in systems], dim=0 + ).to(positions.dtype) + sizes = torch.tensor( + [len(system) for system in systems], device=device, dtype=torch.long + ) + system_indices = torch.repeat_interleave( + torch.arange(len(systems), device=device), sizes + ) + totals = torch.zeros( + len(systems), dtype=positions.dtype, device=device + ).index_add_(0, system_indices, nuclear_charges) + origins = torch.zeros( + (len(systems), 3), dtype=positions.dtype, device=device + ).index_add_(0, system_indices, nuclear_charges.unsqueeze(1) * positions) + origins = origins / totals.unsqueeze(1) + + positions = positions - origins[system_indices] + + # Reorder the axes to match the spherical harmonics convention + # (x, y, z) -> (y, z, x) + positions = positions[:, [1, 2, 0]] + if selected_atoms is not None: + # Here we should filter out the positions only for the selected atoms + raise NotImplementedError( + "Global multipoles hook does not support selected atoms yet." + ) + + return_dict: dict[str, TensorMap] = {} + for out_name in requested_outs: + out_target = self.out_targets[out_name] + + layout = out_target.layout.to(device) + info = layout.info() + max_degree = int(info["max_degree"]) + + # Get the local predictions for each degree in the multipole expansion + input_tmap = inputs[info["input_name"]] + local_values = [ + input_tmap.block(dict(o3_lambda=ell, o3_sigma=1)).values + for ell in range(max_degree + 1) + ] + + # Compute the local contributions to the global multipole for each degree + local_contribs = [] + for ell in range(max_degree + 1): + if ell == 0: + local_contribs.append(local_values[ell]) + elif ell == 1: + local_contribs.append( + torch.einsum( + "sp, sx -> sxp", + local_contribs[ell - 1].squeeze(1), + positions, + ) + + local_values[ell] + ) + + # Build a tensor map with the requested multipole degrees. + local_tmap = TensorMap( + keys=layout.keys, + blocks=[ + TensorBlock( + values=local_contribs[key["o3_lambda"]], + samples=input_tmap.block(0).samples, + components=layout_block.components, + properties=layout_block.properties, + ) + for key, layout_block in layout.items() + ], + ) + + return_dict[out_name] = sum_over_atoms(local_tmap) + + return return_dict diff --git a/src/metatrain/utils/hooks/global_multipoles/tests/__init__.py b/src/metatrain/utils/hooks/global_multipoles/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/metatrain/utils/hooks/global_multipoles/tests/test_basic.py b/src/metatrain/utils/hooks/global_multipoles/tests/test_basic.py new file mode 100644 index 0000000000..66cff968ef --- /dev/null +++ b/src/metatrain/utils/hooks/global_multipoles/tests/test_basic.py @@ -0,0 +1,41 @@ +import pytest + +from metatrain.utils.hooks.testing.hook import HookTests +from metatrain.utils.hooks.testing.inputs import InputTests +from metatrain.utils.hooks.testing.output import OutputTests +from metatrain.utils.hooks.testing.torchscript import TorchscriptTests + + +class GlobalMultipolesTests(HookTests): + hook = "global_multipoles" + + @pytest.fixture + def hypers(self): + return {"outputs": "mtt::global_dipole"} + + def build_hypers(self, input, output): + if input is None: + output_name = output["global_dipole"] + input_name = ( + f"mtt::aux::local_multipoles::{output_name.replace('mtt::', '')}" + ) + return {"outputs": output_name}, [input_name], [output_name] + elif output is None: + ... + else: + input_name = input["local_multipole"] + output_name = output["global_dipole"] + return ( + {"inputs": input_name, "outputs": output_name}, + [input_name], + [output_name], + ) + + +class TestTorchscript(TorchscriptTests, GlobalMultipolesTests): ... + + +class TestInputs(InputTests, GlobalMultipolesTests): ... + + +class TestOutputs(OutputTests, GlobalMultipolesTests): ... diff --git a/src/metatrain/utils/hooks/global_multipoles/tests/test_translation_invariance.py b/src/metatrain/utils/hooks/global_multipoles/tests/test_translation_invariance.py new file mode 100644 index 0000000000..a80fc535eb --- /dev/null +++ b/src/metatrain/utils/hooks/global_multipoles/tests/test_translation_invariance.py @@ -0,0 +1,345 @@ +"""Origin handling of the global multipoles hook. + +The hook builds the global dipole as :math:`\\sum_i q_i \\mathbf{r}_i + \\sum_i +\\mathbf{p}_i`, which is origin-dependent whenever the charges do not sum to zero. With +``origin="center_of_charge"`` the positions are referred to each system's centre of +nuclear charge, which removes that dependence; with ``origin="absolute"`` they are not, +and the dependence is kept. + +These tests feed the hook random charges and local dipoles directly, standing in for the +per-atom output of a translationally invariant architecture: such a model returns the +same per-atom values for a system and for a rigidly translated copy of it, so the same +inputs are reused for both and only the positions change. +""" + +import pytest +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap +from metatomic.torch import ModelOutput, System + +from metatrain.utils.data import DatasetInfo +from metatrain.utils.data.target_info import get_generic_target_info +from metatrain.utils.hooks.helpers import load_hook + + +OUTPUT_NAME = "mtt::global_dipole" +INPUT_NAME = "mtt::aux::local_multipoles::global_dipole" + +ATOMIC_TYPES = [1, 6, 7, 8] +SYSTEM_SIZES = [4, 11, 7] + + +def _dataset_info() -> DatasetInfo: + """A dataset info holding a single global dipole target.""" + global_dipole = get_generic_target_info( + OUTPUT_NAME, + dict( + sample_kind="system", + unit="", + quantity="", + num_subtargets=1, + type=dict(spherical=dict(irreps=[{"o3_lambda": 1, "o3_sigma": 1}])), + ), + ) + return DatasetInfo( + length_unit="angstrom", + atomic_types=ATOMIC_TYPES, + targets={OUTPUT_NAME: global_dipole}, + ) + + +def _random_systems(generator: torch.Generator) -> list[System]: + """Random systems, deliberately far from the origin of their frame.""" + systems = [] + for size in SYSTEM_SIZES: + types = torch.tensor(ATOMIC_TYPES, dtype=torch.int32)[ + torch.randint(len(ATOMIC_TYPES), (size,), generator=generator) + ] + positions = ( + torch.rand((size, 3), generator=generator, dtype=torch.float64) * 4.0 + + torch.tensor([7.0, -3.0, 11.0], dtype=torch.float64) + ) + systems.append( + System( + types=types, + positions=positions, + cell=torch.zeros((3, 3), dtype=torch.float64), + pbc=torch.zeros(3, dtype=torch.bool), + ) + ) + return systems + + +def _random_local_multipoles( + systems: list[System], generator: torch.Generator +) -> dict[str, TensorMap]: + """Random per-atom charges and local dipoles. + + These stand in for what a translationally invariant architecture would + predict: they depend on the systems only through their internal geometry, + so a rigid translation leaves them unchanged. + """ + samples = Labels( + names=["system", "atom"], + values=torch.tensor( + [ + [system_index, atom_index] + for system_index, system in enumerate(systems) + for atom_index in range(len(system)) + ] + ), + ) + total_atoms = sum(len(system) for system in systems) + + # Charges deliberately not centred on zero, so that the sum over each + # system is far from neutral and the origin genuinely matters. + charges = ( + torch.rand((total_atoms, 1, 1), generator=generator, dtype=torch.float64) + 0.5 + ) + local_dipoles = torch.rand( + (total_atoms, 3, 1), generator=generator, dtype=torch.float64 + ) + + keys = Labels( + names=["o3_lambda", "o3_sigma"], values=torch.tensor([[0, 1], [1, 1]]) + ) + blocks = [ + TensorBlock( + values=charges, + samples=samples, + components=[Labels(names=["o3_mu"], values=torch.tensor([[0]]))], + properties=Labels(names=["properties"], values=torch.tensor([[0]])), + ), + TensorBlock( + values=local_dipoles, + samples=samples, + components=[ + Labels(names=["o3_mu"], values=torch.tensor([[-1], [0], [1]])) + ], + properties=Labels(names=["properties"], values=torch.tensor([[0]])), + ), + ] + return {INPUT_NAME: TensorMap(keys=keys, blocks=blocks)} + + +def _translate(systems: list[System], shift: torch.Tensor) -> list[System]: + """Rigidly translate every system by the same vector.""" + return [ + System( + types=system.types, + positions=system.positions + shift, + cell=system.cell, + pbc=system.pbc, + ) + for system in systems + ] + + +def _make_hook(origin: str = "center_of_charge"): + """Build the hook with a given origin convention.""" + hypers = {"outputs": OUTPUT_NAME} + if origin is not None: + hypers["origin"] = origin + return load_hook("global_multipoles")(hypers, _dataset_info()) + + +@pytest.fixture +def hook(): + return _make_hook("center_of_charge") + + +@pytest.fixture +def absolute_hook(): + return _make_hook("absolute") + + +@pytest.fixture +def outputs(): + return {OUTPUT_NAME: ModelOutput(quantity="", unit="", sample_kind="system")} + + +@pytest.fixture +def generator(): + return torch.Generator().manual_seed(0) + + +def test_charges_are_not_neutral(generator): + """The test is only meaningful if the origin actually matters. + + A translation changes the raw sum by ``sum_i q_i * a``, so if the random + charges happened to sum to zero the invariance check below would pass even + without an origin being subtracted. + """ + systems = _random_systems(generator) + inputs = _random_local_multipoles(systems, generator) + charges = inputs[INPUT_NAME].block(0).values.reshape(-1) + + start = 0 + for system in systems: + total = charges[start : start + len(system)].sum() + assert abs(float(total)) > 1.0 + start += len(system) + + +@pytest.mark.parametrize( + "shift", + [ + [100.0, 0.0, 0.0], + [-13.0, 7.0, -21.0], + [1e4, 1e4, 1e4], + ], +) +def test_translation_invariance(hook, outputs, generator, shift): + """Translating every system must leave the global dipole unchanged.""" + systems = _random_systems(generator) + inputs = _random_local_multipoles(systems, generator) + translated = _translate(systems, torch.tensor(shift, dtype=torch.float64)) + + original = hook(systems, outputs, inputs)[OUTPUT_NAME].block(0).values + moved = hook(translated, outputs, inputs)[OUTPUT_NAME].block(0).values + + torch.testing.assert_close(original, moved, rtol=0.0, atol=1e-9) + + +@pytest.mark.parametrize("origin", ["center_of_charge", "absolute"]) +def test_independent_of_batching(outputs, generator, origin): + """The origin is per system, so batching must not change the result.""" + hook = _make_hook(origin) + systems = _random_systems(generator) + inputs = _random_local_multipoles(systems, generator) + + batched = hook(systems, outputs, inputs)[OUTPUT_NAME].block(0).values + + one_at_a_time = [] + start = 0 + for index, system in enumerate(systems): + block = inputs[INPUT_NAME] + stop = start + len(system) + single = { + INPUT_NAME: TensorMap( + keys=block.keys, + blocks=[ + TensorBlock( + values=b.values[start:stop], + samples=Labels( + names=["system", "atom"], + values=torch.tensor( + [[0, atom] for atom in range(len(system))] + ), + ), + components=b.components, + properties=b.properties, + ) + for b in block.blocks() + ], + ) + } + one_at_a_time.append( + hook([system], outputs, single)[OUTPUT_NAME].block(0).values + ) + start = stop + + torch.testing.assert_close( + batched, torch.cat(one_at_a_time, dim=0), rtol=0.0, atol=1e-9 + ) + + +def test_matches_centre_of_nuclear_charge(hook, outputs, generator): + """The global dipole is built about each system's centre of nuclear charge.""" + systems = _random_systems(generator) + inputs = _random_local_multipoles(systems, generator) + + predicted = hook(systems, outputs, inputs)[OUTPUT_NAME].block(0).values + + block = inputs[INPUT_NAME] + charges = block.block(0).values.reshape(-1) + local_dipoles = block.block(1).values.squeeze(-1) + + expected, start = [], 0 + for system in systems: + stop = start + len(system) + positions = system.positions + nuclear_charges = system.types.to(positions.dtype) + origin = (nuclear_charges.unsqueeze(1) * positions).sum(0) / nuclear_charges.sum() + # (x, y, z) -> (y, z, x), the spherical harmonics convention + shifted = (positions - origin)[:, [1, 2, 0]] + expected.append( + (charges[start:stop].unsqueeze(1) * shifted).sum(0) + + local_dipoles[start:stop].sum(0) + ) + start = stop + + torch.testing.assert_close( + predicted.squeeze(-1), torch.stack(expected), rtol=1e-10, atol=1e-10 + ) + + +def test_center_of_charge_is_the_default(outputs, generator): + """Omitting the hyperparameter must give the origin-independent behaviour.""" + systems = _random_systems(generator) + inputs = _random_local_multipoles(systems, generator) + + default = _make_hook(None)(systems, outputs, inputs)[OUTPUT_NAME].block(0).values + explicit = ( + _make_hook("center_of_charge")(systems, outputs, inputs)[OUTPUT_NAME] + .block(0) + .values + ) + + torch.testing.assert_close(default, explicit, rtol=0.0, atol=0.0) + + +def test_absolute_matches_raw_positions(absolute_hook, outputs, generator): + """With ``origin="absolute"`` the positions are used exactly as stored.""" + systems = _random_systems(generator) + inputs = _random_local_multipoles(systems, generator) + + predicted = absolute_hook(systems, outputs, inputs)[OUTPUT_NAME].block(0).values + + block = inputs[INPUT_NAME] + charges = block.block(0).values.reshape(-1) + local_dipoles = block.block(1).values.squeeze(-1) + + expected, start = [], 0 + for system in systems: + stop = start + len(system) + # (x, y, z) -> (y, z, x), the spherical harmonics convention + positions = system.positions[:, [1, 2, 0]] + expected.append( + (charges[start:stop].unsqueeze(1) * positions).sum(0) + + local_dipoles[start:stop].sum(0) + ) + start = stop + + torch.testing.assert_close( + predicted.squeeze(-1), torch.stack(expected), rtol=1e-10, atol=1e-10 + ) + + +def test_absolute_shifts_by_the_total_charge(absolute_hook, outputs, generator): + """Without an origin the prediction moves by exactly ``sum_i q_i * a``. + + This pins down what ``origin="absolute"`` gives up: the residual is not + merely non-zero, it is the total charge times the translation. + """ + shift = torch.tensor([-13.0, 7.0, -21.0], dtype=torch.float64) + systems = _random_systems(generator) + inputs = _random_local_multipoles(systems, generator) + translated = _translate(systems, shift) + + original = absolute_hook(systems, outputs, inputs)[OUTPUT_NAME].block(0).values + moved = absolute_hook(translated, outputs, inputs)[OUTPUT_NAME].block(0).values + + charges = inputs[INPUT_NAME].block(0).values.reshape(-1) + totals, start = [], 0 + for system in systems: + stop = start + len(system) + totals.append(charges[start:stop].sum()) + start = stop + + # (x, y, z) -> (y, z, x), the spherical harmonics convention + expected = torch.stack(totals).unsqueeze(1) * shift[[1, 2, 0]].unsqueeze(0) + + torch.testing.assert_close( + (moved - original).squeeze(-1), expected, rtol=1e-9, atol=1e-9 + ) diff --git a/src/metatrain/utils/hooks/helpers.py b/src/metatrain/utils/hooks/helpers.py new file mode 100644 index 0000000000..205a92c619 --- /dev/null +++ b/src/metatrain/utils/hooks/helpers.py @@ -0,0 +1,274 @@ +import importlib +import importlib.util +import sys +from pathlib import Path +from types import ModuleType +from typing import Optional, cast + +import torch +from omegaconf import OmegaConf + +from metatrain.utils.data import DatasetInfo, TargetInfo +from metatrain.utils.hypers import init_with_defaults + +from .abc import HookInterface + + +def find_all_hooks() -> list[str]: + """ + Returns a list of all hooks in the metatrain.utils.hooks package. + + :return: A list of hook names. + """ + + hooks_root = Path(__file__).parent + + # Get all the directories in the hooks package + return [ + p.name + for p in hooks_root.iterdir() + if (p.is_dir() and (p / "__init__.py").exists() and not p.name == "testing") + ] + + +def load_hook(hook_name: str) -> type[HookInterface]: + """ + Loads a hook from the metatrain.utils.hooks package. + + :param hook_name: The name of the hook to load. + :return: The hook class. + """ + if hook_name not in find_all_hooks(): + raise ValueError(f"Unknown hook: {hook_name}") + + # Import the hook module + hook_module = importlib.import_module(f"metatrain.utils.hooks.{hook_name}") + return hook_module.__hook__ + + +def preload_documentation_module(name: str) -> ModuleType: + """This preloads the documentation module for a given hook. + + It imports the `documentation.py` file in an isolated manner and + adds it to `sys.modules`. + + The reason one might do this is because the documentation module + does not have extra dependencies, so importing it separately is + always possible, while if we didn't preload it, importing the + documentation would trigger the hook's `__init__.py` + which might have extra dependencies that are not installed. + + Doing this preloading is useful especially in the context of + generating the documentation, where we want to be able to + document hooks even if their dependencies are not + installed. + + :param name: Name of the hook + :return: The documentation module for the hook. + """ + file_path = Path(__file__).parent / name / "documentation.py" + if not file_path.exists(): + raise FileNotFoundError( + f"The documentation.py file for hook '{name}' was not found. " + "Cannot load the hook's hyperparameter specification." + ) + spec = importlib.util.spec_from_file_location( + f"metatrain.utils.hooks.{name}.documentation", file_path + ) + assert spec is not None # for mypy + documentation = importlib.util.module_from_spec(spec) + assert spec.loader is not None # for mypy + spec.loader.exec_module(documentation) + sys.modules[f"metatrain.utils.hooks.{name}.documentation"] = documentation + return documentation + + +def get_hypers_class(name: str) -> type: + """ + Returns the hypers classes for a given hook. + + :param name: Name of the hook. + :return: The hypers class for the hook. + """ + documentation = preload_documentation_module(name) + return documentation.Hypers + + +def get_default_hypers(name: str, base_precision: Optional[int] = None) -> dict: + """Returns the default hook hyperparameters. + + When *base_precision* is ``None`` (the default), the returned dict may + contain ``"${base_precision}"`` interpolation strings in precision fields. + These resolve automatically after ``OmegaConf.merge`` with a config that + provides ``base_precision`` at the root level. + + When *base_precision* is an integer (16, 32, or 64), all interpolations are + resolved before returning so callers get a plain dict of concrete values. + + :param name: Name of the hook + :param base_precision: If given, resolve ``${base_precision}`` interpolations + to this value before returning. + :return: Default hyperparameters of the hook + """ + + hypers_class = get_hypers_class(name) + + defaults = init_with_defaults(hypers_class) + + if base_precision is not None: + cfg = OmegaConf.create({"base_precision": base_precision, **defaults}) + container = OmegaConf.to_container(cfg, resolve=True) + container.pop("base_precision") + return container + + return defaults + + +def write_hypers_yaml( + name: str, output_path: Path | str, include_name: bool = False +) -> None: + """Write YAML file with defaults for a given hook. + + Given a hook name, this function imports the corresponding + module, finds out what the hyperparameters are for the hook + and its trainer, and generates a YAML file with the default + hyperparameters. + + :param name: The hook to generate the files for. + :param output_path: The path to write the YAML file to. + :param include_name: If True, the YAML file will contain a top-level + key with the hook name, mimicking the input that needs to be + provided to use the hook. + """ + # Create the dictionary with all default hyperparameters + yaml_defaults = get_default_hypers(name) + + if include_name: + yaml_defaults = {name: yaml_defaults} + + conf = OmegaConf.create(yaml_defaults) + # And write them to a YAML file + with open(output_path, "w") as f: + OmegaConf.save(config=conf, f=f) + + +class UnavailableOutputError(Exception): + """Should be used when a hook is asked to produce an output + but the hook is not able to find the layout of such output. + + This is because the function that sets up hooks will go through + the other hooks and then retry this one. Maybe some other hook + has requested an input that matches this output's name and then + the layout will now be available. + """ + + +def setup_hooks( + dataset_info: DatasetInfo, +) -> tuple[ + list[torch.nn.Module], + dict[str, TargetInfo], +]: + """ + Setup the post-processing hooks to apply to the outputs of the model. + + :param dataset_info: The dataset information, which contains target + information and also the hooks specification. + :return: A tuple containing the list of post-processing hooks, and a + dictionary with the outputs that the model should produce before the + hooks are applied. + """ + model_outputs = dataset_info.targets.copy() + all_hook_outputs = [] + post_hooks = [] + + def _set_up( + post_hooks_hypers: dict, dataset_info: DatasetInfo + ) -> tuple[dict, list]: + unavailable_out_hooks = {} + unavailable_errors = [] + for hook_name, hook_hypers in post_hooks_hypers.items(): + hook_class = load_hook(hook_name) + + if isinstance(hook_hypers, str): + san_hook_hypers: dict[str, dict[str, str] | str] = { + "outputs": hook_hypers, + } + else: + san_hook_hypers = cast(dict[str, dict[str, str] | str], hook_hypers) + + # Get hook and add it to the list of hooks. + try: + hook = hook_class(san_hook_hypers, dataset_info) + except UnavailableOutputError as e: + unavailable_out_hooks[hook_name] = hook_hypers + unavailable_errors.append(e) + continue + post_hooks.append(hook) + + # The model will not need to produce the outputs that this hook + # takes care of, so remove them from the list of model outputs. + hook_outputs = hook.supported_outputs() + for output_name in hook_outputs: + if output_name in all_hook_outputs: + raise ValueError( + f"Output '{output_name}' is already produced by another hook. " + "A given output can only be produced by one hook." + ) + model_outputs.pop(output_name, None) + all_hook_outputs.extend(list(hook_outputs.keys())) + + # Add the inputs that the hook requests to the model outputs, + # unless they are already present somewhere. + requested = hook.requested_target_infos() + extra_data_info = dataset_info.extra_data + for req_name, req_info in requested.items(): + if req_name in model_outputs: + continue + elif req_name in all_hook_outputs: + continue + elif req_name in extra_data_info: + continue + else: + model_outputs[req_name] = req_info + + return unavailable_out_hooks, unavailable_errors + + hooks_to_set_up = dataset_info.hooks.copy() + while True: + new_hooks_to_set_up, errors = _set_up(hooks_to_set_up, dataset_info) + if len(new_hooks_to_set_up) == 0: + break + elif len(new_hooks_to_set_up) == len(hooks_to_set_up): + raise ValueError( + "Some hooks could not be set up because they requested outputs " + "that are not present either in the targets, extra data or " + "in other hooks' inputs. \nThe errors were the following:\n - " + + "\n - ".join(str(e) for e in errors) + ) + # We managed to set up some hooks in this iteration, do another iteration + # to see if we can set up more hooks now that some outputs are available. + hooks_to_set_up = new_hooks_to_set_up + dataset_info.targets.update(model_outputs) + + return post_hooks, model_outputs + + +def restart_hooks( + hooks: list[HookInterface], dataset_info: DatasetInfo +) -> tuple[list[HookInterface], dict[str, TargetInfo]]: + """ + Restart hooks with new dataset information. + + :param hooks: List of existing hooks to restart. + :param dataset_info: New dataset information to use for restarting the hooks. + :return: A tuple containing the list of restarted hooks and a dictionary + with the outputs that the model should produce before the hooks are applied. + + """ + if len(hooks) > 0: + raise ValueError("Restarting from previous hooks is not supported yet.") + + forward_hooks, model_outs = setup_hooks(dataset_info) + + return forward_hooks, model_outs diff --git a/src/metatrain/utils/hooks/identity/__init__.py b/src/metatrain/utils/hooks/identity/__init__.py new file mode 100644 index 0000000000..121b2d7b19 --- /dev/null +++ b/src/metatrain/utils/hooks/identity/__init__.py @@ -0,0 +1,4 @@ +from .hook import Identity + + +__hook__ = Identity diff --git a/src/metatrain/utils/hooks/identity/documentation.py b/src/metatrain/utils/hooks/identity/documentation.py new file mode 100644 index 0000000000..bdfb785277 --- /dev/null +++ b/src/metatrain/utils/hooks/identity/documentation.py @@ -0,0 +1,33 @@ +""" +Identity +======== + +This hook just passes the inputs to the outputs without any modification. + +For someone who wants to understand how hooks work or wants to implement +their own hook, this is the best point to start, as it is the simplest +hook possible. + +It is also very useful for testing and debugging. +""" + +from typing import Optional + +from typing_extensions import TypedDict + + +class Hypers(TypedDict): + """ + Hyperparameters for the identity hook. + """ + + inputs: Optional[str | list[str]] = None + """The names of the inputs to be passed to the outputs. + + If ``None``, they will be set as + ``mtt::aux::identity::{output_name.replace('mtt::', '')}`` + for each output name. + """ + + outputs: Optional[str | list[str]] = None + """The names of the outputs to be produced by the hook.""" diff --git a/src/metatrain/utils/hooks/identity/hook.py b/src/metatrain/utils/hooks/identity/hook.py new file mode 100644 index 0000000000..7f7e36483d --- /dev/null +++ b/src/metatrain/utils/hooks/identity/hook.py @@ -0,0 +1,134 @@ +from typing import Optional + +from metatensor.torch import Labels, TensorMap +from metatomic.torch import ModelOutput, System + +from metatrain.utils.data import DatasetInfo, TargetInfo + +from ..abc import HookInterface +from .documentation import Hypers + + +class Identity(HookInterface[Hypers]): + """ + Passes the inputs to the outputs without any modification. + + :param hypers: A dictionary with the hook's hyper-parameters. + :param dataset_info: Information containing details about the dataset, such as + target quantities and atomic types. + """ + + __checkpoint_version__ = 1 + + def __init__(self, hypers: Hypers, dataset_info: DatasetInfo): + super().__init__(hypers, dataset_info) + + self.hypers = hypers + + # Get the information about the output targets from the dataset info + out_names = hypers["outputs"] + if isinstance(out_names, str): + out_names = [out_names] + if out_names is None: + raise ValueError("Identity hook requires at least one output target") + + # Get the information for the input targets. + input_names = hypers.get("inputs") + if isinstance(input_names, str): + input_names = [input_names] + elif input_names is None: + input_names = [ + f"mtt::aux::identity::{out_name.replace('mtt::', '')}" + for out_name in out_names + ] + + if len(input_names) != len(out_names): + raise ValueError( + f"Identity hook expects the same number of input and output " + f"targets, but got {len(input_names)} inputs and " + f"{len(out_names)} outputs" + ) + + # Build the target infos that the hook will request. + self.out_targets = {} + self._input_target_infos = {} + targets = dataset_info.targets + for in_name, out_name in zip(input_names, out_names, strict=True): + if in_name in targets and out_name in targets: + if targets[in_name] != targets[out_name]: + raise ValueError( + f"Identity hook found both the input '{in_name}' " + f"and the output '{out_name}' in the dataset targets, " + "but they have different layouts. To be able to apply " + "the identity hook, the input and output must have the " + "same layout." + ) + else: + target = targets[in_name] + elif out_name not in targets and in_name not in targets: + raise ValueError( + f"Identity hook expects the output '{out_name}' or the input " + f"'{in_name}' to be present in the dataset targets, but neither " + f"was found." + ) + elif out_name in dataset_info.targets: + target = dataset_info.targets[out_name] + else: + target = dataset_info.targets[in_name] + + self._input_target_infos[in_name] = target + self.out_targets[out_name] = target + + def requested_target_infos(self) -> dict[str, TargetInfo]: + """ + Returns the list of requested target infos for the hook. + + :return: A list of requested target names. + """ + return self._input_target_infos + + def requested_inputs(self) -> dict[str, ModelOutput]: + """ + Returns the list of requested inputs for the hook. + + :return: A list of requested input names. + """ + return { + in_name: ModelOutput( + quantity=target_info.quantity, + unit=target_info.unit, + sample_kind=target_info.sample_kind, + ) + for in_name, target_info in self._input_target_infos.items() + } + + def supported_outputs(self) -> dict[str, ModelOutput]: + """ + Returns the supported outputs for the hook. + + :return: A list of supported output names. + """ + return { + out_name: ModelOutput( + quantity=target_info.quantity, + unit=target_info.unit, + sample_kind=target_info.sample_kind, + ) + for out_name, target_info in self.out_targets.items() + } + + def forward( + self, + systems: list[System], + outputs: dict[str, ModelOutput], + inputs: dict[str, TensorMap], + selected_atoms: Optional[Labels] = None, + ) -> dict[str, TensorMap]: + return_dict: dict[str, TensorMap] = {} + for in_name, out_name in zip( + self._input_target_infos, self.out_targets, strict=True + ): + if out_name in outputs: + return_dict[out_name] = inputs[in_name] + + return return_dict diff --git a/src/metatrain/utils/hooks/identity/tests/__init__.py b/src/metatrain/utils/hooks/identity/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/metatrain/utils/hooks/identity/tests/test_basic.py b/src/metatrain/utils/hooks/identity/tests/test_basic.py new file mode 100644 index 0000000000..7ff0d49732 --- /dev/null +++ b/src/metatrain/utils/hooks/identity/tests/test_basic.py @@ -0,0 +1,41 @@ +import pytest + +from metatrain.utils.hooks.testing.hook import HookTests +from metatrain.utils.hooks.testing.inputs import InputTests +from metatrain.utils.hooks.testing.output import OutputTests +from metatrain.utils.hooks.testing.torchscript import TorchscriptTests + + +class IdentityTests(HookTests): + hook = "identity" + + @pytest.fixture + def hypers(self): + return {"outputs": "mtt::energy"} + + def build_hypers(self, input, output): + if input is None: + output_name = output["energy"] + input_name = f"mtt::aux::identity::{output_name.replace('mtt::', '')}" + return {"outputs": output_name}, [input_name], [output_name] + elif output is None: + input_name = input["energy"] + output_name = f"mtt::aux::identity::{input_name.replace('mtt::', '')}" + return {"inputs": input_name}, [input_name], [output_name] + else: + input_name = input["energy"] + output_name = output["energy"] + return ( + {"inputs": input_name, "outputs": output_name}, + [input_name], + [output_name], + ) + + +class TestTorchscript(TorchscriptTests, IdentityTests): ... + + +class TestInputs(InputTests, IdentityTests): ... + + +class TestOutputs(OutputTests, IdentityTests): ... diff --git a/src/metatrain/utils/hooks/intensive_gap/__init__.py b/src/metatrain/utils/hooks/intensive_gap/__init__.py new file mode 100644 index 0000000000..9e5a69d54d --- /dev/null +++ b/src/metatrain/utils/hooks/intensive_gap/__init__.py @@ -0,0 +1,4 @@ +from .hook import IntensiveGap + + +__hook__ = IntensiveGap diff --git a/src/metatrain/utils/hooks/intensive_gap/documentation.py b/src/metatrain/utils/hooks/intensive_gap/documentation.py new file mode 100644 index 0000000000..5e557ef95a --- /dev/null +++ b/src/metatrain/utils/hooks/intensive_gap/documentation.py @@ -0,0 +1,83 @@ +r""" +Intensive gap +============= + +Does minimum and maximum pooling to get a gap-like intensive property. + +Intensive properties like the band gap can't be computed as a sum of +local contributions due to the simple fact that they do not scale +with the size of the system. This hook implements a simple but powerful +idea that was used in this work :footcite:p:`malosso2026transferable` +to predict gaps from local contributions: + +.. math:: + E_{gap} = \min_i c_i - \max_i v_i + +i.e. it defines the gap as the difference between the maximum and minimum of +two different local contributions :math:`c_i` and :math:`v_i`. The minimum +and maximum can be interpreted as the conduction and valence band edges or +LUMO and HOMO energies, respectively. + +The hook takes care of a global scalar output by requesting two separate +local scalar quantities as inputs and performing the min/max pooling. + +""" + +from typing import Literal, TypedDict + +from typing_extensions import NotRequired + +from metatrain.utils.hypers import init_with_defaults + + +class PoolingHypers(TypedDict): + """Hyperparameters for the per-system pooling step. + + Two pooling types are supported, both controlled by the same ``alpha_max`` + / ``alpha_min`` parameters (sign selects max- vs min-pool, magnitude sets + the sharpness): + + - ``"smoothmax"`` (default): ``E = (1/alpha) log sum_i exp(alpha h_i)``. + Recovers a hard max/min as ``|alpha| -> infinity``. Size-intensive up to a + ``log(N)/|alpha|`` residual. + - ``"softmax"``: ``E = sum_i softmax(alpha * h_i) * h_i``. A self-weighted + softmax pool: the softmax weights are computed from the per-atom *values + themselves*, so the pool attends to the most extreme contributions. + Strictly intensive (softmax weights sum to 1, removing the + ``log(N)/|alpha|`` residual of the smoothmax pool) and recovers a hard + max/min as ``|alpha| -> infinity``. + """ + + type: Literal["smoothmax", "softmax"] = "smoothmax" + """Pooling type. One of ``"smoothmax"`` or ``"softmax"``.""" + + alpha_bottom: float = 20.0 + """Max pooling parameter. ``alpha_bottom > 0`` gives a (smooth/soft) max. + Larger magnitude -> sharper (closer to a hard max). Used by both pooling + types.""" + + alpha_top: float = -20.0 + """Min pooling parameter. ``alpha_top < 0`` gives a (smooth/soft) min. + Larger magnitude -> sharper (closer to a hard min). Used by both pooling + types.""" + + +class HookInputs(TypedDict): + """Inputs for the minmax hook.""" + + bottom: NotRequired[str] + """Name of the target for the bottom of the gap.""" + top: NotRequired[str] + """Name of the target for the top of the gap.""" + + +class Hypers(TypedDict): + """ + Hyperparameters for the global multipole hook. + """ + + pooling: PoolingHypers = init_with_defaults(PoolingHypers) + + inputs: HookInputs | list[HookInputs] = init_with_defaults(HookInputs) + + outputs: str | list[str] diff --git a/src/metatrain/utils/hooks/intensive_gap/hook.py b/src/metatrain/utils/hooks/intensive_gap/hook.py new file mode 100644 index 0000000000..ff8bb8bd49 --- /dev/null +++ b/src/metatrain/utils/hooks/intensive_gap/hook.py @@ -0,0 +1,330 @@ +from typing import Optional + +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap +from metatomic.torch import ModelOutput, System + +from metatrain.utils.data import DatasetInfo, TargetInfo +from metatrain.utils.hypers import init_with_defaults + +from ..abc import HookInterface +from ..helpers import UnavailableOutputError +from .documentation import Hypers, PoolingHypers + + +# ----------------------------------------- +# Functions to do the pooling +# ----------------------------------------- + + +def _scatter_softmax_pool( + values: torch.Tensor, + alpha: float, + system_indices: torch.Tensor, + num_systems: int, +) -> torch.Tensor: + """Per-system self-weighted softmax pool: ``sum_i softmax(alpha * v_i)_i * v_i``. + + Numerically stable: shift ``alpha * v`` by per-system max before exponentiating. + Strictly intensive (softmax weights sum to 1 within each system). The sign of + ``alpha`` selects max- vs min-pool, exactly as in :func:`_scatter_logsumexp`. + + :param values: ``(N,)`` per-atom values. + :param alpha: scalar tensor; sign determines max- vs min-pool. + :param system_indices: ``(N,)`` system index per atom (in ``[0, num_systems)``). + :param num_systems: number of systems ``S`` in the batch. + :return: ``(S,)`` pooled values. + """ + logits = alpha * values # (N,) + neg_inf = torch.full( + (num_systems,), float("-inf"), dtype=values.dtype, device=values.device + ) + sys_max = neg_inf.scatter_reduce( + 0, system_indices, logits, reduce="amax", include_self=True + ) + sys_max = torch.where(torch.isinf(sys_max), torch.zeros_like(sys_max), sys_max) + exps = torch.exp(logits - sys_max[system_indices]) # (N,) + denom = torch.zeros( + num_systems, dtype=values.dtype, device=values.device + ).scatter_add(0, system_indices, exps) + weights = exps / denom[system_indices] # (N,) softmax across each system + weighted = weights * values + pooled = torch.zeros( + num_systems, dtype=values.dtype, device=values.device + ).scatter_add(0, system_indices, weighted) + return pooled + + +def _scatter_logsumexp( + values: torch.Tensor, + alpha: float, + system_indices: torch.Tensor, + num_systems: int, +) -> torch.Tensor: + """Numerically stable per-system ``(1/alpha) * logsumexp(alpha * values)``. + + Works for ``alpha`` of either sign. Implementation: shift by per-system max + of ``alpha * values`` for stability, then scatter-add the exponentials. + + :param values: ``(N,)`` per-atom values. + :param alpha: scalar tensor; sign determines max- vs min-pool. + :param system_indices: ``(N,)`` system index per atom (in ``[0, num_systems)``). + :param num_systems: number of systems ``S`` in the batch. + :return: ``(S,)`` pooled values. + """ + scaled = alpha * values # (N,) + neg_inf = torch.full( + (num_systems,), + float("-inf"), + dtype=values.dtype, + device=values.device, + ) + sys_max = neg_inf.scatter_reduce( + 0, system_indices, scaled, reduce="amax", include_self=True + ) + sys_max = torch.where(torch.isinf(sys_max), torch.zeros_like(sys_max), sys_max) + shifted_exp = torch.exp(scaled - sys_max[system_indices]) + sum_exp = torch.zeros( + num_systems, dtype=values.dtype, device=values.device + ).scatter_add(0, system_indices, shifted_exp) + log_sum_exp = sys_max + torch.log(sum_exp) + return log_sum_exp / alpha + + +# ----------------------------------------- +# The hook itself +# ----------------------------------------- + + +class IntensiveGap(HookInterface[Hypers]): + """Min/max pooling hook for intensive system-level outputs. + + :param hypers: A dictionary with the hook's hyper-parameters. + :param dataset_info: Information containing details about the dataset, such as + target quantities and atomic types. + """ + + __checkpoint_version__ = 1 + + def __init__(self, hypers: Hypers, dataset_info: DatasetInfo): + super().__init__(hypers, dataset_info) + + self.hypers = hypers + + pooling_hypers = hypers.get("pooling", init_with_defaults(PoolingHypers)) + + # Get the information about the output targets from the dataset info + out_names = hypers["outputs"] + if isinstance(out_names, str): + out_names = [out_names] + self.out_targets = {} + for k in out_names: + if k not in dataset_info.targets: + raise UnavailableOutputError( + f"IntensiveGap hook requested output {k} but it is not " + f"available in the dataset info" + ) + self.out_targets[k] = dataset_info.targets[k] + + # Get the information for the input targets. + input_names = hypers.get("inputs") + if isinstance(input_names, dict): + input_names = [input_names] + if input_names is None: + input_names = [{}] * len(out_names) + input_names = [ + { + "bottom": in_names.get( + "bottom", f"mtt::aux::gap_bottom::{out_name.replace('mtt::', '')}" + ), + "top": in_names.get( + "top", f"mtt::aux::gap_top::{out_name.replace('mtt::', '')}" + ), + } + for in_names, out_name in zip(input_names, out_names, strict=True) + ] + + assert len(input_names) == len(out_names), ( + f"IntensiveGap hook expects the same number of input and output " + f"targets, but got {len(input_names)} inputs and {len(out_names)} outputs" + ) + + # Build the target infos that the hook will request. + self._input_target_infos = {} + self._block_shapes = {} + for in_names, (out_name, out_target) in zip( + input_names, self.out_targets.items(), strict=True + ): + if out_target.sample_kind != "system": + raise ValueError( + f"IntensiveGap hook only supports system-level outputs, " + f"but {out_name} has sample kind " + f"{out_target.sample_kind}" + ) + if not out_target.is_scalar: + raise ValueError( + f"IntensiveGap hook only supports scalar outputs, " + f"but {out_name} has components " + f"{out_target.layout.components}" + ) + + layout = out_target.layout + per_atom_layout = TensorMap( + keys=layout.keys, + blocks=[ + TensorBlock( + values=block.values, + samples=Labels( + ["system", "atom"], block.samples.values.reshape((0, 2)) + ), + components=block.components, + properties=block.properties, + ) + for block in layout.blocks() + ], + ) + + for in_name in [in_names["bottom"], in_names["top"]]: + self._input_target_infos[in_name] = TargetInfo( + per_atom_layout, + quantity=out_target.quantity, + unit=out_target.unit, + ) + + layout.set_info("pooling_type", pooling_hypers["type"]) + layout.set_info( + "alpha_bottom", str(pooling_hypers.get("alpha_bottom", 20.0)) + ) + layout.set_info("alpha_top", str(pooling_hypers.get("alpha_top", -20.0))) + layout.set_info("bottom_name", in_names["bottom"]) + layout.set_info("top_name", in_names["top"]) + + self._block_shapes[out_name] = [ + (-1, *block.values.shape[1:]) for block in out_target.layout.blocks() + ] + + def requested_target_infos(self) -> dict[str, TargetInfo]: + """ + Returns the list of requested target infos for the hook. + + :return: A list of requested target names. + """ + return self._input_target_infos + + def requested_inputs(self) -> dict[str, ModelOutput]: + """ + Returns the list of requested inputs for the hook. + + :return: A list of requested input names. + """ + return { + k: ModelOutput( + quantity=target_info.quantity, + unit=target_info.unit, + sample_kind="atom", + ) + for k, target_info in self._input_target_infos.items() + } + + def supported_outputs(self) -> dict[str, ModelOutput]: + """ + Returns the supported outputs for the hook. + + :return: A list of supported output names. + """ + return { + out_name: ModelOutput( + quantity=target_info.quantity, + unit=target_info.unit, + sample_kind="system", + ) + for out_name, target_info in self.out_targets.items() + } + + def forward( + self, + systems: list[System], + outputs: dict[str, ModelOutput], + inputs: dict[str, TensorMap], + selected_atoms: Optional[Labels] = None, + ) -> dict[str, TensorMap]: + requested_outs: list[str] = [] + for out_name in self.out_targets: + if out_name in outputs: + requested_outs.append(out_name) + if not requested_outs: + # Quick exit if no output of this hook is requested + return {} + + device = systems[0].positions.device + + num_systems = len(systems) + system_indices = [] + for i, system in enumerate(systems): + system_indices.append( + torch.full((len(system),), i, dtype=torch.int32, device=device) + ) + system_indices = torch.cat(system_indices, dim=0) + + return_dict: dict[str, TensorMap] = {} + for out_name in requested_outs: + out_target = self.out_targets[out_name] + layout = out_target.layout.to(device) + + # Get the parameters for this target + info = layout.info() + alpha_bottom = float(info["alpha_bottom"]) + alpha_top = float(info["alpha_top"]) + input_bottom = info["bottom_name"] + input_top = info["top_name"] + pooling_type = info["pooling_type"] + block_shapes = self._block_shapes[out_name] + + # Build the output TensorMap by pooling over the per-atom values + # for each block + blocks: list[TensorBlock] = [] + for i, layout_block in enumerate(layout.blocks()): + values_bottom = inputs[input_bottom].block(i).values.ravel() + values_top = inputs[input_top].block(i).values.ravel() + + if pooling_type == "softmax": + # Self-weighted softmax pool: the softmax weights are computed + # directly from the per-atom values themselves, so atoms with the + # most extreme contribution dominate. Strictly intensive (weights + # sum to 1) and recovers a hard max/min as ``|alpha| -> infinity``. + out_bottom = _scatter_softmax_pool( + values_bottom, alpha_bottom, system_indices, num_systems + ) + out_top = _scatter_softmax_pool( + values_top, alpha_top, system_indices, num_systems + ) + else: + out_bottom = _scatter_logsumexp( + values_bottom, alpha_bottom, system_indices, num_systems + ) + out_top = _scatter_logsumexp( + values_top, alpha_top, system_indices, num_systems + ) + + gap = out_top - out_bottom + + blocks.append( + TensorBlock( + values=gap.reshape(block_shapes[i]), + samples=Labels( + names=["system"], + values=torch.arange( + num_systems, dtype=torch.int32, device=device + ).reshape(-1, 1), + ), + components=layout_block.components, + properties=layout_block.properties, + ) + ) + + return_dict[out_name] = TensorMap( + keys=layout.keys, + blocks=blocks, + ) + + return return_dict diff --git a/src/metatrain/utils/hooks/intensive_gap/tests/__init__.py b/src/metatrain/utils/hooks/intensive_gap/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/metatrain/utils/hooks/intensive_gap/tests/test_basic.py b/src/metatrain/utils/hooks/intensive_gap/tests/test_basic.py new file mode 100644 index 0000000000..2e5514367b --- /dev/null +++ b/src/metatrain/utils/hooks/intensive_gap/tests/test_basic.py @@ -0,0 +1,46 @@ +import pytest + +from metatrain.utils.hooks.testing.hook import HookTests +from metatrain.utils.hooks.testing.inputs import InputTests +from metatrain.utils.hooks.testing.output import OutputTests +from metatrain.utils.hooks.testing.torchscript import TorchscriptTests + + +class IntensiveGapTests(HookTests): + hook = "intensive_gap" + + @pytest.fixture + def hypers(self): + return {"outputs": "mtt::energy"} + + def build_hypers(self, input, output): + if input is None: + output_name = output["energy"] + input_names = [ + f"mtt::aux::gap_bottom::{output_name.replace('mtt::', '')}", + f"mtt::aux::gap_top::{output_name.replace('mtt::', '')}", + ] + return {"outputs": output_name}, input_names, [output_name] + elif output is None: + ... + else: + bottom_name = input["local_energy"] + top_name = f"{bottom_name}_1" + output_name = output["energy"] + return ( + { + "inputs": {"bottom": bottom_name, "top": top_name}, + "outputs": output_name, + }, + [bottom_name, top_name], + [output_name], + ) + + +class TestTorchscript(TorchscriptTests, IntensiveGapTests): ... + + +class TestInputs(InputTests, IntensiveGapTests): ... + + +class TestOutputs(OutputTests, IntensiveGapTests): ... diff --git a/src/metatrain/utils/hooks/tensor_basis/__init__.py b/src/metatrain/utils/hooks/tensor_basis/__init__.py new file mode 100644 index 0000000000..528eabebba --- /dev/null +++ b/src/metatrain/utils/hooks/tensor_basis/__init__.py @@ -0,0 +1,4 @@ +from .hook import TensorBasis + + +__hook__ = TensorBasis diff --git a/src/metatrain/utils/hooks/tensor_basis/documentation.py b/src/metatrain/utils/hooks/tensor_basis/documentation.py new file mode 100644 index 0000000000..064e2960d1 --- /dev/null +++ b/src/metatrain/utils/hooks/tensor_basis/documentation.py @@ -0,0 +1,52 @@ +""" +Tensor basis +============ + +Provides a tensor basis in which to predict spherical tensor targets, +following the approach described in this work :footcite:p:`domina2025representing`. + +This hook creates a basis for each target and each angular channel +(``o3_lambda`` block) of the outputs. Then it asks for invariant +coefficients to apply to the basis to produce the target. By using +this hook one can: + +- **Use an architecture that produces only scalar outputs**, and still be + able to predict tensorial targets. +- **Reduce the cost of equivariant models** by asking them to produce only + scalar outputs, and then use this hook to access the angular momentum + channels of the target. + +""" + +from typing import Optional, TypedDict + +from metatrain.soap_bpnn.documentation import SOAPConfig +from metatrain.utils.hypers import init_with_defaults + + +class Hypers(TypedDict): + """ + Hyperparameters for the tensor basis hook. + """ + + soap: SOAPConfig = init_with_defaults(SOAPConfig) + """Hyperparameters used to compute the spherical expansions from + which the vector basis will be built. Higher angular momentum + channels are built by augmenting the order of the vector basis.""" + + inputs: Optional[str | list] = None + """ + Name or names of the targets to use as invariant coefficients + to apply to the tensor basis. + + If ``None``, they will be set as + ``mtt::aux::scalars::{output_name.replace('mtt::', '')}`` + for each output name. + """ + + outputs: Optional[str | list] = None + """ + Name or names of the targets to predict through a tensor basis. + + A separate tensor basis will be built for each target. + """ diff --git a/src/metatrain/utils/hooks/tensor_basis/hook.py b/src/metatrain/utils/hooks/tensor_basis/hook.py new file mode 100644 index 0000000000..2ae7621fbe --- /dev/null +++ b/src/metatrain/utils/hooks/tensor_basis/hook.py @@ -0,0 +1,320 @@ +from typing import Optional + +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap +from metatomic.torch import ModelOutput, System + +from metatrain.soap_bpnn.documentation import SOAPConfig +from metatrain.soap_bpnn.modules.tensor_basis import TensorBasis as TensorBasisModule +from metatrain.utils.data import DatasetInfo, TargetInfo +from metatrain.utils.data.target_info import get_generic_target_info +from metatrain.utils.hypers import init_with_defaults +from metatrain.utils.sum_over_atoms import sum_over_atoms + +from ..abc import HookInterface +from .documentation import Hypers + + +def concatenate_structures( + systems: list[System], +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """ + Concatenate a list of systems into a single batch. + + :param systems: List of systems to concatenate. + :return: A tuple containing the concatenated positions, centers, neighbors, + species, cells, and cell shifts. + """ + positions = [] + centers = [] + neighbors = [] + species = [] + cell_shifts = [] + cells = [] + node_counter = 0 + + for system in systems: + positions.append(system.positions) + species.append(system.types) + + neighbor_list = system.get_neighbor_list(system.known_neighbor_lists()[0]) + nl_values = neighbor_list.samples.values + + centers.append(nl_values[:, 0] + node_counter) + neighbors.append(nl_values[:, 1] + node_counter) + cell_shifts.append(nl_values[:, 2:]) + + cells.append(system.cell) + + node_counter += len(system.positions) + + positions = torch.cat(positions) + centers = torch.cat(centers) + neighbors = torch.cat(neighbors) + species = torch.cat(species) + cells = torch.stack(cells) + cell_shifts = torch.cat(cell_shifts) + + return ( + positions, + centers, + neighbors, + species, + cells, + cell_shifts, + ) + + +class TensorBasis(HookInterface[Hypers]): + """ + Provides a tensor basis in which to predict spherical tensor targets. + + :param hypers: A dictionary with the hook's hyper-parameters. + :param dataset_info: Information containing details about the dataset, such as + target quantities and atomic types. + """ + + __checkpoint_version__ = 1 + + def __init__(self, hypers: Hypers, dataset_info: DatasetInfo): + super().__init__(hypers, dataset_info) + + self.hypers = hypers + + # Helper to map from atomic number to the index of that atomic + # number in the list of atomic types. + species_to_species_index = torch.empty( + max(dataset_info.atomic_types) + 1, dtype=torch.long + ) + species_to_species_index[dataset_info.atomic_types] = torch.arange( + len(dataset_info.atomic_types) + ) + self.register_buffer("species_to_species_index", species_to_species_index) + + # Get the information about the output targets from the dataset info + outputs = hypers["outputs"] + if isinstance(outputs, str): + outputs = [outputs] + if outputs is None: + raise ValueError("TensorBasis hook requires at least one output target") + self.out_targets = {name: dataset_info.targets[name] for name in outputs} + + # Names for the inputs that we will request from the model + self._input_names = [ + f"mtt::aux::scalars::{name.replace('mtt::', '')}" + for name in self.out_targets + ] + + # Build the basis calculators for each target, + # and the output that we have to request from the model + soap_hypers = hypers.get("soap", init_with_defaults(SOAPConfig)) + self.basis_calculators = torch.nn.ModuleDict({}) + self._input_target_infos = {} + for input_name, target_name in zip( + self._input_names, self.out_targets, strict=True + ): + target = self.out_targets[target_name] + # Get one basis calculator for each block of the target, since each block + # has different o3_lambda and o3_sigma values. + self.basis_calculators[target_name] = torch.nn.ModuleList( + [ + TensorBasisModule( + dataset_info.atomic_types, + soap_hypers, + o3_lambda=key["o3_lambda"], + o3_sigma=key["o3_sigma"], + add_lambda_basis=True, + legacy=False, + ) + for key in target.layout.keys + ] + ) + + # Build the input that we will request from the model. + # We will ask for invariant coefficients. For each block we ask for + # 2l+1 coefficients for each property, since the basis will have + # 2l+1 tensors. + # We ask for all the coefficients in a single block, we will untangle + # them in the forward pass. + num_properties = sum( + block.values.shape[1] * block.values.shape[2] + for block in target.layout.blocks() + ) + self._input_target_infos[input_name] = get_generic_target_info( + input_name, + { + "quantity": "_", + "unit": "", + "type": { + "spherical": {"irreps": [{"o3_lambda": 0, "o3_sigma": 1}]} + }, + "num_subtargets": num_properties, + "sample_kind": "atom", + }, + ) + + def requested_target_infos(self) -> dict[str, TargetInfo]: + """ + Returns the list of requested target infos for the hook. + + :return: A list of requested target names. + """ + return self._input_target_infos + + def requested_inputs(self) -> dict[str, ModelOutput]: + """ + Returns the list of requested inputs for the hook. + + :return: A list of requested input names. + """ + return { + name: ModelOutput( + quantity=target.quantity, + unit=target.unit, + sample_kind="atom", + ) + for name, target in self._input_target_infos.items() + } + + def supported_outputs(self) -> dict[str, ModelOutput]: + """ + Returns the supported outputs for the hook. + + :return: A list of supported output names. + """ + return { + out_name: ModelOutput( + quantity=target.quantity, + unit=target.unit, + sample_kind="system", + ) + for out_name, target in self.out_targets.items() + } + + def forward( + self, + systems: list[System], + outputs: dict[str, ModelOutput], + inputs: dict[str, TensorMap], + selected_atoms: Optional[Labels] = None, + ) -> dict[str, TensorMap]: + requested_outs: list[str] = [] + for out_name in self.out_targets: + if out_name in outputs: + requested_outs.append(out_name) + if not requested_outs: + # Quick exit if no output of this hook is requested + return {} + + device = systems[0].positions.device + + # ------------------------------- + # Get structure information + # ------------------------------- + + system_sizes = [len(system) for system in systems] + system_sizes_tensor = torch.tensor(system_sizes, device=device) + system_indices = torch.repeat_interleave( + torch.arange(len(systems), device=device), system_sizes_tensor + ) + atom_indices = torch.cat( + [torch.arange(size, device=device) for size in system_sizes] + ) + sample_values = torch.stack([system_indices, atom_indices], dim=1) + + ( + positions, + centers, + neighbors, + species, + cells, + cell_shifts, + ) = concatenate_structures(systems) + species = self.species_to_species_index[species] + + # somehow the backward of this operation is very slow at evaluation, + # where there is only one cell, therefore we simplify the calculation + # for that case + if len(cells) == 1: + cell_contributions = cell_shifts.to(cells.dtype) @ cells[0] + else: + cell_contributions = torch.einsum( + "ab, abc -> ac", + cell_shifts.to(cells.dtype), + cells[system_indices][centers], + ) + + interatomic_vectors = ( + positions[neighbors] - positions[centers] + cell_contributions + ) + + # ------------------------------------ + # Build the values for each target + # ------------------------------------ + + return_dict: dict[str, TensorMap] = {} + for target_name, basis_calculators in self.basis_calculators.items(): + if target_name in requested_outs: + target_info = self.out_targets[target_name] + target_invariant_coefficients = inputs[target_name].block().values + + offset = 0 + blocks: list[TensorBlock] = [] + for i, basis_calculator in enumerate(basis_calculators): + layout_block = target_info.layout.block(i) + + # Get shapes of the invariant coefficients to retrieve + # for this block. + n_properties = layout_block.properties.values.shape[0] + n_basis = layout_block.values.shape[1] + count = n_properties * n_basis + + # Get those invariant coefficients + invariant_coefficients = target_invariant_coefficients[ + :, 0, offset : offset + count + ].reshape(-1, n_properties, n_basis) + # Update counter for the next block + offset += count + + # Now get the tensor basis. + tensor_basis = basis_calculator( + interatomic_vectors, + centers, + neighbors, + species, + sample_values, + selected_atoms=None, + ) + + # Multiply the invariant coefficients by the tensor basis + # to get the final values for each atom. + atomic_property_tensor = torch.einsum( + "spb, scb -> scp", + invariant_coefficients, + tensor_basis, + ) + + # Build the tensor block. + blocks.append( + TensorBlock( + values=atomic_property_tensor, + samples=Labels( + names=["system", "atom"], values=sample_values + ), + components=layout_block.components, + properties=layout_block.properties, + ) + ) + + tmap = TensorMap( + keys=target_info.layout.keys, + blocks=blocks, + ) + + if target_info.sample_kind == "system": + tmap = sum_over_atoms(tmap) + return_dict[target_name] = tmap + + return return_dict diff --git a/src/metatrain/utils/hooks/tensor_basis/tests/__init__.py b/src/metatrain/utils/hooks/tensor_basis/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/metatrain/utils/hooks/tensor_basis/tests/test_basic.py b/src/metatrain/utils/hooks/tensor_basis/tests/test_basic.py new file mode 100644 index 0000000000..10639bb4ff --- /dev/null +++ b/src/metatrain/utils/hooks/tensor_basis/tests/test_basic.py @@ -0,0 +1,39 @@ +import pytest + +from metatrain.utils.hooks.testing.hook import HookTests +from metatrain.utils.hooks.testing.inputs import InputTests +from metatrain.utils.hooks.testing.output import OutputTests +from metatrain.utils.hooks.testing.torchscript import TorchscriptTests + + +class TensorBasisTests(HookTests): + hook = "tensor_basis" + + @pytest.fixture + def hypers(self): + return {"outputs": "mtt::global_dipole"} + + def build_hypers(self, input, output): + if input is None: + output_name = output["global_dipole"] + input_name = f"mtt::aux::scalars::{output_name.replace('mtt::', '')}" + return {"outputs": output_name}, [input_name], [output_name] + elif output is None: + ... + else: + input_name = input["local_energy"] + output_name = output["global_dipole"] + return ( + {"inputs": input_name, "outputs": output_name}, + [input_name], + [output_name], + ) + + +class TestTorchscript(TorchscriptTests, TensorBasisTests): ... + + +class TestInputs(InputTests, TensorBasisTests): ... + + +class TestOutputs(OutputTests, TensorBasisTests): ... diff --git a/src/metatrain/utils/hooks/testing/__init__.py b/src/metatrain/utils/hooks/testing/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/src/metatrain/utils/hooks/testing/hook.py b/src/metatrain/utils/hooks/testing/hook.py new file mode 100644 index 0000000000..d489fa375b --- /dev/null +++ b/src/metatrain/utils/hooks/testing/hook.py @@ -0,0 +1,187 @@ +from typing import Optional + +import pytest +import torch +from metatensor.torch import Labels, TensorBlock, TensorMap +from metatomic.torch import System + +from metatrain.utils.data import DatasetInfo +from metatrain.utils.data.target_info import ( + get_energy_target_info, + get_generic_target_info, +) +from metatrain.utils.hooks.abc import HookInterface +from metatrain.utils.hooks.helpers import load_hook + + +_target_types = ["energy", "local_energy", "global_dipole", "local_multipole"] + + +class HookTests: + hook: str + """Name of the hook to be tested. + + Based on this, the test suite will find the hook class + as well as the hyperparameters. + """ + + @property + def hook_cls(self) -> type[HookInterface]: + """The hook class to be tested.""" + return load_hook(self.hook) + + def build_hypers( + self, input: Optional[dict], output: Optional[dict] + ) -> tuple[dict, list, list]: + raise NotImplementedError("Subclasses must implement build_hypers method.") + + @pytest.fixture + def dataset_info(self) -> DatasetInfo: + + global_energy = get_energy_target_info( + "energy", + dict( + sample_kind="system", + unit="eV", + quantity="energy", + num_subtargets=1, + ), + ) + local_energy = get_generic_target_info( + "mtt::local_energy", + dict( + sample_kind="atom", + unit="eV", + quantity="energy", + num_subtargets=1, + type="scalar", + ), + ) + global_dipole = get_generic_target_info( + "mtt::global_dipole", + dict( + sample_kind="system", + unit="eV", + quantity="energy", + num_subtargets=1, + type=dict(spherical=dict(irreps=[{"o3_lambda": 1, "o3_sigma": 1}])), + ), + ) + local_multipole = get_generic_target_info( + "mtt::local_multipole", + dict( + sample_kind="atom", + unit="eV", + quantity="energy", + num_subtargets=1, + type=dict( + spherical=dict( + irreps=[ + {"o3_lambda": 0, "o3_sigma": 1}, + {"o3_lambda": 1, "o3_sigma": 1}, + ] + ) + ), + ), + ) + + return DatasetInfo( + length_unit="angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "mtt::energy": global_energy, + "mtt::local_energy": local_energy, + "mtt::local_energy_1": local_energy, + "mtt::global_dipole": global_dipole, + "mtt::local_multipole": local_multipole, + }, + extra_data={ + "mtt::extra_energy": global_energy, + "mtt::extra_local_energy": local_energy, + "mtt::extra_local_energy_1": local_energy, + "mtt::extra_global_dipole": global_dipole, + "mtt::extra_local_multipole": local_multipole, + }, + ) + + @pytest.fixture + def targets(self) -> dict[str, str]: + return {t: f"mtt::{t}" for t in _target_types} + + @pytest.fixture + def extra_data(self) -> dict[str, str]: + return {t: f"mtt::extra_{t}" for t in _target_types} + + @pytest.fixture + def intermediate(self) -> dict[str, str]: + return {t: f"mtt::intermediate::{t}" for t in _target_types} + + def get_random( + self, dataset_info: DatasetInfo, requested: list[str] + ) -> tuple[list[System], dict[str, TensorMap]]: + """Get random values for targets in the dataset info. + + :param dataset_info: The dataset info to use. + :param requested: The requested random tensormaps. + :return: A tuple of a list of systems and a dictionary + of random tensormaps. + """ + return_dict = {} + all_target_infos = {**dataset_info.targets, **dataset_info.extra_data} + + system = System( + positions=torch.zeros((1, 3)), + types=torch.tensor([1]), + cell=torch.eye(3), + pbc=torch.tensor([True, True, True]), + ) + + for name in requested: + target = all_target_infos[name] + + if target.sample_kind == "system": + tmap = TensorMap( + keys=target.layout.keys, + blocks=[ + TensorBlock( + values=torch.rand( + (1, *block.values.shape[1:]), + dtype=block.values.dtype, + device=block.values.device, + ), + samples=Labels( + names=["system"], + values=torch.tensor([[0]], device=block.values.device), + ), + components=block.components, + properties=block.properties, + ) + for block in target.layout.blocks() + ], + ) + elif target.sample_kind == "atom": + tmap = TensorMap( + keys=target.layout.keys, + blocks=[ + TensorBlock( + values=torch.rand( + (1, *block.values.shape[1:]), + dtype=block.values.dtype, + device=block.values.device, + ), + samples=Labels( + names=["system", "atom"], + values=torch.tensor( + [[0, 0]], device=block.values.device + ), + ), + components=block.components, + properties=block.properties, + ) + for block in target.layout.blocks() + ], + ) + + return_dict[name] = tmap + + return [system], return_dict diff --git a/src/metatrain/utils/hooks/testing/inputs.py b/src/metatrain/utils/hooks/testing/inputs.py new file mode 100644 index 0000000000..169f683dd7 --- /dev/null +++ b/src/metatrain/utils/hooks/testing/inputs.py @@ -0,0 +1,176 @@ +from metatomic.torch import ModelOutput + +from metatrain.utils.data import DatasetInfo + +from .hook import HookTests + + +class InputTests(HookTests): + def test_only_output(self, dataset_info: DatasetInfo, targets: dict) -> None: + + input = None + output = targets + + hypers, expected_inputs, expected_outputs = self.build_hypers( + input=input, output=output + ) + + self.check_correctness( + hypers, + dataset_info, + expected_inputs, + expected_outputs, + inputs_in_dataset=False, + ) + + def test_output_and_specific_input( + self, dataset_info: DatasetInfo, intermediate: dict, targets: dict + ) -> None: + + input = intermediate + output = targets + + hypers, expected_inputs, expected_outputs = self.build_hypers( + input=input, output=output + ) + + self.check_correctness( + hypers, + dataset_info, + expected_inputs, + expected_outputs, + inputs_in_dataset=False, + ) + + def test_output_and_extradata_input( + self, dataset_info: DatasetInfo, extra_data: dict, targets: dict + ) -> None: + + input = extra_data + output = targets + + hypers, expected_inputs, expected_outputs = self.build_hypers( + input=input, output=output + ) + + self.check_correctness( + hypers, + dataset_info, + expected_inputs, + expected_outputs, + ) + + def check_correctness( + self, + hypers: dict, + dataset_info: DatasetInfo, + expected_inputs: list[str], + expected_outputs: list[str], + inputs_in_dataset: bool = True, + outputs_in_dataset: bool = True, + ) -> None: + hook = self.hook_cls(hypers, dataset_info) + + all_target_infos = {**dataset_info.targets, **dataset_info.extra_data} + + requested_target_infos = hook.requested_target_infos() + requested_inputs = hook.requested_inputs() + supported_outputs = hook.supported_outputs() + + if inputs_in_dataset: + expected_inp_targets = { + in_name: all_target_infos[in_name] for in_name in expected_inputs + } + + assert requested_target_infos == expected_inp_targets, ( + f"Requested target infos do not match expected inputs.\n" + f"Expected: {expected_inp_targets}\n" + f"Got: {requested_target_infos}" + ) + + expected_model_inputs = { + in_name: ModelOutput( + quantity=target_info.quantity, + unit=target_info.unit, + sample_kind=target_info.sample_kind, + ) + for in_name, target_info in expected_inp_targets.items() + } + + # Check the ModelOutput objects that the hook requests as inputs. + assert len(requested_inputs) == len(expected_model_inputs), ( + f"There are more requested inputs than expected. " + f"Expected: {len(expected_model_inputs)}, Got: {len(requested_inputs)}" + ) + for in_name, model_output in requested_inputs.items(): + assert in_name in expected_model_inputs, ( + f"Requested input '{in_name}' is not in the expected inputs." + ) + expected_output = expected_model_inputs[in_name] + assert model_output.quantity == expected_output.quantity, ( + f"Requested input '{in_name}' has quantity " + f"'{model_output.quantity}', " + f"but expected '{expected_output.quantity}'." + ) + assert model_output.unit == expected_output.unit, ( + f"Requested input '{in_name}' has unit '{model_output.unit}', " + f"but expected '{expected_output.unit}'." + ) + assert model_output.sample_kind == expected_output.sample_kind, ( + f"Requested input '{in_name}' has sample kind " + f"'{model_output.sample_kind}', " + f"but expected '{expected_output.sample_kind}'." + ) + else: + assert len(requested_inputs) == len(expected_model_inputs) + assert set(expected_inputs) == set(requested_inputs) + assert len(requested_target_infos) == len(expected_inputs) + assert set(expected_inputs) == set(requested_target_infos) + + if outputs_in_dataset: + expected_out_targets = { + out_name: all_target_infos[out_name] for out_name in expected_outputs + } + expected_model_outputs = { + out_name: ModelOutput( + quantity=target_info.quantity, + unit=target_info.unit, + sample_kind=target_info.sample_kind, + ) + for out_name, target_info in expected_out_targets.items() + } + + # Check the ModelOutput objects that the hook supports as outputs. + assert len(supported_outputs) == len(expected_model_outputs), ( + f"There are more supported outputs than expected. " + f"Expected: {len(expected_model_outputs)}, " + f"Got: {len(supported_outputs)}" + ) + for out_name, model_output in supported_outputs.items(): + assert out_name in expected_model_outputs, ( + f"Supported output '{out_name}' is not in the expected outputs." + ) + expected_output = expected_model_outputs[out_name] + assert model_output.quantity == expected_output.quantity, ( + f"Supported output '{out_name}' has quantity " + f"'{model_output.quantity}', " + f"but expected '{expected_output.quantity}'." + ) + assert model_output.unit == expected_output.unit, ( + f"Supported output '{out_name}' has unit '{model_output.unit}', " + f"but expected '{expected_output.unit}'." + ) + assert model_output.sample_kind == expected_output.sample_kind, ( + f"Supported output '{out_name}' has sample kind " + f"'{model_output.sample_kind}', " + f"but expected '{expected_output.sample_kind}'." + ) + else: + assert len(supported_outputs) == len(expected_outputs), ( + f"The number of supported outputs does not match the expected number. " + f"Expected: {len(expected_outputs)}, Got: {len(supported_outputs)}" + ) + assert set(expected_outputs) == set(supported_outputs), ( + f"The supported outputs are not as expected. " + f"Expected: {set(expected_outputs)}, Got: {set(supported_outputs)}" + ) diff --git a/src/metatrain/utils/hooks/testing/output.py b/src/metatrain/utils/hooks/testing/output.py new file mode 100644 index 0000000000..2e8df350ea --- /dev/null +++ b/src/metatrain/utils/hooks/testing/output.py @@ -0,0 +1,101 @@ +import torch +from metatomic.torch import System + +from metatrain.utils.data import DatasetInfo + +from .hook import HookTests + + +class OutputTests(HookTests): + def test_empty(self, hypers: dict, dataset_info: DatasetInfo) -> None: + """Test that the hook returns an empty dictionary + when no outputs are requested. + + :param hypers: A dictionary with the hook's hyper-parameters. + :param dataset_info: Information containing details about the dataset, such as + target quantities and atomic types. + """ + hook = self.hook_cls(hypers, dataset_info) + + # We call the forward method of the hook with empty inputs + # and outputs, and a dummy system. + outputs: dict = {} + inputs: dict = {} + system = System( + positions=torch.zeros((1, 3)), + types=torch.tensor([1]), + cell=torch.eye(3), + pbc=torch.tensor([True, True, True]), + ) + + result = hook( + systems=[system], outputs=outputs, inputs=inputs, selected_atoms=None + ) + + # Check that the result is an empty dictionary + assert result == {} + + def test_output( + self, dataset_info: DatasetInfo, extra_data: dict, targets: dict + ) -> None: + """Test that the hook returns the correct output. + + :param dataset_info: Information containing details about the dataset, + such as target quantities and atomic types. + :param extra_data: A dictionary containing the target names that + are present in the dataset's extra_data. + :param targets: A dictionary containing the target names that are + present in the dataset's targets. + """ + + input = extra_data + output = targets + + hypers, expected_inputs, expected_outputs = self.build_hypers( + input=input, output=output + ) + + hook = self.hook_cls(hypers, dataset_info) + + # We call the forward method with random values + systems, inputs = self.get_random(dataset_info, expected_inputs) + + out_targets = { + out_name: dataset_info.targets[out_name] for out_name in expected_outputs + } + + result = hook( + systems=systems, + outputs=hook.supported_outputs(), + inputs=inputs, + selected_atoms=None, + ) + + # Check that the result contains all requested outputs + assert set(result.keys()) == set(expected_outputs) + # Check that each output has the expected metadata + for out_name in expected_outputs: + out_tmap = result[out_name] + layout = out_targets[out_name].layout + + assert out_tmap.keys == layout.keys, ( + f"Output {out_name} has keys {out_tmap.keys}, " + f"but expected {layout.keys}" + ) + + for key, layout_block in layout.items(): + out_block = out_tmap[key] + assert out_block.samples.names == layout_block.samples.names + assert out_block.samples.values.shape[0] == 1 + + assert out_block.components == layout_block.components, ( + f"Output {out_name} block {key} has components " + f"{out_block.components}, " + f"but expected {layout_block.components}" + ) + + assert out_block.properties == layout_block.properties, ( + f"Output {out_name} block {key} has properties " + f"{out_block.properties}, " + f"but expected {layout_block.properties}" + ) diff --git a/src/metatrain/utils/hooks/testing/torchscript.py b/src/metatrain/utils/hooks/testing/torchscript.py new file mode 100644 index 0000000000..c067c4394c --- /dev/null +++ b/src/metatrain/utils/hooks/testing/torchscript.py @@ -0,0 +1,12 @@ +import torch + +from metatrain.utils.data import DatasetInfo + +from .hook import HookTests + + +class TorchscriptTests(HookTests): + def test_torchscript(self, hypers: dict, dataset_info: DatasetInfo) -> None: + hook = self.hook_cls(hypers, dataset_info) + + torch.jit.script(hook) diff --git a/tests/utils/test_hooks.py b/tests/utils/test_hooks.py new file mode 100644 index 0000000000..2ac54f5249 --- /dev/null +++ b/tests/utils/test_hooks.py @@ -0,0 +1,205 @@ +import pytest + +from metatrain.utils.data import DatasetInfo +from metatrain.utils.data.target_info import get_energy_target_info +from metatrain.utils.hooks.helpers import find_all_hooks, setup_hooks + + +@pytest.fixture +def dataset_info(): + + global_energy = get_energy_target_info( + "energy", + dict( + sample_kind="system", + unit="eV", + quantity="energy", + num_subtargets=1, + ), + ) + local_energy = get_energy_target_info( + "mtt::local_energy", + dict( + sample_kind="atom", + unit="eV", + quantity="energy", + num_subtargets=1, + ), + ) + + return DatasetInfo( + length_unit="angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "energy": global_energy, + "mtt::local_energy": local_energy, + "mtt::some_other_energy": local_energy, + }, + extra_data={"mtt::extra_energy": local_energy}, + ) + + +def test_find_all_hooks(): + hooks = find_all_hooks() + assert isinstance(hooks, list) + assert all(isinstance(hook, str) for hook in hooks) + assert len(hooks) == 4 + assert "identity" in hooks + assert "global_multipoles" in hooks + assert "intensive_gap" in hooks + assert "tensor_basis" in hooks + + +class TestSetupPostHooks: + def test_empty(self, dataset_info): + + hypers = {} + + dataset_info = dataset_info.copy() + dataset_info.hooks = hypers + + post_hooks, model_outs = setup_hooks(dataset_info) + + assert len(post_hooks) == 0 + assert model_outs == dataset_info.targets + + def test_single_hook(self, dataset_info): + """Simplest test of a hook owning a single output, and its inputs + are created so that the model can produce them.""" + + hypers = { + "identity": "mtt::local_energy", + } + + dataset_info = dataset_info.copy() + dataset_info.hooks = hypers + + post_hooks, model_outs = setup_hooks(dataset_info) + + assert len(post_hooks) == 1 + assert "mtt::local_energy" not in model_outs + + assert "energy" in model_outs + assert "mtt::some_other_energy" in model_outs + assert len(model_outs) > 2 + + def test_single_hook_with_inputs(self, dataset_info): + """Test that hooks can use a target as input. + + This means that the target will flow to the loss function, + but still we can use it as input to a hook. + """ + hypers = { + "identity": { + "outputs": "mtt::local_energy", + "inputs": "mtt::some_other_energy", + }, + } + + dataset_info = dataset_info.copy() + dataset_info.hooks = hypers + + post_hooks, model_outs = setup_hooks(dataset_info) + + assert len(post_hooks) == 1 + assert "mtt::local_energy" not in model_outs + assert model_outs == { + k: v for k, v in dataset_info.targets.items() if k != "mtt::local_energy" + }, ( + "Hook should not add extra outputs to the model, since its inputs" + " are already present in the model outputs." + ) + + def test_single_hook_with_inputs_from_extra_data(self, dataset_info): + """Test that hooks can get their inputs from the extra data, and + that in that case they don't add an output to the model.""" + + hypers = { + "identity": {"outputs": "mtt::local_energy", "inputs": "mtt::extra_energy"}, + } + + dataset_info = dataset_info.copy() + dataset_info.hooks = hypers + + post_hooks, model_outs = setup_hooks(dataset_info) + + assert len(post_hooks) == 1 + assert "mtt::local_energy" not in model_outs + assert model_outs == { + k: v for k, v in dataset_info.targets.items() if k != "mtt::local_energy" + }, ( + "Hook should not add extra outputs to the model, since its inputs" + " are already present in the extra data." + ) + + def test_chained_hooks(self, dataset_info): + """Check that an output from a previous hook can be used + as an input by the followinf hooks.""" + + hypers = { + "identity": {"outputs": "mtt::gap_bottom", "inputs": "mtt::local_energy"}, + "intensive_gap": { + "outputs": "energy", + "inputs": { + "bottom": "mtt::gap_bottom", + }, + }, + } + + dataset_info = dataset_info.copy() + dataset_info.hooks = hypers + + post_hooks, model_outs = setup_hooks(dataset_info) + + assert len(post_hooks) == 2 + assert "energy" not in model_outs + + assert "mtt::local_energy" in model_outs + assert "mtt::some_other_energy" in model_outs + + # There should be an extra model output for the top of the gap, + # requested by the intensive gap hook. + assert len(model_outs) == 3 + + def test_not_available_output_error(self, dataset_info): + """A hook is asked to use an output that doesn't exist, + then it should raise an error. + """ + + hypers = { + "intensive_gap": { + "outputs": "mtt::intermediate_energy", + }, + } + + dataset_info = dataset_info.copy() + dataset_info.hooks = hypers + + with pytest.raises(ValueError, match="mtt::intermediate_energy"): + post_hooks, model_outs = setup_hooks(dataset_info) + + def test_chained_hooks_with_laterdependency(self, dataset_info): + """Similar to the previous test, but in this case the output will + be known once the second hook is set up.""" + + hypers = { + "intensive_gap": { + "outputs": "mtt::intermediate_energy", + }, + "identity": {"outputs": "energy", "inputs": "mtt::intermediate_energy"}, + } + + dataset_info = dataset_info.copy() + dataset_info.hooks = hypers + + post_hooks, model_outs = setup_hooks(dataset_info) + + assert len(post_hooks) == 2 + assert "energy" not in model_outs + + assert "mtt::local_energy" in model_outs + assert "mtt::some_other_energy" in model_outs + + # There should be two extra model outputs requested by the + # intensive gap hook. + assert len(model_outs) == 4