diff --git a/desc/coils.py b/desc/coils.py index 1b43aae624..77b84a45bb 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -6,8 +6,8 @@ import numpy as np from desc.backend import jnp +from desc.compute import rpz2xyz, xyz2rpz_vec from desc.geometry import FourierPlanarCurve, FourierRZCurve, FourierXYZCurve -from desc.geometry.utils import rpz2xyz, xyz2rpz_vec from desc.grid import Grid from desc.magnetic_fields import MagneticField, biot_savart @@ -46,7 +46,7 @@ def current(self, new): assert jnp.isscalar(new) or new.size == 1 self._current = new - def compute_magnetic_field(self, coords, params={}, basis="rpz"): + def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """Compute magnetic field at a set of points. The coil is discretized into a series of straight line segments, using @@ -64,6 +64,9 @@ def compute_magnetic_field(self, coords, params={}, basis="rpz"): parameters to pass to curve basis : {"rpz", "xyz"} basis for input coordinates and returned magnetic field + grid : Grid, int or None + Grid used to discretize coil. If an integer, uses that many equally spaced + points. Returns ------- @@ -76,8 +79,11 @@ def compute_magnetic_field(self, coords, params={}, basis="rpz"): coords = jnp.atleast_2d(coords) if basis == "rpz": coords = rpz2xyz(coords) - current = params.pop("current", self.current) - coil_coords = self.compute_coordinates(**params, basis="xyz") + if params is None: + current = self.current + else: + current = params.pop("current", self.current) + coil_coords = self.compute("x", params=params, grid=grid, basis="xyz")["x"] B = biot_savart(coords, coil_coords, current) if basis == "rpz": B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) @@ -110,8 +116,6 @@ class FourierRZCoil(Coil, FourierRZCurve): number of field periods sym : bool whether to enforce stellarator symmetry - grid : Grid - default grid for computation name : str name for this coil """ @@ -127,10 +131,9 @@ def __init__( modes_Z=None, NFP=1, sym="auto", - grid=None, name="", ): - super().__init__(current, R_n, Z_n, modes_R, modes_Z, NFP, sym, grid, name) + super().__init__(current, R_n, Z_n, modes_R, modes_Z, NFP, sym, name) class FourierXYZCoil(Coil, FourierXYZCurve): @@ -144,8 +147,6 @@ class FourierXYZCoil(Coil, FourierXYZCurve): fourier coefficients for X, Y, Z modes : array-like mode numbers associated with X_n etc. - grid : Grid - default grid or computation name : str name for this coil @@ -160,10 +161,9 @@ def __init__( Y_n=[0, 0, 0], Z_n=[-2, 0, 0], modes=None, - grid=None, name="", ): - super().__init__(current, X_n, Y_n, Z_n, modes, grid, name) + super().__init__(current, X_n, Y_n, Z_n, modes, name) class FourierPlanarCoil(Coil, FourierPlanarCurve): @@ -185,8 +185,6 @@ class FourierPlanarCoil(Coil, FourierPlanarCurve): fourier coefficients for radius from center as function of polar angle modes : array-like mode numbers associated with r_n - grid : Grid - default grid for computation name : str name for this coil @@ -201,10 +199,9 @@ def __init__( normal=[0, 1, 0], r_n=2, modes=None, - grid=None, name="", ): - super().__init__(current, center, normal, r_n, modes, grid, name) + super().__init__(current, center, normal, r_n, modes, name) class CoilSet(Coil, MutableSequence): @@ -251,35 +248,65 @@ def current(self, new): for coil, cur in zip(self.coils, new): coil.current = cur - @property - def grid(self): - """Grid: nodes for computation.""" - return self.coils[0].grid - - @grid.setter - def grid(self, new): - for coil in self.coils: - coil.grid = new - - def compute_coordinates(self, *args, **kwargs): - """Compute real space coordinates using underlying curve method.""" - return [coil.compute_coordinates(*args, **kwargs) for coil in self.coils] - - def compute_frenet_frame(self, *args, **kwargs): - """Compute Frenet frame using underlying curve method.""" - return [coil.compute_frenet_frame(*args, **kwargs) for coil in self.coils] + def _make_arraylike(self, x): + if isinstance(x, dict): + x = [x] * len(self) + try: + len(x) + except TypeError: + x = [x] * len(self) + assert len(x) == len(self) + return x + + def compute( + self, + names, + grid=None, + params=None, + transforms=None, + data=None, + **kwargs, + ): + """Compute the quantity given by name on grid, for each coil in the coilset. - def compute_curvature(self, *args, **kwargs): - """Compute curvature using underlying curve method.""" - return [coil.compute_curvature(*args, **kwargs) for coil in self.coils] + Parameters + ---------- + names : str or array-like of str + Name(s) of the quantity(s) to compute. + grid : Grid or int or array-like, optional + Grid of coordinates to evaluate at. Defaults to a Linear grid. + If an integer, uses that many equally spaced points. + If array-like, should be 1 value per coil. + params : dict of ndarray or array-like + Parameters from the equilibrium. Defaults to attributes of self. + If array-like, should be 1 value per coil. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + If array-like, should be 1 value per coil. + data : dict of ndarray or array-like + Data computed so far, generally output from other compute functions + If array-like, should be 1 value per coil. - def compute_torsion(self, *args, **kwargs): - """Compute torsion using underlying curve method.""" - return [coil.compute_torsion(*args, **kwargs) for coil in self.coils] + Returns + ------- + data : list of dict of ndarray + Computed quantity and intermediate variables, for each coil in the set. + List entries map to coils in coilset, each dict contains data for an + individual coil. - def compute_length(self, *args, **kwargs): - """Compute the length of the curve using underlying curve method.""" - return [coil.compute_length(*args, **kwargs) for coil in self.coils] + """ + grid = self._make_arraylike(grid) + params = self._make_arraylike(params) + transforms = self._make_arraylike(transforms) + data = self._make_arraylike(data) + return [ + coil.compute( + names, grid=grd, params=par, transforms=tran, data=dat, **kwargs + ) + for (coil, grd, par, tran, dat) in zip( + self.coils, grid, params, transforms, data + ) + ] def translate(self, *args, **kwargs): """Translate the coils along an axis.""" @@ -293,7 +320,7 @@ def flip(self, *args, **kwargs): """Flip the coils across a plane.""" [coil.flip(*args, **kwargs) for coil in self.coils] - def compute_magnetic_field(self, coords, params={}, basis="rpz"): + def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """Compute magnetic field at a set of points. Parameters @@ -305,18 +332,22 @@ def compute_magnetic_field(self, coords, params={}, basis="rpz"): or one for each member basis : {"rpz", "xyz"} basis for input coordinates and returned magnetic field + grid : Grid, int or None or array-like, optional + Grid used to discretize coil, either the same for all coils or one for each + member of the coilset. If an integer, uses that many equally spaced + points. Returns ------- field : ndarray, shape(n,3) magnetic field at specified points, in either rpz or xyz coordinates """ - if isinstance(params, dict): - params = [params] * len(self) - assert len(params) == len(self) + params = self._make_arraylike(params) + grid = self._make_arraylike(grid) + B = 0 - for coil, par in zip(self.coils, params): - B += coil.compute_magnetic_field(coords, par, basis) + for coil, par, grd in zip(self.coils, params, grid): + B += coil.compute_magnetic_field(coords, par, basis, grd) return B diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 5cfba65308..f2e214dcfa 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -30,6 +30,7 @@ _basis_vectors, _bootstrap, _core, + _curve, _equil, _field, _geometry, @@ -37,8 +38,10 @@ _profiles, _qs, _stability, + _surface, ) from .data_index import data_index +from .geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec from .utils import ( arg_order, compute, diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 22b74c55c5..a8a7e1eb32 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -666,6 +666,10 @@ def _b(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta", "|e_theta x e_zeta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _n_rho(params, transforms, profiles, data, **kwargs): # equal to e^rho / |e^rho| but works correctly for surfaces as well that don't have @@ -688,6 +692,10 @@ def _n_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta", "|e_zeta x e_rho|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _n_theta(params, transforms, profiles, data, **kwargs): data["n_theta"] = ( @@ -708,6 +716,10 @@ def _n_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta", "|e_rho x e_theta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _n_zeta(params, transforms, profiles, data, **kwargs): data["n_zeta"] = ( diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py new file mode 100644 index 0000000000..bb26246519 --- /dev/null +++ b/desc/compute/_curve.py @@ -0,0 +1,634 @@ +from desc.backend import jnp + +from .data_index import register_compute_fun +from .geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from .utils import cross, dot + + +@register_compute_fun( + name="s", + label="s", + units="~", + units_long="None", + description="Curve parameter, on [0, 2pi)", + dim=3, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.core.Curve", +) +def _s(params, transforms, profiles, data, **kwargs): + data["s"] = transforms["grid"].nodes[:, 2] + return data + + +def _rotation_matrix_from_normal(normal): + nx, ny, nz = normal + nxny = jnp.sqrt(nx**2 + ny**2) + R = jnp.array( + [ + [ny / nxny, -nx / nxny, 0], + [nx * nx / nxny, ny * nz / nxny, -nxny], + [nx, ny, nz], + ] + ).T + R = jnp.where(nxny == 0, jnp.eye(3), R) + return R + + +@register_compute_fun( + name="x", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=["s"], + parameterization="desc.geometry.curve.FourierPlanarCurve", + basis="basis", +) +def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + # create planar curve at z==0 + r = transforms["r"].transform(params["r_n"], dz=0) + Z = jnp.zeros_like(r) + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + coords = jnp.array([X, Y, Z]).T + # rotate into place + R = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, R.T) + params["center"] + coords = jnp.matmul(coords, transforms["rotmat"].T) + transforms["shift"] + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz(coords) + data["x"] = coords + return data + + +@register_compute_fun( + name="x_s", + label="\\partial_{s} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, first derivative", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0], [0, 0, 1]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=["s"], + parameterization="desc.geometry.curve.FourierPlanarCurve", + basis="basis", +) +def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + r = transforms["r"].transform(params["r_n"], dz=0) + dr = transforms["r"].transform(params["r_n"], dz=1) + dX = dr * jnp.cos(data["s"]) - r * jnp.sin(data["s"]) + dY = dr * jnp.sin(data["s"]) + r * jnp.cos(data["s"]) + dZ = jnp.zeros_like(dX) + coords = jnp.array([dX, dY, dZ]).T + A = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, A.T) + coords = jnp.matmul(coords, transforms["rotmat"].T) + if kwargs.get("basis", "rpz").lower() == "rpz": + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + Z = jnp.zeros_like(X) + xyzcoords = jnp.array([X, Y, Z]).T + xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"] + xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] + x, y, z = xyzcoords.T + coords = xyz2rpz_vec(coords, x=x, y=y) + data["x_s"] = coords + return data + + +@register_compute_fun( + name="x_ss", + label="\\partial_{ss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, second derivative", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=["s"], + parameterization="desc.geometry.curve.FourierPlanarCurve", + basis="basis", +) +def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + r = transforms["r"].transform(params["r_n"], dz=0) + dr = transforms["r"].transform(params["r_n"], dz=1) + d2r = transforms["r"].transform(params["r_n"], dz=2) + d2X = ( + d2r * jnp.cos(data["s"]) - 2 * dr * jnp.sin(data["s"]) - r * jnp.cos(data["s"]) + ) + d2Y = ( + d2r * jnp.sin(data["s"]) + 2 * dr * jnp.cos(data["s"]) - r * jnp.sin(data["s"]) + ) + d2Z = jnp.zeros_like(d2X) + coords = jnp.array([d2X, d2Y, d2Z]).T + A = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, A.T) + coords = jnp.matmul(coords, transforms["rotmat"].T) + if kwargs.get("basis", "rpz").lower() == "rpz": + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + Z = jnp.zeros_like(X) + xyzcoords = jnp.array([X, Y, Z]).T + xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"] + xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] + x, y, z = xyzcoords.T + coords = xyz2rpz_vec(coords, x=x, y=y) + data["x_ss"] = coords + return data + + +@register_compute_fun( + name="x_sss", + label="\\partial_{sss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, third derivative", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=["s"], + parameterization="desc.geometry.curve.FourierPlanarCurve", + basis="basis", +) +def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + r = transforms["r"].transform(params["r_n"], dz=0) + dr = transforms["r"].transform(params["r_n"], dz=1) + d2r = transforms["r"].transform(params["r_n"], dz=2) + d3r = transforms["r"].transform(params["r_n"], dz=3) + d3X = ( + d3r * jnp.cos(data["s"]) + - 3 * d2r * jnp.sin(data["s"]) + - 3 * dr * jnp.cos(data["s"]) + + r * jnp.sin(data["s"]) + ) + d3Y = ( + d3r * jnp.sin(data["s"]) + + 3 * d2r * jnp.cos(data["s"]) + - 3 * dr * jnp.sin(data["s"]) + - r * jnp.cos(data["s"]) + ) + d3Z = jnp.zeros_like(d3X) + coords = jnp.array([d3X, d3Y, d3Z]).T + A = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, A.T) + coords = jnp.matmul(coords, transforms["rotmat"].T) + if kwargs.get("basis", "rpz").lower() == "rpz": + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + Z = jnp.zeros_like(X) + xyzcoords = jnp.array([X, Y, Z]).T + xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"] + xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] + x, y, z = xyzcoords.T + coords = xyz2rpz_vec(coords, x=x, y=y) + data["x_sss"] = coords + return data + + +@register_compute_fun( + name="x", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "grid": [], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZCurve", + basis="basis", +) +def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_n"], dz=0) + Z = transforms["Z"].transform(params["Z_n"], dz=0) + phi = transforms["grid"].nodes[:, 2] + coords = jnp.stack([R, phi, Z], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz(coords) + coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz(coords) + data["x"] = coords + return data + + +@register_compute_fun( + name="x_s", + label="\\partial_{s} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, first derivative", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1]], + "Z": [[0, 0, 1]], + "grid": [], + "rotmat": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZCurve", + basis="basis", +) +def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_n"], dz=0) + dR = transforms["R"].transform(params["R_n"], dz=1) + dZ = transforms["Z"].transform(params["Z_n"], dz=1) + dphi = R0 + coords = jnp.stack([dR, dphi, dZ], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + coords = coords @ transforms["rotmat"].T + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["x_s"] = coords + return data + + +@register_compute_fun( + name="x_ss", + label="\\partial_{ss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, second derivative", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + "Z": [[0, 0, 2]], + "grid": [], + "rotmat": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZCurve", + basis="basis", +) +def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_n"], dz=0) + dR = transforms["R"].transform(params["R_n"], dz=1) + d2R = transforms["R"].transform(params["R_n"], dz=2) + d2Z = transforms["Z"].transform(params["Z_n"], dz=2) + R = d2R - R0 + Z = d2Z + # 2nd derivative wrt phi = 0 + phi = 2 * dR + coords = jnp.stack([R, phi, Z], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + coords = coords @ transforms["rotmat"].T + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["x_ss"] = coords + return data + + +@register_compute_fun( + name="x_sss", + label="\\partial_{sss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, third derivative", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]], + "Z": [[0, 0, 3]], + "grid": [], + "rotmat": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierRZCurve", + basis="basis", +) +def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_n"], dz=0) + dR = transforms["R"].transform(params["R_n"], dz=1) + d2R = transforms["R"].transform(params["R_n"], dz=2) + d3R = transforms["R"].transform(params["R_n"], dz=3) + d3Z = transforms["Z"].transform(params["Z_n"], dz=3) + R = d3R - 3 * dR + Z = d3Z + phi = 3 * d2R - R0 + coords = jnp.stack([R, phi, Z], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + coords = coords @ transforms["rotmat"].T + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["x_sss"] = coords + return data + + +@register_compute_fun( + name="x", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={ + "X": [[0, 0, 0]], + "Y": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierXYZCurve", + basis="basis", +) +def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs): + X = transforms["X"].transform(params["X_n"], dz=0) + Y = transforms["Y"].transform(params["Y_n"], dz=0) + Z = transforms["Z"].transform(params["Z_n"], dz=0) + coords = jnp.stack([X, Y, Z], axis=1) + coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz(coords) + data["x"] = coords + return data + + +@register_compute_fun( + name="x_s", + label="\\partial_{s} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, first derivative", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={ + "X": [[0, 0, 0], [0, 0, 1]], + "Y": [[0, 0, 0], [0, 0, 1]], + "Z": [[0, 0, 1]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierXYZCurve", + basis="basis", +) +def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): + dX = transforms["X"].transform(params["X_n"], dz=1) + dY = transforms["Y"].transform(params["Y_n"], dz=1) + dZ = transforms["Z"].transform(params["Z_n"], dz=1) + coords = jnp.stack([dX, dY, dZ], axis=1) + coords = coords @ transforms["rotmat"].T + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz_vec( + coords, + x=transforms["X"].transform(params["X_n"]) + transforms["shift"][0], + y=transforms["Y"].transform(params["Y_n"]) + transforms["shift"][1], + ) + data["x_s"] = coords + return data + + +@register_compute_fun( + name="x_ss", + label="\\partial_{ss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, second derivative", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={ + "X": [[0, 0, 0], [0, 0, 2]], + "Y": [[0, 0, 0], [0, 0, 2]], + "Z": [[0, 0, 2]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierXYZCurve", + basis="basis", +) +def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): + d2X = transforms["X"].transform(params["X_n"], dz=2) + d2Y = transforms["Y"].transform(params["Y_n"], dz=2) + d2Z = transforms["Z"].transform(params["Z_n"], dz=2) + coords = jnp.stack([d2X, d2Y, d2Z], axis=1) + coords = coords @ transforms["rotmat"].T + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz_vec( + coords, + x=transforms["X"].transform(params["X_n"]) + transforms["shift"][0], + y=transforms["Y"].transform(params["Y_n"]) + transforms["shift"][1], + ) + data["x_ss"] = coords + return data + + +@register_compute_fun( + name="x_sss", + label="\\partial_{sss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, third derivative", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={ + "X": [[0, 0, 0], [0, 0, 3]], + "Y": [[0, 0, 0], [0, 0, 3]], + "Z": [[0, 0, 3]], + "rotmat": [], + "shift": [], + }, + profiles=[], + coordinates="s", + data=[], + parameterization="desc.geometry.curve.FourierXYZCurve", + basis="basis", +) +def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): + d3X = transforms["X"].transform(params["X_n"], dz=3) + d3Y = transforms["Y"].transform(params["Y_n"], dz=3) + d3Z = transforms["Z"].transform(params["Z_n"], dz=3) + coords = jnp.stack([d3X, d3Y, d3Z], axis=1) + coords = coords @ transforms["rotmat"].T + if kwargs.get("basis", "rpz").lower() == "rpz": + coords = xyz2rpz_vec( + coords, + x=transforms["X"].transform(params["X_n"]) + transforms["shift"][0], + y=transforms["Y"].transform(params["Y_n"]) + transforms["shift"][1], + ) + data["x_sss"] = coords + return data + + +@register_compute_fun( + name="frenet_tangent", + label="\\mathbf{T}_{\\mathrm{Frenet-Serret}}", + units="~", + units_long="None", + description="Tangent unit vector to curve in Frenet-Serret frame", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="s", + data=["x_s"], + parameterization="desc.geometry.core.Curve", +) +def _frenet_tangent(params, transforms, profiles, data, **kwargs): + data["frenet_tangent"] = ( + data["x_s"] / jnp.linalg.norm(data["x_s"], axis=-1)[:, None] + ) + return data + + +@register_compute_fun( + name="frenet_normal", + label="\\mathbf{N}_{\\mathrm{Frenet-Serret}}", + units="~", + units_long="None", + description="Normal unit vector to curve in Frenet-Serret frame", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="s", + data=["x_ss"], + parameterization="desc.geometry.core.Curve", +) +def _frenet_normal(params, transforms, profiles, data, **kwargs): + data["frenet_normal"] = ( + data["x_ss"] / jnp.linalg.norm(data["x_ss"], axis=-1)[:, None] + ) + return data + + +@register_compute_fun( + name="frenet_binormal", + label="\\mathbf{B}_{\\mathrm{Frenet-Serret}}", + units="~", + units_long="None", + description="Binormal unit vector to curve in Frenet-Serret frame", + dim=3, + params=[], + transforms={"rotmat": []}, + profiles=[], + coordinates="s", + data=["frenet_tangent", "frenet_normal"], + parameterization="desc.geometry.core.Curve", +) +def _frenet_binormal(params, transforms, profiles, data, **kwargs): + data["frenet_binormal"] = cross( + data["frenet_tangent"], data["frenet_normal"] + ) * jnp.linalg.det(transforms["rotmat"]) + return data + + +@register_compute_fun( + name="curvature", + label="\\kappa", + units="m^{-1}", + units_long="Inverse meters", + description="Scalar curvature of the curve", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="s", + data=["x_s", "x_ss"], + parameterization="desc.geometry.core.Curve", +) +def _curvature(params, transforms, profiles, data, **kwargs): + dxn = jnp.linalg.norm(data["x_s"], axis=-1)[:, jnp.newaxis] + data["curvature"] = jnp.linalg.norm( + cross(data["x_s"], data["x_ss"]) / dxn**3, axis=-1 + ) + return data + + +@register_compute_fun( + name="torsion", + label="\\tau", + units="m^{-1}", + units_long="Inverse meters", + description="Scalar torsion of the curve", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="s", + data=["x_s", "x_ss", "x_sss"], + parameterization="desc.geometry.core.Curve", +) +def _torsion(params, transforms, profiles, data, **kwargs): + dxd2x = cross(data["x_s"], data["x_ss"]) + data["torsion"] = dot(dxd2x, data["x_sss"]) / jnp.linalg.norm(dxd2x, axis=-1) ** 2 + return data + + +@register_compute_fun( + name="length", + label="L", + units="m", + units_long="meters", + description="Length of the curve", + dim=0, + params=[], + transforms={}, + profiles=[], + coordinates="", + data=["s", "x_s"], + parameterization="desc.geometry.core.Curve", +) +def _length(params, transforms, profiles, data, **kwargs): + T = jnp.linalg.norm(data["x_s"], axis=-1) + data["length"] = jnp.trapz(T, data["s"]) + return data diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 62d34f086b..fbc9561a7f 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -22,6 +22,34 @@ def _V(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="V", + label="V", + units="m^{3}", + units_long="cubic meters", + description="Volume", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="", + data=["e_theta", "e_zeta", "x"], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", +) +def _V_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + # divergence theorem: integral(dV div [0, 0, Z]) = integral(dS dot [0, 0, Z]) + data["V"] = jnp.max( # take max in case there are multiple surfaces for some reason + jnp.abs( + surface_integrals( + transforms["grid"], + cross(data["e_theta"], data["e_zeta"])[:, 2] * data["x"][:, 2], + expand_out=False, + ) + ) + ) + return data + + @register_compute_fun( name="V(r)", label="V(\\rho)", @@ -98,6 +126,10 @@ def _V_rr_of_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="", data=["|e_rho x e_theta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _A(params, transforms, profiles, data, **kwargs): data["A"] = jnp.mean( @@ -111,6 +143,37 @@ def _A(params, transforms, profiles, data, **kwargs): return data +# TODO: compute cross section area for toroidal surface using stokes? + + +@register_compute_fun( + name="S", + label="S", + units="m^{2}", + units_long="square meters", + description="Surface area of outermost flux surface", + dim=0, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="", + data=["|e_theta x e_zeta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], +) +def _S(params, transforms, profiles, data, **kwargs): + data["S"] = jnp.max( + surface_integrals( + transforms["grid"], + data["|e_theta x e_zeta|"], + expand_out=False, + ) + ) + return data + + @register_compute_fun( name="S(r)", label="S(\\rho)", @@ -245,6 +308,10 @@ def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_rho", "e_theta_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _L_sff_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -265,6 +332,10 @@ def _L_sff_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_rho", "e_theta_z"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _M_sff_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -285,6 +356,10 @@ def _M_sff_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_rho", "e_zeta_z"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _N_sff_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -305,6 +380,10 @@ def _N_sff_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_tt", "g_tz", "g_zz", "L_sff_rho", "M_sff_rho", "N_sff_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_k1_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -337,6 +416,10 @@ def _curvature_k1_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_tt", "g_tz", "g_zz", "L_sff_rho", "M_sff_rho", "N_sff_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_k2_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -369,6 +452,10 @@ def _curvature_k2_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_rho", "curvature_k2_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_K_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -389,6 +476,10 @@ def _curvature_K_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_rho", "curvature_k2_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_H_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -575,6 +666,10 @@ def _curvature_H_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_zeta", "e_rho_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _L_sff_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -595,6 +690,10 @@ def _L_sff_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_zeta", "e_rho_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _M_sff_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -615,6 +714,10 @@ def _M_sff_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_zeta", "e_theta_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _N_sff_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -635,6 +738,10 @@ def _N_sff_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_rr", "g_rt", "g_tt", "L_sff_zeta", "M_sff_zeta", "N_sff_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_k1_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -667,6 +774,10 @@ def _curvature_k1_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_rr", "g_rt", "g_tt", "L_sff_zeta", "M_sff_zeta", "N_sff_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_k2_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -699,6 +810,10 @@ def _curvature_k2_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_zeta", "curvature_k2_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_K_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -719,6 +834,10 @@ def _curvature_K_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_zeta", "curvature_k2_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_H_zeta(params, transforms, profiles, data, **kwargs): # following notation from diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py index f881c37ce9..a42cb17553 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -54,6 +54,10 @@ def _sqrtg_pest(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs): data["|e_theta x e_zeta|"] = jnp.linalg.norm( @@ -74,6 +78,10 @@ def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _e_zeta_x_e_rho(params, transforms, profiles, data, **kwargs): data["|e_zeta x e_rho|"] = jnp.linalg.norm( @@ -94,6 +102,10 @@ def _e_zeta_x_e_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _e_rho_x_e_theta(params, transforms, profiles, data, **kwargs): data["|e_rho x e_theta|"] = jnp.linalg.norm( @@ -417,6 +429,10 @@ def _sqrtg_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_rr(params, transforms, profiles, data, **kwargs): data["g_rr"] = dot(data["e_rho"], data["e_rho"]) @@ -435,6 +451,10 @@ def _g_sub_rr(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tt(params, transforms, profiles, data, **kwargs): data["g_tt"] = dot(data["e_theta"], data["e_theta"]) @@ -453,6 +473,10 @@ def _g_sub_tt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_zz(params, transforms, profiles, data, **kwargs): data["g_zz"] = dot(data["e_zeta"], data["e_zeta"]) @@ -471,6 +495,10 @@ def _g_sub_zz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_rt(params, transforms, profiles, data, **kwargs): data["g_rt"] = dot(data["e_rho"], data["e_theta"]) @@ -489,6 +517,10 @@ def _g_sub_rt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_rz(params, transforms, profiles, data, **kwargs): data["g_rz"] = dot(data["e_rho"], data["e_zeta"]) @@ -507,6 +539,10 @@ def _g_sub_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tz(params, transforms, profiles, data, **kwargs): data["g_tz"] = dot(data["e_theta"], data["e_zeta"]) @@ -526,6 +562,10 @@ def _g_sub_tz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_theta_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tt_r(params, transforms, profiles, data, **kwargs): data["g_tt_r"] = 2 * dot(data["e_theta"], data["e_theta_r"]) @@ -545,6 +585,10 @@ def _g_sub_tt_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta", "e_theta_r", "e_zeta_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tz_r(params, transforms, profiles, data, **kwargs): data["g_tz_r"] = dot(data["e_theta_r"], data["e_zeta"]) + dot( diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py new file mode 100644 index 0000000000..c6b3981e72 --- /dev/null +++ b/desc/compute/_surface.py @@ -0,0 +1,703 @@ +from desc.backend import jnp + +from .data_index import register_compute_fun +from .geom_utils import rpz2xyz, rpz2xyz_vec + + +@register_compute_fun( + name="x", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"]) + Z = transforms["Z"].transform(params["Z_lmn"]) + phi = transforms["grid"].nodes[:, 2] + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["x"] = coords + return data + + +@register_compute_fun( + name="e_rho", + label="\\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_rho_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho"] = coords + return data + + +@register_compute_fun( + name="e_theta", + label="\\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 1, 0]], + "Z": [[0, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_theta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_theta"] = coords + return data + + +@register_compute_fun( + name="e_zeta", + label="\\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1]], + "Z": [[0, 0, 1]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_zeta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_lmn"], dz=0) + dR = transforms["R"].transform(params["R_lmn"], dz=1) + dZ = transforms["Z"].transform(params["Z_lmn"], dz=1) + dphi = R0 + coords = jnp.stack([dR, dphi, dZ], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_zeta"] = coords + return data + + +@register_compute_fun( + name="e_rho_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho_r"] = coords + return data + + +@register_compute_fun( + name="e_rho_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector, derivative wrt poloidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho_t"] = coords + return data + + +@register_compute_fun( + name="e_rho_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector, derivative wrt toroidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho_z"] = coords + return data + + +@register_compute_fun( + name="e_theta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_theta_r"] = coords + return data + + +@register_compute_fun( + name="e_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt poloidal angle", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 2, 0]], + "Z": [[0, 2, 0]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_theta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=2) + Z = transforms["Z"].transform(params["Z_lmn"], dt=2) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_theta_t"] = coords + return data + + +@register_compute_fun( + name="e_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt toroidal angle", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 1, 1]], + "Z": [[0, 1, 1]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_theta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) + dZ = transforms["Z"].transform(params["Z_lmn"], dt=1, dz=1) + dphi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([dR, dphi, dZ], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_theta_z"] = coords + return data + + +@register_compute_fun( + name="e_zeta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta_r"] = coords + return data + + +@register_compute_fun( + name="e_zeta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector, derivative wrt poloidal angle", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 1, 1]], + "Z": [[0, 1, 1]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_zeta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) + dZ = transforms["Z"].transform(params["Z_lmn"], dt=1, dz=1) + dphi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([dR, dphi, dZ], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_zeta_t"] = coords + return data + + +@register_compute_fun( + name="e_zeta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector, derivative wrt toroidal angle", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + "Z": [[0, 0, 2]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", + basis="basis", +) +def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_lmn"], dz=0) + dR = transforms["R"].transform(params["R_lmn"], dz=1) + d2R = transforms["R"].transform(params["R_lmn"], dz=2) + d2Z = transforms["Z"].transform(params["Z_lmn"], dz=2) + dphi = 2 * dR + coords = jnp.stack([d2R - R0, dphi, d2Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_zeta_z"] = coords + return data + + +@register_compute_fun( + name="x", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _x_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"]) + Z = transforms["Z"].transform(params["Z_lmn"]) + phi = transforms["grid"].nodes[:, 2] + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["x"] = coords + return data + + +@register_compute_fun( + name="e_rho", + label="\\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[1, 0, 0]], + "Z": [[1, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_rho_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=1) + Z = transforms["Z"].transform(params["Z_lmn"], dr=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["e_rho"] = coords + return data + + +@register_compute_fun( + name="e_theta", + label="\\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 1, 0]], + "Z": [[0, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_theta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_theta"] = coords + return data + + +@register_compute_fun( + name="e_zeta", + label="\\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_zeta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta"] = coords + return data + + +@register_compute_fun( + name="e_rho_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector, derivative wrt radial coordinate", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[2, 0, 0]], + "Z": [[2, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=2) + Z = transforms["Z"].transform(params["Z_lmn"], dr=2) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["e_rho_r"] = coords + return data + + +@register_compute_fun( + name="e_rho_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector, derivative wrt poloidal angle", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[1, 1, 0]], + "Z": [[1, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_rho_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=1, dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dr=1, dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["e_rho_t"] = coords + return data + + +@register_compute_fun( + name="e_rho_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector, derivative wrt toroidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho_z"] = coords + return data + + +@register_compute_fun( + name="e_theta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt radial coordinate", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[1, 1, 0]], + "Z": [[1, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=1, dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dr=1, dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["e_theta_r"] = coords + return data + + +@register_compute_fun( + name="e_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt poloidal angle", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 2, 0]], + "Z": [[0, 2, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_theta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=2) + Z = transforms["Z"].transform(params["Z_lmn"], dt=2) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_theta_t"] = coords + return data + + +@register_compute_fun( + name="e_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt toroidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_theta_z"] = coords + return data + + +@register_compute_fun( + name="e_zeta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta_r"] = coords + return data + + +@register_compute_fun( + name="e_zeta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector, derivative wrt poloidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_zeta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta_t"] = coords + return data + + +@register_compute_fun( + name="e_zeta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector, derivative wrt toroidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", + basis="basis", +) +def _e_zeta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta_z"] = coords + return data diff --git a/desc/geometry/utils.py b/desc/compute/geom_utils.py similarity index 100% rename from desc/geometry/utils.py rename to desc/compute/geom_utils.py diff --git a/desc/compute/utils.py b/desc/compute/utils.py index 8a90b3c965..b4f47459e8 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -88,7 +88,7 @@ def compute(parameterization, names, params, transforms, profiles, data=None, ** for name in names: if name not in data_index[p]: raise ValueError(f"Unrecognized value '{name}' for parameterization {p}.") - allowed_kwargs = {"helicity", "M_booz", "N_booz", "gamma"} + allowed_kwargs = {"helicity", "M_booz", "N_booz", "gamma", "basis"} bad_kwargs = kwargs.keys() - allowed_kwargs if len(bad_kwargs) > 0: raise ValueError(f"Unrecognized argument(s): {bad_kwargs}") @@ -315,7 +315,10 @@ def get_params(keys, obj, has_axis=False, **kwargs): params = [] for key in deps: params += data_index[p][key]["dependencies"]["params"] - params = _sort_args(list(set(params))) + if p == "desc.equilibrium.equilibrium.Equilibrium": + # probably need some way to distinguish between params from different instances + # of the same class? + params = _sort_args(list(set(params))) if isinstance(obj, str) or inspect.isclass(obj): return params params = {name: np.atleast_1d(getattr(obj, name)).copy() for name in params} @@ -378,6 +381,11 @@ def get_transforms(keys, obj, grid, **kwargs): build=True, build_pinv=True, ) + elif c == "rotmat": + transforms["rotmat"] = obj.rotmat + elif c == "shift": + transforms["shift"] = obj.shift + return transforms diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 0d9ab0e346..82972ca41a 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -1,19 +1,28 @@ """Base classes for curves and surfaces.""" +import numbers from abc import ABC, abstractmethod import numpy as np from desc.backend import jnp +from desc.compute import compute as compute_fun +from desc.compute import data_index +from desc.compute.geom_utils import reflection_matrix, rotation_matrix +from desc.compute.utils import ( + _parse_parameterization, + get_data_deps, + get_params, + get_transforms, +) +from desc.grid import LinearGrid, QuadratureGrid from desc.io import IOAble -from .utils import reflection_matrix, rotation_matrix - class Curve(IOAble, ABC): """Abstract base class for 1D curves in 3D space.""" - _io_attrs_ = ["_name", "_grid", "shift", "rotmat"] + _io_attrs_ = ["_name", "shift", "rotmat"] def __init__(self, name=""): self.shift = jnp.array([0, 0, 0]) @@ -29,30 +38,91 @@ def name(self): def name(self, new): self._name = new - @property - @abstractmethod - def grid(self): - """Grid: Nodes for computation.""" - - @abstractmethod - def compute_coordinates(self, params=None, grid=None, dt=0): - """Compute real space coordinates on predefined grid.""" - - @abstractmethod - def compute_frenet_frame(self, params=None, grid=None): - """Compute Frenet frame on predefined grid.""" + def compute( + self, + names, + grid=None, + params=None, + transforms=None, + data=None, + **kwargs, + ): + """Compute the quantity given by name on grid. - @abstractmethod - def compute_curvature(self, params=None, grid=None): - """Compute curvature on predefined grid.""" + Parameters + ---------- + names : str or array-like of str + Name(s) of the quantity(s) to compute. + grid : Grid or int, optional + Grid of coordinates to evaluate at. Defaults to a Linear grid. + If an integer, uses that many equally spaced points. + params : dict of ndarray + Parameters from the equilibrium. Defaults to attributes of self. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from grid + data : dict of ndarray + Data computed so far, generally output from other compute functions - @abstractmethod - def compute_torsion(self, params=None, grid=None): - """Compute torsion on predefined grid.""" + Returns + ------- + data : dict of ndarray + Computed quantity and intermediate variables. - @abstractmethod - def compute_length(self, params=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature.""" + """ + if isinstance(names, str): + names = [names] + if grid is None: + NFP = self.NFP if hasattr(self, "NFP") else 1 + grid = LinearGrid(N=2 * self.N + 5, NFP=NFP, endpoint=True) + if isinstance(grid, numbers.Integral): + NFP = self.NFP if hasattr(self, "NFP") else 1 + grid = LinearGrid(N=grid, NFP=NFP, endpoint=True) + + if params is None: + params = get_params(names, obj=self) + if transforms is None: + transforms = get_transforms(names, obj=self, grid=grid, **kwargs) + if data is None: + data = {} + profiles = {} + + p = _parse_parameterization(self) + deps = list(set(get_data_deps(names, obj=p) + names)) + dep0d = [ + dep + for dep in deps + if (data_index[p][dep]["coordinates"] == "") and (dep not in data) + ] + calc0d = bool(len(dep0d)) + # 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): + calc0d = False + + if calc0d: + grid0d = LinearGrid(N=2 * self.N + 5, NFP=NFP, endpoint=True) + data0d = compute_fun( + self, + dep0d, + params=params, + transforms=get_transforms(dep0d, obj=self, grid=grid0d, **kwargs), + profiles={}, + data=None, + **kwargs, + ) + # these should all be 0d quantities so don't need to compress/expand + data0d = {key: val for key, val in data0d.items() if key in dep0d} + data.update(data0d) + + data = compute_fun( + self, + names, + params=params, + transforms=transforms, + profiles=profiles, + data=data, + **kwargs, + ) + return data def translate(self, displacement=[0, 0, 0]): """Translate the curve by a rigid displacement in x, y, z.""" @@ -76,14 +146,14 @@ def __repr__(self): type(self).__name__ + " at " + str(hex(id(self))) - + " (name={}, grid={})".format(self.name, self.grid) + + " (name={})".format(self.name) ) class Surface(IOAble, ABC): """Abstract base class for 2d surfaces in 3d space.""" - _io_attrs_ = ["_name", "_grid", "_sym", "_L", "_M", "_N"] + _io_attrs_ = ["_name", "_sym", "_L", "_M", "_N"] @property def name(self): @@ -147,59 +217,120 @@ def _flip_orientation(self): one[self.Z_basis.modes[:, 1] < 0] *= -1 self.Z_lmn *= one - @property - @abstractmethod - def grid(self): - """Grid: Nodes for computation.""" - @abstractmethod def change_resolution(self, *args, **kwargs): """Change the maximum resolution.""" - @abstractmethod - def compute_coordinates(self, params=None, grid=None, dt=0, dz=0): - """Compute coordinate values at specified nodes.""" - - @abstractmethod - def compute_normal(self, params=None, grid=None): - """Compute normal vectors to the surface on predefined grid.""" - - @abstractmethod - def compute_surface_area(self, params=None, grids=None): - """Compute surface area via quadrature.""" - - def compute_curvature(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute gaussian and mean curvature. + def compute( + self, + names, + grid=None, + params=None, + transforms=None, + data=None, + **kwargs, + ): + """Compute the quantity given by name on grid. Parameters ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) + names : str or array-like of str + Name(s) of the quantity(s) to compute. + grid : Grid, optional + Grid of coordinates to evaluate at. Defaults to a Linear grid for constant + rho surfaces or a Quadrature grid for constant zeta surfaces. + params : dict of ndarray + Parameters from the equilibrium. Defaults to attributes of self. + transforms : dict of Transform + Transforms for R, Z, lambda, etc. Default is to build from grid + data : dict of ndarray + Data computed so far, generally output from other compute functions Returns ------- - K, H, k1, k2 : ndarray, shape(k,) - Gaussian, mean and 2 principle curvatures at points specified in grid. + data : dict of ndarray + Computed quantity and intermediate variables. """ - # following notation from - # https://en.wikipedia.org/wiki/Parametric_surface#Curvature - E, F, G = self._compute_first_fundamental_form(R_lmn, Z_lmn, grid) - L, M, N = self._compute_second_fundamental_form(R_lmn, Z_lmn, grid) - # coeffs of quadratic eqn for determinant - a = E * G - F**2 - b = F * M - L * G - E * N - c = L * N - M**2 - r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) - r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) - k1 = jnp.maximum(r1, r2) - k2 = jnp.minimum(r1, r2) - K = k1 * k2 - H = (k1 + k2) / 2 - return K, H, k1, k2 + if isinstance(names, str): + names = [names] + if grid is None: + if hasattr(self, "rho"): # constant rho surface + grid = LinearGrid( + rho=np.array(self.rho), + M=2 * self.M + 5, + N=2 * self.N + 5, + NFP=self.NFP, + ) + elif hasattr(self, "zeta"): # constant zeta surface + grid = QuadratureGrid(L=2 * self.L + 5, M=2 * self.M + 5, N=0, NFP=1) + grid._nodes[:, 2] = self.zeta + if params is None: + params = get_params(names, obj=self) + if transforms is None: + transforms = get_transforms(names, obj=self, grid=grid, **kwargs) + if data is None: + data = {} + profiles = {} + + p = _parse_parameterization(self) + deps = list(set(get_data_deps(names, obj=p) + names)) + dep0d = [ + dep + for dep in deps + if (data_index[p][dep]["coordinates"] == "") and (dep not in data) + ] + calc0d = bool(len(dep0d)) + # see if the grid we're already using will work for desired qtys + if calc0d and hasattr(self, "rho"): # constant rho surface + if ( + (grid.N >= 2 * self.N + 5) + and (grid.M > 2 * self.M + 5) + and isinstance(grid, LinearGrid) + ): + calc0d = False + else: + grid0d = LinearGrid( + rho=np.array(self.rho), + M=2 * self.M + 5, + N=2 * self.N + 5, + NFP=self.NFP, + ) + elif calc0d and hasattr(self, "zeta"): # constant zeta surface + if ( + (grid.L >= self.L + 1) + and (grid.M > 2 * self.M + 5) + and isinstance(grid, QuadratureGrid) + ): + calc0d = False + else: + grid0d = QuadratureGrid(L=2 * self.L + 5, M=2 * self.M + 5, N=0, NFP=1) + grid0d._nodes[:, 2] = self.zeta + + if calc0d: + data0d = compute_fun( + self, + dep0d, + params=params, + transforms=get_transforms(dep0d, obj=self, grid=grid0d, **kwargs), + profiles={}, + data=None, + **kwargs, + ) + # these should all be 0d quantities so don't need to compress/expand + data0d = {key: val for key, val in data0d.items() if key in dep0d} + data.update(data0d) + + data = compute_fun( + self, + names, + params=params, + transforms=transforms, + profiles=profiles, + data=data, + **kwargs, + ) + return data def __repr__(self): """Get the string form of the object.""" @@ -207,5 +338,5 @@ def __repr__(self): type(self).__name__ + " at " + str(hex(id(self))) - + " (name={}, grid={})".format(self.name, self.grid) + + " (name={})".format(self.name) ) diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index bc00e548b1..2ab0171695 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -6,12 +6,11 @@ from desc.backend import jnp, put from desc.basis import FourierSeries -from desc.grid import Grid, LinearGrid +from desc.grid import LinearGrid from desc.transform import Transform from desc.utils import copy_coeffs from .core import Curve -from .utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec __all__ = [ "FourierRZCurve", @@ -35,8 +34,6 @@ class FourierRZCurve(Curve): Number of field periods. sym : bool Whether to enforce stellarator symmetry. - grid : Grid - Default grid for computation. name : str Name for this curve. @@ -47,8 +44,6 @@ class FourierRZCurve(Curve): "_Z_n", "_R_basis", "_Z_basis", - "_R_transform", - "_Z_transform", "_sym", "_NFP", ] @@ -61,7 +56,6 @@ def __init__( modes_Z=None, NFP=1, sym="auto", - grid=None, name="", ): super().__init__(name) @@ -92,17 +86,12 @@ def __init__( NZ = np.max(abs(modes_Z)) N = max(NR, NZ) self._NFP = NFP - self._R_basis = FourierSeries(NR, NFP, sym="cos" if sym else False) - self._Z_basis = FourierSeries(NZ, NFP, sym="sin" if sym else False) + self._R_basis = FourierSeries(N, NFP, sym="cos" if sym else False) + self._Z_basis = FourierSeries(N, NFP, sym="sin" if sym else False) self._R_n = copy_coeffs(R_n, modes_R, self.R_basis.modes[:, 2]) self._Z_n = copy_coeffs(Z_n, modes_Z, self.Z_basis.modes[:, 2]) - if grid is None: - grid = LinearGrid(N=2 * N, NFP=self.NFP, endpoint=True) - self._grid = grid - self._R_transform, self._Z_transform = self._get_transforms(grid) - @property def sym(self): """Whether this curve has stellarator symmetry.""" @@ -130,26 +119,6 @@ def NFP(self, new): ), f"NFP should be a positive integer, got {type(new)}" self.change_resolution(NFP=new) - @property - def grid(self): - """Default grid for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif jnp.isscalar(new): - self._grid = LinearGrid(N=new, endpoint=True) - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._R_transform.grid = self.grid - self._Z_transform.grid = self.grid - @property def N(self): """Maximum mode number.""" @@ -174,11 +143,6 @@ def change_resolution(self, N=None, NFP=None, sym=None): self.Z_basis.change_resolution( N=N, NFP=self.NFP, sym="sin" if self.sym else self.sym ) - if hasattr(self.grid, "change_resolution"): - self.grid.change_resolution( - self.grid.L, self.grid.M, self.grid.N, self.NFP - ) - self._R_transform, self._Z_transform = self._get_transforms(self.grid) self.R_n = copy_coeffs(self.R_n, R_modes_old, self.R_basis.modes) self.Z_n = copy_coeffs(self.Z_n, Z_modes_old, self.Z_basis.modes) @@ -238,207 +202,6 @@ def Z_n(self, new): + f"basis with {self.Z_basis.num_modes} modes" ) - def _get_transforms(self, grid=None): - if grid is None: - return self._R_transform, self._Z_transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = np.linspace(0, 2 * np.pi, grid) - grid = np.atleast_1d(grid) - if grid.ndim == 1: - grid = np.pad(grid[:, np.newaxis], ((0, 0), (2, 0))) - grid = Grid(grid, sort=False) - R_transform = Transform( - grid, - self.R_basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - Z_transform = Transform( - grid, - self.Z_basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - return R_transform, Z_transform - - def compute_coordinates(self, R_n=None, Z_n=None, grid=None, dt=0, basis="rpz"): - """Compute values using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt: int - derivative order to compute - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - R, phi, Z or x, y, z coordinates of the curve at specified grid locations - in phi. - - """ - assert basis.lower() in ["rpz", "xyz"] - if R_n is None: - R_n = self.R_n - if Z_n is None: - Z_n = self.Z_n - R_transform, Z_transform = self._get_transforms(grid) - if dt == 0: - R = R_transform.transform(R_n, dz=0) - Z = Z_transform.transform(Z_n, dz=0) - phi = R_transform.grid.nodes[:, 2] - coords = jnp.stack([R, phi, Z], axis=1) - elif dt == 1: - R0 = R_transform.transform(R_n, dz=0) - dR = R_transform.transform(R_n, dz=dt) - dZ = Z_transform.transform(Z_n, dz=dt) - dphi = R0 - coords = jnp.stack([dR, dphi, dZ], axis=1) - elif dt == 2: - R0 = R_transform.transform(R_n, dz=0) - dR = R_transform.transform(R_n, dz=1) - d2R = R_transform.transform(R_n, dz=2) - d2Z = Z_transform.transform(Z_n, dz=2) - R = d2R - R0 - Z = d2Z - # 2nd derivative wrt phi = 0 - phi = 2 * dR - coords = jnp.stack([R, phi, Z], axis=1) - elif dt == 3: - R0 = R_transform.transform(R_n, dz=0) - dR = R_transform.transform(R_n, dz=1) - d2R = R_transform.transform(R_n, dz=2) - d3R = R_transform.transform(R_n, dz=3) - d3Z = Z_transform.transform(Z_n, dz=3) - R = d3R - 3 * dR - Z = d3Z - phi = 3 * d2R - R0 - coords = jnp.stack([R, phi, Z], axis=1) - else: - raise NotImplementedError( - "Derivatives higher than 3 have not been implemented in " - + "cylindrical coordinates." - ) - # convert to xyz for displacement and rotation - if dt > 0: - coords = rpz2xyz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = rpz2xyz(coords) - coords = coords @ self.rotmat.T + (self.shift[jnp.newaxis, :] * (dt == 0)) - if basis.lower() == "rpz": - if dt > 0: - coords = xyz2rpz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = xyz2rpz(coords) - return coords - - def compute_frenet_frame(self, R_n=None, Z_n=None, grid=None, basis="rpz"): - """Compute Frenet frame vectors using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - basis : {"rpz", "xyz"} - basis vectors to use for Frenet vector representation - - Returns - ------- - T, N, B : ndarrays, shape(k,3) - tangent, normal, and binormal vectors of the curve at specified grid - locations in phi - - """ - T = self.compute_coordinates(R_n, Z_n, grid, dt=1, basis=basis) - N = self.compute_coordinates(R_n, Z_n, grid, dt=2, basis=basis) - - T = T / jnp.linalg.norm(T, axis=1)[:, jnp.newaxis] - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - B = jnp.cross(T, N, axis=1) * jnp.linalg.det(self.rotmat) - - return T, N, B - - def compute_curvature(self, R_n=None, Z_n=None, grid=None): - """Compute curvature using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - kappa : ndarray, shape(k,) - curvature of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(R_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(R_n, Z_n, grid, dt=2) - dxn = jnp.linalg.norm(dx, axis=1)[:, jnp.newaxis] - kappa = jnp.linalg.norm(jnp.cross(dx, d2x, axis=1) / dxn**3, axis=1) - return kappa - - def compute_torsion(self, R_n=None, Z_n=None, grid=None): - """Compute torsion using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - tau : ndarray, shape(k,) - torsion of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(R_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(R_n, Z_n, grid, dt=2) - d3x = self.compute_coordinates(R_n, Z_n, grid, dt=3) - dxd2x = jnp.cross(dx, d2x, axis=1) - tau = ( - jnp.sum(dxd2x * d3x, axis=1) - / jnp.linalg.norm(dxd2x, axis=1)[:, jnp.newaxis] ** 2 - ) - return tau - - def compute_length(self, R_n=None, Z_n=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. If not given, defaults to values given - by R_n, Z_n attributes - grid : Grid or array-like - Toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - length : float - length of the curve approximated by quadrature - - """ - R_transform, Z_transform = self._get_transforms(grid) - T = self.compute_coordinates(R_n, Z_n, grid, dt=1) - T = jnp.linalg.norm(T, axis=1) - phi = R_transform.grid.nodes[:, 2] - return jnp.trapz(T, phi) - def to_FourierXYZCurve(self, N=None): """Convert to FourierXYZCurve representation. @@ -458,7 +221,7 @@ def to_FourierXYZCurve(self, N=None): N = max(self.R_basis.N, self.Z_basis.N) grid = LinearGrid(N=4 * N, NFP=1, sym=False) basis = FourierSeries(N=N, NFP=1, sym=False) - xyz = self.compute_coordinates(grid=grid, basis="xyz") + xyz = self.compute("x", grid=grid, basis="xyz")["x"] transform = Transform(grid, basis, build_pinv=True) X_n = transform.fit(xyz[:, 0]) Y_n = transform.fit(xyz[:, 1]) @@ -475,14 +238,19 @@ class FourierXYZCurve(Curve): Fourier coefficients for X, Y, Z modes : array-like mode numbers associated with X_n etc. - grid : Grid - default grid for computation name : str name for this curve """ - _io_attrs_ = Curve._io_attrs_ + ["_X_n", "_Y_n", "_Z_n", "_basis", "_transform"] + _io_attrs_ = Curve._io_attrs_ + [ + "_X_n", + "_Y_n", + "_Z_n", + "_X_basis", + "_Y_basis", + "_Z_basis", + ] def __init__( self, @@ -490,7 +258,6 @@ def __init__( Y_n=[0, 0, 0], Z_n=[-2, 0, 0], modes=None, - grid=None, name="", ): super().__init__(name) @@ -503,54 +270,45 @@ def __init__( assert issubclass(modes.dtype.type, np.integer) N = np.max(abs(modes)) - self._basis = FourierSeries(N, NFP=1, sym=False) - self._X_n = copy_coeffs(X_n, modes, self.basis.modes[:, 2]) - self._Y_n = copy_coeffs(Y_n, modes, self.basis.modes[:, 2]) - self._Z_n = copy_coeffs(Z_n, modes, self.basis.modes[:, 2]) + self._X_basis = FourierSeries(N, NFP=1, sym=False) + self._Y_basis = FourierSeries(N, NFP=1, sym=False) + self._Z_basis = FourierSeries(N, NFP=1, sym=False) + self._X_n = copy_coeffs(X_n, modes, self.X_basis.modes[:, 2]) + self._Y_n = copy_coeffs(Y_n, modes, self.Y_basis.modes[:, 2]) + self._Z_n = copy_coeffs(Z_n, modes, self.Z_basis.modes[:, 2]) - if grid is None: - grid = LinearGrid(N=2 * N, endpoint=True) - self._grid = grid - self._transform = self._get_transforms(grid) + @property + def X_basis(self): + """Spectral basis for X Fourier series.""" + return self._X_basis @property - def basis(self): - """Spectral basis for Fourier series.""" - return self._basis + def Y_basis(self): + """Spectral basis for Y Fourier series.""" + return self._Y_basis @property - def grid(self): - """Default grid for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif jnp.isscalar(new): - self._grid = LinearGrid(N=new, endpoint=True) - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._transform.grid = self.grid + def Z_basis(self): + """Spectral basis for Z Fourier series.""" + return self._Z_basis @property def N(self): """Maximum mode number.""" - return self.basis.N + return max(self.X_basis.N, self.Y_basis.N, self.Z_basis.N) def change_resolution(self, N=None): """Change the maximum angular resolution.""" if (N is not None) and (N != self.N): - modes_old = self.basis.modes - self.basis.change_resolution(N=N) - self._transform = self._get_transforms(self.grid) - self.X_n = copy_coeffs(self.X_n, modes_old, self.basis.modes) - self.Y_n = copy_coeffs(self.Y_n, modes_old, self.basis.modes) - self.Z_n = copy_coeffs(self.Z_n, modes_old, self.basis.modes) + Xmodes_old = self.X_basis.modes + Ymodes_old = self.Y_basis.modes + Zmodes_old = self.Z_basis.modes + self.X_basis.change_resolution(N=N) + self.Y_basis.change_resolution(N=N) + self.Z_basis.change_resolution(N=N) + self.X_n = copy_coeffs(self.X_n, Xmodes_old, self.X_basis.modes) + self.Y_n = copy_coeffs(self.Y_n, Ymodes_old, self.Y_basis.modes) + self.Z_n = copy_coeffs(self.Z_n, Zmodes_old, self.Z_basis.modes) def get_coeffs(self, n): """Get Fourier coefficients for given mode number(s).""" @@ -559,11 +317,13 @@ def get_coeffs(self, n): Y = np.zeros_like(n).astype(float) Z = np.zeros_like(n).astype(float) - idx = np.where(n[:, np.newaxis] == self.basis.modes[:, 2]) + Xidx = np.where(n[:, np.newaxis] == self.X_basis.modes[:, 2]) + Yidx = np.where(n[:, np.newaxis] == self.Y_basis.modes[:, 2]) + Zidx = np.where(n[:, np.newaxis] == self.Z_basis.modes[:, 2]) - X[idx[0]] = self.X_n[idx[1]] - Y[idx[0]] = self.Y_n[idx[1]] - Z[idx[0]] = self.Z_n[idx[1]] + X[Xidx[0]] = self.X_n[Xidx[1]] + Y[Yidx[0]] = self.Y_n[Yidx[1]] + Z[Zidx[0]] = self.Z_n[Zidx[1]] return X, Y, Z def set_coeffs(self, n, X=None, Y=None, Z=None): @@ -577,12 +337,18 @@ def set_coeffs(self, n, X=None, Y=None, Z=None): X = np.broadcast_to(X, n.shape) Y = np.broadcast_to(Y, n.shape) Z = np.broadcast_to(Z, n.shape) - for nn, XX, YY, ZZ in zip(n, X, Y, Z): - idx = self.basis.get_idx(0, 0, nn) + for nn, XX in zip(n, X): + idx = self.X_basis.get_idx(0, 0, nn) if XX is not None: self.X_n = put(self.X_n, idx, XX) + + for nn, YY in zip(n, Y): + idx = self.Y_basis.get_idx(0, 0, nn) if YY is not None: self.Y_n = put(self.Y_n, idx, YY) + + for nn, ZZ in zip(n, Z): + idx = self.Z_basis.get_idx(0, 0, nn) if ZZ is not None: self.Z_n = put(self.Z_n, idx, ZZ) @@ -593,12 +359,12 @@ def X_n(self): @X_n.setter def X_n(self, new): - if len(new) == self._basis.num_modes: + if len(new) == self.X_basis.num_modes: self._X_n = jnp.asarray(new) else: raise ValueError( f"X_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.X_basis.num_modes} modes." ) @property @@ -608,12 +374,12 @@ def Y_n(self): @Y_n.setter def Y_n(self, new): - if len(new) == self._basis.num_modes: + if len(new) == self.Y_basis.num_modes: self._Y_n = jnp.asarray(new) else: raise ValueError( f"Y_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.Y_basis.num_modes} modes." ) @property @@ -623,190 +389,14 @@ def Z_n(self): @Z_n.setter def Z_n(self, new): - if len(new) == self._basis.num_modes: + if len(new) == self.Z_basis.num_modes: self._Z_n = jnp.asarray(new) else: raise ValueError( f"Z_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.Z_basis.num_modes} modes." ) - def _get_transforms(self, grid=None): - if grid is None: - return self._transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = np.linspace(0, 2 * np.pi, grid) - grid = np.atleast_1d(grid) - if grid.ndim == 1: - grid = np.pad(grid[:, np.newaxis], ((0, 0), (2, 0))) - grid = Grid(grid, sort=False) - transform = Transform( - grid, - self.basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - return transform - - def compute_coordinates( - self, X_n=None, Y_n=None, Z_n=None, grid=None, dt=0, basis="xyz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt: int - derivative order to compute - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - X, Y, Z or R, phi, Z coordinates of the curve at specified grid locations - in phi. - - """ - assert basis.lower() in ["rpz", "xyz"] - if X_n is None: - X_n = self.X_n - if Y_n is None: - Y_n = self.Y_n - if Z_n is None: - Z_n = self.Z_n - - transform = self._get_transforms(grid) - X = transform.transform(X_n, dz=dt) - Y = transform.transform(Y_n, dz=dt) - Z = transform.transform(Z_n, dz=dt) - - coords = jnp.stack([X, Y, Z], axis=1) - coords = coords @ self.rotmat.T + (self.shift[jnp.newaxis, :] * (dt == 0)) - if basis.lower() == "rpz": - if dt > 0: - coords = xyz2rpz_vec( - coords, - x=coords[:, 1] + self.shift[0], - y=coords[:, 1] + self.shift[1], - ) - else: - coords = xyz2rpz(coords) - return coords - - def compute_frenet_frame( - self, X_n=None, Y_n=None, Z_n=None, grid=None, basis="xyz" - ): - """Compute Frenet frame vectors using specified coefficients. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for Frenet vector representation - - Returns - ------- - T, N, B : ndarrays, shape(k,3) - tangent, normal, and binormal vectors of the curve at specified grid - locations - - """ - T = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1, basis=basis) - N = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=2, basis=basis) - - T = T / jnp.linalg.norm(T, axis=1)[:, jnp.newaxis] - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - B = jnp.cross(T, N, axis=1) * jnp.linalg.det(self.rotmat) - - return T, N, B - - def compute_curvature(self, X_n=None, Y_n=None, Z_n=None, grid=None): - """Compute curvature using specified coefficients. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - kappa : ndarray, shape(k,) - curvature of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=2) - dxn = jnp.linalg.norm(dx, axis=1)[:, jnp.newaxis] - kappa = jnp.linalg.norm(jnp.cross(dx, d2x, axis=1) / dxn**3, axis=1) - return kappa - - def compute_torsion(self, X_n=None, Y_n=None, Z_n=None, grid=None): - """Compute torsion using specified coefficientsnp.empty((0, 3)). - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - tau : ndarray, shape(k,) - torsion of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=2) - d3x = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=3) - dxd2x = jnp.cross(dx, d2x, axis=1) - tau = ( - jnp.sum(dxd2x * d3x, axis=1) - / jnp.linalg.norm(dxd2x, axis=1)[:, jnp.newaxis] ** 2 - ) - return tau - - def compute_length(self, X_n=None, Y_n=None, Z_n=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - length : float - length of the curve approximated by quadrature - - """ - transform = self._get_transforms(grid) - T = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1) - T = jnp.linalg.norm(T, axis=1) - theta = transform.grid.nodes[:, 2] - return jnp.trapz(T, theta) - # TODO: to_rz method for converting to FourierRZCurve representation # (might be impossible to parameterize with toroidal angle phi) @@ -828,8 +418,6 @@ class FourierPlanarCurve(Curve): Fourier coefficients for radius from center as function of polar angle modes : array-like mode numbers associated with r_n - grid : Grid - default grid for computation name : str name for this curve @@ -839,8 +427,7 @@ class FourierPlanarCurve(Curve): "_r_n", "_center", "_normal", - "_basis", - "_transform", + "_r_basis", ] # Reference frame is centered at the origin with normal in the +Z direction. @@ -851,7 +438,6 @@ def __init__( normal=[0, 1, 0], r_n=2, modes=None, - grid=None, name="", ): super().__init__(name) @@ -863,52 +449,28 @@ def __init__( assert issubclass(modes.dtype.type, np.integer) N = np.max(abs(modes)) - self._basis = FourierSeries(N, NFP=1, sym=False) - self._r_n = copy_coeffs(r_n, modes, self.basis.modes[:, 2]) + self._r_basis = FourierSeries(N, NFP=1, sym=False) + self._r_n = copy_coeffs(r_n, modes, self.r_basis.modes[:, 2]) self.normal = normal self.center = center - if grid is None: - grid = LinearGrid(N=2 * self.N, endpoint=True) - self._grid = grid - self._transform = self._get_transforms(grid) @property - def basis(self): + def r_basis(self): """Spectral basis for Fourier series.""" - return self._basis - - @property - def grid(self): - """Default grid for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif jnp.isscalar(new): - self._grid = LinearGrid(N=new, endpoint=True) - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._transform.grid = self.grid + return self._r_basis @property def N(self): """Maximum mode number.""" - return self.basis.N + return self.r_basis.N def change_resolution(self, N=None): """Change the maximum angular resolution.""" if (N is not None) and (N != self.N): - modes_old = self.basis.modes - self.basis.change_resolution(N=N) - self._transform = self._get_transforms(self.grid) - self.r_n = copy_coeffs(self.r_n, modes_old, self.basis.modes) + modes_old = self.r_basis.modes + self.r_basis.change_resolution(N=N) + self.r_n = copy_coeffs(self.r_n, modes_old, self.r_basis.modes) @property def center(self): @@ -945,12 +507,12 @@ def r_n(self): @r_n.setter def r_n(self, new): - if len(np.asarray(new)) == self._basis.num_modes: + if len(np.asarray(new)) == self.r_basis.num_modes: self._r_n = jnp.asarray(new) else: raise ValueError( f"r_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.r_basis.num_modes} modes." ) def get_coeffs(self, n): @@ -958,7 +520,7 @@ def get_coeffs(self, n): n = np.atleast_1d(n).astype(int) r = np.zeros_like(n).astype(float) - idx = np.where(n[:, np.newaxis] == self.basis.modes[:, 2]) + idx = np.where(n[:, np.newaxis] == self.r_basis.modes[:, 2]) r[idx[0]] = self.r_n[idx[1]] return r @@ -968,253 +530,6 @@ def set_coeffs(self, n, r=None): n, r = np.atleast_1d(n), np.atleast_1d(r) r = np.broadcast_to(r, n.shape) for nn, rr in zip(n, r): - idx = self.basis.get_idx(0, 0, nn) + idx = self.r_basis.get_idx(0, 0, nn) if rr is not None: self.r_n = put(self.r_n, idx, rr) - - def _normal_rotmat(self, normal=None): - """Rotation matrix to rotate z axis into plane normal.""" - nx, ny, nz = normal - nxny = jnp.sqrt(nx**2 + ny**2) - - R = jnp.array( - [ - [ny / nxny, -nx / nxny, 0], - [nx * nx / nxny, ny * nz / nxny, -nxny], - [nx, ny, nz], - ] - ).T - return jnp.where(nxny == 0, jnp.eye(3), R) - - def _get_transforms(self, grid=None): - if grid is None: - return self._transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = np.linspace(0, 2 * np.pi, grid) - grid = np.atleast_1d(grid) - if grid.ndim == 1: - grid = np.pad(grid[:, np.newaxis], ((0, 0), (2, 0))) - grid = Grid(grid, sort=False) - transform = Transform( - grid, - self.basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - return transform - - def compute_coordinates( - self, center=None, normal=None, r_n=None, grid=None, dt=0, basis="xyz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt: int - derivative order to compute - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - X, Y, Z or R, phi, Z coordinates of the curve at specified grid locations - in theta. - - """ - assert basis.lower() in ["rpz", "xyz"] - if center is None: - center = self.center - if normal is None: - normal = self.normal - if r_n is None: - r_n = self.r_n - transform = self._get_transforms(grid) - r = transform.transform(r_n, dz=0) - t = transform.grid.nodes[:, -1] - Z = np.zeros_like(r) - - if dt == 0: - X = r * jnp.cos(t) - Y = r * jnp.sin(t) - coords = jnp.array([X, Y, Z]).T - elif dt == 1: - dr = transform.transform(r_n, dz=1) - dX = dr * jnp.cos(t) - r * jnp.sin(t) - dY = dr * jnp.sin(t) + r * jnp.cos(t) - coords = jnp.array([dX, dY, Z]).T - elif dt == 2: - dr = transform.transform(r_n, dz=1) - d2r = transform.transform(r_n, dz=2) - d2X = d2r * jnp.cos(t) - 2 * dr * jnp.sin(t) - r * jnp.cos(t) - d2Y = d2r * jnp.sin(t) + 2 * dr * jnp.cos(t) - r * jnp.sin(t) - coords = jnp.array([d2X, d2Y, Z]).T - elif dt == 3: - dr = transform.transform(r_n, dz=1) - d2r = transform.transform(r_n, dz=2) - d3r = transform.transform(r_n, dz=3) - d3X = ( - d3r * jnp.cos(t) - - 3 * d2r * jnp.sin(t) - - 3 * dr * jnp.cos(t) - + r * jnp.sin(t) - ) - d3Y = ( - d3r * jnp.sin(t) - + 3 * d2r * jnp.cos(t) - - 3 * dr * jnp.sin(t) - - r * jnp.cos(t) - ) - coords = jnp.array([d3X, d3Y, Z]).T - else: - raise NotImplementedError( - "Derivatives higher than 3 have not been implemented for planar curves." - ) - R = self._normal_rotmat(normal) - coords = jnp.matmul(coords, R.T) + (center * (dt == 0)) - coords = jnp.matmul(coords, self.rotmat.T) + (self.shift * (dt == 0)) - - if basis.lower() == "rpz": - X = r * jnp.cos(t) - Y = r * jnp.sin(t) - xyzcoords = jnp.array([X, Y, Z]).T - xyzcoords = jnp.matmul(xyzcoords, R.T) + center - xyzcoords = jnp.matmul(xyzcoords, self.rotmat.T) + self.shift - x, y, z = xyzcoords.T - coords = xyz2rpz_vec(coords, x=x, y=y) - - return coords - - def compute_frenet_frame( - self, center=None, normal=None, r_n=None, grid=None, basis="xyz" - ): - """Compute Frenet frame vectors using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for Frenet vector representation - - Returns - ------- - T, N, B : ndarrays, shape(k,3) - tangent, normal, and binormal vectors of the curve at specified grid - locations in theta. - - """ - T = self.compute_coordinates(center, normal, r_n, grid, dt=1, basis=basis) - N = self.compute_coordinates(center, normal, r_n, grid, dt=2, basis=basis) - - T = T / jnp.linalg.norm(T, axis=1)[:, jnp.newaxis] - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - B = jnp.cross(T, N, axis=1) * jnp.linalg.det(self.rotmat) - - return T, N, B - - def compute_curvature(self, center=None, normal=None, r_n=None, grid=None): - """Compute curvature using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - kappa : ndarray, shape(k,) - curvature of the curve at specified grid locations in theta - - """ - dx = self.compute_coordinates(center, normal, r_n, grid, dt=1) - d2x = self.compute_coordinates(center, normal, r_n, grid, dt=2) - dxn = jnp.linalg.norm(dx, axis=1)[:, jnp.newaxis] - kappa = jnp.linalg.norm(jnp.cross(dx, d2x, axis=1) / dxn**3, axis=1) - return kappa - - def compute_torsion(self, center=None, normal=None, r_n=None, grid=None): - """Compute torsion using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - tau : ndarray, shape(k,) - torsion of the curve at specified grid locations in phi - - """ - # torsion is zero for planar curves - transform = self._get_transforms(grid) - torsion = jnp.zeros_like(transform.grid.nodes[:, -1]) - return torsion - - def compute_length(self, center=None, normal=None, r_n=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - length : float - length of the curve approximated by quadrature - - """ - transform = self._get_transforms(grid) - T = self.compute_coordinates(center, normal, r_n, grid, dt=1) - T = jnp.linalg.norm(T, axis=1) - theta = transform.grid.nodes[:, 2] - return jnp.trapz(T, theta) diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py index 773ac369a5..fe3992ac81 100644 --- a/desc/geometry/surface.py +++ b/desc/geometry/surface.py @@ -7,13 +7,10 @@ from desc.backend import jnp, put, sign from desc.basis import DoubleFourierSeries, ZernikePolynomial -from desc.grid import Grid, LinearGrid from desc.io import InputReader -from desc.transform import Transform from desc.utils import copy_coeffs from .core import Surface -from .utils import rpz2xyz, rpz2xyz_vec __all__ = ["FourierRZToroidalSurface", "ZernikeRZToroidalSection"] @@ -34,10 +31,8 @@ class FourierRZToroidalSurface(Surface): sym : bool whether to enforce stellarator symmetry. Default is "auto" which enforces if modes are symmetric. If True, non-symmetric modes will be truncated. - rho : float (0,1) + rho : float [0,1] flux surface label for the toroidal surface - grid : Grid - default grid for computation name : str name for this surface check_orientation : bool @@ -52,8 +47,6 @@ class FourierRZToroidalSurface(Surface): "_Z_lmn", "_R_basis", "_Z_basis", - "_R_transform", - "_Z_transform", "rho", "_NFP", ] @@ -125,16 +118,6 @@ def __init__( self._flip_orientation() assert self._compute_orientation() == 1 - if grid is None: - grid = LinearGrid( - M=2 * self.M, - N=2 * self.N, - NFP=self.NFP, - rho=np.asarray(self.rho), - endpoint=True, - ) - self._grid = grid - self._R_transform, self._Z_transform = self._get_transforms(grid) self.name = name @property @@ -159,24 +142,6 @@ def Z_basis(self): """DoubleFourierSeries: Spectral basis for Z.""" return self._Z_basis - @property - def grid(self): - """Grid: Nodes for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._R_transform.grid = self.grid - self._Z_transform.grid = self.grid - def change_resolution(self, *args, **kwargs): """Change the maximum poloidal and toroidal resolution.""" assert ( @@ -219,11 +184,6 @@ def change_resolution(self, *args, **kwargs): self.Z_basis.change_resolution( M=M, N=N, NFP=self.NFP, sym="sin" if self.sym else self.sym ) - if hasattr(self.grid, "change_resolution"): - self.grid.change_resolution( - self.grid.L, self.grid.M, self.grid.N, self.NFP - ) - self._R_transform, self._Z_transform = self._get_transforms(self.grid) self.R_lmn = copy_coeffs(self.R_lmn, R_modes_old, self.R_basis.modes) self.Z_lmn = copy_coeffs(self.Z_lmn, Z_modes_old, self.Z_basis.modes) self._M = M @@ -297,195 +257,6 @@ def set_coeffs(self, m, n=0, R=None, Z=None): idxZ = self.Z_basis.get_idx(0, mm, nn) self.Z_lmn = put(self.Z_lmn, idxZ, ZZ) - def _get_transforms(self, grid=None): - if grid is None: - return self._R_transform, self._Z_transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = LinearGrid( - M=grid, N=grid, rho=np.asarray(self.rho), NFP=self.NFP - ) - elif len(grid) == 2: - grid = LinearGrid( - M=grid[0], N=grid[1], rho=np.asarray(self.rho), NFP=self.NFP - ) - elif grid.shape[1] == 2: - grid = np.pad(grid, ((0, 0), (1, 0)), constant_values=self.rho) - grid = Grid(grid, sort=False) - else: - grid = Grid(grid, sort=False) - R_transform = Transform( - grid, - self.R_basis, - derivs=np.array( - [[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 0, 1], [0, 0, 2], [0, 1, 1]] - ), - ) - Z_transform = Transform( - grid, - self.Z_basis, - derivs=np.array( - [[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 0, 1], [0, 0, 2], [0, 1, 1]] - ), - ) - return R_transform, Z_transform - - def _compute_first_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the first fundamental form.""" - rt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - rz = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=1) - E = jnp.sum(rt * rt, axis=-1) - F = jnp.sum(rt * rz, axis=-1) - G = jnp.sum(rz * rz, axis=-1) - return E, F, G - - def _compute_second_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the second fundamental form.""" - rtt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=2) - rtz = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1, dz=1) - rzz = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=2) - n = self.compute_normal(R_lmn, Z_lmn, grid) - L = jnp.sum(rtt * n, axis=-1) - M = jnp.sum(rtz * n, axis=-1) - N = jnp.sum(rzz * n, axis=-1) - return L, M, N - - def compute_coordinates( - self, R_lmn=None, Z_lmn=None, grid=None, dt=0, dz=0, basis="rpz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt, dz: int - derivative order to compute in theta, zeta - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - R, phi, Z or x, y, z coordinates of the surface at points specified in grid - """ - assert basis.lower() in ["rpz", "xyz"] - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - if dz == 0: - R = R_transform.transform(R_lmn, dt=dt, dz=0) - Z = Z_transform.transform(Z_lmn, dt=dt, dz=0) - phi = R_transform.grid.nodes[:, 2] * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - elif dz == 1: - R0 = R_transform.transform(R_lmn, dt=dt, dz=0) - dR = R_transform.transform(R_lmn, dt=dt, dz=1) - dZ = Z_transform.transform(Z_lmn, dt=dt, dz=1) - dphi = R0 * (dt == 0) - coords = jnp.stack([dR, dphi, dZ], axis=1) - elif dz == 2: - R0 = R_transform.transform(R_lmn, dt=dt, dz=0) - dR = R_transform.transform(R_lmn, dt=dt, dz=1) - d2R = R_transform.transform(R_lmn, dt=dt, dz=2) - d2Z = Z_transform.transform(Z_lmn, dt=dt, dz=2) - R = d2R - R0 - Z = d2Z - # 2nd derivative wrt phi = 0 - phi = 2 * dR * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - elif dz == 3: - R0 = R_transform.transform(R_lmn, dt=dt, dz=0) - dR = R_transform.transform(R_lmn, dt=dt, dz=1) - d2R = R_transform.transform(R_lmn, dt=dt, dz=2) - d3R = R_transform.transform(R_lmn, dt=dt, dz=3) - d3Z = Z_transform.transform(Z_lmn, dt=dt, dz=3) - R = d3R - 3 * dR - Z = d3Z - phi = (3 * d2R - R0) * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - else: - raise NotImplementedError( - "Derivatives higher than 3 have not been implemented in " - + "cylindrical coordinates." - ) - if basis.lower() == "xyz": - if (dt > 0) or (dz > 0): - coords = rpz2xyz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = rpz2xyz(coords) - return coords - - def compute_normal(self, R_lmn=None, Z_lmn=None, grid=None, basis="rpz"): - """Compute normal vector to surface on default grid. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for normal vector representation - - Returns - ------- - N : ndarray, shape(k,3) - normal vector to surface in specified coordinates - """ - assert basis.lower() in ["rpz", "xyz"] - - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - r_t = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - r_z = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=1) - - N = jnp.cross(r_t, r_z, axis=1) - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - if basis.lower() == "xyz": - phi = R_transform.grid.nodes[:, 2] - N = rpz2xyz_vec(N, phi=phi) - return N - - def compute_surface_area(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute surface area via quadrature. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - area : float - surface area - """ - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - r_t = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - r_z = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=1) - - N = jnp.cross(r_t, r_z, axis=1) - return jnp.sum(R_transform.grid.weights * jnp.linalg.norm(N, axis=1)) - @classmethod def from_input_file(cls, path): """Create a surface from Fourier coefficients in a DESC or VMEC input file. @@ -593,7 +364,7 @@ class ZernikeRZToroidalSection(Surface): whether to enforce stellarator symmetry. Default is "auto" which enforces if modes are symmetric. If True, non-symmetric modes will be truncated. spectral_indexing : {``'ansi'``, ``'fringe'``} - Indexing method, default value = ``'fringe'`` + Indexing method, default value = ``'ansi'`` For L=0, all methods are equivalent and give a "chevron" shaped basis (only the outer edge of the zernike pyramid of width M). @@ -609,10 +380,8 @@ class ZernikeRZToroidalSection(Surface): decreasing size, ending in a diamond shape for L=2*M where the traditional fringe/U of Arizona indexing is recovered. For L > 2*M, adds chevrons to the bottom, making a hexagonal diamond - zeta : float (0,2pi) + zeta : float [0,2pi) toroidal angle for the section. - grid : Grid - default grid for computation name : str name for this surface check_orientation : bool @@ -627,8 +396,6 @@ class ZernikeRZToroidalSection(Surface): "_Z_lmn", "_R_basis", "_Z_basis", - "_R_transform", - "_Z_transform", "zeta", "_spectral_indexing", ] @@ -639,7 +406,7 @@ def __init__( Z_lmn=None, modes_R=None, modes_Z=None, - spectral_indexing="fringe", + spectral_indexing="ansi", sym="auto", zeta=0.0, grid=None, @@ -708,12 +475,6 @@ def __init__( self._flip_orientation() assert self._compute_orientation() == 1 - if grid is None: - grid = LinearGrid( - L=self.L, M=2 * self.M, zeta=np.asarray(self.zeta), endpoint=True - ) - self._grid = grid - self._R_transform, self._Z_transform = self._get_transforms(grid) self.name = name @property @@ -731,24 +492,6 @@ def Z_basis(self): """ZernikePolynomial: Spectral basis for Z.""" return self._Z_basis - @property - def grid(self): - """Grid: Nodes for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._R_transform.grid = self.grid - self._Z_transform.grid = self.grid - def change_resolution(self, *args, **kwargs): """Change the maximum radial and poloidal resolution.""" assert ( @@ -785,9 +528,6 @@ def change_resolution(self, *args, **kwargs): self.Z_basis.change_resolution( L=L, M=M, sym="sin" if self.sym else self.sym ) - if hasattr(self.grid, "change_resolution"): - self.grid.change_resolution(self.grid.L, self.grid.M, self.grid.N) - self._R_transform, self._Z_transform = self._get_transforms(self.grid) self.R_lmn = copy_coeffs(self.R_lmn, R_modes_old, self.R_basis.modes) self.Z_lmn = copy_coeffs(self.Z_lmn, Z_modes_old, self.Z_basis.modes) self._L = L @@ -860,155 +600,3 @@ def set_coeffs(self, l, m=0, R=None, Z=None): if ZZ is not None: idxZ = self.Z_basis.get_idx(ll, mm, 0) self.Z_lmn = put(self.Z_lmn, idxZ, ZZ) - - def _get_transforms(self, grid=None): - if grid is None: - return self._R_transform, self._Z_transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = LinearGrid(L=grid, M=grid, zeta=np.asarray(self.zeta)) - elif len(grid) == 2: - grid = LinearGrid(L=grid[0], M=grid[1], zeta=np.asarray(self.zeta)) - elif grid.shape[1] == 2: - grid = np.pad(grid, ((0, 0), (0, 1)), constant_values=self.zeta) - grid = Grid(grid, sort=False) - else: - grid = Grid(grid, sort=False) - R_transform = Transform( - grid, - self.R_basis, - derivs=np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [0, 1, 0], [0, 2, 0], [1, 1, 0]] - ), - ) - Z_transform = Transform( - grid, - self.Z_basis, - derivs=np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [0, 1, 0], [0, 2, 0], [1, 1, 0]] - ), - ) - return R_transform, Z_transform - - def _compute_first_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the first fundamental form.""" - rr = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=1) - rt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - E = jnp.sum(rr * rr, axis=-1) - F = jnp.sum(rr * rt, axis=-1) - G = jnp.sum(rt * rt, axis=-1) - return E, F, G - - def _compute_second_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the second fundamental form.""" - rrr = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=2) - rrt = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=1, dt=1) - rtt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=2) - n = self.compute_normal(R_lmn, Z_lmn, grid) - L = jnp.sum(rrr * n, axis=-1) - M = jnp.sum(rrt * n, axis=-1) - N = jnp.sum(rtt * n, axis=-1) - return L, M, N - - def compute_coordinates( - self, R_lmn=None, Z_lmn=None, grid=None, dr=0, dt=0, basis="rpz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - zernike coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,1)x(0,2pi) - dr, dt: int - derivative order to compute in rho, theta - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - R, phi, Z or x, y, z coordinates of the surface at points specified in grid - """ - assert basis.lower() in ["rpz", "xyz"] - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - R = R_transform.transform(R_lmn, dr=dr, dt=dt) - Z = Z_transform.transform(Z_lmn, dr=dr, dt=dt) - phi = R_transform.grid.nodes[:, 2] * (dr == 0) * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - if basis.lower() == "xyz": - if (dt > 0) or (dr > 0): - coords = rpz2xyz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = rpz2xyz(coords) - return coords - - def compute_normal(self, R_lmn=None, Z_lmn=None, grid=None, basis="rpz"): - """Compute normal vector to surface on default grid. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - zernike coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,1)x(0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for normal vector representation - - Returns - ------- - N : ndarray, shape(k,3) - normal vector to surface in specified coordinates - """ - assert basis.lower() in ["rpz", "xyz"] - - R_transform, Z_transform = self._get_transforms(grid) - - phi = R_transform.grid.nodes[:, -1] - - # normal vector is a constant 1*phihat - N = jnp.array([jnp.zeros_like(phi), jnp.ones_like(phi), jnp.zeros_like(phi)]).T - - if basis.lower() == "xyz": - N = rpz2xyz_vec(N, phi=phi) - - return N - - def compute_surface_area(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute surface area via quadrature. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - zernike coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,1)x(0,2pi) - - Returns - ------- - area : float - surface area - - """ - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - r_r = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=1) - r_t = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - - N = jnp.cross(r_r, r_t, axis=1) - return jnp.sum(R_transform.grid.weights * jnp.linalg.norm(N, axis=1)) / ( - 2 * np.pi - ) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 442c10e278..f2315e96a5 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -6,8 +6,8 @@ from netCDF4 import Dataset from desc.backend import jit, jnp, odeint +from desc.compute import rpz2xyz_vec, xyz2rpz from desc.derivatives import Derivative -from desc.geometry.utils import rpz2xyz_vec, xyz2rpz from desc.grid import Grid from desc.interpolate import _approx_df, interp2d, interp3d from desc.io import IOAble diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 2a14ebead8..058134e612 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -6,8 +6,7 @@ from desc.backend import jnp from desc.compute import compute as compute_fun -from desc.compute import get_params, get_profiles, get_transforms -from desc.geometry.utils import rpz2xyz +from desc.compute import get_params, get_profiles, get_transforms, rpz2xyz from desc.grid import LinearGrid, QuadratureGrid from desc.utils import Timer @@ -552,9 +551,9 @@ def build(self, eq=None, use_jit=True, verbose=1): print("Precomputing transforms") timer.start("Precomputing transforms") - self._surface_coords = self._surface.compute_coordinates( - grid=surface_grid, basis="xyz" - ) + self._surface_coords = self._surface.compute( + "x", grid=surface_grid, basis="xyz" + )["x"] self._profiles = get_profiles( self._data_keys, obj=eq, diff --git a/desc/plotting.py b/desc/plotting.py index facf5abf5e..342c71275d 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -2075,7 +2075,7 @@ def flatten_coils(coilset): plot_data["Y"] = [] plot_data["Z"] = [] for i, coil in enumerate(coils_list): - x, y, z = coil.compute_coordinates(grid=grid, basis="xyz").T + x, y, z = coil.compute("x", grid=grid, basis="xyz")["x"].T plot_data["X"].append(x) plot_data["Y"].append(y) plot_data["Z"].append(z) diff --git a/tests/test_coils.py b/tests/test_coils.py index f82e918f59..58391f9ad4 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -20,10 +20,9 @@ def test_biot_savart(self): By_true = 1e-7 * 2 * np.pi * R**2 * I / (y**2 + R**2) ** (3 / 2) B_true = np.array([0, By_true, 0]) coil = FourierXYZCoil(I) - coil.grid = LinearGrid(zeta=100, endpoint=True) - assert coil.grid.num_nodes == 100 + grid = LinearGrid(zeta=100, endpoint=True) B_approx = coil.compute_magnetic_field( - Grid([[10, y, 0], [10, -y, 0]]), basis="xyz" + Grid([[10, y, 0], [10, -y, 0]]), basis="xyz", grid=grid )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @@ -55,9 +54,7 @@ def test_linspaced_linear(self): ) coils.current = I np.testing.assert_allclose(coils.current, I) - coils.grid = 32 - assert coils.grid.N == 32 - B_approx = coils.compute_magnetic_field([0, 0, z[-1]], basis="xyz")[0] + B_approx = coils.compute_magnetic_field([0, 0, z[-1]], basis="xyz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @pytest.mark.unit @@ -71,9 +68,7 @@ def test_linspaced_angular(self): coil = FourierPlanarCoil() coil.current = I coils = CoilSet.linspaced_angular(coil, n=N) - coils.grid = 32 - assert all([coil.grid.N == 32 for coil in coils]) - B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz")[0] + B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @pytest.mark.unit @@ -87,9 +82,7 @@ def test_from_symmetry(self): coil = FourierPlanarCoil() coils = CoilSet.linspaced_angular(coil, angle=np.pi / 2, n=N // 4) coils = CoilSet.from_symmetry(coils, NFP=4) - coils.grid = 32 - assert all([coil.grid.N == 32 for coil in coils]) - B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz")[0] + B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) # with stellarator symmetry @@ -99,10 +92,8 @@ def test_from_symmetry(self): coils = CoilSet.linspaced_angular( coil, I, [0, 0, 1], np.pi / NFP, N // NFP // 2 ) - coils.grid = 32 - assert coils.grid.N == 32 coils2 = CoilSet.from_symmetry(coils, NFP, True) - B_approx = coils2.compute_magnetic_field([10, 0, 0], basis="rpz")[0] + B_approx = coils2.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @pytest.mark.unit @@ -110,9 +101,20 @@ def test_properties(self): """Test getting/setting of CoilSet attributes.""" coil = FourierPlanarCoil() coils = CoilSet.linspaced_linear(coil, n=4) - coils.grid = np.array([[0.0, 0.0, 0.0]]) + data = coils.compute( + [ + "x", + "curvature", + "torsion", + "frenet_tangent", + "frenet_normal", + "frenet_binormal", + ], + grid=0, + basis="xyz", + ) np.testing.assert_allclose( - coils.compute_coordinates(), + [dat["x"] for dat in data], np.array( [ [12, 0, 0], @@ -122,12 +124,11 @@ def test_properties(self): ] ).reshape((4, 1, 3)), ) - np.testing.assert_allclose(coils.compute_curvature(), 1 / 2) - np.testing.assert_allclose(coils.compute_torsion(), 0) - TNB = coils.compute_frenet_frame(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz") - T = [foo[0] for foo in TNB] - N = [foo[1] for foo in TNB] - B = [foo[2] for foo in TNB] + np.testing.assert_allclose([dat["curvature"] for dat in data], 1 / 2) + np.testing.assert_allclose([dat["torsion"] for dat in data], 0) + T = [dat["frenet_tangent"] for dat in data] + N = [dat["frenet_normal"] for dat in data] + B = [dat["frenet_binormal"] for dat in data] np.testing.assert_allclose( T, np.array( @@ -164,16 +165,20 @@ def test_properties(self): ).reshape((4, 1, 3)), atol=1e-12, ) - coils.grid = 32 - np.testing.assert_allclose(coils.compute_length(), 2 * 2 * np.pi) + data = coils.compute("length", grid=32) + np.testing.assert_allclose([dat["length"] for dat in data], 2 * 2 * np.pi) coils.translate([1, 1, 1]) - np.testing.assert_allclose(coils.compute_length(), 2 * 2 * np.pi) + data = coils.compute("length", grid=32) + np.testing.assert_allclose([dat["length"] for dat in data], 2 * 2 * np.pi) coils.flip([1, 0, 0]) - coils.grid = np.array([[0.0, 0.0, 0.0]]) - TNB = coils.compute_frenet_frame(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz") - T = [foo[0] for foo in TNB] - N = [foo[1] for foo in TNB] - B = [foo[2] for foo in TNB] + data = coils.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], + grid=0, + basis="xyz", + ) + T = [dat["frenet_tangent"] for dat in data] + N = [dat["frenet_normal"] for dat in data] + B = [dat["frenet_binormal"] for dat in data] np.testing.assert_allclose( T, np.array( diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index ad80ec7832..911dc7c819 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -9,6 +9,13 @@ from desc.compute import data_index from desc.compute.utils import compress from desc.equilibrium import EquilibriaFamily, Equilibrium +from desc.geometry import ( + FourierPlanarCurve, + FourierRZCurve, + FourierRZToroidalSurface, + FourierXYZCurve, + ZernikeRZToroidalSection, +) from desc.grid import LinearGrid, QuadratureGrid # convolve kernel is reverse of FD coeffs @@ -1121,6 +1128,35 @@ def test_compute_everything(): assert key in data +@pytest.mark.unit +def test_curve_compute_everything(): + """Make sure we can compute every curve thing without errors.""" + curves = { + "desc.geometry.curve.FourierXYZCurve": FourierXYZCurve(), + "desc.geometry.curve.FourierRZCurve": FourierRZCurve(), + "desc.geometry.curve.FourierPlanarCurve": FourierPlanarCurve(), + } + + for p, thing in curves.items(): + for key in data_index[p].keys(): + data = thing.compute(key) + assert key in data + + +@pytest.mark.unit +def test_surface_compute_everything(): + """Make sure we can compute every surface thing without errors.""" + surfaces = { + "desc.geometry.surface.FourierRZToroidalSurface": FourierRZToroidalSurface(), + "desc.geometry.surface.ZernikeRZToroidalSection": ZernikeRZToroidalSection(), + } + + for p, thing in surfaces.items(): + for key in data_index[p].keys(): + data = thing.compute(key) + assert key in data + + @pytest.mark.unit def test_compute_averages(): """Test that computing averages uses the correct grid.""" diff --git a/tests/test_configuration.py b/tests/test_configuration.py index 8fc2173ebe..7b034082e8 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -419,9 +419,7 @@ def test_get_rho_surface(self): rho = 0.5 surf = eq.get_surface_at(rho=rho) assert surf.rho == rho - np.testing.assert_allclose( - surf.compute_surface_area(), 4 * np.pi**2 * R0 * rho - ) + np.testing.assert_allclose(surf.compute("S")["S"], 4 * np.pi**2 * R0 * rho) @pytest.mark.unit def test_get_zeta_surface(self): @@ -430,7 +428,7 @@ def test_get_zeta_surface(self): surf = eq.get_surface_at(zeta=np.pi) assert surf.zeta == np.pi rho = 1 - np.testing.assert_allclose(surf.compute_surface_area(), np.pi * rho**2) + np.testing.assert_allclose(surf.compute("A")["A"], np.pi * rho**2) @pytest.mark.unit def test_get_theta_surface(self): @@ -458,7 +456,7 @@ def test_magnetic_axis(HELIOTRON_vac): grid = LinearGrid(N=3 * eq.N_grid, NFP=eq.NFP, rho=np.array(0.0)) data = eq.compute(["R", "Z"], grid=grid) - coords = axis.compute_coordinates(grid=grid) + coords = axis.compute("x", grid=grid)["x"] np.testing.assert_allclose(coords[:, 0], data["R"]) np.testing.assert_allclose(coords[:, 2], data["Z"]) diff --git a/tests/test_curves.py b/tests/test_curves.py index ab8f767990..1e4e0058ec 100644 --- a/tests/test_curves.py +++ b/tests/test_curves.py @@ -14,46 +14,54 @@ class TestRZCurve: def test_length(self): """Test length of circular curve.""" c = FourierRZCurve() - np.testing.assert_allclose(c.compute_length(grid=20), 10 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 10 * 2 * np.pi + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_length(grid=20), 10 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 10 * 2 * np.pi + ) @pytest.mark.unit def test_curvature(self): """Test curvature of circular curve.""" c = FourierRZCurve() - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 10) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 10) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 10) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 10) @pytest.mark.unit def test_torsion(self): """Test torsion of circular curve.""" c = FourierRZCurve() - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose(c.compute("torsion", grid=20)["torsion"], 0) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose(c.compute("torsion", grid=20)["torsion"], 0) @pytest.mark.unit def test_frenet(self): """Test frenet-seret frame of circular curve.""" c = FourierRZCurve() - c.grid = 0 - T, N, B = c.compute_frenet_frame(basis="rpz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="rpz", 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) c.rotate(angle=np.pi) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - c.grid = np.array([[0, 0, 0]]) - T, N, B = c.compute_frenet_frame(basis="xyz") + 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) @@ -62,14 +70,14 @@ def test_frenet(self): def test_coords(self): """Test lab frame coordinates of circular curve.""" c = FourierRZCurve() - x, y, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz").T + x, y, z = c.compute("x", grid=0, basis="xyz")["x"].T np.testing.assert_allclose(x, 10) np.testing.assert_allclose(y, 0) np.testing.assert_allclose(z, 0) c.rotate(angle=np.pi / 2) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - r, p, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="rpz").T + r, p, z = c.compute("x", grid=0, basis="rpz")["x"].T np.testing.assert_allclose(r, np.sqrt(1**2 + 9**2)) np.testing.assert_allclose(p, np.arctan2(-9, 1)) np.testing.assert_allclose(z, 1) @@ -78,9 +86,6 @@ def test_coords(self): def test_misc(self): """Test getting/setting misc attributes of FourierRZCurve.""" c = FourierRZCurve() - grid = LinearGrid(M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) R, Z = c.get_coeffs(0) np.testing.assert_allclose(R, 10) @@ -121,18 +126,12 @@ def test_misc(self): assert c.NFP == 3 assert c.R_basis.NFP == 3 assert c.Z_basis.NFP == 3 - assert c.grid.NFP == 3 @pytest.mark.unit def test_asserts(self): """Test error checking when creating FourierRZCurve.""" with pytest.raises(ValueError): - c = FourierRZCurve(R_n=[]) - c = FourierRZCurve() - with pytest.raises(NotImplementedError): - c.compute_coordinates(dt=4) - with pytest.raises(TypeError): - c.grid = [1, 2, 3] + _ = FourierRZCurve(R_n=[]) @pytest.mark.unit def test_to_FourierXYZCurve(self): @@ -140,18 +139,23 @@ def test_to_FourierXYZCurve(self): rz = FourierRZCurve(R_n=[0, 10, 1], Z_n=[-1, 0, 0]) xyz = rz.to_FourierXYZCurve(N=2) + grid = LinearGrid(N=20, endpoint=True) + np.testing.assert_allclose( - rz.compute_curvature(), xyz.compute_curvature(grid=rz.grid) + rz.compute("curvature", grid=grid)["curvature"], + xyz.compute("curvature", grid=grid)["curvature"], ) np.testing.assert_allclose( - rz.compute_torsion(), xyz.compute_torsion(grid=rz.grid) + rz.compute("torsion", grid=grid)["torsion"], + xyz.compute("torsion", grid=grid)["torsion"], ) np.testing.assert_allclose( - rz.compute_length(), xyz.compute_length(grid=rz.grid) + rz.compute("length", grid=grid)["length"], + xyz.compute("length", grid=grid)["length"], ) np.testing.assert_allclose( - rz.compute_coordinates(basis="rpz"), - xyz.compute_coordinates(basis="rpz", grid=rz.grid), + rz.compute("x", grid=grid, basis="rpz")["x"], + xyz.compute("x", basis="rpz", grid=grid)["x"], atol=1e-12, ) @@ -163,46 +167,58 @@ class TestXYZCurve: def test_length(self): """Test length of circular curve.""" c = FourierXYZCurve() - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) @pytest.mark.unit def test_curvature(self): """Test curvature of circular curve.""" c = FourierXYZCurve() - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) @pytest.mark.unit def test_torsion(self): """Test torsion of circular curve.""" c = FourierXYZCurve(modes=[-1, 0, 1]) - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) @pytest.mark.unit def test_frenet(self): """Test frenet-seret frame of circular curve.""" c = FourierXYZCurve() - c.grid = 0 - T, N, B = c.compute_frenet_frame(basis="rpz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="rpz", 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.rotate(angle=np.pi) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - c.grid = np.array([0, 0, 0]) - T, N, B = c.compute_frenet_frame(basis="xyz") + 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) @@ -211,14 +227,14 @@ def test_frenet(self): def test_coords(self): """Test lab frame coordinates of circular curve.""" c = FourierXYZCurve() - x, y, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz").T + x, y, z = c.compute("x", grid=0, basis="xyz")["x"].T np.testing.assert_allclose(x, 12) np.testing.assert_allclose(y, 0) np.testing.assert_allclose(z, 0) c.rotate(angle=np.pi / 2) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - r, p, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="rpz").T + r, p, z = c.compute("x", grid=0, basis="rpz")["x"].T np.testing.assert_allclose(r, np.sqrt(1**2 + 11**2)) np.testing.assert_allclose(p, np.arctan2(-11, 1)) np.testing.assert_allclose(z, 1) @@ -227,9 +243,6 @@ def test_coords(self): def test_misc(self): """Test getting/setting misc attributes of FourierXYZCurve.""" c = FourierXYZCurve() - grid = LinearGrid(M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) X, Y, Z = c.get_coeffs(0) np.testing.assert_allclose(X, 10) @@ -252,15 +265,6 @@ def test_misc(self): with pytest.raises(ValueError): c.Z_n = s.Z_n - @pytest.mark.unit - def test_asserts(self): - """Test error checking when creating FourierXYZCurve.""" - c = FourierXYZCurve() - with pytest.raises(ValueError): - c.compute_coordinates(dt=4) - with pytest.raises(TypeError): - c.grid = [1, 2, 3] - class TestPlanarCurve: """Tests for FourierPlanarCurve class.""" @@ -269,46 +273,58 @@ class TestPlanarCurve: def test_length(self): """Test length of circular curve.""" c = FourierPlanarCurve(modes=[0]) - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) @pytest.mark.unit def test_curvature(self): """Test curvature of circular curve.""" c = FourierPlanarCurve() - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) @pytest.mark.unit def test_torsion(self): """Test torsion of circular curve.""" c = FourierPlanarCurve() - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) @pytest.mark.unit def test_frenet(self): """Test frenet-seret frame of circular curve.""" c = FourierPlanarCurve() - c.grid = 0 - T, N, B = c.compute_frenet_frame(basis="xyz") + 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.rotate(angle=np.pi) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - c.grid = np.array([0, 0, 0]) - T, N, B = c.compute_frenet_frame(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz") + 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) @@ -317,20 +333,18 @@ def test_frenet(self): def test_coords(self): """Test lab frame coordinates of circular curve.""" c = FourierPlanarCurve() - r, p, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="rpz").T + r, p, z = c.compute("x", grid=0, basis="rpz")["x"].T np.testing.assert_allclose(r, 12) np.testing.assert_allclose(p, 0) np.testing.assert_allclose(z, 0) - dr, dp, dz = c.compute_coordinates( - grid=np.array([[0.0, 0.0, 0.0]]), dt=3, basis="rpz" - ).T + dr, dp, dz = c.compute("x_sss", grid=0, basis="rpz")["x_sss"].T np.testing.assert_allclose(dr, 0) np.testing.assert_allclose(dp, 0) np.testing.assert_allclose(dz, 2) c.rotate(angle=np.pi / 2) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - x, y, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz").T + x, y, z = c.compute("x", grid=0, basis="xyz")["x"].T np.testing.assert_allclose(x, 1) np.testing.assert_allclose(y, -11) np.testing.assert_allclose(z, 1) @@ -339,9 +353,6 @@ def test_coords(self): def test_misc(self): """Test getting/setting misc attributes of FourierPlanarCurve.""" c = FourierPlanarCurve() - grid = LinearGrid(M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) r = c.get_coeffs(0) np.testing.assert_allclose(r, 2) @@ -369,10 +380,6 @@ def test_misc(self): def test_asserts(self): """Test error checking when creating FourierPlanarCurve.""" c = FourierPlanarCurve() - with pytest.raises(NotImplementedError): - c.compute_coordinates(dt=4) - with pytest.raises(TypeError): - c.grid = [1, 2, 3] with pytest.raises(ValueError): c.center = [4] with pytest.raises(ValueError): diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 18622ad3c4..79c3f88040 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from desc.geometry.utils import ( +from desc.compute.geom_utils import ( rotation_matrix, rpz2xyz, rpz2xyz_vec, diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py index 5fe5a44285..010e4d72d3 100644 --- a/tests/test_surfaces.py +++ b/tests/test_surfaces.py @@ -17,33 +17,24 @@ def test_area(self): """Test calculation of surface area.""" s = FourierRZToroidalSurface() grid = LinearGrid(M=24, N=24) - s.grid = grid - area = 4 * np.pi**2 * 10 - np.testing.assert_allclose(s.compute_surface_area(), area) - np.testing.assert_allclose(s.compute_surface_area(grid=10), area) - np.testing.assert_allclose(s.compute_surface_area(grid=(10, 15)), area) + np.testing.assert_allclose(s.compute("S", grid=grid)["S"], area) @pytest.mark.unit def test_normal(self): """Test calculation of surface normal vector.""" s = FourierRZToroidalSurface() grid = LinearGrid(theta=np.pi / 2, zeta=np.pi) - s.grid = grid - N = s.compute_normal() + N = s.compute("n_rho", grid=grid)["n_rho"] np.testing.assert_allclose(N[0], [0, 0, -1], atol=1e-14) grid = LinearGrid(theta=0.0, zeta=0.0) - s.grid = grid - N = s.compute_normal(basis="xyz") + N = s.compute("n_rho", grid=grid)["n_rho"] np.testing.assert_allclose(N[0], [1, 0, 0], atol=1e-12) @pytest.mark.unit def test_misc(self): """Test getting/setting attributes of surface.""" c = FourierRZToroidalSurface() - grid = LinearGrid(L=0, M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) R, Z = c.get_coeffs(0, 0) np.testing.assert_allclose(R, 10) @@ -82,7 +73,6 @@ def test_misc(self): assert c.NFP == 3 assert c.R_basis.NFP == 3 assert c.Z_basis.NFP == 3 - assert c.grid.NFP == 3 @pytest.mark.unit def test_from_input_file(self): @@ -133,12 +123,19 @@ def test_curvature(self): """Tests for gaussian, mean, principle curvatures.""" s = FourierRZToroidalSurface() grid = LinearGrid(theta=np.pi / 2, zeta=np.pi) - s.grid = grid - K, H, k1, k2 = s.compute_curvature() - np.testing.assert_allclose(K, 0) - np.testing.assert_allclose(H, -1 / 2) - np.testing.assert_allclose(k1, 0) - np.testing.assert_allclose(k2, -1) + data = s.compute( + [ + "curvature_K_rho", + "curvature_H_rho", + "curvature_k1_rho", + "curvature_k2_rho", + ], + grid=grid, + ) + np.testing.assert_allclose(data["curvature_K_rho"], 0) + np.testing.assert_allclose(data["curvature_H_rho"], -1 / 2) + np.testing.assert_allclose(data["curvature_k1_rho"], 0) + np.testing.assert_allclose(data["curvature_k2_rho"], -1) class TestZernikeRZToroidalSection: @@ -149,29 +146,21 @@ def test_area(self): """Test calculation of surface area.""" s = ZernikeRZToroidalSection() grid = LinearGrid(L=10, M=10) - s.grid = grid - area = np.pi * 1**2 - np.testing.assert_allclose(s.compute_surface_area(), area) - np.testing.assert_allclose(s.compute_surface_area(grid=15), area) - np.testing.assert_allclose(s.compute_surface_area(grid=(5, 5)), area) + np.testing.assert_allclose(s.compute("A", grid=grid)["A"], area) @pytest.mark.unit def test_normal(self): """Test calculation of surface normal vector.""" s = ZernikeRZToroidalSection() - grid = LinearGrid(L=8, M=4, N=0) - s.grid = grid - N = s.compute_normal(basis="xyz") + grid = LinearGrid(L=8, M=4, N=0, axis=False) + N = s.compute("n_zeta", grid=grid)["n_zeta"] np.testing.assert_allclose(N, np.broadcast_to([0, 1, 0], N.shape), atol=1e-12) @pytest.mark.unit def test_misc(self): """Test getting/setting surface attributes.""" c = ZernikeRZToroidalSection() - grid = LinearGrid(L=2, M=2, N=0) - c.grid = grid - assert grid.eq(c.grid) R, Z = c.get_coeffs(0, 0) np.testing.assert_allclose(R, 10) @@ -212,12 +201,19 @@ def test_curvature(self): """ s = ZernikeRZToroidalSection() grid = LinearGrid(theta=np.pi / 2, rho=0.5) - s.grid = grid - K, H, k1, k2 = s.compute_curvature() - np.testing.assert_allclose(K, 0) - np.testing.assert_allclose(H, 0) - np.testing.assert_allclose(k1, 0) - np.testing.assert_allclose(k2, 0) + data = s.compute( + [ + "curvature_K_zeta", + "curvature_H_zeta", + "curvature_k1_zeta", + "curvature_k2_zeta", + ], + grid=grid, + ) + np.testing.assert_allclose(data["curvature_K_zeta"], 0) + np.testing.assert_allclose(data["curvature_H_zeta"], 0) + np.testing.assert_allclose(data["curvature_k1_zeta"], 0) + np.testing.assert_allclose(data["curvature_k2_zeta"], 0) @pytest.mark.unit