diff --git a/pyproject.toml b/pyproject.toml index 7e3f2d703..bec22e930 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,9 +10,9 @@ authors = [{name = "metatrain developers"}] dependencies = [ "ase < 3.23.0", - "metatensor-learn==0.2.3", - "metatensor-operations==0.2.3", - "metatensor-torch==0.5.5", + "metatensor-learn==0.3.0", + "metatensor-operations==0.3.0", + "metatensor-torch==0.6.0", "jsonschema", "omegaconf", "python-hostlist", @@ -59,7 +59,7 @@ build-backend = "setuptools.build_meta" [project.optional-dependencies] soap-bpnn = [ - "rascaline-torch @ git+https://github.com/luthaf/rascaline@d181b28#subdirectory=python/rascaline-torch", + "rascaline-torch @ git+https://github.com/luthaf/rascaline@5326b6e#subdirectory=python/rascaline-torch", ] alchemical-model = [ "torch_alchemical @ git+https://github.com/abmazitov/torch_alchemical.git@51ff519", @@ -68,7 +68,7 @@ pet = [ "pet @ git+https://github.com/lab-cosmo/pet@7eddb2e", ] gap = [ - "rascaline-torch @ git+https://github.com/luthaf/rascaline@d181b28#subdirectory=python/rascaline-torch", + "rascaline-torch @ git+https://github.com/luthaf/rascaline@5326b6e#subdirectory=python/rascaline-torch", "skmatter", "metatensor-learn", "scipy", diff --git a/src/metatrain/experimental/alchemical_model/model.py b/src/metatrain/experimental/alchemical_model/model.py index 8ffbd8dc7..4e985dc75 100644 --- a/src/metatrain/experimental/alchemical_model/model.py +++ b/src/metatrain/experimental/alchemical_model/model.py @@ -78,6 +78,7 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]: NeighborListOptions( cutoff=self.cutoff, full_list=True, + strict=True, ) ] diff --git a/src/metatrain/experimental/alchemical_model/tests/test_exported.py b/src/metatrain/experimental/alchemical_model/tests/test_exported.py index 891983693..8b06c2d74 100644 --- a/src/metatrain/experimental/alchemical_model/tests/test_exported.py +++ b/src/metatrain/experimental/alchemical_model/tests/test_exported.py @@ -33,6 +33,7 @@ def test_to(device, dtype): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) requested_neighbor_lists = get_requested_neighbor_lists(exported) system = get_system_with_neighbor_lists(system, requested_neighbor_lists) diff --git a/src/metatrain/experimental/alchemical_model/tests/test_functionality.py b/src/metatrain/experimental/alchemical_model/tests/test_functionality.py index 7ee3331af..9f998f925 100644 --- a/src/metatrain/experimental/alchemical_model/tests/test_functionality.py +++ b/src/metatrain/experimental/alchemical_model/tests/test_functionality.py @@ -27,6 +27,7 @@ def test_prediction_subset_elements(): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) requested_neighbor_lists = get_requested_neighbor_lists(model) system = get_system_with_neighbor_lists(system, requested_neighbor_lists) diff --git a/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py b/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py index 03b7ef1df..7b560e786 100644 --- a/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py +++ b/src/metatrain/experimental/alchemical_model/tests/test_torch_alchemical_compatibility.py @@ -29,6 +29,7 @@ nl_options = NeighborListOptions( cutoff=5.0, full_list=True, + strict=True, ) systems = [get_system_with_neighbor_lists(system, [nl_options]) for system in systems] diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py index bf567dee1..d0d21d73f 100644 --- a/src/metatrain/experimental/pet/model.py +++ b/src/metatrain/experimental/pet/model.py @@ -73,6 +73,7 @@ def requested_neighbor_lists( NeighborListOptions( cutoff=self.cutoff, full_list=True, + strict=True, ) ] @@ -111,9 +112,7 @@ def forward( values=predictions, ) if selected_atoms is not None: - block = metatensor.torch.slice_block( - block, axis="samples", labels=selected_atoms - ) + block = metatensor.torch.slice_block(block, "samples", selected_atoms) output_tmap = TensorMap(keys=empty_labels, blocks=[block]) if not outputs[output_name].per_atom: output_tmap = metatensor.torch.sum_over_samples(output_tmap, "atom") diff --git a/src/metatrain/experimental/pet/tests/test_exported.py b/src/metatrain/experimental/pet/tests/test_exported.py index f67a15e4c..5ed9ecf6e 100644 --- a/src/metatrain/experimental/pet/tests/test_exported.py +++ b/src/metatrain/experimental/pet/tests/test_exported.py @@ -61,6 +61,7 @@ def test_to(device): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) requested_neighbor_lists = get_requested_neighbor_lists(exported) system = get_system_with_neighbor_lists(system, requested_neighbor_lists) diff --git a/src/metatrain/experimental/pet/tests/test_functionality.py b/src/metatrain/experimental/pet/tests/test_functionality.py index ddf527603..b0f5771ea 100644 --- a/src/metatrain/experimental/pet/tests/test_functionality.py +++ b/src/metatrain/experimental/pet/tests/test_functionality.py @@ -76,6 +76,7 @@ def test_prediction(): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) requested_neighbor_lists = get_requested_neighbor_lists(model) system = get_system_with_neighbor_lists(system, requested_neighbor_lists) @@ -126,6 +127,7 @@ def test_per_atom_predictions_functionality(): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) requested_neighbor_lists = get_requested_neighbor_lists(model) system = get_system_with_neighbor_lists(system, requested_neighbor_lists) @@ -177,6 +179,7 @@ def test_selected_atoms_functionality(): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) requested_neighbor_lists = get_requested_neighbor_lists(model) system = get_system_with_neighbor_lists(system, requested_neighbor_lists) diff --git a/src/metatrain/experimental/pet/tests/test_pet_compatibility.py b/src/metatrain/experimental/pet/tests/test_pet_compatibility.py index 04ecae1bc..4a2bdd71e 100644 --- a/src/metatrain/experimental/pet/tests/test_pet_compatibility.py +++ b/src/metatrain/experimental/pet/tests/test_pet_compatibility.py @@ -59,7 +59,7 @@ def test_batch_dicts_compatibility(cutoff): structure = ase.io.read(DATASET_PATH) atomic_types = sorted(set(structure.numbers)) system = systems_to_torch(structure) - options = NeighborListOptions(cutoff=cutoff, full_list=True) + options = NeighborListOptions(cutoff=cutoff, full_list=True, strict=True) system = get_system_with_neighbor_lists(system, [options]) ARCHITECTURAL_HYPERS = Hypers(DEFAULT_HYPERS["model"]) @@ -121,7 +121,7 @@ def test_predictions_compatibility(cutoff): model.set_trained_model(raw_pet) system = systems_to_torch(structure) - options = NeighborListOptions(cutoff=cutoff, full_list=True) + options = NeighborListOptions(cutoff=cutoff, full_list=True, strict=True) system = get_system_with_neighbor_lists(system, [options]) evaluation_options = ModelEvaluationOptions( diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py index 63242161e..349e89cfd 100644 --- a/src/metatrain/experimental/soap_bpnn/tests/test_exported.py +++ b/src/metatrain/experimental/soap_bpnn/tests/test_exported.py @@ -33,6 +33,7 @@ def test_to(device, dtype): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) requested_neighbor_lists = get_requested_neighbor_lists(exported) system = get_system_with_neighbor_lists(system, requested_neighbor_lists) diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py index 25dd250b6..b1fb4f1d3 100644 --- a/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py +++ b/src/metatrain/experimental/soap_bpnn/tests/test_functionality.py @@ -28,6 +28,7 @@ def test_prediction_subset_elements(): types=torch.tensor([6, 6]), positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) model( [system], @@ -56,6 +57,7 @@ def test_prediction_subset_atoms(): [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0]], ), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) energy_monomer = model( @@ -76,6 +78,7 @@ def test_prediction_subset_atoms(): ], ), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) selection_labels = metatensor.torch.Labels( @@ -119,6 +122,7 @@ def test_output_last_layer_features(): [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], ), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) # last-layer features per atom: @@ -190,6 +194,7 @@ def test_output_per_atom(): [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]], ), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) outputs = model( diff --git a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py index 4d6b89898..6be11f582 100644 --- a/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py +++ b/src/metatrain/experimental/soap_bpnn/tests/test_torchscript.py @@ -26,6 +26,7 @@ def test_torchscript(): [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]] ), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) model( [system], @@ -52,6 +53,7 @@ def test_torchscript_with_identity(): [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 2.0], [0.0, 0.0, 3.0]] ), cell=torch.zeros(3, 3), + pbc=torch.tensor([False, False, False]), ) model( [system], diff --git a/src/metatrain/utils/additive/zbl.py b/src/metatrain/utils/additive/zbl.py index edf2ec7b2..cec2ee96e 100644 --- a/src/metatrain/utils/additive/zbl.py +++ b/src/metatrain/utils/additive/zbl.py @@ -251,6 +251,7 @@ def requested_neighbor_lists(self) -> List[NeighborListOptions]: NeighborListOptions( cutoff=self.cutoff_radius, full_list=True, + strict=True, ) ] diff --git a/src/metatrain/utils/evaluate_model.py b/src/metatrain/utils/evaluate_model.py index 48447b879..5b4500a21 100644 --- a/src/metatrain/utils/evaluate_model.py +++ b/src/metatrain/utils/evaluate_model.py @@ -278,6 +278,7 @@ def _prepare_system( positions=system.positions @ strain, cell=system.cell @ strain, types=system.types, + pbc=system.pbc, ) else: if positions_grad: @@ -285,6 +286,7 @@ def _prepare_system( positions=system.positions.detach().clone().requires_grad_(True), cell=system.cell, types=system.types, + pbc=system.pbc, ) strain = None else: @@ -292,6 +294,7 @@ def _prepare_system( positions=system.positions, cell=system.cell, types=system.types, + pbc=system.pbc, ) strain = None diff --git a/src/metatrain/utils/metrics.py b/src/metatrain/utils/metrics.py index 357b567d0..29ace5104 100644 --- a/src/metatrain/utils/metrics.py +++ b/src/metatrain/utils/metrics.py @@ -8,7 +8,7 @@ class RMSEAccumulator: """Accumulates the RMSE between predictions and targets for an arbitrary number of keys, each corresponding to one target.""" - def __init__(self): + def __init__(self) -> None: """Initialize the accumulator.""" self.information: Dict[str, Tuple[float, int]] = {} @@ -91,7 +91,7 @@ class MAEAccumulator: """Accumulates the MAE between predictions and targets for an arbitrary number of keys, each corresponding to one target.""" - def __init__(self): + def __init__(self) -> None: """Initialize the accumulator.""" self.information: Dict[str, Tuple[float, int]] = {} diff --git a/tests/utils/data/test_system_to_ase.py b/tests/utils/data/test_system_to_ase.py index e6062b52d..3d66f7e7c 100644 --- a/tests/utils/data/test_system_to_ase.py +++ b/tests/utils/data/test_system_to_ase.py @@ -11,6 +11,7 @@ def test_system_to_ase(): positions=torch.tensor([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]), types=torch.tensor([1, 8]), cell=torch.tensor([[10.0, 0.0, 0.0], [0.0, 10.0, 0.0], [0.0, 0.0, 10.0]]), + pbc=torch.tensor([True, True, True]), ) # Convert the system to an ASE atoms object diff --git a/tests/utils/data/test_target_writers.py b/tests/utils/data/test_target_writers.py index 190d37fbe..525ccd446 100644 --- a/tests/utils/data/test_target_writers.py +++ b/tests/utils/data/test_target_writers.py @@ -20,6 +20,7 @@ def systems_capabilities_predictions( types=torch.tensor([1, 1]), positions=torch.tensor([[0, 0, 0], [0, 0, 0.74]]), cell=cell, + pbc=torch.logical_not(torch.all(cell == 0, dim=1)), ), ] diff --git a/tests/utils/test_additive.py b/tests/utils/test_additive.py index fd2179e5e..3daf3c980 100644 --- a/tests/utils/test_additive.py +++ b/tests/utils/test_additive.py @@ -33,6 +33,7 @@ def test_composition_model_train(): positions=torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64), types=torch.tensor([8]), cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), ), System( positions=torch.tensor( @@ -40,6 +41,7 @@ def test_composition_model_train(): ), types=torch.tensor([1, 1, 8]), cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), ), System( positions=torch.tensor( @@ -55,6 +57,7 @@ def test_composition_model_train(): ), types=torch.tensor([1, 1, 8, 1, 1, 8]), cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), ), ] energies = [1.0, 5.0, 10.0] @@ -200,6 +203,7 @@ def test_composition_model_torchscript(tmpdir): positions=torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64), types=torch.tensor([8]), cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), ) composition_model = CompositionModel( @@ -289,6 +293,7 @@ def test_composition_model_missing_types(): positions=torch.tensor([[0.0, 0.0, 0.0]], dtype=torch.float64), types=torch.tensor([8]), cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), ), System( positions=torch.tensor( @@ -296,6 +301,7 @@ def test_composition_model_missing_types(): ), types=torch.tensor([1, 1, 8]), cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), ), System( positions=torch.tensor( @@ -311,6 +317,7 @@ def test_composition_model_missing_types(): ), types=torch.tensor([1, 1, 8, 1, 1, 8]), cell=torch.eye(3, dtype=torch.float64), + pbc=torch.tensor([True, True, True]), ), ] energies = [1.0, 5.0, 10.0] diff --git a/tests/utils/test_neighbor_list.py b/tests/utils/test_neighbor_list.py index 9c35e05d5..768114d6c 100644 --- a/tests/utils/test_neighbor_list.py +++ b/tests/utils/test_neighbor_list.py @@ -14,9 +14,9 @@ def test_attach_neighbor_lists(): systems = read_systems_ase(filename) requested_neighbor_lists = [ - NeighborListOptions(cutoff=4.0, full_list=True), - NeighborListOptions(cutoff=5.0, full_list=False), - NeighborListOptions(cutoff=6.0, full_list=True), + NeighborListOptions(cutoff=4.0, full_list=True, strict=True), + NeighborListOptions(cutoff=5.0, full_list=False, strict=True), + NeighborListOptions(cutoff=6.0, full_list=True, strict=True), ] new_system = get_system_with_neighbor_lists(systems[0], requested_neighbor_lists) @@ -25,7 +25,7 @@ def test_attach_neighbor_lists(): assert requested_neighbor_lists[1] in new_system.known_neighbor_lists() assert requested_neighbor_lists[2] in new_system.known_neighbor_lists() - extraneous_nl = NeighborListOptions(cutoff=5.0, full_list=True) + extraneous_nl = NeighborListOptions(cutoff=5.0, full_list=True, strict=True) assert extraneous_nl not in new_system.known_neighbor_lists() for nl_options in new_system.known_neighbor_lists(): diff --git a/tests/utils/test_output_gradient.py b/tests/utils/test_output_gradient.py index 7bad9279e..513d43930 100644 --- a/tests/utils/test_output_gradient.py +++ b/tests/utils/test_output_gradient.py @@ -32,6 +32,7 @@ def test_forces(is_training): positions=system.positions.requires_grad_(True), cell=system.cell, types=system.types, + pbc=system.pbc, ) for system in systems ] @@ -50,6 +51,7 @@ def test_forces(is_training): positions=system.positions.requires_grad_(True), cell=system.cell, types=system.types, + pbc=system.pbc, ) for system in systems ] @@ -96,6 +98,7 @@ def test_virial(is_training): positions=system.positions @ strain, cell=system.cell @ strain, types=system.types, + pbc=system.pbc, ) for system, strain in zip(systems, strains) ] @@ -121,6 +124,7 @@ def test_virial(is_training): positions=system.positions @ strain, cell=system.cell @ strain, types=system.types, + pbc=system.pbc, ) for system, strain in zip(systems, strains) ] @@ -170,6 +174,7 @@ def test_both(is_training): positions=system.positions @ strain, cell=system.cell @ strain, types=system.types, + pbc=system.pbc, ) for system, strain in zip(systems, strains) ] @@ -193,6 +198,7 @@ def test_both(is_training): positions=system.positions @ strain, cell=system.cell @ strain, types=system.types, + pbc=system.pbc, ) for system, strain in zip(systems, strains) ] diff --git a/tests/utils/test_per_atom.py b/tests/utils/test_per_atom.py index 3f4f34945..0f84180e7 100644 --- a/tests/utils/test_per_atom.py +++ b/tests/utils/test_per_atom.py @@ -13,16 +13,19 @@ def test_average_by_num_atoms(): positions=torch.tensor([[0.0, 0.0, 0.0]]), cell=torch.eye(3), types=torch.tensor([0]), + pbc=torch.tensor([True, True, True]), ), System( positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), cell=torch.eye(3), types=torch.tensor([0, 0]), + pbc=torch.tensor([True, True, True]), ), System( positions=torch.tensor([[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]), cell=torch.eye(3), types=torch.tensor([0, 0, 0]), + pbc=torch.tensor([True, True, True]), ), ] diff --git a/tests/utils/test_transfer.py b/tests/utils/test_transfer.py index b669c0c29..283842779 100644 --- a/tests/utils/test_transfer.py +++ b/tests/utils/test_transfer.py @@ -14,6 +14,7 @@ def test_systems_and_targets_to_dtype(): positions=torch.tensor([[1.0, 1.0, 1.0]]), cell=torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]), types=torch.tensor([1]), + pbc=torch.tensor([True, True, True]), ) targets = TensorMap( keys=Labels.single(), @@ -39,6 +40,7 @@ def test_systems_and_targets_to_dtype_and_device(): positions=torch.tensor([[1.0, 1.0, 1.0]]), cell=torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]), types=torch.tensor([1]), + pbc=torch.tensor([True, True, True]), ) targets = TensorMap( keys=Labels.single(),