Basic formatting

Author: frostedoyster

docs/requirements.txt

@@ -6,4 +6,4 @@ sphinx-gallery
 sphinx-toggleprompt
 setuptools  # required for sphinxcontrib-bibtex together with Python 3.13
 tomli
-sphericart-torch
+sphericart-torch  # for SOAP-BPNN + tensor basis example

examples/programmatic/tensor_basis/save_tensor_map.py

@@ -5,9 +5,7 @@
 from featomic.torch.clebsch_gordan import cartesian_to_spherical
 
 
-# %%
-#
-# Read a subset of 1000 molecules from the QM7x dataset in the XYZ format decorated with
+# Read a subset of 100 molecules from the QM7x dataset in the XYZ format decorated with
 # the polarizability (Cartesian) tensor.
 # Extract the polarizability from the ase.Atoms.info dictionary.
 #
@@ -16,11 +14,8 @@
     [molecule.info["polarizability"].reshape(3, 3) for molecule in molecules]
 )
 
-# %%
-#
 # Create a ``metatensor.torch.TensorMap`` containing the Cartesian polarizability tensor
 # values and the respective metadata
-
 cartesian_tensormap = mts.TensorMap(
     keys=mts.Labels.single(),
     blocks=[
@@ -33,24 +28,18 @@
     ],
 )
 
-# %%
-#
 # Extract from the Cartesian polarizability tensor its irreducible spherical components
-
 spherical_tensormap = mts.remove_dimension(
     cartesian_to_spherical(cartesian_tensormap, components=["xyz_1", "xyz_2"]),
     "keys",
     "_",
 )
 
-# %%
-#
 # We drop the block with ``o3_sigma=-1``, as polarizability should be symmetric and
 # therefore any non-zero pseudo-vector component is spurious.
 spherical_tensormap = mts.drop_blocks(
     spherical_tensormap, mts.Labels(["o3_sigma"], torch.tensor([[-1]]))
 )
-# %%
-#
-# Let's save the spherical components of the polarizability tensor to disk
+
+# Save the spherical components of the polarizability tensor to disk
 mts.save("spherical_polarizability.mts", spherical_tensormap)

examples/programmatic/tensor_basis/tensor_basis.py

@@ -18,6 +18,14 @@
 import numpy as np
 
 
+# %%
+#
+# First, we need to prepare the dataset by saving the polarizability in
+# the form of spherical tensors. This is done using the ``metatensor`` library.
+#
+# .. literalinclude:: save_tensor_map.py
+#   :language: python
+
 # %%
 #
 # Write the metatrain ``options.yaml`` file for the training of the polarizability
@@ -57,7 +65,8 @@
 
 # %%
 #
-# In case you need to export a specific checkpoints, you can do so using:
+# In case you need to export a specific checkpoint (most likely from the ``outputs``
+# directory), you can do so using:
 #
 # .. code:: bash
 #