diff --git a/src/probeinterface/probe.py b/src/probeinterface/probe.py index 5e205fa4..b334b4bd 100644 --- a/src/probeinterface/probe.py +++ b/src/probeinterface/probe.py @@ -809,7 +809,10 @@ def move(self, translation_vector: np.ndarray | list): self.probe_planar_contour += translation_vector def rotate( - self, theta: float, center: list | np.ndarray | None = None, axis: Literal["xy", "yz", "xz"] | None = None + self, + theta: float, + center: list | np.ndarray | None = None, + axis: Literal["xy", "yz", "xz"] | list | np.ndarray | None = None, ): """ Rotate the probe around a specified axis. @@ -820,10 +823,13 @@ def rotate( In degrees, anticlockwise/counterclockwise center : array | list | None, default: None Center of rotation. If None, the center of probe is used - axis : "xy" | "yz" | "xz" | None, default: None + axis : "xy" | "yz" | "xz" | array | list | None, default: None Axis of rotation. It must be None for 2D probes - It must be given for 3D probes + It must be given for 3D probes. + A plane name selects the axis normal to that plane, so "xy" rotates + about z, "yz" about x, and "xz" about y. A 3-element vector can be + given instead to rotate about an arbitrary axis. """ @@ -841,7 +847,7 @@ def rotate( R = _rotation_matrix_2d(theta) elif self.ndim == 3: assert axis is not None, "axis must be specified for 3d probes" - R = _rotation_matrix_3d(axis, theta).T + R = _rotation_matrix_3d(_axis_to_vector(axis), theta).T new_positions = (self.contact_positions - center) @ R + center @@ -1600,6 +1606,42 @@ def _rotation_matrix_2d(theta: float) -> np.ndarray: return R +_plane_to_rotation_axis = { + "xy": np.array([0.0, 0.0, 1.0]), + "yz": np.array([1.0, 0.0, 0.0]), + "xz": np.array([0.0, 1.0, 0.0]), +} + + +def _axis_to_vector(axis: str | np.ndarray | list) -> np.ndarray: + """ + Normalize the ``axis`` argument of :meth:`Probe.rotate` to a 3D vector. + + A plane name is mapped to the unit vector normal to that plane, so rotating + "in the xy plane" means rotating about z. + + Parameters + ---------- + axis : "xy" | "yz" | "xz" | np.array | list + Plane name or 3D axis of rotation + + Returns + ------- + axis : np.array + 3D axis of rotation + + """ + if isinstance(axis, str): + if axis not in _plane_to_rotation_axis: + raise ValueError(f"axis must be one of {list(_plane_to_rotation_axis)} or a 3-element vector, not {axis!r}") + return _plane_to_rotation_axis[axis] + + axis = np.asarray(axis, dtype="float64") + if axis.shape != (3,): + raise ValueError(f"axis must be a 3-element vector, not an array of shape {axis.shape}") + return axis + + def _rotation_matrix_3d(axis: np.ndarray | list, theta: float) -> np.ndarray: """ Returns 3D rotation matrix diff --git a/tests/test_probe.py b/tests/test_probe.py index 538e8837..136953ea 100644 --- a/tests/test_probe.py +++ b/tests/test_probe.py @@ -320,6 +320,49 @@ def test_copy_is_independent(): assert probe.contact_ids[0] == "c0" +def _probe_3d_for_rotation(): + probe = Probe(ndim=2) + positions = np.array([[0.0, 0.0], [10.0, 0.0], [0.0, 20.0], [10.0, 20.0]]) + probe.set_contacts(positions=positions, shapes="circle", shape_params={"radius": 5}) + return probe.to_3d(axes="xy") + + +@pytest.mark.parametrize( + "plane, vector", + [("xy", [0, 0, 1]), ("yz", [1, 0, 0]), ("xz", [0, 1, 0])], +) +def test_rotate_3d_accepts_plane_name(plane, vector): + """A plane name must rotate about the axis normal to that plane (issue #273).""" + from_plane = _probe_3d_for_rotation() + from_vector = _probe_3d_for_rotation() + + from_plane.rotate(theta=37, center=[0, 0, 0], axis=plane) + from_vector.rotate(theta=37, center=[0, 0, 0], axis=vector) + + assert np.allclose(from_plane.contact_positions, from_vector.contact_positions) + assert np.allclose(from_plane.contact_plane_axes, from_vector.contact_plane_axes) + + +def test_rotate_3d_plane_name_keeps_normal_coordinate_fixed(): + """Rotating in the xy plane leaves z untouched.""" + probe_3d = _probe_3d_for_rotation() + original_z = probe_3d.contact_positions[:, 2].copy() + + probe_3d.rotate(theta=90, center=[0, 0, 0], axis="xy") + + assert np.allclose(probe_3d.contact_positions[:, 2], original_z) + + +def test_rotate_3d_rejects_unknown_axis(): + probe_3d = _probe_3d_for_rotation() + + with pytest.raises(ValueError, match="axis must be one of"): + probe_3d.rotate(theta=90, center=[0, 0, 0], axis="zz") + + with pytest.raises(ValueError, match="axis must be a 3-element vector"): + probe_3d.rotate(theta=90, center=[0, 0, 0], axis=[0, 1]) + + if __name__ == "__main__": import tempfile