spatial sampling: option to limit sampling along X-direction to a given range (#1316)

Co-authored-by: Sylwester Arabas <[email protected]>

Author: jtbuch

PySDM/initialisation/sampling/spatial_sampling.py

@@ -8,23 +8,36 @@
 
 class Pseudorandom:  # pylint: disable=too-few-public-methods
     @staticmethod
-    def sample(*, backend, grid, n_sd, z_part=None):
-        dimension = len(grid)
-        n_elements = dimension * n_sd
+    def sample(*, backend, grid, n_sd, z_part=None, x_part=None):
+        n_dims = len(grid)
+        scale_factors = []
+        affine_factors = []
 
-        storage = backend.Storage.empty(n_elements, dtype=float)
-        backend.Random(seed=backend.formulae.seed, size=n_elements)(storage)
-        positions = storage.to_ndarray().reshape(dimension, n_sd)
+        storage = backend.Storage.empty(n_dims * n_sd, dtype=float)
+        backend.Random(seed=backend.formulae.seed, size=n_dims * n_sd)(storage)
+        positions = storage.to_ndarray().reshape(n_dims, n_sd)
 
         if z_part is None:
-            for dim in range(dimension):
-                positions[dim, :] *= grid[dim]
+            scale_factors.append(grid[0])
+            affine_factors.append(0)
         else:
-            assert dimension == 1
             iz_min = int(grid[0] * z_part[0])
             iz_max = int(grid[0] * z_part[1])
-            for dim in range(dimension):
-                positions[dim, :] *= iz_max - iz_min
-                positions[dim, :] += iz_min
+            scale_factors.append(iz_max - iz_min)
+            affine_factors.append(iz_min)
+
+        if x_part is not None:
+            ix_min = int(grid[1] * x_part[0])
+            ix_max = int(grid[1] * x_part[1])
+            scale_factors.append(ix_max - ix_min)
+            affine_factors.append(ix_min)
+        else:
+            if n_dims == 2:
+                scale_factors.append(grid[1])
+                affine_factors.append(0)
+
+        for dim in range(n_dims):
+            positions[dim, :] *= scale_factors[dim]
+            positions[dim, :] += affine_factors[dim]
 
         return positions

tests/unit_tests/initialisation/test_spatial_discretisation.py

@@ -55,3 +55,33 @@ def test_pseudorandom_zrange(z_range, backend_class):
     # assert
     np.testing.assert_array_less(positions, comp * z_range[1] * grid[0])
     np.testing.assert_array_less(comp * z_range[0] * grid[0], positions)
+
+
+@pytest.mark.parametrize(
+    "z_range, x_range",
+    (
+        pytest.param((0.0, 1.0), (0.0, 1.0), id="full range"),
+        pytest.param((0.5, 0.75), (0.5, 0.75), id="partial range"),
+    ),
+)
+def test_pseudorandom_x_z_range(z_range, x_range, backend_class):
+    # arrange
+    assert len(z_range) == 2
+    assert len(x_range) == 2
+    backend = backend_class(Formulae())
+    grid = (8, 8)
+    n_sd = 100
+
+    # act
+    positions = Pseudorandom.sample(
+        backend=backend, grid=grid, n_sd=n_sd, z_part=z_range, x_part=x_range
+    )
+
+    for droplet in range(n_sd):
+        # assert z positions
+        np.testing.assert_array_less(positions[0][droplet], z_range[1] * grid[0])
+        np.testing.assert_array_less(z_range[0] * grid[0], positions[0][droplet])
+
+        # assert x positions
+        np.testing.assert_array_less(positions[1][droplet], x_range[1] * grid[1])
+        np.testing.assert_array_less(x_range[0] * grid[1], positions[1][droplet])