lint: ruff safe fixes

mace/__init__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from .__version__ import __version__

mace/__version__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 __version__ = "0.3.6"

mace/calculators/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .foundations_models import mace_anicc, mace_mp, mace_off
 from .lammps_mace import LAMMPS_MACE
 from .mace import MACECalculator

mace/calculators/foundations_models.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import os
 import urllib.request
 from pathlib import Path

mace/calculators/lammps_mace.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Dict, List, Optional
 
 import torch

mace/calculators/mace.py

@@ -3,7 +3,7 @@
 # Authors: Ilyes Batatia, David Kovacs
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
-
+from __future__ import annotations
 
 from glob import glob
 from pathlib import Path

mace/cli/active_learning_md.py

@@ -1,4 +1,5 @@
 """Demonstrates active learning molecular dynamics with constant temperature."""
+from __future__ import annotations
 
 import argparse
 import os

mace/cli/create_lammps_model.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import sys
 
 import torch

mace/cli/eval_configs.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import argparse
 

mace/cli/plot_train.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import argparse
 import dataclasses
 import glob
@@ -49,7 +51,7 @@ def parse_path(path: str) -> RunInfo:
 def parse_training_results(path: str) -> List[dict]:
     run_info = parse_path(path)
     results = []
-    with open(path, mode="r", encoding="utf-8") as f:
+    with open(path, encoding="utf-8") as f:
         for line in f:
             d = json.loads(line)
             d["name"] = run_info.name

mace/cli/preprocess_data.py

@@ -1,5 +1,6 @@
 # This file loads an xyz dataset and prepares
 # new hdf5 file that is ready for training with on-the-fly dataloading
+from __future__ import annotations
 
 import argparse
 import ast
@@ -77,7 +78,7 @@ def split_array(a: np.ndarray, max_size: int):
     factors = get_prime_factors(len(a))
     max_factor = 1
     for i in range(1, len(factors) + 1):
-        for j in range(0, len(factors) - i + 1):
+        for j in range(len(factors) - i + 1):
             if np.prod(factors[j : j + i]) <= max_size:
                 test = np.prod(factors[j : j + i])
                 max_factor = max(test, max_factor)
@@ -208,7 +209,7 @@ def run(args: argparse.Namespace):
             [atomic_energies_dict[z] for z in z_table.zs]
         )
         logging.info(f"Atomic energies: {atomic_energies.tolist()}")
-        _inputs = [args.h5_prefix+'train', z_table, args.r_max, atomic_energies, args.batch_size, args.num_process]
+        _inputs = [args.h5_prefix+"train", z_table, args.r_max, atomic_energies, args.batch_size, args.num_process]
         avg_num_neighbors, mean, std=pool_compute_stats(_inputs)
         logging.info(f"Average number of neighbors: {avg_num_neighbors}")
         logging.info(f"Mean: {mean}")

mace/cli/run_train.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm, David Kovacs
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import argparse
 import ast
@@ -60,7 +61,7 @@ def run(args: argparse.Namespace) -> None:
         try:
             distr_env = DistributedEnvironment()
         except Exception as e:  # pylint: disable=W0703
-            logging.error(f"Failed to initialize distributed environment: {e}")
+            logging.exception(f"Failed to initialize distributed environment: {e}")
             return
         world_size = distr_env.world_size
         local_rank = distr_env.local_rank
@@ -69,7 +70,7 @@ def run(args: argparse.Namespace) -> None:
             print(distr_env)
         torch.distributed.init_process_group(backend="nccl")
     else:
-        rank = int(0)
+        rank = 0
 
     # Setup
     tools.set_seeds(args.seed)
@@ -113,7 +114,7 @@ def run(args: argparse.Namespace) -> None:
         args.r_max = model_foundation.r_max.item()
 
     if args.statistics_file is not None:
-        with open(args.statistics_file, "r") as f:  # pylint: disable=W1514
+        with open(args.statistics_file) as f:  # pylint: disable=W1514
             statistics = json.load(f)
         logging.info("Using statistics json file")
         args.r_max = statistics["r_max"] if args.foundation_model is None else args.r_max
@@ -182,11 +183,10 @@ def run(args: argparse.Namespace) -> None:
                 z: model_foundation.atomic_energies_fn.atomic_energies[z_table_foundation.z_to_index(z)].item()
                 for z in z_table.zs
             }
+        elif args.train_file.endswith(".xyz"):
+            atomic_energies_dict = get_atomic_energies(args.E0s, collections.train, z_table)
         else:
-            if args.train_file.endswith(".xyz"):
-                atomic_energies_dict = get_atomic_energies(args.E0s, collections.train, z_table)
-            else:
-                atomic_energies_dict = get_atomic_energies(args.E0s, None, z_table)
+            atomic_energies_dict = get_atomic_energies(args.E0s, None, z_table)
 
     if args.model == "AtomicDipolesMACE":
         atomic_energies = None
@@ -540,13 +540,12 @@ def run(args: argparse.Namespace) -> None:
         swas.append(True)
         if args.start_swa is None:
             args.start_swa = max(1, args.max_num_epochs // 4 * 3)
-        else:
-            if args.start_swa > args.max_num_epochs:
-                logging.info(
-                    f"Start Stage Two must be less than max_num_epochs, got {args.start_swa} > {args.max_num_epochs}"
-                )
-                args.start_swa = max(1, args.max_num_epochs // 4 * 3)
-                logging.info(f"Setting start Stage Two to {args.start_swa}")
+        elif args.start_swa > args.max_num_epochs:
+            logging.info(
+                f"Start Stage Two must be less than max_num_epochs, got {args.start_swa} > {args.max_num_epochs}"
+            )
+            args.start_swa = max(1, args.max_num_epochs // 4 * 3)
+            logging.info(f"Setting start Stage Two to {args.start_swa}")
         if args.loss == "forces_only":
             raise ValueError("Can not select Stage Two with forces only loss.")
         if args.loss == "virials":
@@ -712,7 +711,7 @@ def run(args: argparse.Namespace) -> None:
             )
         try:
             drop_last = test_set.drop_last
-        except AttributeError as e:  # pylint: disable=W0612
+        except AttributeError:  # pylint: disable=W0612
             drop_last = False
         test_loader = torch_geometric.dataloader.DataLoader(
             test_set,
@@ -775,7 +774,7 @@ def run(args: argparse.Namespace) -> None:
                         path_complied,
                         _extra_files=extra_files,
                     )
-                except Exception as e:  # pylint: disable=W0703
+                except Exception:  # pylint: disable=W0703
                     pass
             else:
                 torch.save(model, Path(args.model_dir) / (args.name + ".model"))
@@ -788,7 +787,7 @@ def run(args: argparse.Namespace) -> None:
                         path_complied,
                         _extra_files=extra_files,
                     )
-                except Exception as e:  # pylint: disable=W0703
+                except Exception:  # pylint: disable=W0703
                     pass
 
         if args.distributed:

mace/data/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .atomic_data import AtomicData
 from .hdf5_dataset import HDF5Dataset, dataset_from_sharded_hdf5
 from .neighborhood import get_neighborhood

mace/data/atomic_data.py

@@ -3,18 +3,13 @@
 # Authors: Ilyes Batatia, Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 from typing import Optional, Sequence
 
 import torch.utils.data
 
-from mace.tools import (
-    AtomicNumberTable,
-    atomic_numbers_to_indices,
-    to_one_hot,
-    torch_geometric,
-    voigt_to_matrix,
-)
+from mace.tools import AtomicNumberTable, atomic_numbers_to_indices, to_one_hot, torch_geometric, voigt_to_matrix
 
 from .neighborhood import get_neighborhood
 from .utils import Configuration
@@ -107,7 +102,7 @@ def __init__(
         super().__init__(**data)
 
     @classmethod
-    def from_config(cls, config: Configuration, z_table: AtomicNumberTable, cutoff: float) -> "AtomicData":
+    def from_config(cls, config: Configuration, z_table: AtomicNumberTable, cutoff: float) -> AtomicData:
         edge_index, shifts, unit_shifts = get_neighborhood(
             positions=config.positions, cutoff=cutoff, pbc=config.pbc, cell=config.cell
         )

mace/data/hdf5_dataset.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from glob import glob
 from typing import List
 

mace/data/neighborhood.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Optional, Tuple
 
 import numpy as np

mace/data/utils.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm and David Kovacs
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
@@ -226,7 +227,7 @@ def load_from_xyz(
         for atoms in atoms_list:
             try:
                 atoms.info["REF_stress"] = atoms.get_stress()
-            except Exception as e:  # pylint: disable=W0703
+            except Exception:  # pylint: disable=W0703
                 atoms.info["REF_stress"] = None
     if not isinstance(atoms_list, list):
         atoms_list = [atoms_list]

mace/modules/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Callable, Dict, Optional, Type
 
 import torch
@@ -30,14 +32,7 @@
     WeightedForcesLoss,
     WeightedHuberEnergyForcesStressLoss,
 )
-from .models import (
-    MACE,
-    AtomicDipolesMACE,
-    BOTNet,
-    EnergyDipolesMACE,
-    ScaleShiftBOTNet,
-    ScaleShiftMACE,
-)
+from .models import MACE, AtomicDipolesMACE, BOTNet, EnergyDipolesMACE, ScaleShiftBOTNet, ScaleShiftMACE
 from .radial import BesselBasis, GaussianBasis, PolynomialCutoff, ZBLBasis
 from .symmetric_contraction import SymmetricContraction
 from .utils import (

mace/modules/blocks.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 from abc import abstractmethod
 from typing import Callable, List, Optional, Tuple, Union
@@ -15,19 +16,8 @@
 from mace.tools.compile import simplify_if_compile
 from mace.tools.scatter import scatter_sum
 
-from .irreps_tools import (
-    linear_out_irreps,
-    reshape_irreps,
-    tp_out_irreps_with_instructions,
-)
-from .radial import (
-    AgnesiTransform,
-    BesselBasis,
-    ChebychevBasis,
-    GaussianBasis,
-    PolynomialCutoff,
-    SoftTransform,
-)
+from .irreps_tools import linear_out_irreps, reshape_irreps, tp_out_irreps_with_instructions
+from .radial import AgnesiTransform, BesselBasis, ChebychevBasis, GaussianBasis, PolynomialCutoff, SoftTransform
 from .symmetric_contraction import SymmetricContraction
 
 

mace/modules/irreps_tools.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 from typing import List, Tuple
 

mace/modules/loss.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import torch
 

mace/modules/models.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 

mace/modules/radial.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import ase
 import numpy as np

mace/modules/symmetric_contraction.py

@@ -4,6 +4,7 @@
 # Authors: Ilyes Batatia
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 from typing import Dict, Optional, Union
 

mace/modules/utils.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm and David Kovacs
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import logging
 from typing import List, Optional, Tuple

mace/tools/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .arg_parser import build_default_arg_parser, build_preprocess_arg_parser
 from .cg import U_matrix_real
 from .checkpoint import CheckpointHandler, CheckpointIO, CheckpointState

mace/tools/arg_parser.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia, Gregor Simm, David Kovacs
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import argparse
 import os

mace/tools/cg.py

@@ -3,6 +3,7 @@
 # Authors: Ilyes Batatia
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import collections
 from typing import List, Union

mace/tools/checkpoint.py

@@ -3,6 +3,7 @@
 # Authors: Gregor Simm
 # This program is distributed under the MIT License (see MIT.md)
 ###########################################################################################
+from __future__ import annotations
 
 import dataclasses
 import logging

mace/tools/compile.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from contextlib import contextmanager
 from functools import wraps
 from typing import Callable, Tuple

mace/tools/finetuning_utils.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import torch
 
 from mace.tools.utils import AtomicNumberTable