torchmd-net 2.5.0__tar.gz

torchmd_net-2.5.0/.github/workflows/docs.yml

@@ -0,0 +1,39 @@
+name: Build Documentation
+
+on:
+  workflow_dispatch:
+  workflow_call:
+
+jobs:
+  build-docs:
+    runs-on: ubuntu-latest
+
+    defaults: # Needed for conda
+      run:
+        shell: bash -l {0}
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          python-version: "3.10"
+          channels: conda-forge
+          conda-remove-defaults: "true"
+
+      - name: Install compiler
+        run: conda install gxx --channel conda-forge --override-channels
+
+      - name: Install docs dependencies
+        run: |
+          pip install build
+          python -m build
+          pip install dist/*.whl
+          pip install -r docs/requirements.txt
+        env:
+          ACCELERATOR: "cpu"
+
+      - name: Build Sphinx Documentation
+        run: |
+          cd docs
+          make html

torchmd_net-2.5.0/.github/workflows/publish.yml

@@ -0,0 +1,74 @@
+name: Build wheels and publish
+
+on:
+  workflow_dispatch:
+  workflow_call:
+
+jobs:
+  build:
+    name: Create source distribution
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v5
+
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.13"
+
+      - name: Install cibuildwheel
+        run: pip install build
+
+      - name: Build pypi package
+        run: python -m build --sdist
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: source_dist
+          path: dist/*.tar.gz
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    needs:
+      - build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/torchmd-net
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          path: dist/
+          merge-multiple: true
+
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          password: ${{ secrets.TMDNET_PYPI_API_TOKEN }}
+          skip_existing: true
+
+  github-release:
+    name: >-
+      Create GitHub Release
+    needs:
+      - build
+    runs-on: ubuntu-latest
+
+    permissions:
+      contents: write # IMPORTANT: mandatory for making GitHub Releases
+      id-token: write # IMPORTANT: mandatory for sigstore
+
+    steps:
+      - name: Create GitHub Release
+        env:
+          GITHUB_TOKEN: ${{ github.token }}
+        run: >-
+          gh release create
+          "$GITHUB_REF_NAME"
+          --repo "$GITHUB_REPOSITORY"
+          --notes ""

torchmd_net-2.5.0/.github/workflows/test.yml

@@ -0,0 +1,44 @@
+name: Test
+
+on:
+  workflow_dispatch:
+  workflow_call:
+  schedule:
+    - cron: "0 0 * * 0" # Run every Sunday at 0:00
+
+jobs:
+  test:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      fail-fast: false
+      matrix:
+        os:
+          ["ubuntu-latest", "ubuntu-22.04-arm", "macos-latest", "windows-2022"]
+        python-version: ["3.13"]
+
+    steps:
+      - name: Check out
+        uses: actions/checkout@v5
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install the project
+        run: uv sync --all-extras --dev
+
+      - name: Lint with flake8
+        run: |
+          # stop the build if there are Python syntax errors or undefined names
+          uv run flake8 ./torchmdnet --count --select=E9,F63,F7,F82 --show-source --statistics
+          # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+          uv run flake8 ./torchmdnet --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+
+      - name: Run tests
+        # For example, using `pytest`
+        run: uv run pytest tests
+
+      - name: Test torchmd-train utility
+        run: uv run torchmd-train --help

torchmd_net-2.5.0/.github/workflows/workflow.yml

@@ -0,0 +1,25 @@
+name: Whole
+
+on: [push]
+
+jobs:
+  test:
+    uses: ./.github/workflows/test.yml
+    secrets: inherit
+
+  publish:
+    if: startsWith(github.event.ref, 'refs/tags/v')
+    needs: test
+    uses: ./.github/workflows/publish.yml
+    secrets: inherit
+
+  docs:
+    if: startsWith(github.event.ref, 'refs/tags/')
+    needs: publish
+    uses: ./.github/workflows/docs.yml
+    secrets: inherit
+
+  docs_update:
+    if: startsWith(github.event.ref, 'refs/tags/d')
+    uses: ./.github/workflows/docs.yml
+    secrets: inherit

torchmd_net-2.5.0/.gitignore

@@ -0,0 +1,143 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# VS Code folder
+.vscode/
+
+# temp directories
+logs/
+
+# Docs
+docs/build/
+docs/source/generated
+
+# Extra
+
+*~

torchmd_net-2.5.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021-2023 Universitat Pompeu Fabra,  https://www.compscience.org
+
+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.

torchmd_net-2.5.0/MANIFEST.in

@@ -0,0 +1,9 @@
+prune benchmarks
+prune cibuildwheel_support
+prune docs
+prune examples
+prune scripts
+prune tests
+exclude .readthedocs.yaml
+exclude environment.yml
+exclude Makefile

torchmd_net-2.5.0/PKG-INFO

@@ -0,0 +1,256 @@
+Metadata-Version: 2.4
+Name: torchmd-net
+Version: 2.5.0
+Summary: TorchMD-NET provides state-of-the-art neural networks potentials for biomolecular systems
+Author-email: Acellera <info@acellera.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/torchmd/torchmd-net
+Project-URL: Bug Tracker, https://github.com/torchmd/torchmd-net/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: POSIX :: Linux
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: h5py
+Requires-Dist: torch
+Requires-Dist: torch_geometric
+Requires-Dist: lightning
+Requires-Dist: tqdm
+Requires-Dist: numpy
+Requires-Dist: triton; sys_platform == "linux" and platform_machine != "aarch64"
+Requires-Dist: triton-windows; sys_platform == "win32"
+Requires-Dist: ase
+Dynamic: license-file
+
+[![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black)
+[![CI](https://github.com/torchmd/torchmd-net/actions/workflows/workflow.yml/badge.svg)](https://github.com/torchmd/torchmd-net/actions/workflows/workflow.yml)
+[![Documentation Status](https://readthedocs.org/projects/torchmd-net/badge/?version=latest)](https://torchmd-net.readthedocs.io/en/latest/?badge=latest)  
+
+# TorchMD-NET
+
+TorchMD-NET provides state-of-the-art neural networks potentials (NNPs) and a mechanism to train them. It offers efficient and fast implementations if several NNPs and it is integrated in GPU-accelerated molecular dynamics code like [ACEMD](https://www.acellera.com/products/molecular-dynamics-software-gpu-acemd/), [OpenMM](https://www.openmm.org) and [TorchMD](https://github.com/torchmd/torchmd). TorchMD-NET exposes its NNPs as [PyTorch](https://pytorch.org) modules.
+
+
+## Documentation
+
+Documentation is available at https://torchmd-net.readthedocs.io  
+
+## Available architectures
+
+- [Equivariant Transformer (ET)](https://arxiv.org/abs/2202.02541)
+- [Transformer (T)](https://arxiv.org/abs/2202.02541)
+- [Graph Neural Network (GN)](https://arxiv.org/abs/2212.07492)
+- [TensorNet](https://arxiv.org/abs/2306.06482)
+
+
+
+## Installation  
+TorchMD-Net is available as a pip package as well as in [conda-forge](https://conda-forge.org/)
+
+As TorchMD-Net depends on PyTorch we need to add additional index URLs to the installation command as per [pytorch](https://pytorch.org/get-started/locally/)
+
+```sh
+# The following will install TorchMD-Net with PyTorch CUDA 11.8 version
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cu118
+# The following will install TorchMD-Net with PyTorch CUDA 12.4 version
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cu124
+# The following will install TorchMD-Net with PyTorch CPU only version (not recommended)
+pip install torchmd-net --extra-index-url https://download.pytorch.org/whl/cpu
+```
+
+Alternatively it can be installed with conda or mamba with one of the following commands.
+We recommend using [Miniforge](https://github.com/conda-forge/miniforge/) instead of anaconda.
+
+```shell
+mamba install torchmd-net cuda-version=11.8
+mamba install torchmd-net cuda-version=12.4
+```
+
+### Install from source  
+
+TorchMD-Net is installed using pip with `pip install -e .` to create an editable install.  
+
+## Usage
+Specifying training arguments can either be done via a configuration yaml file or through command line arguments directly. Several examples of architectural and training specifications for some models and datasets can be found in [examples/](https://github.com/torchmd/torchmd-net/tree/main/examples). Note that if a parameter is present both in the yaml file and the command line, the command line version takes precedence.
+GPUs can be selected by setting the `CUDA_VISIBLE_DEVICES` environment variable. Otherwise, the argument `--ngpus` can be used to select the number of GPUs to train on (-1, the default, uses all available GPUs or the ones specified in `CUDA_VISIBLE_DEVICES`). Keep in mind that the [GPU ID reported by nvidia-smi might not be the same as the one `CUDA_VISIBLE_DEVICES` uses](https://stackoverflow.com/questions/26123252/inconsistency-of-ids-between-nvidia-smi-l-and-cudevicegetname).  
+For example, to train the Equivariant Transformer on the QM9 dataset with the architectural and training hyperparameters described in the paper, one can run:
+```shell
+mkdir output
+CUDA_VISIBLE_DEVICES=0 torchmd-train --conf torchmd-net/examples/ET-QM9.yaml --log-dir output/
+``` 
+Run `torchmd-train --help` to see all available options and their descriptions.
+
+## Inference with Pretrained models
+
+### AceFF
+Trained [AceFF models](https://huggingface.co/collections/Acellera/aceff-machine-learning-potentials) can be loaded and used for inference.
+please see [here](https://github.com/torchmd/torchmd-net/tree/main/examples/aceff_examples) 
+
+
+To load your own trained models see [here](https://github.com/torchmd/torchmd-net/tree/main/examples#loading-checkpoints) for instructions on how to load pretrained models.
+
+## Creating a new dataset
+If you want to train on custom data, first have a look at `torchmdnet.datasets.Custom`, which provides functionalities for 
+loading a NumPy dataset consisting of atom types and coordinates, as well as energies, forces or both as the labels.
+Alternatively, you can implement a custom class according to the torch-geometric way of implementing a dataset. That is, 
+derive the `Dataset` or `InMemoryDataset` class and implement the necessary functions (more info [here](https://pytorch-geometric.readthedocs.io/en/latest/notes/create_dataset.html#creating-your-own-datasets)). The dataset must return torch-geometric `Data` 
+objects, containing at least the keys `z` (atom types) and `pos` (atomic coordinates), as well as `y` (label), `neg_dy` (negative derivative of the label w.r.t atom coordinates) or both.
+
+### Custom prior models
+In addition to implementing a custom dataset class, it is also possible to add a custom prior model to the model. This can be
+done by implementing a new prior model class in `torchmdnet.priors` and adding the argument `--prior-model <PriorModelName>`.
+As an example, have a look at `torchmdnet.priors.Atomref`.
+
+## Multi-Node Training
+
+In order to train models on multiple nodes some environment variables have to be set, which provide all necessary information to PyTorch Lightning. In the following we provide an example bash script to start training on two machines with two GPUs each. The script has to be started once on each node. Once `torchmd-train` is started on all nodes, a network connection between the nodes will be established using NCCL.
+
+In addition to the environment variables the argument `--num-nodes` has to be specified with the number of nodes involved during training.
+
+```shell
+export NODE_RANK=0
+export MASTER_ADDR=hostname1
+export MASTER_PORT=12910
+
+mkdir -p output
+CUDA_VISIBLE_DEVICES=0,1 torchmd-train --conf torchmd-net/examples/ET-QM9.yaml.yaml --num-nodes 2 --log-dir output/
+```
+
+- `NODE_RANK` : Integer indicating the node index. Must be `0` for the main node and incremented by one for each additional node.
+- `MASTER_ADDR` : Hostname or IP address of the main node. The same for all involved nodes.
+- `MASTER_PORT` : A free network port for communication between nodes. PyTorch Lightning suggests port `12910` as a default.
+
+
+### Known Limitations
+- Due to the way PyTorch Lightning calculates the number of required DDP processes, all nodes must use the same number of GPUs. Otherwise training will not start or crash.
+- We observe a 50x decrease in performance when mixing nodes with different GPU architectures (tested with RTX 2080 Ti and RTX 3090).
+- Some CUDA systems might hang during a multi-GPU parallel training. Try `export NCCL_P2P_DISABLE=1`, which disables direct peer to peer GPU communication.
+
+
+## Cite
+If you use TorchMD-NET in your research, please cite the following papers:
+
+#### Main reference
+
+```
+@misc{pelaez2024torchmdnet,
+title={TorchMD-Net 2.0: Fast Neural Network Potentials for Molecular Simulations}, 
+author={Raul P. Pelaez and Guillem Simeon and Raimondas Galvelis and Antonio Mirarchi and Peter Eastman and Stefan Doerr and Philipp Thölke and Thomas E. Markland and Gianni De Fabritiis},
+year={2024},
+eprint={2402.17660},
+archivePrefix={arXiv},
+primaryClass={cs.LG}
+}
+```
+
+#### TensorNet
+
+```
+@inproceedings{simeon2023tensornet,
+title={TensorNet: Cartesian Tensor Representations for Efficient Learning of Molecular Potentials},
+author={Guillem Simeon and Gianni De Fabritiis},
+booktitle={Thirty-seventh Conference on Neural Information Processing Systems},
+year={2023},
+url={https://openreview.net/forum?id=BEHlPdBZ2e}
+}
+```
+
+#### Equivariant Transformer
+```
+@inproceedings{
+tholke2021equivariant,
+title={Equivariant Transformers for Neural Network based Molecular Potentials},
+author={Philipp Th{\"o}lke and Gianni De Fabritiis},
+booktitle={International Conference on Learning Representations},
+year={2022},
+url={https://openreview.net/forum?id=zNHzqZ9wrRB}
+}
+```
+
+#### Graph Network 
+
+```
+@article{Majewski2023,
+  title = {Machine learning coarse-grained potentials of protein thermodynamics},
+  volume = {14},
+  ISSN = {2041-1723},
+  url = {http://dx.doi.org/10.1038/s41467-023-41343-1},
+  DOI = {10.1038/s41467-023-41343-1},
+  number = {1},
+  journal = {Nature Communications},
+  publisher = {Springer Science and Business Media LLC},
+  author = {Majewski,  Maciej and Pérez,  Adrià and Th\"{o}lke,  Philipp and Doerr,  Stefan and Charron,  Nicholas E. and Giorgino,  Toni and Husic,  Brooke E. and Clementi,  Cecilia and Noé,  Frank and De Fabritiis,  Gianni},
+  year = {2023},
+  month = sep 
+}
+```
+
+## Developer guide
+
+### Implementing a new architecture
+
+To implement a new architecture, you need to follow these steps:  
+ **1.** Create a new class in `torchmdnet.models` that inherits from `torch.nn.Model`. Follow TorchMD_ET as a template. This is a minimum implementation of a model:  
+```python 
+class MyModule(nn.Module):
+  def __init__(self, parameter1, parameter2):
+	super(MyModule, self).__init__()
+	# Define your model here
+	self.layer1 = nn.Linear(10, 10)
+	...
+	# Initialize your model parameters here
+	self.reset_parameters()
+
+    def reset_parameters(self):
+      # Initialize your model parameters here
+	  nn.init.xavier_uniform_(self.layer1.weight)
+	...
+	
+  def forward(self,
+        z: Tensor, # Atomic numbers, shape (n_atoms, 1)
+        pos: Tensor, # Atomic positions, shape (n_atoms, 3)
+        batch: Tensor, # Batch vector, shape (n_atoms, 1). All atoms in the same molecule have the same value and are contiguous.
+        q: Optional[Tensor] = None, # Atomic charges, shape (n_atoms, 1)
+        s: Optional[Tensor] = None, # Atomic spins, shape (n_atoms, 1)
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:
+	# Define your forward pass here
+	scalar_features = ...
+	vector_features = ...
+	# Return the scalar and vector features, as well as the atomic numbers, positions and batch vector
+	return scalar_features, vector_features, z, pos, batch
+```
+ **2.** Add the model to the `__all__` list in `torchmdnet.models.__init__.py`. This will make the tests pick your model up.  
+ **3.** Tell models.model.create_model how to initialize your module by adding a new entry, for instance:  
+ ```python
+     elif args["model"] == "mymodule":
+        from torchmdnet.models.torchmd_mymodule import MyModule
+        is_equivariant = False # Set to True if your model is equivariant
+        representation_model = MyModule(
+            parameter1=args["parameter1"],
+            parameter2=args["parameter2"],
+            **shared_args, # Arguments typically shared by all models
+        )
+ ```
+ 
+ **4.** Add any new parameters required to initialize your module to scripts.train.get_args. For instance:  
+ ```python
+   parser.add_argument('--parameter1', type=int, default=32, help='Parameter1 required by MyModule')
+   ...
+ ```
+ **5.** Add an example configuration file to `torchmd-net/examples` that uses your model.  
+ **6.** Make tests use your configuration file by adding a case to tests.utils.load_example_args. For instance:  
+ ```python
+ if model_name == "mymodule":
+        config_file = join(dirname(dirname(__file__)), "examples", "MyModule-QM9.yaml")
+ ```
+
+At this point, if your module is missing some feature the tests will let you know, and you can add it. If you add a new feature to the package, please add a test for it.  
+
+### Code style
+
+We use [black](https://https://black.readthedocs.io/en/stable/). Please run `black` on your modified files before committing.  
+
+### Testing
+
+To run the tests, install the package and run `pytest` in the root directory of the repository. Tests are a good source of knowledge on how to use the different components of the package.  
+