Proof of concept for hooks

Author: pfebrer

src/metatrain/experimental/mace/documentation.py

@@ -121,6 +121,7 @@
 
 from metatrain.composition.documentation import FixedCompositionWeights
 from metatrain.utils.finetuning import FullFinetuneHypers
+from metatrain.utils.hooks import PostHooksHypers
 from metatrain.utils.loss import LossSpecification
 from metatrain.utils.scaler import FixedScalerWeights
 
@@ -295,6 +296,9 @@ class ModelHypers(TypedDict):
     use_agnostic_product: bool = False
     """Use element agnostic product"""
 
+    post_hooks: PostHooksHypers = {}
+    """Post-processing hooks to apply to the outputs of the model."""
+
 
 class TrainerHypers(TypedDict):
     # Optimizer hypers (directly using MACE's scripts)

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 setup_post_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,15 @@ def __init__(self, hypers: ModelHypers, dataset_info: DatasetInfo) -> None:
         # the model during training.
         train_dataset_info = self._train_dataset_info(dataset_info)
 
+        post_hooks, model_outs = setup_post_hooks(
+            self.hypers["post_hooks"], train_dataset_info
+        )
+        self.post_hooks = torch.nn.ModuleList(post_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 +285,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
         }
 
         # ---------------------------
@@ -358,6 +365,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.post_hooks:
+            requested_inputs = hook.requested_inputs()
+            outputs.update(requested_inputs)
+
         # --------------------------
         # Moving to device and dtype
         # --------------------------
@@ -460,6 +475,13 @@ def forward(
                 else sum_over_atoms(per_atom_output)
             )
 
+        # -----------------------------------
+        #            Apply hooks
+        # -----------------------------------
+
+        for hook in self.post_hooks:
+            return_dict.update(hook(systems, return_dict))
+
         # -----------------------------------------
         #   Undo data preprocessing (eval only)
         # -----------------------------------------

src/metatrain/utils/hooks/__init__.py

@@ -0,0 +1,68 @@
+import torch
+
+from metatrain.utils.data import DatasetInfo, TargetInfo
+
+from .global_multipole import GlobalMultipole
+
+
+KNOWN_POST_HOOKS = {
+    "global_multipoles": GlobalMultipole,
+}
+
+
+# For documentation purposes
+PostHooksHypers = dict[str, dict[str, str] | str]
+
+
+# To help models work with hooks.
+def setup_post_hooks(
+    post_hooks_hypers: PostHooksHypers, 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 post_hooks_hypers: The hyperparameters for the post-processing hooks.
+    :param dataset_info: The dataset information.
+    :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()
+    post_hooks = []
+    for hook_name, hook_hypers in post_hooks_hypers.items():
+        if hook_name not in KNOWN_POST_HOOKS:
+            raise ValueError(f"Unknown post-processing hook: {hook_name}")
+
+        # Get outputs of the hook and remove them from the model outputs.
+        if isinstance(hook_hypers, str):
+            san_hook_hypers: dict[str, dict[str, str] | str] = {
+                "outputs": hook_hypers,
+                "inputs": {},
+            }
+        else:
+            san_hook_hypers = hook_hypers
+
+        hook_outputs: str | dict = san_hook_hypers["outputs"]
+        if isinstance(hook_outputs, str):
+            hook_output_names = [hook_outputs]
+        else:
+            hook_output_names = list(hook_outputs.values())
+
+        for output_name in hook_output_names:
+            model_outputs.pop(output_name, None)
+
+        # Get hook and add it to the list of hooks.
+        hook_class = KNOWN_POST_HOOKS[hook_name]
+        hook = hook_class(san_hook_hypers, dataset_info)
+        post_hooks.append(hook)
+
+        # Add the inputs that the hook requests to the model outputs.
+        # TODO: Here we should not add the inputs that will be provided
+        # as extra_data.
+        requested = hook.requested_target_infos()
+        model_outputs.update(requested)
+
+    return post_hooks, model_outputs

src/metatrain/utils/hooks/global_multipole.py

@@ -0,0 +1,141 @@
+import torch
+from metatensor.torch import TensorBlock, TensorMap
+from metatomic.torch import ModelOutput, System
+from typing_extensions import TypedDict
+
+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
+
+
+class HookHypers(TypedDict):
+    """
+    Hyperparameters for the global multipole hook.
+    """
+
+    inputs: str
+
+    outputs: str
+
+
+class GlobalMultipole(torch.nn.Module):
+    """
+    Computes a global multipole from local predictions.
+    """
+
+    def __init__(self, hypers: HookHypers, dataset_info: DatasetInfo):
+        super().__init__()
+
+        self.hypers = hypers
+
+        # Get the information about the output target from the dataset info
+        self.out_name = hypers["outputs"]
+        self.out_target = dataset_info.targets[self.out_name]
+
+        self.degrees = self.out_target.layout.keys["o3_lambda"]
+        self.max_degree = self.degrees.max().item()
+
+        if self.max_degree > 1:
+            raise ValueError(
+                f"Global multipoles hook only supports multipoles up "
+                f"to l=1 for now, but {self.out_name} has max degree "
+                f"{self.max_degree}"
+            )
+
+        # Build the input target that we will request from the model,
+        # which is the local multipoles
+        self._input_name = "mtt::aux::local_multipoles"
+
+        self._input_target_info = get_generic_target_info(
+            self._input_name,
+            {
+                "quantity": "",
+                "unit": "",
+                "type": {
+                    "spherical": {
+                        "irreps": [
+                            {"o3_lambda": i, "o3_sigma": 1}
+                            for i in range(self.max_degree + 1)
+                        ]
+                    }
+                },
+                "num_subtargets": 1,
+                "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_name: self._input_target_info}
+
+    def requested_inputs(self) -> dict[str, ModelOutput]:
+        """
+        Returns the list of requested inputs for the hook.
+
+        :return: A list of requested input names.
+        """
+        return {
+            self._input_name: ModelOutput(
+                quantity="",
+                unit="",
+                sample_kind="atom",
+            )
+        }
+
+    def forward(
+        self, systems: list[System], inputs: dict[str, TensorMap]
+    ) -> dict[str, TensorMap]:
+        """
+        Computes the global multipole from the local predictions.
+        """
+        # Get the concatenated positions of all atoms in the systems,
+        # and reorder the axes to match the spherical harmonics convention
+        # (x, y, z) -> (y, z, x)
+        positions = torch.cat([s.positions for s in systems], dim=0)
+        positions = positions[:, [1, 2, 0]]
+
+        # Get the local predictions for each degree in the multipole expansion
+        input_tmap = inputs[self._input_name]
+        local_values = [
+            input_tmap.block(dict(o3_lambda=ell, o3_sigma=1)).values
+            for ell in range(self.max_degree + 1)
+        ]
+
+        # Compute the local contributions to the global multipole for each degree
+        local_contribs = []
+        for ell in range(self.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]
+                )
+            else:
+                raise ValueError(
+                    "Global multipoles hook only supports multipoles up to l=1 for now"
+                    f", but {self.out_name} has degree {ell}"
+                )
+
+        # Build a tensor map with the requested multipole degrees.
+        local_tmap = TensorMap(
+            keys=self.out_target.layout.keys,
+            blocks=[
+                TensorBlock(
+                    values=local_contribs[degree],
+                    samples=input_tmap.block(0).samples,
+                    components=self.out_target.component_labels[i],
+                    properties=self.out_target.property_labels[i],
+                )
+                for i, degree in enumerate(self.degrees)
+            ],
+        )
+
+        # Return the global multipole by summing over the atoms in the system
+        return {self.out_name: sum_over_atoms(local_tmap)}