gradient

python/metatensor-operations/metatensor/operations/sort.py

@@ -70,14 +70,10 @@ def _sort_single_gradient_block(
  _dispatch.to_index_array(sample_values[:, 0])
  ]
 
- # I AM NOT SURE WHAT THESE LINES DO THE OLD TEST PASS EVEN
- # IF THEY ARE COMMENTED
-
  # sort the samples in gradient regularly moving the rows considering all columns
- # sorted_idx = _dispatch.argsort_labels_values(sample_values,
- # reverse=descending)
- # sample_values = sample_values[sorted_idx]
- # values = values[sorted_idx]
+ sorted_idx = _dispatch.argsort_labels_values(sample_values, reverse=descending)
+ sample_values = sample_values[sorted_idx]
+ values = values[sorted_idx]
  if "components" in axes:
  for i, _ in enumerate(gradient_block.components):
  sorted_idx = _dispatch.argsort_labels_values(
@@ -254,24 +250,6 @@ def sort_block(
  >>> sorted_block = metatensor.sort_block(block)
  >>> np.all(sorted_block.values == block_sorted_stepwise.values)
  True
- >>> # You can also choose along which axis of "samples“ you sort
- >>> block2 = TensorBlock(
- ... values=np.arange(12).reshape(4, 3),
- ... samples=Labels(
- ... ["system", "atom"], np.array([[0, 2], [1, 0], [2, 5], [2, 1]])
- ... ),
- ... components=[],
- ... properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
- ... )
- >>> block_sorted_2_sample = metatensor.sort_block(
- ... block2, axes=["samples"], name="atom"
- ... )
- >>> # samples (first dimension of the array) are sorted
- >>> block_sorted_2_sample.values
- array([[ 3, 4, 5],
- [ 9, 10, 11],
- [ 0, 1, 2],
- [ 6, 7, 8]])
  >>> # This function can also sort gradients:
  >>> sorted_block.add_gradient(
  ... parameter="g",
@@ -297,7 +275,24 @@ def sort_block(
  <BLANKLINE>
  [[ 6, 7, 8],
  [ 9, 10, 11]]])
-
+ >>> # You can also choose along which axis of "samples“ you sort
+ >>> block2 = TensorBlock(
+ ... values=np.arange(12).reshape(4, 3),
+ ... samples=Labels(
+ ... ["system", "atom"], np.array([[0, 2], [1, 0], [2, 5], [2, 1]])
+ ... ),
+ ... components=[],
+ ... properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
+ ... )
+ >>> block_sorted_2_sample = metatensor.sort_block(
+ ... block2, axes=["samples"], name="atom"
+ ... )
+ >>> # samples (first dimension of the array) are sorted
+ >>> block_sorted_2_sample.values
+ array([[ 3, 4, 5],
+ [ 9, 10, 11],
+ [ 0, 1, 2],
+ [ 6, 7, 8]])
  """
  if isinstance(axes, str):
  if axes == "all":

python/metatensor-operations/tests/sort.py

@@ -341,26 +341,26 @@ def tensor_two_samples_ascending_a():
  gradient=TensorBlock(
  values=np.array(
  [
- [[1, 2], [3, 4]],
  [[5, 6], [7, 8]],
+ [[17, 18], [19, 20]],
  [[9, 10], [11, 12]],
  [[13, 14], [15, 16]],
- [[17, 18], [19, 20]],
- [[21, 22], [23, 24]],
+ [[1, 2], [3, 4]],
  [[25, 26], [27, 28]],
+ [[21, 22], [23, 24]],
  ]
  ),
  samples=Labels(
  ["sample", "g"],
  np.array(
  [
- [3, 1],
  [0, 1],
+ [1, 1],
  [2, 1],
  [2, 2],
- [1, 1],
- [4, 1],
+ [3, 1],
  [4, 0],
+ [4, 1],
  ]
  ),
  ),
@@ -451,6 +451,14 @@ def test_raise_error(tensor, tensor_sorted_ascending):
 
 
 def test_sort_two_sample(tensor_two_samples, tensor_two_samples_ascending_a):
+ print(
+ "ii",
+ metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a")
+ .gradient("g")
+ .samples,
+ )
+ print("jj", tensor_two_samples_ascending_a.block(0).gradient("g").samples)
+
  metatensor.allclose_block_raise(
  metatensor.sort_block(tensor_two_samples.block(0), axes="samples", name="a"),
  tensor_two_samples_ascending_a.block(0),