diff --git a/src/constellaration/geometry/surface_rz_fourier.py b/src/constellaration/geometry/surface_rz_fourier.py index 33d39c3..f8209a6 100644 --- a/src/constellaration/geometry/surface_rz_fourier.py +++ b/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 diff --git a/tests/geometry/surface_rz_fourier_test.py b/tests/geometry/surface_rz_fourier_test.py index 9c03814..f43d0f2 100644 --- a/tests/geometry/surface_rz_fourier_test.py +++ b/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)