Documentation suggestion

lab-cosmo · Jan 7, 2025 · 5fdaca7 · 5fdaca7
1 parent df64c88
commit 5fdaca7
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/examples/metatensor/example.py b/examples/metatensor/example.py
@@ -31,6 +31,7 @@
 calculator = sphericart.metatensor.SphericalHarmonics(l_max)
 
 spherical_harmonics = calculator.compute(xyz)
+# for each block, the samples are the same as those of the `xyz` input
 
 for single_l in range(l_max + 1):
     spherical_single_l = spherical_harmonics.block({"o3_lambda": single_l})

diff --git a/python/src/sphericart/metatensor.py b/python/src/sphericart/metatensor.py
@@ -61,7 +61,8 @@ def compute(self, xyz: TensorMap) -> TensorMap:
             there must be one component named ``"xyz"`` with 3 values, and one property.
 
         :return: The spherical harmonics and their metadata as a
-            :py:class:`metatensor.TensorMap`
+            :py:class:`metatensor.TensorMap`. All ``samples`` in the output
+            ``TensorMap`` will be the same as those of the ``xyz`` input.
         """
         _check_xyz_tensor_map(xyz)
         sh_values = self.raw_calculator.compute(xyz.block().values.squeeze(-1))
@@ -85,7 +86,9 @@ def compute_with_gradients(self, xyz: TensorMap) -> TensorMap:
 
         :return: The spherical harmonics and their metadata as a
             :py:class:`metatensor.TensorMap`. Each ``TensorBlock`` in the output
-            ``TensorMap`` will have a gradient with respect to the Cartesian positions.
+            ``TensorMap`` will have a gradient block with respect to the Cartesian
+            positions. All ``samples`` in the output ``TensorMap`` will be the same as
+            those of the ``xyz`` input.
         """
         _check_xyz_tensor_map(xyz)
         sh_values, sh_gradients = self.raw_calculator.compute_with_gradients(
@@ -113,8 +116,10 @@ def compute_with_hessians(self, xyz: TensorMap) -> TensorMap:
 
         :return: The spherical harmonics and their metadata as a
             :py:class:`metatensor.TensorMap`. Each ``TensorBlock`` in the output
-            ``TensorMap`` will have a gradient with respect to the Cartesian positions,
-            which will itself have a gradient with respect to the Cartesian positions.
+            ``TensorMap`` will have a gradient block  with respect to the Cartesian
+            positions, which will itself have a gradient with respect to the Cartesian
+            positions. All ``samples`` in the output ``TensorMap`` will be the same as
+            those of the ``xyz`` input.
         """
         _check_xyz_tensor_map(xyz)
         sh_values, sh_gradients, sh_hessians = (