.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/programmatic/use_architectures_outside/use_outside.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_examples_programmatic_use_architectures_outside_use_outside.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_programmatic_use_architectures_outside_use_outside.py:


Using metatrain architectures outside of metatrain
==================================================

This tutorial demonstrates how to use one of metatrain's implemented architectures
outside of metatrain. This will be done by taking internal representations of a
NanoPET model (as an example) and using them inside a user-defined torch ``Module``.

Only architectures which can output internal representations ("features" output) can
be used in this way.

.. GENERATED FROM PYTHON SOURCE LINES 15-28

.. code-block:: Python


    import torch
    from metatensor.torch.atomistic import ModelOutput

    from metatrain.experimental.nanopet import NanoPET
    from metatrain.utils.architectures import get_default_hypers
    from metatrain.utils.data import DatasetInfo, read_systems
    from metatrain.utils.neighbor_lists import (
        get_requested_neighbor_lists,
        get_system_with_neighbor_lists,
    )









.. GENERATED FROM PYTHON SOURCE LINES 29-31

Read some sample systems. Metatrain always reads systems in float64, while torch
uses float32 by default. We will convert the systems to float32.

.. GENERATED FROM PYTHON SOURCE LINES 32-37

.. code-block:: Python


    systems = read_systems("qm9_reduced_100.xyz")
    systems = [s.to(torch.float32) for s in systems]









.. GENERATED FROM PYTHON SOURCE LINES 38-41

Define the custom model using the NanoPET architecture as a building block.
The dummy architecture here adds a linear layer and a tanh activation function
on top of the NanoPET model.

.. GENERATED FROM PYTHON SOURCE LINES 42-68

.. code-block:: Python



    class NanoPETWithTanh(torch.nn.Module):
        def __init__(self):
            super().__init__()
            self.nanopet = NanoPET(
                get_default_hypers("experimental.nanopet")["model"],
                DatasetInfo(
                    length_unit="angstrom",
                    atomic_types=[1, 6, 7, 8, 9],
                    targets={},
                ),
            )
            self.linear = torch.nn.Linear(128, 1)
            self.tanh = torch.nn.Tanh()

        def forward(self, systems):
            model_outputs = self.nanopet(
                systems,
                {"features": ModelOutput()},
                # ModelOutput(per_atom=True) would give per-atom features
            )
            features = model_outputs["features"].block().values
            return self.tanh(self.linear(features))









.. GENERATED FROM PYTHON SOURCE LINES 69-71

Now we can train the custom model. Here is one training step executed with
some random targets.

.. GENERATED FROM PYTHON SOURCE LINES 72-91

.. code-block:: Python

    my_targets = torch.randn(100, 1)

    # instantiate the model
    model = NanoPETWithTanh()

    # all metatrain models require neighbor lists to be present in the input systems
    systems = [
        get_system_with_neighbor_lists(sys, get_requested_neighbor_lists(model))
        for sys in systems
    ]

    # define an optimizer
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)

    # this is one training step
    predictions = model(systems)
    loss = torch.nn.functional.mse_loss(predictions, my_targets)
    loss.backward()
    optimizer.step()








.. _sphx_glr_download_examples_programmatic_use_architectures_outside_use_outside.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: use_outside.ipynb <use_outside.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: use_outside.py <use_outside.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: use_outside.zip <use_outside.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_