Merge pull request #204 from apax-hub/docs

Docs update

.readthedocs.yaml

@@ -0,0 +1,20 @@
+# .readthedocs.yaml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+version: 2
+
+# Set the version of Python and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
+  jobs:
+    post_install:
+      - pip install poetry
+      - poetry config virtualenvs.create false
+      - poetry install --with=docs
+
+# Build documentation in the docs/ directory with Sphinx
+sphinx:
+   configuration: docs/source/conf.py

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)

README.md

@@ -1,4 +1,5 @@
 # `apax`: Atomstic learned Potentials in JAX!
+[![Read the Docs](https://readthedocs.org/projects/apax/badge/)](https://apax.readthedocs.io/en/latest/)
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.10040711.svg)](https://doi.org/10.5281/zenodo.10040711)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/python/black)

apax/cli/apax_app.py

@@ -90,8 +90,8 @@ def docs():
     """
     Opens the documentation website in your browser.
     """
-    console.print("Opening apax's docs at https://github.com/apax-hub/apax")
-    typer.launch("https://github.com/apax-hub/apax")
+    console.print("Opening apax's docs at https://apax.readthedocs.io/en/latest/")
+    typer.launch("https://apax.readthedocs.io/en/latest/")
 
 
 @validate_app.command("train")

apax/utils/data.py

@@ -21,22 +21,24 @@ def make_minimal_input():
 
 def load_data(data_path):
     """Non ASE compatible parameters have to be saved in an exta file that has the same
-        name as the datapath but with the extension '_labels.npz'.
-        example for the npz-file:
-                dipole = np.random.rand(3, 1)
-                charge = np.random.rand(3, 2)
-                mat = np.random.rand(3, 1)
-                shape = ['ragged', 'ragged', 'fixed']
-
-                np.savez(
-                    "data_path_labels.npz",
-                    dipole=dipole,
-                    charge=charge,
-                    mat=mat,
-                    shape=shape
-                )
-
-        shape has to be in the same order than the parameters
+    name as the datapath but with the extension `_labels.npz`.
+    example for the npz-file:
+
+    dipole = np.random.rand(3, 1)
+    charge = np.random.rand(3, 2)
+    mat = np.random.rand(3, 1)
+    shape = ['ragged', 'ragged', 'fixed']
+
+    np.savez(
+        "data_path_labels.npz",
+        dipole=dipole,
+        charge=charge,
+        mat=mat,
+        shape=shape,
+    )
+
+    shape has to be in the same order than the parameters
+
     Parameters
     ----------
     data_path :
@@ -46,6 +48,7 @@ def load_data(data_path):
     -------
     atoms_list
         List of all structures where entries are ASE atoms objects.
+
     """
     external_labels = {}
     # TODO external labels can be included via hdf5 files only? this would clean up a lot

docs/source/_tutorials/index.rst

@@ -0,0 +1,10 @@
+Tutorials
+=========
+
+.. toctree::
+   :maxdepth: 2
+
+   md_with_ase
+   molecular_dynamics
+   training_a_model
+   transfer_learning

docs/source/_tutorials/md_with_ase.rst

@@ -1,4 +1,5 @@
-# The ASE calculator
+The ASE calculator
+==================
 
 If you require some ASE features during your simulation, we provide an alternative to the JaxMD interface.
 

docs/source/_tutorials/molecular_dynamics.rst

@@ -1,4 +1,5 @@
-# Molecular Dynamics with JaxMD
+Molecular Dynamics with JaxMD
+=============================
 
 
 JaxMD LINK is a high performance molecular dynamics engine built on top of Jax LINK.

docs/source/_tutorials/training_a_model.rst

@@ -1,4 +1,5 @@
-# Training
+Training
+========
 
 ## Acquiring a dataset
 
@@ -65,7 +66,8 @@ callbacks:
 
 If training is interrupted for any reason, re-running the above `train` command will resume training from the latest checkpoint.
 
-## Evaluation
+Evaluation
+==========
 
 After the training is completed and we are satisfied with our choice of hyperparameters and vadliation loss, we can evaluate the model on the test set.
 We provide a separate command for test set evaluation:

docs/source/_tutorials/transfer_learning.rst

@@ -1,4 +1,5 @@
-# Transfer Learning
+Transfer Learning
+=================
 
 apax comes with discriminative transfer learning capabilities out of the box.
 In this tutorial we are going to fine tune a model trained on benzene data at the DFT level of theory to CCSDT.

docs/source/conf.py

@@ -12,7 +12,7 @@
 sys.path.insert(0, os.path.abspath("."))
 sys.path.insert(0, os.path.abspath("../../"))
 
-project = "apax"
+project = "Apax"
 copyright = "2023, Moritz Schäfer, Nico Segreto, Johannes Kästner"
 author = "Moritz Schäfer, Nico Segreto, Johannes Kästner"
 release = "0.1.0"
@@ -35,5 +35,6 @@
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
-html_theme = "sphinx_rtd_theme"
-html_static_path = ["_static"]
+html_theme = "furo"
+
+html_theme_options = {}

docs/source/getting_started/index.rst

@@ -0,0 +1,7 @@
+Getting Started
+===============
+
+.. toctree::
+   :maxdepth: 2
+
+   install

docs/source/getting_started/install.rst

@@ -1,4 +1,3 @@
-============
 Installation
 ============
 

docs/source/index.rst

@@ -1,47 +1,14 @@
-.. apax documentation master file, created by
-   sphinx-quickstart on Wed Jan 25 21:12:03 2023.
-   You can adapt this file completely to your liking, but it should at least
-   contain the root `toctree` directive.
-
 ===================================================
-apax: Gaussian Moment Neural Networks - in Jax!
+apax: Atomistic learned Potentials in JAX!
 ===================================================
 
-`apax` is an implementation of Gaussian Moment Neural Networks in Jax. 
-It facilitates the training and evaluation of atomistic neural networks, notably GM-NN, for molecular simulations.
-Building on top of Jax allows for easy integration with the high-performance molecular dynamics engine `JaxMD <https://github.com/jax-md/jax-md>`_
-
-Contents
-========
-.. toctree::
-   :maxdepth: 1
-   :caption: Getting Started
-
-   getting_started/install
-
-.. .. toctree:: TODO
-..    :maxdepth: 1
-..    :caption: Tutorials
 
-..    tutorials/...
+`apax` is a high-performance, extendable package for training of and inference with atomistic neural networks., most prominently the Gaussian Moment Neural Network.
+It is based on [JAX](https://jax.readthedocs.io/en/latest/) and uses [JaxMD](https://github.com/jax-md/jax-md) as a molecular dynamics engine.
 
-Package
-=======
 .. toctree::
-   :maxdepth: 2
-   :caption: Package
-
-Config
-=====================
-.. automodule:: apax.config
-   :members:
-
-   .. modules/index
-
-
-Indices and Tables
-==================
+   :hidden:
 
-* :ref:`genindex`
-* :ref:`modindex`
-* :ref:`search`
+   getting_started/index
+   _tutorials/index
+   modules/index

docs/source/modules/bal.rst

@@ -0,0 +1,17 @@
+Batch Active Learning
+=====================
+
+.. automodule:: apax.bal.api
+    :members:
+
+.. automodule:: apax.bal.feature_maps
+    :members:
+
+.. automodule:: apax.bal.kernel
+    :members:
+
+.. automodule:: apax.bal.selection
+    :members:
+
+.. automodule:: apax.bal.transforms
+    :members:

docs/source/modules/cli.rst

@@ -0,0 +1,5 @@
+Command Line Interface
+======================
+
+.. automodule:: apax.cli.apax_app
+    :members:

docs/source/modules/data.rst

@@ -1,2 +1,16 @@
 Data Pipeline
-==============
+==============
+
+:mod:`apax.data`
+
+.. automodule:: apax.data.preprocessing
+    :members:
+
+.. automodule:: apax.data.initialization
+    :members:
+
+.. automodule:: apax.data.input_pipeline
+    :members:
+
+.. automodule:: apax.data.statistics
+    :members:

docs/source/modules/index.rst

@@ -2,13 +2,16 @@ Modules
 =======
 
 .. toctree::
-   :maxdepth: 1
+   :maxdepth: 2
 
+   cli
    config
    data
    optimizer
    layers
    model
    train
    md
-   utils
+   transfer_learning
+   bal
+   utils

docs/source/modules/layers.rst

@@ -1,2 +1,27 @@
 Layers
 ======
+
+
+.. automodule:: apax.layers.activation
+    :members:
+
+.. automodule:: apax.layers.empirical
+    :members:
+
+.. automodule:: apax.layers.initializers
+    :members:
+
+.. automodule:: apax.layers.masking
+    :members:
+
+.. automodule:: apax.layers.ntk_linear
+    :members:
+
+.. automodule:: apax.layers.properties
+    :members:
+
+.. automodule:: apax.layers.readout
+    :members:
+
+.. automodule:: apax.layers.scaling
+    :members:

docs/source/modules/md.rst

@@ -1,2 +1,14 @@
 Molecular Dynamics
 ==================
+
+.. automodule:: apax.md.ase_calc
+    :members:
+
+.. automodule:: apax.md.function_transformations
+    :members:
+
+.. automodule:: apax.md.io
+    :members:
+
+.. automodule:: apax.md.nvt
+    :members:

docs/source/modules/model.rst

@@ -1,2 +1,8 @@
 Model
-=====
+=====
+
+.. automodule:: apax.model.builder
+    :members:
+
+.. automodule:: apax.model.gmnn
+    :members:

docs/source/modules/optimizer.rst

@@ -1,2 +1,5 @@
 Optimizers
 ==========
+
+.. automodule:: apax.optimizer.get_optimizer
+    :members:

docs/source/modules/train.rst

@@ -1,2 +1,23 @@
 Training
 ========
+
+.. automodule:: apax.train.callbacks
+    :members:
+
+.. automodule:: apax.train.checkpoints
+    :members:
+
+.. automodule:: apax.train.eval
+    :members:
+
+.. automodule:: apax.train.loss
+    :members:
+
+.. automodule:: apax.train.metrics
+    :members:
+
+.. automodule:: apax.train.run
+    :members:
+
+.. automodule:: apax.train.trainer
+    :members: