Lint format

Author: jwa7

src/metatrain/pet/model.py

@@ -410,7 +410,10 @@ def forward(
                                 )
                             else:  # per_pair
                                 sample_values.append(
-                                    torch.tensor([int(A), int(i), int(i), 0, 0, 0], dtype=torch.int32)
+                                    torch.tensor(
+                                        [int(A), int(i), int(i), 0, 0, 0],
+                                        dtype=torch.int32,
+                                    )
                                 )
                     if self.sample_kinds[output_name] == "per_atom":
                         sample_labels = Labels(
@@ -419,7 +422,14 @@ def forward(
                         )
                     else:
                         sample_labels = Labels(
-                            ["system", "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"],
+                            [
+                                "system",
+                                "first_atom",
+                                "second_atom",
+                                "cell_shift_a",
+                                "cell_shift_b",
+                                "cell_shift_c",
+                            ],
                             torch.vstack(sample_values),
                         )
                     X = self._compute_X_per_atom(
@@ -550,7 +560,7 @@ def _include_key(key: LabelsEntry) -> bool:
           if values are 0 and 1 respectively (indicating an invariant block of a
           spherical target).
         - If the key has names ["o3_lambda", "o3_sigma", "s2_pi], it is included
-          if values are 0, 1, 0 respectively, indicating an unsymmetrized 
+          if values are 0, 1, 0 respectively, indicating an unsymmetrized
           invariant block of a per-pair target.
     """
     include_key = False

src/metatrain/utils/additive/_base_composition.py

@@ -8,9 +8,9 @@
 from metatensor.torch.operations._multiply import _multiply_block_constant
 from metatomic.torch import System
 
+from ..basis import get_onsite_samples_mask
 from ..data import TargetInfo
 from ..evaluate_model import evaluate_model
-from ..basis import get_onsite_samples_mask
 
 
 def remove_additive(
@@ -78,7 +78,7 @@ def remove_additive(
                 values[onsite_sample_idxs] = old_block.values.detach()
             else:
                 values = old_block.values.detach()
-            
+
             device = targets[target_key].block(block_key).values.device
             block = metatensor.torch.TensorBlock(
                 values=values.to(device=device),

src/metatrain/utils/additive/remove.py

@@ -157,8 +157,7 @@ def _apply_wigner_D_matrices(
                 block.samples.values[:, 0], return_inverse=True
             )
             split_values = [
-                values[inverse_indices == i]
-                for i in range(len(unique_system_ids))
+                values[inverse_indices == i] for i in range(len(unique_system_ids))
             ]
         else:
             if "atom" in block.samples.names:

src/metatrain/utils/augmentation.py

@@ -1,6 +1,7 @@
-import torch
-from metatensor.torch import LabelsEntry, Labels, TensorMap
+from typing import List
 
+import torch
+from metatensor.torch import Labels, TensorMap
 
 
 def is_spherical_atomic_basis(target: TensorMap) -> bool:
@@ -10,12 +11,9 @@ def is_spherical_atomic_basis(target: TensorMap) -> bool:
     :param target: The target tensor map to check.
     :returns: True if the target is a spherical atomic basis, False otherwise.
     """
-    if not (
-        "o3_lambda" in target.keys.names
-        and "o3_sigma" in target.keys.names
-    ): 
+    if not ("o3_lambda" in target.keys.names and "o3_sigma" in target.keys.names):
         return False
-    
+
     # i.e. electron density on basis
     if (
         target.sample_names == ["system", "atom"]
@@ -25,19 +23,27 @@ def is_spherical_atomic_basis(target: TensorMap) -> bool:
     ):
         return True
 
-
     # i.e. hamiltonian
     if (
-        target.sample_names == ["system", "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c"]
+        target.sample_names
+        == [
+            "system",
+            "first_atom",
+            "second_atom",
+            "cell_shift_a",
+            "cell_shift_b",
+            "cell_shift_c",
+        ]
         and len(target[0].components) == 1
         and target[0].components[0].names == ["o3_mu"]
-        and target.property_names == ['first_atom_type', 'second_atom_type', 'l_1', 'l_2', 'n_1', 'n_2']
+        and target.property_names
+        == ["first_atom_type", "second_atom_type", "l_1", "l_2", "n_1", "n_2"]
     ):
         return True
 
-
     return False
 
+
 def get_block_sample_idxs_per_atom(
     sample_labels_with_types: Labels,
     atomic_basis: Labels,
@@ -49,14 +55,10 @@ def get_block_sample_idxs_per_atom(
     respectively.
     """
     assert sample_labels_with_types.names == ["system", "atom", "center_type"]
-    
+
     # find the center types that have basis functions with the given o3_lambda
     valid_center_types = torch.sort(
-        torch.unique(
-            atomic_basis.values[
-                atomic_basis.values[:, 0] == o3_lambda
-            ][:, 1]
-        )
+        torch.unique(atomic_basis.values[atomic_basis.values[:, 0] == o3_lambda][:, 1])
     )[0]
 
     # find the sample indices that have these center types
@@ -83,11 +85,7 @@ def get_block_sample_idxs_per_pair(
     """
 
     if node_or_edge == "node":
-        assert sample_labels_with_types.names == [
-            "system",
-            "atom",
-            "center_type"
-        ]
+        assert sample_labels_with_types.names == ["system", "atom", "center_type"]
         # Find the pairs of atom types that can be in this block
         atomic_type_pairs = atomic_basis.values[
             atomic_basis.select(
@@ -104,7 +102,8 @@ def get_block_sample_idxs_per_pair(
         ][:, 3]
         sample_idxs = sample_labels_with_types.select(
             Labels(
-                ["center_type"], center_types.reshape(-1, 1),
+                ["center_type"],
+                center_types.reshape(-1, 1),
             )
         )
 
@@ -136,12 +135,9 @@ def get_block_sample_idxs_per_pair(
                 torch.hstack(
                     [
                         atomic_type_pairs[:, 3:],
-                        torch.full(
-                            (atomic_type_pairs.shape[0], 1),
-                            s2_pi
-                        )
+                        torch.full((atomic_type_pairs.shape[0], 1), s2_pi),
                     ]
-                )
+                ),
             )
         )
 
@@ -297,4 +293,3 @@ def symmetrize_edge_features(
             edge_features_sym_m1,
         ]
     )
-

src/metatrain/utils/basis.py

@@ -72,9 +72,7 @@ def read_generic(target: DictConfig) -> Tuple[List[TensorMap], TargetInfo]:
     # actual properties of the tensor maps
     target_info.layout = _empty_tensor_map_like(tensor_map)
 
-    system_ids = metatensor.torch.unique_metadata(
-        tensor_map, "samples", "system"
-    )
+    system_ids = metatensor.torch.unique_metadata(tensor_map, "samples", "system")
     selections = [
         Labels(
             names=["system"],

src/metatrain/utils/data/readers/metatensor.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Union
+from typing import List, Union
 
 import torch
 from metatensor.torch import Labels, TensorBlock, TensorMap
@@ -49,12 +49,9 @@ def gradients(self) -> List[str]:
     @property
     def per_atom(self) -> bool:
         """Whether the target is per atom."""
-        return (
-            "atom" in self.layout.block(0).samples.names
-            or (
-                "first_atom" in self.layout.block(0).samples.names
-                and "second_atom" in self.layout.block(0).samples.names
-            )
+        return "atom" in self.layout.block(0).samples.names or (
+            "first_atom" in self.layout.block(0).samples.names
+            and "second_atom" in self.layout.block(0).samples.names
         )
 
     def __repr__(self):
@@ -180,7 +177,7 @@ def _check_layout(self, layout: TensorMap) -> None:
                     o3_lambda, o3_sigma, s2_pi = (
                         int(key.values[0].item()),
                         int(key.values[1].item()),
-                        None
+                        None,
                     )
                 else:
                     assert len(key.names) == 3
@@ -191,33 +188,33 @@ def _check_layout(self, layout: TensorMap) -> None:
                     )
                 if o3_sigma not in [-1, 1]:
                     raise ValueError(
-                        "The layout ``TensorMap`` of a spherical tensor target should "
-                        "have key dimension 'o3_sigma' that is either -1 or 1."
-                        f"Found '{o3_sigma}' instead."
+                        "The layout ``TensorMap`` of a spherical tensor "
+                        "target should have key dimension 'o3_sigma' that "
+                        f"is either -1 or 1. Found '{o3_sigma}' instead."
                     )
                 if o3_lambda < 0:
                     raise ValueError(
-                        "The layout ``TensorMap`` of a spherical tensor target should "
-                        "have key dimension 'o3_lambda' that is non-negative."
-                        f"Found '{o3_lambda}' instead."
+                        "The layout ``TensorMap`` of a spherical tensor "
+                        "target should have key dimension 'o3_lambda' that "
+                        f"is non-negative. Found '{o3_lambda}' instead."
                     )
                 if s2_pi is not None:
                     if s2_pi not in [-1, 0, +1]:
                         raise ValueError(
-                            "The layout ``TensorMap`` of a spherical tensor target should "
-                            "have key dimension 's2_pi' that is either -1, 0, or +1."
-                            f"Found '{s2_pi}' instead."
+                            "The layout ``TensorMap`` of a spherical tensor "
+                            "target should have key dimension 's2_pi' that "
+                            f"is either -1, 0, or +1. Found '{s2_pi}' instead."
                         )
                 components = block.components
                 if len(components) != 1:
                     raise ValueError(
-                        "The layout ``TensorMap`` of a spherical tensor target should "
-                        "have a single component."
+                        "The layout ``TensorMap`` of a spherical tensor "
+                        "target should have a single component."
                     )
                 if len(components[0]) != 2 * o3_lambda + 1:
                     raise ValueError(
-                        "Each ``TensorBlock`` of a spherical tensor target should have "
-                        "a component with 2*o3_lambda + 1 elements."
+                        "Each ``TensorBlock`` of a spherical tensor target"
+                        "should have a component with 2*o3_lambda + 1 elements."
                         f"Found '{len(components[0])}' elements instead."
                     )
                 if len(block.gradients_list()) > 0:
@@ -414,7 +411,6 @@ def _get_cartesian_target_info(target: DictConfig) -> TargetInfo:
 
 
 def _get_spherical_target_info(target: DictConfig) -> TargetInfo:
-    
     irreps = target["type"]["spherical"]["irreps"]
     atomic_basis = target["type"]["spherical"].get("atomic_basis", None)
 
@@ -427,14 +423,11 @@ def _get_spherical_target_info(target: DictConfig) -> TargetInfo:
         atomic_basis = Labels(
             names=atomic_basis_names,
             values=torch.tensor(
-                [
-                    [i[name] for name in atomic_basis_names]
-                    for i in atomic_basis
-                ],
+                [[i[name] for name in atomic_basis_names] for i in atomic_basis],
                 dtype=torch.int32,
             ),
         )
-    
+
     # Infer the sample names
     if target["per_atom"]:
         if atomic_basis is None:

src/metatrain/utils/data/target_info.py

@@ -98,7 +98,9 @@ def __call__(
                     "TensorMapLoss requires the two TensorMaps to have the same "
                     "components."
                 )
-            if not torch.all(block_1.samples.values[:, 1:] == block_2.samples.values[:, 1:]):
+            if not torch.all(
+                block_1.samples.values[:, 1:] == block_2.samples.values[:, 1:]
+            ):
                 raise ValueError(
                     "TensorMapLoss requires the two TensorMaps "
                     "to have the same samples."
@@ -158,7 +160,9 @@ def __call__(
                 assert values_1.shape[0] == 0
             else:
                 loss += (
-                    self.weight * self.losses["values"](values_1, values_2) / sliding_weight
+                    self.weight
+                    * self.losses["values"](values_1, values_2)
+                    / sliding_weight
                 )
             for gradient_name in block_2.gradients_list():
                 gradient_weight = self.gradient_weights[gradient_name]