diff --git a/CHANGELOG.md b/CHANGELOG.md index 2229f4d75b..2f9776723b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ Performance Improvements - Improves memory management to reduce the base memory used during optimization while using `lsq-exact`, `lsq-auglag` and `fmin-auglag` optimizers. - Speeds up ``field_line_integrate`` and ``trace_particles`` for filamentary coils (``Coil``, ``CoilSet``, ``MixedCoilSet``) by precomputing the constant source information, so that the ODE right hand side only evaluates a single fused Biot-Savart kernel instead of recomputing the coil geometry at every solver step. - Improves the non-singular Biot-Savart kernel which should give a speed/memory improvement to objectives that compute magnetic field from coils such as ``QuadraticFlux``. +- Adds classes ``SurfaceCurve``, ``FourierRZSurfaceCurve``, and ``FourierRZSurfaceCoil`` to describe curves and coils constrained to lie on given FourierRZToroidalSurface objects. A new objective ``SurfaceCurveConsistency`` facilitates joint optimization of the curve and underlying surface. Bug Fixes diff --git a/desc/coils.py b/desc/coils.py index ee7c9a051d..5ab54057f3 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -23,9 +23,11 @@ ) from desc.compute import get_params from desc.compute.utils import _compute as compute_fun +from desc.equilibrium.coords import map_coordinates from desc.geometry import ( FourierPlanarCurve, FourierRZCurve, + FourierRZSurfaceCurve, FourierXYCurve, FourierXYZCurve, SplineXYZCurve, @@ -889,6 +891,132 @@ def to_FourierXY(self, N=10, grid=None, s=None, basis="xyz", name="", **kwargs): self.current, coords, N=N, s=s, basis=basis, name=name ) + def to_FourierRZSurface( + self, + surface=None, + equilibrium=None, + NFP=1, + N_theta=10, + N_zeta=10, + secular_theta=None, + secular_zeta=None, + sym_theta=None, + sym_zeta=None, + grid=None, + s=None, + basis="xyz", + name="", + **kwargs, + ): + """Convert Coil to FourierRZSurface representation. + + Coils of this type are constraint to lie on a FourierRZToroidalSurface. + This method therefore requires either a surface or equilibrium. Moreover, + the secular terms in the representations theta(s) = secular_theta*..., + zeta(s) = secular_zeta*..., are required. + + Parameters + ---------- + surface: FourierRZToroidalSurface + Surface on which curve lies. Only one of "surface" and "equilibrium" + can be provided. + equilibrium: Equilibrium + Provided in place of "surface" when the desired surface is the rho=1 + flux surface of an equilibrium. Only one of "surface" and "equilibrium" + can be provided. + NFP: int, optional + Field period symmetry. Defaults to 1. + N_theta: int or None, optional + Max Fourier mode number for theta(s). Set to None to indicate no basis + for theta should be included. Default is 10. + N_zeta: int or None, optional + Max Fourier mode number for zeta(s). Set to None to indicate no basis + for zeta should be included. Default is 10. + secular_theta: int or None, optional + Coefficient of secular term in theta. Required for this method. + secular_zeta: int or None, optional + Coefficient of secular term in zeta. Required for this method. + sym_theta: str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for theta(s). Default is None. + sym_zeta: str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for zeta(s). Default is None. + grid: Grid, int or None, optional + Grid used to evaluate curve coordinates on for fitting. + If an integer, uses that many equally spaced points. + s: ndarray + Arbitrary curve parameter to use for the fitting. + Should be monotonic, 1D array of same length as + coords. if None, defaults linearly spaced in [0,2pi) + basis : {'xyz', 'rpz'} + Coordinate system in which the curve is evaluated before being mapped + onto the surface. Default = 'xyz'. + name: str, optional + Name for this coil. + + Returns + ------- + FourierRZSurfaceCoil + FourierRZSurfaceCoil object fit to the given coords. + """ + errorif( + surface is not None and equilibrium is not None, + ValueError, + "Either surface or equilibrium is required", + ) + errorif( + secular_theta is None or secular_zeta is None, + ValueError, + "Must provide both secular_theta and secular_zeta", + ) + source_R_basis = ( + surface.R_basis if surface is not None else equilibrium.surface.R_basis + ) + source_Z_basis = ( + surface.Z_basis if surface is not None else equilibrium.surface.Z_basis + ) + M_source = max(source_R_basis.M, source_Z_basis.M) + N_source = max(source_R_basis.N, source_Z_basis.N) + if equilibrium is None: + from .equilibrium import Equilibrium + + equilibrium = Equilibrium(surface=surface, L=1, M=M_source, N=N_source) + + # Compute rpz values along curve + if grid is None and s is not None: + grid = LinearGrid(zeta=s) + elif grid is None: + N = max(N_theta, N_zeta) + N_eff = N * NFP + M_source * secular_theta + N_source * secular_zeta + grid = LinearGrid(N=2 * N_eff + 5) + + coords_RpZ = self.compute("x", grid=grid, basis=basis)["x"] + if basis.lower() == "xyz": + coords_RpZ = xyz2rpz(coords_RpZ) + coords_rtz = map_coordinates(equilibrium, coords_RpZ, inbasis=("R", "phi", "Z")) + coords_rtz = coords_rtz.at[:, 1:].set(np.unwrap(coords_rtz[:, 1:], axis=0)) + + if s is None: + s = np.linspace(0, 2 * np.pi, coords_rtz.shape[0], endpoint=False) + + coords = np.hstack([s[:, np.newaxis], coords_rtz[:, 1:]]) + + return FourierRZSurfaceCoil.from_values( + self.current, + coords=coords, + surface=surface, + equilibrium=equilibrium if surface is None else None, + NFP=NFP, + N_theta=N_theta, + N_zeta=N_zeta, + secular_theta=secular_theta, + secular_zeta=secular_zeta, + sym_theta=sym_theta, + sym_zeta=sym_zeta, + name=name, + ) + class FourierRZCoil(_Coil, FourierRZCurve): """Coil parameterized by fourier series for R,Z in terms of toroidal angle phi. @@ -1589,6 +1717,210 @@ def from_values( ) +class FourierRZSurfaceCoil(_Coil, FourierRZSurfaceCurve): + r"""Coil that is contained in a FourierRZToroidalSurface. + + Parameterized by a pair of integers (secular terms for toroidal + and poloidal winding) and Fourier series for theta(s), zeta(s). + + Parameters + ---------- + current : float + Current through the coil, in Amperes. + surface: FourierRZToroidalSurface, optional + Underlying surface which the coil lies on. + If None, must pass a value for "equilibrium." + equilibrium: Equilibrium, optional + Used in place of "surface" when the coil is intended + to live on the rho=1 surface of an equilibrium which + may be optimized alongside the coil. + Must pass exactly one of "surface" or "equilibrium." + secular_theta: non-negative int, optional + Net winding of theta(s), equivalent to + 1/2pi \int_{0}^{2pi} theta'(s) ds. Default + value is 0. Nonzero values result in coils + linking the torus poloidally. + Note: gcd(secular_theta, secular_zeta) should equal 1. + secular_zeta: non-negative int, optional + Net winding of zeta(s), equivalent to + 1/2pi \int_{0}^{2pi} zeta'(s) ds. Default + value is 1. Nonzero values result in coils + linking the torus toroidally. + Note: gcd(secular_theta, secular_zeta) should equal 1. + theta_n: array-like, optional + Coefficients of Fourier modes defining + theta(s). + zeta_n: array-like, optional + Coefficients of Fourier modes defining + zeta(s). + modes_theta: array-like, optional + Mode numbers associated with theta. If + not given, defaults to [-N_theta:N_theta], + where N_theta = len(theta_n)//2. + modes_zeta: array-like, optional + Mode numbers associated with zeta. If + not given, defaults to [-N_theta:N_theta], + where N_theta = len(theta_n)//2. + sym_theta: str, optional + If "sin"/"cos", retains only the corresponding + modes in the series for theta(s). Defaults + to None, i.e. all modes are kept. + sym_zeta: bool, optional + If "sin"/"cos", retains only the corresponding + modes in the series for zeta(s). Defaults + to None, i.e. all modes are kept. + NFP: positive int, optional + Field period symmetry of the coil. Defaults to 1. + The curve's image is invariant + under zeta -> zeta + 2pi/NFP, and Fourier modes + are multiples of NFP in s. + Note: If changed, must satisfy (a) NFP | surface.NFP, + (b) NFP | secular_theta, (c) secular_zeta \neq 0. + name: str, optional + Name for this coil. + """ + + _io_attrs_ = _Coil._io_attrs_ + FourierRZSurfaceCurve._io_attrs_ + _static_attrs = _Coil._static_attrs + FourierRZSurfaceCurve._static_attrs + + def __init__( + self, + current=1, + surface=None, + equilibrium=None, + secular_theta=0, + secular_zeta=1, + theta_n=[0.0], + zeta_n=[0.0], + modes_theta=None, + modes_zeta=None, + sym_theta=False, + sym_zeta=False, + NFP=1, + name="", + ): + super().__init__( + current=current, + surface=surface, + equilibrium=equilibrium, + secular_theta=secular_theta, + secular_zeta=secular_zeta, + theta_n=theta_n, + zeta_n=zeta_n, + modes_theta=modes_theta, + modes_zeta=modes_zeta, + sym_theta=sym_theta, + sym_zeta=sym_zeta, + NFP=NFP, + name=name, + ) + + def _compute_A_or_B( + self, + coords, + params=None, + basis="rpz", + source_grid=None, + transforms=None, + compute_A_or_B="B", + chunk_size=None, + ): + """Compute magnetic field or vector potential at a set of points.""" + if source_grid is None: + source_grid = LinearGrid(N=2 * self.N_effective + 5) + return super()._compute_A_or_B( + coords, params, basis, source_grid, transforms, compute_A_or_B, chunk_size + ) + + @classmethod + def from_values( + cls, + current, + coords, + surface=None, + equilibrium=None, + NFP=1, + N_theta=10, + N_zeta=10, + secular_theta=None, + secular_zeta=None, + sym_theta=None, + sym_zeta=None, + name="", + ): + """Fit coordinates to FourierRZSurfaceCoil object. + + Parameters + ---------- + current : float + Current through the coil, in Amperes. + coords: array-like, shape (N,3) + Default shape (N,3) corresponds to coordinates (s,theta(s),zeta(s)) at N + points along curve. Assumed to be ordered by increasing curve parameter. + surface: FourierRZToroidalSurface + Surface on which curve lies. Only one of "surface" and "equilibrium" + can be provided. + equilibrium: Equilibrium + Provided in place of "surface" when the desired surface is the rho=1 + flux surface of an equilibrium. Only one of "surface" and "equilibrium" + can be provided. + NFP: int, optional + Field period symmetry. Defaults to 1. + N_theta: int or None, optional + Max Fourier mode number for theta(s). Set to None to indicate no basis + for theta should be included. Default is 10. + N_zeta: int or None, optional + Max Fourier mode number for zeta(s). Set to None to indicate no basis + for zeta should be included. Default is 10. + secular_theta: int or None, optional + Coefficient of secular term in theta. If None, estimated automatically. + secular_zeta: int or None, optional + Coefficient of secular term in zeta. If None, estimated automatically. + sym_theta: str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for theta(s). Default is None. + sym_zeta: str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for zeta(s). Default is None. + name: str, optional + Name for this curve. + + Returns + ------- + curve : FourierRZSurfaceCoil + New representation of the coil parameterized by a Fourier series + for theta(s), zeta(s). + """ + curve = super().from_values( + coords, + surface=surface, + equilibrium=equilibrium, + NFP=NFP, + N_theta=N_theta, + N_zeta=N_zeta, + secular_theta=secular_theta, + secular_zeta=secular_zeta, + sym_theta=sym_theta, + sym_zeta=sym_zeta, + name="", + ) + return cls( + current=current, + surface=surface, + equilibrium=equilibrium, + secular_theta=curve.secular_theta, + secular_zeta=curve.secular_zeta, + theta_n=curve.theta_n, + zeta_n=curve.zeta_n, + modes_theta=curve.theta_basis.modes[:, 2], + modes_zeta=curve.zeta_basis.modes[:, 2], + sym_theta=sym_theta, + sym_zeta=sym_zeta, + NFP=NFP, + name=name, + ) + + def _check_type(coil0, coil): errorif( not isinstance(coil, coil0.__class__), @@ -1612,6 +1944,15 @@ def _check_type(coil0, coil): FourierXYCoil: ["X_basis", "Y_basis"], FourierXYZCoil: ["X_basis", "Y_basis", "Z_basis"], SplineXYZCoil: ["method", "N", "knots"], + FourierRZSurfaceCoil: [ + "theta_basis", + "zeta_basis", + "R_basis", + "Z_basis", + "NFP", + "secular_theta", + "secular_zeta", + ], } for attr in attrs[coil0.__class__]: @@ -2765,6 +3106,101 @@ def to_SplineXYZ( check_intersection=check_intersection, ) + def to_FourierRZSurface( + self, + surface=None, + equilibrium=None, + NFP=1, + N_theta=10, + N_zeta=10, + secular_theta=None, + secular_zeta=None, + sym_theta=None, + sym_zeta=None, + grid=None, + s=None, + basis="xyz", + name="", + check_intersection=False, + ): + """Convert all coils to FourierRZSurfaceCoil representation. + + Note that every coil in the set must lie on the given surface. + + Parameters + ---------- + surface: FourierRZToroidalSurface + Surface on which the coils lie. Only one of "surface" and "equilibrium" + can be provided. + equilibrium: Equilibrium + Provided in place of "surface" when the desired surface is the rho=1 + flux surface of an equilibrium. Only one of "surface" and "equilibrium" + can be provided. + NFP : int, optional + Field period symmetry of the new coils. Defaults to 1. + N_theta : int or None, optional + Max Fourier mode number for theta(s). Set to None to indicate no basis + for theta should be included. Default is 10. + N_zeta : int or None, optional + Max Fourier mode number for zeta(s). Set to None to indicate no basis + for zeta should be included. Default is 10. + secular_theta : int + Net poloidal winding of the coils. Required. + secular_zeta : int + Net toroidal winding of the coils. Required. + sym_theta : str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for theta(s). Default is None. + sym_zeta : str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for zeta(s). Default is None. + grid : Grid, int or None + Grid used to evaluate curve coordinates on to fit with + FourierRZSurfaceCoil. If an integer, uses that many equally spaced points. + s : ndarray + Arbitrary curve parameter to use for the fitting. Should be monotonic, + 1D array of same length as coords. If None, defaults to linearly spaced + in [0, 2pi). + basis : {'xyz', 'rpz'} + Coordinate system for the curve coordinates. Default = 'xyz'. + name : str + Name for this coilset. + check_intersection: bool + Whether or not to check the coils in the new coilset for intersections. + Defaults to False. + + Returns + ------- + coilset : CoilSet + New representation of the coilset parameterized by Fourier series for + theta(s), zeta(s) on the given surface. + + """ + coils = [ + coil.to_FourierRZSurface( + surface=surface, + equilibrium=equilibrium, + NFP=NFP, + N_theta=N_theta, + N_zeta=N_zeta, + secular_theta=secular_theta, + secular_zeta=secular_zeta, + sym_theta=sym_theta, + sym_zeta=sym_zeta, + grid=grid, + s=s, + basis=basis, + ) + for coil in self + ] + return self.__class__( + *coils, + NFP=self.NFP, + sym=self.sym, + name=name, + check_intersection=check_intersection, + ) + def is_self_intersecting(self, grid=None, tol=None): """Check if any coils in the CoilSet intersect. @@ -3431,6 +3867,96 @@ def to_SplineXYZ( ] return self.__class__(*coils, name=name, check_intersection=check_intersection) + def to_FourierRZSurface( + self, + surface=None, + equilibrium=None, + NFP=1, + N_theta=10, + N_zeta=10, + secular_theta=None, + secular_zeta=None, + sym_theta=None, + sym_zeta=None, + grid=None, + s=None, + basis="xyz", + name="", + check_intersection=False, + ): + """Convert all coils to FourierRZSurfaceCoil representation. + + Note that every coil in the set must lie on the given surface. + + Parameters + ---------- + surface: FourierRZToroidalSurface + Surface on which the coils lie. Only one of "surface" and "equilibrium" + can be provided. + equilibrium: Equilibrium + Provided in place of "surface" when the desired surface is the rho=1 + flux surface of an equilibrium. Only one of "surface" and "equilibrium" + can be provided. + NFP : int, optional + Field period symmetry of the new coils. Defaults to 1. + N_theta : int or None, optional + Max Fourier mode number for theta(s). Set to None to indicate no basis + for theta should be included. Default is 10. + N_zeta : int or None, optional + Max Fourier mode number for zeta(s). Set to None to indicate no basis + for zeta should be included. Default is 10. + secular_theta : int + Net poloidal winding of the coils. Required. + secular_zeta : int + Net toroidal winding of the coils. Required. + sym_theta : str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for theta(s). Default is None. + sym_zeta : str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for zeta(s). Default is None. + grid : Grid, int or None + Grid used to evaluate curve coordinates on to fit with + FourierRZSurfaceCoil. If an integer, uses that many equally spaced points. + s : ndarray + Arbitrary curve parameter to use for the fitting. Should be monotonic, + 1D array of same length as coords. If None, defaults to linearly spaced + in [0, 2pi). + basis : {'xyz', 'rpz'} + Coordinate system for the curve coordinates. Default = 'xyz'. + name : str + Name for this coilset. + check_intersection: bool + Whether or not to check the coils in the new coilset for intersections. + Defaults to False. + + Returns + ------- + coilset : MixedCoilSet + New representation of the coilset parameterized by Fourier series for + theta(s), zeta(s) on the given surface. + + """ + coils = [ + coil.to_FourierRZSurface( + surface=surface, + equilibrium=equilibrium, + NFP=NFP, + N_theta=N_theta, + N_zeta=N_zeta, + secular_theta=secular_theta, + secular_zeta=secular_zeta, + sym_theta=sym_theta, + sym_zeta=sym_zeta, + grid=grid, + s=s, + basis=basis, + check_intersection=check_intersection, + ) + for coil in self + ] + return self.__class__(*coils, name=name, check_intersection=check_intersection) + def __add__(self, other): if isinstance(other, (CoilSet, MixedCoilSet)): return MixedCoilSet(*self.coils, *other.coils) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index ef2c5b3aca..0c509c4b59 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -2,6 +2,7 @@ from desc.backend import jnp, sign +from ..grid import Grid from ..utils import ( cross, dot, @@ -1237,3 +1238,532 @@ def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs): # but also works if grid.endpoint is False data["length"] = jnp.sum(T * data["ds"]) return data + + +@register_compute_fun( + name="center", + label="\\langle\\mathbf{x}\\rangle", + units="m", + units_long="meters", + description="Centroid of the curve", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="s", + data=["x", "x_s", "ds"], + parameterization="desc.geometry.core.SurfaceCurve", +) +def _center_SurfaceCurve(params, transforms, profiles, data, **kwargs): + # weight by arclength + xyz = rpz2xyz(data["x"]) + w = jnp.linalg.norm(data["x_s"], axis=-1) * data["ds"] + center = jnp.sum(xyz * w[:, jnp.newaxis], axis=0) / jnp.sum(w) + data["center"] = xyz2rpz(center) * jnp.ones_like(data["x"]) + return data + + +@register_compute_fun( + name="x", + label="\\mathbf{x}", + units="~", + units_long="not applicable", + description="Coordinate triplet. " + "This is not a position vector unless basis is cartesian. " + "When basis is cartesian, the units are meters.", + dim=3, + params=["R_lmn", "Z_lmn", "rotmat", "shift"], + transforms={"surface": []}, + profiles=[], + coordinates="s", + data=["theta", "zeta"], + parameterization="desc.geometry.core.SurfaceCurve", +) +def _x_SurfaceCurve(params, transforms, profiles, data, **kwargs): + nodes = jnp.vstack([jnp.ones_like(data["theta"]), data["theta"], data["zeta"]]).T + grid = Grid(nodes, sort=False, jitable=True) + params_temp = transforms["surface"].params_dict.copy() + params_temp["R_lmn"] = params["R_lmn"] + params_temp["Z_lmn"] = params["Z_lmn"] + + data_surf = transforms["surface"].compute( + ["R", "Z"], + grid=grid, + method="jitable", + params=params_temp, + ) + + # Assuming zeta=phi + coords = jnp.stack([data_surf["R"], data["zeta"], data_surf["Z"]], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz(coords) + coords = ( + coords @ params["rotmat"].reshape((3, 3)).T + params["shift"][jnp.newaxis, :] + ) + # convert back to rpz + coords = xyz2rpz(coords) + data["x"] = coords + return data + + +@register_compute_fun( + name="x_s", + label="\\partial_{s} \\mathbf{x}", + units="m", + units_long="meters", + description="Position vector along curve, first derivative", + dim=3, + params=["R_lmn", "Z_lmn", "rotmat"], + transforms={"surface": []}, + profiles=[], + coordinates="s", + data=["theta", "zeta", "theta_s", "zeta_s", "phi"], + parameterization="desc.geometry.core.SurfaceCurve", +) +def _x_s_SurfaceCurve(params, transforms, profiles, data, **kwargs): + nodes = jnp.vstack([jnp.ones_like(data["theta"]), data["theta"], data["zeta"]]).T + grid = Grid(nodes, sort=False, jitable=True) + params_temp = transforms["surface"].params_dict.copy() + params_temp["R_lmn"] = params["R_lmn"] + params_temp["Z_lmn"] = params["Z_lmn"] + names = ["R", "R_t", "R_z", "Z_t", "Z_z"] + data_surf = transforms["surface"].compute( + names, + grid=grid, + method="jitable", + params=params_temp, + ) + + xs_R = data_surf["R_t"] * data["theta_s"] + data_surf["R_z"] * data["zeta_s"] + xs_p = data_surf["R"] * data["zeta_s"] + xs_Z = data_surf["Z_t"] * data["theta_s"] + data_surf["Z_z"] * data["zeta_s"] + + coords = jnp.stack([xs_R, xs_p, xs_Z], axis=1) + coords = rpz2xyz_vec(coords, phi=data["zeta"]) + coords = coords @ params["rotmat"].reshape((3, 3)).T + coords = xyz2rpz_vec(coords, phi=data["phi"]) + data["x_s"] = coords + return data + + +@register_compute_fun( + name="x_ss", + label="\\partial_{ss} \\mathbf{x}", + units="m", + units_long="meters", + description="Position vector along curve, second derivative", + dim=3, + params=["R_lmn", "Z_lmn", "rotmat"], + transforms={"surface": []}, + profiles=[], + coordinates="s", + data=["theta", "zeta", "theta_s", "zeta_s", "theta_ss", "zeta_ss", "phi"], + parameterization="desc.geometry.core.SurfaceCurve", +) +def _x_ss_SurfaceCurve(params, transforms, profiles, data, **kwargs): + nodes = jnp.vstack([jnp.ones_like(data["theta"]), data["theta"], data["zeta"]]).T + grid = Grid(nodes, sort=False, jitable=True) + params_temp = transforms["surface"].params_dict.copy() + params_temp["R_lmn"] = params["R_lmn"] + params_temp["Z_lmn"] = params["Z_lmn"] + names = [ + "R", + "R_t", + "R_z", + "Z_t", + "Z_z", + "R_tt", + "R_tz", + "R_zz", + "Z_tt", + "Z_tz", + "Z_zz", + ] + data_surf = transforms["surface"].compute( + names, + grid=grid, + method="jitable", + params=params_temp, + ) + R = data_surf["R"] + R_t = data_surf["R_t"] + R_z = data_surf["R_z"] + R_tt = data_surf["R_tt"] + R_tz = data_surf["R_tz"] + R_zz = data_surf["R_zz"] + + Z_t = data_surf["Z_t"] + Z_z = data_surf["Z_z"] + Z_tt = data_surf["Z_tt"] + Z_tz = data_surf["Z_tz"] + Z_zz = data_surf["Z_zz"] + + t_s = data["theta_s"] + t_ss = data["theta_ss"] + z_s = data["zeta_s"] + z_ss = data["zeta_ss"] + + # Notation: xs_R is the R hat component of x'(s), + # and dz_xs_R its partial zeta derivative, etc. + xs_R = R_t * t_s + R_z * z_s + xs_p = R * z_s + dt_xs_R = R_tt * t_s + R_tz * z_s + dz_xs_R = R_tz * t_s + R_zz * z_s + ds_xs_R = R_t * t_ss + R_z * z_ss + dt_xs_p = R_t * z_s + dz_xs_p = R_z * z_s + ds_xs_p = R * z_ss + dt_xs_Z = Z_tt * t_s + Z_tz * z_s + dz_xs_Z = Z_tz * t_s + Z_zz * z_s + ds_xs_Z = Z_t * t_ss + Z_z * z_ss + + # As vectors, x_s = xs_R(theta(s),zeta(s)) R_hat(zeta(s)) + # + xs_p(theta(s),zeta(s)) p_hat(zeta(s)) + # + xs_Z(theta(s),zeta(s)) Z_hat + # We now apply the chain rule with the derivatives computed above + xss_R = dt_xs_R * t_s + dz_xs_R * z_s + ds_xs_R - xs_p * z_s + xss_p = dt_xs_p * t_s + dz_xs_p * z_s + ds_xs_p + xs_R * z_s + xss_Z = dt_xs_Z * t_s + dz_xs_Z * z_s + ds_xs_Z + + coords = jnp.stack([xss_R, xss_p, xss_Z], axis=1) + coords = rpz2xyz_vec(coords, phi=data["zeta"]) + coords = coords @ params["rotmat"].reshape((3, 3)).T + coords = xyz2rpz_vec(coords, phi=data["phi"]) + data["x_ss"] = coords + return data + + +@register_compute_fun( + name="x_sss", + label="\\partial_{sss} \\mathbf{x}", + units="m", + units_long="meters", + description="Position vector along curve, third derivative", + dim=3, + params=["R_lmn", "Z_lmn", "rotmat"], + transforms={"surface": []}, + profiles=[], + coordinates="s", + data=[ + "theta", + "zeta", + "theta_s", + "zeta_s", + "theta_ss", + "zeta_ss", + "theta_sss", + "zeta_sss", + "phi", + ], + parameterization="desc.geometry.core.SurfaceCurve", +) +def _x_sss_SurfaceCurve(params, transforms, profiles, data, **kwargs): + nodes = jnp.vstack([jnp.ones_like(data["theta"]), data["theta"], data["zeta"]]).T + grid = Grid(nodes, sort=False, jitable=True) + params_temp = transforms["surface"].params_dict.copy() + params_temp["R_lmn"] = params["R_lmn"] + params_temp["Z_lmn"] = params["Z_lmn"] + names = [ + "R", + "R_t", + "R_z", + "Z_t", + "Z_z", + "R_tt", + "R_tz", + "R_zz", + "Z_tt", + "Z_tz", + "Z_zz", + "R_ttt", + "R_ttz", + "R_tzz", + "R_zzz", + "Z_ttt", + "Z_ttz", + "Z_tzz", + "Z_zzz", + ] + data_surf = transforms["surface"].compute( + names, + grid=grid, + method="jitable", + params=params_temp, + ) + R = data_surf["R"] + R_t = data_surf["R_t"] + R_z = data_surf["R_z"] + R_tt = data_surf["R_tt"] + R_tz = data_surf["R_tz"] + R_zz = data_surf["R_zz"] + R_ttt = data_surf["R_ttt"] + R_ttz = data_surf["R_ttz"] + R_tzz = data_surf["R_tzz"] + R_zzz = data_surf["R_zzz"] + + Z_t = data_surf["Z_t"] + Z_z = data_surf["Z_z"] + Z_tt = data_surf["Z_tt"] + Z_tz = data_surf["Z_tz"] + Z_zz = data_surf["Z_zz"] + Z_ttt = data_surf["Z_ttt"] + Z_ttz = data_surf["Z_ttz"] + Z_tzz = data_surf["Z_tzz"] + Z_zzz = data_surf["Z_zzz"] + + t_s = data["theta_s"] + t_ss = data["theta_ss"] + t_sss = data["theta_sss"] + z_s = data["zeta_s"] + z_ss = data["zeta_ss"] + z_sss = data["zeta_sss"] + + # rebuilt here rather than read from data["x_ss"], which carries the rigid + # transform and is expressed in the basis at phi, not at zeta + xss_R = ( + R_tt * (t_s**2) + + 2 * R_tz * t_s * z_s + + R_zz * (z_s**2) + + R_t * t_ss + + R_z * z_ss + - R * (z_s**2) + ) + xss_p = 2 * R_t * t_s * z_s + 2 * R_z * (z_s**2) + R * z_ss + + dt_xss_R = ( + R_ttt * (t_s**2) + + 2 * R_ttz * t_s * z_s + + R_tzz * (z_s**2) + + R_tt * t_ss + + R_tz * z_ss + - R_t * (z_s**2) + ) + dz_xss_R = ( + R_ttz * (t_s**2) + + 2 * R_tzz * t_s * z_s + + R_zzz * (z_s**2) + + R_tz * t_ss + + R_zz * z_ss + - R_z * (z_s**2) + ) + ds_xss_R = ( + 2 * R_tt * t_s * t_ss + + 2 * R_tz * (t_ss * z_s + t_s * z_ss) + + 2 * R_zz * z_s * z_ss + + R_t * t_sss + + R_z * z_sss + - 2 * R * z_s * z_ss + ) + + dt_xss_p = 2 * R_tt * t_s * z_s + 2 * R_tz * (z_s**2) + R_t * z_ss + dz_xss_p = 2 * R_tz * t_s * z_s + 2 * R_zz * (z_s**2) + R_z * z_ss + ds_xss_p = 2 * R_t * (t_ss * z_s + t_s * z_ss) + 4 * R_z * z_s * z_ss + R * z_sss + + dt_xss_Z = ( + Z_ttt * (t_s**2) + + 2 * Z_ttz * t_s * z_s + + Z_tzz * (z_s**2) + + Z_tt * t_ss + + Z_tz * z_ss + ) + dz_xss_Z = ( + Z_ttz * (t_s**2) + + 2 * Z_tzz * t_s * z_s + + Z_zzz * (z_s**2) + + Z_tz * t_ss + + Z_zz * z_ss + ) + ds_xss_Z = ( + 2 * Z_tt * t_s * t_ss + + 2 * Z_tz * (t_ss * z_s + t_s * z_ss) + + 2 * Z_zz * z_s * z_ss + + Z_t * t_sss + + Z_z * z_sss + ) + + xsss_R = dt_xss_R * t_s + dz_xss_R * z_s + ds_xss_R - xss_p * z_s + xsss_p = dt_xss_p * t_s + dz_xss_p * z_s + ds_xss_p + xss_R * z_s + xsss_Z = dt_xss_Z * t_s + dz_xss_Z * z_s + ds_xss_Z + + coords = jnp.stack([xsss_R, xsss_p, xsss_Z], axis=1) + coords = rpz2xyz_vec(coords, phi=data["zeta"]) + coords = coords @ params["rotmat"].reshape((3, 3)).T + coords = xyz2rpz_vec(coords, phi=data["phi"]) + data["x_sss"] = coords + return data + + +@register_compute_fun( + name="theta", + label="\\mathbf{\\theta}", + units="rad", + units_long="radian", + description="Poloidal angle along curve", + dim=1, + params=["theta_n"], + transforms={ + "theta": [[0, 0, 0]], + }, + profiles=[], + coordinates="s", + data=["s"], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", + secular_theta="Net poloidal winding number of the curve.", +) +def _theta_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + secular_theta = kwargs["secular_theta"] + theta_0 = secular_theta * data["s"] + theta_1 = transforms["theta"].transform(params["theta_n"], dz=0) + data["theta"] = theta_0 + theta_1 + return data + + +@register_compute_fun( + name="zeta", + label="\\mathbf{\\zeta}", + units="rad", + units_long="radians", + description="Toroidal angle along curve", + dim=1, + params=["zeta_n"], + transforms={"zeta": [[0, 0, 0]]}, + profiles=[], + coordinates="s", + data=["s"], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", + secular_zeta="Net toroidal winding number of the curve.", +) +def _zeta_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + secular_zeta = kwargs["secular_zeta"] + zeta_0 = secular_zeta * data["s"] + zeta_1 = transforms["zeta"].transform(params["zeta_n"], dz=0) + data["zeta"] = zeta_0 + zeta_1 + return data + + +@register_compute_fun( + name="theta_s", + label="\\partial_{s} \\mathbf{\\theta}", + units="rad", + units_long="radians", + description="Poloidal angle along curve, first derivative", + dim=1, + params=["theta_n"], + transforms={ + "theta": [[0, 0, 1]], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", + secular_theta="Net poloidal winding number of the curve.", +) +def _theta_s_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + secular_theta = kwargs["secular_theta"] + theta_0 = secular_theta + theta_1 = transforms["theta"].transform(params["theta_n"], dz=1) + data["theta_s"] = theta_0 + theta_1 + return data + + +@register_compute_fun( + name="theta_ss", + label="\\partial_{ss} \\mathbf{\\theta}", + units="rad", + units_long="radians", + description="Poloidal angle along curve, second derivative", + dim=1, + params=["theta_n"], + transforms={"theta": [[0, 0, 2]]}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", +) +def _theta_ss_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + theta = transforms["theta"].transform(params["theta_n"], dz=2) + data["theta_ss"] = theta + return data + + +@register_compute_fun( + name="theta_sss", + label="\\partial_{sss} \\mathbf{\\theta}", + units="rad", + units_long="radians", + description="Poloidal angle along curve, third derivative", + dim=1, + params=["theta_n"], + transforms={"theta": [[0, 0, 3]]}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", +) +def _theta_sss_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + theta = transforms["theta"].transform(params["theta_n"], dz=3) + data["theta_sss"] = theta + return data + + +@register_compute_fun( + name="zeta_s", + label="\\partial_{s} \\mathbf{\\zeta}", + units="rad", + units_long="radians", + description="Toroidal angle along curve, first derivative", + dim=1, + params=["zeta_n"], + transforms={"zeta": [[0, 0, 1]]}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", + secular_zeta="Net toroidal winding number of the curve.", +) +def _zeta_s_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + secular_zeta = kwargs["secular_zeta"] + zeta_0 = secular_zeta + zeta_1 = transforms["zeta"].transform(params["zeta_n"], dz=1) + data["zeta_s"] = zeta_0 + zeta_1 + return data + + +@register_compute_fun( + name="zeta_ss", + label="\\partial_{ss} \\mathbf{\\zeta}", + units="rad", + units_long="radians", + description="Toroidal angle along curve, second derivative", + dim=1, + params=["zeta_n"], + transforms={"zeta": [[0, 0, 2]]}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", +) +def _zeta_ss_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + zeta = transforms["zeta"].transform(params["zeta_n"], dz=2) + data["zeta_ss"] = zeta + return data + + +@register_compute_fun( + name="zeta_sss", + label="\\partial_{sss} \\mathbf{\\zeta}", + units="rad", + units_long="radians", + description="Toroidal angle along curve, third derivative", + dim=1, + params=["zeta_n"], + transforms={"zeta": [[0, 0, 3]]}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZSurfaceCurve", +) +def _zeta_sss_FourierRZSurfaceCurve(params, transforms, profiles, data, **kwargs): + zeta = transforms["zeta"].transform(params["zeta_n"], dz=3) + data["zeta_sss"] = zeta + return data diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index d52386b0cf..7ca69ff3f0 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -249,6 +249,10 @@ def _decorator(func): "desc.geometry.curve.SplineXYZCurve": [ "desc.geometry.core.Curve", ], + "desc.geometry.curve.FourierRZSurfaceCurve": [ + "desc.geometry.core.SurfaceCurve", + "desc.geometry.core.Curve", + ], "desc.geometry.surface.FourierRZToroidalSurface": [ "desc.geometry.core.Surface", ], @@ -271,6 +275,11 @@ def _decorator(func): "desc.geometry.curve.FourierXYCurve", "desc.geometry.core.Curve", ], + "desc.coils.FourierRZSurfaceCoil": [ + "desc.geometry.curve.FourierRZSurfaceCurve", + "desc.geometry.core.SurfaceCurve", + "desc.geometry.core.Curve", + ], "desc.magnetic_fields._current_potential.CurrentPotentialField": [ "desc.geometry.surface.FourierRZToroidalSurface", "desc.geometry.core.Surface", diff --git a/desc/geometry/__init__.py b/desc/geometry/__init__.py index 8404ca7fb9..16fd0d4db1 100644 --- a/desc/geometry/__init__.py +++ b/desc/geometry/__init__.py @@ -1,9 +1,10 @@ """Classes for representing geometric objects like curves and surfaces.""" -from .core import Curve, Surface +from .core import Curve, Surface, SurfaceCurve from .curve import ( FourierPlanarCurve, FourierRZCurve, + FourierRZSurfaceCurve, FourierXYCurve, FourierXYZCurve, SplineXYZCurve, diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 38c70f764b..709a31a334 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -148,12 +148,15 @@ def compute( if (data_index[p][dep]["coordinates"] == "") and (dep not in data) ] calc0d = bool(len(dep0d)) + # some curves (e.g. SurfaceCurve) have a resolution + # requirement which can exceed self.N + N = getattr(self, "N_effective", self.N) # see if the grid we're already using will work for desired qtys - if calc0d and (grid.N >= 2 * self.N + 5) and isinstance(grid, LinearGrid): + if calc0d and (grid.N >= 2 * N + 5) and isinstance(grid, LinearGrid): calc0d = False if calc0d and override_grid: - grid0d = LinearGrid(N=2 * self.N * getattr(self, "NFP", 1) + 5) + grid0d = LinearGrid(N=2 * N * getattr(self, "NFP", 1) + 5) data0d = compute_fun( self, dep0d, @@ -621,3 +624,85 @@ def __repr__(self): + str(hex(id(self))) + " (name={})".format(self.name) ) + + +class SurfaceCurve(Curve): + r"""Curve which lies in a toroidal surface. + + Parameterized in terms of poloidal/toroidal + angles: + (theta,zeta)= (theta(s), zeta(s)), s in [0,2pi). + Translated to lab coordinates via the surface: + (R,Z)=(R(theta(s),zeta(s)), Z(theta(s),zeta(s))). + + Note: objects carry a read-only copy of + the surface and its underlying parameters. + The surface's params are optimizable, so + optimization with fixed surface should include + a FixParameters call. If this surface + appears across multiple objectives, should use + the objective SurfaceCurveConsistency. + + Note: Optimizing a + SurfaceCurve should generally include a + FixParameters(curve, {"rotmat": True, "shift": True}). + + Parameters + ---------- + surface: FourierRZToroidalSurface + Underlying surface which the curve lies on. + name: str, optional + Name for this curve. + """ + + _io_attrs_ = Curve._io_attrs_ + ["_surface"] + _static_attrs = Curve._static_attrs + + def __init__( + self, + surface, + name="", + ): + assert surface is not None, "Surface cannot be None" + super().__init__(name=name) + self._surface = surface.copy() + + @property + def surface(self): + """The curve's own copy of the surface it lies on.""" + return self._surface + + @optimizable_parameter + @property + def R_lmn(self): + """Spectral coefficients for R of the underlying surface.""" + return self._surface.R_lmn + + @R_lmn.setter + def R_lmn(self, new): + self._surface.R_lmn = new + + @optimizable_parameter + @property + def Z_lmn(self): + """Spectral coefficients for Z of the underlying surface.""" + return self._surface.Z_lmn + + @Z_lmn.setter + def Z_lmn(self, new): + self._surface.Z_lmn = new + + @property + def R_basis(self): + """Spectral basis for R of the underlying surface.""" + return self._surface.R_basis + + @property + def Z_basis(self): + """Spectral basis for Z of the underlying surface.""" + return self._surface.Z_basis + + @property + def NFP_surface(self): + """Number of field periods of the underlying surface.""" + return self._surface.NFP diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index 54532ff336..7fe64b8796 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -20,12 +20,13 @@ rotation_matrix, rpz2xyz, rpz2xyz_vec, + setdefault, warnif, xyz2rpz, xyz2rpz_vec, ) -from .core import Curve +from .core import Curve, SurfaceCurve __all__ = [ "FourierPlanarCurve", @@ -1692,3 +1693,501 @@ def from_values(cls, coords, knots=None, method="cubic", basis="xyz", name=""): method=method, name=name, ) + + +class FourierRZSurfaceCurve(SurfaceCurve): + r"""Fourier parameterized SurfaceCurve. + + Poloidal and toroidal angles parameterized as follows:: + + theta(s) = theta_secular*s + sum_{0}^{N_theta} theta_n cos(NFP*n*s) + + sum_{-N_theta}^{-1} theta_n sin(NFP*|n|*s) + zeta(s) = zeta_secular*s + sum_{0}^{N_zeta} theta_n cos(NFP*n*s) + + sum_{-N_zeta}^{-1} theta_n sin(NFP*|n|*s) + + Parameters + ---------- + surface: FourierRZToroidalSurface, optional + Underlying surface which the curve lies on. + If None, must pass a value for "equilibrium." + equilibrium: Equilibrium, optional + Used in place of "surface" when the curve is intended + to live on the rho=1 surface of an equilibrium which + may be optimized alongside the curve. + Must pass exactly one of "surface" or "equilibrium." + secular_theta: non-negative int, optional + Net winding of theta(s), equivalent to + 1/2pi \int_{0}^{2pi} theta'(s) ds. Default + value is 0. Nonzero values result in curves + linking the torus poloidally. + Note: gcd(secular_theta, secular_zeta) should equal 1. + secular_zeta: non-negative int, optional + Net winding of zeta(s), equivalent to + 1/2pi \int_{0}^{2pi} zeta'(s) ds. Default + value is 1. Nonzero values result in curves + linking the torus toroidally. + Note: gcd(secular_theta, secular_zeta) should equal 1. + theta_n: array-like, optional + Coefficients of Fourier modes defining + theta(s). + zeta_n: array-like, optional + Coefficients of Fourier modes defining + zeta(s). + modes_theta: array-like, optional + Mode numbers associated with theta. If + not given, defaults to [-N_theta:N_theta], + where N_theta = len(theta_n)//2. If given + alongside theta_n, must have the same length. + modes_zeta: array-like, optional + Mode numbers associated with zeta. If + not given, defaults to [-N_theta:N_theta], + where N_theta = len(theta_n)//2. If given + alongside zeta_n, must have the same length. + sym_theta: str, optional + If "sin"/"cos", retains only the corresponding + modes in the series for theta(s). In this case, + modes_theta should also be given. Defaults + to None, i.e. all modes are kept. + sym_zeta: bool, optional + If "sin"/"cos", retains only the corresponding + modes in the series for zeta(s). In this case, + modes_zeta should also be given. Defaults + to None, i.e. all modes are kept. + NFP: positive int, optional + Field period symmetry of the curve. Defaults to 1. + The curve's image is invariant + under zeta -> zeta + 2pi/NFP, and Fourier modes + are multiples of NFP in s. + Note: If changed, must satisfy (a) NFP | surface.NFP, + (b) NFP | secular_theta, (c) secular_zeta \neq 0. + name: str, optional + Name for this curve. + """ + + _io_attrs_ = SurfaceCurve._io_attrs_ + [ + "_theta_basis", + "_zeta_basis", + "_secular_theta", + "_secular_zeta", + "_theta_n", + "_zeta_n", + "_modes_theta", + "_modes_zeta", + "_sym_theta", + "_sym_zeta", + "_NFP", + ] + _static_attrs = SurfaceCurve._static_attrs + [ + "_theta_basis", + "_zeta_basis", + "_theta_n_fixed", + "_zeta_n_fixed", + "_secular_theta", + "_secular_zeta", + "_modes_theta", + "_modes_zeta", + "_sym_theta", + "_sym_zeta", + "_NFP", + ] + + def __init__( + self, + surface=None, + equilibrium=None, + secular_theta=0, + secular_zeta=1, + theta_n=[0.0], + zeta_n=[0.0], + modes_theta=None, + modes_zeta=None, + sym_theta=False, + sym_zeta=False, + NFP=1, + name="", + ): + theta_n = theta_n if theta_n is None else np.atleast_1d(theta_n) + zeta_n = zeta_n if zeta_n is None else np.atleast_1d(zeta_n) + self._sym_theta = sym_theta + self._sym_zeta = sym_zeta + self._secular_theta = secular_theta + self._secular_zeta = secular_zeta + self._NFP = check_posint(NFP, "NFP", False) + assert int(secular_theta) == secular_theta, "secular_theta must be an integer" + assert int(secular_zeta) == secular_zeta, "secular_zeta must be an integer" + + errorif( + all([surface, equilibrium]) or not any([surface, equilibrium]), + ValueError, + "Exactly one of surface or equilibrium is required", + ) + if equilibrium is not None: + surface = equilibrium.surface + + errorif( + np.gcd(abs(secular_theta), abs(secular_zeta)) != 1, + ValueError, + "secular_theta and secular_zeta should have a gcd of 1", + ) + errorif(surface.NFP % NFP != 0, ValueError, "NFP must divide surface.NFP") + errorif( + NFP > 1 and secular_theta % NFP != 0, + ValueError, + "NFP must divide secular_theta", + ) + errorif( + NFP > 1 and np.gcd(abs(secular_zeta), NFP) != 1, + ValueError, + "secular_zeta and NFP should have a gcd of 1", + ) + errorif( + NFP > 1 and secular_zeta == 0, + ValueError, + "a modular curve (secular_zeta=0) requires NFP=1", + ) + errorif( + (sym_theta and modes_theta is None) or (sym_zeta and modes_zeta is None), + ValueError, + "sym option given without corresponding mode numbers", + ) + + super().__init__(surface=surface, name=name) + + self._theta_n_fixed = (theta_n is None) and (modes_theta is None) + self._zeta_n_fixed = (zeta_n is None) and (modes_zeta is None) + + # we're doing funny things with symmetry here, potentially ignoring user inputs + if self._theta_n_fixed: + # sym="sin" with N=0 gives a basis with zero modes. + self._theta_n = jnp.array([]) + self._modes_theta = np.array([], dtype=int) + self._theta_basis = FourierSeries(0, sym="sin", NFP=NFP) + else: + if modes_theta is None: + self._modes_theta = np.arange( + -(len(theta_n) // 2), len(theta_n) // 2 + 1 + ) + else: + self._modes_theta = np.array(modes_theta) + if theta_n is None: + self._theta_n = jnp.zeros_like(self._modes_theta, dtype=float) + else: + self._theta_n = jnp.array(theta_n) + errorif( + len(self._theta_n) != len(self._modes_theta), + ValueError, + "theta_n and modes_theta must have same length", + ) + + N = np.max(np.abs(self._modes_theta)) + self._theta_basis = FourierSeries(N, sym=sym_theta, NFP=NFP) + self._theta_n = copy_coeffs( + self._theta_n, self._modes_theta, self._theta_basis.modes[:, 2] + ) + + if self._zeta_n_fixed: + self._zeta_n = jnp.array([]) + self._modes_zeta = np.array([], dtype=int) + self._zeta_basis = FourierSeries(0, sym="sin", NFP=NFP) + else: + if modes_zeta is None: + self._modes_zeta = np.arange(-(len(zeta_n) // 2), len(zeta_n) // 2 + 1) + else: + self._modes_zeta = np.array(modes_zeta) + if zeta_n is None: + self._zeta_n = jnp.zeros_like(self._modes_zeta, dtype=float) + else: + self._zeta_n = jnp.array(zeta_n) + errorif( + len(self._zeta_n) != len(self._modes_zeta), + ValueError, + "zeta_n and modes_zeta must have same length", + ) + N = np.max(np.abs(self._modes_zeta)) + self._zeta_basis = FourierSeries(N, NFP=NFP, sym=sym_zeta) + self._zeta_n = copy_coeffs( + self._zeta_n, self._modes_zeta, self._zeta_basis.modes[:, 2] + ) + + @property + def secular_theta(self): + return self._secular_theta + + @secular_theta.setter + def secular_theta(self, new): + errorif(int(new) != new, ValueError, "secular_theta should be an integer") + self._secular_theta = new + + @property + def secular_zeta(self): + return self._secular_zeta + + @secular_zeta.setter + def secular_zeta(self, new): + errorif(int(new) != new, ValueError, "secular_zeta should be an integer") + self._secular_zeta = new + + @property + def theta_basis(self): + return self._theta_basis + + @property + def zeta_basis(self): + return self._zeta_basis + + @optimizable_parameter + @property + def theta_n(self): + return self._theta_n + + @theta_n.setter + def theta_n(self, new): + num_modes = self.theta_basis.num_modes + if len(new) == num_modes: + self._theta_n = jnp.asarray(new) + else: + raise ValueError( + f"theta_n should have the same size as the basis, got {len(new)} for " + + f"basis with {num_modes} modes." + ) + + @optimizable_parameter + @property + def zeta_n(self): + return self._zeta_n + + @zeta_n.setter + def zeta_n(self, new): + num_modes = self.zeta_basis.num_modes + if len(new) == num_modes: + self._zeta_n = jnp.asarray(new) + else: + raise ValueError( + f"zeta_n should have the same size as the basis, got {len(new)} for " + + f"basis with {num_modes} modes." + ) + + @property + def sym_theta(self): + return self._sym_theta + + @property + def sym_zeta(self): + return self._sym_zeta + + @property + def NFP(self): + return self._NFP + + @property + def N_theta(self): + return self._theta_basis.N + + @property + def N_zeta(self): + return self._zeta_basis.N + + @property + def N(self): + return max(self.N_theta, self.N_zeta) + + @property + def N_effective(self): + r"""Frequency necessary to resolve surface quantities. + + Differs from self.N, as the latter only captures the + resolution of the curve. Note this is an approximation, + most accurate when ``|theta'(s)|*M_surf``, ``|zeta'(s)|*N_surf`` + are not too large. + """ + surface = self.surface + M_surf = max(surface.R_basis.M, surface.Z_basis.M) + N_surf = max(surface.R_basis.N, surface.Z_basis.N) + return ( + self.N * self.NFP + + M_surf * abs(self.secular_theta) + + N_surf * surface.NFP * abs(self.secular_zeta) + ) + + def change_resolution( + self, + N_theta=None, + N_zeta=None, + sym_theta=None, + sym_zeta=None, + ): + # Limitation: this cannot eliminate a basis, since N_theta=None just gets + # defaulted to self.N_theta. Build the curve with theta_n=None instead. + # NFP is also not adjustable here. + + N_theta = setdefault(N_theta, self.N_theta) + N_zeta = setdefault(N_zeta, self.N_zeta) + sym_theta = setdefault(sym_theta, self.sym_theta) + sym_zeta = setdefault(sym_zeta, self.sym_zeta) + NFP = self._NFP + + if (N_theta != self.N_theta) or (sym_theta != self.sym_theta): + modes_theta_old = self.theta_basis.modes[:, 2] + self.theta_basis.change_resolution(N_theta, NFP, sym_theta) + self.theta_n = copy_coeffs( + self.theta_n, modes_theta_old, self.theta_basis.modes[:, 2] + ) + + if (N_zeta != self.N_zeta) or (sym_zeta != self.sym_zeta): + modes_zeta_old = self.zeta_basis.modes[:, 2] + self.zeta_basis.change_resolution(N_zeta, NFP, sym_zeta) + self.zeta_n = copy_coeffs( + self.zeta_n, modes_zeta_old, self.zeta_basis.modes[:, 2] + ) + + self._sym_theta = sym_theta + self._sym_zeta = sym_zeta + + def get_coeffs(self, n): + """Get Fourier coefficients of theta(s), zeta(s) for given mode numbers.""" + n = np.atleast_1d(n).astype(int) + theta_n = np.zeros_like(n).astype(float) + zeta_n = np.zeros_like(n).astype(float) + + idx_t = np.where(n[:, np.newaxis] == self.theta_basis.modes[:, 2]) + idx_z = np.where(n[:, np.newaxis] == self.zeta_basis.modes[:, 2]) + + theta_n[idx_t[0]] = self.theta_n[idx_t[1]] + zeta_n[idx_z[0]] = self.zeta_n[idx_z[1]] + return theta_n, zeta_n + + def set_coeffs(self, n, theta_n=None, zeta_n=None): + """Set specific Fourier coefficients of theta(s), zeta(s).""" + n, theta_n, zeta_n = ( + np.atleast_1d(n), + np.atleast_1d(theta_n), + np.atleast_1d(zeta_n), + ) + theta_n = np.broadcast_to(theta_n, n.shape) + zeta_n = np.broadcast_to(zeta_n, n.shape) + for nn, tt, zz in zip(n, theta_n, zeta_n): + if tt is not None: + idx_t = self.theta_basis.get_idx(0, 0, nn) + self.theta_n = put(self.theta_n, idx_t, tt) + if zz is not None: + idx_z = self.zeta_basis.get_idx(0, 0, nn) + self.zeta_n = put(self.zeta_n, idx_z, zz) + + def compute( + self, names, grid=None, params=None, transforms=None, data=None, **kwargs + ): + kwargs.setdefault("secular_theta", self.secular_theta) + kwargs.setdefault("secular_zeta", self.secular_zeta) + if grid is None: + grid = LinearGrid(N=2 * self.N_effective + 5) + return super().compute(names, grid, params, transforms, data, **kwargs) + + @classmethod + def from_values( + cls, + coords, + surface=None, + equilibrium=None, + NFP=1, + N_theta=10, + N_zeta=10, + secular_theta=None, + secular_zeta=None, + sym_theta=None, + sym_zeta=None, + name="", + ): + """Fit a FourierRZSurfaceCurve to a set of (theta(s), zeta(s)) coordinates. + + Parameters + ---------- + coords: array-like, shape (N,3) + Default shape (N,3) corresponds to coordinates (s,theta(s),zeta(s)) at N + points along curve. Assumed to be ordered by increasing curve parameter. + surface: FourierRZToroidalSurface + Surface on which curve lies. Only one of "surface" and "equilibrium" + can be provided. + equilibrium: Equilibrium + Provided in place of "surface" when the desired surface is the rho=1 + flux surface of an equilibrium. Only one of "surface" and "equilibrium" + can be provided. + NFP: int, optional + Field period symmetry. Defaults to 1. + N_theta: int or None, optional + Max Fourier mode number for theta(s). Set to None to indicate no basis + for theta should be included. Default is 10. + N_zeta: int or None, optional + Max Fourier mode number for zeta(s). Set to None to indicate no basis + for zeta should be included. Default is 10. + secular_theta: int or None, optional + Coefficient of secular term in theta. If None, estimated automatically. + secular_zeta: int or None, optional + Coefficient of secular term in zeta. If None, estimated automatically. + sym_theta: str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for theta(s). Default is None. + sym_zeta: str or None, optional + Whether to use "cos", "sin", or no (None) symmetry in the series + for zeta(s). Default is None. + name: str, optional + Name for this curve. + + Returns + ------- + FourierRZSurfaceCurve + FourierRZSurfaceCurve object fit to the given coords. + + """ + s = coords[:, 0] + theta = coords[:, 1] + zeta = coords[:, 2] + + if secular_theta is None: + secular_theta = np.asarray((theta[-1] - theta[0]) / (s[-1] - s[0])) + secular_theta = np.round(secular_theta).astype(int) + if secular_zeta is None: + secular_zeta = np.asarray((zeta[-1] - zeta[0]) / (s[-1] - s[0])) + secular_zeta = np.round(secular_zeta).astype(int) + + theta_single_val = theta - secular_theta * s + zeta_single_val = zeta - secular_zeta * s + + # Sort and remove duplicate values of s. Single valued parts are not modded. + s, idx = np.unique(s, return_index=True) + theta_single_val = theta_single_val[idx] + zeta_single_val = zeta_single_val[idx] + + if N_theta is None: + theta_n = None + modes_theta = None + else: + grid_theta = LinearGrid(zeta=s) + theta_basis = FourierSeries(N=N_theta, NFP=NFP, sym=sym_theta) + transform_theta = Transform(grid_theta, theta_basis, build_pinv=True) + theta_n = transform_theta.fit(theta_single_val) + modes_theta = theta_basis.modes[:, 2] + + if N_zeta is None: + zeta_n = None + modes_zeta = None + else: + grid_zeta = LinearGrid(zeta=s) + zeta_basis = FourierSeries(N=N_zeta, NFP=NFP, sym=sym_zeta) + transform_zeta = Transform(grid_zeta, zeta_basis, build_pinv=True) + zeta_n = transform_zeta.fit(zeta_single_val) + modes_zeta = zeta_basis.modes[:, 2] + + return cls( + surface=surface, + equilibrium=equilibrium, + secular_theta=secular_theta, + secular_zeta=secular_zeta, + theta_n=theta_n, + zeta_n=zeta_n, + modes_theta=modes_theta, + modes_zeta=modes_zeta, + sym_theta=sym_theta, + sym_zeta=sym_zeta, + NFP=NFP, + name=name, + ) diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index 52ca059e98..e84d410653 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -104,5 +104,6 @@ FixSumModesZ, FixThetaSFL, ShareParameters, + SurfaceCurveConsistency, ) from .objective_funs import ObjectiveFunction diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index eaafe93a65..112ebbc010 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -200,7 +200,10 @@ def expand(t, idx=0): if grid is None: grid = [] for c in coils: - grid.append(LinearGrid(N=2 * c.N * getattr(c, "NFP", 1) + 5)) + # some coils (e.g. FourierRZSurfaceCoil) have a different + # resolution requirement than c.N + N = getattr(c, "N_effective", c.N * getattr(c, "NFP", 1)) + grid.append(LinearGrid(N=2 * N + 5)) if isinstance(grid, numbers.Integral): grid = LinearGrid(N=self._grid) if isinstance(grid, _Grid): diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index be657f4df4..ca5e5f27ec 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -20,7 +20,7 @@ tree_structure, ) from desc.basis import zernike_radial -from desc.geometry import FourierRZCurve +from desc.geometry import FourierRZCurve, Surface from desc.utils import broadcast_tree, errorif, setdefault from .normalization import compute_scaling_factors @@ -855,6 +855,157 @@ def compute(self, params, constants=None): return f +class SurfaceCurveConsistency(_Objective): + """Ensures a SurfaceCurve shares params with underlying Surface. + + Objective is only needed when both the curve and surface objects + are both being optimized. + + Parameters + ---------- + source: FourierRZToroidalSurface or Equilibrium + Surface on which the curve input should lie. If Equilibrium is passed, + surface is set to the equilibrium's rho=1 surface. + curve: SurfaceCurve, FourierRZWindingCoil CoilSet + Curve or collection of curves carrying a copy of surface params. + name: str, optional + Name of the objective function. + """ + + __doc__ = __doc__.rstrip() + collect_docs( + overwrite={ + "target": "", + "bounds": "", + "normalize": "", + "normalize_target": "", + "weight": "", + } + ) + _scalar = False + _linear = True + _fixed = False + _units = "(m)" + _print_value_fmt = "SurfaceCurve consistency error: " + + _static_attrs = _Objective._static_attrs + ["_src_params"] + + def __init__( + self, + source, + curve, + name="SurfaceCurve consistency", + ): + # local import to avoid circular dependency + from desc.equilibrium import Equilibrium + + errorif( + not isinstance(source, (Surface, Equilibrium)), + ValueError, + "Source must be a surface or equilibrium.", + ) + + things = [source, curve] + + super().__init__( + things=things, + target=0, + weight=1, + name=name, + normalize=False, + normalize_target=False, + ) + + @execute_on_cpu + def build(self, use_jit=False, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + """ + from desc.coils import CoilSet + from desc.equilibrium import Equilibrium + + def _expand(curve): + if isinstance(curve, CoilSet): + return [c for coil in curve.coils for c in _expand(coil)] + return [curve] + + # in case the curve is a coilset, expand to a list of all individual curves + curve_expanded = _expand(self.things[1]) + source = self.things[0] + + if isinstance(source, Equilibrium): + source_R_basis = source.surface.R_basis + source_Z_basis = source.surface.Z_basis + self._src_params = {"R": "Rb_lmn", "Z": "Zb_lmn"} + else: + source_R_basis = source.R_basis + source_Z_basis = source.Z_basis + self._src_params = {"R": "R_lmn", "Z": "Z_lmn"} + + AR = [] + AZ = [] + for curve in curve_expanded: + nR = len(curve.surface.R_basis.modes) + nZ = len(curve.surface.Z_basis.modes) + + # The curve carries a copy of the source surface, so the bases must agree. + errorif( + not np.array_equal(curve.R_basis.modes, source_R_basis.modes) + or not np.array_equal(curve.Z_basis.modes, source_Z_basis.modes), + ValueError, + "The curve's surface and the source must share the same R and Z bases.", + ) + # The weights are just the identity matrices. If rho<1 surfaces + # are supported at any point, this will change. + AR.append(np.eye(nR)) + AZ.append(np.eye(nZ)) + + self._dim_f = sum(a.shape[0] for a in AR) + sum(a.shape[0] for a in AZ) + self._A = {"R": AR, "Z": AZ} + super().build(use_jit=use_jit, verbose=verbose) + + def compute(self, params_source, params_curve, constants=None): + """Compute SurfaceCurveConsistency error. + + Measures the mismatch between the surface params attached to the curve, + and the params attached to the underlying surface (or equilibrium.surface). + + Parameters + ---------- + params_source: dict + Dictionary of surface degrees of freedom, or if surface=equilibrium.surface, + then the equilibrium degrees of freedom + params_curve : dict + Dictionary of curve (or coilset) degrees of freedom + constants : dict + Dictionary of constant data, eg transforms, profiles etc. Defaults to + self.constants. (Deprecated) + + Returns + ------- + f : ndarray + SurfaceCurveConsistency errors. + """ + params_curve = tree_leaves(params_curve, is_leaf=lambda x: isinstance(x, dict)) + return jnp.concatenate( + [ + block + for i in range(len(params_curve)) + for block in ( + jnp.dot(self._A["R"][i], params_source[self._src_params["R"]]) + - params_curve[i]["R_lmn"], + jnp.dot(self._A["Z"][i], params_source[self._src_params["Z"]]) + - params_curve[i]["Z_lmn"], + ) + ] + ) + + class FixBoundaryR(FixParameters): """Boundary condition on the R boundary parameters. diff --git a/docs/api.rst b/docs/api.rst index c6ca14b9a3..2b6869e04b 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -31,6 +31,7 @@ Coils desc.coils.CoilSet desc.coils.FourierPlanarCoil desc.coils.FourierRZCoil + desc.coils.FourierRZSurfaceCoil desc.coils.FourierXYCoil desc.coils.FourierXYZCoil desc.coils.MixedCoilSet @@ -120,10 +121,12 @@ Geometry desc.geometry.FourierPlanarCurve desc.geometry.FourierRZCurve + desc.geometry.FourierRZSurfaceCurve desc.geometry.FourierRZToroidalSurface desc.geometry.FourierXYCurve desc.geometry.FourierXYZCurve desc.geometry.SplineXYZCurve + desc.geometry.SurfaceCurve desc.geometry.ZernikeRZToroidalSection Grid diff --git a/docs/api_fields.rst b/docs/api_fields.rst index 648499e1b7..465132e2ce 100644 --- a/docs/api_fields.rst +++ b/docs/api_fields.rst @@ -96,6 +96,7 @@ classes, can also use the same ``Curve`` conversion methods to convert between c :template: class.rst desc.coils.FourierRZCoil + desc.coils.FourierRZSurfaceCoil desc.coils.FourierXYZCoil desc.coils.FourierPlanarCoil desc.coils.FourierXYCoil diff --git a/docs/api_objectives.rst b/docs/api_objectives.rst index 72b6468e78..3b99837474 100644 --- a/docs/api_objectives.rst +++ b/docs/api_objectives.rst @@ -189,6 +189,7 @@ Fixing degrees of freedom desc.objectives.FixSumCoilCurrent desc.objectives.FixParameters desc.objectives.ShareParameters + desc.objectives.SurfaceCurveConsistency User defined objectives diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index a4ac2ba3aa..3720bde38c 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/test_coils.py b/tests/test_coils.py index 0a89921a2b..2cde972260 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -12,6 +12,7 @@ CoilSet, FourierPlanarCoil, FourierRZCoil, + FourierRZSurfaceCoil, FourierXYCoil, FourierXYZCoil, MixedCoilSet, @@ -189,6 +190,41 @@ def test_biot_savart_all_coils(self): B_true_rpz_phi, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" ) + # FourierRZSurfaceCoil + surf = FourierRZToroidalSurface( + R_lmn=[R + 1, 1], modes_R=[[0, 0], [1, 0]], Z_lmn=[-1], modes_Z=[[-1, 0]] + ) + coil = FourierRZSurfaceCoil( + I, surface=surf, secular_theta=0, secular_zeta=1, theta_n=[np.pi] + ) + B_xyz = coil.compute_magnetic_field( + grid_xyz, basis="xyz", source_grid=coil_grid + ) + B_rpz = coil.compute_magnetic_field( + grid_rpz, basis="rpz", source_grid=coil_grid + ) + np.testing.assert_allclose( + B_true_xyz, + B_xyz, + rtol=1e-3, + atol=1e-10, + err_msg="Using FourierRZSurfaceCoil", + ) + np.testing.assert_allclose( + B_true_rpz_xy, + B_rpz, + rtol=1e-3, + atol=1e-10, + err_msg="Using FourierRZSurfaceCoil", + ) + np.testing.assert_allclose( + B_true_rpz_phi, + B_rpz, + rtol=1e-3, + atol=1e-10, + err_msg="Using FourierRZSurfaceCoil", + ) + @pytest.mark.unit def test_biot_savart_vector_potential_all_coils(self): """Test biot-savart vec potential implementation against analytic formula.""" @@ -250,6 +286,15 @@ def test(coil, grid_xyz, grid_rpz): coil3[1].current = 0 test(coil3, grid_xyz, grid_rpz) + # FourierRZSurfaceCoil + surf = FourierRZToroidalSurface( + R_lmn=[R + 1, 1], modes_R=[[0, 0], [1, 0]], Z_lmn=[-1], modes_Z=[[-1, 0]] + ) + coil = FourierRZSurfaceCoil( + I, surface=surf, secular_theta=0, secular_zeta=1, theta_n=[np.pi] + ) + test(coil, grid_xyz, grid_rpz) + @pytest.mark.unit def test_biot_savart_vector_potential_integral_all_coils(self): """Test analytic expression of flux integral for all coils.""" @@ -399,6 +444,28 @@ def test( atol=1e-12, ) + # FourierRZSurfaceCoil + # surface has major radius R+1, minor radius 1, so a coil on inner + # midboard has the correct dimensions + surf = FourierRZToroidalSurface( + R_lmn=[R + 1, 1], + modes_R=[[0, 0], [1, 0]], + Z_lmn=[-1], + modes_Z=[[-1, 0]], + ) + coil = FourierRZSurfaceCoil( + I, surface=surf, secular_theta=0, secular_zeta=1, theta_n=[np.pi] + ) + test( + coil, + grid_xyz, + grid_rpz, + A_true_rpz, + correct_flux, + rtol=1e-8, + atol=1e-12, + ) + @pytest.mark.unit def test_properties(self): """Test getting/setting attributes for Coil class.""" @@ -465,6 +532,15 @@ def test_convert_type(self): coil4 = coil1.to_FourierRZ(N=coil1.N) coil5 = coil1.to_FourierXY(N=10, basis="rpz") coil6 = coil1.to_FourierPlanar(N=10, basis="rpz") + # coil1 traces R = 10 + cos(zeta), Z = 0, which is the theta=0 curve of this + # surface. + surf = FourierRZToroidalSurface( + R_lmn=[9, 1, 1], + modes_R=[[0, 0], [0, 1], [1, 0]], + Z_lmn=[-1], + modes_Z=[[-1, 0]], + ) + coil7 = coil1.to_FourierRZSurface(surface=surf, secular_theta=0, secular_zeta=1) grid = LinearGrid(zeta=s) x1 = coil1.compute("x", grid=grid, basis="xyz")["x"] @@ -478,6 +554,7 @@ def test_convert_type(self): ) # use Grid instead of LinearGrid to prevent node sorting grid_planar = Grid(np.array([np.zeros_like(zeta), np.zeros_like(zeta), zeta]).T) x6 = coil6.compute("x", grid=grid_planar, basis="xyz")["x"] + x7 = coil7.compute("x", grid=grid, basis="xyz")["x"] B1 = coil1.compute_magnetic_field( np.zeros((1, 3)), source_grid=grid, basis="xyz" @@ -497,17 +574,22 @@ def test_convert_type(self): B6 = coil6.compute_magnetic_field( np.zeros((1, 3)), source_grid=grid, basis="xyz" ) + B7 = coil7.compute_magnetic_field( + np.zeros((1, 3)), source_grid=grid, basis="xyz" + ) np.testing.assert_allclose(x1, x2, atol=1e-12) np.testing.assert_allclose(x1, x3, atol=1e-12) np.testing.assert_allclose(x1, x4, atol=1e-12) np.testing.assert_allclose(x1, x5, atol=1e-12) np.testing.assert_allclose(x1, x6, atol=1e-10) + np.testing.assert_allclose(x1, x7, atol=1e-6) # looser tolerance np.testing.assert_allclose(B1, B2, rtol=1e-8, atol=1e-8) np.testing.assert_allclose(B1, B3, rtol=1e-3, atol=1e-8) np.testing.assert_allclose(B1, B4, rtol=1e-8, atol=1e-8) np.testing.assert_allclose(B1, B5, rtol=1e-6, atol=1e-7) np.testing.assert_allclose(B1, B6, rtol=1e-6, atol=1e-7) + np.testing.assert_allclose(B1, B7, rtol=1e-6, atol=1e-7) class TestCoilSet: @@ -833,14 +915,38 @@ def test_convert_type(self): coils2 = coils0.to_FourierXYZ(grid=grid, check_intersection=False) coils3 = coils0.to_FourierXY(grid=grid, check_intersection=False) coils4 = coils0.to_FourierPlanar(grid=grid, check_intersection=False) + # to test conversion to surface coils, need a common surface + surf = FourierRZToroidalSurface( + R_lmn=[9, 1, 1], + modes_R=[[0, 0], [0, 1], [1, 0]], + Z_lmn=[-1], + modes_Z=[[-1, 0]], + ) + surf_coils = MixedCoilSet( + FourierRZSurfaceCoil( + 1e6, surface=surf, secular_theta=0, secular_zeta=1, theta_n=[0.0] + ), + FourierRZSurfaceCoil( + 1e6, surface=surf, secular_theta=0, secular_zeta=1, theta_n=[np.pi] + ), + ) + coils5 = surf_coils.to_FourierXYZ(grid=grid).to_FourierRZSurface( + grid=grid, + check_intersection=False, + surface=surf, + secular_theta=0, + secular_zeta=1, + ) assert isinstance(coils1, MixedCoilSet) assert isinstance(coils2, MixedCoilSet) assert isinstance(coils3, MixedCoilSet) assert isinstance(coils4, MixedCoilSet) + assert isinstance(coils5, MixedCoilSet) assert all(isinstance(coil, SplineXYZCoil) for coil in coils1) assert all(isinstance(coil, FourierXYZCoil) for coil in coils2) assert all(isinstance(coil, FourierXYCoil) for coil in coils3) assert all(isinstance(coil, FourierPlanarCoil) for coil in coils4) + assert all(isinstance(coil, FourierRZSurfaceCoil) for coil in coils5) x0 = coils0.compute("x", grid=grid, basis="xyz") x1 = coils1.compute("x", grid=grid, basis="xyz") x2 = coils2.compute("x", grid=grid, basis="xyz") @@ -863,6 +969,11 @@ def test_convert_type(self): np.testing.assert_allclose( [xi["x"] for xi in x0], [xi["x"] for xi in x4], atol=1e-12 ) + x5_orig = surf_coils.compute("x", grid=grid, basis="xyz") + x5_rt = coils5.compute("x", grid=grid, basis="xyz") + np.testing.assert_allclose( + [xi["x"] for xi in x5_orig], [xi["x"] for xi in x5_rt], atol=1e-6 + ) B0 = coils0.compute_magnetic_field(np.array([[5, 2, 1]]), source_grid=grid) B1 = coils1.compute_magnetic_field(np.array([[5, 2, 1]]), source_grid=grid) B2 = coils2.compute_magnetic_field(np.array([[5, 2, 1]]), source_grid=grid) diff --git a/tests/test_compute_everything.py b/tests/test_compute_everything.py index 0889b46bfc..b3541ad14a 100644 --- a/tests/test_compute_everything.py +++ b/tests/test_compute_everything.py @@ -11,6 +11,7 @@ from desc.coils import ( FourierPlanarCoil, FourierRZCoil, + FourierRZSurfaceCoil, FourierXYCoil, FourierXYZCoil, SplineXYZCoil, @@ -21,6 +22,7 @@ from desc.geometry import ( FourierPlanarCurve, FourierRZCurve, + FourierRZSurfaceCurve, FourierRZToroidalSurface, FourierXYCurve, FourierXYZCurve, @@ -142,6 +144,13 @@ def test_compute_everything(): "desc.geometry.curve.SplineXYZCurve": FourierXYZCurve( X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] ).to_SplineXYZ(grid=LinearGrid(N=50)), + "desc.geometry.curve.FourierRZSurfaceCurve": FourierRZSurfaceCurve( + surface=FourierRZToroidalSurface(**elliptic_cross_section_with_torsion), + secular_theta=1, + secular_zeta=2, + theta_n=[0.1, 0.2, 0.3], + zeta_n=[-0.1, 0.2, -0.3], + ), # surfaces "desc.geometry.surface.FourierRZToroidalSurface": FourierRZToroidalSurface( **elliptic_cross_section_with_torsion @@ -198,6 +207,14 @@ def test_compute_everything(): "desc.coils.SplineXYZCoil": SplineXYZCoil( current=5, X=[5, 10, 2, 5], Y=[1, 2, 3, 1], Z=[-4, -5, -6, -4] ), + "desc.coils.FourierRZSurfaceCoil": FourierRZSurfaceCoil( + current=5, + surface=FourierRZToroidalSurface(**elliptic_cross_section_with_torsion), + secular_theta=1, + secular_zeta=2, + theta_n=[0.1, 0.2, 0.3], + zeta_n=[-0.1, 0.2, -0.3], + ), } assert things.keys() == data_index.keys(), ( f"Missing the parameterization {data_index.keys() - things.keys()}" @@ -229,6 +246,8 @@ def test_compute_everything(): "desc.geometry.curve.FourierPlanarCurve": {"grid": curvegrid1}, "desc.geometry.curve.FourierXYCurve": {"grid": curvegrid1}, "desc.geometry.curve.SplineXYZCurve": {"grid": curvegrid1}, + "desc.geometry.curve.FourierRZSurfaceCurve": {"grid": curvegrid1}, + "desc.coils.FourierRZSurfaceCoil": {"grid": curvegrid1}, "desc.magnetic_fields._core.OmnigenousField": {"grid": fieldgrid}, } diff --git a/tests/test_curves.py b/tests/test_curves.py index 5c089b1ed6..f17774d096 100644 --- a/tests/test_curves.py +++ b/tests/test_curves.py @@ -7,6 +7,8 @@ from desc.geometry import ( FourierPlanarCurve, FourierRZCurve, + FourierRZSurfaceCurve, + FourierRZToroidalSurface, FourierXYCurve, FourierXYZCurve, SplineXYZCurve, @@ -1270,3 +1272,244 @@ def test_compute_ndarray_error(self): c = SplineXYZCurve(X=R * np.cos(phi), Y=R * np.sin(phi), Z=np.zeros_like(phi)) with pytest.raises(TypeError): c.compute("length", grid=np.linspace(0, 1, 10)) + + +class TestFourierRZSurfaceCurve: + """Tests for FourierRZSurfaceCurve class.""" + + @pytest.mark.unit + def test_center(self): + """Test center of curve.""" + surf = FourierRZToroidalSurface() + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=1, + secular_zeta=0, + zeta_n=np.array([2]), + modes_zeta=np.array([0]), + ) + + # Default surface has major radius 10 + np.testing.assert_allclose( + c.compute("center")["center"][0], [10, 2, 0], atol=1e-12 + ) + + @pytest.mark.unit + def test_length(self): + """Test length of circular curve.""" + surf = FourierRZToroidalSurface() + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=1, + secular_zeta=0, + ) + + # Default surface has minor radius 1 + np.testing.assert_allclose(c.compute("length")["length"], 2 * np.pi) + + # Next, test a curve closing toroidally + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=0, + secular_zeta=1, + ) + + # Curve lies on the outboard midplane, tracing out a + # circle with radius 11 + np.testing.assert_allclose(c.compute("length")["length"], 2 * np.pi * 11) + + @pytest.mark.unit + def test_coords(self): + """Test lab frame coordinates of circular curve.""" + surf = FourierRZToroidalSurface() + + # curve closing after 2 toroidal turns + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=1, + secular_zeta=2, + ) + grid = LinearGrid(zeta=np.array([0, np.pi / 2, np.pi, 2 * np.pi])) + xyz = c.compute("x", grid=grid, basis="xyz")["x"] + + np.testing.assert_allclose(xyz[0], [11, 0, 0], atol=1e-12) + np.testing.assert_allclose(xyz[1], [-10, 0, -1], atol=1e-12) + np.testing.assert_allclose(xyz[2], [9, 0, 0], atol=1e-12) + np.testing.assert_allclose(xyz[3], [11, 0, 0], atol=1e-12) + + @pytest.mark.unit + def test_curvature(self): + """Test curvature of circular curve.""" + surf = FourierRZToroidalSurface() + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=1, + secular_zeta=0, + ) + + np.testing.assert_allclose(c.compute("curvature")["curvature"], 1) + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=0, + theta_n=[np.pi / 2], + secular_zeta=1, + ) + + np.testing.assert_allclose(c.compute("curvature")["curvature"], 1 / 10) + + @pytest.mark.unit + def test_torsion(self): + """Test torsion of circular curve.""" + surf = FourierRZToroidalSurface() + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=1, + secular_zeta=0, + ) + np.testing.assert_allclose(c.compute("torsion")["torsion"], 0) + + @pytest.mark.unit + def test_frenet(self): + """Test frenet-serret frame of circular curve.""" + surf = FourierRZToroidalSurface() + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=1, + secular_zeta=0, + ) + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] + np.testing.assert_allclose(T, np.array([[0, 0, -1]]), atol=1e-12) + np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12) + np.testing.assert_allclose(B, np.array([[0, 1, 0]]), atol=1e-12) + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=0, + secular_zeta=1, + ) + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] + np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12) + np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12) + np.testing.assert_allclose(B, np.array([[0, 0, 1]]), atol=1e-12) + + @pytest.mark.unit + def test_to_FourierXYZCurve(self): + """Test converting FourierRZSurfaceCurve to FourierXYZCurve object.""" + surf = FourierRZToroidalSurface() + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=1, + secular_zeta=0, + ) + c2 = c.to_FourierXYZ() + np.testing.assert_allclose( + c.compute("length")["length"], c2.compute("length")["length"] + ) + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=0, + secular_zeta=1, + ) + c2 = c.to_FourierXYZ() + np.testing.assert_allclose( + c.compute("length")["length"], c2.compute("length")["length"] + ) + + @pytest.mark.unit + def test_asserts_and_errors(self): + """Test asserts and errors of FourierRZSurfaceCurve.""" + surf = FourierRZToroidalSurface(NFP=8) + with pytest.raises(ValueError): + _ = FourierRZSurfaceCurve(surface=surf, secular_theta=2, secular_zeta=4) + with pytest.raises(ValueError): + _ = FourierRZSurfaceCurve( + surface=surf, secular_theta=2, secular_zeta=3, NFP=4 + ) + with pytest.raises(ValueError): + _ = FourierRZSurfaceCurve( + surface=surf, secular_theta=2, secular_zeta=3, NFP=5 + ) + with pytest.raises(ValueError): + _ = FourierRZSurfaceCurve(surface=surf, theta_n=[0, 1, 2], sym_theta="sin") + + @pytest.mark.unit + def test_modes(self): + """Test bases, mode numbering, and resolution for FourierRZSurfaceCurve.""" + surf = FourierRZToroidalSurface() + c = FourierRZSurfaceCurve( + surface=surf, secular_theta=2, secular_zeta=3, theta_n=[2, -1, 3, 0, 5] + ) + + assert c.N == 2 + assert c.theta_basis.N == 2 + assert c.zeta_basis.N == 0 + + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=2, + secular_zeta=3, + theta_n=[2, -1, 3, 0, 5], + modes_theta=[-8, -6, -4, -3, -2], + sym_theta="sin", + ) + assert c.sym_theta == "sin" + assert c.N == 8 + + np.testing.assert_allclose(c.theta_basis.modes[:, 2], np.arange(-8, 0)) + np.testing.assert_allclose(c.get_coeffs([-6, -5])[0], [-1, 0]) + c.change_resolution(N_theta=12) + c.set_coeffs(-9, theta_n=3) + assert c.N == 12 + np.testing.assert_allclose(c.theta_basis.modes[:, 2], np.arange(-12, 0)) + np.testing.assert_allclose(c.get_coeffs([-9, -8, -7])[0], [3, 2, 0]) + + @pytest.mark.unit + def test_from_values(self): + """Test fitting FourierRZSurfaceCurve from values.""" + surf = FourierRZToroidalSurface(NFP=4) + c = FourierRZSurfaceCurve( + surface=surf, + secular_theta=2, + secular_zeta=3, + theta_n=[2, -1, 3, 0, 5], + modes_theta=[-8, -6, -4, -3, -2], + zeta_n=[0, 2, 4], + sym_theta="sin", + NFP=2, + ) + grid = LinearGrid(zeta=np.linspace(0, 2 * np.pi, 201, endpoint=True)) + curve_data = c.compute(names=["s", "theta", "zeta"], grid=grid) + s, theta, zeta = curve_data["s"], curve_data["theta"], curve_data["zeta"] + coords = np.vstack([s, theta, zeta]).T + + with pytest.warns( + UserWarning, match="Unequal number of field periods for grid 1 and basis 2" + ): + c2 = FourierRZSurfaceCurve.from_values( + coords=coords, + surface=surf, + N_theta=c.N_theta, + N_zeta=c.N_zeta, + sym_theta=c.sym_theta, + sym_zeta=c.sym_zeta, + NFP=c.NFP, + ) + + assert c2.secular_theta == c.secular_theta + assert c2.secular_zeta == c.secular_zeta + np.testing.assert_allclose(c.theta_n, c2.theta_n, atol=1e-12) + np.testing.assert_allclose(c.zeta_n, c2.zeta_n, atol=1e-12) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index de7b703a2a..388f184264 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -6,9 +6,9 @@ import desc.examples from desc.backend import jnp, put -from desc.coils import CoilSet, FourierXYZCoil +from desc.coils import CoilSet, FourierRZSurfaceCoil, FourierXYZCoil, MixedCoilSet from desc.equilibrium import Equilibrium -from desc.geometry import FourierRZToroidalSurface +from desc.geometry import FourierRZSurfaceCurve, FourierRZToroidalSurface from desc.io import load from desc.magnetic_fields import OmnigenousField from desc.objectives import ( @@ -51,6 +51,7 @@ LinearObjectiveFromUser, ObjectiveFunction, ShareParameters, + SurfaceCurveConsistency, get_equilibrium_objective, get_fixed_axis_constraints, get_fixed_boundary_constraints, @@ -1235,3 +1236,62 @@ def test_NAE_asym_with_sym_axis(): conZ.build() assert conR._A.shape[0] == conR.dim_f assert conZ._A.shape[0] == conZ.dim_f + + +@pytest.mark.unit +def test_surface_curve_consistency(): + """Test SurfaceCurveConsistency on curves.""" + surf = FourierRZToroidalSurface( + R_lmn=[10, 1, 0.2], + modes_R=[[0, 0], [1, 0], [0, 1]], + Z_lmn=[-1, -0.2], + modes_Z=[[-1, 0], [0, -1]], + NFP=2, + ) + curve = FourierRZSurfaceCurve(surface=surf, secular_theta=1, secular_zeta=2) + obj = SurfaceCurveConsistency(surf, curve) + obj.build() + + nR = surf.R_basis.num_modes + nZ = surf.Z_basis.num_modes + assert obj.dim_f == nR + nZ + np.testing.assert_allclose(obj.compute(surf.params_dict, curve.params_dict), 0) + + # move the surface, checking it computes the new residual + surf_params = dict(surf.params_dict) + surf_params["R_lmn"] = surf_params["R_lmn"] + 0.1 + obj_value = obj.compute(surf_params, curve.params_dict) + np.testing.assert_allclose(obj_value[:nR], 0.1) + np.testing.assert_allclose(obj_value[nR:], 0) + + +@pytest.mark.unit +def test_surface_curve_consistency_coilset(): + """Test SurfaceCurveConsistency with coilsets.""" + surf = FourierRZToroidalSurface( + R_lmn=[10, 1, 0.2], + modes_R=[[0, 0], [1, 0], [0, 1]], + Z_lmn=[-1, -0.2], + modes_Z=[[-1, 0], [0, -1]], + NFP=2, + ) + + coil1 = FourierRZSurfaceCoil(1e5, surface=surf, secular_theta=1, secular_zeta=2) + coil2 = FourierRZSurfaceCoil(2e5, surface=surf, secular_theta=1, secular_zeta=3) + coil3 = FourierRZSurfaceCoil(3e5, surface=surf, secular_theta=1, secular_zeta=4) + coilset = MixedCoilSet((coil1, coil2, coil3), check_intersection=False) + + obj = SurfaceCurveConsistency(surf, coilset) + obj.build() + + nR = surf.R_basis.num_modes + nZ = surf.Z_basis.num_modes + assert obj.dim_f == 3 * (nR + nZ) + np.testing.assert_allclose(obj.compute(surf.params_dict, coilset.params_dict), 0) + + # move the surface, checking it computes the new residual for each coil + surf_params = dict(surf.params_dict) + surf_params["R_lmn"] = surf_params["R_lmn"] + 0.1 + obj_value = obj.compute(surf_params, coilset.params_dict).reshape(3, nR + nZ) + np.testing.assert_allclose(obj_value[:, :nR], 0.1) + np.testing.assert_allclose(obj_value[:, nR:], 0) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 4ab143d47b..9c9f19aeb7 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -28,7 +28,12 @@ from desc.compute import get_transforms from desc.equilibrium import Equilibrium from desc.examples import get -from desc.geometry import FourierPlanarCurve, FourierRZToroidalSurface, FourierXYZCurve +from desc.geometry import ( + FourierPlanarCurve, + FourierRZSurfaceCurve, + FourierRZToroidalSurface, + FourierXYZCurve, +) from desc.grid import ConcentricGrid, Grid, LinearGrid, QuadratureGrid from desc.integrals import Bounce2D from desc.io import load @@ -88,6 +93,7 @@ RotationalTransform, Shear, SurfaceCurrentRegularization, + SurfaceCurveConsistency, SurfaceQuadraticFlux, ToroidalCurrent, ToroidalFlux, @@ -3357,6 +3363,7 @@ class TestComputeScalarResolution: ToroidalFlux, SurfaceCurrentRegularization, VacuumBoundaryError, + SurfaceCurveConsistency, # no grid dependence for DeflationOperator DeflationOperator, # need to avoid blowup near the axis @@ -3876,6 +3883,7 @@ class TestObjectiveNaNGrad: PlasmaVesselDistance, QuadraticFlux, SurfaceCurrentRegularization, + SurfaceCurveConsistency, SurfaceQuadraticFlux, ToroidalFlux, VacuumBoundaryError, @@ -4148,6 +4156,16 @@ def test_objective_no_nangrad_surface_current_reg(self): g = obj.grad(obj.x(field)) assert not np.any(np.isnan(g)), "surface current regularization" + @pytest.mark.unit + def test_objective_no_nangrad_surface_curve_consistency(self): + """SurfaceCurveConsistency.""" + surf = FourierRZToroidalSurface() + curve = FourierRZSurfaceCurve(surface=surf, secular_theta=1, secular_zeta=2) + obj = ObjectiveFunction(SurfaceCurveConsistency(surf, curve), use_jit=False) + obj.build() + g = obj.grad(obj.x(surf, curve)) + assert not np.any(np.isnan(g)), "surface curve consistency" + @pytest.mark.unit @pytest.mark.parametrize( "objective", sorted(other_objectives, key=lambda x: str(x.__name__))