From e886a8de1a7b7f8f0050fcac26b848fa95ba59a3 Mon Sep 17 00:00:00 2001 From: frostedoyster Date: Tue, 7 Jan 2025 08:46:26 +0100 Subject: [PATCH] Documentation --- python/src/sphericart/metatensor.py | 64 ++++++++++++++++++- .../python/sphericart/torch/metatensor.py | 42 +++++++++--- 2 files changed, 94 insertions(+), 12 deletions(-) diff --git a/python/src/sphericart/metatensor.py b/python/src/sphericart/metatensor.py index 6665e83cb..197054c7e 100644 --- a/python/src/sphericart/metatensor.py +++ b/python/src/sphericart/metatensor.py @@ -16,6 +16,14 @@ class SphericalHarmonics: + """ + ``metatensor``-based wrapper around the + :py:meth:`sphericart.SphericalHarmonics` class. + + :param l_max: the maximum degree of the spherical harmonics to be calculated + + :return: a spherical harmonics calculator object + """ def __init__(self, l_max: int): self.l_max = l_max @@ -44,6 +52,18 @@ def __init__(self, l_max: int): self.precomputed_properties = Labels.single() def compute(self, xyz: TensorMap) -> TensorMap: + """ + Computes the spherical harmonics for the given Cartesian coordinates, up to + the maximum degree ``l_max`` specified during initialization. + + :param xyz: a :py:class:`metatensor.TensorMap` containing the Cartesian + coordinates of the 3D points. This ``TensorMap`` should have only one + ``TensorBlock``. In this ``TensorBlock``, the samples are arbitrary, + 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` + """ _check_xyz_tensor_map(xyz) sh_values = self.raw_calculator.compute(xyz.block().values.squeeze(-1)) return _wrap_into_tensor_map( @@ -57,6 +77,17 @@ def compute(self, xyz: TensorMap) -> TensorMap: ) def compute_with_gradients(self, xyz: TensorMap) -> TensorMap: + """ + Computes the spherical harmonics for the given Cartesian coordinates, up to + the maximum degree ``l_max`` specified during initialization, + together with their gradients with respect to the Cartesian coordinates. + + :param xyz: see :py:meth:`compute` + + :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. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients = self.raw_calculator.compute_with_gradients( xyz.block().values.squeeze(-1) @@ -73,6 +104,19 @@ def compute_with_gradients(self, xyz: TensorMap) -> TensorMap: ) def compute_with_hessians(self, xyz: TensorMap) -> TensorMap: + """ + Computes the spherical harmonics for the given Cartesian coordinates, up to + the maximum degree ``l_max`` specified during initialization, + together with their gradients and Hessians with respect to the Cartesian + coordinates. + + :param xyz: see :py:meth:`compute` + + :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. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients, sh_hessians = ( self.raw_calculator.compute_with_hessians(xyz.block().values.squeeze(-1)) @@ -91,6 +135,11 @@ def compute_with_hessians(self, xyz: TensorMap) -> TensorMap: class SolidHarmonics: + """ + ``metatensor``-based wrapper around the :py:meth:`sphericart.SolidHarmonics` class. + + See :py:class:`SphericalHarmonics` for more details. + """ def __init__(self, l_max: int): self.l_max = l_max @@ -118,7 +167,10 @@ def __init__(self, l_max: int): ) self.precomputed_properties = Labels.single() - def compute(self, xyz: np.ndarray) -> TensorMap: + def compute(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute`. + """ _check_xyz_tensor_map(xyz) sh_values = self.raw_calculator.compute(xyz.block().values.squeeze(-1)) return _wrap_into_tensor_map( @@ -131,7 +183,10 @@ def compute(self, xyz: np.ndarray) -> TensorMap: self.precomputed_properties, ) - def compute_with_gradients(self, xyz: np.ndarray) -> TensorMap: + def compute_with_gradients(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute_with_gradients`. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients = self.raw_calculator.compute_with_gradients( xyz.block().values.squeeze(-1) @@ -147,7 +202,10 @@ def compute_with_gradients(self, xyz: np.ndarray) -> TensorMap: sh_gradients, ) - def compute_with_hessians(self, xyz: np.ndarray) -> TensorMap: + def compute_with_hessians(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute_with_hessians`. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients, sh_hessians = ( self.raw_calculator.compute_with_hessians(xyz.block().values.squeeze(-1)) diff --git a/sphericart-torch/python/sphericart/torch/metatensor.py b/sphericart-torch/python/sphericart/torch/metatensor.py index 552d17c49..8f34b42ae 100644 --- a/sphericart-torch/python/sphericart/torch/metatensor.py +++ b/sphericart-torch/python/sphericart/torch/metatensor.py @@ -20,6 +20,10 @@ class SphericalHarmonics: """ ``metatensor``-based wrapper around the :py:meth:`sphericart.torch.SphericalHarmonics` class. + + See :py:class:`sphericart.metatensor.SphericalHarmonics` for more details. + ``backward_second_derivatives`` has the same meaning as in + :py:class:`sphericart.torch.SphericalHarmonics`. """ def __init__( @@ -53,6 +57,9 @@ def __init__( self.precomputed_properties = Labels.single() def compute(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute`. + """ _check_xyz_tensor_map(xyz) sh_values = self.raw_calculator.compute(xyz.block().values.squeeze(-1)) return _wrap_into_tensor_map( @@ -63,10 +70,12 @@ def compute(self, xyz: TensorMap) -> TensorMap: self.precomputed_xyz_components, self.precomputed_xyz_2_components, self.precomputed_properties, - metatensor_module=metatensor.torch, ) def compute_with_gradients(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute_with_gradients`. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients = self.raw_calculator.compute_with_gradients( xyz.block().values.squeeze(-1) @@ -80,10 +89,12 @@ def compute_with_gradients(self, xyz: TensorMap) -> TensorMap: self.precomputed_xyz_2_components, self.precomputed_properties, sh_gradients, - metatensor_module=metatensor.torch, ) def compute_with_hessians(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute_with_hessians`. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients, sh_hessians = ( self.raw_calculator.compute_with_hessians(xyz.block().values.squeeze(-1)) @@ -98,11 +109,18 @@ def compute_with_hessians(self, xyz: TensorMap) -> TensorMap: self.precomputed_xyz_2_components, sh_gradients, sh_hessians, - metatensor_module=metatensor.torch, ) class SolidHarmonics: + """ + ``metatensor``-based wrapper around the + :py:meth:`sphericart.torch.SolidHarmonics` class. + + See :py:class:`sphericart.metatensor.SphericalHarmonics` for more details. + ``backward_second_derivatives`` has the same meaning as in + :py:class:`sphericart.torch.SphericalHarmonics`. + """ def __init__( self, @@ -134,7 +152,10 @@ def __init__( ) self.precomputed_properties = Labels.single() - def compute(self, xyz: torch.Tensor) -> TensorMap: + def compute(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute`. + """ _check_xyz_tensor_map(xyz) sh_values = self.raw_calculator.compute(xyz.block().values.squeeze(-1)) return _wrap_into_tensor_map( @@ -145,10 +166,12 @@ def compute(self, xyz: torch.Tensor) -> TensorMap: self.precomputed_xyz_components, self.precomputed_xyz_2_components, self.precomputed_properties, - metatensor_module=metatensor.torch, ) - def compute_with_gradients(self, xyz: torch.Tensor) -> TensorMap: + def compute_with_gradients(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute_with_gradients`. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients = self.raw_calculator.compute_with_gradients( xyz.block().values.squeeze(-1) @@ -162,10 +185,12 @@ def compute_with_gradients(self, xyz: torch.Tensor) -> TensorMap: self.precomputed_xyz_2_components, self.precomputed_properties, sh_gradients, - metatensor_module=metatensor.torch, ) - def compute_with_hessians(self, xyz: torch.Tensor) -> TensorMap: + def compute_with_hessians(self, xyz: TensorMap) -> TensorMap: + """ + See :py:meth:`sphericart.metatensor.SphericalHarmonics.compute_with_hessians`. + """ _check_xyz_tensor_map(xyz) sh_values, sh_gradients, sh_hessians = ( self.raw_calculator.compute_with_hessians(xyz.block().values.squeeze(-1)) @@ -180,7 +205,6 @@ def compute_with_hessians(self, xyz: torch.Tensor) -> TensorMap: self.precomputed_properties, sh_gradients, sh_hessians, - metatensor_module=metatensor.torch, )