aimnet 0.0.1__py3-none-any.whl → 0.1.0__py3-none-any.whl

aimnet/__init__.py

@@ -0,0 +1,7 @@
+# Version is managed by hatch-vcs from git tags
+try:
+    from importlib.metadata import version
+
+    __version__ = version("aimnet")
+except Exception:
+    __version__ = "0.0.0+unknown"

aimnet/base.py

@@ -1,4 +1,4 @@
-from typing import ClassVar, Dict, Final
+from typing import ClassVar, Final
 
 import torch
 from torch import Tensor, nn
@@ -6,30 +6,46 @@ from torch import Tensor, nn
 from aimnet import nbops
 
 
-class AIMNet2Base(nn.Module):  # pylint: disable=abstract-method
+class AIMNet2Base(nn.Module):
     """Base class for AIMNet2 models. Implements pre-processing data:
     converting to right dtype and device, setting nb mode, calculating masks.
     """
 
     _required_keys: Final = ["coord", "numbers", "charge"]
     _required_keys_dtype: Final = [torch.float32, torch.int64, torch.float32]
-    _optional_keys: Final = ["mult", "nbmat", "nbmat_lr", "mol_idx", "shifts", "cell"]
-    _optional_keys_dtype: Final = [torch.float32, torch.int64, torch.int64, torch.int64, torch.float32, torch.float32]
+    _optional_keys: Final = [
+        "mult",
+        "nbmat",
+        "nbmat_lr",
+        "mol_idx",
+        "shifts",
+        "shifts_lr",
+        "cell",
+    ]
+    _optional_keys_dtype: Final = [
+        torch.float32,  # mult
+        torch.int64,  # nbmat
+        torch.int64,  # nbmat_lr
+        torch.int64,  # mol_idx
+        torch.float32,  # shifts
+        torch.float32,  # shifts_lr
+        torch.float32,  # cell
+    ]
     __constants__: ClassVar = ["_required_keys", "_required_keys_dtype", "_optional_keys", "_optional_keys_dtype"]
 
     def __init__(self):
         super().__init__()
 
-    def _prepare_dtype(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
-        for k, d in zip(self._required_keys, self._required_keys_dtype):
+    def _prepare_dtype(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        for k, d in zip(self._required_keys, self._required_keys_dtype, strict=False):
             assert k in data, f"Key {k} is required"
             data[k] = data[k].to(d)
-        for k, d in zip(self._optional_keys, self._optional_keys_dtype):
+        for k, d in zip(self._optional_keys, self._optional_keys_dtype, strict=False):
             if k in data:
                 data[k] = data[k].to(d)
         return data
 
-    def prepare_input(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]:
+    def prepare_input(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
         """Some sommon operations"""
         data = self._prepare_dtype(data)
         data = nbops.set_nb_mode(data)

aimnet/calculators/__init__.py

@@ -5,11 +5,11 @@ from .calculator import AIMNet2Calculator
 __all__ = ["AIMNet2Calculator"]
 
 if importlib.util.find_spec("ase") is not None:
-    from .aimnet2ase import AIMNet2ASE
+    from .aimnet2ase import AIMNet2ASE  # noqa: F401
 
-    __all__.append(AIMNet2ASE)  # type: ignore
+    __all__.append("AIMNet2ASE")
 
 if importlib.util.find_spec("pysisyphus") is not None:
-    from .aimnet2pysis import AIMNet2Pysis
+    from .aimnet2pysis import AIMNet2Pysis  # noqa: F401
 
-    __all__.append(AIMNet2Pysis)  # type: ignore
+    __all__.append("AIMNet2Pysis")

aimnet/calculators/aimnet2ase.py

@@ -12,16 +12,24 @@ from .calculator import AIMNet2Calculator
 class AIMNet2ASE(Calculator):
     from typing import ClassVar
 
-    implemented_properties: ClassVar[list[str]] = ["energy", "forces", "free_energy", "charges", "stress"]
+    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__()
         if isinstance(base_calc, str):
             base_calc = AIMNet2Calculator(base_calc)
         self.base_calc = base_calc
+        self.reset()
         self.charge = charge
         self.mult = mult
-        self.reset()
+        self.update_tensors()
         # list of implemented species
         if hasattr(base_calc, "implemented_species"):
             self.implemented_species = base_calc.implemented_species.cpu().numpy()  # type: ignore
@@ -33,8 +41,6 @@ class AIMNet2ASE(Calculator):
         self._t_numbers = None
         self._t_charge = None
         self._t_mult = None
-        self.charge = 0.0
-        self.mult = 1.0
 
     def set_atoms(self, atoms):
         if self.implemented_species is not None and not np.in1d(atoms.numbers, self.implemented_species).all():
@@ -53,13 +59,18 @@ class AIMNet2ASE(Calculator):
         self.update_tensors()
 
     def update_tensors(self):
-        if self._t_numbers is None:
+        if self._t_numbers is None and getattr(self, "atoms", None):
             self._t_numbers = torch.tensor(self.atoms.numbers, dtype=torch.int64, device=self.base_calc.device)
         if self._t_charge is None:
             self._t_charge = torch.tensor(self.charge, dtype=torch.float32, device=self.base_calc.device)
         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):
+        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"]
@@ -90,8 +101,10 @@ class AIMNet2ASE(Calculator):
                 v = v.squeeze(0)
             results[k] = v.detach().cpu().numpy()  # type: ignore
 
-        self.results["energy"] = results["energy"]
+        self.results["energy"] = results["energy"].item()
         self.results["charges"] = results["charges"]
+        self.results["dipole_moment"] = self.get_dipole_moment(self.atoms)
+
         if "forces" in properties:
             self.results["forces"] = results["forces"]
         if "stress" in properties: