atomistic-machine-learning · stefaanhessmann · Jul 23, 2026 · Feb 21, 2026 · Feb 22, 2026 · Feb 22, 2026
diff --git a/.gitignore b/.gitignore
@@ -128,4 +128,4 @@ interfaces/lammps/examples/*/*.dat
 interfaces/lammps/examples/*/deployed_model
 
 # batchwise optimizer examples
-examples/howtos/howto_batchwise_relaxations_outputs/*
+examples/howtos/howto_batchwise_relaxations_outputs/*
diff --git a/docs/api/data.rst b/docs/api/data.rst
@@ -10,24 +10,11 @@ Atoms data
     :recursive:
     :template: classtemplate.rst
 
-    BaseAtomsData
     ASEAtomsData
+    DownloadableASEAtomsData
     AtomsLoader
-    resolve_format
-    AtomsDataFormat
     StratifiedSampler
 
-
-Creation
---------
-.. autosummary::
-    :toctree: generated
-    :nosignatures:
-    :template: classtemplate.rst
-
-    create_dataset
-    load_dataset
-
 Data modules
 ------------
 
@@ -47,5 +34,7 @@ Statistics
     :template: classtemplate.rst
 
     calculate_stats
+    estimate_atomrefs
+    StatsAtomrefProvider
     NumberOfAtomsCriterion
     PropertyCriterion
diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst
@@ -16,7 +16,6 @@ Molecules
     ISO17
     MD22
     QM7X
-    TMQM
 
 Materials
 ---------

diff --git a/docs/userguide/overview.rst b/docs/userguide/overview.rst
@@ -31,8 +31,11 @@ and calculation of neighbor lists.
 
 Furthermore, we support PyTorch Lightning datamodules through :class:`AtomsDataModule`,
 which combines :class:`ASEAtomsData` with code for preparation, setup and partitioning
-into train/validation/test splits. We provide specific implementations of
-:class:`AtomsDataModule` for several benchmark datasets.
+into train/validation/test splits. The datamodule also owns the training statistics
+(per-property mean/std and atom references): it exposes them via ``get_stats`` and
+``get_atomrefs``, initializes all transforms that require them, and persists computed
+values next to the split file so reruns read instead of recompute. We provide specific
+implementations of :class:`AtomsDataModule` for several benchmark datasets.
 
 
 Model

diff --git a/examples/howtos/howto_ensemble_calculation.ipynb b/examples/howtos/howto_ensemble_calculation.ipynb
@@ -42,7 +42,11 @@
     "from ase.md.velocitydistribution import MaxwellBoltzmannDistribution\n",
     "from ase.md.langevin import Langevin\n",
     "\n",
-    "from schnetpack.interfaces.ase_interface import SpkEnsembleCalculator, AbsoluteUncertainty, RelativeUncertainty\n",
+    "from schnetpack.interfaces.ase_interface import (\n",
+    "    SpkEnsembleCalculator,\n",
+    "    AbsoluteUncertainty,\n",
+    "    RelativeUncertainty,\n",
+    ")\n",
     "import schnetpack.transform as trn\n",
     "from schnetpack.datasets import MD17\n",
     "import torch\n",
@@ -78,11 +82,13 @@
    },
    "outputs": [],
    "source": [
-    "model_path_list = ['../trained_models/rmd17_ethanol/painn_1/best_model',\n",
-    "                   '../trained_models/rmd17_ethanol/painn_2/best_model',\n",
-    "                   '../trained_models/rmd17_ethanol/painn_3/best_model',\n",
-    "                   '../trained_models/rmd17_ethanol/painn_4/best_model',\n",
-    "                   '../trained_models/rmd17_ethanol/painn_5/best_model']"
+    "model_path_list = [\n",
+    "    \"../trained_models/rmd17_ethanol/painn_1/best_model\",\n",
+    "    \"../trained_models/rmd17_ethanol/painn_2/best_model\",\n",
+    "    \"../trained_models/rmd17_ethanol/painn_3/best_model\",\n",
+    "    \"../trained_models/rmd17_ethanol/painn_4/best_model\",\n",
+    "    \"../trained_models/rmd17_ethanol/painn_5/best_model\",\n",
+    "]"
    ]
   },
   {
@@ -113,7 +119,7 @@
    },
    "outputs": [],
    "source": [
-    "uncertainty_abs = AbsoluteUncertainty(energy_weight=0.5,force_weight=1.0)\n",
+    "uncertainty_abs = AbsoluteUncertainty(energy_weight=0.5, force_weight=1.0)\n",
     "uncertainty_rel = RelativeUncertainty(energy_weight=1.0, force_weight=2.0)\n",
     "\n",
     "uncertainty = [uncertainty_abs, uncertainty_rel]\n",
@@ -125,7 +131,8 @@
     "    force_key=MD17.forces,\n",
     "    energy_unit=\"kcal/mol\",\n",
     "    position_unit=\"Ang\",\n",
-    "    uncertainty_fn=uncertainty)"
+    "    uncertainty_fn=uncertainty,\n",
+    ")"
    ]
   },
   {
@@ -146,8 +153,8 @@
    },
    "outputs": [],
    "source": [
-    "#load data into atoms object\n",
-    "atoms = read('../../tests/testdata/md_ethanol.xyz', index=0)\n",
+    "# load data into atoms object\n",
+    "atoms = read(\"../../tests/testdata/md_ethanol.xyz\", index=0)\n",
     "# specify atoms calculator\n",
     "atoms.calc = ensemble_calculator"
    ]
@@ -290,22 +297,22 @@
     "fig, ax1 = plt.subplots(figsize=(8, 6))\n",
     "\n",
     "# Plot absolute uncertainty on left y-axis\n",
-    "ax1.plot(steps, abs_vals, label=\"Absolute Uncertainty\", marker='o', color='tab:blue')\n",
+    "ax1.plot(steps, abs_vals, label=\"Absolute Uncertainty\", marker=\"o\", color=\"tab:blue\")\n",
     "ax1.set_xlabel(\"Optimization Step\")\n",
-    "ax1.set_ylabel(\"Absolute Uncertainty\", color='tab:blue')\n",
-    "ax1.tick_params(axis='y', labelcolor='tab:blue')\n",
+    "ax1.set_ylabel(\"Absolute Uncertainty\", color=\"tab:blue\")\n",
+    "ax1.tick_params(axis=\"y\", labelcolor=\"tab:blue\")\n",
     "ax1.grid(True)\n",
     "\n",
     "# Create second y-axis for relative uncertainty\n",
     "ax2 = ax1.twinx()\n",
-    "ax2.plot(steps, rel_vals, label=\"Relative Uncertainty\", marker='x', color='tab:red')\n",
-    "ax2.set_ylabel(\"Relative Uncertainty\", color='tab:red')\n",
-    "ax2.tick_params(axis='y', labelcolor='tab:red')\n",
+    "ax2.plot(steps, rel_vals, label=\"Relative Uncertainty\", marker=\"x\", color=\"tab:red\")\n",
+    "ax2.set_ylabel(\"Relative Uncertainty\", color=\"tab:red\")\n",
+    "ax2.tick_params(axis=\"y\", labelcolor=\"tab:red\")\n",
     "\n",
     "# Title and layout\n",
     "plt.title(\"Uncertainty during Optimization\")\n",
     "fig.tight_layout()\n",
-    "plt.show()\n"
+    "plt.show()"
    ]
   },
   {
@@ -344,7 +351,7 @@
    },
    "outputs": [],
    "source": [
-    "uncertainty_abs = AbsoluteUncertainty(energy_weight=0.5,force_weight=1.0)\n",
+    "uncertainty_abs = AbsoluteUncertainty(energy_weight=0.5, force_weight=1.0)\n",
     "\n",
     "abs_ensemble_calculator = SpkEnsembleCalculator(\n",
     "    models=model_path_list,\n",
@@ -353,7 +360,8 @@
     "    force_key=MD17.forces,\n",
     "    energy_unit=\"kcal/mol\",\n",
     "    position_unit=\"Ang\",\n",
-    "    uncertainty_fn=uncertainty_abs)"
+    "    uncertainty_fn=uncertainty_abs,\n",
+    ")"
    ]
   },
   {
@@ -367,13 +375,13 @@
    },
    "outputs": [],
    "source": [
-    "target_temperatures = [_ for _ in range(50, 800, 100)]  \n",
-    "n_steps = 1000          \n",
-    "sampling_interval = 10  \n",
-    "step_size = 0.5   \n",
+    "target_temperatures = [_ for _ in range(50, 800, 100)]\n",
+    "n_steps = 1000\n",
+    "sampling_interval = 10\n",
+    "step_size = 0.5\n",
     "\n",
     "# setting up initial atoms\n",
-    "atoms = read('../../tests/testdata/md_ethanol.xyz', index=0)\n",
+    "atoms = read(\"../../tests/testdata/md_ethanol.xyz\", index=0)\n",
     "atoms.calc = abs_ensemble_calculator\n",
     "\n",
     "MaxwellBoltzmannDistribution(atoms, temperature_K=target_temperatures[0])\n",
@@ -385,19 +393,19 @@
     "for target_temperature in target_temperatures:\n",
     "    print(f\"Temp: {target_temperature:.2f} K\")\n",
     "    for step in tqdm(range(n_steps // sampling_interval)):\n",
-    "    \n",
+    "\n",
     "        dyn = Langevin(\n",
-    "            atoms, \n",
-    "            timestep=step_size * units.fs, \n",
+    "            atoms,\n",
+    "            timestep=step_size * units.fs,\n",
     "            temperature_K=target_temperature,\n",
-    "            friction=0.01 / units.fs\n",
+    "            friction=0.01 / units.fs,\n",
     "        )\n",
-    "    \n",
+    "\n",
     "        dyn.run(sampling_interval)\n",
-    "    \n",
+    "\n",
     "        temp.append(atoms.get_temperature())\n",
     "        uncertainties.append(abs_ensemble_calculator.get_uncertainty(atoms))\n",
-    "        \n",
+    "\n",
     "        ats_traj.append(atoms.copy())"
    ]
   },
@@ -409,22 +417,22 @@
    "source": [
     "fig, ax1 = plt.subplots(figsize=(8, 6))\n",
     "\n",
-    "ax1.plot(uncertainties, marker='o', color='blue', label='Uncertainty')\n",
+    "ax1.plot(uncertainties, marker=\"o\", color=\"blue\", label=\"Uncertainty\")\n",
     "ax1.set_xlabel(\"MD Step\")\n",
-    "ax1.set_ylabel(\"Uncertainty\", color='blue')\n",
-    "ax1.tick_params(axis='y', labelcolor='blue')\n",
+    "ax1.set_ylabel(\"Uncertainty\", color=\"blue\")\n",
+    "ax1.tick_params(axis=\"y\", labelcolor=\"blue\")\n",
     "\n",
     "ax2 = ax1.twinx()\n",
-    "ax2.plot(temp, marker='x', color='red', label='Temperature')\n",
-    "ax2.set_ylabel(\"Temperature (K)\", color='red')\n",
-    "ax2.tick_params(axis='y', labelcolor='red')\n",
+    "ax2.plot(temp, marker=\"x\", color=\"red\", label=\"Temperature\")\n",
+    "ax2.set_ylabel(\"Temperature (K)\", color=\"red\")\n",
+    "ax2.tick_params(axis=\"y\", labelcolor=\"red\")\n",
     "\n",
     "plt.title(\"Molecular Dynamics: Uncertainty and Temperature Profile\")\n",
     "ax1.grid(True)\n",
     "\n",
     "lines_1, labels_1 = ax1.get_legend_handles_labels()\n",
     "lines_2, labels_2 = ax2.get_legend_handles_labels()\n",
-    "ax1.legend(lines_1 + lines_2, labels_1 + labels_2, loc='upper right')\n",
+    "ax1.legend(lines_1 + lines_2, labels_1 + labels_2, loc=\"upper right\")\n",
     "\n",
     "plt.tight_layout()\n",
     "plt.show()"
@@ -444,6 +452,7 @@
    "outputs": [],
    "source": [
     "from ase.visualize import view\n",
+    "\n",
     "view(ats_traj)"
    ]
   }

diff --git a/examples/tutorials/tutorial_02_qm9.ipynb b/examples/tutorials/tutorial_02_qm9.ipynb
@@ -377,9 +377,7 @@
     "from ase import Atoms\n",
     "from schnetpack.utils.compatibility import load_model\n",
     "\n",
-    "best_model = load_model(\n",
-    "    os.path.join(qm9tut, \"best_inference_model\"), device=\"cpu\"\n",
-    ")"
+    "best_model = load_model(os.path.join(qm9tut, \"best_inference_model\"), device=\"cpu\")"
    ]
   },
   {

diff --git a/src/schnetpack/__init__.py b/src/schnetpack/__init__.py
@@ -16,5 +16,4 @@
 from schnetpack.task import *
 from schnetpack import md
 
-
 __version__ = "2.2.0"
diff --git a/src/schnetpack/cli.py b/src/schnetpack/cli.py
@@ -17,11 +17,10 @@
 import schnetpack as spk
 from schnetpack.utils import str2class
 from schnetpack.utils.script import log_hyperparameters, print_config
-from schnetpack.data import BaseAtomsData, AtomsLoader
+from schnetpack.data import ASEAtomsData, AtomsLoader
 from schnetpack.train import PredictionWriter
 from schnetpack import properties
-from schnetpack.utils import load_model
-
+from schnetpack.utils import load_model, load_task_from_checkpoint
 
 log = logging.getLogger(__name__)
 
@@ -176,16 +175,20 @@ def train(config: DictConfig):
     log.info("Starting training.")
     trainer.fit(model=task, datamodule=datamodule, ckpt_path=config.run.ckpt_path)
 
-    # Evaluate model on test set after training
-    log.info("Starting testing.")
-    trainer.test(model=task, datamodule=datamodule, ckpt_path="best")
-
-    # Store best model
+    # Load the best checkpoint through the compatibility helper (it handles
+    # `weights_only` across PL versions) and test that task directly, instead
+    # of having Lightning re-load the checkpoint internally via
+    # ckpt_path="best", whose weights_only handling is version-dependent.
     best_path = trainer.checkpoint_callback.best_model_path
     log.info(f"Best checkpoint path:\n{best_path}")
+    best_task = load_task_from_checkpoint(type(task), best_path)
 
+    # Evaluate best model on test set after training
+    log.info("Starting testing.")
+    trainer.test(model=best_task, datamodule=datamodule)
+
+    # Store best model
     log.info(f"Store best model")
-    best_task = type(task).load_from_checkpoint(best_path)
     torch.save(best_task, config.globals.model_path + ".task")
 
     best_task.save_model(config.globals.model_path, do_postprocessing=True)
@@ -195,7 +198,7 @@ def train(config: DictConfig):
 @hydra.main(config_path="configs", config_name="predict", version_base="1.2")
 def predict(config: DictConfig):
     log.info(f"Load data from `{config.data.datapath}`")
-    dataset: BaseAtomsData = hydra.utils.instantiate(config.data)
+    dataset: ASEAtomsData = hydra.utils.instantiate(config.data)
     loader = AtomsLoader(dataset, batch_size=config.batch_size, num_workers=8)
 
     model = load_model("best_model")

diff --git a/src/schnetpack/configs/data/ani1.yaml b/src/schnetpack/configs/data/ani1.yaml
@@ -1,16 +1,16 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.ANI1
+dataset:
+  _target_: schnetpack.datasets.ANI1
+  datapath: ${run.data_dir}/ani1.db # data_dir is specified in train.yaml
+  num_heavy_atoms: 8
+  high_energies: false
+  distance_unit: Ang
+  property_units:
+    energy: eV
 
-datapath: ${run.data_dir}/ani1.db  # data_dir is specified in train.yaml
 batch_size: 32
 num_train: 10000000
 num_val: 100000
-num_heavy_atoms: 8
-high_energies: False
-
-# convert to typically used units
-distance_unit: Ang
-property_units:
-  energy: eV
diff --git a/src/schnetpack/configs/data/custom.yaml b/src/schnetpack/configs/data/custom.yaml
@@ -1,12 +1,26 @@
+# @package data
 _target_: schnetpack.data.AtomsDataModule
 
-datapath: ???
-data_workdir: null
+dataset:
+  _target_: schnetpack.data.ASEAtomsData
+  datapath: ???
+  load_properties: null
+  # null: keep the dataset's native units — setting a unit here converts on
+  # load.
+  distance_unit: null
+  property_units: null
+  transforms: null
+  train_transforms: null
+  val_transforms: null
+  test_transforms: null
+
 batch_size: 10
 num_train: ???
 num_val: ???
 num_test: null
+split_file: ${run.data_dir}/split.npz
+splitting: null
 num_workers: 8
-num_val_workers: null
-num_test_workers: null
-train_sampler_cls: null
+train_sampler_cls: null
+train_sampler_args: {}
+pin_memory: false
-Original file line number
+Diff line change
@@ Expand Up / @@ -16,7 +16,6 @@ Molecules @@
         ISO17
         MD22
         QM7X
-        TMQM
     Materials
     ---------
@@ Expand Down @@
Original file line number	Diff line number	Diff line change
Expand Up		@@ -16,5 +16,4 @@
		from schnetpack.task import *
		from schnetpack import md


		__version__ = "2.2.0"