diff --git a/autoarray/structures/triangles/abstract.py b/autoarray/structures/triangles/abstract.py index 463d1017..c0efd2fe 100644 --- a/autoarray/structures/triangles/abstract.py +++ b/autoarray/structures/triangles/abstract.py @@ -252,7 +252,7 @@ def containing_indices(self, shape: Shape) -> np.ndarray: Returns ------- - The triangles that intersect the shape. + The indices of triangles that intersect the shape. """ @abstractmethod diff --git a/autoarray/structures/triangles/abstract_coordinate_array.py b/autoarray/structures/triangles/abstract_coordinate_array.py index 0a5a36b2..72267275 100644 --- a/autoarray/structures/triangles/abstract_coordinate_array.py +++ b/autoarray/structures/triangles/abstract_coordinate_array.py @@ -10,7 +10,7 @@ class AbstractCoordinateArray(ABC): def __init__( self, coordinates: np.ndarray, - side_length: float, + side_length: float = 1.0, x_offset: float = 0.0, y_offset: float = 0.0, flipped: bool = False, @@ -140,23 +140,18 @@ def for_limits_and_scale( scale: float = 1.0, **_, ): + x_shift = int(2 * x_min / scale) + y_shift = int(y_min / (HEIGHT_FACTOR * scale)) + coordinates = [] - x = x_min - while x < x_max: - y = y_min - while y < y_max: - coordinates.append([x, y]) - y += scale - x += scale - x_mean = (x_min + x_max) / 2 - y_mean = (y_min + y_max) / 2 + for x in range(x_shift, int(2 * x_max / scale) + 1): + for y in range(y_shift - 1, int(y_max / (HEIGHT_FACTOR * scale)) + 2): + coordinates.append([x, y]) return cls( coordinates=np.array(coordinates), side_length=scale, - x_offset=x_mean, - y_offset=y_mean, ) @property diff --git a/autoarray/structures/triangles/coordinate_array.py b/autoarray/structures/triangles/coordinate_array.py index e45f3eed..9551885e 100644 --- a/autoarray/structures/triangles/coordinate_array.py +++ b/autoarray/structures/triangles/coordinate_array.py @@ -5,6 +5,7 @@ AbstractCoordinateArray, ) from autoarray.structures.triangles.array import ArrayTriangles +from autoarray.structures.triangles.shape import Shape from autoconf import cached_property @@ -88,7 +89,9 @@ def neighborhood(self) -> "CoordinateArrayTriangles": def _vertices_and_indices(self): flat_triangles = self.triangles.reshape(-1, 2) vertices, inverse_indices = np.unique( - flat_triangles, axis=0, return_inverse=True + flat_triangles, + axis=0, + return_inverse=True, ) indices = inverse_indices.reshape(-1, 3) return vertices, indices @@ -131,3 +134,18 @@ def for_indexes(self, indexes: np.ndarray) -> "CoordinateArrayTriangles": x_offset=self.x_offset, flipped=self.flipped, ) + + def containing_indices(self, shape: Shape) -> np.ndarray: + """ + Find the triangles that insect with a given shape. + + Parameters + ---------- + shape + The shape + + Returns + ------- + The indices of triangles that intersect the shape. + """ + return self.with_vertices(self.vertices).containing_indices(shape) diff --git a/autoarray/structures/triangles/jax_coordinate_array.py b/autoarray/structures/triangles/jax_coordinate_array.py index ce3e9e0d..2e1dd3c0 100644 --- a/autoarray/structures/triangles/jax_coordinate_array.py +++ b/autoarray/structures/triangles/jax_coordinate_array.py @@ -146,6 +146,8 @@ def _vertices_and_indices(self): axis=0, return_inverse=True, size=3 * self.coordinates.shape[0], + equal_nan=True, + fill_value=np.nan, ) indices = inverse_indices.reshape(-1, 3) return vertices, indices diff --git a/test_autoarray/structures/triangles/test_coordinate_implementation.py b/test_autoarray/structures/triangles/test_coordinate_implementation.py index d7c7a338..96b040b2 100644 --- a/test_autoarray/structures/triangles/test_coordinate_implementation.py +++ b/test_autoarray/structures/triangles/test_coordinate_implementation.py @@ -4,6 +4,7 @@ from autoarray.structures.triangles.abstract import HEIGHT_FACTOR from autoarray.structures.triangles.coordinate_array import CoordinateArrayTriangles +from autoarray.structures.triangles.shape import Point @pytest.fixture @@ -269,3 +270,47 @@ def test_for_limits_and_scale(): def test_means(one_triangle): assert np.all(one_triangle.means == [[0.0, -0.14433756729740643]]) + + +@pytest.mark.parametrize( + "x, y", + [ + (0.0, 0.0), + (-0.5, -HEIGHT_FACTOR / 2), + (0.5, -HEIGHT_FACTOR / 2), + (0.0, HEIGHT_FACTOR / 2), + ], +) +def test_containment(one_triangle, x, y): + assert one_triangle.containing_indices(Point(x, y)) == [0] + + +def test_triangles_touch(): + triangles = CoordinateArrayTriangles( + np.array([[0, 0], [2, 0]]), + ) + + assert max(triangles.triangles[0][:, 0]) == min(triangles.triangles[1][:, 0]) + + triangles = CoordinateArrayTriangles( + np.array([[0, 0], [0, 1]]), + ) + assert max(triangles.triangles[0][:, 1]) == min(triangles.triangles[1][:, 1]) + + +def test_from_grid_regression(): + triangles = CoordinateArrayTriangles.for_limits_and_scale( + x_min=-4.75, + x_max=4.75, + y_min=-4.75, + y_max=4.75, + scale=0.5, + ) + + x = triangles.vertices[:, 0] + assert min(x) <= -4.75 + assert max(x) >= 4.75 + + y = triangles.vertices[:, 1] + assert min(y) <= -4.75 + assert max(y) >= 4.75