0.2.0 Release

.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

README.md

@@ -1,8 +1,11 @@
 # `apax`: Atomistic 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)
 [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/python/black)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
 
-`apax` is a high-performance, extendable package for training of and inference with atomistic neural networks.
+
+`apax`[1] is a high-performance, extendable package for training of and inference with atomistic neural networks.
 It implements the Gaussian Moment Neural Network model [2, 3].
 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.
 
@@ -39,12 +42,12 @@ pip install --upgrade pip
 
 CUDA 12 installation. Wheels only available on linux.
 ```bash
-pip install --upgrade "jax[cuda12_pip]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 CUDA 11 installation. Wheels only available on linux.
 ```bash
-pip install --upgrade "jax[cuda11_pip]==0.4.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 ```
 
 See the [Jax installation instructions](https://github.com/google/jax#installation) for more details.
@@ -91,14 +94,21 @@ The second way is to use the ASE calculator provided in `apax.md`.
 
 Under the supervion of Johannes Kästner
 
-## References
-* [1] DOI PLACEHOLDER
-* [2] V. Zaverkin and J. Kästner, [“Gaussian Moments as Physically Inspired Molecular Descriptors for Accurate and Scalable Machine Learning Potentials,”](https://doi.org/10.1021/acs.jctc.0c00347) J. Chem. Theory Comput. **16**, 5410–5421 (2020).
-* [3] V. Zaverkin, D. Holzmüller, I. Steinwart,  and J. Kästner, [“Fast and Sample-Efficient Interatomic Neural Network Potentials for Molecules and Materials Based on Gaussian Moments,”](https://pubs.acs.org/doi/10.1021/acs.jctc.1c00527) J. Chem. Theory Comput. **17**, 6658–6670 (2021).
-
 
 ## Contributing
 
 We are happy to receive your issues and pull requests!
 
 Do not hesitate to contact any of the authors above if you have any further questions.
+
+
+## Acknowledgements
+
+The creation of Apax was supported by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) in the framework of the priority program SPP 2363, “Utilization and Development of Machine Learning for Molecular Applications - Molecular Machine Learning” Project No. 497249646 and the Ministry of Science, Research and the Arts Baden-Württemberg in the Artificial Intelligence Software Academy (AISA).
+Further funding though the DFG under Germany's Excellence Strategy - EXC 2075 - 390740016 and the Stuttgart Center for Simulation Science (SimTech) was provided.
+
+
+## References
+* [1] 10.5281/zenodo.10040711
+* [2] V. Zaverkin and J. Kästner, [“Gaussian Moments as Physically Inspired Molecular Descriptors for Accurate and Scalable Machine Learning Potentials,”](https://doi.org/10.1021/acs.jctc.0c00347) J. Chem. Theory Comput. **16**, 5410–5421 (2020).
+* [3] V. Zaverkin, D. Holzmüller, I. Steinwart,  and J. Kästner, [“Fast and Sample-Efficient Interatomic Neural Network Potentials for Molecules and Materials Based on Gaussian Moments,”](https://pubs.acs.org/doi/10.1021/acs.jctc.1c00527) J. Chem. Theory Comput. **17**, 6658–6670 (2021).

apax/bal/api.py

@@ -8,11 +8,16 @@
 from tqdm import trange
 
 from apax.bal import feature_maps, kernel, selection, transforms
-from apax.data.input_pipeline import TFPipeline
+from apax.data.initialization import RawDataset
+from apax.data.input_pipeline import AtomisticDataset
 from apax.model.builder import ModelBuilder
 from apax.model.gmnn import EnergyModel
-from apax.train.checkpoints import restore_parameters
-from apax.train.run import RawDataset, initialize_dataset
+from apax.train.checkpoints import (
+    canonicalize_energy_model_parameters,
+    check_for_ensemble,
+    restore_parameters,
+)
+from apax.train.run import initialize_dataset
 
 
 def create_feature_fn(
@@ -43,11 +48,11 @@ def create_feature_fn(
     return feature_fn
 
 
-def compute_features(feature_fn, dataset: TFPipeline, processing_batch_size: int):
+def compute_features(feature_fn, dataset: AtomisticDataset):
     """Compute the features of a dataset."""
     features = []
     n_data = dataset.n_data
-    ds = dataset.batch(processing_batch_size)
+    ds = dataset.batch()
 
     pbar = trange(n_data, desc="Computing features", ncols=100, leave=True)
     for i, (inputs, _) in enumerate(ds):
@@ -70,29 +75,32 @@ def kernel_selection(
     selection_batch_size: int = 10,
     processing_batch_size: int = 64,
 ):
-    n_models = 1 if isinstance(model_dir, (Path, str)) else len(model_dir)
-    is_ensemble = n_models > 1
-
     selection_fn = {
         "max_dist": selection.max_dist_selection,
     }[selection_method]
 
     base_feature_map = feature_maps.FeatureMapOptions(base_fm_options)
 
     config, params = restore_parameters(model_dir)
+    params = canonicalize_energy_model_parameters(params)
+    n_models = check_for_ensemble(params)
+    is_ensemble = n_models > 1
 
     n_train = len(train_atoms)
-    dataset = initialize_dataset(config, RawDataset(atoms_list=train_atoms + pool_atoms))
+    dataset = initialize_dataset(
+        config, RawDataset(atoms_list=train_atoms + pool_atoms), calc_stats=False
+    )
+    dataset.set_batch_size(processing_batch_size)
 
-    init_box = dataset.init_input()["box"][0]
+    _, init_box = dataset.init_input()
 
     builder = ModelBuilder(config.model.get_dict(), n_species=119)
     model = builder.build_energy_model(apply_mask=True, init_box=init_box)
 
     feature_fn = create_feature_fn(
         model, params, base_feature_map, feature_transforms, is_ensemble
     )
-    g = compute_features(feature_fn, dataset, processing_batch_size)
+    g = compute_features(feature_fn, dataset)
     km = kernel.KernelMatrix(g, n_train)
     new_indices = selection_fn(km, selection_batch_size)
 

apax/cli/apax_app.py

@@ -1,5 +1,6 @@
 import importlib.metadata
 import importlib.resources as pkg_resources
+import sys
 from pathlib import Path
 
 import typer
@@ -34,15 +35,14 @@ def train(
     train_config_path: Path = typer.Argument(
         ..., help="Training configuration YAML file."
     ),
-    log_level: str = typer.Option("error", help="Sets the training logging level."),
-    log_file: str = typer.Option("train.log", help="Specifies the name of the log file"),
+    log_level: str = typer.Option("info", help="Sets the training logging level."),
 ):
     """
     Starts the training of a model with parameters provided by a configuration file.
     """
     from apax.train.run import run
 
-    run(train_config_path, log_file, log_level)
+    run(train_config_path, log_level)
 
 
 @app.command()
@@ -51,16 +51,15 @@ def md(
         ..., help="Configuration YAML file that was used to train a model."
     ),
     md_config_path: Path = typer.Argument(..., help="MD configuration YAML file."),
-    log_level: str = typer.Option("error", help="Sets the training logging level."),
-    log_file: str = typer.Option("md.log", help="Specifies the name of the log file"),
+    log_level: str = typer.Option("info", help="Sets the training logging level."),
 ):
     """
     Starts performing a molecular dynamics simulation (currently only NHC thermostat)
     with parameters provided by a configuration file.
     """
     from apax.md import run_md
 
-    run_md(train_config_path, md_config_path, log_file, log_level)
+    run_md(train_config_path, md_config_path, log_level)
 
 
 @app.command()
@@ -90,8 +89,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")
@@ -211,6 +210,7 @@ def template_train_config(
 
     if Path(config_path).is_file():
         console.print("There is already a config file in the working directory.")
+        sys.exit(1)
     else:
         with open(config_path, "w") as config:
             config.write(template_content)
@@ -229,6 +229,7 @@ def template_md_config():
 
     if Path(config_path).is_file():
         console.print("There is already a config file in the working directory.")
+        sys.exit(1)
     else:
         with open(config_path, "w") as config:
             config.write(template_content)
@@ -249,7 +250,3 @@ def main(
 ):
     # Taken from https://github.com/zincware/dask4dvc/blob/main/dask4dvc/cli/main.py
     _ = version
-
-
-if __name__ == "__main__":
-    app()

apax/cli/templates/md_config_minimal.yaml

@@ -1,8 +1,11 @@
-temperature: <T> # K
-duration: <DURATION> #fs
-n_inner: 100 # compiled innner steps
-sampling_rate: 10 # dump interval
-dt: 0.5 # fs time step
+ensemble:
+  name: nvt
+  dt: 0.5 # fs time step
+  temperature: <T> # K
+
+duration: <DURATION> # fs
+n_inner: 1 # compiled innner steps
+sampling_rate: 1  # dump interval
 dr_threshold: 0.5 # Neighborlist skin
 
 sim_dir: md

apax/cli/templates/train_config_full.yaml

@@ -79,5 +79,3 @@ checkpoints:
 progress_bar:
   disable_epoch_pbar: false
   disable_nl_pbar: false
-
-maximize_l2_cache: true

apax/config/common.py

@@ -18,7 +18,6 @@ def parse_config(config: Union[str, os.PathLike, dict], mode: str = "train") ->
         config: Path to the config file or a dictionary
         containing the config.
     """
-    log.info("Loading user config")
     if isinstance(config, (str, os.PathLike)):
         with open(config, "r") as stream:
             config = yaml.safe_load(stream)

apax/config/md_config.py

@@ -46,7 +46,9 @@ class MDConfig(BaseModel, frozen=True, extra="forbid"):
     n_inner: Number of compiled simulation steps (i.e. number of iterations of the
         `jax.lax.fori_loop` loop). Also determines atoms buffer size.
     sampling_rate:
-        Trajectory dumping interval.
+        Interval between saving frames.
+    buffer_size:
+        Number of collected frames to be dumped at once.
     dr_threshold: Skin of the neighborlist.
     extra_capacity: JaxMD allocates a maximal number of neighbors.
         This argument lets you add additional capacity to avoid recompilation.
@@ -56,6 +58,9 @@ class MDConfig(BaseModel, frozen=True, extra="forbid"):
     traj_name: Name of the trajectory file.
     restart: Whether the simulation should restart from the latest configuration
         in `traj_name`.
+    checkpoint_interval: Number of time steps between saving
+        full simulation state checkpoints. These will be loaded
+        with the `restart` option.
     disable_pbar: Disables the MD progressbar.
     """
 
@@ -69,6 +74,7 @@ class MDConfig(BaseModel, frozen=True, extra="forbid"):
     duration: PositiveFloat
     n_inner: PositiveInt = 100
     sampling_rate: PositiveInt = 10
+    buffer_size: PositiveInt = 100
     dr_threshold: PositiveFloat = 0.5
     extra_capacity: PositiveInt = 0
 
@@ -77,6 +83,7 @@ class MDConfig(BaseModel, frozen=True, extra="forbid"):
     sim_dir: str = "."
     traj_name: str = "md.h5"
     restart: bool = True
+    checkpoint_interval: int = 50_000
     disable_pbar: bool = False
 
     def dump_config(self):

apax/config/train_config.py

@@ -112,12 +112,14 @@ def validate_shift_scale_methods(self):
 
         return self
 
+    @property
     def model_version_path(self):
         version_path = Path(self.directory) / self.experiment
         return version_path
 
+    @property
     def best_model_path(self):
-        return self.model_version_path() / "best"
+        return self.model_version_path / "best"
 
 
 class ModelConfig(BaseModel, extra="forbid"):
@@ -286,12 +288,12 @@ class Config(BaseModel, frozen=True, extra="forbid"):
     callbacks: List of :class: `callback` <config.CallbackConfig> configurations.
     progress_bar: Progressbar configuration.
     checkpoints: Checkpoint configuration.
-    maximize_l2_cache: Whether or not to maximize GPU L2 cache.
     """
 
     n_epochs: PositiveInt
     patience: Optional[PositiveInt] = None
     seed: int = 1
+    n_models: int = 1
 
     data: DataConfig
     model: ModelConfig = ModelConfig()
@@ -301,7 +303,6 @@ class Config(BaseModel, frozen=True, extra="forbid"):
     callbacks: List[CallbackConfig] = [CallbackConfig(name="csv")]
     progress_bar: TrainProgressbarConfig = TrainProgressbarConfig()
     checkpoints: CheckpointConfig = CheckpointConfig()
-    maximize_l2_cache: bool = False
 
     def dump_config(self, save_path):
         """