diff --git a/.readthedocs.yml b/.readthedocs.yml index ca1e81dc7..10f1df129 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -26,5 +26,5 @@ sphinx: python: install: - method: pip - path: . + path: .[soap-bpnn] - requirements: docs/requirements.txt diff --git a/pyproject.toml b/pyproject.toml index 8bf4a98ed..beecac57f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,6 @@ dependencies = [ "ase", "torch", "hydra-core", - "rascaline-torch @ git+https://github.com/luthaf/rascaline#subdirectory=python/rascaline-torch", "metatensor-operations==0.2.1", "metatensor-torch==0.3.0", "metatensor-learn==0.2.1", @@ -55,7 +54,9 @@ requires = [ build-backend = "setuptools.build_meta" [project.optional-dependencies] -soap-bpnn = [] +soap-bpnn = [ + "rascaline-torch @ git+https://github.com/luthaf/rascaline@ae05064#subdirectory=python/rascaline-torch", +] alchemical-model = [ "torch_alchemical @ git+https://github.com/abmazitov/torch_alchemical.git@357a01f", ] diff --git a/src/metatensor/models/experimental/alchemical_model/tests/test_functionality.py b/src/metatensor/models/experimental/alchemical_model/tests/test_functionality.py index f7d602f63..d3d09f589 100644 --- a/src/metatensor/models/experimental/alchemical_model/tests/test_functionality.py +++ b/src/metatensor/models/experimental/alchemical_model/tests/test_functionality.py @@ -1,5 +1,4 @@ import ase -import rascaline.torch import torch from metatensor.torch.atomistic import ( MetatensorAtomisticModel, @@ -7,6 +6,7 @@ ModelEvaluationOptions, ModelMetadata, ModelOutput, + systems_to_torch, ) from metatensor.models.experimental.alchemical_model import DEFAULT_HYPERS, Model @@ -31,7 +31,7 @@ def test_prediction_subset(): alchemical_model = Model(capabilities, DEFAULT_HYPERS["model"]) system = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = rascaline.torch.systems_to_torch(system).to(torch.get_default_dtype()) + system = systems_to_torch(system, dtype=torch.get_default_dtype()) system = get_system_with_neighbors_lists( system, alchemical_model.requested_neighbors_lists() ) diff --git a/src/metatensor/models/experimental/alchemical_model/tests/test_invariance.py b/src/metatensor/models/experimental/alchemical_model/tests/test_invariance.py index 8ac96ae25..caaff9f6e 100644 --- a/src/metatensor/models/experimental/alchemical_model/tests/test_invariance.py +++ b/src/metatensor/models/experimental/alchemical_model/tests/test_invariance.py @@ -1,7 +1,6 @@ import copy import ase.io -import rascaline.torch import torch from metatensor.torch.atomistic import ( MetatensorAtomisticModel, @@ -9,6 +8,7 @@ ModelEvaluationOptions, ModelMetadata, ModelOutput, + systems_to_torch, ) from metatensor.models.experimental.alchemical_model import DEFAULT_HYPERS, Model @@ -35,13 +35,11 @@ def test_rotational_invariance(): system = ase.io.read(DATASET_PATH) original_system = copy.deepcopy(system) system.rotate(48, "y") - original_system = rascaline.torch.systems_to_torch(original_system).to( - torch.get_default_dtype() - ) + original_system = systems_to_torch(original_system, dtype=torch.get_default_dtype()) original_system = get_system_with_neighbors_lists( original_system, alchemical_model.requested_neighbors_lists() ) - system = rascaline.torch.systems_to_torch(system).to(torch.get_default_dtype()) + system = systems_to_torch(system, dtype=torch.get_default_dtype()) system = get_system_with_neighbors_lists( system, alchemical_model.requested_neighbors_lists() ) diff --git a/src/metatensor/models/experimental/alchemical_model/tests/test_regression.py b/src/metatensor/models/experimental/alchemical_model/tests/test_regression.py index de3ded498..817432432 100644 --- a/src/metatensor/models/experimental/alchemical_model/tests/test_regression.py +++ b/src/metatensor/models/experimental/alchemical_model/tests/test_regression.py @@ -2,7 +2,6 @@ import ase.io import numpy as np -import rascaline.torch import torch from metatensor.learn.data import Dataset from metatensor.torch.atomistic import ( @@ -11,6 +10,7 @@ ModelEvaluationOptions, ModelMetadata, ModelOutput, + systems_to_torch, ) from omegaconf import OmegaConf @@ -46,8 +46,7 @@ def test_regression_init(): # Predict on the first five systems systems = ase.io.read(DATASET_PATH, ":5") systems = [ - rascaline.torch.systems_to_torch(system).to(torch.get_default_dtype()) - for system in systems + systems_to_torch(system, dtype=torch.get_default_dtype()) for system in systems ] systems = [ get_system_with_neighbors_lists( diff --git a/src/metatensor/models/experimental/alchemical_model/tests/test_torchscript.py b/src/metatensor/models/experimental/alchemical_model/tests/test_torchscript.py index 057f657c0..56d6da87c 100644 --- a/src/metatensor/models/experimental/alchemical_model/tests/test_torchscript.py +++ b/src/metatensor/models/experimental/alchemical_model/tests/test_torchscript.py @@ -20,7 +20,7 @@ def test_torchscript(): ) }, ) - alchemical_model = Model(capabilities, DEFAULT_HYPERS["model"]).to(torch.float64) + alchemical_model = Model(capabilities, DEFAULT_HYPERS["model"]) torch.jit.script( alchemical_model, {"energy": alchemical_model.capabilities.outputs["energy"]} ) @@ -39,7 +39,7 @@ def test_torchscript_save(): ) }, ) - alchemical_model = Model(capabilities, DEFAULT_HYPERS["model"]).to(torch.float64) + alchemical_model = Model(capabilities, DEFAULT_HYPERS["model"]) torch.jit.save( torch.jit.script( alchemical_model, diff --git a/src/metatensor/models/experimental/pet/tests/test_functionality.py b/src/metatensor/models/experimental/pet/tests/test_functionality.py index e14c3c188..93aba4c05 100644 --- a/src/metatensor/models/experimental/pet/tests/test_functionality.py +++ b/src/metatensor/models/experimental/pet/tests/test_functionality.py @@ -1,12 +1,11 @@ import ase -import rascaline.torch -import torch from metatensor.torch.atomistic import ( MetatensorAtomisticModel, ModelCapabilities, ModelEvaluationOptions, ModelMetadata, ModelOutput, + systems_to_torch, ) from metatensor.models.experimental.pet import DEFAULT_HYPERS, Model @@ -29,11 +28,9 @@ def test_prediction_subset(): supported_devices=["cpu"], ) - model = Model(capabilities, DEFAULT_HYPERS["ARCHITECTURAL_HYPERS"]).to( - torch.float64 - ) + model = Model(capabilities, DEFAULT_HYPERS["ARCHITECTURAL_HYPERS"]) structure = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) - system = rascaline.torch.systems_to_torch(structure) + system = systems_to_torch(structure) system = get_system_with_neighbors_lists(system, model.requested_neighbors_lists()) evaluation_options = ModelEvaluationOptions( diff --git a/src/metatensor/models/experimental/soap_bpnn/tests/test_functionality.py b/src/metatensor/models/experimental/soap_bpnn/tests/test_functionality.py index 9bd85fa35..84ca34049 100644 --- a/src/metatensor/models/experimental/soap_bpnn/tests/test_functionality.py +++ b/src/metatensor/models/experimental/soap_bpnn/tests/test_functionality.py @@ -1,8 +1,7 @@ import ase import metatensor.torch -import rascaline.torch import torch -from metatensor.torch.atomistic import ModelCapabilities, ModelOutput +from metatensor.torch.atomistic import ModelCapabilities, ModelOutput, systems_to_torch from metatensor.models.experimental.soap_bpnn import DEFAULT_HYPERS, Model @@ -26,7 +25,7 @@ def test_prediction_subset_elements(): system = ase.Atoms("O2", positions=[[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) soap_bpnn( - [rascaline.torch.systems_to_torch(system).to(torch.get_default_dtype())], + [systems_to_torch(system, dtype=torch.get_default_dtype())], {"energy": soap_bpnn.capabilities.outputs["energy"]}, ) @@ -56,11 +55,7 @@ def test_prediction_subset_atoms(): ) energy_monomer = soap_bpnn( - [ - rascaline.torch.systems_to_torch(system_monomer).to( - torch.get_default_dtype() - ) - ], + [systems_to_torch(system_monomer).to(torch.get_default_dtype())], {"energy": soap_bpnn.capabilities.outputs["energy"]}, ) @@ -82,20 +77,12 @@ def test_prediction_subset_atoms(): ) energy_dimer = soap_bpnn( - [ - rascaline.torch.systems_to_torch(system_far_away_dimer).to( - torch.get_default_dtype() - ) - ], + [systems_to_torch(system_far_away_dimer).to(torch.get_default_dtype())], {"energy": soap_bpnn.capabilities.outputs["energy"]}, ) energy_monomer_in_dimer = soap_bpnn( - [ - rascaline.torch.systems_to_torch(system_far_away_dimer).to( - torch.get_default_dtype() - ) - ], + [systems_to_torch(system_far_away_dimer).to(torch.get_default_dtype())], {"energy": soap_bpnn.capabilities.outputs["energy"]}, selected_atoms=selection_labels, ) diff --git a/src/metatensor/models/experimental/soap_bpnn/tests/test_invariance.py b/src/metatensor/models/experimental/soap_bpnn/tests/test_invariance.py index f63df4466..ac4544c53 100644 --- a/src/metatensor/models/experimental/soap_bpnn/tests/test_invariance.py +++ b/src/metatensor/models/experimental/soap_bpnn/tests/test_invariance.py @@ -1,9 +1,8 @@ import copy import ase.io -import rascaline.torch import torch -from metatensor.torch.atomistic import ModelCapabilities, ModelOutput +from metatensor.torch.atomistic import ModelCapabilities, ModelOutput, systems_to_torch from metatensor.models.experimental.soap_bpnn import DEFAULT_HYPERS, Model @@ -23,18 +22,18 @@ def test_rotational_invariance(): ) }, ) - soap_bpnn = Model(capabilities, DEFAULT_HYPERS["model"]).to(torch.float64) + soap_bpnn = Model(capabilities, DEFAULT_HYPERS["model"]) system = ase.io.read(DATASET_PATH) original_system = copy.deepcopy(system) system.rotate(48, "y") original_output = soap_bpnn( - [rascaline.torch.systems_to_torch(original_system)], + [systems_to_torch(original_system)], {"energy": soap_bpnn.capabilities.outputs["energy"]}, ) rotated_output = soap_bpnn( - [rascaline.torch.systems_to_torch(system)], + [systems_to_torch(system)], {"energy": soap_bpnn.capabilities.outputs["energy"]}, ) diff --git a/src/metatensor/models/experimental/soap_bpnn/tests/test_regression.py b/src/metatensor/models/experimental/soap_bpnn/tests/test_regression.py index a78b0c75a..bce7df48e 100644 --- a/src/metatensor/models/experimental/soap_bpnn/tests/test_regression.py +++ b/src/metatensor/models/experimental/soap_bpnn/tests/test_regression.py @@ -2,10 +2,9 @@ import ase.io import numpy as np -import rascaline.torch import torch from metatensor.learn.data import Dataset -from metatensor.torch.atomistic import ModelCapabilities, ModelOutput +from metatensor.torch.atomistic import ModelCapabilities, ModelOutput, systems_to_torch from omegaconf import OmegaConf from metatensor.models.experimental.soap_bpnn import DEFAULT_HYPERS, Model, train @@ -41,7 +40,7 @@ def test_regression_init(): output = soap_bpnn( [ - rascaline.torch.systems_to_torch(system).to(torch.get_default_dtype()) + systems_to_torch(system, dtype=torch.get_default_dtype()) for system in systems ], {"U0": soap_bpnn.capabilities.outputs["U0"]}, diff --git a/src/metatensor/models/experimental/soap_bpnn/tests/test_torchscript.py b/src/metatensor/models/experimental/soap_bpnn/tests/test_torchscript.py index 2e0f23471..910ae5660 100644 --- a/src/metatensor/models/experimental/soap_bpnn/tests/test_torchscript.py +++ b/src/metatensor/models/experimental/soap_bpnn/tests/test_torchscript.py @@ -17,7 +17,7 @@ def test_torchscript(): ) }, ) - soap_bpnn = Model(capabilities, DEFAULT_HYPERS["model"]).to(torch.float64) + soap_bpnn = Model(capabilities, DEFAULT_HYPERS["model"]) torch.jit.script(soap_bpnn, {"energy": soap_bpnn.capabilities.outputs["energy"]}) @@ -34,7 +34,7 @@ def test_torchscript_save(): ) }, ) - soap_bpnn = Model(capabilities, DEFAULT_HYPERS["model"]).to(torch.float64) + soap_bpnn = Model(capabilities, DEFAULT_HYPERS["model"]) torch.jit.save( torch.jit.script( soap_bpnn, {"energy": soap_bpnn.capabilities.outputs["energy"]} diff --git a/src/metatensor/models/utils/composition.py b/src/metatensor/models/utils/composition.py index 50682a3c3..26b352f93 100644 --- a/src/metatensor/models/utils/composition.py +++ b/src/metatensor/models/utils/composition.py @@ -1,6 +1,5 @@ -from typing import List, Tuple +from typing import List, Tuple, Union -import rascaline.torch import torch from metatensor.learn.data.dataset import _BaseDataset from metatensor.torch import Labels, TensorBlock, TensorMap @@ -9,7 +8,7 @@ def calculate_composition_weights( - datasets: _BaseDataset, property: str + datasets: Union[_BaseDataset, List[_BaseDataset]], property: str ) -> Tuple[torch.Tensor, List[int]]: """Calculate the composition weights for a dataset. @@ -19,34 +18,33 @@ def calculate_composition_weights( :returns: Composition weights for the dataset, as well as the list of species that the weights correspond to. """ + if not isinstance(datasets, list): + datasets = [datasets] + + species = get_all_species(datasets) + # note that this is sorted, and the composition weights are sorted + # as well, because the species are sorted in the composition features - # Get the target for each system in the dataset - # TODO: the dataset will be iterable once metatensor PR #500 merged. targets = torch.stack( [ - dataset[sample_id]._asdict()[property].block().values + sample._asdict()[property].block().values for dataset in datasets - for sample_id in range(len(dataset)) + for sample in dataset ] ) + targets = targets.squeeze(dim=(1, 2)) # remove component and property dimensions - # Get the composition for each system in the dataset - composition_calculator = rascaline.torch.AtomicComposition(per_system=True) - # TODO: the dataset will be iterable once metatensor PR #500 merged. - composition_features = composition_calculator.compute( - [ - dataset[sample_id]._asdict()["system"] - for dataset in datasets - for sample_id in range(len(dataset)) - ] - ) - composition_features = composition_features.keys_to_properties("center_type") - composition_features = composition_features.block().values + structure_list = [ + sample._asdict()["system"] for dataset in datasets for sample in dataset + ] - targets = targets.squeeze(dim=(1, 2)) # remove component and property dimensions + dtype = structure_list[0].positions.dtype + composition_features = torch.empty((len(structure_list), len(species)), dtype=dtype) + for i, structure in enumerate(structure_list): + for j, s in enumerate(species): + composition_features[i, j] = torch.sum(structure.types == s) regularizer = 1e-20 - while regularizer: if regularizer > 1e5: raise RuntimeError( @@ -69,10 +67,6 @@ def calculate_composition_weights( except torch._C._LinAlgError: regularizer *= 10.0 - species = get_all_species(datasets) - # note that this is sorted, and the composition weights are sorted - # as well, because the species are sorted in the composition features - return solution, species diff --git a/src/metatensor/models/utils/neighbors_lists.py b/src/metatensor/models/utils/neighbors_lists.py index ac27ef093..003679433 100644 --- a/src/metatensor/models/utils/neighbors_lists.py +++ b/src/metatensor/models/utils/neighbors_lists.py @@ -1,6 +1,6 @@ from typing import List -import ase +import ase.neighborlist import torch from metatensor.torch import Labels, TensorBlock from metatensor.torch.atomistic import ( diff --git a/tests/utils/test_composition.py b/tests/utils/test_composition.py new file mode 100644 index 000000000..9513ead81 --- /dev/null +++ b/tests/utils/test_composition.py @@ -0,0 +1,71 @@ +from pathlib import Path + +import torch +from metatensor.learn import Dataset +from metatensor.torch import Labels, TensorBlock, TensorMap +from metatensor.torch.atomistic import System + +from metatensor.models.utils.composition import calculate_composition_weights + + +RESOURCES_PATH = Path(__file__).parent.resolve() / ".." / "resources" + + +def test_calculate_composition_weights(): + """Test the calculation of composition weights.""" + + # Here we use three synthetic structures: + # - O atom, with an energy of 1.0 + # - H2O molecule, with an energy of 5.0 + # - H4O2 molecule, with an energy of 10.0 + # The expected composition weights are 2.0 for H and 1.0 for O. + + systems = [ + System( + positions=torch.tensor([[0.0, 0.0, 0.0]]), + types=torch.tensor([8]), + cell=torch.eye(3), + ), + System( + positions=torch.tensor([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]), + types=torch.tensor([1, 1, 8]), + cell=torch.eye(3), + ), + System( + positions=torch.tensor( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [1.0, 0.0, 1.0], + [0.0, 1.0, 1.0], + ] + ), + types=torch.tensor([1, 1, 8, 1, 1, 8]), + cell=torch.eye(3), + ), + ] + energies = [1.0, 5.0, 10.0] + energies = [ + TensorMap( + keys=Labels(names=["_"], values=torch.tensor([[0]])), + blocks=[ + TensorBlock( + values=torch.tensor([[e]]), + samples=Labels(names=["system"], values=torch.tensor([[i]])), + components=[], + properties=Labels(names=["energy"], values=torch.tensor([[0]])), + ) + ], + ) + for i, e in enumerate(energies) + ] + dataset = Dataset(system=systems, energy=energies) + + weights, species = calculate_composition_weights(dataset, "energy") + + assert len(weights) == len(species) + assert len(weights) == 2 + assert species == [1, 8] + assert torch.allclose(weights, torch.tensor([2.0, 1.0])) diff --git a/tests/utils/test_model_io.py b/tests/utils/test_model_io.py index 33bfbc853..443c1a59f 100644 --- a/tests/utils/test_model_io.py +++ b/tests/utils/test_model_io.py @@ -3,7 +3,7 @@ import metatensor.torch import pytest -import rascaline.torch +import torch from metatensor.torch.atomistic import ModelCapabilities, ModelOutput from metatensor.models.experimental import soap_bpnn @@ -35,10 +35,12 @@ def test_save_load_checkpoint(monkeypatch, tmp_path): ) model = soap_bpnn.Model(capabilities) - systems = read_systems(RESOURCES_PATH / "qm9_reduced_100.xyz") + systems = read_systems( + RESOURCES_PATH / "qm9_reduced_100.xyz", dtype=torch.get_default_dtype() + ) output_before_save = model( - rascaline.torch.systems_to_torch(systems), + systems, {"energy": model.capabilities.outputs["energy"]}, ) @@ -46,7 +48,7 @@ def test_save_load_checkpoint(monkeypatch, tmp_path): loaded_model = load_checkpoint("test_model.ckpt") output_after_load = loaded_model( - rascaline.torch.systems_to_torch(systems), + systems, {"energy": model.capabilities.outputs["energy"]}, ) diff --git a/tests/utils/test_output_gradient.py b/tests/utils/test_output_gradient.py index 98866a5a6..0159499a5 100644 --- a/tests/utils/test_output_gradient.py +++ b/tests/utils/test_output_gradient.py @@ -2,9 +2,8 @@ import metatensor.torch import pytest -import rascaline.torch import torch -from metatensor.torch.atomistic import ModelCapabilities, ModelOutput +from metatensor.torch.atomistic import ModelCapabilities, ModelOutput, System from metatensor.models.experimental import soap_bpnn from metatensor.models.utils.data import read_systems @@ -33,7 +32,14 @@ def test_forces(is_training): systems = read_systems( RESOURCES_PATH / "qm9_reduced_100.xyz", dtype=torch.get_default_dtype() )[:5] - systems = rascaline.torch.systems_to_torch(systems, positions_requires_grad=True) + systems = [ + System( + positions=system.positions.requires_grad_(True), + cell=system.cell, + types=system.types, + ) + for system in systems + ] output = model(systems, {"energy": model.capabilities.outputs["energy"]}) position_gradients = compute_gradient( output["energy"].block().values, @@ -43,7 +49,17 @@ def test_forces(is_training): forces = [-position_gradient for position_gradient in position_gradients] jitted_model = torch.jit.script(model) - systems = rascaline.torch.systems_to_torch(systems, positions_requires_grad=True) + systems = read_systems( + RESOURCES_PATH / "qm9_reduced_100.xyz", dtype=torch.get_default_dtype() + )[:5] + systems = [ + System( + positions=system.positions.requires_grad_(True), + cell=system.cell, + types=system.types, + ) + for system in systems + ] output = jitted_model(systems, {"energy": model.capabilities.outputs["energy"]}) jitted_position_gradients = compute_gradient( output["energy"].block().values, diff --git a/tox.ini b/tox.ini index 8de85355f..25755607c 100644 --- a/tox.ini +++ b/tox.ini @@ -61,6 +61,7 @@ deps = pytest pytest-cov changedir = tests +extras = soap-bpnn # this model is used in the package tests commands = pytest \ --cov={env_site_packages_dir}/metatensor --cov-append --cov-report= --import-mode=append \ @@ -108,7 +109,7 @@ description = Run PET tests with pytest passenv = * deps = pytest - pet @ git+https://github.com/serfg/pet.git@5668bda +extras = pet changedir = src/metatensor/models/experimental/pet/tests/ commands = pytest {[testenv]warning_options} {posargs} @@ -120,6 +121,8 @@ deps = allowlist_externals = bash commands = + # SOAP-BPNN is required for the documentation + pip install .[soap-bpnn] # Run .sh scripts in the example folder. #bash -c "set -e && cd {toxinidir}/examples/basic_usage && bash usage.sh" # Disable slow training at runtime for ase example to speed up CI for now.