Models
======

`mlcg.nn` currently implements the SchNet graph neural network, as well as several utility classes for computing distance expansions and cutoffs. The `nn` subpackage also contains several useful classes for extracting other properties from energy predictions or aggregating the predictions from several different model types.

SchNet Utilities
----------------

These classes are used to define a SchNet graph neural network. For "typical" SchNet models, users
may find the class `StandardSchNet` to be helpful in getting started quickly.

.. autoclass:: mlcg.nn.StandardSchNet
.. autoclass:: mlcg.nn.schnet.SchNet
.. autoclass:: mlcg.nn.schnet.InteractionBlock
.. autoclass:: mlcg.nn.schnet.CFConv

MACE Utilities
--------------

`mlcg` provides wrapper utitilies for [MACE]_ models, the base implementation of which is available `here <https://github.com/felixmusil/mace.git>`_ in the `develop` branch.

.. autoclass:: mlcg.nn.mace_interface.MACEInterface

Radial Basis Functions
----------------------

Sets of radial basis functions are used to expand the distances (or other molecular features) between atoms on a fixed-sized vector. For instance, this is the main transformation of the distances in the `SchNet` model.

.. autoclass:: mlcg.nn.radial_basis.GaussianBasis
.. autoclass:: mlcg.nn.radial_basis.ExpNormalBasis
.. autoclass:: mlcg.nn.radial_basis.RIGTOBasis
.. autoclass:: mlcg.nn.radial_basis.SpacedExpBasis
.. autoclass:: mlcg.nn.angular_basis.SphericalHarmonics

Cutoff Functions
----------------

Cutoff functions are used to enforce the smoothness of the models w.r.t. neighbor insertion/removal from an atomic environment. Some are also used to damp the signal from a neighbor's displacement that is "far" from the central atom, e.g. `CosineCutoff`. Cutoff functions are also used in the construction of radial basis functions.

.. autoclass:: mlcg.nn.cutoff.IdentityCutoff
.. autoclass:: mlcg.nn.cutoff.CosineCutoff
.. autoclass:: mlcg.nn.cutoff.ShiftedCosineCutoff

Model Building Utilities
------------------------

These classes are used to build more complicated models.

.. autoclass:: mlcg.nn.gradients.GradientsOut
.. autoclass:: mlcg.nn.gradients.SumOut

Loss Functions
--------------

These classes define loss functions for model optimization, as well as generalized losses that combine several losses of different types.

.. autoclass:: mlcg.nn.losses.Loss
.. autoclass:: mlcg.nn.losses.ForceRMSE
.. autoclass:: mlcg.nn.losses.ForceMSE