Add GAP model

.github/workflows/gap-tests.yml

@@ -0,0 +1,32 @@
+name: GAP tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    # Check all PR
+
+jobs:
+  tests:
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-22.04
+            python-version: "3.11"
+
+    steps:
+    - uses: actions/checkout@v3
+
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+    - run: pip install tox
+
+    - name: run SparseGAP tests
+      run: tox -e gap-tests
+      env:
+        # Use the CPU only version of torch when building/running the code
+        PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu
+

CODEOWNERS

@@ -5,3 +5,4 @@
 **/soap_bpnn            @frostedoyster
 **/alchemical_model     @abmazitov
 **/pet                  @spozdn @abmazitov
+**/gap                  @DavideTisi

docs/src/architectures/gap.rst

@@ -0,0 +1,120 @@
+.. _architecture-sparse-gap:
+
+GAP
+===
+
+This is an implementation of the sparse `Gaussian Approximation Potential
+<GAP_>`_ (GAP) using `Smooth Overlap of Atomic Positions <SOAP_>`_ (SOAP)
+implemented in `rascaline <RASCALINE_>`_.
+
+
+.. _SOAP: https://doi.org/10.1103/PhysRevB.87.184115
+.. _GAP:  https://doi.org/10.1002/qua.24927
+.. _RASCALINE: https://github.com/Luthaf/rascaline
+
+The GAP model in metatensor-models can only train on CPU, but evaluation
+is also supported on GPU.
+
+
+Installation
+------------
+
+To install the package, you can run the following command in the root directory
+of the repository:
+
+.. code-block:: bash
+
+    pip install .[gap]
+
+This will install the package with the GAP dependencies.
+
+
+Architecture Hyperparameters
+----------------------------
+
+:param name: ``experimental.gap``
+
+model
+#####
+soap
+^^^^
+:param cutoff: Spherical cutoff (Å) to use for atomic environments. Default 5.0
+:param max_radial: Number of radial basis function to use. Default 8
+:param max_angular: Number of angular basis function to use also denoted by the  maximum
+    degree of spherical harmonics. Default 6
+:param atomic_gaussian_width: Width of the atom-centered gaussian creating the atomic
+    density. Default 0.3
+:param center_atom_weight: Weight of the central atom contribution to the features. If
+    1.0 the center atom contribution is weighted the same as any other contribution. If
+    0.0 the central atom does not contribute to the features at all. Default 1.0
+:param cutoff_function: cutoff function used to smooth the behavior around the cutoff
+    radius. The supported cutoff function are
+
+    - ``Step``: Step function, 1 if ``r < cutoff`` and 0 if ``r >= cutoff``. This cutoff
+      function takes no additional parameters and can set as in ``.yaml`` file:
+
+      .. code-block:: yaml
+
+        cutoff_function:
+          Step:
+
+    - ``ShiftedCosine`` (Default value): Shifted cosine switching function
+      ``f(r) = 1/2 * (1 + cos(π (r- cutoff + width) / width ))``.
+      This cutoff function takes the ``width``` as
+      additional parameter and can set as in ``options.yaml`` file as:
+
+      .. code-block:: yaml
+
+        cutoff_function:
+          ShiftedCosine:
+            width: 1.0
+
+:param radial_scaling: Radial scaling can be used to reduce the importance of neighbor
+    atoms further away from the center, usually improving the performance of the model.
+    The supported radial scaling functions are
+
+    - ``None``: No radial scaling.
+
+      .. code-block:: yaml
+
+        radial_scaling:
+          None:
+
+    - ``Willatt2018`` (Default value): Use a long-range algebraic decay and
+      smooth behavior at :math:`r
+      \rightarrow 0`: as introduced by :footcite:t:`willatt_feature_2018` as ``f(r) =
+      rate / (rate + (r / scale) ^ exponent)`` This radial scaling function can be set
+      in the ``options.yaml`` file as.
+
+      .. code-block:: yaml
+
+        radial_scaling:
+          Willatt2018:
+            rate: 1.0
+            scale: 2.0
+            exponent: 7.0
+
+.. note::
+
+  Currently, we only support a Gaussian type orbitals (GTO) as radial basis functions
+  and radial integrals.
+
+krr
+^^^^
+:param degree: degree of the polynomial kernel. Default 2
+:param num_sparse_points: number of pseudo points to select
+    (by farthest point sampling). Default 500
+
+training:
+^^^^^^^^^
+:param regularizer: value of the energy regularizer. Default 0.001
+:param regularizer_forces: value of the forces regularizer. Default null
+
+
+Default Hyperparameters
+-----------------------
+The default hyperparameters for the GAP model are:
+
+.. literalinclude:: ../../../src/metatensor/models/experimental/gap/default-hypers.yaml
+   :language: yaml
+

pyproject.toml

@@ -63,6 +63,12 @@ alchemical-model = [
 pet = [
   "pet @ git+https://github.com/spozdn/pet.git@9f6119d",
 ]
+gap = [
+    "rascaline-torch @ git+https://github.com/luthaf/rascaline@5348132#subdirectory=python/rascaline-torch",
+    "skmatter",
+    "metatensor-learn",
+    "scipy",
+]
 
 [tool.setuptools.packages.find]
 where = ["src"]

src/metatensor/models/cli/eval.py

@@ -1,4 +1,5 @@
 import argparse
+import itertools
 import logging
 from pathlib import Path
 from typing import Dict, List, Optional, Union
@@ -158,7 +159,7 @@ def _eval_targets(
         get_system_with_neighbor_lists(system, model.requested_neighbor_lists())
 
     # Infer the device from the model
-    device = next(model.parameters()).device
+    device = next(itertools.chain(model.parameters(), model.buffers())).device
 
     # Create a dataloader
     dataloader = torch.utils.data.DataLoader(

src/metatensor/models/experimental/gap/__init__.py

@@ -0,0 +1,14 @@
+from .model import GAP
+from .trainer import Trainer
+
+__model__ = GAP
+__trainer__ = Trainer
+
+__authors__ = [
+    ("Alexander Goscinski <alex.goscinski@posteo.de>", "@agosckinski"),
+    ("Davide Tisi <davide.tisi@epfl.ch>", "@DavideTisi"),
+]
+
+__maintainers__ = [
+    ("Davide Tisi <davide.tisi@epfl.ch>", "@DavideTisi"),
+]

src/metatensor/models/experimental/gap/default-hypers.yaml

@@ -0,0 +1,28 @@
+# default hyperparameters for the SparseGAP model
+name: gap
+
+model:
+  soap:
+    cutoff: 5.0
+    max_radial: 8
+    max_angular: 6
+    atomic_gaussian_width: 0.3
+    radial_basis:
+      Gto: {}
+    center_atom_weight: 1.0
+    cutoff_function:
+      ShiftedCosine:
+        width: 1.0
+    radial_scaling:
+      Willatt2018:
+        rate: 1.0
+        scale: 2.0
+        exponent: 7.0
+
+  krr:
+    degree: 2
+    num_sparse_points: 500
+
+training:
+  regularizer: 0.001
+  regularizer_forces: null