Updates to use numpy 2.0 and ase >= 3.23

aimnet/calculators/aimnet2ase.py

@@ -12,7 +12,14 @@
 class AIMNet2ASE(Calculator):
     from typing import ClassVar
 
-    implemented_properties: ClassVar[list[str]] = ["energy", "forces", "free_energy", "charges", "stress", "dipole_moment"]
+    implemented_properties: ClassVar[list[str]] = [
+        "energy",
+        "forces",
+        "free_energy",
+        "charges",
+        "stress",
+        "dipole_moment",
+    ]
 
     def __init__(self, base_calc: AIMNet2Calculator | str = "aimnet2", charge=0, mult=1):
         super().__init__()
@@ -59,12 +66,11 @@ def update_tensors(self):
         if self._t_mult is None:
             self._t_mult = torch.tensor(self.mult, dtype=torch.float32, device=self.base_calc.device)
 
-    def get_dipole_moment(self,atoms):
+    def get_dipole_moment(self, atoms):
         charges = self.get_charges()[:, np.newaxis]
         positions = atoms.get_positions()
         return np.sum(charges * positions, axis=0)
 
-
     def calculate(self, atoms=None, properties=None, system_changes=all_changes):
         if properties is None:
             properties = ["energy"]
@@ -97,7 +103,7 @@ def calculate(self, atoms=None, properties=None, system_changes=all_changes):
 
         self.results["energy"] = results["energy"].item()
         self.results["charges"] = results["charges"]
-        self.results['dipole_moment'] = self.get_dipole_moment(self.atoms)
+        self.results["dipole_moment"] = self.get_dipole_moment(self.atoms)
 
         if "forces" in properties:
             self.results["forces"] = results["forces"]

aimnet/calculators/calculator.py

@@ -146,7 +146,7 @@ def mol_flatten(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
         elif ndim == 3:
             # batched input
             B, N = data["coord"].shape[:2]
-            if N > self.nb_threshold or self.device == "cpu":
+            if self.nb_threshold < N or self.device == "cpu":
                 self._batch = B
                 data["mol_idx"] = torch.repeat_interleave(
                     torch.arange(0, B, device=self.device), torch.full((B,), N, device=self.device)

pyproject.toml

@@ -18,16 +18,16 @@ aimnet = "aimnet.cli:cli"
 python = ">=3.10,<4.0"
 torch = ">=2.4"
 pyyaml = "^6.0.2"
-numpy = "<2.0"
-numba = "^0.60.0"
+numpy = ">=2.0"
+numba = ">=0.61.0"
 requests = "^2.32.3"
 click = "^8.1.7"
 omegaconf = "^2.3.0"
 wandb = "^0.18.5"
 jinja2 = "^3.1.4"
 h5py = "^3.12.1"
 pytorch-ignite = "^0.5.1"
-ase = "3.22.1"
+ase = ">=3.23"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.2.0"

tests/test_ase.py

@@ -4,22 +4,22 @@
 
 from aimnet.calculators import AIMNet2ASE
 
-MODELS = ('aimnet2', 'aimnet2_b973c')
+MODELS = ("aimnet2", "aimnet2_b973c")
 
 file = os.path.join(os.path.dirname(__file__), "data", "caffeine.xyz")
 
+
 def test_calculator():
     for model in MODELS:
-
         atoms = read(file)
         atoms.calc = AIMNet2ASE(model)
         e = atoms.get_potential_energy()
         assert isinstance(e, float)
 
-        assert hasattr(atoms, 'get_charges')
+        assert hasattr(atoms, "get_charges")
         q = atoms.get_charges()
         assert q.shape == (len(atoms),)
 
-        assert hasattr(atoms, 'get_dipole_moment')
+        assert hasattr(atoms, "get_dipole_moment")
         dm = atoms.get_dipole_moment()
         assert dm.shape == (3,)

tests/test_model.py

@@ -44,9 +44,9 @@ def test_aimnet2():
     model.load_state_dict(model_from_zoo.state_dict(), strict=False)
     model = Forces(model)
     atoms = ase.io.read(os.path.join(os.path.dirname(__file__), "data", "caffeine.xyz"))
-    ref_e = atoms.info["energy"]  # type: ignore
-    ref_f = atoms.arrays["forces"]  # type: ignore
-    ref_q = atoms.arrays["initial_charges"]  # type: ignore
+    ref_e = atoms.get_total_energy()  # type: ignore
+    ref_f = atoms.get_forces()  # type: ignore
+    ref_q = atoms.get_charges()  # type: ignore
     _in = {
         "coord": torch.as_tensor(atoms.get_positions()).unsqueeze(0),  # type: ignore
         "numbers": torch.as_tensor(atoms.get_atomic_numbers()).unsqueeze(0),  # type: ignore