Add tutorial

Author: frostedoyster

docs/generate_examples/conf.py

@@ -18,13 +18,15 @@
         os.path.join(ROOT, "examples", "ase"),
         os.path.join(ROOT, "examples", "programmatic", "llpr"),
         os.path.join(ROOT, "examples", "zbl"),
-        os.path.join(ROOT, "examples", "programmatic", "use_architectures_outside")
+        os.path.join(ROOT, "examples", "programmatic", "use_architectures_outside"),
+        os.path.join(ROOT, "examples", "programmatic", "disk_dataset"),
     ],
     "gallery_dirs": [
         os.path.join(ROOT, "docs", "src", "examples", "ase"),
         os.path.join(ROOT, "docs", "src", "examples", "programmatic", "llpr"),
         os.path.join(ROOT, "docs", "src", "examples", "zbl"),
-        os.path.join(ROOT, "docs", "src", "examples", "programmatic", "use_architectures_outside")
+        os.path.join(ROOT, "docs", "src", "examples", "programmatic", "use_architectures_outside"),
+        os.path.join(ROOT, "docs", "src", "examples", "programmatic", "disk_dataset"),
     ],
     "min_reported_time": 5,
     "matplotlib_animations": True,

docs/src/tutorials/index.rst

@@ -13,3 +13,4 @@ This sections includes some more advanced tutorials on the usage of the
    ../examples/zbl/dimers
    ../examples/programmatic/llpr/llpr
    ../examples/programmatic/use_architectures_outside/use_outside
+   ../examples/programmatic/disk_dataset/disk_dataset

examples/programmatic/disk-dataset/README.rst renamed to examples/programmatic/disk_dataset/README.rst

@@ -4,15 +4,15 @@
 
 Large datasets may not fit into memory. In such cases, it is useful to save the
 dataset to disk and load it on the fly during training. This example demonstrates
-how to save a ``DiskDataset`` for this purpose.
+how to save a ``DiskDataset`` for this purpose. Metatrain will then be able to load
+``DiskDataset`` objects saved in this way to execute on-the-fly data loading.
 """
 
 # %%
 #
 
 import ase.io
 import torch
-import tqdm
 from metatensor.torch import Labels, TensorBlock, TensorMap
 from metatensor.torch.atomistic import NeighborListOptions, systems_to_torch
 
@@ -22,11 +22,12 @@
 
 # %%
 #
-# Read some sample systems. Metatrain always reads systems in float64, while torch
-# uses float32 by default. We will convert the systems to float32.
+# As an example, we will use 100 structures from the QM9 dataset. In addition to the
+# systems and targets (here the energy), we also need to save the neighbor lists that
+# the model will use during training.
 
 disk_dataset_writer = DiskDatasetWriter("qm9_reduced_100.zip")
-for i in tqdm.tqdm(range(100)):
+for i in range(100):
     frame = ase.io.read("qm9_reduced_100.xyz", index=i)
     system = systems_to_torch(frame, dtype=torch.float64)
     system = get_system_with_neighbor_lists(
@@ -48,4 +49,10 @@
         ],
     )
     disk_dataset_writer.write_sample(system, {"energy": energy})
-del disk_dataset_writer
+del disk_dataset_writer  # not necessary if the file ends here, but good in general
+
+# %%
+#
+# The dataset is saved to disk. You can now provide it to ``metatrain`` as a
+# dataset to train from, simply by replacing your ``.xyz`` file with the newly created
+# zip file (e.g. ``read_from: qm9_reduced_100.zip``).