From f8a62ce0f30357d24179bd666ee75dacad05e3fd Mon Sep 17 00:00:00 2001
From: David Zwicker <david.zwicker@ds.mpg.de>
Date: Wed, 2 Aug 2023 09:05:01 +0200
Subject: [PATCH] Allow generating random emulsions for grids

---
 droplets/emulsions.py  | 35 +++++++++++++++++++----------------
 tests/test_emulsion.py |  2 +-
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/droplets/emulsions.py b/droplets/emulsions.py
index 6e98498..59f3998 100644
--- a/droplets/emulsions.py
+++ b/droplets/emulsions.py
@@ -100,7 +100,7 @@ def __init__(
     def from_random(
         cls,
         num: int,
-        bounds: Union[Sequence[Tuple[float, float]], GridBase],
+        grid_or_bounds: Union[GridBase, Sequence[Tuple[float, float]]],
         radius: Union[float, Tuple[float, float]],
         *,
         remove_overlapping: bool = True,
@@ -113,12 +113,11 @@ def from_random(
         Args:
             num (int):
                 The (maximal) number of droplets to generate
-            bounds:
-                Boundaries of the space in which droplets are placed. This is either a
-                :class:`~pde.grids.base.GridBase` from which the bounds are taken or a
-                sequence of tuples with lower and upper bounds for each axes. The length
-                of the sequence determines the dimension of the space in which droplets
-                are placed.
+            grid_or_bounds (:class:`~pde.grids.base.GridBase` or list of float tuples):
+                Determines the space in which droplets are placed. This is either a
+                :class:`~pde.grids.base.GridBase` describing the geometry or a sequence
+                of tuples with lower and upper bounds for each axes, so the length of
+                the sequence determines the space dimension.
             radius (float or tuple of float):
                 Radius of the droplets that are created. If two numbers are given, they
                 specify the bounds of a uniform distribution from which the radius of
@@ -134,11 +133,18 @@ def from_random(
         if rng is None:
             rng = np.random.default_rng()
 
-        # extract information about possible positions
-        if isinstance(bounds, GridBase):
-            bounds = bounds.axes_bounds
-        bnds = np.atleast_2d(bounds)
-        assert bnds.ndim == 2 and bnds.shape[0] > 0 and bnds.shape[1] == 2
+        # determine how to get random droplet positions
+        if isinstance(grid_or_bounds, GridBase):
+
+            def get_position():
+                return grid_or_bounds.get_random_point(rng=rng)
+
+        else:
+            bnds = np.atleast_2d(grid_or_bounds)
+            assert bnds.ndim == 2 and bnds.shape[0] > 0 and bnds.shape[1] == 2
+
+            def get_position():
+                return rng.uniform(bnds[:, 0], bnds[:, 1])
 
         # extract information about radius
         try:
@@ -147,10 +153,7 @@ def from_random(
             r0 = r1 = float(radius)  # type: ignore
 
         # create the emulsion from a list of droplets
-        drops = [
-            droplet_class(rng.uniform(bnds[:, 0], bnds[:, 1]), rng.uniform(r0, r1))
-            for _ in range(num)
-        ]
+        drops = [droplet_class(get_position(), rng.uniform(r0, r1)) for _ in range(num)]
         emulsion = cls(drops)
 
         if remove_overlapping:
diff --git a/tests/test_emulsion.py b/tests/test_emulsion.py
index 6e96db7..e518975 100644
--- a/tests/test_emulsion.py
+++ b/tests/test_emulsion.py
@@ -244,7 +244,7 @@ def test_emulsion_random(dim, grid):
         bounds = CartesianGrid([(10, 30)] * dim, 1)
     else:
         bounds = [(10, 30)] * dim
-    em = Emulsion.from_random(num=10, bounds=bounds, radius=(1, 2), rng=rng)
+    em = Emulsion.from_random(10, bounds, radius=(1, 2), rng=rng)
     assert 1 < len(em) < 10
     assert em.dim == dim
     assert np.all(em.data["position"] > 10) and np.all(em.data["position"] < 30)