Add testing utilities to perform finite difference for operations

python/metatensor-operations/metatensor/operations/_dispatch.py

@@ -58,6 +58,20 @@ def _check_all_np_ndarray(arrays):
  )
 
 
+def sum(array, axis: Optional[int] = None):
+ """
+ Returns the sum of the elements in the array at the axis.
+
+ It is equivalent of np.sum(array, axis=axis) and torch.sum(tensor, dim=axis)
+ """
+ if isinstance(array, TorchTensor):
+ return torch.sum(array, dim=axis)
+ elif isinstance(array, np.ndarray):
+ return np.sum(array, axis=axis).astype(array.dtype)
+ else:
+ raise TypeError(UNKNOWN_ARRAY_TYPE)
+
+
 def abs(array):
  """
  Returns the absolute value of the elements in the array.

python/metatensor-operations/tests/_gradcheck.py

@@ -0,0 +1,152 @@
+import numpy as np
+
+import metatensor
+from metatensor import Labels, TensorBlock, TensorMap
+from metatensor.operations import _dispatch
+
+
+def check_finite_differences(
+ function,
+ array,
+ *,
+ parameter: str,
+ displacement: float = 1e-6,
+ rtol: float = 1e-5,
+ atol: float = 1e-15,
+) -> None:
+ """
+ Check that analytical gradients with respect to ``parameter`` in the
+ :py:class:`TensorMap` returned by ``function`` agree with finite differences.
+
+ The ``function`` must take an array (either torch or numpy) and return a
+ :py:class:`TensorMap`. All the blocks in the returned TensorMap should have one
+ sample per row of the ``array``, and the gradient-specific components must match the
+ other dimensions of the ``array``.
+ """
+ n_samples = array.shape[0]
+ n_grad_components = array.shape[1:]
+
+ reference = function(array)
+
+ values_components = reference.block(0).components
+ grad_components = reference.block(0).gradient(parameter).components
+
+ assert len(grad_components) == len(values_components) + len(n_grad_components)
+
+ for sample_i in range(n_samples):
+ for grad_components_i in np.ndindex(n_grad_components):
+ array_pos = _dispatch.copy(array)
+ index = (sample_i,) + grad_components_i
+ array_pos[index] += displacement / 2
+ updated_pos = function(array_pos)
+
+ array_neg = _dispatch.copy(array)
+ array_neg[index] -= displacement / 2
+ updated_neg = function(array_neg)
+
+ assert updated_pos.keys == reference.keys
+ assert updated_neg.keys == reference.keys
+
+ for key, block in reference.items():
+ gradients = block.gradient(parameter)
+
+ block_pos = updated_pos.block(key)
+ block_neg = updated_neg.block(key)
+
+ for gradient_i, gradient_sample in enumerate(gradients.samples):
+ current_sample_i = gradient_sample[0]
+ if current_sample_i != sample_i:
+ continue
+
+ assert block_pos.samples[sample_i] == block.samples[sample_i]
+ assert block_neg.samples[sample_i] == block.samples[sample_i]
+
+ value_pos = block_pos.values[sample_i]
+ value_neg = block_neg.values[sample_i]
+
+ grad_index = (gradient_i,) + grad_components_i
+ gradient = gradients.values[grad_index]
+
+ assert value_pos.shape == gradient.shape
+ assert value_neg.shape == gradient.shape
+
+ finite_difference = (value_pos - value_neg) / displacement
+
+ np.testing.assert_allclose(
+ finite_difference,
+ gradient,
+ rtol=rtol,
+ atol=atol,
+ )
+
+
+def cartesian_cubic(array) -> TensorMap:
+ """
+ Creates a TensorMap from a set of cartesian vectors according to the function:
+
+ .. math::
+
+ f(x, y, z) = x^3 + y^3 + z^3
+
+ \\nabla f = (3x^2, 3y^2, 3z^2)
+
+ """
+ n_samples = array.shape[0]
+ assert array.shape == (n_samples, 3)
+
+ values = _dispatch.sum(array**3, axis=1)
+ values_grad = 3 * array**2
+
+ block = metatensor.block_from_array(values.reshape(n_samples, 1))
+ block.add_gradient(
+ parameter="g",
+ gradient=TensorBlock(
+ values=values_grad.reshape(n_samples, 3, 1),
+ samples=Labels.range("sample", len(values)),
+ components=[Labels.range("xyz", 3)],
+ properties=block.properties,
+ ),
+ )
+
+ return TensorMap(Labels.range("_", 1), [block])
+
+
+def cartesian_linear(array) -> TensorMap:
+ """
+ Creates a TensorMap from a set of cartesian vectors according to the function:
+
+ .. math::
+
+ f(x, y, z) = 3x + 2y + 8*z + 4
+
+ \\nabla f = (3, 2, 8)
+
+ """
+ n_samples = array.shape[0]
+ assert array.shape == (n_samples, 3)
+
+ values = 3 * array[:, 0] + 2 * array[:, 1] + 8 * array[:, 2] + 4
+
+ values_grad = _dispatch.zeros_like(array, (n_samples, 3, 1))
+ values_grad[:, 0] = 3 * _dispatch.ones_like(array, (n_samples, 1))
+ values_grad[:, 1] = 2 * _dispatch.ones_like(array, (n_samples, 1))
+ values_grad[:, 2] = 8 * _dispatch.ones_like(array, (n_samples, 1))
+
+ block = metatensor.block_from_array(values.reshape(-1, 1))
+ block.add_gradient(
+ parameter="g",
+ gradient=TensorBlock(
+ values=values_grad,
+ samples=Labels.range("sample", len(values)),
+ components=[Labels.range("xyz", 3)],
+ properties=block.properties,
+ ),
+ )
+
+ return TensorMap(Labels.range("_", 1), [block])
+
+
+def test_basic_functions():
+ array = np.random.rand(42, 3)
+ check_finite_differences(cartesian_cubic, array, parameter="g")
+ check_finite_differences(cartesian_linear, array, parameter="g")

python/metatensor-operations/tests/add.py

@@ -4,6 +4,16 @@
 import metatensor
 from metatensor import Labels, TensorBlock, TensorMap
 
+from . import _gradcheck
+
+
+try:
+ import torch
+
+ HAS_TORCH = True
+except ImportError:
+ HAS_TORCH = False
+
 
 @pytest.fixture
 def keys():
@@ -258,3 +268,24 @@ def test_self_add_error():
  )
  with pytest.raises(TypeError, match=message):
  metatensor.add(tensor, np.ones((3, 4)))
+
+
+def test_add_finite_difference():
+ def function(array):
+ tensor_1 = _gradcheck.cartesian_linear(array)
+ tensor_2 = _gradcheck.cartesian_cubic(array)
+ return metatensor.add(tensor_1, tensor_2)
+
+ array = np.random.rand(5, 3)
+ _gradcheck.check_finite_differences(function, array, parameter="g")
+
+
+@pytest.mark.skipif(not HAS_TORCH, reason="requires torch")
+def test_torch_add_finite_difference():
+ def function(array):
+ tensor_1 = _gradcheck.cartesian_linear(array)
+ tensor_2 = _gradcheck.cartesian_cubic(array)
+ return metatensor.add(tensor_1, tensor_2)
+
+ array = torch.rand(5, 3, dtype=torch.float64)
+ _gradcheck.check_finite_differences(function, array, parameter="g")