Skip to content

Commit

Permalink
Documentation
Browse files Browse the repository at this point in the history
  • Loading branch information
@frostedoyster
frostedoyster committed Jan 7, 2025
1 parent 6d017c1 commit e886a8d
Show file tree
Hide file tree
Showing 2 changed files with 94 additions and 12 deletions.
64 changes: 61 additions & 3 deletions python/src/sphericart/metatensor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -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))
Expand All @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -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))
Expand Down
42 changes: 33 additions & 9 deletions sphericart-torch/python/sphericart/torch/metatensor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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__(
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -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))
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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)
Expand All @@ -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))
Expand All @@ -180,7 +205,6 @@ def compute_with_hessians(self, xyz: torch.Tensor) -> TensorMap:
self.precomputed_properties,
sh_gradients,
sh_hessians,
metatensor_module=metatensor.torch,
)


Expand Down

0 comments on commit e886a8d

Please sign in to comment.