Add hooks for indirect learning.

.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/

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``.

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

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/*",
 ]
 
 

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 <https://jinja.palletsprojects.com/en/stable/templates>`_ , 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.<hook_name>.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")

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/*

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.

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 %}

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.

docs/src/hooks/templates/references.rst

@@ -0,0 +1,6 @@
+.. _hook-{{hook}}_references:
+
+References
+----------
+
+.. footbibliography::

docs/src/index.rst

@@ -21,6 +21,7 @@
    installation
    getting-started/index
    architectures/index
+   hooks/index
    generated_examples/index
    concepts/index
    faq

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}
+}
+

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 {},
     )
 
     ###########################