diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..bf8a436 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +aimnet2/_version.py export-subst diff --git a/.github/workflows/CI.yaml b/.github/workflows/CI.yaml new file mode 100644 index 0000000..5f5afa6 --- /dev/null +++ b/.github/workflows/CI.yaml @@ -0,0 +1,47 @@ +name: CI + +on: + push: + branches: + - main + pull_request: + branches: + - main + schedule: + - cron: "0 0 * * *" + workflow_dispatch: + +defaults: + run: + shell: bash -l {0} + +jobs: + test: + name: Test on ${{ matrix.os }}, Python ${{ matrix.python-version }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: + - macOS-latest + - ubuntu-latest + python-version: + - "3.10" + + steps: + - uses: actions/checkout@v4 + + - name: Install conda environment + uses: mamba-org/setup-micromamba@v1 + with: + environment-file: devtools/conda-envs/test_env.yaml + create-args: >- + python=${{ matrix.python-version }} + + - name: Install package + run: | + python -m pip install . + + - name: Run tests + run: | + pytest -v --color=yes pyaimnet2/tests/ \ No newline at end of file diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..185efc1 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +recursive-include pyaimnet2 * \ No newline at end of file diff --git a/README.md b/README.md index 9c8f4d2..e76bd1e 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,47 @@ Output is a dictionary with the following keys: energy: shape (m, ) - energy in eV charges: shape (m, n) - partial atomic charges ``` + +## Installation +To make accessing the models simpler this package can be installed into a python environment. It is recommended that you +create a new conda environment with the required dependencies to run the AIMNET2 model and provided calculator interfaces. + You can create the environment using the following command +```bash +mamba create -n aimnet2 -c conda-forge 'pytorch=2' numpy ase +``` +you should then activate the environment using +```bash +conda activate aimnet2 +``` +to use the `pysisyphus` calculator you will need to install the package from PyPI using pip +```bash +pip install pysisyphus +``` +you should then clone this package and install from source via +```bash +git clone https://github.com/isayevlab/AIMNet2.git +cd AIMNet2 +pip install . +``` + +## Loading models via python +The aimnet2 package provides a convenience function to load the ensemble models + +```python +from pyaimnet2 import load_model + +model = load_model("wb97m-d3") # can also load b973c +``` +the model can then be used by providing the inputs as described in [Models](#models) or used with one of the provided +calculators like ASE + +```python +from pyaimnet2.calculators.aimnet2ase import AIMNet2Calculator + +calculator = AIMNet2Calculator(model=model) +``` + + ## Calculators We provide example code for AIMNet2 calculators for [ASE](https://wiki.fysik.dtu.dk/ase) and [pysisyphus](https://pysisyphus.readthedocs.io/) Python libraries. The code shows an example use of the AIMNet2 models. diff --git a/calculators/aimnet2ase.py b/calculators/aimnet2ase.py deleted file mode 100644 index 0fe73c0..0000000 --- a/calculators/aimnet2ase.py +++ /dev/null @@ -1,70 +0,0 @@ -import torch -import torch.nn.functional as F -import ase.calculators.calculator -import numpy as np - - -class AIMNet2Calculator(ase.calculators.calculator.Calculator): - """ ASE calculator for AIMNet2 model - Arguments: - model (:class:`torch.nn.Module`): AIMNet2 model - charge (int or float): molecular charge. Default: 0 - """ - - implemented_properties = ['energy', 'forces', 'free_energy', 'charges'] - - def __init__(self, model, charge=0): - super().__init__() - self.model = model - self.charge = charge - self.device = next(model.parameters()).device - cutoff = max(v.item() for k, v in model.state_dict().items() if k.endswith('aev.rc_s')) - self.cutoff = float(cutoff) - self._t_numbers = None - self._t_charge = None - - def do_reset(self): - self._t_numbers = None - self._t_charge = None - self.charge = 0.0 - - def set_charge(self, charge): - self.charge = float(charge) - - def _make_input(self): - coord = torch.as_tensor(self.atoms.positions).to(torch.float).to(self.device).unsqueeze(0) - if self._t_numbers is None: - self._t_numbers = torch.as_tensor(self.atoms.numbers).to(torch.long).to(self.device).unsqueeze(0) - self._t_charge = torch.tensor([self.charge], dtype=torch.float, device=self.device) - d = dict(coord=coord, numbers=self._t_numbers, charge=self._t_charge) - return d - - def _eval_model(self, d, forces=True): - prev = torch.is_grad_enabled() - torch._C._set_grad_enabled(forces) - if forces: - d['coord'].requires_grad_(True) - _out = self.model(d) - ret = dict(energy=_out['energy'].item(), charges=_out['charges'].detach()[0].cpu().numpy()) - if forces: - if 'forces' in _out: - f = _out['forces'][0] - else: - f = - torch.autograd.grad(_out['energy'], d['coord'])[0][0] - ret['forces'] = f.detach().cpu().numpy() - torch._C._set_grad_enabled(prev) - return ret - - def calculate(self, atoms=None, properties=['energy'], - system_changes=ase.calculators.calculator.all_changes): - super().calculate(atoms, properties, system_changes) - _in = self._make_input() - do_forces = 'forces' in properties - _out = self._eval_model(_in, do_forces) - - self.results['energy'] = _out['energy'] - self.results['charges'] = _out['charges'] - if do_forces: - self.results['forces'] = _out['forces'] - - diff --git a/devtools/conda-envs/test_env.yaml b/devtools/conda-envs/test_env.yaml new file mode 100644 index 0000000..855d6e5 --- /dev/null +++ b/devtools/conda-envs/test_env.yaml @@ -0,0 +1,12 @@ +name: pyaminet2 +channels: + - conda-forge +dependencies: + - python + - pip + - ase + - numpy + - pytorch==2.0.0 + + # testing + - pytest diff --git a/models/aimnet2_b973c_0.jpt b/models/aimnet2_b973c_0.jpt deleted file mode 100644 index 94ae30d..0000000 Binary files a/models/aimnet2_b973c_0.jpt and /dev/null differ diff --git a/models/aimnet2_b973c_1.jpt b/models/aimnet2_b973c_1.jpt deleted file mode 100644 index f752abe..0000000 Binary files a/models/aimnet2_b973c_1.jpt and /dev/null differ diff --git a/models/aimnet2_b973c_2.jpt b/models/aimnet2_b973c_2.jpt deleted file mode 100644 index 22a962f..0000000 Binary files a/models/aimnet2_b973c_2.jpt and /dev/null differ diff --git a/models/aimnet2_b973c_3.jpt b/models/aimnet2_b973c_3.jpt deleted file mode 100644 index 482c91e..0000000 Binary files a/models/aimnet2_b973c_3.jpt and /dev/null differ diff --git a/models/aimnet2_wb97m-d3_0.jpt b/models/aimnet2_wb97m-d3_0.jpt deleted file mode 100644 index baf3d23..0000000 Binary files a/models/aimnet2_wb97m-d3_0.jpt and /dev/null differ diff --git a/models/aimnet2_wb97m-d3_1.jpt b/models/aimnet2_wb97m-d3_1.jpt deleted file mode 100644 index 1ed25e7..0000000 Binary files a/models/aimnet2_wb97m-d3_1.jpt and /dev/null differ diff --git a/models/aimnet2_wb97m-d3_2.jpt b/models/aimnet2_wb97m-d3_2.jpt deleted file mode 100644 index 86d0b20..0000000 Binary files a/models/aimnet2_wb97m-d3_2.jpt and /dev/null differ diff --git a/models/aimnet2_wb97m-d3_3.jpt b/models/aimnet2_wb97m-d3_3.jpt deleted file mode 100644 index 1c358b9..0000000 Binary files a/models/aimnet2_wb97m-d3_3.jpt and /dev/null differ diff --git a/pyaimnet2/__init__.py b/pyaimnet2/__init__.py new file mode 100644 index 0000000..d2232ad --- /dev/null +++ b/pyaimnet2/__init__.py @@ -0,0 +1,7 @@ +from . import _version +from pyaimnet2.utils import load_model + + +__all__ = ["load_model"] + +__version__ = _version.get_versions()["version"] diff --git a/pyaimnet2/_version.py b/pyaimnet2/_version.py new file mode 100644 index 0000000..8cdee39 --- /dev/null +++ b/pyaimnet2/_version.py @@ -0,0 +1,716 @@ +# This file helps to compute a version number in source trees obtained from +# git-archive tarball (such as those provided by githubs download-from-tag +# feature). Distribution tarballs (built by setup.py sdist) and build +# directories (produced by setup.py build) will contain a much shorter file +# that just contains the computed version number. + +# This file is released into the public domain. +# Generated by versioneer-0.29 +# https://github.com/python-versioneer/python-versioneer + +"""Git implementation of _version.py.""" + +import errno +import os +import re +import subprocess +import sys +from typing import Any, Callable, Dict, List, Optional, Tuple +import functools + + +def get_keywords() -> Dict[str, str]: + """Get the keywords needed to look up the version information.""" + # these strings will be replaced by git during git-archive. + # setup.py/versioneer.py will grep for the variable names, so they must + # each be defined on a line of their own. _version.py will just call + # get_keywords(). + git_refnames = "$Format:%d$" + git_full = "$Format:%H$" + git_date = "$Format:%ci$" + keywords = {"refnames": git_refnames, "full": git_full, "date": git_date} + return keywords + + +class VersioneerConfig: + """Container for Versioneer configuration parameters.""" + + VCS: str + style: str + tag_prefix: str + parentdir_prefix: str + versionfile_source: str + verbose: bool + + +def get_config() -> VersioneerConfig: + """Create, populate and return the VersioneerConfig() object.""" + # these strings are filled in when 'setup.py versioneer' creates + # _version.py + cfg = VersioneerConfig() + cfg.VCS = "git" + cfg.style = "pep440" + cfg.tag_prefix = "" + cfg.parentdir_prefix = "pyaimnet2-" + cfg.versionfile_source = "pyaimnet2/_version.py" + cfg.verbose = False + return cfg + + +class NotThisMethod(Exception): + """Exception raised if a method is not valid for the current scenario.""" + + +LONG_VERSION_PY: Dict[str, str] = {} +HANDLERS: Dict[str, Dict[str, Callable]] = {} + + +def register_vcs_handler(vcs: str, method: str) -> Callable: # decorator + """Create decorator to mark a method as the handler of a VCS.""" + + def decorate(f: Callable) -> Callable: + """Store f in HANDLERS[vcs][method].""" + if vcs not in HANDLERS: + HANDLERS[vcs] = {} + HANDLERS[vcs][method] = f + return f + + return decorate + + +def run_command( + commands: List[str], + args: List[str], + cwd: Optional[str] = None, + verbose: bool = False, + hide_stderr: bool = False, + env: Optional[Dict[str, str]] = None, +) -> Tuple[Optional[str], Optional[int]]: + """Call the given command(s).""" + assert isinstance(commands, list) + process = None + + popen_kwargs: Dict[str, Any] = {} + if sys.platform == "win32": + # This hides the console window if pythonw.exe is used + startupinfo = subprocess.STARTUPINFO() + startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW + popen_kwargs["startupinfo"] = startupinfo + + for command in commands: + try: + dispcmd = str([command] + args) + # remember shell=False, so use git.cmd on windows, not just git + process = subprocess.Popen( + [command] + args, + cwd=cwd, + env=env, + stdout=subprocess.PIPE, + stderr=(subprocess.PIPE if hide_stderr else None), + **popen_kwargs, + ) + break + except OSError as e: + if e.errno == errno.ENOENT: + continue + if verbose: + print("unable to run %s" % dispcmd) + print(e) + return None, None + else: + if verbose: + print("unable to find command, tried %s" % (commands,)) + return None, None + stdout = process.communicate()[0].strip().decode() + if process.returncode != 0: + if verbose: + print("unable to run %s (error)" % dispcmd) + print("stdout was %s" % stdout) + return None, process.returncode + return stdout, process.returncode + + +def versions_from_parentdir( + parentdir_prefix: str, + root: str, + verbose: bool, +) -> Dict[str, Any]: + """Try to determine the version from the parent directory name. + + Source tarballs conventionally unpack into a directory that includes both + the project name and a version string. We will also support searching up + two directory levels for an appropriately named parent directory + """ + rootdirs = [] + + for _ in range(3): + dirname = os.path.basename(root) + if dirname.startswith(parentdir_prefix): + return { + "version": dirname[len(parentdir_prefix) :], + "full-revisionid": None, + "dirty": False, + "error": None, + "date": None, + } + rootdirs.append(root) + root = os.path.dirname(root) # up a level + + if verbose: + print( + "Tried directories %s but none started with prefix %s" + % (str(rootdirs), parentdir_prefix) + ) + raise NotThisMethod("rootdir doesn't start with parentdir_prefix") + + +@register_vcs_handler("git", "get_keywords") +def git_get_keywords(versionfile_abs: str) -> Dict[str, str]: + """Extract version information from the given file.""" + # the code embedded in _version.py can just fetch the value of these + # keywords. When used from setup.py, we don't want to import _version.py, + # so we do it with a regexp instead. This function is not used from + # _version.py. + keywords: Dict[str, str] = {} + try: + with open(versionfile_abs, "r") as fobj: + for line in fobj: + if line.strip().startswith("git_refnames ="): + mo = re.search(r'=\s*"(.*)"', line) + if mo: + keywords["refnames"] = mo.group(1) + if line.strip().startswith("git_full ="): + mo = re.search(r'=\s*"(.*)"', line) + if mo: + keywords["full"] = mo.group(1) + if line.strip().startswith("git_date ="): + mo = re.search(r'=\s*"(.*)"', line) + if mo: + keywords["date"] = mo.group(1) + except OSError: + pass + return keywords + + +@register_vcs_handler("git", "keywords") +def git_versions_from_keywords( + keywords: Dict[str, str], + tag_prefix: str, + verbose: bool, +) -> Dict[str, Any]: + """Get version information from git keywords.""" + if "refnames" not in keywords: + raise NotThisMethod("Short version file found") + date = keywords.get("date") + if date is not None: + # Use only the last line. Previous lines may contain GPG signature + # information. + date = date.splitlines()[-1] + + # git-2.2.0 added "%cI", which expands to an ISO-8601 -compliant + # datestamp. However we prefer "%ci" (which expands to an "ISO-8601 + # -like" string, which we must then edit to make compliant), because + # it's been around since git-1.5.3, and it's too difficult to + # discover which version we're using, or to work around using an + # older one. + date = date.strip().replace(" ", "T", 1).replace(" ", "", 1) + refnames = keywords["refnames"].strip() + if refnames.startswith("$Format"): + if verbose: + print("keywords are unexpanded, not using") + raise NotThisMethod("unexpanded keywords, not a git-archive tarball") + refs = {r.strip() for r in refnames.strip("()").split(",")} + # starting in git-1.8.3, tags are listed as "tag: foo-1.0" instead of + # just "foo-1.0". If we see a "tag: " prefix, prefer those. + TAG = "tag: " + tags = {r[len(TAG) :] for r in refs if r.startswith(TAG)} + if not tags: + # Either we're using git < 1.8.3, or there really are no tags. We use + # a heuristic: assume all version tags have a digit. The old git %d + # expansion behaves like git log --decorate=short and strips out the + # refs/heads/ and refs/tags/ prefixes that would let us distinguish + # between branches and tags. By ignoring refnames without digits, we + # filter out many common branch names like "release" and + # "stabilization", as well as "HEAD" and "master". + tags = {r for r in refs if re.search(r"\d", r)} + if verbose: + print("discarding '%s', no digits" % ",".join(refs - tags)) + if verbose: + print("likely tags: %s" % ",".join(sorted(tags))) + for ref in sorted(tags): + # sorting will prefer e.g. "2.0" over "2.0rc1" + if ref.startswith(tag_prefix): + r = ref[len(tag_prefix) :] + # Filter out refs that exactly match prefix or that don't start + # with a number once the prefix is stripped (mostly a concern + # when prefix is '') + if not re.match(r"\d", r): + continue + if verbose: + print("picking %s" % r) + return { + "version": r, + "full-revisionid": keywords["full"].strip(), + "dirty": False, + "error": None, + "date": date, + } + # no suitable tags, so version is "0+unknown", but full hex is still there + if verbose: + print("no suitable tags, using unknown + full revision id") + return { + "version": "0+unknown", + "full-revisionid": keywords["full"].strip(), + "dirty": False, + "error": "no suitable tags", + "date": None, + } + + +@register_vcs_handler("git", "pieces_from_vcs") +def git_pieces_from_vcs( + tag_prefix: str, root: str, verbose: bool, runner: Callable = run_command +) -> Dict[str, Any]: + """Get version from 'git describe' in the root of the source tree. + + This only gets called if the git-archive 'subst' keywords were *not* + expanded, and _version.py hasn't already been rewritten with a short + version string, meaning we're inside a checked out source tree. + """ + GITS = ["git"] + if sys.platform == "win32": + GITS = ["git.cmd", "git.exe"] + + # GIT_DIR can interfere with correct operation of Versioneer. + # It may be intended to be passed to the Versioneer-versioned project, + # but that should not change where we get our version from. + env = os.environ.copy() + env.pop("GIT_DIR", None) + runner = functools.partial(runner, env=env) + + _, rc = runner(GITS, ["rev-parse", "--git-dir"], cwd=root, hide_stderr=not verbose) + if rc != 0: + if verbose: + print("Directory %s not under git control" % root) + raise NotThisMethod("'git rev-parse --git-dir' returned error") + + # if there is a tag matching tag_prefix, this yields TAG-NUM-gHEX[-dirty] + # if there isn't one, this yields HEX[-dirty] (no NUM) + describe_out, rc = runner( + GITS, + [ + "describe", + "--tags", + "--dirty", + "--always", + "--long", + "--match", + f"{tag_prefix}[[:digit:]]*", + ], + cwd=root, + ) + # --long was added in git-1.5.5 + if describe_out is None: + raise NotThisMethod("'git describe' failed") + describe_out = describe_out.strip() + full_out, rc = runner(GITS, ["rev-parse", "HEAD"], cwd=root) + if full_out is None: + raise NotThisMethod("'git rev-parse' failed") + full_out = full_out.strip() + + pieces: Dict[str, Any] = {} + pieces["long"] = full_out + pieces["short"] = full_out[:7] # maybe improved later + pieces["error"] = None + + branch_name, rc = runner(GITS, ["rev-parse", "--abbrev-ref", "HEAD"], cwd=root) + # --abbrev-ref was added in git-1.6.3 + if rc != 0 or branch_name is None: + raise NotThisMethod("'git rev-parse --abbrev-ref' returned error") + branch_name = branch_name.strip() + + if branch_name == "HEAD": + # If we aren't exactly on a branch, pick a branch which represents + # the current commit. If all else fails, we are on a branchless + # commit. + branches, rc = runner(GITS, ["branch", "--contains"], cwd=root) + # --contains was added in git-1.5.4 + if rc != 0 or branches is None: + raise NotThisMethod("'git branch --contains' returned error") + branches = branches.split("\n") + + # Remove the first line if we're running detached + if "(" in branches[0]: + branches.pop(0) + + # Strip off the leading "* " from the list of branches. + branches = [branch[2:] for branch in branches] + if "master" in branches: + branch_name = "master" + elif not branches: + branch_name = None + else: + # Pick the first branch that is returned. Good or bad. + branch_name = branches[0] + + pieces["branch"] = branch_name + + # parse describe_out. It will be like TAG-NUM-gHEX[-dirty] or HEX[-dirty] + # TAG might have hyphens. + git_describe = describe_out + + # look for -dirty suffix + dirty = git_describe.endswith("-dirty") + pieces["dirty"] = dirty + if dirty: + git_describe = git_describe[: git_describe.rindex("-dirty")] + + # now we have TAG-NUM-gHEX or HEX + + if "-" in git_describe: + # TAG-NUM-gHEX + mo = re.search(r"^(.+)-(\d+)-g([0-9a-f]+)$", git_describe) + if not mo: + # unparsable. Maybe git-describe is misbehaving? + pieces["error"] = "unable to parse git-describe output: '%s'" % describe_out + return pieces + + # tag + full_tag = mo.group(1) + if not full_tag.startswith(tag_prefix): + if verbose: + fmt = "tag '%s' doesn't start with prefix '%s'" + print(fmt % (full_tag, tag_prefix)) + pieces["error"] = "tag '%s' doesn't start with prefix '%s'" % ( + full_tag, + tag_prefix, + ) + return pieces + pieces["closest-tag"] = full_tag[len(tag_prefix) :] + + # distance: number of commits since tag + pieces["distance"] = int(mo.group(2)) + + # commit: short hex revision ID + pieces["short"] = mo.group(3) + + else: + # HEX: no tags + pieces["closest-tag"] = None + out, rc = runner(GITS, ["rev-list", "HEAD", "--left-right"], cwd=root) + pieces["distance"] = len(out.split()) # total number of commits + + # commit date: see ISO-8601 comment in git_versions_from_keywords() + date = runner(GITS, ["show", "-s", "--format=%ci", "HEAD"], cwd=root)[0].strip() + # Use only the last line. Previous lines may contain GPG signature + # information. + date = date.splitlines()[-1] + pieces["date"] = date.strip().replace(" ", "T", 1).replace(" ", "", 1) + + return pieces + + +def plus_or_dot(pieces: Dict[str, Any]) -> str: + """Return a + if we don't already have one, else return a .""" + if "+" in pieces.get("closest-tag", ""): + return "." + return "+" + + +def render_pep440(pieces: Dict[str, Any]) -> str: + """Build up version string, with post-release "local version identifier". + + Our goal: TAG[+DISTANCE.gHEX[.dirty]] . Note that if you + get a tagged build and then dirty it, you'll get TAG+0.gHEX.dirty + + Exceptions: + 1: no tags. git_describe was just HEX. 0+untagged.DISTANCE.gHEX[.dirty] + """ + if pieces["closest-tag"]: + rendered = pieces["closest-tag"] + if pieces["distance"] or pieces["dirty"]: + rendered += plus_or_dot(pieces) + rendered += "%d.g%s" % (pieces["distance"], pieces["short"]) + if pieces["dirty"]: + rendered += ".dirty" + else: + # exception #1 + rendered = "0+untagged.%d.g%s" % (pieces["distance"], pieces["short"]) + if pieces["dirty"]: + rendered += ".dirty" + return rendered + + +def render_pep440_branch(pieces: Dict[str, Any]) -> str: + """TAG[[.dev0]+DISTANCE.gHEX[.dirty]] . + + The ".dev0" means not master branch. Note that .dev0 sorts backwards + (a feature branch will appear "older" than the master branch). + + Exceptions: + 1: no tags. 0[.dev0]+untagged.DISTANCE.gHEX[.dirty] + """ + if pieces["closest-tag"]: + rendered = pieces["closest-tag"] + if pieces["distance"] or pieces["dirty"]: + if pieces["branch"] != "master": + rendered += ".dev0" + rendered += plus_or_dot(pieces) + rendered += "%d.g%s" % (pieces["distance"], pieces["short"]) + if pieces["dirty"]: + rendered += ".dirty" + else: + # exception #1 + rendered = "0" + if pieces["branch"] != "master": + rendered += ".dev0" + rendered += "+untagged.%d.g%s" % (pieces["distance"], pieces["short"]) + if pieces["dirty"]: + rendered += ".dirty" + return rendered + + +def pep440_split_post(ver: str) -> Tuple[str, Optional[int]]: + """Split pep440 version string at the post-release segment. + + Returns the release segments before the post-release and the + post-release version number (or -1 if no post-release segment is present). + """ + vc = str.split(ver, ".post") + return vc[0], int(vc[1] or 0) if len(vc) == 2 else None + + +def render_pep440_pre(pieces: Dict[str, Any]) -> str: + """TAG[.postN.devDISTANCE] -- No -dirty. + + Exceptions: + 1: no tags. 0.post0.devDISTANCE + """ + if pieces["closest-tag"]: + if pieces["distance"]: + # update the post release segment + tag_version, post_version = pep440_split_post(pieces["closest-tag"]) + rendered = tag_version + if post_version is not None: + rendered += ".post%d.dev%d" % (post_version + 1, pieces["distance"]) + else: + rendered += ".post0.dev%d" % (pieces["distance"]) + else: + # no commits, use the tag as the version + rendered = pieces["closest-tag"] + else: + # exception #1 + rendered = "0.post0.dev%d" % pieces["distance"] + return rendered + + +def render_pep440_post(pieces: Dict[str, Any]) -> str: + """TAG[.postDISTANCE[.dev0]+gHEX] . + + The ".dev0" means dirty. Note that .dev0 sorts backwards + (a dirty tree will appear "older" than the corresponding clean one), + but you shouldn't be releasing software with -dirty anyways. + + Exceptions: + 1: no tags. 0.postDISTANCE[.dev0] + """ + if pieces["closest-tag"]: + rendered = pieces["closest-tag"] + if pieces["distance"] or pieces["dirty"]: + rendered += ".post%d" % pieces["distance"] + if pieces["dirty"]: + rendered += ".dev0" + rendered += plus_or_dot(pieces) + rendered += "g%s" % pieces["short"] + else: + # exception #1 + rendered = "0.post%d" % pieces["distance"] + if pieces["dirty"]: + rendered += ".dev0" + rendered += "+g%s" % pieces["short"] + return rendered + + +def render_pep440_post_branch(pieces: Dict[str, Any]) -> str: + """TAG[.postDISTANCE[.dev0]+gHEX[.dirty]] . + + The ".dev0" means not master branch. + + Exceptions: + 1: no tags. 0.postDISTANCE[.dev0]+gHEX[.dirty] + """ + if pieces["closest-tag"]: + rendered = pieces["closest-tag"] + if pieces["distance"] or pieces["dirty"]: + rendered += ".post%d" % pieces["distance"] + if pieces["branch"] != "master": + rendered += ".dev0" + rendered += plus_or_dot(pieces) + rendered += "g%s" % pieces["short"] + if pieces["dirty"]: + rendered += ".dirty" + else: + # exception #1 + rendered = "0.post%d" % pieces["distance"] + if pieces["branch"] != "master": + rendered += ".dev0" + rendered += "+g%s" % pieces["short"] + if pieces["dirty"]: + rendered += ".dirty" + return rendered + + +def render_pep440_old(pieces: Dict[str, Any]) -> str: + """TAG[.postDISTANCE[.dev0]] . + + The ".dev0" means dirty. + + Exceptions: + 1: no tags. 0.postDISTANCE[.dev0] + """ + if pieces["closest-tag"]: + rendered = pieces["closest-tag"] + if pieces["distance"] or pieces["dirty"]: + rendered += ".post%d" % pieces["distance"] + if pieces["dirty"]: + rendered += ".dev0" + else: + # exception #1 + rendered = "0.post%d" % pieces["distance"] + if pieces["dirty"]: + rendered += ".dev0" + return rendered + + +def render_git_describe(pieces: Dict[str, Any]) -> str: + """TAG[-DISTANCE-gHEX][-dirty]. + + Like 'git describe --tags --dirty --always'. + + Exceptions: + 1: no tags. HEX[-dirty] (note: no 'g' prefix) + """ + if pieces["closest-tag"]: + rendered = pieces["closest-tag"] + if pieces["distance"]: + rendered += "-%d-g%s" % (pieces["distance"], pieces["short"]) + else: + # exception #1 + rendered = pieces["short"] + if pieces["dirty"]: + rendered += "-dirty" + return rendered + + +def render_git_describe_long(pieces: Dict[str, Any]) -> str: + """TAG-DISTANCE-gHEX[-dirty]. + + Like 'git describe --tags --dirty --always -long'. + The distance/hash is unconditional. + + Exceptions: + 1: no tags. HEX[-dirty] (note: no 'g' prefix) + """ + if pieces["closest-tag"]: + rendered = pieces["closest-tag"] + rendered += "-%d-g%s" % (pieces["distance"], pieces["short"]) + else: + # exception #1 + rendered = pieces["short"] + if pieces["dirty"]: + rendered += "-dirty" + return rendered + + +def render(pieces: Dict[str, Any], style: str) -> Dict[str, Any]: + """Render the given version pieces into the requested style.""" + if pieces["error"]: + return { + "version": "unknown", + "full-revisionid": pieces.get("long"), + "dirty": None, + "error": pieces["error"], + "date": None, + } + + if not style or style == "default": + style = "pep440" # the default + + if style == "pep440": + rendered = render_pep440(pieces) + elif style == "pep440-branch": + rendered = render_pep440_branch(pieces) + elif style == "pep440-pre": + rendered = render_pep440_pre(pieces) + elif style == "pep440-post": + rendered = render_pep440_post(pieces) + elif style == "pep440-post-branch": + rendered = render_pep440_post_branch(pieces) + elif style == "pep440-old": + rendered = render_pep440_old(pieces) + elif style == "git-describe": + rendered = render_git_describe(pieces) + elif style == "git-describe-long": + rendered = render_git_describe_long(pieces) + else: + raise ValueError("unknown style '%s'" % style) + + return { + "version": rendered, + "full-revisionid": pieces["long"], + "dirty": pieces["dirty"], + "error": None, + "date": pieces.get("date"), + } + + +def get_versions() -> Dict[str, Any]: + """Get version information or return default if unable to do so.""" + # I am in _version.py, which lives at ROOT/VERSIONFILE_SOURCE. If we have + # __file__, we can work backwards from there to the root. Some + # py2exe/bbfreeze/non-CPython implementations don't do __file__, in which + # case we can only use expanded keywords. + + cfg = get_config() + verbose = cfg.verbose + + try: + return git_versions_from_keywords(get_keywords(), cfg.tag_prefix, verbose) + except NotThisMethod: + pass + + try: + root = os.path.realpath(__file__) + # versionfile_source is the relative path from the top of the source + # tree (where the .git directory might live) to this file. Invert + # this to find the root from __file__. + for _ in cfg.versionfile_source.split("/"): + root = os.path.dirname(root) + except NameError: + return { + "version": "0+unknown", + "full-revisionid": None, + "dirty": None, + "error": "unable to find root of source tree", + "date": None, + } + + try: + pieces = git_pieces_from_vcs(cfg.tag_prefix, root, verbose) + return render(pieces, cfg.style) + except NotThisMethod: + pass + + try: + if cfg.parentdir_prefix: + return versions_from_parentdir(cfg.parentdir_prefix, root, verbose) + except NotThisMethod: + pass + + return { + "version": "0+unknown", + "full-revisionid": None, + "dirty": None, + "error": "unable to compute version", + "date": None, + } diff --git a/pyaimnet2/calculators/__init__.py b/pyaimnet2/calculators/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/calculators/aimnet2_ase_opt.py b/pyaimnet2/calculators/aimnet2_ase_opt.py similarity index 70% rename from calculators/aimnet2_ase_opt.py rename to pyaimnet2/calculators/aimnet2_ase_opt.py index c0d8f8e..a35b555 100644 --- a/calculators/aimnet2_ase_opt.py +++ b/pyaimnet2/calculators/aimnet2_ase_opt.py @@ -27,7 +27,7 @@ def update_mol(mol, atoms, align=True): mol_old = pybel.Molecule(pybel.ob.OBMol(mol.OBMol)) # update coord for i, c in enumerate(atoms.get_positions()): - mol.OBMol.GetAtom(i+1).SetVector(*c.tolist()) + mol.OBMol.GetAtom(i + 1).SetVector(*c.tolist()) # align if align: aligner = pybel.ob.OBAlign(False, False) @@ -36,16 +36,16 @@ def update_mol(mol, atoms, align=True): aligner.Align() rmsd = aligner.GetRMSD() aligner.UpdateCoords(mol.OBMol) - print(f'RMSD: {rmsd:.2f} Angs') + print(f"RMSD: {rmsd:.2f} Angs") def guess_pybel_type(filename): - assert '.' in filename + assert "." in filename return os.path.splitext(filename)[1][1:] def guess_charge(mol): - m = re.search('charge: (-?\d+)', mol.title) + m = re.search("charge: (-?\d+)", mol.title) if m: charge = int(m.group(1)) else: @@ -53,31 +53,38 @@ def guess_charge(mol): return charge -if __name__ == '__main__': +if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() - parser.add_argument('--charge', type=int, default=None, help='Molecular charge (default: check molecule title for "charge: {int}" or get total charge from OpenBabel).') - parser.add_argument('--traj', help='Trajectory file', type=str, default=None) - parser.add_argument('--fmax', type=float, default=5e-3, help='Optimization threshold.') - parser.add_argument('model') - parser.add_argument('in_file') - parser.add_argument('out_file') + parser.add_argument( + "--charge", + type=int, + default=None, + help='Molecular charge (default: check molecule title for "charge: {int}" or get total charge from OpenBabel).', + ) + parser.add_argument("--traj", help="Trajectory file", type=str, default=None) + parser.add_argument( + "--fmax", type=float, default=5e-3, help="Optimization threshold." + ) + parser.add_argument("model") + parser.add_argument("in_file") + parser.add_argument("out_file") args = parser.parse_args() torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - print('Loading AIMNet2 model from file', args.model) + print("Loading AIMNet2 model from file", args.model) model = torch.jit.load(args.model, map_location=device) calc = AIMNet2Calculator(model) in_format = guess_pybel_type(args.in_file) out_format = guess_pybel_type(args.out_file) - with open(args.out_file, 'w') as f: + with open(args.out_file, "w") as f: for mol in pybel.readfile(in_format, args.in_file): atoms = pybel2atoms(mol) charge = args.charge if args.charge is not None else guess_charge(mol) diff --git a/calculators/aimnet2_sph_opt.py b/pyaimnet2/calculators/aimnet2_sph_opt.py similarity index 65% rename from calculators/aimnet2_sph_opt.py rename to pyaimnet2/calculators/aimnet2_sph_opt.py index 6d1013e..881e419 100644 --- a/calculators/aimnet2_sph_opt.py +++ b/pyaimnet2/calculators/aimnet2_sph_opt.py @@ -10,7 +10,7 @@ from pysisyphus.helpers import do_final_hessian -def pybel2geom(mol, coord_type='dlc'): +def pybel2geom(mol, coord_type="dlc"): coord = np.array([a.coords for a in mol.atoms]).flatten() * ANG2BOHR atoms = [INV_ATOMIC_NUMBERS[a.atomicnum].lower() for a in mol.atoms] geom = Geometry(atoms, coord, coord_type=coord_type) @@ -23,7 +23,7 @@ def update_mol(mol, geom, align=True): # update coord coord = geom.coords3d / ANG2BOHR for i, c in enumerate(coord): - mol.OBMol.GetAtom(i+1).SetVector(*c.tolist()) + mol.OBMol.GetAtom(i + 1).SetVector(*c.tolist()) # align if align: aligner = pybel.ob.OBAlign(False, False) @@ -32,16 +32,16 @@ def update_mol(mol, geom, align=True): aligner.Align() rmsd = aligner.GetRMSD() aligner.UpdateCoords(mol.OBMol) - print(f'RMSD: {rmsd:.2f} Angs') + print(f"RMSD: {rmsd:.2f} Angs") def guess_pybel_type(filename): - assert '.' in filename + assert "." in filename return os.path.splitext(filename)[1][1:] def guess_charge(mol): - m = re.search('charge: (-?\d+)', mol.title) + m = re.search("charge: (-?\d+)", mol.title) if m: charge = int(m.group(1)) else: @@ -49,25 +49,37 @@ def guess_charge(mol): return charge -if __name__ == '__main__': +if __name__ == "__main__": import argparse parser = argparse.ArgumentParser() - parser.add_argument('--charge', type=int, default=None, help='Molecular charge (default: check molecule title for "charge: {int}" or get total charge from OpenBabel).') - parser.add_argument('--ts', action='store_true', help='Do TS optimization') - parser.add_argument('--coord', type=str, help='Coordinates type, e.g. cart, dlc (default) or redund') - parser.add_argument('--thresh', type=str, default='gau_loose', help='Optimization threshold, one of aaug_loose (default), gau, gau_tight, gau_vtight, baker.') - parser.add_argument('model') - parser.add_argument('in_file') - parser.add_argument('out_file') + parser.add_argument( + "--charge", + type=int, + default=None, + help='Molecular charge (default: check molecule title for "charge: {int}" or get total charge from OpenBabel).', + ) + parser.add_argument("--ts", action="store_true", help="Do TS optimization") + parser.add_argument( + "--coord", type=str, help="Coordinates type, e.g. cart, dlc (default) or redund" + ) + parser.add_argument( + "--thresh", + type=str, + default="gau_loose", + help="Optimization threshold, one of aaug_loose (default), gau, gau_tight, gau_vtight, baker.", + ) + parser.add_argument("model") + parser.add_argument("in_file") + parser.add_argument("out_file") args = parser.parse_args() torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cudnn.allow_tf32 = False - device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - print('Loading AIMNet2 model from file', args.model) + print("Loading AIMNet2 model from file", args.model) model = torch.jit.load(args.model, map_location=device) calc = AIMNet2Calculator(model) @@ -76,24 +88,28 @@ def guess_charge(mol): if args.ts: from pysisyphus.tsoptimizers import RSPRFOptimizer as Opt + opt_kwargs = dict(assert_neg_eigval=True, hessian_recalc=30) else: from pysisyphus.optimizers.RFOptimizer import RFOptimizer as Opt + opt_kwargs = dict() - with open(args.out_file, 'w') as f: + with open(args.out_file, "w") as f: for mol in pybel.readfile(in_format, args.in_file): geom = pybel2geom(mol) charge = args.charge if args.charge is not None else guess_charge(mol) calc.charge = charge geom.set_calculator(calc) opt = Opt(geom, thresh=args.thresh, max_cycles=1000, **opt_kwargs) - + with torch.jit.optimized_execution(False): opt.run() if args.ts: - do_final_hessian(geom, save_hessian=False, is_ts=args.ts, print_thermo=True) + do_final_hessian( + geom, save_hessian=False, is_ts=args.ts, print_thermo=True + ) update_mol(mol, geom, align=False) f.write(mol.write(out_format)) diff --git a/pyaimnet2/calculators/aimnet2ase.py b/pyaimnet2/calculators/aimnet2ase.py new file mode 100644 index 0000000..4db6a3f --- /dev/null +++ b/pyaimnet2/calculators/aimnet2ase.py @@ -0,0 +1,89 @@ +import torch +import torch.nn.functional as F +import ase.calculators.calculator +import numpy as np + + +class AIMNet2Calculator(ase.calculators.calculator.Calculator): + """ASE calculator for AIMNet2 model + Arguments: + model (:class:`torch.nn.Module`): AIMNet2 model + charge (int or float): molecular charge. Default: 0 + """ + + implemented_properties = ["energy", "forces", "free_energy", "charges"] + + def __init__(self, model, charge=0): + super().__init__() + self.model = model + self.charge = charge + self.device = next(model.parameters()).device + cutoff = max( + v.item() for k, v in model.state_dict().items() if k.endswith("aev.rc_s") + ) + self.cutoff = float(cutoff) + self._t_numbers = None + self._t_charge = None + + def do_reset(self): + self._t_numbers = None + self._t_charge = None + self.charge = 0.0 + + def set_charge(self, charge): + self.charge = float(charge) + + def _make_input(self): + coord = ( + torch.as_tensor(self.atoms.positions) + .to(torch.float) + .to(self.device) + .unsqueeze(0) + ) + if self._t_numbers is None: + self._t_numbers = ( + torch.as_tensor(self.atoms.numbers) + .to(torch.long) + .to(self.device) + .unsqueeze(0) + ) + self._t_charge = torch.tensor( + [self.charge], dtype=torch.float, device=self.device + ) + d = dict(coord=coord, numbers=self._t_numbers, charge=self._t_charge) + return d + + def _eval_model(self, d, forces=True): + prev = torch.is_grad_enabled() + torch._C._set_grad_enabled(forces) + if forces: + d["coord"].requires_grad_(True) + _out = self.model(d) + ret = dict( + energy=_out["energy"].item(), + charges=_out["charges"].detach()[0].cpu().numpy(), + ) + if forces: + if "forces" in _out: + f = _out["forces"][0] + else: + f = -torch.autograd.grad(_out["energy"], d["coord"])[0][0] + ret["forces"] = f.detach().cpu().numpy() + torch._C._set_grad_enabled(prev) + return ret + + def calculate( + self, + atoms=None, + properties=["energy"], + system_changes=ase.calculators.calculator.all_changes, + ): + super().calculate(atoms, properties, system_changes) + _in = self._make_input() + do_forces = "forces" in properties + _out = self._eval_model(_in, do_forces) + + self.results["energy"] = _out["energy"] + self.results["charges"] = _out["charges"] + if do_forces: + self.results["forces"] = _out["forces"] diff --git a/calculators/aimnet2sph.py b/pyaimnet2/calculators/aimnet2sph.py similarity index 57% rename from calculators/aimnet2sph.py rename to pyaimnet2/calculators/aimnet2sph.py index 17471c9..188ca31 100644 --- a/calculators/aimnet2sph.py +++ b/pyaimnet2/calculators/aimnet2sph.py @@ -12,9 +12,9 @@ class AIMNet2Calculator(Calculator): def __init__(self, model=None, charge=0, **kwargs): super().__init__(charge=charge, **kwargs) if model is None: - model = os.environ.get('AIMNET_MODEL') + model = os.environ.get("AIMNET_MODEL") if isinstance(model, str): - device = 'cuda' if torch.cuda.is_available() else 'cpu' + device = "cuda" if torch.cuda.is_available() else "cpu" model = torch.jit.load(model, map_location=device) assert isinstance(model, torch.nn.Module) self.model = model @@ -22,8 +22,15 @@ def __init__(self, model=None, charge=0, **kwargs): self.device = next(self.model.parameters()).device def _prepere_input(self, atoms, coord): - numbers = torch.as_tensor([[ATOMIC_NUMBERS[a.lower()] for a in atoms]], device=self.device) - coord = torch.as_tensor(coord, dtype=torch.float, device=self.device).view(1, numbers.shape[1], 3) * BOHR2ANG + numbers = torch.as_tensor( + [[ATOMIC_NUMBERS[a.lower()] for a in atoms]], device=self.device + ) + coord = ( + torch.as_tensor(coord, dtype=torch.float, device=self.device).view( + 1, numbers.shape[1], 3 + ) + * BOHR2ANG + ) charge = torch.as_tensor([self.charge], dtype=torch.float, device=self.device) return dict(coord=coord, numbers=numbers, charge=charge) @@ -31,44 +38,53 @@ def get_energy(self, atoms, coords): _in = self._prepere_input(atoms, coords) with torch.no_grad(), torch.jit.optimized_execution(False): _out = self.model(_in) - energy = _out['energy'].item() * EV2AU + energy = _out["energy"].item() * EV2AU return dict(energy=energy) def get_forces(self, atoms, coords): _in = self._prepere_input(atoms, coords) with torch.jit.optimized_execution(False): - _in['coord'].requires_grad_(True) + _in["coord"].requires_grad_(True) _out = self.model(_in) - e = _out['energy'] - f = - torch.autograd.grad(e, _in['coord'])[0] + e = _out["energy"] + f = -torch.autograd.grad(e, _in["coord"])[0] energy = e.item() * EV2AU forces = (f * (EV2AU / ANG2BOHR))[0].flatten().cpu().numpy() return dict(energy=energy, forces=forces) - + def get_hessian(self, atoms, coords): _in = self._prepere_input(atoms, coords) with torch.jit.optimized_execution(False): - _in['coord'].requires_grad_(True) + _in["coord"].requires_grad_(True) _out = self.model(_in) - e = _out['energy'] - f = -_get_derivatives_not_none(_in['coord'], e, create_graph=True) - h = - torch.stack([ - _get_derivatives_not_none(_in['coord'], _f, retain_graph=True)[0] + e = _out["energy"] + f = -_get_derivatives_not_none(_in["coord"], e, create_graph=True) + h = -torch.stack( + [ + _get_derivatives_not_none(_in["coord"], _f, retain_graph=True)[0] for _f in f.flatten().unbind() - ]) + ] + ) energy = e.item() * EV2AU - forces = (f.detach() * (EV2AU / ANG2BOHR))[0].flatten().to(torch.double).cpu().numpy() - hessian = (h.flatten(-2, -1) * (EV2AU / ANG2BOHR / ANG2BOHR)).to(torch.double).cpu().numpy() + forces = ( + (f.detach() * (EV2AU / ANG2BOHR))[0] + .flatten() + .to(torch.double) + .cpu() + .numpy() + ) + hessian = ( + (h.flatten(-2, -1) * (EV2AU / ANG2BOHR / ANG2BOHR)) + .to(torch.double) + .cpu() + .numpy() + ) return dict(energy=energy, forces=forces, hessian=hessian) - + def _get_derivatives_not_none(x, y, retain_graph=None, create_graph=False): - ret = torch.autograd.grad([y.sum()], [x], retain_graph=retain_graph, create_graph=create_graph)[0] + ret = torch.autograd.grad( + [y.sum()], [x], retain_graph=retain_graph, create_graph=create_graph + )[0] assert ret is not None return ret - - - - - - diff --git a/pyaimnet2/models/__init__.py b/pyaimnet2/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/models/aimnet2_b973c_ens.jpt b/pyaimnet2/models/aimnet2_b973c_ens.jpt similarity index 100% rename from models/aimnet2_b973c_ens.jpt rename to pyaimnet2/models/aimnet2_b973c_ens.jpt diff --git a/models/aimnet2_wb97m-d3_ens.jpt b/pyaimnet2/models/aimnet2_wb97m-d3_ens.jpt similarity index 100% rename from models/aimnet2_wb97m-d3_ens.jpt rename to pyaimnet2/models/aimnet2_wb97m-d3_ens.jpt diff --git a/models/ensemble.py b/pyaimnet2/models/ensemble.py similarity index 75% rename from models/ensemble.py rename to pyaimnet2/models/ensemble.py index 4067653..82f96e7 100644 --- a/models/ensemble.py +++ b/pyaimnet2/models/ensemble.py @@ -4,9 +4,15 @@ class EnsembledModel(nn.Module): - """ Create ensemble of AIMNet2 models. - """ - def __init__(self, models: List, x=['coord', 'numbers', 'charge'], out=['energy', 'forces', 'charges'], detach=True): + """Create ensemble of AIMNet2 models.""" + + def __init__( + self, + models: List, + x=["coord", "numbers", "charge"], + out=["energy", "forces", "charges"], + detach=True, + ): super().__init__() self.models = nn.ModuleList(models) self.x = x @@ -14,7 +20,7 @@ def __init__(self, models: List, x=['coord', 'numbers', 'charge'], out=['energy' self.detach = detach def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]: - res : List[Dict[str, Tensor]] = [] + res: List[Dict[str, Tensor]] = [] for model in self.models: _in = dict() for k in data: @@ -35,6 +41,6 @@ def forward(self, data: Dict[str, Tensor]) -> Dict[str, Tensor]: v.append(x[k]) vv = torch.stack(v, dim=0) data[k] = vv.mean(dim=0) - data[k + '_std'] = vv.std(dim=0) + data[k + "_std"] = vv.std(dim=0) return data diff --git a/pyaimnet2/tests/conftest.py b/pyaimnet2/tests/conftest.py new file mode 100644 index 0000000..670c92b --- /dev/null +++ b/pyaimnet2/tests/conftest.py @@ -0,0 +1,21 @@ +import pytest +import torch + + +@pytest.fixture(scope="package") +def water() -> dict: + """Return a dict of a water molecule which can be used with the aimnet2 model to check the energy""" + return { + "coord": torch.tensor( + [ + [ + [0.06112021, 0.38865671, 0.05890042], + [0.72378355, -0.31162935, -0.03823509], + [-0.78490371, -0.07702733, -0.02066534], + ], + ], + dtype=torch.float64, + ), + "numbers": torch.tensor([[8, 1, 1]], dtype=torch.long), + "charge": torch.tensor([0], dtype=torch.float64), + } diff --git a/pyaimnet2/tests/test_models.py b/pyaimnet2/tests/test_models.py new file mode 100644 index 0000000..abd2502 --- /dev/null +++ b/pyaimnet2/tests/test_models.py @@ -0,0 +1,34 @@ +import pytest +from pyaimnet2 import load_model + + +@pytest.mark.parametrize( + "model", + [pytest.param("b973c", id="b973c"), pytest.param("wb97m-d3", id="wb97m-d3")], +) +def test_loading_models(model): + """test loading different models using the package.""" + + _ = load_model(model_name=model) + + +def test_missing_models(): + """Make sure an error is raised if we try and load a model not supported""" + + with pytest.raises(FileNotFoundError): + _ = load_model(model_name="b3lyp") + + +@pytest.mark.parametrize( + "model, expected_energy", + [ + pytest.param("wb97m-d3", -2081.0415, id="wb97m-d3"), + pytest.param("b973c", -2078.9055, id="b973c"), + ], +) +def test_model_energies(model, expected_energy, water): + """Test the energies calculated for a water molecule match the expected values.""" + + aimnet2 = load_model(model_name=model) + return_data = aimnet2(water) + assert return_data["energy"] == pytest.approx(expected_energy) diff --git a/pyaimnet2/utils.py b/pyaimnet2/utils.py new file mode 100644 index 0000000..254858c --- /dev/null +++ b/pyaimnet2/utils.py @@ -0,0 +1,28 @@ +from typing import Literal + +available_models = ["b973c", "wb97m-d3"] +MODELS = Literal["b973c", "wb97m-d3"] + + +def load_model(model_name: MODELS): + """ + Load the specified AIMNET2 model. + + Args: + model_name: The name of the ensemble model which should be loaded (`b973c`, `wb97m-d3`). + + Returns: + A aimnet2 model ready for evaluation. + """ + from importlib.resources import files + import pathlib + import torch.jit + + if model_name.lower() not in available_models: + raise FileNotFoundError( + f"The model {model_name} is not available chose from {available_models}" + ) + + package_path: pathlib.Path = files("pyaimnet2") + model_path = package_path.joinpath("models", f"aimnet2_{model_name}_ens.jpt") + return torch.jit.load(model_path.as_posix()) diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..643424e --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,59 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel", "versioneer"] +build-backend = "setuptools.build_meta" + +[project] +name = "pyaimnet2" +description = "AIMNet2: a general-purpose neural network potential for organic and element-organic molecules." +authors = [ {name = "Dylan Anstine ,Roman Zubatyuk ,Olexandr Isayev"} ] +license = { text = "MIT" } +dynamic = ["version"] +readme = "README.md" +requires-python = ">=3.10" +classifiers = ["Programming Language :: Python :: 3"] + +[tool.setuptools] +zip-safe = false +include-package-data = true + +[tool.setuptools.dynamic] +version = {attr = "pyaimnet2.__version__"} + +[tool.setuptools.packages.find] +namespaces = true +where = ["."] + +[tool.setuptools.package-data] + + +[tool.versioneer] +VCS = "git" +style = "pep440" +versionfile_source = "pyaimnet2/_version.py" +versionfile_build = "pyaimnet2/_version.py" +tag_prefix = "" +parentdir_prefix = "pyaimnet2-" + +[tool.black] +line-length = 88 + +[tool.isort] +profile = "black" + +[tool.flake8] +max-line-length = 88 +ignore = ["E203", "E266", "E501", "W503"] +select = ["B","C","E","F","W","T4","B9"] + +[tool.coverage.run] +omit = ["**/tests/*", "**/_version.py"] + +[tool.coverage.report] +exclude_lines = [ + "@overload", + "pragma: no cover", + "raise NotImplementedError", + "if __name__ = .__main__.:", + "if TYPE_CHECKING:", + "if typing.TYPE_CHECKING:", +] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 9ecedeb..0000000 --- a/requirements.txt +++ /dev/null @@ -1,4 +0,0 @@ -torch==2.0.0 -ase -numpy -pysisyphus diff --git a/aimnet2 b/scripts/aimnet2 similarity index 100% rename from aimnet2 rename to scripts/aimnet2 diff --git a/calculators/pysis_mod b/scripts/pysis_mod similarity index 100% rename from calculators/pysis_mod rename to scripts/pysis_mod