From 7334795b978fab666f41dab98b2c26ba8210550a Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Wed, 29 May 2024 17:55:21 +0200 Subject: [PATCH 01/16] Fix typo --- docs/src/conf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/src/conf.py b/docs/src/conf.py index e9373df1c..c7e5abacb 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -9,7 +9,7 @@ # When importing metatensor-torch, this will change the definition of the classes -# to include the documentation +# to include in the documentation os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1" os.environ["RASCALINE_IMPORT_FOR_SPHINX"] = "1" From e7b293fbf2e5eb59e707c84823b9ca3de2ea39ec Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 30 May 2024 13:57:11 +0200 Subject: [PATCH 02/16] Prototype --- docs/src/conf.py | 2 +- docs/src/dev-docs/new-architecture.rst | 4 ++-- src/metatensor/models/__main__.py | 15 --------------- src/metatensor/models/cli/eval.py | 2 +- src/metatensor/models/cli/export.py | 2 +- src/metatensor/models/cli/train.py | 2 +- src/metatensor/models/utils/evaluate_model.py | 3 ++- 7 files changed, 8 insertions(+), 22 deletions(-) diff --git a/docs/src/conf.py b/docs/src/conf.py index c7e5abacb..e9373df1c 100644 --- a/docs/src/conf.py +++ b/docs/src/conf.py @@ -9,7 +9,7 @@ # When importing metatensor-torch, this will change the definition of the classes -# to include in the documentation +# to include the documentation os.environ["METATENSOR_IMPORT_FOR_SPHINX"] = "1" os.environ["RASCALINE_IMPORT_FOR_SPHINX"] = "1" diff --git a/docs/src/dev-docs/new-architecture.rst b/docs/src/dev-docs/new-architecture.rst index 2921a59a1..7d744d294 100644 --- a/docs/src/dev-docs/new-architecture.rst +++ b/docs/src/dev-docs/new-architecture.rst @@ -32,10 +32,10 @@ to these lines checkpoint_dir="path", ) - model.save_checkpoint("final.ckpt") + model.save_checkpoint("model.ckpt") mts_atomistic_model = model.export() - mts_atomistic_model.export("path", collect_extensions="extensions-dir/") + mts_atomistic_model.export("model.pt", collect_extensions="extensions/") In order to follow this, a new architectures has two define two classes diff --git a/src/metatensor/models/__main__.py b/src/metatensor/models/__main__.py index 1c2157ae0..a4bc25c02 100644 --- a/src/metatensor/models/__main__.py +++ b/src/metatensor/models/__main__.py @@ -20,21 +20,6 @@ from .utils.logging import setup_logging -# This import is necessary to avoid errors when loading an -# exported alchemical model, which depends on sphericart-torch. -# TODO: Remove this when https://github.com/lab-cosmo/metatensor/issues/512 -# is ready -try: - import sphericart.torch # noqa: F401 -except ImportError: - pass - -try: - import rascaline.torch # noqa: F401 -except ImportError: - pass - - logger = logging.getLogger(__name__) diff --git a/src/metatensor/models/cli/eval.py b/src/metatensor/models/cli/eval.py index b6f5ad5e8..7633d554e 100644 --- a/src/metatensor/models/cli/eval.py +++ b/src/metatensor/models/cli/eval.py @@ -58,7 +58,7 @@ def _add_eval_model_parser(subparser: argparse._SubParsersAction) -> None: ) parser.add_argument( "-e", - "--extdir", + "--extensions-dir", type=str, required=False, dest="extensions_directory", diff --git a/src/metatensor/models/cli/export.py b/src/metatensor/models/cli/export.py index 5c452b425..121cc5a2c 100644 --- a/src/metatensor/models/cli/export.py +++ b/src/metatensor/models/cli/export.py @@ -63,4 +63,4 @@ def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> torch.jit.save(model, path) else: mts_atomistic_model = model.export() - mts_atomistic_model.export(path) + mts_atomistic_model.export(path, collect_extensions="extensions") diff --git a/src/metatensor/models/cli/train.py b/src/metatensor/models/cli/train.py index 57683dfc6..2e55f79e3 100644 --- a/src/metatensor/models/cli/train.py +++ b/src/metatensor/models/cli/train.py @@ -389,7 +389,7 @@ def train_model( raise ArchitectureError(e) mts_atomistic_model = model.export() - mts_atomistic_model.export(str(output_checked)) + mts_atomistic_model.export(str(output_checked), collect_extensions="extensions") ########################### # EVALUATE FINAL MODEL #### diff --git a/src/metatensor/models/utils/evaluate_model.py b/src/metatensor/models/utils/evaluate_model.py index 22e5689b1..aaf5e8fe2 100644 --- a/src/metatensor/models/utils/evaluate_model.py +++ b/src/metatensor/models/utils/evaluate_model.py @@ -22,7 +22,8 @@ "ignore", category=UserWarning, message="neighbor", -) # TODO: this is not filtering out the warning for some reason +) # TODO: this is not filtering out the warning for some reason, therefore: +warnings.filterwarnings("ignore") # ignore all warnings if not in debug mode def evaluate_model( From 0555c971562cacce42c32e5e469803c6ebe25216 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 30 May 2024 14:00:15 +0200 Subject: [PATCH 03/16] Add GAP implementation --- .github/workflows/gap-tests.yml | 36 + docs/src/architectures/gap.rst | 109 ++ pyproject.toml | 6 + src/metatensor/models/cli/eval.py | 3 +- .../models/experimental/gap/__init__.py | 14 + .../experimental/gap/default-hypers.yaml | 28 + .../models/experimental/gap/model.py | 1248 +++++++++++++++++ .../models/experimental/gap/tests/__init__.py | 11 + .../experimental/gap/tests/test_errors.py | 79 ++ .../experimental/gap/tests/test_regression.py | 186 +++ .../gap/tests/test_torchscript.py | 84 ++ .../models/experimental/gap/trainer.py | 146 ++ .../models/experimental/soap_bpnn/trainer.py | 2 +- tox.ini | 10 + 14 files changed, 1960 insertions(+), 2 deletions(-) create mode 100644 .github/workflows/gap-tests.yml create mode 100644 docs/src/architectures/gap.rst create mode 100644 src/metatensor/models/experimental/gap/__init__.py create mode 100644 src/metatensor/models/experimental/gap/default-hypers.yaml create mode 100644 src/metatensor/models/experimental/gap/model.py create mode 100644 src/metatensor/models/experimental/gap/tests/__init__.py create mode 100644 src/metatensor/models/experimental/gap/tests/test_errors.py create mode 100644 src/metatensor/models/experimental/gap/tests/test_regression.py create mode 100644 src/metatensor/models/experimental/gap/tests/test_torchscript.py create mode 100644 src/metatensor/models/experimental/gap/trainer.py diff --git a/.github/workflows/gap-tests.yml b/.github/workflows/gap-tests.yml new file mode 100644 index 000000000..a097c2a4a --- /dev/null +++ b/.github/workflows/gap-tests.yml @@ -0,0 +1,36 @@ +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 + + - name: Upload codecoverage + uses: codecov/codecov-action@v3 + with: + files: ./tests/coverage.xml diff --git a/docs/src/architectures/gap.rst b/docs/src/architectures/gap.rst new file mode 100644 index 000000000..631a5c7cf --- /dev/null +++ b/docs/src/architectures/gap.rst @@ -0,0 +1,109 @@ +.. _architecture-sparse-gap: + +GAP +=== + +This is an implementation of the sparse `Gaussian Approximation Potential +`_ (GAP) using `Smooth Overlap of Atomic Positions `_ (SOAP) +implemented in `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. + + +Hyperparameters +--------------- + +:param name: ``experimental.gap`` + +model +##### +soap +^^^^ +:param cutoff: Spherical cutoff (Å) to use for atomic environments +:param max_radial: Number of radial basis function to use +:param max_angular: Number of angular basis function to use also denoted by the maximum + degree of spherical harmonics +:param atomic_gaussian_width: Width of the atom-centered gaussian creating the atomic + density +: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. +: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``: 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`` 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 +:param num_sparse_points: number of pseudo points to select (by farthest point sampling) + +training: +^^^^^^^^^ +:param regularizer: value of the energy regularizer +:param regularizer_forces: value of the forces regularizer + diff --git a/pyproject.toml b/pyproject.toml index c75eb3d95..81f6b40e5 100644 --- a/pyproject.toml +++ b/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"] diff --git a/src/metatensor/models/cli/eval.py b/src/metatensor/models/cli/eval.py index 7633d554e..4732c11de 100644 --- a/src/metatensor/models/cli/eval.py +++ b/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 @@ -157,7 +158,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( diff --git a/src/metatensor/models/experimental/gap/__init__.py b/src/metatensor/models/experimental/gap/__init__.py new file mode 100644 index 000000000..93f81b611 --- /dev/null +++ b/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 ", "@agosckinski"), + ("Davide Tisi ", "@DavideTisi"), +] + +__maintainers__ = [ + ("Davide Tisi ", "@DavideTisi"), +] diff --git a/src/metatensor/models/experimental/gap/default-hypers.yaml b/src/metatensor/models/experimental/gap/default-hypers.yaml new file mode 100644 index 000000000..ad275a0c0 --- /dev/null +++ b/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 # degree of the polynomial kernel + num_sparse_points: 500 # number of pseudo points to select, farthest point sampling is used + +training: + regularizer: 0.001 + regularizer_forces: null diff --git a/src/metatensor/models/experimental/gap/model.py b/src/metatensor/models/experimental/gap/model.py new file mode 100644 index 000000000..b1d8f46fd --- /dev/null +++ b/src/metatensor/models/experimental/gap/model.py @@ -0,0 +1,1248 @@ +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Type, Union + +import metatensor.torch +import numpy as np +import rascaline +import rascaline.torch +import scipy +import skmatter +import torch +from metatensor.torch import Labels as TorchLabels +from metatensor.torch import TensorBlock as TorchTensorBlock +from metatensor.torch import TensorMap as TorchTensorMap +from metatensor.torch.atomistic import ( + MetatensorAtomisticModel, + ModelCapabilities, + ModelOutput, + System, +) +from skmatter._selection import _FPS + +from metatensor import Labels, TensorBlock, TensorMap +from metatensor.models.utils.data.dataset import DatasetInfo + +from ...utils.export import export + + +class GAP(torch.nn.Module): + __supported_devices__ = ["cpu"] + __supported_dtypes__ = [torch.float64] + + def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None: + super().__init__() + + if len(dataset_info.targets) > 1: + raise NotImplementedError("GAP only supports a single output") + + # Check capabilities + for target in dataset_info.targets.values(): + if target.quantity != "energy": + raise ValueError( + "GAP only supports energy-like outputs, " + f"but a {target.quantity} was provided" + ) + if target.per_atom: + raise ValueError( + "GAP only supports per-structure outputs, " + "but a per-atom output was provided" + ) + target_name = next(iter(dataset_info.targets.keys())) + if dataset_info.targets[target_name].quantity != "energy": + raise ValueError("GAP only supports energies as target") + if dataset_info.targets[target_name].per_atom: + raise ValueError("GAP does not support per-atom energies") + + self.dataset_info = dataset_info + + self.outputs = { + key: ModelOutput( + quantity=value.quantity, + unit=value.unit, + per_atom=False, + ) + for key, value in dataset_info.targets.items() + } + + self.all_species = dataset_info.atomic_types + self.hypers = model_hypers + + # creates a composition weight tensor that can be directly indexed by species, + # this can be left as a tensor of zero or set from the outside using + # set_composition_weights (recommended for better accuracy) + n_outputs = len(dataset_info.targets) + self.register_buffer( + "composition_weights", torch.zeros((n_outputs, max(self.all_species) + 1)) + ) + # buffers cannot be indexed by strings (torchscript), so we create a single + # tensor for all output. Due to this, we need to slice the tensor when we use + # it and use the output name to select the correct slice via a dictionary + self.output_to_index = { + output_name: i for i, output_name in enumerate(dataset_info.targets.keys()) + } + + self.register_buffer( + "kernel_weights", torch.zeros((model_hypers["krr"]["num_sparse_points"])) + ) + self._soap_torch_calculator = rascaline.torch.SoapPowerSpectrum( + **model_hypers["soap"] + ) + self._soap_calculator = rascaline.SoapPowerSpectrum(**model_hypers["soap"]) + + kernel_kwargs = { + "degree": model_hypers["krr"]["degree"], + "aggregate_names": ["atom", "center_type"], + } + self._subset_of_regressors = SubsetOfRegressors( + kernel_type=model_hypers["krr"].get("kernel", "polynomial"), + kernel_kwargs=kernel_kwargs, + ) + + self._sampler = FPS(n_to_select=model_hypers["krr"]["num_sparse_points"]) + + # set it do dummy keys, these are properly set during training + self._keys = TorchLabels.empty("_") + + dummy_weights = TorchTensorMap( + TorchLabels(["_"], torch.tensor([[0]])), + [metatensor.torch.block_from_array(torch.empty(1, 1))], + ) + dummy_X_pseudo = TorchTensorMap( + TorchLabels(["_"], torch.tensor([[0]])), + [metatensor.torch.block_from_array(torch.empty(1, 1))], + ) + self._subset_of_regressors_torch = TorchSubsetofRegressors( + dummy_weights, + dummy_X_pseudo, + kernel_kwargs={ + "aggregate_names": ["atom", "center_type"], + }, + ) + self._species_labels: TorchLabels = TorchLabels.empty("_") + + def restart(self, dataset_info: DatasetInfo) -> "GAP": + raise ValueError("GAP does not allow restarting training") + + def forward( + self, + systems: List[System], + outputs: Dict[str, ModelOutput], + selected_atoms: Optional[TorchLabels] = None, + ) -> Dict[str, TorchTensorMap]: + soap_features = self._soap_torch_calculator( + systems, selected_samples=selected_atoms + ) + # move keys and species labels to device + self._keys = self._keys.to(systems[0].device) + self._species_labels = self._species_labels.to(systems[0].device) + + new_blocks: List[TorchTensorBlock] = [] + # HACK: to add a block of zeros if there are missing species + # which were present at training time + # (with samples "system", "atom" = 0, 0) + # given the values are all zeros, it does not introduce an error + dummyblock: TorchTensorBlock = TorchTensorBlock( + values=torch.zeros( + (1, len(soap_features[0].properties)), + dtype=systems[0].positions.dtype, + device=systems[0].device, + ), + samples=TorchLabels( + ["system", "atom"], + torch.tensor([[0, 0]], dtype=torch.int, device=systems[0].device), + ), + properties=soap_features[0].properties, + components=[], + ) + if len(soap_features[0].gradients_list()) > 0: + for idx, grad in enumerate(soap_features[0].gradients_list()): + dummyblock_grad: TorchTensorBlock = TorchTensorBlock( + values=torch.zeros( + ( + 1, + soap_features[0].gradient(grad).values.shape[1 + idx], + len(soap_features[0].gradient(grad).properties), + ), + dtype=systems[0].positions.dtype, + device=systems[0].device, + ), + samples=TorchLabels( + ["sample", "system", "atom"], + torch.tensor( + [[0, 0, 0]], dtype=torch.int, device=systems[0].device + ), + ), + components=soap_features[0].gradient(grad).components, + properties=soap_features[0].gradient(grad).properties, + ) + dummyblock.add_gradient(grad, dummyblock_grad) + + for idx_key in range(len(self._species_labels)): + key = self._species_labels.entry(idx_key) + if soap_features.keys.position(key) is not None: + new_blocks.append(soap_features.block(key)) + else: + new_blocks.append(dummyblock) + soap_features = TorchTensorMap(keys=self._species_labels, blocks=new_blocks) + soap_features = soap_features.keys_to_samples("center_type") + # here, we move to properties to use metatensor operations to aggregate + # later on. Perhaps we could retain the sparsity all the way to the kernels + # of the soap features with a lot more implementation effort + soap_features = soap_features.keys_to_properties( + ["neighbor_1_type", "neighbor_2_type"] + ) + soap_features = TorchTensorMap(self._keys, soap_features.blocks()) + output_key = list(outputs.keys())[0] + energies = self._subset_of_regressors_torch(soap_features) + out_tensor = self.apply_composition_weights(systems, energies) + return {output_key: out_tensor} + + def save_checkpoint(self, path: Union[str, Path]): + # GAP will not save checkpoints, as it does not allow + # restarting training + return + + @classmethod + def load_checkpoint(cls, path: Union[str, Path]) -> "GAP": + raise ValueError("GAP does not allow restarting training") + + def export(self) -> MetatensorAtomisticModel: + capabilities = ModelCapabilities( + outputs=self.outputs, + atomic_types=self.dataset_info.atomic_types, + interaction_range=self.hypers["soap"]["cutoff"], + length_unit=self.dataset_info.length_unit, + supported_devices=["cuda", "cpu"], + dtype="float64", + ) + + # we export a torch scriptable regressor TorchSubsetofRegressors + # that is used in the forward path + self._subset_of_regressors_torch = ( + self._subset_of_regressors.export_torch_script_model() + ) + + return export(model=self, model_capabilities=capabilities) + + def set_composition_weights( + self, + output_name: str, + input_composition_weights: torch.Tensor, + species: List[int], + ) -> None: + """Set the composition weights for a given output.""" + # all species that are not present retain their weight of zero + self.composition_weights[self.output_to_index[output_name]][ # type: ignore + species + ] = input_composition_weights.to( + dtype=self.composition_weights.dtype, # type: ignore + device=self.composition_weights.device, # type: ignore + ) + + def apply_composition_weights( + self, systems: List[System], energies: TorchTensorMap + ) -> TorchTensorMap: + """Apply the composition weights to the energies.""" + new_blocks: List[TorchTensorBlock] = [] + for block in energies.blocks(): + atomic_species = [system.types for system in systems] + new_values = block.values + for i in range(len(new_values)): + for s in atomic_species[i]: + new_values[i] += self.composition_weights[0, s] + new_blocks.append( + TorchTensorBlock( + values=new_values, + samples=block.samples, + components=block.components, + properties=block.properties, + ) + ) + + return TorchTensorMap(energies.keys, new_blocks) + + +######################################################################################## +# All classes and methods will be moved to metatensor-operations and metatensor-learn] # +######################################################################################## + + +class _SorKernelSolver: + """ + A few quick implementation notes, docs to be done. + + This is meant to solve the subset of regressors (SoR) problem:: + + .. math:: + + w = (KNM.T@KNM + reg*KMM)^-1 @ KNM.T@y + + The inverse needs to be stabilized with application of a numerical jitter, + that is expressed as a fraction of the largest eigenvalue of KMM + + :param KMM: + KNM matrix + :param regularizer: + regularizer + :param jitter: + numerical jitter to stabilize fit + :param solver: + Method to solve the sparse KRR equations. + + * RKHS-QR: TBD + * RKHS: Compute first the reproducing kernel features by diagonalizing K_MM and + computing `P_NM = K_NM @ U_MM @ Lam_MM^(-1.2)` and then solves the linear + problem for those (which is usually better conditioned):: + + (P_NM.T@P_NM + 1)^(-1) P_NM.T@Y + + * QR: TBD + * solve: Uses `scipy.linalg.solve` for the normal equations:: + + (K_NM.T@K_NM + K_MM)^(-1) K_NM.T@Y + + * lstsq: require rcond value. Use `numpy.linalg.solve(rcond=rcond)` for the + normal equations:: + + (K_NM.T@K_NM + K_MM)^(-1) K_NM.T@Y + + Reference + --------- + Foster, L., Waagen, A., Aijaz, N., Hurley, M., Luis, A., Rinsky, J., ... & + Srivastava, A. (2009). Stable and Efficient Gaussian Process Calculations. Journal + of Machine Learning Research, 10(4). + """ + + def __init__( + self, + KMM: np.ndarray, + regularizer: float = 1.0, + jitter: float = 0.0, + solver: str = "RKHS", + relative_jitter: bool = True, + ): + self.solver = solver + self.KMM = KMM + self.relative_jitter = relative_jitter + + self._nM = len(KMM) + if self.solver == "RKHS" or self.solver == "RKHS-QR": + self._vk, self._Uk = scipy.linalg.eigh(KMM) + self._vk = self._vk[::-1] + self._Uk = self._Uk[:, ::-1] + elif self.solver == "QR" or self.solver == "solve" or self.solver == "lstsq": + # gets maximum eigenvalue of KMM to scale the numerical jitter + self._KMM_maxeva = scipy.sparse.linalg.eigsh( + KMM, k=1, return_eigenvectors=False + )[0] + else: + raise ValueError( + f"Solver {solver} not supported. Possible values " + "are 'RKHS', 'RKHS-QR', 'QR', 'solve' or lstsq." + ) + if relative_jitter: + if self.solver == "RKHS" or self.solver == "RKHS-QR": + self._jitter_scale = self._vk[0] + elif ( + self.solver == "QR" or self.solver == "solve" or self.solver == "lstsq" + ): + self._jitter_scale = self._KMM_maxeva + else: + self._jitter_scale = 1.0 + self.set_regularizers(regularizer, jitter) + + def set_regularizers(self, regularizer=1.0, jitter=0.0): + self.regularizer = regularizer + self.jitter = jitter + if self.solver == "RKHS" or self.solver == "RKHS-QR": + self._nM = len(np.where(self._vk > self.jitter * self._jitter_scale)[0]) + self._PKPhi = self._Uk[:, : self._nM] * 1 / np.sqrt(self._vk[: self._nM]) + elif self.solver == "QR": + self._VMM = scipy.linalg.cholesky( + self.regularizer * self.KMM + + np.eye(self._nM) * self._jitter_scale * self.jitter + ) + self._Cov = np.zeros((self._nM, self._nM)) + self._KY = None + + def fit(self, KNM, Y, rcond=None): + if len(Y.shape) == 1: + Y = Y[:, np.newaxis] + if self.solver == "RKHS": + Phi = KNM @ self._PKPhi + self._weights = self._PKPhi @ scipy.linalg.solve( + Phi.T @ Phi + np.eye(self._nM) * self.regularizer, + Phi.T @ Y, + assume_a="pos", + ) + elif self.solver == "RKHS-QR": + A = np.vstack( + [KNM @ self._PKPhi, np.sqrt(self.regularizer) * np.eye(self._nM)] + ) + Q, R = np.linalg.qr(A) + self._weights = self._PKPhi @ scipy.linalg.solve_triangular( + R, Q.T @ np.vstack([Y, np.zeros((self._nM, Y.shape[1]))]) + ) + elif self.solver == "QR": + A = np.vstack([KNM, self._VMM]) + Q, R = np.linalg.qr(A) + self._weights = scipy.linalg.solve_triangular( + R, Q.T @ np.vstack([Y, np.zeros((KNM.shape[1], Y.shape[1]))]) + ) + elif self.solver == "solve": + self._weights = scipy.linalg.solve( + KNM.T @ KNM + + self.regularizer * self.KMM + + np.eye(self._nM) * self.jitter * self._jitter_scale, + KNM.T @ Y, + assume_a="pos", + ) + elif self.solver == "lstsq": + self._weights = np.linalg.lstsq( + KNM.T @ KNM + + self.regularizer * self.KMM + + np.eye(self._nM) * self.jitter * self._jitter_scale, + KNM.T @ Y, + rcond=rcond, + )[0] + else: + ValueError("solver not implemented") + + def partial_fit(self, KNM, Y, accumulate_only=False, rcond=None): + if len(Y) > 0: + # only accumulate if we are passing data + if len(Y.shape) == 1: + Y = Y[:, np.newaxis] + if self.solver == "RKHS": + Phi = KNM @ self._PKPhi + elif self.solver == "solve" or self.solver == "lstsq": + Phi = KNM + else: + raise ValueError( + "Partial fit can only be realized with " + "solver = 'RKHS' or 'solve'" + ) + if self._KY is None: + self._KY = np.zeros((self._nM, Y.shape[1])) + + self._Cov += Phi.T @ Phi + self._KY += Phi.T @ Y + + # do actual fit if called with empty array or if asked + if len(Y) == 0 or (not accumulate_only): + if self.solver == "RKHS": + self._weights = self._PKPhi @ scipy.linalg.solve( + self._Cov + np.eye(self._nM) * self.regularizer, + self._KY, + assume_a="pos", + ) + elif self.solver == "solve": + self._weights = scipy.linalg.solve( + self._Cov + + self.regularizer * self.KMM + + np.eye(self.KMM.shape[0]) * self.jitter * self._jitter_scale, + self._KY, + assume_a="pos", + ) + elif self.solver == "lstsq": + self._weights = np.linalg.lstsq( + self._Cov + + self.regularizer * self.KMM + + np.eye(self.KMM.shape[0]) * self.jitter * self._jitter_scale, + self._KY, + rcond=rcond, + )[0] + + @property + def weights(self): + return self._weights + + def predict(self, KTM): + return KTM @ self._weights + + +class AggregateKernel(torch.nn.Module): + """ + A kernel that aggregates values in a kernel over :param aggregate_names: using + a aggregaten function given by :param aggregate_type: + + :param aggregate_names: + + :param aggregate_type: + """ + + def __init__( + self, + aggregate_names: Union[str, List[str]], + aggregate_type: str = "sum", + structurewise_aggregate: bool = False, + ): + super().__init__() + valid_aggregate_types = ["sum", "mean"] + if aggregate_type not in valid_aggregate_types: + raise ValueError( + f"Given aggregate_type {aggregate_type!r} but only " + f"{aggregate_type!r} are supported." + ) + if structurewise_aggregate: + raise NotImplementedError( + "structurewise aggregation has not been implemented." + ) + + self._aggregate_names = aggregate_names + self._aggregate_type = aggregate_type + self._structurewise_aggregate = structurewise_aggregate + + def aggregate_features(self, tensor: TensorMap) -> TensorMap: + if self._aggregate_type == "sum": + return metatensor.sum_over_samples( + tensor, sample_names=self._aggregate_names + ) + elif self._aggregate_type == "mean": + return metatensor.mean_over_samples( + tensor, sample_names=self._aggregate_names + ) + else: + raise NotImplementedError( + f"aggregate_type {self._aggregate_type!r} has not been implemented." + ) + + def aggregate_kernel( + self, kernel: TensorMap, are_pseudo_points: Tuple[bool, bool] = (False, False) + ) -> TensorMap: + if self._aggregate_type == "sum": + if not are_pseudo_points[0]: + kernel = metatensor.sum_over_samples(kernel, self._aggregate_names) + if not are_pseudo_points[1]: + raise NotImplementedError( + "properties dimenson cannot be aggregated for the moment" + ) + return kernel + elif self._aggregate_type == "mean": + if not are_pseudo_points[0]: + kernel = metatensor.mean_over_samples(kernel, self._aggregate_names) + if not are_pseudo_points[1]: + raise NotImplementedError( + "properties dimenson cannot be aggregated for the moment" + ) + return kernel + else: + raise NotImplementedError( + f"aggregate_type {self._aggregate_type!r} has not been implemented." + ) + + def forward( + self, + tensor1: TensorMap, + tensor2: TensorMap, + are_pseudo_points: Tuple[bool, bool] = (False, False), + ) -> TensorMap: + return self.aggregate_kernel( + self.compute_kernel(tensor1, tensor2), are_pseudo_points + ) + + def compute_kernel(self, tensor1: TensorMap, tensor2: TensorMap) -> TensorMap: + raise NotImplementedError("compute_kernel needs to be implemented.") + + +class AggregateLinear(AggregateKernel): + def __init__( + self, + aggregate_names: Union[str, List[str]], + aggregate_type: str = "sum", + structurewise_aggregate: bool = False, + ): + super().__init__(aggregate_names, aggregate_type, structurewise_aggregate) + + def forward( + self, + tensor1: TensorMap, + tensor2: TensorMap, + are_pseudo_points: Tuple[bool, bool] = (False, False), + ) -> TensorMap: + # we overwrite default behavior because for linear kernels we can do it more + # memory efficient + if not are_pseudo_points[0]: + tensor1 = self.aggregate_features(tensor1) + if not are_pseudo_points[1]: + tensor2 = self.aggregate_features(tensor2) + return self.compute_kernel(tensor1, tensor2) + + def compute_kernel(self, tensor1: TensorMap, tensor2: TensorMap) -> TensorMap: + return metatensor.dot(tensor1, tensor2) + + +class AggregatePolynomial(AggregateKernel): + def __init__( + self, + aggregate_names: Union[str, List[str]], + aggregate_type: str = "sum", + structurewise_aggregate: bool = False, + degree: int = 2, + ): + super().__init__(aggregate_names, aggregate_type, structurewise_aggregate) + self._degree = degree + + def compute_kernel(self, tensor1: TensorMap, tensor2: TensorMap): + return metatensor.pow(metatensor.dot(tensor1, tensor2), self._degree) + + +class TorchAggregateKernel(torch.nn.Module): + """ + A kernel that aggregates values in a kernel over :param aggregate_names: using + a aggregaten function given by :param aggregate_type: + + :param aggregate_names: + + :param aggregate_type: + """ + + def __init__( + self, + aggregate_names: Union[str, List[str]], + aggregate_type: str = "sum", + structurewise_aggregate: bool = False, + ): + super().__init__() + valid_aggregate_types = ["sum", "mean"] + if aggregate_type not in valid_aggregate_types: + raise ValueError( + f"Given aggregate_type {aggregate_type} but only " + f"{aggregate_type} are supported." + ) + if structurewise_aggregate: + raise NotImplementedError( + "structurewise aggregation has not been implemented." + ) + + self._aggregate_names = aggregate_names + self._aggregate_type = aggregate_type + self._structurewise_aggregate = structurewise_aggregate + + def aggregate_features(self, tensor: TorchTensorMap) -> TorchTensorMap: + if self._aggregate_type == "sum": + return metatensor.torch.sum_over_samples( + tensor, sample_names=self._aggregate_names + ) + elif self._aggregate_type == "mean": + return metatensor.torch.mean_over_samples( + tensor, sample_names=self._aggregate_names + ) + else: + raise NotImplementedError( + f"aggregate_type {self._aggregate_type} has not been implemented." + ) + + def aggregate_kernel( + self, + kernel: TorchTensorMap, + are_pseudo_points: Tuple[bool, bool] = (False, False), + ) -> TorchTensorMap: + if self._aggregate_type == "sum": + if not are_pseudo_points[0]: + kernel = metatensor.torch.sum_over_samples( + kernel, self._aggregate_names + ) + if not are_pseudo_points[1]: + raise NotImplementedError( + "properties dimenson cannot be aggregated for the moment" + ) + return kernel + elif self._aggregate_type == "mean": + if not are_pseudo_points[0]: + kernel = metatensor.torch.mean_over_samples( + kernel, self._aggregate_names + ) + if not are_pseudo_points[1]: + raise NotImplementedError( + "properties dimenson cannot be aggregated for the moment" + ) + return kernel + else: + raise NotImplementedError( + f"aggregate_type {self._aggregate_type} has not been implemented." + ) + + def forward( + self, + tensor1: TorchTensorMap, + tensor2: TorchTensorMap, + are_pseudo_points: Tuple[bool, bool] = (False, False), + ) -> TorchTensorMap: + return self.aggregate_kernel( + self.compute_kernel(tensor1, tensor2), are_pseudo_points + ) + + def compute_kernel( + self, tensor1: TorchTensorMap, tensor2: TorchTensorMap + ) -> TorchTensorMap: + raise NotImplementedError("compute_kernel needs to be implemented.") + + +class TorchAggregateLinear(TorchAggregateKernel): + def __init__( + self, + aggregate_names: Union[str, List[str]], + aggregate_type: str = "sum", + structurewise_aggregate: bool = False, + ): + super().__init__(aggregate_names, aggregate_type, structurewise_aggregate) + + def forward( + self, + tensor1: TorchTensorMap, + tensor2: TorchTensorMap, + are_pseudo_points: Tuple[bool, bool] = (False, False), + ) -> TorchTensorMap: + # we overwrite default behavior because for linear kernels we can do it more + # memory efficient + if not are_pseudo_points[0]: + tensor1 = self.aggregate_features(tensor1) + if not are_pseudo_points[1]: + tensor2 = self.aggregate_features(tensor2) + return self.compute_kernel(tensor1, tensor2) + + def compute_kernel( + self, tensor1: TorchTensorMap, tensor2: TorchTensorMap + ) -> TorchTensorMap: + return metatensor.torch.dot(tensor1, tensor2) + + +class TorchAggregatePolynomial(TorchAggregateKernel): + def __init__( + self, + aggregate_names: Union[str, List[str]], + aggregate_type: str = "sum", + structurewise_aggregate: bool = False, + degree: int = 2, + ): + super().__init__(aggregate_names, aggregate_type, structurewise_aggregate) + self._degree = degree + + def compute_kernel(self, tensor1: TorchTensorMap, tensor2: TorchTensorMap): + return metatensor.torch.pow( + metatensor.torch.dot(tensor1, tensor2), self._degree + ) + + +class GreedySelector: + """Wraps :py:class:`skmatter._selection.GreedySelector` for a TensorMap. + + The class creates a selector for each block. The selection will be done based + the values of each :py:class:`TensorBlock`. Gradients will not be considered for + the selection. + """ + + def __init__( + self, + selector_class: Type[skmatter._selection.GreedySelector], + selection_type: str, + **selector_arguments, + ) -> None: + self._selector_class = selector_class + self._selection_type = selection_type + self._selector_arguments = selector_arguments + + self._selector_arguments["selection_type"] = self._selection_type + self._support = None + + @property + def selector_class(self) -> Type[skmatter._selection.GreedySelector]: + """The class to perform the selection.""" + return self._selector_class + + @property + def selection_type(self) -> str: + """Whether to choose a subset of columns ('feature') or rows ('sample').""" + return self._selection_type + + @property + def selector_arguments(self) -> dict: + """Arguments passed to the ``selector_class``.""" + return self._selector_arguments + + @property + def support(self) -> TensorMap: + """TensorMap containing the support.""" + if self._support is None: + raise ValueError("No selections. Call fit method first.") + + return self._support + + def fit(self, X: TensorMap, warm_start: bool = False): # -> GreedySelector: + """Learn the features to select. + + :param X: + Training vectors. + :param warm_start: + Whether the fit should continue after having already run, after increasing + `n_to_select`. Assumes it is called with the same X. + """ + if len(X.component_names) != 0: + raise ValueError("Only blocks with no components are supported.") + + blocks = [] + for _, block in X.items(): + selector = self.selector_class(**self.selector_arguments) + selector.fit(block.values, warm_start=warm_start) + mask = selector.get_support() + + if self._selection_type == "feature": + samples = Labels.single() + properties = Labels( + names=block.properties.names, values=block.properties.values[mask] + ) + elif self._selection_type == "sample": + samples = Labels( + names=block.samples.names, values=block.samples.values[mask] + ) + properties = Labels.single() + + blocks.append( + TensorBlock( + values=np.zeros([len(samples), len(properties)], dtype=np.int32), + samples=samples, + components=[], + properties=properties, + ) + ) + + self._support = TensorMap(X.keys, blocks) + + return self + + def transform(self, X: TensorMap) -> TensorMap: + """Reduce X to the selected features. + + :param X: + The input tensor. + :returns: + The selected subset of the input. + """ + blocks = [] + for key, block in X.items(): + block_support = self.support.block(key) + + if self._selection_type == "feature": + new_block = metatensor.slice_block( + block, "properties", block_support.properties + ) + elif self._selection_type == "sample": + new_block = metatensor.slice_block( + block, "samples", block_support.samples + ) + blocks.append(new_block) + + return TensorMap(X.keys, blocks) + + def fit_transform(self, X: TensorMap, warm_start: bool = False) -> TensorMap: + """Fit to data, then transform it. + + :param X: + Training vectors. + :param warm_start: + Whether the fit should continue after having already run, after increasing + `n_to_select`. Assumes it is called with the same X. + """ + return self.fit(X, warm_start=warm_start).transform(X) + + +class FPS(GreedySelector): + """ + Transformer that performs Greedy Sample Selection using Farthest Point Sampling. + + Refer to :py:class:`skmatter.sample_selection.FPS` for full documentation. + """ + + def __init__( + self, + initialize=0, + n_to_select=None, + score_threshold=None, + score_threshold_type="absolute", + progress_bar=False, + full=False, + random_state=0, + ): + super().__init__( + selector_class=_FPS, + selection_type="sample", + initialize=initialize, + n_to_select=n_to_select, + score_threshold=score_threshold, + score_threshold_type=score_threshold_type, + progress_bar=progress_bar, + full=full, + random_state=random_state, + ) + self._n_to_select = n_to_select + + +def torch_tensor_map_to_core(torch_tensor: TorchTensorMap): + torch_blocks = [] + for _, torch_block in torch_tensor.items(): + torch_blocks.append(torch_tensor_block_to_core(torch_block)) + torch_keys = torch_labels_to_core(torch_tensor.keys) + return TensorMap(torch_keys, torch_blocks) + + +def torch_tensor_block_to_core(torch_block: TorchTensorBlock): + """Transforms a tensor block from metatensor-torch to metatensor-torch + :param torch_block: + tensor block from metatensor-torch + :returns torch_block: + tensor block from metatensor-torch + """ + block = TensorBlock( + values=torch_block.values.detach().cpu().numpy(), + samples=torch_labels_to_core(torch_block.samples), + components=[ + torch_labels_to_core(component) for component in torch_block.components + ], + properties=torch_labels_to_core(torch_block.properties), + ) + for parameter, gradient in torch_block.gradients(): + block.add_gradient( + parameter=parameter, + gradient=TensorBlock( + values=gradient.values.detach().cpu().numpy(), + samples=torch_labels_to_core(gradient.samples), + components=[ + torch_labels_to_core(component) for component in gradient.components + ], + properties=torch_labels_to_core(gradient.properties), + ), + ) + return block + + +def torch_labels_to_core(torch_labels: TorchLabels): + """Transforms labels from metatensor-torch to metatensor-torch + :param torch_block: + tensor block from metatensor-torch + :returns torch_block: + labels from metatensor-torch + """ + return Labels(torch_labels.names, torch_labels.values.detach().cpu().numpy()) + + +### + + +def core_tensor_map_to_torch(core_tensor: TensorMap): + """Transforms a tensor map from metatensor-core to metatensor-torch + :param core_tensor: + tensor map from metatensor-core + :returns torch_tensor: + tensor map from metatensor-torch + """ + + torch_blocks = [] + for _, core_block in core_tensor.items(): + torch_blocks.append(core_tensor_block_to_torch(core_block)) + torch_keys = core_labels_to_torch(core_tensor.keys) + return TorchTensorMap(torch_keys, torch_blocks) + + +def core_tensor_block_to_torch(core_block: TensorBlock): + """Transforms a tensor block from metatensor-core to metatensor-torch + :param core_block: + tensor block from metatensor-core + :returns torch_block: + tensor block from metatensor-torch + """ + block = TorchTensorBlock( + values=torch.tensor(core_block.values), + samples=core_labels_to_torch(core_block.samples), + components=[ + core_labels_to_torch(component) for component in core_block.components + ], + properties=core_labels_to_torch(core_block.properties), + ) + for parameter, gradient in core_block.gradients(): + block.add_gradient( + parameter=parameter, + gradient=TensorBlock( + values=gradient.values.detach().cpu().numpy(), + samples=torch_labels_to_core(gradient.samples), + components=[ + torch_labels_to_core(component) for component in gradient.components + ], + properties=torch_labels_to_core(gradient.properties), + ), + ) + return block + + +def core_labels_to_torch(core_labels: Labels): + """Transforms labels from metatensor-core to metatensor-torch + :param core_block: + tensor block from metatensor-core + :returns torch_block: + labels from metatensor-torch + """ + return TorchLabels(core_labels.names, torch.tensor(core_labels.values)) + + +class SubsetOfRegressors: + def __init__( + self, + kernel_type: Union[str, AggregateKernel] = "linear", + kernel_kwargs: Optional[dict] = None, + ): + if kernel_kwargs is None: + kernel_kwargs = {} + self._set_kernel(kernel_type, **kernel_kwargs) + self._kernel_type = kernel_type + self._kernel_kwargs = kernel_kwargs + self._X_pseudo = None + self._weights = None + + def _set_kernel(self, kernel: Union[str, AggregateKernel], **kernel_kwargs): + valid_kernels = ["linear", "polynomial", "precomputed"] + aggregate_type = kernel_kwargs.get("aggregate_type", "sum") + if aggregate_type != "sum": + raise ValueError( + f'aggregate_type={aggregate_type!r} found but must be "sum"' + ) + self._kernel: Union[AggregateKernel, None] = None + if kernel == "linear": + self._kernel = AggregateLinear(aggregate_type="sum", **kernel_kwargs) + elif kernel == "polynomial": + self._kernel = AggregatePolynomial(aggregate_type="sum", **kernel_kwargs) + elif kernel == "precomputed": + self._kernel = None + elif isinstance(kernel, AggregateKernel): + self._kernel = kernel + else: + raise ValueError( + f"kernel type {kernel!r} is not supported. Please use one " + f"of the valid kernels {valid_kernels!r}" + ) + + def fit( + self, + X: TensorMap, + X_pseudo: TensorMap, + y: TensorMap, + alpha: float = 1.0, + alpha_forces: Optional[float] = None, + solver: str = "RKHS-QR", + rcond: Optional[float] = None, + ): + r""" + :param X: + features + if kernel type "precomputed" is used, the kernel k_nm is assumed + :param X_pseudo: + pseudo points + if kernel type "precomputed" is used, the kernel k_mm is assumed + :param y: + targets + :param alpha: + regularizationfor the energies, it must be a float + :param alpha_forces: + regularization for the forces, it must be a float. If None is set + equal to alpha + :param solver: + determines which solver to use + :param rcond: + argument for the solver lstsq + + + Derivation + ---------- + + We take equation the mean expression + + .. math:: + + \sigma^{-2} K_{tm}\Sigma K_{MN}y + + we put in the $\Sigma$ + + .. math:: + + \sigma^{-2} K_{tm}(\sigma^{-2}K_{mn}K_{mn}+K_{mm})^{-1} K_{mn}y + + We can move around the $\sigma's$ + + .. math:: + + K_{tm}((K_{mn}\sigma^{-1})(\sigma^{-1}K_{mn)}+K_{mm})^{-1} + (K_{mn}\sigma^{-1})(y\sigma^{-1}) + + you can see the building blocks in the code are $K_{mn}\sigma^{-1}$ and + $y\sigma^{-1}$ + """ + if isinstance(alpha, float): + alpha_energy = alpha + else: + raise ValueError("alpha must either be a float") + + if alpha_forces is None: + alpha_forces = alpha_energy + else: + if not isinstance(alpha_forces, float): + raise ValueError("alpha must either be a float") + + X = X.to(arrays="numpy") + X_pseudo = X_pseudo.to(arrays="numpy") + y = y.to(arrays="numpy") + + if self._kernel is None: + # _set_kernel only returns None if kernel type is precomputed + k_nm = X + k_mm = X_pseudo + else: + k_mm = self._kernel(X_pseudo, X_pseudo, are_pseudo_points=(True, True)) + k_nm = self._kernel(X, X_pseudo, are_pseudo_points=(False, True)) + + # solve + # TODO: allow for different regularizer for energies and forces + weight_blocks = [] + for key, y_block in y.items(): + k_nm_block = k_nm.block(key) + k_mm_block = k_mm.block(key) + X_block = X.block(key) + structures = metatensor.operations._dispatch.unique( + k_nm_block.samples["system"] + ) + n_atoms_per_structure = [] + for structure in structures: + n_atoms = np.sum(X_block.samples["system"] == structure) + n_atoms_per_structure.append(float(n_atoms)) + + n_atoms_per_structure = np.array(n_atoms_per_structure) + normalization = metatensor.operations._dispatch.sqrt(n_atoms_per_structure) + + if not (np.allclose(alpha_energy, 0.0)): + normalization /= alpha_energy + normalization = normalization[:, None] + + k_nm_reg = k_nm_block.values * normalization + y_reg = (y_block.values) * normalization + if len(k_nm_block.gradients_list()) > 0: + grad_shape = k_nm_block.gradient("positions").values.shape + k_nm_reg = np.vstack( + [ + k_nm_reg, + k_nm_block.gradient("positions").values.reshape( + grad_shape[0] * grad_shape[1], + grad_shape[2], + ) + / alpha_forces, + ] + ) + grad_shape = y_block.gradient("positions").values.shape + y_reg = np.vstack( + [ + y_reg, + y_block.gradient("positions").values.reshape( + grad_shape[0] * grad_shape[1], + grad_shape[2], + ) + / alpha_forces, + ] + ) + self._solver = _SorKernelSolver( + k_mm_block.values, regularizer=1, jitter=0, solver=solver + ) + + if rcond is None: + rcond_ = max(k_nm_reg.shape) * np.finfo(k_nm_reg.dtype.char.lower()).eps + else: + rcond_ = rcond + self._solver.fit(k_nm_reg, y_reg, rcond=rcond_) + + weight_block = TensorBlock( + values=self._solver.weights.T, + samples=y_block.properties, + components=k_nm_block.components, + properties=k_nm_block.properties, + ) + weight_blocks.append(weight_block) + + self._weights = TensorMap(y.keys, weight_blocks) + + self._X_pseudo = X_pseudo.copy() + + def predict(self, T: TensorMap) -> TensorMap: + """ + :param T: + features + if kernel type "precomputed" is used, the kernel k_tm is assumed + """ + if self._weights is None: + raise ValueError( + "The weights are not defined. Are you sure you" + " have run the `fit` method?" + ) + if self._kernel_type == "precomputed": + k_tm = T + else: + k_tm = self._kernel(T, self._X_pseudo, are_pseudo_points=(False, True)) + return metatensor.dot(k_tm, self._weights) + + def export_torch_script_model(self): + return TorchSubsetofRegressors( + core_tensor_map_to_torch(self._weights), + core_tensor_map_to_torch(self._X_pseudo), + self._kernel_type, + self._kernel_kwargs, + ) + + +class TorchSubsetofRegressors(torch.nn.Module): + def __init__( + self, + weights: TorchTensorMap, + X_pseudo: TorchTensorMap, + kernel_type: Union[str, AggregateKernel] = "linear", + kernel_kwargs: Optional[dict] = None, + ): + super().__init__() + self._weights = weights + self._X_pseudo = X_pseudo + if kernel_kwargs is None: + kernel_kwargs = {} + self._set_kernel(kernel_type, **kernel_kwargs) + + def forward(self, T: TorchTensorMap) -> TorchTensorMap: + """ + :param T: + features + if kernel type "precomputed" is used, the kernel k_tm is assumed + """ + # move weights and X_pseudo to the same device as T + self._weights = self._weights.to(T.device) + self._X_pseudo = self._X_pseudo.to(T.device) + + k_tm = self._kernel(T, self._X_pseudo, are_pseudo_points=(False, True)) + return metatensor.torch.dot(k_tm, self._weights) + + def _set_kernel(self, kernel: Union[str, TorchAggregateKernel], **kernel_kwargs): + valid_kernels = ["linear", "polynomial", "precomputed"] + aggregate_type = kernel_kwargs.get("aggregate_type", "sum") + if aggregate_type != "sum": + raise ValueError( + f'aggregate_type={aggregate_type!r} found but must be "sum"' + ) + if kernel == "linear": + self._kernel = TorchAggregateLinear(aggregate_type="sum", **kernel_kwargs) + elif kernel == "polynomial": + self._kernel = TorchAggregatePolynomial( + aggregate_type="sum", **kernel_kwargs + ) + elif kernel == "precomputed": + raise NotImplementedError( + "precomputed kernels are note supported in torch" + "version of subset of regressor." + ) + elif isinstance(kernel, TorchAggregateKernel): + self._kernel = kernel + else: + raise ValueError( + f"kernel type {kernel!r} is not supported. Please use one " + f"of the valid kernels {valid_kernels!r}" + ) diff --git a/src/metatensor/models/experimental/gap/tests/__init__.py b/src/metatensor/models/experimental/gap/tests/__init__.py new file mode 100644 index 000000000..48d5a22ec --- /dev/null +++ b/src/metatensor/models/experimental/gap/tests/__init__.py @@ -0,0 +1,11 @@ +from pathlib import Path + +DATASET_PATH = str( + Path(__file__).parent.resolve() + / "../../../../../../tests/resources/qm9_reduced_100.xyz" +) + +DATASET_ETHANOL_PATH = str( + Path(__file__).parent.resolve() + / "../../../../../../tests/resources/ethanol_reduced_100.xyz" +) diff --git a/src/metatensor/models/experimental/gap/tests/test_errors.py b/src/metatensor/models/experimental/gap/tests/test_errors.py new file mode 100644 index 000000000..fc384f1a8 --- /dev/null +++ b/src/metatensor/models/experimental/gap/tests/test_errors.py @@ -0,0 +1,79 @@ +import copy +import random +import re + +import numpy as np +import pytest +import torch +from omegaconf import OmegaConf + +from metatensor.models.experimental.gap import GAP, Trainer +from metatensor.models.utils.architectures import get_default_hypers +from metatensor.models.utils.data import Dataset, DatasetInfo, TargetInfo +from metatensor.models.utils.data.readers import read_systems, read_targets + +from . import DATASET_ETHANOL_PATH + + +DEFAULT_HYPERS = get_default_hypers("experimental.gap") + + +torch.set_default_dtype(torch.float64) # GAP only supports float64 + + +# reproducibility +random.seed(0) +np.random.seed(0) +torch.manual_seed(0) + + +def test_ethanol_regression_train_and_invariance(): + """test the error if the number of sparse point + is bigger than the number of environments + """ + + systems = read_systems(DATASET_ETHANOL_PATH, dtype=torch.float64) + + conf = { + "energy": { + "quantity": "energy", + "read_from": DATASET_ETHANOL_PATH, + "file_format": ".xyz", + "key": "energy", + "forces": { + "read_from": DATASET_ETHANOL_PATH, + "file_format": ".xyz", + "key": "forces", + }, + "stress": False, + "virial": False, + } + } + + targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + dataset = Dataset({"system": systems[:2], "energy": targets["energy"][:2]}) + + hypers = copy.deepcopy(DEFAULT_HYPERS) + hypers["model"]["krr"]["num_sparse_points"] = 30 + + dataset_info = DatasetInfo( + length_unit="Angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "energy": TargetInfo( + quantity="energy", + unit="kcal/mol", + ), + }, + ) + + gap = GAP(hypers["model"], dataset_info) + trainer = Trainer(hypers["training"]) + with pytest.raises( + ValueError, + match=re.escape( + """number of sparse points (30) + should be smaller than the number of environments (18)""" + ), + ): + trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".") diff --git a/src/metatensor/models/experimental/gap/tests/test_regression.py b/src/metatensor/models/experimental/gap/tests/test_regression.py new file mode 100644 index 000000000..ab6d14369 --- /dev/null +++ b/src/metatensor/models/experimental/gap/tests/test_regression.py @@ -0,0 +1,186 @@ +import copy +import random + +import ase.io +import metatensor.torch +import numpy as np +import torch +from omegaconf import OmegaConf + +from metatensor.models.experimental.gap import GAP, Trainer +from metatensor.models.utils.architectures import get_default_hypers +from metatensor.models.utils.data import Dataset, DatasetInfo, TargetInfo +from metatensor.models.utils.data.readers import read_systems, read_targets + +from . import DATASET_ETHANOL_PATH, DATASET_PATH + + +DEFAULT_HYPERS = get_default_hypers("experimental.gap") + + +torch.set_default_dtype(torch.float64) # GAP only supports float64 + + +# reproducibility +random.seed(0) +np.random.seed(0) +torch.manual_seed(0) + + +def test_regression_init(): + """Perform a regression test on the model at initialization""" + + dataset_info = DatasetInfo( + length_unit="Angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "mtm::U0": TargetInfo( + quantity="energy", + unit="eV", + ), + }, + ) + GAP(DEFAULT_HYPERS["model"], dataset_info) + + +def test_regression_train_and_invariance(): + """Perform a regression test on the model when trained for 2 epoch on a small + dataset. We perform also the invariance test here because one needs a trained model + for this. + """ + + systems = read_systems(DATASET_PATH, dtype=torch.float64) + + conf = { + "mtm::U0": { + "quantity": "energy", + "read_from": DATASET_PATH, + "file_format": ".xyz", + "key": "U0", + "forces": False, + "stress": False, + "virial": False, + } + } + targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + dataset = Dataset({"system": systems, "mtm::U0": targets["mtm::U0"]}) + + dataset_info = DatasetInfo( + length_unit="Angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "mtm::U0": TargetInfo( + quantity="energy", + unit="eV", + ), + }, + ) + gap = GAP(DEFAULT_HYPERS["model"], dataset_info) + trainer = Trainer(DEFAULT_HYPERS["training"]) + trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".") + + # Predict on the first five systems + output = gap(systems[:5], {"mtm::U0": gap.outputs["mtm::U0"]}) + + expected_output = torch.tensor( + [[-40.5891], [-56.7122], [-76.4146], [-77.3364], [-93.4905]] + ) + + assert torch.allclose(output["mtm::U0"].block().values, expected_output, rtol=0.3) + + # Tests that the model is rotationally invariant + system = ase.io.read(DATASET_PATH) + system.numbers = np.ones(len(system.numbers)) + + original_system = copy.deepcopy(system) + system.rotate(48, "y") + + original_output = gap( + [metatensor.torch.atomistic.systems_to_torch(original_system)], + {"mtm::U0": gap.outputs["mtm::U0"]}, + ) + rotated_output = gap( + [metatensor.torch.atomistic.systems_to_torch(system)], + {"mtm::U0": gap.outputs["mtm::U0"]}, + ) + + assert torch.allclose( + original_output["mtm::U0"].block().values, + rotated_output["mtm::U0"].block().values, + ) + + +def test_ethanol_regression_train_and_invariance(): + """Perform a regression test on the model when trained for 2 epoch on a small + dataset. We perform also the invariance test here because one needs a trained model + for this. + """ + + systems = read_systems(DATASET_ETHANOL_PATH, dtype=torch.float64) + + conf = { + "energy": { + "quantity": "energy", + "read_from": DATASET_ETHANOL_PATH, + "file_format": ".xyz", + "key": "energy", + "forces": { + "read_from": DATASET_ETHANOL_PATH, + "file_format": ".xyz", + "key": "forces", + }, + "stress": False, + "virial": False, + } + } + + targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + dataset = Dataset({"system": systems, "energy": targets["energy"]}) + + hypers = copy.deepcopy(DEFAULT_HYPERS) + hypers["model"]["krr"]["num_sparse_points"] = 900 + + dataset_info = DatasetInfo( + length_unit="Angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "energy": TargetInfo( + quantity="energy", + unit="kcal/mol", + ), + }, + ) + + gap = GAP(hypers["model"], dataset_info) + trainer = Trainer(hypers["training"]) + trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".") + + # Predict on the first five systems + output = gap(systems[:5], {"energy": gap.outputs["energy"]}) + data = ase.io.read(DATASET_ETHANOL_PATH, ":5", format="extxyz") + + expected_output = torch.tensor([[i.info["energy"]] for i in data]) + assert torch.allclose(output["energy"].block().values, expected_output, rtol=0.1) + + # TODO: check accuracy of training forces + # expected_forces = torch.vstack([torch.Tensor(i.arrays["forces"]) for i in data]) + + # Tests that the model is rotationally invariant + system = ase.io.read(DATASET_ETHANOL_PATH) + + original_system = copy.deepcopy(system) + system.rotate(48, "y") + + original_output = gap( + [metatensor.torch.atomistic.systems_to_torch(original_system)], + {"energy": gap.outputs["energy"]}, + ) + rotated_output = gap( + [metatensor.torch.atomistic.systems_to_torch(system)], + {"energy": gap.outputs["energy"]}, + ) + + assert torch.allclose( + original_output["energy"].block().values, + rotated_output["energy"].block().values, + ) diff --git a/src/metatensor/models/experimental/gap/tests/test_torchscript.py b/src/metatensor/models/experimental/gap/tests/test_torchscript.py new file mode 100644 index 000000000..a134f7e37 --- /dev/null +++ b/src/metatensor/models/experimental/gap/tests/test_torchscript.py @@ -0,0 +1,84 @@ +import torch +from omegaconf import OmegaConf + +from metatensor.models.experimental.gap import GAP, Trainer +from metatensor.models.utils.architectures import get_default_hypers +from metatensor.models.utils.data import Dataset, DatasetInfo, TargetInfo +from metatensor.models.utils.data.readers import read_systems, read_targets + +from . import DATASET_PATH + + +DEFAULT_HYPERS = get_default_hypers("experimental.gap") + + +torch.set_default_dtype(torch.float64) # GAP only supports float64 + + +def test_torchscript(): + """Tests that the model can be jitted.""" + + dataset_info = DatasetInfo( + length_unit="Angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "mtm::U0": TargetInfo( + quantity="energy", + unit="eV", + ), + }, + ) + conf = { + "mtm::U0": { + "quantity": "energy", + "read_from": DATASET_PATH, + "file_format": ".xyz", + "key": "U0", + "forces": False, + "stress": False, + "virial": False, + } + } + targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + systems = read_systems(DATASET_PATH, dtype=torch.float64) + + # for system in systems: + # system.types = torch.ones(len(system.types), dtype=torch.int32) + dataset = Dataset({"system": systems, "mtm::U0": targets["mtm::U0"]}) + + hypers = DEFAULT_HYPERS.copy() + gap = GAP(DEFAULT_HYPERS["model"], dataset_info) + trainer = Trainer(hypers["training"]) + gap = trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".") + scripted_gap = torch.jit.script(gap) + + ref_output = gap.forward(systems[:5], {"mtm::U0": gap.outputs["mtm::U0"]}) + scripted_output = scripted_gap.forward( + systems[:5], {"mtm::U0": gap.outputs["mtm::U0"]} + ) + + assert torch.allclose( + ref_output["mtm::U0"].block().values, + scripted_output["mtm::U0"].block().values, + ) + + +def test_torchscript_save(): + """Tests that the model can be jitted and saved.""" + + dataset_info = DatasetInfo( + length_unit="Angstrom", + atomic_types=[1, 6, 7, 8], + targets={ + "mtm::U0": TargetInfo( + quantity="energy", + unit="eV", + ), + }, + ) + gap = GAP(DEFAULT_HYPERS["model"], dataset_info) + torch.jit.save( + torch.jit.script(gap), + "gap.pt", + ) + torch.jit.load("gap.pt") diff --git a/src/metatensor/models/experimental/gap/trainer.py b/src/metatensor/models/experimental/gap/trainer.py new file mode 100644 index 000000000..236b83932 --- /dev/null +++ b/src/metatensor/models/experimental/gap/trainer.py @@ -0,0 +1,146 @@ +import logging +from typing import List, Union + +import metatensor.torch +import torch +from metatensor.torch import TensorMap + +import metatensor +from metatensor.models.utils.data import Dataset + +from ...utils.composition import calculate_composition_weights +from ...utils.data import check_datasets +from . import GAP +from .model import torch_tensor_map_to_core + + +logger = logging.getLogger(__name__) + + +class Trainer: + def __init__(self, train_hypers): + self.hypers = train_hypers + + def train( + self, + model: GAP, + devices: List[torch.device], + train_datasets: List[Union[Dataset, torch.utils.data.Subset]], + validation_datasets: List[Union[Dataset, torch.utils.data.Subset]], + checkpoint_dir: str, + ): + # checks + assert devices == [torch.device("cpu")] + dtype = train_datasets[0][0]["system"].positions.dtype + if dtype != torch.float64: + raise ValueError("GAP only supports float64") + target_name = next(iter(model.dataset_info.targets.keys())) + if len(train_datasets) != 1: + raise ValueError("GAP only supports a single training dataset") + if len(validation_datasets) != 1: + raise ValueError("GAP only supports a single validation dataset") + + # Perform checks on the datasets: + logger.info("Checking datasets for consistency") + check_datasets(train_datasets, validation_datasets) + + logger.info("Training on device cpu") + + outputs_dict = model.dataset_info.targets + if len(outputs_dict.keys()) > 1: + raise NotImplementedError("More than one output is not supported yet.") + output_name = next(iter(outputs_dict.keys())) + + # Calculate and set the composition weights: + logger.info("Calculating composition weights") + composition_weights, species = calculate_composition_weights( + train_datasets, target_name + ) + model.set_composition_weights(target_name, composition_weights, species) + + logger.info("Setting up data loaders") + + if len(train_datasets[0][0][output_name].keys) > 1: + raise NotImplementedError( + "Found more than 1 key in targets. Assuming " + "equivariant learning which is not supported yet." + ) + train_dataset = train_datasets[0] + train_y = metatensor.torch.join( + [sample[output_name] for sample in train_dataset], + axis="samples", + remove_tensor_name=True, + ) + model._keys = train_y.keys + train_structures = [sample["system"] for sample in train_dataset] + composition_energies = torch.zeros(len(train_y.block().values), dtype=dtype) + for i, structure in enumerate(train_structures): + for j, s in enumerate(species): + composition_energies[i] += ( + torch.sum(structure.types == s) * composition_weights[j] + ) + train_y_values = train_y.block().values + train_y_values = train_y_values - composition_energies.reshape(-1, 1) + train_block = metatensor.torch.TensorBlock( + values=train_y_values, + samples=train_y.block().samples, + components=train_y.block().components, + properties=train_y.block().properties, + ) + if len(train_y[0].gradients_list()) > 0: + train_block.add_gradient("positions", train_y[0].gradient("positions")) + + train_y = metatensor.torch.TensorMap( + train_y.keys, + [train_block], + ) + + if len(train_y[0].gradients_list()) > 0: + train_tensor = model._soap_torch_calculator.compute( + train_structures, gradients=["positions"] + ) + else: + train_tensor = model._soap_torch_calculator.compute(train_structures) + model._species_labels = train_tensor.keys + train_tensor = train_tensor.keys_to_samples("center_type") + # here, we move to properties to use metatensor operations to aggregate + # later on. Perhaps we could retain the sparsity all the way to the kernels + # of the soap features with a lot more implementation effort + train_tensor = train_tensor.keys_to_properties( + ["neighbor_1_type", "neighbor_2_type"] + ) + # change backend + train_tensor = TensorMap(train_y.keys, train_tensor.blocks()) + train_tensor = torch_tensor_map_to_core(train_tensor) + train_y = torch_tensor_map_to_core(train_y) + + lens = len(train_tensor[0].values) + if model._sampler._n_to_select > lens: + raise ValueError( + f"""number of sparse points ({model._sampler._n_to_select}) + should be smaller than the number of environments ({lens})""" + ) + sparse_points = model._sampler.fit_transform(train_tensor) + sparse_points = metatensor.operations.remove_gradients(sparse_points) + alpha_energy = self.hypers["regularizer"] + if self.hypers["regularizer_forces"] is None: + alpha_forces = alpha_energy + else: + alpha_forces = self.hypers["regularizer_forces"] + + logger.info(f"Training on device cpu with dtype {dtype}") + logger.info("Fitting GAP") + + model._subset_of_regressors.fit( + train_tensor, + sparse_points, + train_y, + alpha=alpha_energy, + alpha_forces=alpha_forces, + ) + + model._subset_of_regressors_torch = ( + model._subset_of_regressors.export_torch_script_model() + ) + + return model diff --git a/src/metatensor/models/experimental/soap_bpnn/trainer.py b/src/metatensor/models/experimental/soap_bpnn/trainer.py index 9fdc97985..3dbd3443b 100644 --- a/src/metatensor/models/experimental/soap_bpnn/trainer.py +++ b/src/metatensor/models/experimental/soap_bpnn/trainer.py @@ -46,7 +46,7 @@ def train( assert len(devices) == 1 device = devices[0] - logger.info(f"training on device {device} with dtype {dtype}") + logger.info(f"Training on device {device} with dtype {dtype}") model.to(device=device, dtype=dtype) # Calculate and set the composition weights for all targets: diff --git a/tox.ini b/tox.ini index 821212271..597d2c7a4 100644 --- a/tox.ini +++ b/tox.ini @@ -117,6 +117,16 @@ changedir = src/metatensor/models/experimental/pet/tests/ commands = pytest {[testenv]warning_options} {posargs} +[testenv:gap-tests] +description = Run GAP tests with pytest +passenv = * +deps = + pytest +extras = gap +changedir = src/metatensor/models/experimental/gap/tests/ +commands = + pytest {[testenv]warning_options} {posargs} + [testenv:docs] description = builds the documentation with sphinx deps = From 958a773c3f1e9b0fb231dab5dc7b891b091b47f5 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 30 May 2024 15:50:06 +0200 Subject: [PATCH 04/16] Adapt to new read_targets --- src/metatensor/models/experimental/gap/tests/test_errors.py | 3 ++- .../models/experimental/gap/tests/test_regression.py | 6 ++++-- .../models/experimental/gap/tests/test_torchscript.py | 3 ++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/metatensor/models/experimental/gap/tests/test_errors.py b/src/metatensor/models/experimental/gap/tests/test_errors.py index fc384f1a8..6cb53df65 100644 --- a/src/metatensor/models/experimental/gap/tests/test_errors.py +++ b/src/metatensor/models/experimental/gap/tests/test_errors.py @@ -40,6 +40,7 @@ def test_ethanol_regression_train_and_invariance(): "read_from": DATASET_ETHANOL_PATH, "file_format": ".xyz", "key": "energy", + "unit": "kcal/mol", "forces": { "read_from": DATASET_ETHANOL_PATH, "file_format": ".xyz", @@ -50,7 +51,7 @@ def test_ethanol_regression_train_and_invariance(): } } - targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + targets, _ = read_targets(OmegaConf.create(conf), dtype=torch.float64) dataset = Dataset({"system": systems[:2], "energy": targets["energy"][:2]}) hypers = copy.deepcopy(DEFAULT_HYPERS) diff --git a/src/metatensor/models/experimental/gap/tests/test_regression.py b/src/metatensor/models/experimental/gap/tests/test_regression.py index ab6d14369..c1208528b 100644 --- a/src/metatensor/models/experimental/gap/tests/test_regression.py +++ b/src/metatensor/models/experimental/gap/tests/test_regression.py @@ -57,12 +57,13 @@ def test_regression_train_and_invariance(): "read_from": DATASET_PATH, "file_format": ".xyz", "key": "U0", + "unit": "kcal/mol", "forces": False, "stress": False, "virial": False, } } - targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + targets, _ = read_targets(OmegaConf.create(conf), dtype=torch.float64) dataset = Dataset({"system": systems, "mtm::U0": targets["mtm::U0"]}) dataset_info = DatasetInfo( @@ -129,12 +130,13 @@ def test_ethanol_regression_train_and_invariance(): "file_format": ".xyz", "key": "forces", }, + "unit": "kcal/mol", "stress": False, "virial": False, } } - targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + targets, _ = read_targets(OmegaConf.create(conf), dtype=torch.float64) dataset = Dataset({"system": systems, "energy": targets["energy"]}) hypers = copy.deepcopy(DEFAULT_HYPERS) diff --git a/src/metatensor/models/experimental/gap/tests/test_torchscript.py b/src/metatensor/models/experimental/gap/tests/test_torchscript.py index a134f7e37..2c69f6464 100644 --- a/src/metatensor/models/experimental/gap/tests/test_torchscript.py +++ b/src/metatensor/models/experimental/gap/tests/test_torchscript.py @@ -34,12 +34,13 @@ def test_torchscript(): "read_from": DATASET_PATH, "file_format": ".xyz", "key": "U0", + "unit": "kcal/mol", "forces": False, "stress": False, "virial": False, } } - targets = read_targets(OmegaConf.create(conf), dtype=torch.float64) + targets, _ = read_targets(OmegaConf.create(conf), dtype=torch.float64) systems = read_systems(DATASET_PATH, dtype=torch.float64) # for system in systems: From dd9b52300713f18bad8c1954ec7cd9c9ec68bb75 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 30 May 2024 22:22:53 +0200 Subject: [PATCH 05/16] Fix this test --- tests/cli/test_eval_model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/cli/test_eval_model.py b/tests/cli/test_eval_model.py index 64f3e5bdb..d1ea226eb 100644 --- a/tests/cli/test_eval_model.py +++ b/tests/cli/test_eval_model.py @@ -35,6 +35,8 @@ def test_eval_cli(monkeypatch, tmp_path): "eval", str(MODEL_PATH), str(EVAL_OPTIONS_PATH), + "-e", + str(RESOURCES_PATH / "extensions"), ] output = subprocess.check_output(command, stderr=subprocess.STDOUT) From 5f6f7fed5558454a6691b4e910291b6869c51007 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Thu, 30 May 2024 22:33:07 +0200 Subject: [PATCH 06/16] Log destination of checkpoint, exported model and extensions at the end of training --- src/metatensor/models/cli/export.py | 2 +- src/metatensor/models/cli/train.py | 12 ++++++++++-- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/metatensor/models/cli/export.py b/src/metatensor/models/cli/export.py index 121cc5a2c..a0fe576b2 100644 --- a/src/metatensor/models/cli/export.py +++ b/src/metatensor/models/cli/export.py @@ -63,4 +63,4 @@ def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> torch.jit.save(model, path) else: mts_atomistic_model = model.export() - mts_atomistic_model.export(path, collect_extensions="extensions") + mts_atomistic_model.export(path, collect_extensions="extensions/") diff --git a/src/metatensor/models/cli/train.py b/src/metatensor/models/cli/train.py index 6bdc44de3..5350341cc 100644 --- a/src/metatensor/models/cli/train.py +++ b/src/metatensor/models/cli/train.py @@ -366,15 +366,23 @@ def train_model( # SAVE FINAL MODEL ######## ########################### - logger.info("Training finished; save final checkpoint and model") output_checked = check_suffix(filename=output, suffix=".pt") + logger.info( + "Training finished, saving final checkpoint " + f"to {str(Path(output_checked).stem)}.ckpt" + ) try: model.save_checkpoint(f"{Path(output_checked).stem}.ckpt") except Exception as e: raise ArchitectureError(e) mts_atomistic_model = model.export() - mts_atomistic_model.export(str(output_checked), collect_extensions="extensions") + extensions_path = "extensions/" + + logger.info( + f"Exporting model to {output_checked} and extensions to {extensions_path}" + ) + mts_atomistic_model.export(str(output_checked), collect_extensions=extensions_path) ########################### # EVALUATE FINAL MODEL #### From 028b6c50828a79ac273d443c3bc0a4ac91cbc0e9 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Mon, 3 Jun 2024 13:52:11 +0200 Subject: [PATCH 07/16] Fix ASE --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c75eb3d95..b34c82e91 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -9,7 +9,7 @@ description = "Atomistic models using metatensor" authors = [{name = "metatensor-models developers"}] dependencies = [ - "ase", + "ase < 3.23.0", "metatensor-learn==0.2.2", "metatensor-operations==0.2.1", "metatensor-torch==0.5.1", From c9c1357312b61c7d73d32e3acca21b4f78b40438 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Mon, 3 Jun 2024 14:01:54 +0200 Subject: [PATCH 08/16] Add tests that extensions are saved --- tests/cli/test_export_model.py | 9 +++++++++ tests/cli/test_train_model.py | 4 ++++ 2 files changed, 13 insertions(+) diff --git a/tests/cli/test_export_model.py b/tests/cli/test_export_model.py index a659e1b7c..c9941e9ed 100644 --- a/tests/cli/test_export_model.py +++ b/tests/cli/test_export_model.py @@ -3,6 +3,7 @@ Actual unit tests for the function are performed in `tests/utils/test_export`. """ +import glob import subprocess from pathlib import Path @@ -33,6 +34,10 @@ def test_export(monkeypatch, tmp_path, path): model = __model__(model_hypers=MODEL_HYPERS, dataset_info=dataset_info) export_model(model, path) + # Test if extensions are saved + extensions_glob = glob.glob("extensions/") + assert len(extensions_glob) == 1 + assert Path(path).is_file() @@ -55,6 +60,10 @@ def test_export_cli(monkeypatch, tmp_path, output): subprocess.check_call(command) assert Path(output).is_file() + # Test if extensions are saved + extensions_glob = glob.glob("extensions/") + assert len(extensions_glob) == 1 + def test_reexport(monkeypatch, tmp_path): """Test that an already exported model can be loaded and again exported.""" diff --git a/tests/cli/test_train_model.py b/tests/cli/test_train_model.py index 6c07fb15c..78bc98226 100644 --- a/tests/cli/test_train_model.py +++ b/tests/cli/test_train_model.py @@ -57,6 +57,10 @@ def test_train(capfd, monkeypatch, tmp_path, output): log_glob = glob.glob("outputs/*/*/train.log") assert len(log_glob) == 1 + # Test if extensions are saved + extensions_glob = glob.glob("extensions/") + assert len(extensions_glob) == 1 + # Open the log file and check if the logging is correct with open(log_glob[0]) as f: file_log = f.read() From 1260d467c2ba84ebe00a75cd3f0605d507119821 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Mon, 3 Jun 2024 14:05:19 +0200 Subject: [PATCH 09/16] Also log message when exporting model --- src/metatensor/models/cli/export.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/metatensor/models/cli/export.py b/src/metatensor/models/cli/export.py index a0fe576b2..0166e1f0c 100644 --- a/src/metatensor/models/cli/export.py +++ b/src/metatensor/models/cli/export.py @@ -1,4 +1,5 @@ import argparse +import logging from pathlib import Path from typing import Any, Union @@ -9,6 +10,9 @@ from .formatter import CustomHelpFormatter +logger = logging.getLogger(__name__) + + def _add_export_model_parser(subparser: argparse._SubParsersAction) -> None: """Add `export_model` paramaters to an argparse (sub)-parser.""" @@ -60,7 +64,10 @@ def export_model(model: Any, output: Union[Path, str] = "exported-model.pt") -> path = str(check_suffix(filename=output, suffix=".pt")) if is_exported(model): + logger.info(f"The model is already exported. Saving it to `{path}`.") torch.jit.save(model, path) else: + extensions_path = "extensions/" + logger.info(f"Exporting model to {path} and extensions to {extensions_path}") mts_atomistic_model = model.export() - mts_atomistic_model.export(path, collect_extensions="extensions/") + mts_atomistic_model.export(path, collect_extensions=extensions_path) From edac40c605e4de9f423a7302d14f029ef48cbb6a Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Tue, 4 Jun 2024 09:11:26 +0200 Subject: [PATCH 10/16] Add Davide to codeowners --- CODEOWNERS | 1 + 1 file changed, 1 insertion(+) diff --git a/CODEOWNERS b/CODEOWNERS index 340908947..7da30f095 100644 --- a/CODEOWNERS +++ b/CODEOWNERS @@ -5,3 +5,4 @@ **/soap_bpnn @frostedoyster **/alchemical_model @abmazitov **/pet @spozdn @abmazitov +**/gap @DavideTisi From c6ce54e644a9e1d4febb6255db58e6e1f037bee7 Mon Sep 17 00:00:00 2001 From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com> Date: Tue, 4 Jun 2024 10:15:32 +0200 Subject: [PATCH 11/16] added suggestions --- .github/workflows/gap-tests.yml | 4 --- docs/src/architectures/gap.rst | 26 +++++++++---------- .../experimental/gap/default-hypers.yaml | 4 +-- 3 files changed, 15 insertions(+), 19 deletions(-) diff --git a/.github/workflows/gap-tests.yml b/.github/workflows/gap-tests.yml index a097c2a4a..f23f834aa 100644 --- a/.github/workflows/gap-tests.yml +++ b/.github/workflows/gap-tests.yml @@ -30,7 +30,3 @@ jobs: # Use the CPU only version of torch when building/running the code PIP_EXTRA_INDEX_URL: https://download.pytorch.org/whl/cpu - - name: Upload codecoverage - uses: codecov/codecov-action@v3 - with: - files: ./tests/coverage.xml diff --git a/docs/src/architectures/gap.rst b/docs/src/architectures/gap.rst index 631a5c7cf..e0539b6c1 100644 --- a/docs/src/architectures/gap.rst +++ b/docs/src/architectures/gap.rst @@ -29,8 +29,8 @@ of the repository: This will install the package with the GAP dependencies. -Hyperparameters ---------------- +Architecture Hyperparameters +---------------------------- :param name: ``experimental.gap`` @@ -38,15 +38,15 @@ model ##### soap ^^^^ -:param cutoff: Spherical cutoff (Å) to use for atomic environments -:param max_radial: Number of radial basis function to use +: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 + degree of spherical harmonics. Default 6 :param atomic_gaussian_width: Width of the atom-centered gaussian creating the atomic - density + 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. + 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 @@ -58,7 +58,7 @@ soap cutoff_function: Step: - - ``ShiftedCosine``: Shifted cosine switching function ``f(r) = 1/2 * (1 + cos(π (r + - ``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: @@ -79,7 +79,7 @@ soap radial_scaling: None: - - ``Willatt2018`` Use a long-range algebraic decay and smooth behavior at :math:`r + - ``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. @@ -99,11 +99,11 @@ soap krr ^^^^ -:param degree: degree of the polynomial kernel -:param num_sparse_points: number of pseudo points to select (by farthest point sampling) +: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 -:param regularizer_forces: value of the forces regularizer +:param regularizer: value of the energy regularizer. Default 0.001 +:param regularizer_forces: value of the forces regularizer. Default null diff --git a/src/metatensor/models/experimental/gap/default-hypers.yaml b/src/metatensor/models/experimental/gap/default-hypers.yaml index ad275a0c0..e106b8ccf 100644 --- a/src/metatensor/models/experimental/gap/default-hypers.yaml +++ b/src/metatensor/models/experimental/gap/default-hypers.yaml @@ -20,8 +20,8 @@ model: exponent: 7.0 krr: - degree: 2 # degree of the polynomial kernel - num_sparse_points: 500 # number of pseudo points to select, farthest point sampling is used + degree: 2 + num_sparse_points: 500 training: regularizer: 0.001 From 822e02fb20d9e6dc941d1986d4ef1a0878e9af4b Mon Sep 17 00:00:00 2001 From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com> Date: Tue, 4 Jun 2024 10:37:02 +0200 Subject: [PATCH 12/16] linter + comments --- docs/src/architectures/gap.rst | 17 ++++++++++------- .../models/experimental/gap/tests/__init__.py | 8 ++------ .../experimental/gap/tests/test_torchscript.py | 2 +- .../models/experimental/gap/trainer.py | 5 ++--- 4 files changed, 15 insertions(+), 17 deletions(-) diff --git a/docs/src/architectures/gap.rst b/docs/src/architectures/gap.rst index e0539b6c1..155297243 100644 --- a/docs/src/architectures/gap.rst +++ b/docs/src/architectures/gap.rst @@ -29,8 +29,8 @@ of the repository: This will install the package with the GAP dependencies. -Architecture Hyperparameters ----------------------------- +Architecture Hyperparameters +---------------------------- :param name: ``experimental.gap`` @@ -41,7 +41,7 @@ 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 + 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 @@ -58,8 +58,9 @@ soap 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 + - ``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 @@ -79,7 +80,8 @@ soap radial_scaling: None: - - ``Willatt2018`` (Default value): Use a long-range algebraic decay and smooth behavior at :math:`r + - ``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. @@ -100,7 +102,8 @@ soap 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 +:param num_sparse_points: number of pseudo points to select + (by farthest point sampling). Default 500 training: ^^^^^^^^^ diff --git a/src/metatensor/models/experimental/gap/tests/__init__.py b/src/metatensor/models/experimental/gap/tests/__init__.py index 48d5a22ec..66cee20c1 100644 --- a/src/metatensor/models/experimental/gap/tests/__init__.py +++ b/src/metatensor/models/experimental/gap/tests/__init__.py @@ -1,11 +1,7 @@ from pathlib import Path -DATASET_PATH = str( - Path(__file__).parent.resolve() - / "../../../../../../tests/resources/qm9_reduced_100.xyz" -) +DATASET_PATH = str(Path(__file__).parents[6] / "tests/resources/qm9_reduced_100.xyz") DATASET_ETHANOL_PATH = str( - Path(__file__).parent.resolve() - / "../../../../../../tests/resources/ethanol_reduced_100.xyz" + Path(__file__).parents[6] / "tests/resources/ethanol_reduced_100.xyz" ) diff --git a/src/metatensor/models/experimental/gap/tests/test_torchscript.py b/src/metatensor/models/experimental/gap/tests/test_torchscript.py index 2c69f6464..a4392fdb3 100644 --- a/src/metatensor/models/experimental/gap/tests/test_torchscript.py +++ b/src/metatensor/models/experimental/gap/tests/test_torchscript.py @@ -50,7 +50,7 @@ def test_torchscript(): hypers = DEFAULT_HYPERS.copy() gap = GAP(DEFAULT_HYPERS["model"], dataset_info) trainer = Trainer(hypers["training"]) - gap = trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".") + trainer.train(gap, [torch.device("cpu")], [dataset], [dataset], ".") scripted_gap = torch.jit.script(gap) ref_output = gap.forward(systems[:5], {"mtm::U0": gap.outputs["mtm::U0"]}) diff --git a/src/metatensor/models/experimental/gap/trainer.py b/src/metatensor/models/experimental/gap/trainer.py index 236b83932..cfaae9edd 100644 --- a/src/metatensor/models/experimental/gap/trainer.py +++ b/src/metatensor/models/experimental/gap/trainer.py @@ -32,8 +32,7 @@ def train( # checks assert devices == [torch.device("cpu")] dtype = train_datasets[0][0]["system"].positions.dtype - if dtype != torch.float64: - raise ValueError("GAP only supports float64") + assert dtype == torch.float64 target_name = next(iter(model.dataset_info.targets.keys())) if len(train_datasets) != 1: raise ValueError("GAP only supports a single training dataset") @@ -143,4 +142,4 @@ def train( model._subset_of_regressors.export_torch_script_model() ) - return model + # return model From 2c73d01bf0644b714ac6553c9fa32e731345ca1a Mon Sep 17 00:00:00 2001 From: Davide Tisi Date: Tue, 4 Jun 2024 10:46:37 +0200 Subject: [PATCH 13/16] Update docs/src/architectures/gap.rst Co-authored-by: Philip Loche --- docs/src/architectures/gap.rst | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/src/architectures/gap.rst b/docs/src/architectures/gap.rst index 155297243..476ab434b 100644 --- a/docs/src/architectures/gap.rst +++ b/docs/src/architectures/gap.rst @@ -109,4 +109,9 @@ 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 From 2ca5ae8605af153cc1b52091462a99d2b0023687 Mon Sep 17 00:00:00 2001 From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com> Date: Tue, 4 Jun 2024 11:07:35 +0200 Subject: [PATCH 14/16] fix doc --- docs/src/architectures/gap.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/src/architectures/gap.rst b/docs/src/architectures/gap.rst index 476ab434b..2f7347d24 100644 --- a/docs/src/architectures/gap.rst +++ b/docs/src/architectures/gap.rst @@ -109,9 +109,13 @@ 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 + + From 6bed07abcbba81cd36603b3d952dbf377d2b02f8 Mon Sep 17 00:00:00 2001 From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com> Date: Tue, 4 Jun 2024 11:18:02 +0200 Subject: [PATCH 15/16] linter --- docs/src/architectures/gap.rst | 1 - src/metatensor/models/experimental/gap/trainer.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/docs/src/architectures/gap.rst b/docs/src/architectures/gap.rst index 2f7347d24..002f47320 100644 --- a/docs/src/architectures/gap.rst +++ b/docs/src/architectures/gap.rst @@ -118,4 +118,3 @@ The default hyperparameters for the GAP model are: .. literalinclude:: ../../../src/metatensor/models/experimental/gap/default-hypers.yaml :language: yaml - diff --git a/src/metatensor/models/experimental/gap/trainer.py b/src/metatensor/models/experimental/gap/trainer.py index cfaae9edd..0419d67e8 100644 --- a/src/metatensor/models/experimental/gap/trainer.py +++ b/src/metatensor/models/experimental/gap/trainer.py @@ -141,5 +141,3 @@ def train( model._subset_of_regressors_torch = ( model._subset_of_regressors.export_torch_script_model() ) - - # return model From 2cdd5638ce276c0b9a7cf4fd71de85591bfe3c4a Mon Sep 17 00:00:00 2001 From: Davide Tisi <47503434+DavideTisi@users.noreply.github.com> Date: Tue, 4 Jun 2024 11:21:10 +0200 Subject: [PATCH 16/16] atomic_types --- src/metatensor/models/experimental/gap/model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/metatensor/models/experimental/gap/model.py b/src/metatensor/models/experimental/gap/model.py index b1d8f46fd..500d092ef 100644 --- a/src/metatensor/models/experimental/gap/model.py +++ b/src/metatensor/models/experimental/gap/model.py @@ -64,7 +64,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None: for key, value in dataset_info.targets.items() } - self.all_species = dataset_info.atomic_types + self.atomic_types = sorted(dataset_info.atomic_types) self.hypers = model_hypers # creates a composition weight tensor that can be directly indexed by species, @@ -72,7 +72,7 @@ def __init__(self, model_hypers: Dict, dataset_info: DatasetInfo) -> None: # set_composition_weights (recommended for better accuracy) n_outputs = len(dataset_info.targets) self.register_buffer( - "composition_weights", torch.zeros((n_outputs, max(self.all_species) + 1)) + "composition_weights", torch.zeros((n_outputs, max(self.atomic_types) + 1)) ) # buffers cannot be indexed by strings (torchscript), so we create a single # tensor for all output. Due to this, we need to slice the tensor when we use