Accept plane names for the axis argument of Probe.rotate

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

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