Add to_standard_orientation for SurfaceRZFourier

src/constellaration/geometry/surface_rz_fourier.py

@@ -1,3 +1,4 @@
+import logging
 from collections.abc import Mapping
 from typing import Sequence
 
@@ -13,6 +14,8 @@
 from constellaration.geometry import surface_utils
 from constellaration.utils.types import NpOrJaxArray, ScalarFloat
 
+logger = logging.getLogger(__name__)
+
 FourierCoefficients = jt.Float[np.ndarray, "n_poloidal_modes n_toroidal_modes"]
 FourierModes = jt.Int[np.ndarray, "n_poloidal_modes n_toroidal_modes"]
 
@@ -1513,3 +1516,102 @@ def from_points(
         surface.z_cos = split[3].reshape(shape)
 
     return surface, residual
+
+
+def _is_theta_external(surface: SurfaceRZFourier) -> bool:
+    """Determine if a toroidal surface has theta=0 on the outside.
+
+    "Outside" means the radial coordinate at theta=0 is greater than at theta=pi.
+    """
+    rz_theta0 = evaluate_points_rz(surface, np.array([[0.0, 0.0]]))
+    rz_thetapi = evaluate_points_rz(surface, np.array([[np.pi, 0.0]]))
+    return bool(rz_theta0[0, 0] >= rz_thetapi[0, 0])
+
+
+def _is_theta_counterclockwise(surface: SurfaceRZFourier) -> bool:
+    """Determine if theta increases in the counterclockwise direction in an r-z plane.
+
+    Evaluated at the outboard side: counterclockwise means dz/dtheta >= 0.
+    """
+    theta_outside = 0.0 if _is_theta_external(surface) else np.pi
+    dz_dtheta = _evaluate_dz_dtheta(surface, np.array([[theta_outside, 0.0]]))
+    return bool(dz_dtheta[0] >= 0)
+
+
+def _surface_shift_theta_by_pi(surface: SurfaceRZFourier) -> SurfaceRZFourier:
+    """Return a copy of the surface where theta=0 is shifted by pi.
+
+    Due to Fourier properties, shifting theta by pi multiplies each coefficient
+    by (-1)^m where m is the poloidal mode number.
+    """
+    sign = (-1) ** surface.poloidal_modes
+    r_cos = sign * surface.r_cos
+    z_sin = sign * surface.z_sin
+    if not surface.is_stellarator_symmetric:
+        assert surface.r_sin is not None
+        assert surface.z_cos is not None
+        r_sin = sign * surface.r_sin
+        z_cos = sign * surface.z_cos
+    else:
+        r_sin = None
+        z_cos = None
+    return surface.model_copy(
+        update=dict(
+            r_cos=r_cos,
+            z_sin=z_sin,
+            r_sin=r_sin,
+            z_cos=z_cos,
+        )
+    )
+
+
+def _surface_change_theta_sign(surface: SurfaceRZFourier) -> SurfaceRZFourier:
+    """Return a copy of the surface where the direction of theta is reversed.
+
+    Flipping theta is equivalent to replacing n -> -n in the Fourier representation,
+    which reverses the toroidal mode array. For sine terms the sign is also flipped.
+    The m=0 modes are unaffected by the theta flip and must be restored.
+    """
+    r_cos = np.asarray(surface.r_cos)[:, ::-1]
+    z_sin = -np.asarray(surface.z_sin)[:, ::-1]
+    if surface.is_stellarator_symmetric:
+        # m=0 modes are not affected by theta sign flip; undo the reversal
+        r_cos = r_cos.copy()
+        z_sin = z_sin.copy()
+        r_cos[0] = r_cos[0, ::-1]
+        z_sin[0] = -z_sin[0, ::-1]
+    if not surface.is_stellarator_symmetric:
+        assert surface.r_sin is not None
+        assert surface.z_cos is not None
+        r_sin = -np.asarray(surface.r_sin)[:, ::-1]
+        z_cos = np.asarray(surface.z_cos)[:, ::-1]
+    else:
+        r_sin = None
+        z_cos = None
+    return surface.model_copy(
+        update=dict(
+            r_cos=r_cos,
+            z_sin=z_sin,
+            r_sin=r_sin,
+            z_cos=z_cos,
+        )
+    )
+
+
+def to_standard_orientation(surface: SurfaceRZFourier) -> SurfaceRZFourier:
+    """Return a copy of the surface where theta=0 is on the outside and theta increases
+    in the counterclockwise direction (in an r-z cross-section).
+
+    This normalization only applies to stellarator-symmetric surfaces. Non-symmetric
+    surfaces are returned unchanged.
+    """
+    if not surface.is_stellarator_symmetric:
+        logger.warning(
+            "The surface is not stellarator symmetric, therefore it was not modified."
+        )
+        return surface
+    if not _is_theta_external(surface):
+        surface = _surface_shift_theta_by_pi(surface)
+    if not _is_theta_counterclockwise(surface):
+        surface = _surface_change_theta_sign(surface)
+    return surface

tests/geometry/surface_rz_fourier_test.py

@@ -1090,3 +1090,97 @@ def test_compute_normal_displacement_raises_for_mismatched_parameters() -> None:
             n_poloidal_points=3,
             n_toroidal_points=2,
         )
+
+
+# ---------- to_standard_orientation tests ----------
+
+
+def test_to_standard_orientation_torus_external_ccw() -> None:
+    """Surface already in standard orientation is returned unchanged."""
+    major_radius = 1.5
+    minor_radius = 0.5
+    r_cos = np.array([[major_radius], [minor_radius]])
+    z_sin = np.array([[0.0], [minor_radius]])
+    surface = surface_rz_fourier.SurfaceRZFourier(
+        r_cos=r_cos, z_sin=z_sin, is_stellarator_symmetric=True
+    )
+    assert surface_rz_fourier._is_theta_external(surface)
+    assert surface_rz_fourier._is_theta_counterclockwise(surface)
+    surface_std = surface_rz_fourier.to_standard_orientation(surface)
+    np.testing.assert_allclose(surface_std.r_cos, surface.r_cos)
+    np.testing.assert_allclose(surface_std.z_sin, surface.z_sin)
+    assert surface_rz_fourier._is_theta_external(surface_std)
+    assert surface_rz_fourier._is_theta_counterclockwise(surface_std)
+
+
+def test_to_standard_orientation_torus_internal_ccw() -> None:
+    """theta=0 on the inside, counterclockwise: needs a theta shift by pi."""
+    major_radius = 1.5
+    minor_radius = 0.5
+    r_cos = np.array([[major_radius], [-minor_radius]])
+    z_sin = np.array([[0.0], [-minor_radius]])
+    surface = surface_rz_fourier.SurfaceRZFourier(
+        r_cos=r_cos, z_sin=z_sin, is_stellarator_symmetric=True
+    )
+    assert not surface_rz_fourier._is_theta_external(surface)
+    assert surface_rz_fourier._is_theta_counterclockwise(surface)
+    surface_std = surface_rz_fourier.to_standard_orientation(surface)
+    assert surface_rz_fourier._is_theta_external(surface_std)
+    assert surface_rz_fourier._is_theta_counterclockwise(surface_std)
+    # Verify geometric equivalence: original at theta evaluates to same xyz
+    # as standardized at theta+pi
+    rng = np.random.default_rng(42)
+    theta_phi = rng.random((7, 8, 2)) * 2 * np.pi
+    points = surface_rz_fourier.evaluate_points_xyz(surface, theta_phi)
+    theta_phi_std = theta_phi.copy()
+    theta_phi_std[:, :, 0] += np.pi
+    points_std = surface_rz_fourier.evaluate_points_xyz(surface_std, theta_phi_std)
+    np.testing.assert_allclose(points, points_std, atol=1e-12)
+
+
+def test_to_standard_orientation_torus_external_cw() -> None:
+    """theta=0 on the outside, clockwise: needs theta sign flip."""
+    major_radius = 1.5
+    minor_radius = 0.5
+    r_cos = np.array([[major_radius], [minor_radius]])
+    z_sin = np.array([[0.0], [-minor_radius]])
+    surface = surface_rz_fourier.SurfaceRZFourier(
+        r_cos=r_cos, z_sin=z_sin, is_stellarator_symmetric=True
+    )
+    assert surface_rz_fourier._is_theta_external(surface)
+    assert not surface_rz_fourier._is_theta_counterclockwise(surface)
+    surface_std = surface_rz_fourier.to_standard_orientation(surface)
+    assert surface_rz_fourier._is_theta_external(surface_std)
+    assert surface_rz_fourier._is_theta_counterclockwise(surface_std)
+    # Verify geometric equivalence: original at theta = standardized at -theta
+    rng = np.random.default_rng(42)
+    theta_phi = rng.random((7, 8, 2)) * 2 * np.pi
+    points = surface_rz_fourier.evaluate_points_xyz(surface, theta_phi)
+    theta_phi_std = theta_phi.copy()
+    theta_phi_std[:, :, 0] = -theta_phi_std[:, :, 0]
+    points_std = surface_rz_fourier.evaluate_points_xyz(surface_std, theta_phi_std)
+    np.testing.assert_allclose(points, points_std, atol=1e-12)
+
+
+def test_to_standard_orientation_torus_internal_cw() -> None:
+    """theta=0 on the inside, clockwise: needs both shift and sign flip."""
+    major_radius = 1.5
+    minor_radius = 0.5
+    r_cos = np.array([[major_radius], [-minor_radius]])
+    z_sin = np.array([[0.0], [minor_radius]])
+    surface = surface_rz_fourier.SurfaceRZFourier(
+        r_cos=r_cos, z_sin=z_sin, is_stellarator_symmetric=True
+    )
+    assert not surface_rz_fourier._is_theta_external(surface)
+    assert not surface_rz_fourier._is_theta_counterclockwise(surface)
+    surface_std = surface_rz_fourier.to_standard_orientation(surface)
+    assert surface_rz_fourier._is_theta_external(surface_std)
+    assert surface_rz_fourier._is_theta_counterclockwise(surface_std)
+    # Verify geometric equivalence: original at theta = standardized at -theta+pi
+    rng = np.random.default_rng(42)
+    theta_phi = rng.random((7, 8, 2)) * 2 * np.pi
+    points = surface_rz_fourier.evaluate_points_xyz(surface, theta_phi)
+    theta_phi_std = theta_phi.copy()
+    theta_phi_std[:, :, 0] = -theta_phi_std[:, :, 0] + np.pi
+    points_std = surface_rz_fourier.evaluate_points_xyz(surface_std, theta_phi_std)
+    np.testing.assert_allclose(points, points_std, atol=1e-12)