aimnet 0.0.1__tar.gz

aimnet-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024, Roman Zubatyuk
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

aimnet-0.0.1/PKG-INFO

@@ -0,0 +1,78 @@
+Metadata-Version: 2.1
+Name: aimnet
+Version: 0.0.1
+Summary: AIMNet Machine Learned Interatomic Potential
+Home-page: https://github.com/isayevlab/aimnetcentral
+Author: Roman Zubatyuk
+Author-email: zubatyuk@gmail.com
+Requires-Python: >=3.10,<4.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: click (>=8.1.7,<9.0.0)
+Requires-Dist: h5py (>=3.12.1,<4.0.0)
+Requires-Dist: ignite (>=1.1.0,<2.0.0)
+Requires-Dist: jinja2 (>=3.1.4,<4.0.0)
+Requires-Dist: numba (>=0.60.0,<0.61.0)
+Requires-Dist: numpy (<2.0)
+Requires-Dist: omegaconf (>=2.3.0,<3.0.0)
+Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
+Requires-Dist: requests (>=2.32.3,<3.0.0)
+Requires-Dist: torch (>=2.5.0,<3.0.0)
+Requires-Dist: wandb (>=0.18.5,<0.19.0)
+Project-URL: Documentation, https://isayevlab.github.io/aimnetcentral/
+Project-URL: Repository, https://github.com/isayevlab/aimnetcentral
+Description-Content-Type: text/markdown
+
+# aimnetcentral
+
+[![Release](https://img.shields.io/github/v/release/isayevlab/aimnetcentral)](https://img.shields.io/github/v/release/isayevlab/aimnetcentral)
+[![Build status](https://img.shields.io/github/actions/workflow/status/isayevlab/aimnetcentral/main.yml?branch=main)](https://github.com/isayevlab/aimnetcentral/actions/workflows/main.yml?query=branch%3Amain)
+[![codecov](https://codecov.io/gh/isayevlab/aimnetcentral/branch/main/graph/badge.svg)](https://codecov.io/gh/isayevlab/aimnetcentral)
+[![Commit activity](https://img.shields.io/github/commit-activity/m/isayevlab/aimnetcentral)](https://img.shields.io/github/commit-activity/m/isayevlab/aimnetcentral)
+[![License](https://img.shields.io/github/license/isayevlab/aimnetcentral)](https://img.shields.io/github/license/isayevlab/aimnetcentral)
+
+AIMNet Machine Learned Interatomic Potential
+
+- **Github repository**: <https://github.com/isayevlab/aimnetcentral/>
+- **Documentation** <https://isayevlab.github.io/aimnetcentral/>
+
+## Getting started with your project
+
+First, create a repository on GitHub with the same name as this project, and then run the following commands:
+
+```bash
+git init -b main
+git add .
+git commit -m "init commit"
+git remote add origin git@github.com:isayevlab/aimnetcentral.git
+git push -u origin main
+```
+
+Finally, install the environment and the pre-commit hooks with
+
+```bash
+make install
+```
+
+You are now ready to start development on your project!
+The CI/CD pipeline will be triggered when you open a pull request, merge to main, or when you create a new release.
+
+To finalize the set-up for publishing to PyPI or Artifactory, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/publishing/#set-up-for-pypi).
+For activating the automatic documentation with MkDocs, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/mkdocs/#enabling-the-documentation-on-github).
+To enable the code coverage reports, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/codecov/).
+
+## Releasing a new version
+
+- Create an API Token on [PyPI](https://pypi.org/).
+- Add the API Token to your projects secrets with the name `PYPI_TOKEN` by visiting [this page](https://github.com/isayevlab/aimnetcentral/settings/secrets/actions/new).
+- Create a [new release](https://github.com/isayevlab/aimnetcentral/releases/new) on Github.
+- Create a new tag in the form `*.*.*`.
+- For more details, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/cicd/#how-to-trigger-a-release).
+
+---
+
+Repository initiated with [fpgmaas/cookiecutter-poetry](https://github.com/fpgmaas/cookiecutter-poetry).
+

aimnet-0.0.1/README.md

@@ -0,0 +1,49 @@
+# aimnetcentral
+
+[![Release](https://img.shields.io/github/v/release/isayevlab/aimnetcentral)](https://img.shields.io/github/v/release/isayevlab/aimnetcentral)
+[![Build status](https://img.shields.io/github/actions/workflow/status/isayevlab/aimnetcentral/main.yml?branch=main)](https://github.com/isayevlab/aimnetcentral/actions/workflows/main.yml?query=branch%3Amain)
+[![codecov](https://codecov.io/gh/isayevlab/aimnetcentral/branch/main/graph/badge.svg)](https://codecov.io/gh/isayevlab/aimnetcentral)
+[![Commit activity](https://img.shields.io/github/commit-activity/m/isayevlab/aimnetcentral)](https://img.shields.io/github/commit-activity/m/isayevlab/aimnetcentral)
+[![License](https://img.shields.io/github/license/isayevlab/aimnetcentral)](https://img.shields.io/github/license/isayevlab/aimnetcentral)
+
+AIMNet Machine Learned Interatomic Potential
+
+- **Github repository**: <https://github.com/isayevlab/aimnetcentral/>
+- **Documentation** <https://isayevlab.github.io/aimnetcentral/>
+
+## Getting started with your project
+
+First, create a repository on GitHub with the same name as this project, and then run the following commands:
+
+```bash
+git init -b main
+git add .
+git commit -m "init commit"
+git remote add origin git@github.com:isayevlab/aimnetcentral.git
+git push -u origin main
+```
+
+Finally, install the environment and the pre-commit hooks with
+
+```bash
+make install
+```
+
+You are now ready to start development on your project!
+The CI/CD pipeline will be triggered when you open a pull request, merge to main, or when you create a new release.
+
+To finalize the set-up for publishing to PyPI or Artifactory, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/publishing/#set-up-for-pypi).
+For activating the automatic documentation with MkDocs, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/mkdocs/#enabling-the-documentation-on-github).
+To enable the code coverage reports, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/codecov/).
+
+## Releasing a new version
+
+- Create an API Token on [PyPI](https://pypi.org/).
+- Add the API Token to your projects secrets with the name `PYPI_TOKEN` by visiting [this page](https://github.com/isayevlab/aimnetcentral/settings/secrets/actions/new).
+- Create a [new release](https://github.com/isayevlab/aimnetcentral/releases/new) on Github.
+- Create a new tag in the form `*.*.*`.
+- For more details, see [here](https://fpgmaas.github.io/cookiecutter-poetry/features/cicd/#how-to-trigger-a-release).
+
+---
+
+Repository initiated with [fpgmaas/cookiecutter-poetry](https://github.com/fpgmaas/cookiecutter-poetry).

aimnet-0.0.1/aimnet/base.py

@@ -0,0 +1,41 @@
+from typing import ClassVar, Dict, Final
+
+import torch
+from torch import Tensor, nn
+
+from aimnet import nbops
+
+
+class AIMNet2Base(nn.Module):  # pylint: disable=abstract-method
+    """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]
+    __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):
+            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):
+            if k in data:
+                data[k] = data[k].to(d)
+        return data
+
+    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)
+        data = nbops.calc_masks(data)
+
+        assert data["charge"].ndim == 1, "Charge should be 1D tensor."
+        if "mult" in data:
+            assert data["mult"].ndim == 1, "Mult should be 1D tensor."
+        return data

aimnet-0.0.1/aimnet/calculators/__init__.py

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

aimnet-0.0.1/aimnet/calculators/aimnet2ase.py

@@ -0,0 +1,98 @@
+import numpy as np
+import torch
+
+try:
+    from ase.calculators.calculator import Calculator, all_changes  # type: ignore
+except ImportError:
+    raise ImportError("ASE is not installed. Please install ASE to use this module.") from None
+
+from .calculator import AIMNet2Calculator
+
+
+class AIMNet2ASE(Calculator):
+    from typing import ClassVar
+
+    implemented_properties: ClassVar[list[str]] = ["energy", "forces", "free_energy", "charges", "stress"]
+
+    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.charge = charge
+        self.mult = mult
+        self.reset()
+        # list of implemented species
+        if hasattr(base_calc, "implemented_species"):
+            self.implemented_species = base_calc.implemented_species.cpu().numpy()  # type: ignore
+        else:
+            self.implemented_species = None
+
+    def reset(self):
+        super().reset()
+        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():
+            raise ValueError("Some species are not implemented in the AIMNet2Calculator")
+        self.reset()
+        self.atoms = atoms
+
+    def set_charge(self, charge):
+        self.charge = charge
+        self._t_charge = None
+        self.update_tensors()
+
+    def set_mult(self, mult):
+        self.mult = mult
+        self._t_mult = None
+        self.update_tensors()
+
+    def update_tensors(self):
+        if self._t_numbers is 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 calculate(self, atoms=None, properties=None, system_changes=all_changes):
+        if properties is None:
+            properties = ["energy"]
+        super().calculate(atoms, properties, system_changes)
+        self.update_tensors()
+
+        cell = self.atoms.cell.array if self.atoms.cell is not None and self.atoms.pbc.any() else None
+
+        _in = {
+            "coord": torch.tensor(self.atoms.positions, dtype=torch.float32, device=self.base_calc.device),
+            "numbers": self._t_numbers,
+            "charge": self._t_charge,
+            "mult": self._t_mult,
+        }
+
+        _unsqueezed = False
+        if cell is not None:
+            _in["cell"] = cell
+        else:
+            for k, v in _in.items():
+                _in[k] = v.unsqueeze(0)
+            _unsqueezed = True
+
+        results = self.base_calc(_in, forces="forces" in properties, stress="stress" in properties)
+
+        for k, v in results.items():
+            if _unsqueezed:
+                v = v.squeeze(0)
+            results[k] = v.detach().cpu().numpy()  # type: ignore
+
+        self.results["energy"] = results["energy"]
+        self.results["charges"] = results["charges"]
+        if "forces" in properties:
+            self.results["forces"] = results["forces"]
+        if "stress" in properties:
+            self.results["stress"] = results["stress"]

aimnet-0.0.1/aimnet/calculators/aimnet2pysis.py

@@ -0,0 +1,76 @@
+from typing import ClassVar
+
+import torch
+
+from .calculator import AIMNet2Calculator
+
+try:
+    import pysisyphus.run  # type: ignore
+    from pysisyphus.calculators.Calculator import Calculator  # type: ignore
+    from pysisyphus.constants import ANG2BOHR, AU2EV, BOHR2ANG  # type: ignore
+    from pysisyphus.elem_data import ATOMIC_NUMBERS  # type: ignore
+except ImportError:
+    raise ImportError("Pysisyphus is not installed. Please install Pysisyphus to use this module.") from None
+
+EV2AU = 1 / AU2EV
+
+
+class AIMNet2Pysis(Calculator):
+    implemented_properties: ClassVar = ["energy", "forces", "free_energy", "charges", "stress"]
+
+    def __init__(self, model: AIMNet2Calculator | str = "aimnet2", charge=0, mult=1, **kwargs):
+        super().__init__(charge=charge, mult=mult, **kwargs)
+        if isinstance(model, str):
+            model = AIMNet2Calculator(model)
+        self.model = model
+
+    def _prepere_input(self, atoms, coord):
+        device = self.model.device
+        numbers = torch.as_tensor([ATOMIC_NUMBERS[a.lower()] for a in atoms], device=device)
+        coord = torch.as_tensor(coord, dtype=torch.float, device=device).view(-1, 3) * BOHR2ANG
+        charge = torch.as_tensor([self.charge], dtype=torch.float, device=device)
+        mult = torch.as_tensor([self.mult], dtype=torch.float, device=device)
+        return {"coord": coord, "numbers": numbers, "charge": charge, "mult": mult}
+
+    @staticmethod
+    def _results_get_energy(results):
+        return results["energy"].item() * EV2AU
+
+    @staticmethod
+    def _results_get_forces(results):
+        return (results["forces"].detach() * (EV2AU / ANG2BOHR)).flatten().to(torch.double).cpu().numpy()
+
+    @staticmethod
+    def _results_get_hessian(results):
+        return (
+            (results["hessian"].flatten(0, 1).flatten(-2, -1) * (EV2AU / ANG2BOHR / ANG2BOHR))
+            .to(torch.double)
+            .cpu()
+            .numpy()
+        )
+
+    def get_energy(self, atoms, coords):
+        _in = self._prepere_input(atoms, coords)
+        res = self.model(_in)
+        energy = self._results_get_energy(res)
+        return {"energy": energy}
+
+    def get_forces(self, atoms, coords):
+        _in = self._prepere_input(atoms, coords)
+        res = self.model(_in, forces=True)
+        energy = self._results_get_energy(res)
+        forces = self._results_get_forces(res)
+        return {"energy": energy, "forces": forces}
+
+    def get_hessian(self, atoms, coords):
+        _in = self._prepere_input(atoms, coords)
+        res = self.model(_in, forces=True, hessian=True)
+        energy = self._results_get_energy(res)
+        forces = self._results_get_forces(res)
+        hessian = self._results_get_hessian(res)
+        return {"energy": energy, "forces": forces, "hessian": hessian}
+
+
+def run_pysis():
+    pysisyphus.run.CALC_DICT["aimnet"] = AIMNet2Pysis
+    pysisyphus.run.run()