Pre-commit formatting

nequip/data/AtomicDataDict.py

@@ -5,6 +5,7 @@
 
 Authors: Albert Musaelian
 """
+
 from typing import Dict, Any
 
 import torch

nequip/data/_dataset/_ase_dataset.py

@@ -51,10 +51,12 @@ def _ase_dataset_reader(
         datas.append(
             (
                 global_index,
-                AtomicData.from_ase(atoms=atoms, **atomicdata_kwargs)
-                if global_index in include_frames
-                # in-memory dataset will ignore this later, but needed for indexing to work out
-                else None,
+                (
+                    AtomicData.from_ase(atoms=atoms, **atomicdata_kwargs)
+                    if global_index in include_frames
+                    # in-memory dataset will ignore this later, but needed for indexing to work out
+                    else None
+                ),
             )
         )
     # Save to a tempfile---

nequip/data/_keys.py

@@ -2,6 +2,7 @@
 
 This is a seperate module to compensate for a TorchScript bug that can only recognize constants when they are accessed as attributes of an imported module.
 """
+
 import sys
 from typing import List
 

nequip/data/dataloader.py

@@ -84,6 +84,7 @@ class PartialSampler(Sampler[int]):
             If `None`, defaults to `len(data_source)`.
         generator (Generator): Generator used in sampling.
     """
+
     data_source: Dataset
     num_samples_per_epoch: int
     shuffle: bool

nequip/nn/_convnetlayer.py

@@ -149,9 +149,9 @@ def __init__(
         # updated with whatever the convolution outputs (which is a full graph module)
         self.irreps_out.update(self.conv.irreps_out)
         # but with the features updated by the nonlinearity
-        self.irreps_out[
-            AtomicDataDict.NODE_FEATURES_KEY
-        ] = self.equivariant_nonlin.irreps_out
+        self.irreps_out[AtomicDataDict.NODE_FEATURES_KEY] = (
+            self.equivariant_nonlin.irreps_out
+        )
 
     def forward(self, data: AtomicDataDict.Type) -> AtomicDataDict.Type:
         # save old features for resnet

nequip/nn/_grad_output.py

@@ -20,6 +20,7 @@ class GradientOutput(GraphModuleMixin, torch.nn.Module):
         out_field: the field in which to return the computed gradients. Defaults to ``f"d({of})/d({wrt})"`` for each field in ``wrt``.
         sign: either 1 or -1; the returned gradient is multiplied by this.
     """
+
     sign: float
     _negate: bool
     skip: bool
@@ -119,6 +120,7 @@ class PartialForceOutput(GraphModuleMixin, torch.nn.Module):
         vectorize: the vectorize option to ``torch.autograd.functional.jacobian``,
             false by default since it doesn't work well.
     """
+
     vectorize: bool
 
     def __init__(
@@ -183,6 +185,7 @@ class StressOutput(GraphModuleMixin, torch.nn.Module):
         func: the energy model to wrap
         do_forces: whether to compute forces as well
     """
+
     do_forces: bool
 
     def __init__(

nequip/nn/_interaction_block.py

@@ -1,4 +1,5 @@
 """ Interaction Block """
+
 from typing import Optional, Dict, Callable
 
 import torch

nequip/scripts/evaluate.py

@@ -383,9 +383,11 @@ def main(args=None, running_as_script: bool = True):
         if do_metrics:
             display_bar = context_stack.enter_context(
                 tqdm(
-                    bar_format=""
-                    if prog.disable  # prog.ncols doesn't exist if disabled
-                    else ("{desc:." + str(prog.ncols) + "}"),
+                    bar_format=(
+                        ""
+                        if prog.disable  # prog.ncols doesn't exist if disabled
+                        else ("{desc:." + str(prog.ncols) + "}")
+                    ),
                     disable=None,
                 )
             )

nequip/scripts/train.py

@@ -1,4 +1,5 @@
 """ Train a network."""
+
 import logging
 import argparse
 import warnings
@@ -24,7 +25,8 @@
 from nequip.scripts._logger import set_up_script_logger
 
 warnings.filterwarnings(  # unnecessary e3nn-related JIT warning
-    "ignore", message="The TorchScript type system doesn't support instance-level annotations"
+    "ignore",
+    message="The TorchScript type system doesn't support instance-level annotations",
 )
 default_config = dict(
     root="./",

nequip/utils/config.py

@@ -34,6 +34,7 @@
     If a parameter is updated, the updated value will be formatted back to the same type.
 
 """
+
 from typing import Set, Dict, Any, List
 
 import inspect

nequip/utils/savenload.py

@@ -1,6 +1,7 @@
 """
 utilities that involve file searching and operations (i.e. save/load)
 """
+
 from typing import Union, List, Tuple, Optional, Callable
 import sys
 import logging

nequip/utils/test.py

@@ -45,9 +45,11 @@ def assert_permutation_equivariant(
 
     if tolerance is None:
         atol = PERMUTATION_FLOAT_TOLERANCE[
-            func.model_dtype
-            if isinstance(func, GraphModel)
-            else torch.get_default_dtype()
+            (
+                func.model_dtype
+                if isinstance(func, GraphModel)
+                else torch.get_default_dtype()
+            )
         ]
     else:
         atol = tolerance

nequip/utils/unittests/model_tests.py

@@ -449,10 +449,12 @@ def test_partial_forces(self, config, atomic_batch, device, strict_locality):
                 assert torch.allclose(
                     output[k],
                     output_partial[k],
-                    atol=1e-8
-                    if k == AtomicDataDict.TOTAL_ENERGY_KEY
-                    and torch.get_default_dtype() == torch.float64
-                    else 1e-5,
+                    atol=(
+                        1e-8
+                        if k == AtomicDataDict.TOTAL_ENERGY_KEY
+                        and torch.get_default_dtype() == torch.float64
+                        else 1e-5
+                    ),
                 )
             else:
                 assert torch.equal(output[k], output_partial[k])

tests/integration/test_evaluate.py

@@ -171,9 +171,11 @@ def runit(params: dict):
             assert np.allclose(
                 err,
                 0.0,
-                atol=1e-8
-                if true_identity
-                else (1e-2 if metric.startswith("e") else 1e-4),
+                atol=(
+                    1e-8
+                    if true_identity
+                    else (1e-2 if metric.startswith("e") else 1e-4)
+                ),
             ), f"Metric `{metric}` wasn't zero!"
     elif builder == ConstFactorModel:
         # TODO: check comperable to naive numpy compute

tests/unit/model/test_pair/test_zbl.py

@@ -81,7 +81,7 @@ def test_lammps_repro(self, config):
         # $ lmp -in zbl_data.lmps
         # $ python -c "import numpy as np; d = np.loadtxt('zbl.dat', skiprows=1); np.save('zbl.npy', d)"
         refdata = np.load(Path(__file__).parent / "zbl.npy")
-        for (r, Zi, Zj, pe, fxi, fxj) in refdata:
+        for r, Zi, Zj, pe, fxi, fxj in refdata:
             if r >= r_max:
                 continue
             atoms.positions[1, 0] = r

tests/unit/utils/test_config.py

@@ -1,6 +1,7 @@
 """
 Config tests
 """
+
 import pytest
 
 from os import remove

tests/unit/utils/test_output.py

@@ -1,6 +1,7 @@
 """
 Config tests
 """
+
 import pytest
 import tempfile