Improve documentation (#224)

* Improve TorchMD_NET documentation Add some notes about charge/spin * Add charge and spin to the example yamls, to highlight their existence * Add some docstrings * Add first ReadTheDocs version * Remove unused lines in conf * Move some stuff to Usage * Update docs * Move options to another file so they can be autodocumented * Update docs * Update docs * Hide utility classes in transformer * Update some docstrings * Update conf.py * Add style.css * Document ACE Dataset * Document ANI * Revert _MultiHeadAttention to MultiHeadAttention * Update ace documentation * Update ani and atomref docs * Update ace docs * Add extension mock import * Update comment * Add docstring to extension * Add docstring for extensions * Revert "Move options to another file so they can be autodocumented" This reverts commit c61c922. * Store model names in model.__all_models__ * Set __all__ for some files * Fix External docstring * Make docs respect __all__ * Fix docstring * Update docs * Update doc build * Move readthedocs to root * Update readthedocs * Add environemnt.yml * Install local master version for docs * Update requirements * Remove env * Fix typo * Hide classes with :meta private: * Add intersphinx * Add version with git * Blacken * Change title * Small changes * Small changes * Default tag to master * Add priors * Add workflow to test documentation build * Update ci * Update ci * Update ci * Update ci * Update ci * Add documentation badge to README * Update README
torchmd · Nov 16, 2023 · be6a214 · be6a214
1 parent c3b46f6
commit be6a214
Show file tree

Hide file tree

Showing 49 changed files with 1,442 additions and 176 deletions.
diff --git a/.github/workflows/docs_build.yaml b/.github/workflows/docs_build.yaml
@@ -0,0 +1,36 @@
+name: Build Documentation
+
+on:
+  push:
+    branches:
+      - "main"
+  pull_request:
+    branches:
+      - "main"
+
+
+jobs:
+  build-docs:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+
+    - name: Set up Env
+      uses: mamba-org/setup-micromamba@v1
+      with:
+        environment-file: environment.yml
+        init-shell: bash
+        generate-run-shell: true
+
+    - name: Install docs dependencies
+      run: |
+        pip install -vv .
+        pip install -r docs/requirements.txt
+      shell: bash -el {0}
+
+    - name: Build Sphinx Documentation
+      run: |
+        cd docs
+        make html
+      shell: bash -el {0}
diff --git a/.gitignore b/.gitignore
@@ -133,3 +133,11 @@ dmypy.json
 
 # temp directories
 logs/
+
+# Docs
+docs/build/
+docs/source/generated
+
+# Extra
+
+*~
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
@@ -0,0 +1,17 @@
+version: "2"
+
+build:
+  os: "ubuntu-22.04"
+  tools:
+     python: "mambaforge-22.9"
+  jobs:
+     post_create_environment:
+       - pip install -r docs/requirements.txt
+       - pip install .
+
+conda:
+  environment: environment.yml
+
+
+sphinx:
+  configuration: docs/source/conf.py
diff --git a/README.md b/README.md
@@ -1,10 +1,16 @@
 [![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/CI.yml/badge.svg)](https://github.com/torchmd/torchmd-net/actions/workflows/CI.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)
@@ -13,6 +19,7 @@ TorchMD-NET provides state-of-the-art neural networks potentials (NNPs) and a me
 - [TensorNet](https://arxiv.org/abs/2306.06482)
 
 
+
 ## Installation  
 TorchMD-Net is available in [conda-forge](https://conda-forge.org/) and can be installed with:  
 ```shell

diff --git a/docs/Makefile b/docs/Makefile
@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
diff --git a/docs/_static/style.css b/docs/_static/style.css
@@ -0,0 +1,59 @@
+@import 'theme.css';
+
+.rst-content dl:not(.docutils) dt:first-child {
+  margin-top: 0;
+}
+
+.rst-content dl:not(.docutils) dl dt {
+    margin-bottom: 4px;
+    border: none;
+    border-left: solid 3px #ccc;
+    background: #f0f0f0;
+    color: #555;
+}
+
+.rst-content dl table,
+.rst-content dl ul,
+.rst-content dl ol,
+.rst-content dl p {
+    margin-bottom: 8px !important;
+}
+
+.rst-content dl:not(.docutils) dt {
+    display: table;
+    margin: 6px 0;
+    font-size: 90%;
+    line-height: normal;
+    background: #e7f2fa;
+    color: #2980B9;
+    border-top: solid 3px #6ab0de;
+    padding: 6px;
+    position: relative;
+}
+
+html.writer-html5 .rst-content dl.field-list {
+    display: initial;
+}
+
+html.writer-html5 .rst-content dl.field-list>dd,
+html.writer-html5 .rst-content dl.field-list>dt {
+    margin-bottom: 4px;
+    padding-left: 6px;
+}
+
+p {
+    line-height: 20px;
+    font-size: 14px;
+}
+
+html.writer-html5 .rst-content dl.field-list>dt:after {
+    content: initial;
+}
+
+dt.field-even {
+    text-transform: uppercase;
+}
+
+dt.field-odd {
+    text-transform: uppercase;
+}
diff --git a/docs/make.bat b/docs/make.bat
@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.http://sphinx-doc.org/
+	exit /b 1
+)
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd
diff --git a/docs/requirements.txt b/docs/requirements.txt
@@ -0,0 +1,5 @@
+sphinx==7.2.6
+sphinx-rtd-theme==1.3.0
+sphinxcontrib-autoprogram==0.1.8
+sphinxcontrib-napoleon==0.7
+gitpython
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -0,0 +1,10 @@
+API Reference
+=============
+
+.. autosummary::
+   :toctree: generated
+   :recursive:
+
+   torchmdnet
+
+
diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -0,0 +1,88 @@
+# Configuration file for the Sphinx documentation builder.
+
+# -- Project information
+project = "TorchMD-Net"
+author = "RaulPPelaez"
+
+import git
+
+
+def get_latest_git_tag(repo_path="."):
+    repo = git.Repo(repo_path)
+    tags = sorted(repo.tags, key=lambda t: t.commit.committed_datetime)
+    return tags[-1].name if tags else None
+
+
+current_tag = get_latest_git_tag("../../")
+if current_tag is None:
+    current_tag = "master"
+release = current_tag
+version = current_tag
+
+# -- General configuration
+extensions = [
+    "sphinx.ext.duration",
+    "sphinx.ext.doctest",
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+    "sphinxcontrib.autoprogram",
+]
+napoleon_google_docstring = True
+napoleon_numpy_docstring = True
+napoleon_include_init_with_doc = False
+napoleon_include_private_with_doc = False
+napoleon_include_special_with_doc = True
+napoleon_use_admonition_for_examples = True
+napoleon_use_admonition_for_notes = True
+napoleon_use_admonition_for_references = True
+napoleon_use_ivar = False
+napoleon_use_param = True
+napoleon_use_rtype = True
+napoleon_preprocess_types = True
+napoleon_type_aliases = None
+napoleon_attr_annotations = True
+autosummary_ignore_module_all = False
+
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3/", None),
+    "torch": ("https://pytorch.org/docs/stable/", None),
+    "sphinx": ("https://www.sphinx-doc.org/en/master/", None),
+}
+intersphinx_disabled_domains = ["std"]
+
+templates_path = ["_templates"]
+
+# -- Options for HTML output
+
+html_theme = "sphinx_rtd_theme"
+
+# -- Options for EPUB output
+epub_show_urls = "footnote"
+
+autoclass_content = "both"
+autodoc_typehints = "none"
+autodoc_inherit_docstrings = False
+sphinx_autodoc_typehints = True
+html_show_sourcelink = True
+autodoc_default_options = {
+    "members": True,
+    "member-order": "bysource",
+    "exclude-members": "__weakref__",
+    "undoc-members": False,
+    "show-inheritance": True,
+    "inherited-members": False,
+}
+# Exclude all torchmdnet.datasets.*.rst files in source/generated/
+exclude_patterns = [
+    "generated/torchmdnet.datasets.*.rst",
+    "generated/torchmdnet.scripts.*rst",
+]
+html_static_path = ["../_static"]
+html_css_files = [
+    "style.css",
+]
+
+autodoc_mock_imports = ["torchmdnet.extensions.torchmdnet_extensions"]
diff --git a/docs/source/datasets.rst b/docs/source/datasets.rst
@@ -0,0 +1,81 @@
+Datasets
+========
+.. toctree::
+   :maxdepth: 3
+
+
+
+Using a Dataset
+---------------
+
+There are two ways of using a dataset in TorchMD-Net, depending on whether you are using the Python API (for instance to run inferences) or the command line interface (for training).
+
+Via a configuration file
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+You can make use of any of the available datasets via the :ref:`configuration file <configuration-file>` if you are using the :ref:`torchmd-train utility <torchmd-train>`.
+Take a look at one of the example configuration files. As an example lets set up the :py:mod:`QM9` dataset, which allows for an additional argument `labels` specifying the subset to be provided. The part of the yaml configuration file for the dataset would look like this:
+
+.. code:: yaml
+
+   dataset: QM9
+   dataset_arg:
+     label: energy_U0
+
+
+
+Via the Python API
+~~~~~~~~~~~~~~~~~~
+
+TorchMD-Net datasets are inherited from `torch Geometric datasets <https://pytorch-geometric.readthedocs.io/en/latest/tutorial/create_dataset.html>`_, you may use them whenever the PyTorch Geometric datasets are used. For instance, to load the QM9 dataset, you may use the following code:
+
+.. code:: python
+
+   from torchmdnet.datasets import QM9
+   from torchmdnet.data import DataModule
+   dataset = QM9(root='data', labels='energy_U0')
+   
+   print(dataset[0])
+   print(len(dataset))
+   # Some arbitrary parameters for the DataModule
+   params = {'batch_size': 32,
+	     'inference_batch_size': 32,
+             'num_workers': 4,
+	     'train_size': 0.8,
+	     'val_size': 0.1,
+	     'test_size': 0.1,
+	     'seed': 42,
+	     'log_dir': 'logs',
+	     'splits': None,
+	     'standardize': False,}
+   
+   dataloader = DataModule(params, dataset)
+   dataloader.prepare_data()
+   dataloader.setup("fit")
+
+   # You can use this directly with PyTorch Lightning
+   # trainer.fit(model, dataloader)
+
+   
+Adding a new Dataset
+--------------------
+
+If you want to train on custom data, first have a look at :py:mod:`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.
+
+To add a new dataset, you need to:
+
+1. Write a new class inheriting from :py:mod:`torch_geometric.data.Dataset`, see `here <https://pytorch-geometric.readthedocs.io/en/latest/tutorial/create_dataset.html>`_ for a tutorial on that.
+2. Add the new class to the :py:mod:`torchmdnet.datasets` module (by listing it in the `__all__` variable in the `datasets/__init__.py` file) so that the :ref:`configuration file <configuration-file>` recognizes it.
+
+.. note::
+
+   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.
+
+Available Datasets
+------------------
+
+.. automodule:: torchmdnet.datasets
+   :noindex:
+
+   .. include:: generated/torchmdnet.datasets.rst
+      :start-line: 5