### Setup configuration directory structure

Source: https://schnetpack.readthedocs.io/en/latest/_sources/userguide/configs.rst.txt

Commands to create a standard configuration directory structure and run an experiment.

```bash
# create main config directory and subdirectory for experiment config group
$ mkdir configs
$ mkdir configs/experiment

# create experiment and edit config
$ vim my_configs/experiments/my_experiment.yaml

# explicitly stating config dir is not required since it is in current working dir
$ spktrain experiment=my_experiment
```

--------------------------------

### Install Tensorboard

Source: https://schnetpack.readthedocs.io/en/latest/_sources/getstarted.rst.txt

Install Tensorboard, the default logging backend for SchNetPack, using pip.

```bash
pip install tensorboard
```

--------------------------------

### Create working directory for training

Source: https://schnetpack.readthedocs.io/en/latest/getstarted.html

Set up a directory to store data and training runs before starting the training process.

```bash
mkdir spk_workdir
cd spk_workdir
```

--------------------------------

### Install SchNetPack using pip

Source: https://schnetpack.readthedocs.io/en/latest/_sources/getstarted.rst.txt

Use this command for the simplest installation of SchNetPack, which automatically fetches the source code from PyPI.

```bash
pip install --upgrade schnetpack
```

--------------------------------

### Install SchNetPack from Source

Source: https://schnetpack.readthedocs.io/en/latest/_sources/getstarted.rst.txt

Install the latest code directly from the GitHub repository. This involves cloning the repository and then installing it using pip.

```bash
git clone https://github.com/atomistic-machine-learning/schnetpack.git <dest_dir>
cd <dest_dir>
pip install .
```

--------------------------------

### Initialize QM9 dataset

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_01_preparing_data.ipynb.txt

Instantiate the QM9 benchmark dataset, which triggers automatic download and setup of data partitions.

```python
from schnetpack.datasets import QM9
from schnetpack.transform import ASENeighborList

qm9data = QM9(
    "./qm9.db",
    batch_size=10,
    num_train=110000,
    num_val=10000,
    split_file="./split_qm9.npz",
    transforms=[ASENeighborList(cutoff=5.0)],
)
qm9data.prepare_data()
qm9data.setup()
```

--------------------------------

### Start QM9 Training with Default Settings

Source: https://schnetpack.readthedocs.io/en/latest/_sources/getstarted.rst.txt

Initiate training for a SchNet model on the QM9 dataset using the `spktrain` CLI with default configurations. The dataset will be downloaded automatically.

```bash
mkdir spk_workdir
cd spk_workdir
spktrain experiment=qm9_atomwise
```

--------------------------------

### Train PaiNN on QM9 Example Command

Source: https://schnetpack.readthedocs.io/en/latest/_sources/userguide/configs.rst.txt

This command initiates the training of a PaiNN model on the QM9 dataset using a predefined experiment configuration. Ensure the 'spktrain' executable is in your PATH and the experiment configuration is accessible.

```bash
spktrain experiment=qm9_atomwise
```

--------------------------------

### Import Libraries and Create Directory

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_03_force_models.ipynb.txt

Imports necessary PyTorch and SchNetPack libraries and creates a directory for tutorial experiments. Ensure these libraries are installed before running.

```python
import torch
import torchmetrics
import schnetpack as spk
import schnetpack.transform as trn
import pytorch_lightning as pl
import os
import matplotlib.pyplot as plt
import numpy as np

forcetut = "./forcetut"
if not os.path.exists(forcetut):
    os.makedirs(forcetut)
```

--------------------------------

### Initialize OrcaPropertyParser

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/md.parsers.OrcaPropertyParser.html

Initializes the OrcaPropertyParser with start and stop flags. Optionally accepts formatters for data processing.

```python
OrcaPropertyParser(_start : str_, _stop : str | List[str]_, _formatters : OrcaFormatter | List[OrcaFormatter] | None = None_)
```

--------------------------------

### Configure SchNetPack experiment with stress and neighbor caching

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_05_materials.html

Example YAML configuration for a SchNetPack experiment, demonstrating the inclusion of stress prediction, neighbor list caching, and atom selection constraints.

```yaml
# @package _global_

defaults:
  - override /model: nnp

run.path: runs

globals:
  cutoff: 5.
  lr: 5e-4
  energy_key: energy
  forces_key: forces
  stress_key: stress
data:
  distance_unit: Ang
  property_units:
    energy: eV
    forces: eV/Ang
    stress: eV/Ang/Ang/Ang
  transforms:
    - _target_: schnetpack.transform.RemoveOffsets
      property: energy
      remove_mean: True
    - _target_: schnetpack.transform.CachedNeighborList
      cache_path: ${run.work_dir}/cache
      keep_cache: False
      neighbor_list:
        _target_: schnetpack.transform.MatScipyNeighborList
        cutoff: ${globals.cutoff}
      nbh_transforms:
        - _target_: schnetpack.transform.FilterNeighbors
          selection_name: filtered_out_neighbors
    - _target_: schnetpack.transform.CastTo32
  test_transforms:
    - _target_: schnetpack.transform.RemoveOffsets
      property: energy
      remove_mean: True
    - _target_: schnetpack.transform.MatScipyNeighborList
      cutoff: ${globals.cutoff}
    - _target_: schnetpack.transform.FilterNeighbors
      selection_name: filtered_out_neighbors
    - _target_: schnetpack.transform.CastTo32

model:
  input_modules:
    - _target_: schnetpack.atomistic.Strain
    - _target_: schnetpack.atomistic.PairwiseDistances
  output_modules:
    - _target_: schnetpack.atomistic.Atomwise
      output_key: ${globals.energy_key}
      n_in: ${model.representation.n_atom_basis}
      aggregation_mode: sum
    - _target_: schnetpack.atomistic.Forces
      energy_key: ${globals.energy_key}
      force_key: ${globals.forces_key}
      stress_key: ${globals.stress_key}
      calc_stress: True
  postprocessors:
    - _target_: schnetpack.transform.CastTo64
    - _target_: schnetpack.transform.AddOffsets
      property: energy
      add_mean: True

task:
  optimizer_cls: torch.optim.AdamW
  optimizer_args:
    lr: ${globals.lr}
    weight_decay: 0.01
  outputs:
    - _target_: schnetpack.task.ModelOutput
      name: ${globals.energy_key}
      loss_fn:
        _target_: torch.nn.MSELoss
      metrics:
        mae:
          _target_: torchmetrics.regression.MeanAbsoluteError
        mse:
          _target_: torchmetrics.regression.MeanSquaredError
      loss_weight: 0.0
    - _target_: schnetpack.task.ModelOutput
      name: ${globals.forces_key}
      constraints:
        - _target_: schnetpack.task.ConsiderOnlySelectedAtoms
          selection_name: considered_atoms
      loss_fn:
        _target_: torch.nn.MSELoss
      metrics:
        mae:
          _target_: torchmetrics.regression.MeanAbsoluteError
        mse:
          _target_: torchmetrics.regression.MeanSquaredError
      loss_weight: 1.0
    - _target_: schnetpack.task.ModelOutput
      name: ${globals.stress}
      loss_fn:
        _target_: torch.nn.MSELoss
      metrics:
        mae:
          _target_: torchmetrics.regression.MeanAbsoluteError
        mse:
          _target_: torchmetrics.regression.MeanSquaredError
      loss_weight: 100.
```

--------------------------------

### Full Configuration Structure

Source: https://schnetpack.readthedocs.io/en/latest/userguide/configs.html

Example of a complete hierarchical Hydra configuration for training a PaiNN model on the QM9 dataset.

```yaml
⚙ Running with the following config:
├── run
│   └── work_dir: ${hydra:runtime.cwd}
│       data_dir: ${run.work_dir}/data
│       path: runs
│       id: ${uuid:1}
│
├── globals
│   └── model_path: best_model
│       cutoff: 5.0
│       lr: 0.0005
│       property: energy_U0
│
├── data
│   └── _target_: schnetpack.datasets.QM9
│       datapath: ${run.data_dir}/qm9.db
│       data_workdir: null
│       batch_size: 100
│       num_train: 110000
│       num_val: 10000
│       num_test: null
│       num_workers: 8
│       num_val_workers: null
│       num_test_workers: null
│       remove_uncharacterized: false
│       distance_unit: Ang
│       property_units:
│         energy_U0: eV
│         energy_U: eV
│         enthalpy_H: eV
│         free_energy: eV
│         homo: eV
│         lumo: eV
│         gap: eV
│         zpve: eV
│       transforms:
│       - _target_: schnetpack.transform.SubtractCenterOfMass
│       - _target_: schnetpack.transform.RemoveOffsets
│         property: ${globals.property}
│         remove_atomrefs: true
│         remove_mean: true
│       - _target_: schnetpack.transform.MatScipyNeighborList
│         cutoff: ${globals.cutoff}
│       - _target_: schnetpack.transform.CastTo32
│
├── model
│   └── representation:
│         _target_: schnetpack.representation.PaiNN
│         n_atom_basis: 128
│         n_interactions: 3
│         shared_interactions: false
│         shared_filters: false
│         radial_basis:
│           _target_: schnetpack.nn.radial.GaussianRBF
│           n_rbf: 20
│           cutoff: ${globals.cutoff}
│         cutoff_fn:
│           _target_: schnetpack.nn.cutoff.CosineCutoff
│           cutoff: ${globals.cutoff}
│       _target_: schnetpack.model.NeuralNetworkPotential
│       input_modules:
│       - _target_: schnetpack.atomistic.PairwiseDistances
│       output_modules:
│       - _target_: schnetpack.atomistic.Atomwise
│         output_key: ${globals.property}
│         n_in: ${model.representation.n_atom_basis}
│         aggregation_mode: sum
│       postprocessors:
│       - _target_: schnetpack.transform.CastTo64
│       - _target_: schnetpack.transform.AddOffsets
│         property: ${globals.property}
│         add_mean: true
│         add_atomrefs: true
│
├── task
│   └── optimizer_cls: torch.optim.AdamW
│       optimizer_args:
│         lr: ${globals.lr}
│         weight_decay: 0.0
│       scheduler_cls: schnetpack.train.ReduceLROnPlateau
│       scheduler_monitor: val_loss
│       scheduler_args:
│         mode: min
│         factor: 0.8
│         patience: 80
│         threshold: 0.0001
│         threshold_mode: rel
│         cooldown: 10
│         min_lr: 0.0
│         smoothing_factor: 0.0
│       _target_: schnetpack.AtomisticTask
│       outputs:
│       - _target_: schnetpack.task.ModelOutput
│         name: ${globals.property}
│         loss_fn:
│           _target_: torch.nn.MSELoss
│         metrics:
│           mae:
│             _target_: torchmetrics.regression.MeanAbsoluteError
│           mse:
│             _target_: torchmetrics.regression.MeanSquaredError
│         loss_weight: 1.0
│       warmup_steps: 0
│
├── trainer
│   └── _target_: pytorch_lightning.Trainer
│       devices: 1
│       min_epochs: null
│       max_epochs: 100000
│       enable_model_summary: true
│       profiler: null
│       gradient_clip_val: 0
│       accumulate_grad_batches: 1
│       val_check_interval: 1.0
│       check_val_every_n_epoch: 1
│       num_sanity_val_steps: 0
│       fast_dev_run: false
│       overfit_batches: 0
│       limit_train_batches: 1.0
│       limit_val_batches: 1.0
│       limit_test_batches: 1.0
│       track_grad_norm: -1
│       detect_anomaly: false
│       amp_backend: native
│       amp_level: null
│       precision: 32
│       accelerator: auto
│       num_nodes: 1
│       tpu_cores: null
│       deterministic: false
│       resume_from_checkpoint: null
│
├── callbacks
│   └── model_checkpoint:
│         _target_: schnetpack.train.ModelCheckpoint
│         monitor: val_loss
│         save_top_k: 1
│         save_last: true
│         mode: min
│         verbose: false
│         dirpath: checkpoints/
│         filename: '{epoch:02d}'
│         model_path: ${globals.model_path}
│       early_stopping:
│         _target_: pytorch_lightning.callbacks.EarlyStopping
│         monitor: val_loss
│         patience: 1000
│         mode: min
│         min_delta: 0.0
```

--------------------------------

### View TensorBoard Logs

Source: https://schnetpack.readthedocs.io/en/latest/_sources/getstarted.rst.txt

If TensorBoard is installed, you can view the training results by running this command in your terminal. Replace '<rundir>' with the actual run directory.

```bash
$ tensorboard --logdir=<rundir>
```

--------------------------------

### Configure Loss Weights for Energy and Forces

Source: https://schnetpack.readthedocs.io/en/latest/_sources/getstarted.rst.txt

Adjust the loss weights for energy and forces during training. The example shows how to set specific weights for the first and second outputs in the task configuration.

```bash
$ spktrain experiment=md17 data.molecule=uracil task.outputs.0.loss_weight=0.005 \
    task.outputs.1.loss_weight=0.995
```

--------------------------------

### Initialize Environment and Imports

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_02_qm9.ipynb.txt

Imports necessary modules and sets up the working directory for the QM9 tutorial.

```python
import os
import schnetpack as spk
from schnetpack.datasets import QM9
import schnetpack.transform as trn

import torch
import torchmetrics
import pytorch_lightning as pl

qm9tut = "./qm9tut"
if not os.path.exists("qm9tut"):
    os.makedirs(qm9tut)
```

--------------------------------

### Assemble and Configure Simulator

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Initializes the simulator with the system, integrator, and hooks, then moves it to the target device and precision.

```python
# Assemble the simulator
rpmd_simulator = Simulator(
    rpmd_system, rpmd_integrator, md_calculator, simulator_hooks=rpmd_hooks
)

rpmd_simulator = rpmd_simulator.to(md_device)
rpmd_simulator = rpmd_simulator.to(md_precision)
```

--------------------------------

### SchNet Representation Configuration

Source: https://schnetpack.readthedocs.io/en/latest/_sources/userguide/configs.rst.txt

Example structure of a representation subconfig.

```yaml
...
├── model
│   └── representation:
│         _target_: schnetpack.representation.SchNet
│         n_atom_basis: 128
│         n_interactions: 6
│         radial_basis:
│           _target_: schnetpack.nn.radial.GaussianRBF
│           n_rbf: 20
│           cutoff: ${globals.cutoff}
│         cutoff_fn:
│           _target_: schnetpack.nn.cutoff.CosineCutoff
│           cutoff: ${globals.cutoff}
...
```

--------------------------------

### Initialize Configuration Directory

Source: https://schnetpack.readthedocs.io/en/latest/userguide/configs.html

Create the necessary directory structure for storing experiment configuration files.

```bash
$ mkdir configs
$ mkdir configs/experiment
```

--------------------------------

### Manage Results Directory

Source: https://schnetpack.readthedocs.io/en/latest/_sources/howtos/howto_batchwise_relaxations.ipynb.txt

Cleans up existing result directories before starting a new optimization run.

```python
# if os.path.exists(results_dir):
#    shutil.rmtree(results_dir)
```

--------------------------------

### Override Arguments via CLI

Source: https://schnetpack.readthedocs.io/en/latest/_sources/userguide/configs.rst.txt

Examples of overriding configuration values or subconfigs directly from the command line.

```bash
$ spktrain experiment=qm9_atomwise globals.lr=1e-4
```

```bash
$ spktrain experiment=qm9_atomwise model/representation=schnet
```

--------------------------------

### Initializing RPMD System

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html

Sets up a system with multiple replicas for Ring Polymer Molecular Dynamics.

```python
# Number of beads in RPMD simulation
n_replicas = 4

# Set up the system, load structures
rpmd_system = System()
rpmd_system.load_molecules(molecule, n_replicas, position_unit_input="Angstrom")
```

--------------------------------

### Initializing an RPMD System

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Sets up a system with multiple replicas (beads) for Ring Polymer Molecular Dynamics.

```python
# Number of beads in RPMD simulation
n_replicas = 4

# Set up the system, load structures
rpmd_system = System()
rpmd_system.load_molecules(molecule, n_replicas, position_unit_input="Angstrom")

# reuse initializer for momenta
md_initializer.initialize_system(rpmd_system)
```

--------------------------------

### Launch TensorBoard

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_03_force_models.ipynb.txt

Command to visualize training logs generated during the process.

```bash
tensorboard --logdir=forcetut/lightning_logs
```

--------------------------------

### Define Experiment Configuration

Source: https://schnetpack.readthedocs.io/en/latest/_sources/userguide/configs.rst.txt

An example of an experiment configuration file that overrides default model and data settings for a specific dataset.

```yaml
# @package _global_

defaults:
  - override /model: nnp
  - override /data: qm9
```

--------------------------------

### Get Power Spectrum Frequencies and Intensities

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html

Retrieves the computed frequencies and intensities from the PowerSpectrum object. These can then be used for plotting or further analysis.

```python
# Get frequencies and intensities
frequencies, intensities = spectrum.get_spectrum()
```

--------------------------------

### Initialize Simulator

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html

Combines the system, integrator, and calculator into a simulator instance to manage the MD simulation loop.

```python
from schnetpack.md import Simulator

md_simulator = Simulator(md_system, md_integrator, md_calculator)

```

--------------------------------

### Initialize MD System

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Create a system instance and load molecular structures from ASE atoms objects.

```python
from schnetpack.md import System
from ase.io import read

# Load atoms with ASE
molecule = read(molecule_path)

# Number of molecular replicas
n_replicas = 1

# Create system instance and load molecule
md_system = System()
md_system.load_molecules(molecule, n_replicas, position_unit_input="Angstrom")
```

--------------------------------

### Get Total Simulation Steps

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Retrieve the total number of simulation steps performed so far using the `step` attribute of the Simulator object.

```python
print("Total number of steps:", md_simulator.step)
```

--------------------------------

### Set Up Optimizer and Calculator

Source: https://schnetpack.readthedocs.io/en/latest/_sources/howtos/howto_ensemble_calculation.ipynb.txt

Initializes the LBFGS optimizer and assigns the ensemble calculator to the atoms object for structure optimization.

```python
optimizer = LBFGS(atoms)

```

```python
atoms.calc = ensemble_calculator
```

--------------------------------

### Class: md.parsers.OrcaPropertyParser

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/md.parsers.OrcaPropertyParser.html

The OrcaPropertyParser class is a utility for parsing ORCA output files by identifying start and stop markers and applying formatters to the extracted data.

```APIDOC
## Class: md.parsers.OrcaPropertyParser

### Description
Basic property parser for ORCA output files. It collects data entries between a specified start flag and one or more stop flags, optionally applying formatters to the retrieved data.

### Parameters
- **start** (str) - Required - The string that marks the beginning of the data collection.
- **stop** (str | List[str]) - Required - The string or list of strings that mark the end of the data collection.
- **formatters** (OrcaFormatter | List[OrcaFormatter]) - Optional - Formatter(s) used to convert the collected data upon retrieval.
```

--------------------------------

### Load SchNetPack Model and Structure

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Loads a pre-trained SchNetPack model and a molecular structure file (e.g., XYZ format) for use in MD simulations. Requires the `schnetpack` library to be installed.

```python
import schnetpack as spk

# Get the parent directory of SchNetPack
spk_path = os.path.abspath(os.path.join(os.path.dirname(spk.__file__), "../.."))

# Get the path to the test data
test_path = os.path.join(spk_path, "tests/testdata")

# Load model and structure
model_path = os.path.join(test_path, "md_ethanol.model")
molecule_path = os.path.join(test_path, "md_ethanol.xyz")
```

--------------------------------

### Run Training Experiment

Source: https://schnetpack.readthedocs.io/en/latest/userguide/configs.html

Execute a training experiment using the command line interface with a specified configuration.

```bash
$ spktrain experiment=qm9_atomwise
```

--------------------------------

### Creating a SchNetPack Database

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_01_preparing_data.ipynb.txt

Initializes a new SchNetPack database and populates it with the prepared systems and properties.

```python
%rm './new_dataset.db'
new_dataset = ASEAtomsData.create(
    "./new_dataset.db",
    distance_unit="Ang",
    property_unit_dict={"energy": "kcal/mol", "forces": "kcal/mol/Ang"},
)
new_dataset.add_systems(property_list, atoms_list)
```

--------------------------------

### Running the optimization

Source: https://schnetpack.readthedocs.io/en/latest/_sources/howtos/howto_batchwise_relaxations.ipynb.txt

Initializes the ASEBatchwiseLBFGS optimizer and executes the relaxation.

```python
# results_dir = "./howto_batchwise_relaxations_outputs"
# if not os.path.exists(results_dir):
#    os.makedirs(results_dir)

## Initialize optimizer
# optimizer = ASEBatchwiseLBFGS(
#    calculator=calculator,
#    atoms=ats,

```

--------------------------------

### Run Training via CLI

Source: https://schnetpack.readthedocs.io/en/latest/userguide/configs.html

Execute the training process using the specified experiment configuration.

```bash
$ spktrain experiment=my_experiment
```

--------------------------------

### SchNetPack Calculator Configuration for MD

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_05_materials.ipynb.txt

YAML configuration for SchNetPackCalculator used in MD simulations. It specifies required properties, model, units, and a default neighbor list setup with a skin buffer.

```yaml
_target_: schnetpack.md.calculators.SchNetPackCalculator
required_properties:
  - energy
  - forces
model: ???
force_label: forces
energy_units: kcal / mol
position_units: Angstrom
energy_label: energy
stress_label: null
script_model: false

defaults:
  - neighbor_list: 
    _target_: schnetpack.transform.SkinNeighborList
    cutoff_skin: 2.0
    neighbor_list:
      _target_: schnetpack.transform.ASENeighborList
      cutoff: 5.0
    nbh_transforms:
      - _target_: schnetpack.transform.FilteredNeighbors
        selection_name: filtered_out_neighbors

```

--------------------------------

### Initialize System Momenta

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Assign initial momenta to the system using a uniform distribution at a specified temperature.

```python
from schnetpack.md import UniformInit

system_temperature = 300  # Kelvin

# Set up the initializer
md_initializer = UniformInit(
    system_temperature,
    remove_center_of_mass=True,
    remove_translation=True,
    remove_rotation=True,
)

# Initialize the system momenta
md_initializer.initialize_system(md_system)
```

--------------------------------

### Perform Inference on a Batch of Molecules

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_02_qm9.ipynb.txt

Iterate through the test dataloader to get a batch of molecules and pass it to the loaded SchNetPack model to obtain predictions. The result is a dictionary containing various prediction outputs.

```python
for batch in qm9data.test_dataloader():
    result = best_model(batch)
    print("Result dictionary:", result)
    break
```

--------------------------------

### Import Dependencies and Initialize Environment

Source: https://schnetpack.readthedocs.io/en/latest/_sources/howtos/howto_ensemble_calculation.ipynb.txt

Imports necessary libraries for numerical operations, visualization, ASE simulation tools, and SchNetPack ensemble components.

```python
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt

from ase import units
from ase.io import read
from ase.optimize.lbfgs import LBFGS
from ase.md.velocitydistribution import MaxwellBoltzmannDistribution
from ase.md.langevin import Langevin

from schnetpack.interfaces.ase_interface import SpkEnsembleCalculator, AbsoluteUncertainty, RelativeUncertainty
import schnetpack.transform as trn
from schnetpack.datasets import MD17
import torch


np.random.seed(42)
```

--------------------------------

### Convert ASE Atoms and Predict with SchNetPack

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_02_qm9.ipynb.txt

Convert the ASE Atoms object into SchNetPack input format using the initialized converter. Then, pass the converted input to the loaded SchNetPack model to get predictions, such as the U0 energy.

```python
inputs = converter(atoms)

print("Keys:", list(inputs.keys()))

pred = best_model(inputs)

print("Prediction:", pred[QM9.U0])
```

--------------------------------

### Initialize Simulator for MD

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Create a Simulator instance to perform MD simulations. It takes the molecular system, integrator, and calculator as arguments.

```python
from schnetpack.md import Simulator

md_simulator = Simulator(md_system, md_integrator, md_calculator)
```

--------------------------------

### Initialize RPMD Integrator

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Sets up the RingPolymer integrator with specified time step, number of replicas, and system temperature.

```python
# Here, we set it to the system temperature
rpmd_integrator = RingPolymer(
    rpmd_time_step,
    n_replicas,
    system_temperature,
)
```

--------------------------------

### Base Transform Class

Source: https://schnetpack.readthedocs.io/en/latest/_modules/transform/base.html

Base class for all transforms. Ensures initialization of data and datamodule attributes. Override the `forward` method for custom logic. Preprocessors act on single examples, postprocessors on batches. All transforms must return a modified `inputs` dictionary.

```python
from typing import Optional, Dict

import torch
import torch.nn as nn

import schnetpack as spk

__all__ = [
    "Transform",
    "TransformException",
]


class TransformException(Exception):
    pass



class Transform(nn.Module):
    """
    Base class for all transforms.
    The base class ensures that the reference to the data and datamodule attributes are
    initialized.
    Transforms can be used as pre- or post-processing layers.
    They can also be used for other parts of a model, that need to be
    initialized based on data.

    To implement a new transform, override the forward method. Preprocessors are applied
    to single examples, while postprocessors operate on batches. All transforms should
    return a modified `inputs` dictionary.

    """

    def datamodule(self, value):
        """
        Extract all required information from data module automatically when using
        PyTorch Lightning integration. The transform should also implement a way to
        set these things manually, to make it usable independent of PL.

        Do not store the datamodule, as this does not work with torchscript conversion!
        """
        pass

    def forward(
        self,
        inputs: Dict[str, torch.Tensor],
    ) -> Dict[str, torch.Tensor]:
        raise NotImplementedError

    def teardown(self):
        pass

```

--------------------------------

### Create Experiment Configuration

Source: https://schnetpack.readthedocs.io/en/latest/userguide/configs.html

Create and edit a new experiment configuration file using a text editor.

```bash
$ vim my_configs/experiments/my_experiment.yaml
```

--------------------------------

### Configure File Logger

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Sets up a FileLogger to store simulation data like positions, momenta, and energy in an HDF5 database.

```python
log_file = os.path.join(md_workdir, "simulation.hdf5")

# Size of the buffer
buffer_size = 100

# Set up data streams to store positions, momenta and the energy
data_streams = [
    callback_hooks.MoleculeStream(store_velocities=True),
    callback_hooks.PropertyStream(target_properties=[properties.energy]),
]

# Create the file logger
file_logger = callback_hooks.FileLogger(
    log_file,
    buffer_size,
    data_streams=data_streams,
    every_n_steps=1,  # logging frequency
    precision=32,  # floating point precision used in hdf5 database
)

# Update the simulation hooks
simulation_hooks.append(file_logger)
```

--------------------------------

### md.initial_conditions.UniformInit

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/md.initial_conditions.UniformInit.html

Initializes the system momenta according to a uniform distribution scaled to the given temperature.

```APIDOC
## md.initial_conditions.UniformInit

### Description
Initializes the system momenta according to a uniform distribution scaled to the given temperature.

### Method
Constructor

### Endpoint
N/A (Class definition)

### Parameters
#### Path Parameters
N/A

#### Query Parameters
N/A

#### Request Body
N/A

### Request Example
N/A

### Response
#### Success Response (200)
N/A

#### Response Example
N/A

### Parameters
- **temperature** (float | List[float]) - Target temperature in Kelvin.
- **remove_center_of_mass** (bool) - Whether to remove the center of mass motion (default: True).
- **remove_translation** (bool) - Remove the translational components of the momenta after initialization. Will stop molecular drift for NVE simulations and NVT simulations with deterministic thermostat (default: False).
- **remove_rotation** (bool) - Remove the rotational components of the momenta after initialization. Will reduce molecular rotation for NVE simulations and NVT simulations with deterministic thermostat (default: False).
- **wrap_positions** (bool) - Wrap atom positions back to box when using periodic boundary conditions.
```

--------------------------------

### Initialize SchNetPackCalculator

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Instantiate a SchNetPackCalculator for MD simulations. Requires model path, force key, energy and length units, and neighbor list. Optional parameters include energy_key and required_properties.

```python
md_calculator = SchNetPackCalculator(
    model_path,  # path to stored model
    "forces",  # force key
    "kcal/mol",  # energy units
    "Angstrom",  # length units
    md_neighborlist,  # neighbor list
    energy_key="energy",  # name of potential energies
    required_properties=[],  # additional properties extracted from the model
)
```

--------------------------------

### Initialize Langevin Thermostat

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Configures a Langevin thermostat with a specified bath temperature and time constant for NVT simulations.

```python
from schnetpack.md.simulation_hooks import LangevinThermostat

# Set temperature and thermostat constant
bath_temperature = 300  # K
time_constant = 100  # fs

# Initialize the thermostat
langevin = LangevinThermostat(bath_temperature, time_constant)
```

--------------------------------

### Initialize Simulator with Hooks

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Constructs the MD simulator using the configured system, integrator, calculator, and accumulated simulation hooks.

```python
md_simulator = Simulator(
    md_system, md_integrator, md_calculator, simulator_hooks=simulation_hooks
)

md_simulator = md_simulator.to(md_precision)
md_simulator = md_simulator.to(md_device)
```

--------------------------------

### Initialize SchNetPack Calculator

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Imports necessary modules for setting up the SchNetPack calculator interface.

```python
from schnetpack.md.calculators import SchNetPackCalculator
from schnetpack import properties


```

--------------------------------

### Class: md.initial_conditions.MaxwellBoltzmannInit

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/md.initial_conditions.MaxwellBoltzmannInit.html

Initializes system momenta according to a Maxwell-Boltzmann distribution at a specified temperature.

```APIDOC
## Class: md.initial_conditions.MaxwellBoltzmannInit

### Description
Initializes the system momenta according to a Maxwell–Boltzmann distribution at the given temperature.

### Parameters
- **temperature** (float | List[float]) - Required - Target temperature in Kelvin.
- **remove_center_of_mass** (bool) - Optional - Default: True.
- **remove_translation** (bool) - Optional - Default: True. Remove the translational components of the momenta after initialization.
- **remove_rotation** (bool) - Optional - Default: False. Remove the rotational components of the momenta after initialization.
- **wrap_positions** (bool) - Optional - Default: False. Wrap atom positions back to box when using periodic boundary conditions.
```

--------------------------------

### Initialize FileLogger for Data Collection

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html

Sets up a FileLogger to collect molecule positions, velocities, and predicted energies during simulation. Configure buffer size, logging frequency, and floating-point precision.

```python
from schnetpack.md.simulation_hooks import callback_hooks

# Path to database
log_file = os.path.join(md_workdir, "simulation.hdf5")

# Size of the buffer
buffer_size = 100

# Set up data streams to store positions, momenta and the energy
data_streams = [
    callback_hooks.MoleculeStream(store_velocities=True),
    callback_hooks.PropertyStream(target_properties=[properties.energy]),
]

# Create the file logger
file_logger = callback_hooks.FileLogger(
    log_file,
    buffer_size,
    data_streams=data_streams,
    every_n_steps=1,  # logging frequency
    precision=32,  # floating point precision used in hdf5 database
)

# Update the simulation hooks
simulation_hooks.append(file_logger)
```

--------------------------------

### Configure SchNetPackCalculator

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html

Initializes the calculator by providing the model path, unit definitions, and the previously configured neighbor list.

```python
from schnetpack.md.calculators import SchNetPackCalculator
from schnetpack import properties

md_calculator = SchNetPackCalculator(
    model_path,  # path to stored model
    "forces",  # force key
    "kcal/mol",  # energy units
    "Angstrom",  # length units
    md_neighborlist,  # neighbor list
    energy_key="energy",  # name of potential energies
    required_properties=[],  # additional properties extracted from the model
)

```

--------------------------------

### Initialize SpkCalculator with Model and Neighbor List

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_02_qm9.ipynb.txt

Instantiate the SpkCalculator by providing the path to a trained model, a neighbor list configuration, and property details. Specify the energy key, unit, and computation device.

```python
calculator = spk.interfaces.SpkCalculator(
    model=os.path.join(qm9tut, "best_inference_model"),  # path to model
    neighbor_list=trn.ASENeighborList(cutoff=5.0),  # neighbor list
    energy_key=QM9.U0,  # name of energy property in model
    force_key=None,
    energy_unit="eV",  # units of energy property
    device="cpu",  # device for computation
)
atoms.calc = calculator
print("Prediction:", atoms.get_total_energy())
```

--------------------------------

### View TensorBoard Logs

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Command to launch the TensorBoard interface for visualizing logged simulation data.

```bash
tensorboard --logdir=mdtut/logs
```

--------------------------------

### Initialize RPMD System

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_04_molecular_dynamics.html

Reuse the initializer for momenta on the RPMD system.

```python
md_initializer.initialize_system(rpmd_system)
```

--------------------------------

### Initialize EnergyCoulomb

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/atomistic.EnergyCoulomb.html

Instantiate the EnergyCoulomb module with specified units, potential function, and keys for input and output.

```python
EnergyCoulomb(
        energy_unit=u.eV,
        position_unit=u.angstrom,
        coulomb_potential=DampedCoulombPotential(cutoff=5.0 * u.angstrom),
        output_key="energy",
        charges_key="charges",
        use_neighbors_lr=True,
        cutoff=None,
    )
```

--------------------------------

### Load and Configure QM9 Dataset

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_02_qm9.ipynb.txt

Configures the QM9 dataset with batching, neighbor lists, and property offsets. The split file is managed via the local directory.

```python
%rm split.npz

qm9data = QM9(
    "./qm9.db",
    batch_size=100,
    num_train=1000,
    num_val=1000,
    transforms=[
        trn.ASENeighborList(cutoff=5.0),
        trn.RemoveOffsets(QM9.U0, remove_mean=True, remove_atomrefs=True),
        trn.CastTo32(),
    ],
    property_units={QM9.U0: "eV"},
    num_workers=1,
    split_file=os.path.join(qm9tut, "split.npz"),
    pin_memory=True,  # set to false, when not using a GPU
    load_properties=[QM9.U0],  # only load U0 property
)
qm9data.prepare_data()
qm9data.setup()
```

--------------------------------

### Configure and Run Optimization

Source: https://schnetpack.readthedocs.io/en/latest/_sources/howtos/howto_batchwise_relaxations.ipynb.txt

Defines the trajectory output path and executes the optimization process with specified force convergence criteria.

```python
#    trajectory="./howto_batchwise_relaxations_outputs/relax_traj",
# )

## run optimization
# optimizer.run(fmax=0.0005, steps=1000)
```

--------------------------------

### Configure Simulation Hooks

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Sets up logging, checkpointing, and TensorBoard monitoring for the RPMD simulation.

```python
# Logging
rpmd_log_file = os.path.join(md_workdir, "rpmd_simulation.hdf5")
rpmd_file_logger = callback_hooks.FileLogger(
    rpmd_log_file,
    buffer_size,
    data_streams=[
        callback_hooks.MoleculeStream(
            store_velocities=True,
        ),
        callback_hooks.PropertyStream(["energy"]),
    ],
)

# Checkpoints
rpmd_chk_file = os.path.join(md_workdir, "rpmd_simulation.chk")
rpmd_checkpoint = callback_hooks.Checkpoint(rpmd_chk_file, every_n_steps=100)

# Tensorboard
rpmd_tensorboard_dir = os.path.join(md_workdir, "rpmd_logs")
rpmd_tensorboard_logger = callback_hooks.TensorBoardLogger(
    rpmd_tensorboard_dir,
    ["energy", "temperature"],
)

# Assemble the hooks:
rpmd_hooks = [
    pile,
    rpmd_file_logger,
    rpmd_checkpoint,
    rpmd_tensorboard_logger,
]
```

--------------------------------

### Initialize VibrationalSpectrum

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/md.data.VibrationalSpectrum.html

Instantiate the VibrationalSpectrum class with a dataset, resolution, and an optional window function. The 'data' parameter must be an HDF5Loader object.

```python
VibrationalSpectrum(data: ~schnetpack.md.data.hdf5_data.HDF5Loader, resolution: int = 4096, window: callable = <function cosine_sq_window>)
```

--------------------------------

### Initializing the BatchwiseCalculator

Source: https://schnetpack.readthedocs.io/en/latest/_sources/howtos/howto_batchwise_relaxations.ipynb.txt

Configures the model, neighbor list, and atoms converter to prepare the calculator for force and energy calculations.

```python
# model_path = "../../tests/testdata/md_ethanol.model"

## set device
# device = torch.device("cpu")

## load model
# model = torch.load(model_path, map_location=device, weights_only=False)

## define neighbor list
# cutoff = model.representation.cutoff.item()
# nbh_list=spk.transform.MatScipyNeighborList(cutoff=cutoff)

## build atoms converter
# atoms_converter = AtomsConverter(
#    neighbor_list=nbh_list,
#    device=device,
# )

## build calculator
# calculator = BatchwiseCalculator(
#    model=model_path,
#    atoms_converter=atoms_converter,
#    device=device,
#    energy_unit="kcal/mol",
#    position_unit="Ang",
# )
```

--------------------------------

### Importing necessary modules

Source: https://schnetpack.readthedocs.io/en/latest/_sources/howtos/howto_batchwise_relaxations.ipynb.txt

Imports required for setting up the batch-wise optimization environment.

```python
# import os
# import shutil

# import torch
# from ase.io import read

# import schnetpack as spk
# from schnetpack import properties
# from schnetpack.interfaces.ase_interface import AtomsConverter
# from schnetpack.interfaces.batchwise_optimization import ASEBatchwiseLBFGS, BatchwiseCalculator
```

--------------------------------

### Import ASEAtomsData

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_01_preparing_data.ipynb.txt

Import the class used for interacting with ASE database files.

```python
from schnetpack.data import ASEAtomsData
```

--------------------------------

### Configure TensorBoard Logger

Source: https://schnetpack.readthedocs.io/en/latest/_sources/tutorials/tutorial_04_molecular_dynamics.ipynb.txt

Sets up a TensorBoardLogger to monitor specific simulation properties like energy and temperature.

```python
# directory where tensorboard log will be stored to
tensorboard_dir = os.path.join(md_workdir, "logs")

tensorboard_logger = callback_hooks.TensorBoardLogger(
    tensorboard_dir,
    ["energy", "temperature"],  # properties to log
)

# update simulation hooks
simulation_hooks.append(tensorboard_logger)
```

--------------------------------

### NPTRingPolymer Class Initialization

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/md.integrators.NPTRingPolymer.html

Initializes the NPTRingPolymer integrator for constant pressure dynamics. Requires time step, number of beads, temperature, and a barostat.

```python
class _md.integrators.NPTRingPolymer(_* args: Any_, _** kwargs: Any_)
```

--------------------------------

### Class: datasets.MD17

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/datasets.MD17.html

Initializes the MD17 benchmark dataset for molecular dynamics of small molecules containing molecular forces.

```APIDOC
## Class: datasets.MD17

### Description
MD17 benchmark data set for molecular dynamics of small molecules containing molecular forces.

### Parameters
- **datapath** (string) - Required - Path to dataset
- **batch_size** (int) - Optional - (train) batch size
- **num_train** (int) - Optional - Number of training examples
- **num_val** (int) - Optional - Number of validation examples
- **num_test** (int) - Optional - Number of test examples
- **split_file** (string) - Optional - Path to npz file with data partitions
- **format** (string) - Optional - Dataset format
- **load_properties** (list) - Optional - Subset of properties to load
- **val_batch_size** (int) - Optional - Validation batch size
- **test_batch_size** (int) - Optional - Test batch size
- **transforms** (callable) - Optional - Transform applied to each system separately before batching
- **train_transforms** (callable) - Optional - Overrides transform_fn for training
- **val_transforms** (callable) - Optional - Overrides transform_fn for validation
- **test_transforms** (callable) - Optional - Overrides transform_fn for testing
- **num_workers** (int) - Optional - Number of data loader workers
- **num_val_workers** (int) - Optional - Number of validation data loader workers
- **num_test_workers** (int) - Optional - Number of test data loader workers
- **distance_unit** (string) - Optional - Unit of the atom positions and cell
- **data_workdir** (string) - Optional - Copy data here as part of setup
```

--------------------------------

### Initialize OrcaMainFileParser

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/md.parsers.OrcaMainFileParser.html

Instantiate the parser with an optional list of target properties to extract from the ORCA output.

```python
class md.parsers.OrcaMainFileParser(target_properties : List[str] | None = None)
```

--------------------------------

### Create ASE Atoms Object with Periodic Boundary Conditions

Source: https://schnetpack.readthedocs.io/en/latest/tutorials/tutorial_05_materials.html

Use this to create an ASE Atoms object for systems with periodic boundary conditions. Ensure `pbc=True` is set during initialization.

```python
from ase import Atoms

atoms = Atoms(
    symbols=symbols,
    positions=positions,
    cell=cell,
    pbc=True,
)
data = dict(
    energy=np.array(energy),
    forces=np.array(forces),
    stress=np.array(stress),
    considered_atoms=considered_atoms,
    filtered_out_neighbors=filtered_out_neighbors,
)
```

--------------------------------

### Class: datasets.ISO17

Source: https://schnetpack.readthedocs.io/en/latest/api/generated/datasets.ISO17.html

Initializes the ISO17 benchmark dataset for molecular dynamics simulations.

```APIDOC
## Class datasets.ISO17

### Description
ISO17 benchmark data set for molecular dynamics of C7O2H10 isomers containing molecular forces.

### Parameters
- **datapath** (str) - Required - Path to dataset
- **fold** (str) - Optional - Select a specific dataset of iso17
- **batch_size** (int) - Optional - (train) batch size
- **num_train** (int) - Optional - Number of training examples
- **num_val** (int) - Optional - Number of validation examples
- **num_test** (int) - Optional - Number of test examples
- **split_file** (str) - Optional - Path to npz file with data partitions
- **format** (str) - Optional - Dataset format
- **load_properties** (list) - Optional - Subset of properties to load
- **val_batch_size** (int) - Optional - Validation batch size
- **test_batch_size** (int) - Optional - Test batch size
- **transforms** (list) - Optional - Transform applied to each system separately before batching
- **train_transforms** (list) - Optional - Overrides transform_fn for training
- **val_transforms** (list) - Optional - Overrides transform_fn for validation
- **test_transforms** (list) - Optional - Overrides transform_fn for testing
- **num_workers** (int) - Optional - Number of data loader workers
- **num_val_workers** (int) - Optional - Number of validation data loader workers
- **num_test_workers** (int) - Optional - Number of test data loader workers
- **distance_unit** (str) - Optional - Unit of the atom positions and cell
```