From abf5919658ab06ee784c6c7979930d9c19d28b04 Mon Sep 17 00:00:00 2001
From: fdrgsp <federico.gasparoli@gmail.com>
Date: Mon, 4 Sep 2023 15:09:18 -0400
Subject: [PATCH] fix: return np.ndarray

---
 src/useq/_grid.py      | 22 ++++++++++++----------
 tests/test_grid.py     |  4 ++--
 tests/test_sequence.py |  6 +++---
 3 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/useq/_grid.py b/src/useq/_grid.py
index 2c0f9918..5b03263e 100644
--- a/src/useq/_grid.py
+++ b/src/useq/_grid.py
@@ -420,7 +420,7 @@ def __iter__(self) -> Iterator[GridPosition]:  # type: ignore
                 break
         else:
             warnings.warn(
-                f"Unable to generate {self.num_points} points non-overlapping points."
+                f"Unable to generate {self.num_points} non-overlapping points. "
                 f"Only {len(points)} points were found.",
                 stacklevel=2,
             )
@@ -448,27 +448,29 @@ def _is_a_valid_point(
 
 def _random_points_in_ellipse(
     seed: np.random.RandomState, n_points: int, max_width: float, max_height: float
-) -> Iterator[Tuple[float, float]]:
+) -> np.ndarray:
     """Generate a random point around a circle with center (0, 0).
 
     The point is within +/- radius_x and +/- radius_y at a random angle.
     """
-    for x0, y0, angle in seed.uniform(0, 1, size=(n_points, 3)):
-        yield (
-            np.sqrt(x0) * (max_width / 2) * np.cos(angle * 2 * np.pi),
-            np.sqrt(y0) * (max_height / 2) * np.sin(angle * 2 * np.pi),
-        )
+    xy = np.sqrt(seed.uniform(0, 1, size=(n_points, 2)))
+    angle = seed.uniform(0, 2 * np.pi, size=n_points)
+    xy[:, 0] *= (max_width / 2) * np.cos(angle)
+    xy[:, 1] *= (max_height / 2) * np.sin(angle)
+    return xy
 
 
 def _random_points_in_rectangle(
     seed: np.random.RandomState, n_points: int, max_width: float, max_height: float
-) -> Iterator[Tuple[float, float]]:
+) -> np.ndarray:
     """Generate a random point around a rectangle with center (0, 0).
 
     The point is within the bounding box (-width/2, -height/2, width, height).
     """
-    for x0, y0 in seed.uniform(0, 1, size=(n_points, 2)):
-        yield (x0 * max_width) - (max_width / 2), (y0 * max_height) - (max_height / 2)
+    xy = seed.uniform(0, 1, size=(n_points, 2))
+    xy[:, 0] = (xy[:, 0] * max_width) - (max_width / 2)
+    xy[:, 1] = (xy[:, 1] * max_height) - (max_height / 2)
+    return xy
 
 
 PointGenerator = Callable[[np.random.RandomState, int, float, float], np.ndarray]
diff --git a/tests/test_grid.py b/tests/test_grid.py
index 056d6d27..42903807 100644
--- a/tests/test_grid.py
+++ b/tests/test_grid.py
@@ -104,7 +104,7 @@ def test_num_position_error() -> None:
 
 
 expected_rectangle = [(0.2, 1.1), (0.4, 0.2), (-0.3, 0.7)]
-expected_ellipse = [(-1.2, -1.3), (1.3, -0.5), (-1.2, -0.4)]
+expected_ellipse = [(-0.0, -2.1), (0.7, 1.7), (-1.0, 1.3)]
 
 
 @pytest.mark.parametrize("n_points", [3, 100])
@@ -130,5 +130,5 @@ def test_random_points(n_points: int, shape: str, seed: Optional[int]) -> None:
         else:
             assert values == expected
     else:
-        with pytest.raises(UserWarning, match="Max number of iterations reached"):
+        with pytest.raises(UserWarning, match="Unable to generate"):
             list(rp)
diff --git a/tests/test_sequence.py b/tests/test_sequence.py
index 57f24045..e1d478d3 100644
--- a/tests/test_sequence.py
+++ b/tests/test_sequence.py
@@ -116,9 +116,9 @@
             random_seed=0,
         ),
         [
-            GridPosition(x=-1.2, y=-1.3, row=0, col=0, is_relative=True),
-            GridPosition(x=1.3, y=-0.5, row=0, col=0, is_relative=True),
-            GridPosition(x=-1.2, y=-0.4, row=0, col=0, is_relative=True),
+            GridPosition(x=-0.0, y=-2.1, row=0, col=0, is_relative=True),
+            GridPosition(x=0.7, y=1.7, row=0, col=0, is_relative=True),
+            GridPosition(x=-1.0, y=1.3, row=0, col=0, is_relative=True),
         ],
     ),
 ]