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

aimnet/train/pt2jpt.py

@@ -1,81 +0,0 @@
-import os
-from typing import List, Optional
-
-import click
-import torch
-from torch import nn
-
-from aimnet.config import build_module, load_yaml
-
-
-def set_eval(model: nn.Module) -> torch.nn.Module:
-    for p in model.parameters():
-        p.requires_grad_(False)
-    return model.eval()
-
-
-def add_cutoff(
-    model: nn.Module, cutoff: Optional[float] = None, cutoff_lr: Optional[float] = float("inf")
-) -> nn.Module:
-    if cutoff is None:
-        cutoff = max(v.item() for k, v in model.state_dict().items() if k.endswith("aev.rc_s"))
-    model.cutoff = cutoff  # type: ignore[assignment]
-    if cutoff_lr is not None:
-        model.cutoff_lr = cutoff_lr  # type: ignore[assignment]
-    return model
-
-
-def add_sae_to_shifts(model: nn.Module, sae_file: str) -> nn.Module:
-    sae = load_yaml(sae_file)
-    if not isinstance(sae, dict):
-        raise TypeError("SAE file must contain a dictionary.")
-    model.outputs.atomic_shift.double()
-    for k, v in sae.items():
-        model.outputs.atomic_shift.shifts.weight[k] += v
-    return model
-
-
-def mask_not_implemented_species(model: nn.Module, species: List[int]) -> nn.Module:
-    weight = model.afv.weight
-    for i in range(1, weight.shape[0]):
-        if i not in species:
-            weight[i, :] = torch.nan
-    return model
-
-
-_default_aimnet2_config = os.path.join(os.path.dirname(__file__), "..", "models", "aimnet2.yaml")
-
-
-@click.command(short_help="Compile PyTorch model to TorchScript.")
-@click.argument("pt", type=str)  # , help='Path to the input PyTorch weights file.')
-@click.argument("jpt", type=str)  # , help='Path to the output TorchScript file.')
-@click.option("--model", type=str, default=_default_aimnet2_config, help="Path to model definition YAML file")
-@click.option("--sae", type=str, default=None, help="Path to the energy shift YAML file.")
-@click.option("--species", type=str, default=None, help="Comma-separated list of parametrized atomic numbers.")
-@click.option("--no-lr", is_flag=True, help="Do not add LR cutoff for model")
-def jitcompile(model: str, pt: str, jpt: str, sae=None, species=None, no_lr=False):  # type: ignore
-    """Build model from YAML config, load weight from PT file and write JIT-compiled JPT file.
-    Plus some modifications to work with aimnet2calc.
-    """
-    model: nn.Module = build_module(model)  # type: ignore[annotation-unchecked]
-    model = set_eval(model)
-    cutoff_lr = None if no_lr else float("inf")
-    model = add_cutoff(model, cutoff_lr=cutoff_lr)
-    sd = torch.load(pt, map_location="cpu", weights_only=True)
-    print(model.load_state_dict(sd, strict=False))
-    if sae:
-        model = add_sae_to_shifts(model, sae)
-    numbers = None
-    if species:
-        numbers = list(map(int, species.split(",")))
-    elif sae:
-        numbers = list(load_yaml(sae).keys())  # type: ignore[union-attr]
-    if numbers:
-        model = mask_not_implemented_species(model, numbers)  # type: ignore[call-arg]
-        model.register_buffer("impemented_species", torch.tensor(numbers, dtype=torch.int64))
-    model_jit = torch.jit.script(model)
-    model_jit.save(jpt)
-
-
-if __name__ == "__main__":
-    jitcompile()

aimnet-0.0.1.dist-info/METADATA

@@ -1,78 +0,0 @@
-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.dist-info/RECORD

@@ -1,41 +0,0 @@
-aimnet/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aimnet/base.py,sha256=v-8ukDanaxnhYP0IbugKoixssiJTQMR6KyGtnriH1RI,1663
-aimnet/calculators/__init__.py,sha256=V-x9HIkk_wDkBI0W88A0St-9dmTYPZXErGT_hWLkgng,383
-aimnet/calculators/aimnet2ase.py,sha256=omyzJKzTpLvIltvYKOvikQ6VuSx6IViWQNZlceMrV_A,3536
-aimnet/calculators/aimnet2pysis.py,sha256=TDoDl-TExCQkg8EXfsTKjOu-z6Kd25xXqSH98gF8iGM,2876
-aimnet/calculators/calculator.py,sha256=9QbhAfw7EBi5CRhP0rT-2XVx243Ub3cq-XrqJI2LLb0,13092
-aimnet/calculators/model_registry.py,sha256=TaQdTUiKqKWt7fSOxODd69QgnmCuNVwHL3qBOcyMsZI,1921
-aimnet/calculators/model_registry.yaml,sha256=Yejl2ls8DWvwRaBIwpGfc3WsHB9jZdIH_3B-i5klmks,1410
-aimnet/calculators/nb_kernel_cpu.py,sha256=a6ClFEzfc7-Fr4m_TNGOprjdq0GiFZ8wheqmKeDjqjM,7309
-aimnet/calculators/nb_kernel_cuda.py,sha256=0FHpGJ6-MwQgoX-GiBLGZqdhN4P2uqX4CPxQIi_-S2c,6552
-aimnet/calculators/nbmat.py,sha256=_3LKqoHEVsM1dXBMICQYyciPoqIBUgVq6KGM7XgIMh4,7785
-aimnet/cli.py,sha256=2A51wv5N2i9HdmiCcam_JDY03WKeOTTn8dTxTtp1QXM,415
-aimnet/config.py,sha256=Ld5wyApDhI8OovodC_ibpcDL9rU-VJ8gkQZ3BaiDf0c,5920
-aimnet/constants.py,sha256=T6eb_CG5dkxuUg22Xf8WiF0iGB-JLoLNzHq9R8SRzrQ,8760
-aimnet/data/__init__.py,sha256=rD84W2qazjXUbdcYBuHkMmYp44fGRFMS3qnv1OKPsQs,87
-aimnet/data/sgdataset.py,sha256=N3Nwwn1GicoohcRh-DJWJFdibmV2o6YZjKaTlWyZZx4,18926
-aimnet/dftd3_data.pt,sha256=g9HQ_2nekMs97UnIDIFReGEG8q6n6zQlq-Yl5LGQnOI,2710647
-aimnet/models/__init__.py,sha256=gb6kA4dHknYOzllaz02-Bx5kBkWaSjj9wLtY0Fx9zdo,87
-aimnet/models/aimnet2.py,sha256=O0CcP21sJcSeZCE-H0klKfFPuqDOmxWdrX_ZbX1aFm8,7006
-aimnet/models/aimnet2.yaml,sha256=VmI0ub7UbiJm833zwCtCKsrBIaXZKjJsCXxXxutwdLI,1062
-aimnet/models/aimnet2_dftd3_wb97m.yaml,sha256=Gk869_2RQIXzlHwRNbfsai4u9wPX3oAyINK1I71eBFI,1219
-aimnet/models/base.py,sha256=uLcOaTiXEZPpLgBVEnT5QFcIZrA8KgVNLe_guR5o99o,1815
-aimnet/modules/__init__.py,sha256=c4tcXbNfRRvFgIk9RBk6X_5iKmCy6za7ooeuJfqlAHc,205
-aimnet/modules/aev.py,sha256=QgnFsz_k-qd_d0sufBc0B_r2iCCJezl2Pn9wbvDIeIU,7801
-aimnet/modules/core.py,sha256=kDhU50plXKXNXnr5BgKJREwHa2F6aawL6wI2O74cMb0,8236
-aimnet/modules/lr.py,sha256=x6AAdyF_Qgkp7OKcS9dk8FDU5ENoVkKYQ2aQgUajz6Y,8835
-aimnet/nbops.py,sha256=iR9MR3MSts5Ow6j-_PtwylqHSYEQyJDwDEvF7FHNSJc,5450
-aimnet/ops.py,sha256=p6_jwC_xjEw3FaNvwmNkkIOIttFRsd9WF03ovUkX3cU,7888
-aimnet/train/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-aimnet/train/calc_sae.py,sha256=h-y8jQVU6O7Rmw8xJWX3DHOpFv1gtjFwtZzZyUOhz9g,1459
-aimnet/train/default_train.yaml,sha256=FAG3f3pFC2dkFUuVO_XE_UB6rFr-R3GkACddn7FH_58,4783
-aimnet/train/loss.py,sha256=BqBaMBCPBkFcm1QnpczFJY6ZU_5KfAQEBo-LCYdJ0m0,3175
-aimnet/train/metrics.py,sha256=HEfIoFhOc6eD8yrgxqXUGI3CfWnXfm6G6AUpncGQljs,6867
-aimnet/train/pt2jpt.py,sha256=XMc95PrXvlHBUqtAVNJ0JzJP0ZJK-WrJ8rbSg4yxHCU,3129
-aimnet/train/train.py,sha256=S34BWOhfhe8fg9qGkp-xw6O4-sgHM6R2oQ0-e8dk6hg,5711
-aimnet/train/utils.py,sha256=2pYZbY1L9KdmuHqmfdQDId-TWdl-WMeNoeatcGZWPPE,14730
-aimnet-0.0.1.dist-info/LICENSE,sha256=73sk-zg2yVRrOZQDeVbPlVB7ScZc1iK0kyCBMwNwQgA,1072
-aimnet-0.0.1.dist-info/METADATA,sha256=RAVsz_d3HL2QFGV9H9M2J57riF1W9VMmhsjUWvTOSaw,3717
-aimnet-0.0.1.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-aimnet-0.0.1.dist-info/entry_points.txt,sha256=3SPZbjuQFS1V1xryBwVoo8ix0DNGYzHeoDbvw0z2Mdk,164
-aimnet-0.0.1.dist-info/RECORD,,

aimnet-0.0.1.dist-info/entry_points.txt

@@ -1,5 +0,0 @@
-[console_scripts]
-aimnet=aimnet.cli:cli
-aimnet-clear-models-cache=aimnet.calculators.model_registry:clear_assets
-aimnet2pysis=aimnet.calculators.aimnet2pysis:main
-