diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 2e96e7a76..d7fef3f04 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -415,6 +415,37 @@ def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="x", + label="\\mathbf{x}", + units="m", + units_long="meters", + description="Position vector along curve", + dim=3, + params=["A_n", "NFP_umbilic_facror", "NFP"], + transforms={"A": [[0, 0, 0]], "grid": []}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.UmbilicCurve", +) +def _x_UmbilicCurve(params, transforms, profiles, data, **kwargs): + A = transforms["A"].transform(params["A_n"], dz=0) + phi = transforms["grid"].nodes[:, 2] + NFP_umbilic_factor = params["NFP_umbilic_factor"] + NFP = params["NFP"] + theta = (A - NFP * phi) / NFP_umbilic_factor + rho = jnp.zeros_like(phi) + coords = jnp.stack([rho, theta, phi], axis=1) + # MUST CALCULATE XYZ on DESC grid + # convert to xyz for displacement and rotation + coords = rpz2xyz(coords) + # convert back to rpz + coords = xyz2rpz(coords) + data["x"] = coords + return data + + @register_compute_fun( name="x_s", label="\\partial_{s} \\mathbf{x}", diff --git a/desc/geometry/__init__.py b/desc/geometry/__init__.py index 2806afc37..638e93aeb 100644 --- a/desc/geometry/__init__.py +++ b/desc/geometry/__init__.py @@ -1,5 +1,11 @@ """Classes for representing geometric objects like curves and surfaces.""" from .core import Curve, Surface -from .curve import FourierPlanarCurve, FourierRZCurve, FourierXYZCurve, SplineXYZCurve +from .curve import ( + FourierPlanarCurve, + FourierRZCurve, + FourierUmbilicCurve, + FourierXYZCurve, + SplineXYZCurve, +) from .surface import FourierRZToroidalSurface, ZernikeRZToroidalSection diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index cc05f884b..1205fcd7b 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -16,7 +16,13 @@ from .core import Curve -__all__ = ["FourierRZCurve", "FourierXYZCurve", "FourierPlanarCurve", "SplineXYZCurve"] +__all__ = [ + "FourierRZCurve", + "FourierUmbilicCurve", + "FourierXYZCurve", + "FourierPlanarCurve", + "SplineXYZCurve", +] class FourierRZCurve(Curve): @@ -621,6 +627,192 @@ def from_values(cls, coords, N=10, s=None, basis="xyz", name=""): ) +class FourierUmbilicCurve(Curve): + """Curve parameterized by Fourier series for A in terms of toroidal angle phi. + + Parameters + ---------- + A_n: array-like + Fourier coefficients fo Z. + modes_A : array-like, optional + Mode numbers associated with Z_n, If not given defaults to [-n:n]]. + NFP : int + Number of field periods. + NFP_umbilic_factor : float + Rational number of the form 1/integer with integer>=1. + This is needed for the umbilic torus design. + sym : bool + Whether to enforce stellarator symmetry. + name : str + Name for this curve. + + """ + + _io_attrs_ = Curve._io_attrs_ + [ + "_A_n", + "_A_basis", + "_sym", + "_NFP", + "_NFP_umbilic_factor", + ] + + def __init__( + self, + A_n=0, + modes_A=None, + NFP=1, + NFP_umbilic_factor=1, + sym="auto", + name="", + ): + super().__init__(name) + A_n = np.atleast_1d(A_n) + if modes_A is None: + modes_A = np.arange(-(A_n.size // 2), A_n.size // 2 + 1) + + if A_n.size == 0: + A_n = np.array([0.0]) + modes_A = np.array([0]) + + modes_A = np.asarray(modes_A) + + assert A_n.size == modes_A.size, "A_n size and modes_A must be the same size" + + assert issubclass(modes_A.dtype.type, np.integer) + + if sym == "auto": + if np.all(A_n[modes_A >= 0] == 0): + sym = True + else: + sym = False + self._sym = sym + NA = np.max(abs(modes_A)) + N = NA + self._NFP = check_posint(NFP, "NFP", False) + self._NFP_umbilic_factor = check_posint( + NFP_umbilic_factor, "NFP_umbilic_factor", False + ) + self._A_basis = FourierSeries( + N, + int(NFP), + NFP_umbilic_factor=int(NFP_umbilic_factor), + sym="sin" if sym else False, + ) + + self._A_n = copy_coeffs(A_n, modes_A, self.A_basis.modes[:, 2]) + + @property + def sym(self): + """bool: Whether or not the curve is stellarator symmetric.""" + return self._sym + + @property + def A_basis(self): + """Spectral basis for A_Fourier series.""" + return self._A_basis + + @property + def NFP(self): + """Number of field periods.""" + return self._NFP + + @property + def NFP_umbilic_factor(self): + """NFP umbilic factor. Effective NFP -> NFP/NFP_umbilic_factor.""" + return self.__dict__.setdefault("_NFP_umbilic_factor", 1) + + def _NFP_umbilic_factor(self): + """NFP umbilic factor. Effective NFP -> NFP/NFP_umbilic_factor.""" + self._NFP_umbilic_factor = self.NFP_umbilic_factor + + @property + def N(self): + """Maximum mode number.""" + return self.A_basis.N + + def get_coeffs(self, n): + """Get Fourier coefficients for given mode number(s).""" + n = np.atleast_1d(n).astype(int) + A = np.zeros_like(n).astype(float) + + idxA = np.where(n[:, np.newaxis] == self.A_basis.modes[:, 2]) + + A[idxA[0]] = self.A_n[idxA[1]] + return A + + def set_coeffs(self, n, A=None): + """Set specific Fourier coefficients.""" + n, A = np.atleast_1d(n), np.atleast_1d(A) + A = np.broadcast_to(A, n.shape) + for nn, AA in zip(n, A): + if AA is not None: + idxA = self.A_basis.get_idx(0, 0, nn) + self.A_n = put(self.A_n, idxA, AA) + + @optimizable_parameter + @property + def A_n(self): + """Spectral coefficients for Z.""" + return self._A_n + + @A_n.setter + def A_n(self, new): + if len(new) == self.A_basis.num_modes: + self._A_n = jnp.asarray(new) + else: + raise ValueError( + f"A_n should have the same size as the basis, got {len(new)} for " + + f"basis with {self.A_basis.num_modes} modes" + ) + + @classmethod + def from_values(cls, coords, N=10, NFP=1, NFP_umbilic_factor=1, name="", sym=False): + """Fit coordinates to FourierRZCurve representation. + + Parameters + ---------- + coords: ndarray, shape (num_coords,2) + coordinates theta, zeta, the different of which is fit with a FourierSeries + N : int + Fourier resolution of the new R,Z representation. + NFP : int + Number of field periods, the curve will have a discrete toroidal symmetry + according to NFP. + NFP_umbilic_factor : int + Umbilic factor to fit curves that go around multiple times toroidally before + closing on themselves. + sym : bool + Whether to enforce stellarator symmetry. + basis : {"rpz", "xyz"} + basis for input coordinates. Defaults to "rpz" + name : str + Name for this curve. + + Returns + ------- + curve : FourierZSeries + New representation of the curve parameterized by Fourier series for Z. + + """ + theta = coords[:, 0] + phi = coords[:, 1] + A = NFP_umbilic_factor * theta + phi + + grid = LinearGrid(zeta=phi, NFP=1, sym=sym) + basis = FourierSeries(N=N, NFP=NFP, sym=sym) + transform = Transform(grid, basis, build_pinv=True) + A_n = transform.fit(A) + + return FourierUmbilicCurve( + A_n=A_n, + NFP=NFP, + NFP_umbilic_factor=NFP_umbilic_factor, + modes_A=basis.modes[:, 2], + sym=sym, + name=name, + ) + + class FourierPlanarCurve(Curve): """Curve that lies in a plane. diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index b17697b5c..d089785dc 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -1558,7 +1558,7 @@ class UmbilicHighCurvature(_Objective): _coordinates = "rtz" _units = "(m^-1)" - _print_value_fmt = "Umbilic high curvature: {:10.3e} " + _print_value_fmt = "Umbilic high curvature: " def __init__( self,