Address open reviewer's comments in test_constructors

GridTools · Nov 14, 2023 · bfbc65c · bfbc65c
1 parent 739604e
commit bfbc65c
Showing 1 changed file with 83 additions and 72 deletions.
diff --git a/tests/next_tests/unit_tests/test_constructors.py b/tests/next_tests/unit_tests/test_constructors.py
@@ -30,74 +30,13 @@
 sizes = {I: 10, J: 10, K: 10}
 
 
-def test_as_field():
-    a = np.random.rand(sizes[I]).astype(gtx.float32)
-    ref = gtx.as_field([I], a)
-    assert np.allclose(ref.ndarray, a)
-
-
-def test_as_field_domain():
-    a = np.random.rand(sizes[I] - 1, sizes[J] - 1).astype(gtx.float32)
-    domain = common.Domain(
-        dims=(I, J),
-        ranges=(common.UnitRange(0, sizes[I] - 1), common.UnitRange(0, sizes[J] - 1)),
-    )
-    ref = gtx.as_field(domain, a)
-    assert np.allclose(ref.ndarray, a)
-
-
-def test_as_field_origin():
-    a = np.random.rand(sizes[I], sizes[J]).astype(gtx.float32)
-    ref = gtx.as_field([I, J], a, origin={I: 1, J: 2})
-    domain_range = [(val.start, val.stop) for val in ref.domain.ranges]
-    assert np.allclose(domain_range, [(-1, 9), (-2, 8)])
-
-
-# for as_field, check that the domain is correct depending on data origin and domain itself
-
-
-def test_field_wrong_dims():
-    with pytest.raises(
-        ValueError,
-        match=(r"Cannot construct `Field` from array of shape"),
-    ):
-        gtx.as_field([I, J], np.random.rand(sizes[I]).astype(gtx.float32))
-
-
-def test_field_wrong_domain():
-    with pytest.raises(
-        ValueError,
-        match=(r"Cannot construct `Field` from array of shape"),
-    ):
-        domain = common.Domain(
-            dims=(I, J),
-            ranges=(common.UnitRange(0, sizes[I] - 1), common.UnitRange(0, sizes[J] - 1)),
-        )
-        gtx.as_field(domain, np.random.rand(sizes[I], sizes[J]).astype(gtx.float32))
-
-
-def test_field_wrong_origin():
-    with pytest.raises(
-        ValueError,
-        match=(r"Origin keys {'J'} not in domain"),
-    ):
-        gtx.as_field([I], np.random.rand(sizes[I]).astype(gtx.float32), origin={"J": 0})
-
-    with pytest.raises(
-        ValueError,
-        match=(r"Cannot specify origin for domain I"),
-    ):
-        gtx.as_field("I", np.random.rand(sizes[J]).astype(gtx.float32), origin={"J": 0})
-
-
-@pytest.mark.xfail(reason="aligned_index not supported yet")
-def test_aligned_index():
-    gtx.as_field([I], np.random.rand(sizes[I]).astype(gtx.float32), aligned_index=[I, 0])
-
-
+# TODO: parametrize with gpu backend and compare with cupy array
 @pytest.mark.parametrize(
     "allocator, device",
-    [[roundtrip.backend, None], [None, core_defs.Device(core_defs.DeviceType.CPU, 0)]],
+    [
+        [next_allocators.StandardCPUFieldBufferAllocator(), None],
+        [None, core_defs.Device(core_defs.DeviceType.CPU, 0)],
+    ],
 )
 def test_empty(allocator, device):
     a = np.empty([sizes[I], sizes[J]]).astype(gtx.float32)
@@ -110,9 +49,13 @@ def test_empty(allocator, device):
     assert ref.shape, a.shape
 
 
+# TODO: parametrize with gpu backend and compare with cupy array
 @pytest.mark.parametrize(
     "allocator, device",
-    [[roundtrip.backend, None], [None, core_defs.Device(core_defs.DeviceType.CPU, 0)]],
+    [
+        [next_allocators.StandardCPUFieldBufferAllocator(), None],
+        [None, core_defs.Device(core_defs.DeviceType.CPU, 0)],
+    ],
 )
 def test_zeros(allocator, device):
     ref = gtx.zeros(
@@ -128,12 +71,13 @@ def test_zeros(allocator, device):
     assert np.array_equal(ref.ndarray, a)
 
 
-# parametrize with gpu backend and compare with cupy array
-
-
+# TODO: parametrize with gpu backend and compare with cupy array
 @pytest.mark.parametrize(
     "allocator, device",
-    [[roundtrip.backend, None], [None, core_defs.Device(core_defs.DeviceType.CPU, 0)]],
+    [
+        [next_allocators.StandardCPUFieldBufferAllocator(), None],
+        [None, core_defs.Device(core_defs.DeviceType.CPU, 0)],
+    ],
 )
 def test_ones(allocator, device):
     ref = gtx.ones(
@@ -147,9 +91,13 @@ def test_ones(allocator, device):
     assert np.array_equal(ref.ndarray, a)
 
 
+# TODO: parametrize with gpu backend and compare with cupy array
 @pytest.mark.parametrize(
     "allocator, device",
-    [[roundtrip.backend, None], [None, core_defs.Device(core_defs.DeviceType.CPU, 0)]],
+    [
+        [next_allocators.StandardCPUFieldBufferAllocator(), None],
+        [None, core_defs.Device(core_defs.DeviceType.CPU, 0)],
+    ],
 )
 def test_full(allocator, device):
     ref = gtx.full(
@@ -162,3 +110,66 @@ def test_full(allocator, device):
     a = np.full((sizes[I] - 2, sizes[J] - 2), 42.0).astype(gtx.float32)
 
     assert np.array_equal(ref.ndarray, a)
+
+
+def test_as_field():
+    a = np.random.rand(sizes[I]).astype(gtx.float32)
+    ref = gtx.as_field([I], a)
+    assert np.allclose(ref.ndarray, a)
+
+
+def test_as_field_domain():
+    a = np.random.rand(sizes[I] - 1, sizes[J] - 1).astype(gtx.float32)
+    domain = common.Domain(
+        dims=(I, J),
+        ranges=(common.UnitRange(0, sizes[I] - 1), common.UnitRange(0, sizes[J] - 1)),
+    )
+    ref = gtx.as_field(domain, a)
+    assert np.allclose(ref.ndarray, a)
+
+
+def test_as_field_origin():
+    a = np.random.rand(sizes[I], sizes[J]).astype(gtx.float32)
+    ref = gtx.as_field([I, J], a, origin={I: 1, J: 2})
+    domain_range = [(val.start, val.stop) for val in ref.domain.ranges]
+    assert np.allclose(domain_range, [(-1, 9), (-2, 8)])
+
+
+# check that `as_field()` domain is correct depending on data origin and domain itself
+def test_field_wrong_dims():
+    with pytest.raises(
+        ValueError,
+        match=(r"Cannot construct `Field` from array of shape"),
+    ):
+        gtx.as_field([I, J], np.random.rand(sizes[I]).astype(gtx.float32))
+
+
+def test_field_wrong_domain():
+    with pytest.raises(
+        ValueError,
+        match=(r"Cannot construct `Field` from array of shape"),
+    ):
+        domain = common.Domain(
+            dims=(I, J),
+            ranges=(common.UnitRange(0, sizes[I] - 1), common.UnitRange(0, sizes[J] - 1)),
+        )
+        gtx.as_field(domain, np.random.rand(sizes[I], sizes[J]).astype(gtx.float32))
+
+
+def test_field_wrong_origin():
+    with pytest.raises(
+        ValueError,
+        match=(r"Origin keys {'J'} not in domain"),
+    ):
+        gtx.as_field([I], np.random.rand(sizes[I]).astype(gtx.float32), origin={"J": 0})
+
+    with pytest.raises(
+        ValueError,
+        match=(r"Cannot specify origin for domain I"),
+    ):
+        gtx.as_field("I", np.random.rand(sizes[J]).astype(gtx.float32), origin={"J": 0})
+
+
+@pytest.mark.xfail(reason="aligned_index not supported yet")
+def test_aligned_index():
+    gtx.as_field([I], np.random.rand(sizes[I]).astype(gtx.float32), aligned_index=[I, 0])