merge tensor creation in numpy and nn utilities

the function random_single_block_no_components_tensor_map is now in one
global utilities

tests/equisolve_tests/nn/test_module_tensor.py

@@ -1,8 +1,7 @@
-import functools
-
-import numpy as np
 import pytest
 
+from ..utilities import random_single_block_no_components_tensor_map
+
 
 try:
     import torch
@@ -15,79 +14,15 @@
     from equisolve.nn import LinearTensorMap
 
 try:
-    from metatensor.torch import Labels, TensorBlock, TensorMap, allclose_raise
+    from metatensor.torch import allclose_raise
 
     HAS_METATENSOR_TORCH = True
 except ImportError:
-    from metatensor import Labels, TensorBlock, TensorMap, allclose_raise
+    from metatensor import allclose_raise
 
     HAS_METATENSOR_TORCH = False
 
 
-def random_single_block_no_components_tensor_map():
-    """
-    Create a dummy tensor map to be used in tests. This is the same one as the
-    tensor map used in `tensor.rs` tests.
-    """
-    if HAS_TORCH:
-        create_random_array = torch.rand
-    else:
-        create_random_array = np.random.rand
-
-    if HAS_METATENSOR_TORCH:
-        create_int32_array = functools.partial(torch.tensor, dtype=torch.int32)
-    else:
-        create_int32_array = functools.partial(np.array, dtype=np.int32)
-
-    block_1 = TensorBlock(
-        values=create_random_array(4, 2),
-        samples=Labels(
-            ["sample", "structure"],
-            create_int32_array([[0, 0], [1, 1], [2, 2], [3, 3]]),
-        ),
-        components=[],
-        properties=Labels(["properties"], create_int32_array([[0], [1]])),
-    )
-    positions_gradient = TensorBlock(
-        values=create_random_array(7, 3, 2),
-        samples=Labels(
-            ["sample", "structure", "center"],
-            create_int32_array(
-                [
-                    [0, 0, 1],
-                    [0, 0, 2],
-                    [1, 1, 0],
-                    [1, 1, 1],
-                    [1, 1, 2],
-                    [2, 2, 0],
-                    [3, 3, 0],
-                ],
-            ),
-        ),
-        components=[Labels(["direction"], create_int32_array([[0], [1], [2]]))],
-        properties=block_1.properties,
-    )
-    block_1.add_gradient("positions", positions_gradient)
-
-    cell_gradient = TensorBlock(
-        values=create_random_array(4, 6, 2),
-        samples=Labels(
-            ["sample", "structure"],
-            create_int32_array([[0, 0], [1, 1], [2, 2], [3, 3]]),
-        ),
-        components=[
-            Labels(
-                ["direction_xx_yy_zz_yz_xz_xy"],
-                create_int32_array([[0], [1], [2], [3], [4], [5]]),
-            )
-        ],
-        properties=block_1.properties,
-    )
-    block_1.add_gradient("cell", cell_gradient)
-
-    return TensorMap(Labels.single(), [block_1])
-
-
 @pytest.mark.skipif(not (HAS_TORCH), reason="requires torch to be run")
 class TestModuleTensorMap:
     @pytest.fixture(autouse=True)
@@ -101,7 +36,9 @@ def set_random_generator(self):
     @pytest.mark.parametrize(
         "tensor",
         [
-            random_single_block_no_components_tensor_map(),
+            random_single_block_no_components_tensor_map(
+                HAS_TORCH, HAS_METATENSOR_TORCH
+            ),
         ],
     )
     def test_linear_module(self, tensor):
@@ -128,7 +65,9 @@ def test_linear_module(self, tensor):
     @pytest.mark.parametrize(
         "tensor",
         [
-            random_single_block_no_components_tensor_map(),
+            random_single_block_no_components_tensor_map(
+                HAS_TORCH, HAS_METATENSOR_TORCH
+            ),
         ],
     )
     @pytest.mark.skipif(

tests/equisolve_tests/numpy/utilities.py

@@ -1,61 +1,17 @@
+import functools
+
 import metatensor
 import numpy as np
-from metatensor import Labels, TensorBlock, TensorMap
-
+from metatensor import Labels, TensorMap
 
-def random_single_block_no_components_tensor_map():
-    """
-    Create a dummy tensor map to be used in tests. This is the same one as the
-    tensor map used in `tensor.rs` tests.
-    """
-    block_1 = TensorBlock(
-        values=np.random.rand(4, 2),
-        samples=Labels(
-            ["sample", "structure"],
-            np.array([[0, 0], [1, 1], [2, 2], [3, 3]], dtype=np.int32),
-        ),
-        components=[],
-        properties=Labels(["properties"], np.array([[0], [1]], dtype=np.int32)),
-    )
-    positions_gradient = TensorBlock(
-        values=np.random.rand(7, 3, 2),
-        samples=Labels(
-            ["sample", "structure", "center"],
-            np.array(
-                [
-                    [0, 0, 1],
-                    [0, 0, 2],
-                    [1, 1, 0],
-                    [1, 1, 1],
-                    [1, 1, 2],
-                    [2, 2, 0],
-                    [3, 3, 0],
-                ],
-                dtype=np.int32,
-            ),
-        ),
-        components=[Labels(["direction"], np.array([[0], [1], [2]], dtype=np.int32))],
-        properties=block_1.properties,
-    )
-    block_1.add_gradient("positions", positions_gradient)
+from ..utilities import random_single_block_no_components_tensor_map
 
-    cell_gradient = TensorBlock(
-        values=np.random.rand(4, 6, 2),
-        samples=Labels(
-            ["sample", "structure"],
-            np.array([[0, 0], [1, 1], [2, 2], [3, 3]], dtype=np.int32),
-        ),
-        components=[
-            Labels(
-                ["direction_xx_yy_zz_yz_xz_xy"],
-                np.array([[0], [1], [2], [3], [4], [5]], dtype=np.int32),
-            )
-        ],
-        properties=block_1.properties,
-    )
-    block_1.add_gradient("cell", cell_gradient)
 
-    return TensorMap(Labels.single(), [block_1])
+random_single_block_no_components_tensor_map = functools.partial(
+    random_single_block_no_components_tensor_map,
+    use_torch=False,
+    use_metatensor_torch=False,
+)
 
 
 def tensor_to_tensormap(a: np.ndarray, key_name: str = "keys") -> TensorMap:

tests/equisolve_tests/utilities.py

@@ -0,0 +1,79 @@
+import functools
+
+
+def random_single_block_no_components_tensor_map(use_torch, use_metatensor_torch):
+    """
+    Create a dummy tensor map to be used in tests. This is the same one as the
+    tensor map used in `tensor.rs` tests.
+    """
+    if not use_torch and use_metatensor_torch:
+        raise ValueError(
+            "torch.TensorMap cannot be created without torch.Tensor block values."
+        )
+    if use_metatensor_torch:
+        import torch
+        from metatensor.torch import Labels, TensorBlock, TensorMap
+
+        create_int32_array = functools.partial(torch.tensor, dtype=torch.int32)
+    else:
+        import numpy as np
+        from metatensor import Labels, TensorBlock, TensorMap
+
+        create_int32_array = functools.partial(np.array, dtype=np.int32)
+
+    if use_torch:
+        import torch
+
+        create_random_array = torch.rand
+    else:
+        import numpy as np
+
+        create_random_array = np.random.rand
+
+    block_1 = TensorBlock(
+        values=create_random_array(4, 2),
+        samples=Labels(
+            ["sample", "structure"],
+            create_int32_array([[0, 0], [1, 1], [2, 2], [3, 3]]),
+        ),
+        components=[],
+        properties=Labels(["properties"], create_int32_array([[0], [1]])),
+    )
+    positions_gradient = TensorBlock(
+        values=create_random_array(7, 3, 2),
+        samples=Labels(
+            ["sample", "structure", "center"],
+            create_int32_array(
+                [
+                    [0, 0, 1],
+                    [0, 0, 2],
+                    [1, 1, 0],
+                    [1, 1, 1],
+                    [1, 1, 2],
+                    [2, 2, 0],
+                    [3, 3, 0],
+                ],
+            ),
+        ),
+        components=[Labels(["direction"], create_int32_array([[0], [1], [2]]))],
+        properties=block_1.properties,
+    )
+    block_1.add_gradient("positions", positions_gradient)
+
+    cell_gradient = TensorBlock(
+        values=create_random_array(4, 6, 2),
+        samples=Labels(
+            ["sample", "structure"],
+            create_int32_array([[0, 0], [1, 1], [2, 2], [3, 3]]),
+        ),
+        components=[
+            Labels(
+                ["direction_xx_yy_zz_yz_xz_xy"],
+                create_int32_array([[0], [1], [2], [3], [4], [5]]),
+            )
+        ],
+        properties=block_1.properties,
+    )
+    block_1.add_gradient("cell", cell_gradient)
+
+    return TensorMap(Labels.single(), [block_1])