Document the possibility of having a dataset with a per-sample box

docs/source/usage.rst

@@ -39,7 +39,12 @@ Periodic Boundary Conditions
 
 TorchMD-Net supports periodic boundary conditions with arbitrary triclinic boxes.
 
-Periodic boundary conditions can be enabled by passing the `box-vecs` option in the :ref:`configuration file <configuration-file>` or by passing the ``--box-vecs`` argument to the :ref:`torchmd-train <torchmd-train>` utility. You may also send the box vectors directly to a :ref:`neural network potential <neural-network-potentials>` as an argument when running inference, e.g. ``model(z, pos, batch, box=box_vecs)``.
+Periodic boundary conditions can be enabled in several ways, depending on how you are using TorchMD-Net:
+#. Pass the `box-vecs` option in the :ref:`configuration file <configuration-file>`.
+#. Pass the ``--box-vecs`` argument to the :ref:`torchmd-train <torchmd-train>` utility.
+#. Choose or write a dataset that provides a box for each sample. See for instance the :py:mod:`torchmdnet.datasets.WaterBox` dataset.
+#. You may also send the box vectors directly to a :ref:`neural network potential <neural-network-potentials>` as an argument when running inference, e.g. ``model(z, pos, batch, box=box_vecs)``.
+
 
 For a given cutoff, :math:`r_c`, the box vectors :math:`\vec{a},\vec{b},\vec{c}` must satisfy certain requirements:
 

torchmdnet/datasets/water.py

@@ -41,6 +41,14 @@ class WaterBox(InMemoryDataset):
     The dataset consists of 1593 water molecules in a cubic box with periodic boundary conditions.
     The molecules are sampled from a molecular dynamics simulation of liquid water.
 
+    Each sample in the dataset contains the following properties:
+
+    - z (LongTensor): Atomic numbers of the atoms in the molecule.
+    - pos (FloatTensor): Positions of the atoms in the molecule.
+    - y (FloatTensor): Total energy of the molecule.
+    - neg_dy (FloatTensor): Negative of the forces on the atoms in the molecule.
+    - box (FloatTensor): Box vectors of the simulation cell.
+
     Parameters
     ----------
     root : str