diff --git a/predicators/utils.py b/predicators/utils.py index 8584cc8e64..c3012e53d7 100644 --- a/predicators/utils.py +++ b/predicators/utils.py @@ -34,6 +34,7 @@ import pathos.multiprocessing as mp from gym.spaces import Box from matplotlib import patches +from numpy.typing import NDArray from pyperplan.heuristics.heuristic_base import \ Heuristic as _PyperplanBaseHeuristic from pyperplan.planner import HEURISTICS as _PYPERPLAN_HEURISTICS @@ -502,6 +503,21 @@ def from_center(center_x: float, center_y: float, width: float, return norm_rect.rotate_about_point(center_x, center_y, rotation_about_center) + @functools.cached_property + def rotation_matrix(self) -> NDArray[np.float64]: + """Get the rotation matrix.""" + return np.array([[np.cos(self.theta), -np.sin(self.theta)], + [np.sin(self.theta), + np.cos(self.theta)]]) + + @functools.cached_property + def inverse_rotation_matrix(self) -> NDArray[np.float64]: + """Get the inverse rotation matrix.""" + return np.array([[np.cos(self.theta), + np.sin(self.theta)], + [-np.sin(self.theta), + np.cos(self.theta)]]) + @functools.cached_property def vertices(self) -> List[Tuple[float, float]]: """Get the four vertices for the rectangle.""" @@ -509,9 +525,6 @@ def vertices(self) -> List[Tuple[float, float]]: [self.width, 0], [0, self.height], ]) - rotate_matrix = np.array([[np.cos(self.theta), -np.sin(self.theta)], - [np.sin(self.theta), - np.cos(self.theta)]]) translate_vector = np.array([self.x, self.y]) vertices = np.array([ (0, 0), @@ -520,7 +533,7 @@ def vertices(self) -> List[Tuple[float, float]]: (1, 0), ]) vertices = vertices @ scale_matrix.T - vertices = vertices @ rotate_matrix.T + vertices = vertices @ self.rotation_matrix.T vertices = translate_vector + vertices # Convert to a list of tuples. Slightly complicated to appease both # type checking and linting. @@ -549,26 +562,22 @@ def circumscribed_circle(self) -> Circle: return Circle(x, y, radius) def contains_point(self, x: float, y: float) -> bool: - rotate_matrix = np.array([[np.cos(self.theta), - np.sin(self.theta)], - [-np.sin(self.theta), - np.cos(self.theta)]]) - rx, ry = np.array([x - self.x, y - self.y]) @ rotate_matrix.T + # First invert translation, then invert rotation. + rx, ry = np.array([x - self.x, y - self.y + ]) @ self.inverse_rotation_matrix.T return 0 <= rx <= self.width and \ 0 <= ry <= self.height def sample_random_point(self, rng: np.random.Generator) -> Tuple[float, float]: - rotate_matrix = np.array([[np.cos(self.theta), - np.sin(self.theta)], - [-np.sin(self.theta), - np.cos(self.theta)]]) rand_width = rng.uniform(0, self.width) rand_height = rng.uniform(0, self.height) - rx, ry = np.array([self.x + rand_width, self.y + rand_height - ]) @ rotate_matrix.T - assert self.contains_point(rx, ry) - return (rx, ry) + # First rotate, then translate. + rx, ry = np.array([rand_width, rand_height]) @ self.rotation_matrix.T + x = rx + self.x + y = ry + self.y + assert self.contains_point(x, y) + return (x, y) def rotate_about_point(self, x: float, y: float, rot: float) -> Rectangle: """Create a new rectangle that is this rectangle, but rotated CCW by diff --git a/tests/test_utils.py b/tests/test_utils.py index 54636029b2..2a295b158e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -418,10 +418,10 @@ def test_rectangle(): assert rect7.center == (1, 2) rng = np.random.default_rng(0) - for _ in range(10): - p7 = rect7.sample_random_point(rng) - assert rect7.contains_point(p7[0], p7[1]) - plt.plot(p7[0], p7[1], 'bo') + for _ in range(100): + p5 = rect5.sample_random_point(rng) + assert rect5.contains_point(p5[0], p5[1]) + plt.plot(p5[0], p5[1], 'bo') # Uncomment for debugging. # plt.savefig("/tmp/rectangle_unit_test.png")