diff --git a/desc/coils.py b/desc/coils.py index 8009246ee7..ab7f0dc2b3 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -458,17 +458,19 @@ class FourierPlanarCoil(_Coil, FourierPlanarCurve): Parameters ---------- current : float - current through the coil, in Amperes + Current through the coil, in Amperes. center : array-like, shape(3,) - x,y,z coordinates of center of coil + Coordinates of center of curve, in system determined by basis. normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface + Components of normal vector to planar surface, in system determined by basis. r_n : array-like - fourier coefficients for radius from center as function of polar angle + Fourier coefficients for radius from center as function of polar angle modes : array-like mode numbers associated with r_n + basis : {'xyz', 'rpz'} + Coordinate system for center and normal vectors. Default = 'xyz'. name : str - name for this coil + Name for this coil. Examples -------- @@ -514,9 +516,10 @@ def __init__( normal=[0, 1, 0], r_n=2, modes=None, + basis="xyz", name="", ): - super().__init__(current, center, normal, r_n, modes, name) + super().__init__(current, center, normal, r_n, modes, basis, name) class SplineXYZCoil(_Coil, SplineXYZCurve): diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 85d03ca31f..40d77174b3 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -168,16 +168,21 @@ def _Z_Curve(params, transforms, profiles, data, **kwargs): description="Position vector along curve", dim=3, params=["r_n", "center", "normal", "rotmat", "shift"], - transforms={ - "r": [[0, 0, 0]], - }, + transforms={"r": [[0, 0, 0]]}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.FourierPlanarCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'", ) def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + if kwargs.get("basis_in", "xyz").lower() == "rpz": + center = rpz2xyz(params["center"]) + normal = rpz2xyz_vec(params["normal"], phi=params["center"][1]) + else: + center = params["center"] + normal = params["normal"] # create planar curve at Z==0 r = transforms["r"].transform(params["r_n"], dz=0) Z = jnp.zeros_like(r) @@ -186,10 +191,10 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([X, Y, Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, params["normal"]) - angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"]))) + axis = cross(Zaxis, normal) + angle = jnp.arccos(dot(Zaxis, safenormalize(normal))) A = rotation_matrix(axis=axis, angle=angle) - coords = jnp.matmul(coords, A.T) + params["center"] + coords = jnp.matmul(coords, A.T) + center coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) + params["shift"] if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz(coords) @@ -205,16 +210,21 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, first derivative", dim=3, params=["r_n", "center", "normal", "rotmat", "shift"], - transforms={ - "r": [[0, 0, 0], [0, 0, 1]], - }, + transforms={"r": [[0, 0, 0], [0, 0, 1]]}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.FourierPlanarCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'", ) def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + if kwargs.get("basis_in", "xyz").lower() == "rpz": + center = rpz2xyz(params["center"]) + normal = rpz2xyz_vec(params["normal"], phi=params["center"][1]) + else: + center = params["center"] + normal = params["normal"] 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"]) @@ -223,8 +233,8 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([dX, dY, dZ]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, params["normal"]) - angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"]))) + axis = cross(Zaxis, normal) + angle = jnp.arccos(dot(Zaxis, safenormalize(normal))) A = rotation_matrix(axis=axis, angle=angle) coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) @@ -233,7 +243,7 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): 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, A.T) + center xyzcoords = ( jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"] ) @@ -251,16 +261,21 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, second derivative", dim=3, params=["r_n", "center", "normal", "rotmat", "shift"], - transforms={ - "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], - }, + transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]]}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.FourierPlanarCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'", ) def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + if kwargs.get("basis_in", "xyz").lower() == "rpz": + center = rpz2xyz(params["center"]) + normal = rpz2xyz_vec(params["normal"], phi=params["center"][1]) + else: + center = params["center"] + normal = params["normal"] 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) @@ -274,8 +289,8 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([d2X, d2Y, d2Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, params["normal"]) - angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"]))) + axis = cross(Zaxis, normal) + angle = jnp.arccos(dot(Zaxis, safenormalize(normal))) A = rotation_matrix(axis=axis, angle=angle) coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) @@ -284,7 +299,7 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): 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, A.T) + center xyzcoords = ( jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"] ) @@ -302,16 +317,21 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, third derivative", dim=3, params=["r_n", "center", "normal", "rotmat", "shift"], - transforms={ - "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]], - }, + transforms={"r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.FourierPlanarCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + basis_in="{'rpz', 'xyz'}: Basis for input params vectors, Default 'xyz'", ) def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): + if kwargs.get("basis_in", "xyz").lower() == "rpz": + center = rpz2xyz(params["center"]) + normal = rpz2xyz_vec(params["normal"], phi=params["center"][1]) + else: + center = params["center"] + normal = params["normal"] 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) @@ -332,8 +352,8 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.array([d3X, d3Y, d3Z]).T # rotate into place Zaxis = jnp.array([0.0, 0.0, 1.0]) # 2D curve in X-Y plane has normal = +Z axis - axis = cross(Zaxis, params["normal"]) - angle = jnp.arccos(dot(Zaxis, safenormalize(params["normal"]))) + axis = cross(Zaxis, normal) + angle = jnp.arccos(dot(Zaxis, safenormalize(normal))) A = rotation_matrix(axis=axis, angle=angle) coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, params["rotmat"].reshape((3, 3)).T) @@ -342,7 +362,7 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): 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, A.T) + center xyzcoords = ( jnp.matmul(xyzcoords, params["rotmat"].reshape((3, 3)).T) + params["shift"] ) @@ -360,11 +380,7 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve", dim=3, params=["R_n", "Z_n", "rotmat", "shift"], - transforms={ - "R": [[0, 0, 0]], - "Z": [[0, 0, 0]], - "grid": [], - }, + transforms={"R": [[0, 0, 0]], "Z": [[0, 0, 0]], "grid": []}, profiles=[], coordinates="s", data=[], @@ -395,11 +411,7 @@ def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, first derivative", dim=3, params=["R_n", "Z_n", "rotmat"], - transforms={ - "R": [[0, 0, 0], [0, 0, 1]], - "Z": [[0, 0, 1]], - "grid": [], - }, + transforms={"R": [[0, 0, 0], [0, 0, 1]], "Z": [[0, 0, 1]], "grid": []}, profiles=[], coordinates="s", data=[], @@ -429,11 +441,7 @@ def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, second derivative", dim=3, params=["R_n", "Z_n", "rotmat"], - transforms={ - "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], - "Z": [[0, 0, 2]], - "grid": [], - }, + transforms={"R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], "Z": [[0, 0, 2]], "grid": []}, profiles=[], coordinates="s", data=[], @@ -505,11 +513,7 @@ def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve", dim=3, params=["X_n", "Y_n", "Z_n", "rotmat", "shift"], - transforms={ - "X": [[0, 0, 0]], - "Y": [[0, 0, 0]], - "Z": [[0, 0, 0]], - }, + transforms={"X": [[0, 0, 0]], "Y": [[0, 0, 0]], "Z": [[0, 0, 0]]}, profiles=[], coordinates="s", data=[], @@ -643,9 +647,7 @@ def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve", dim=3, params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={ - "method": [], - }, + transforms={"method": []}, profiles=[], coordinates="s", data=["s"], @@ -698,9 +700,7 @@ def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, first derivative", dim=3, params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={ - "method": [], - }, + transforms={"method": []}, profiles=[], coordinates="s", data=["s"], @@ -784,9 +784,7 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, second derivative", dim=3, params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={ - "method": [], - }, + transforms={"method": []}, profiles=[], coordinates="s", data=["s"], @@ -869,9 +867,7 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, third derivative", dim=3, params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={ - "method": [], - }, + transforms={"method": []}, profiles=[], coordinates="s", data=["s"], diff --git a/desc/geometry/core.py b/desc/geometry/core.py index dfc3b53ba1..a637ff529e 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -178,17 +178,17 @@ def compute( return data def translate(self, displacement=[0, 0, 0]): - """Translate the curve by a rigid displacement in X, Y, Z.""" + """Translate the curve by a rigid displacement in X,Y,Z coordinates.""" self.shift = self.shift + jnp.asarray(displacement) def rotate(self, axis=[0, 0, 1], angle=0): - """Rotate the curve by a fixed angle about axis in X, Y, Z coordinates.""" + """Rotate the curve by a fixed angle about axis in X,Y,Z coordinates.""" R = rotation_matrix(axis=axis, angle=angle) self.rotmat = (R @ self.rotmat.reshape(3, 3)).flatten() self.shift = self.shift @ R.T def flip(self, normal=[0, 0, 1]): - """Flip the curve about the plane with specified normal.""" + """Flip the curve about the plane with specified normal in X,Y,Z coordinates.""" F = reflection_matrix(normal) self.rotmat = (F @ self.rotmat.reshape(3, 3)).flatten() self.shift = self.shift @ F.T diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index 0c4edbb90d..3f670656e7 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -556,24 +556,21 @@ class FourierPlanarCurve(Curve): Parameters ---------- center : array-like, shape(3,) - x,y,z coordinates of center of curve + Coordinates of center of curve, in system determined by basis. normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface + Components of normal vector to planar surface, in system determined by basis. r_n : array-like Fourier coefficients for radius from center as function of polar angle modes : array-like mode numbers associated with r_n + basis : {'xyz', 'rpz'} + Coordinate system for center and normal vectors. Default = 'xyz'. name : str - name for this curve + Name for this curve. """ - _io_attrs_ = Curve._io_attrs_ + [ - "_r_n", - "_center", - "_normal", - "_r_basis", - ] + _io_attrs_ = Curve._io_attrs_ + ["_r_n", "_center", "_normal", "_r_basis", "_basis"] # Reference frame is centered at the origin with normal in the +Z direction. # The curve is computed in this frame and then shifted/rotated to the correct frame. @@ -583,6 +580,7 @@ def __init__( normal=[0, 1, 0], r_n=2, modes=None, + basis="xyz", name="", ): super().__init__(name) @@ -593,6 +591,7 @@ def __init__( modes = np.asarray(modes) assert issubclass(modes.dtype.type, np.integer) assert r_n.size == modes.size, "r_n size and modes must be the same size" + assert basis.lower() in ["xyz", "rpz"] N = np.max(abs(modes)) self._r_basis = FourierSeries(N, NFP=1, sym=False) @@ -600,6 +599,7 @@ def __init__( self.normal = normal self.center = center + self._basis = basis @property def r_basis(self): @@ -632,7 +632,9 @@ def center(self, new): self._center = np.asarray(new) else: raise ValueError( - "center should be a 3 element vector [cx, cy, cz], got {}".format(new) + "center should be a 3 element vector in " + + self._basis + + " coordinates, got {}".format(new) ) @optimizable_parameter @@ -647,7 +649,9 @@ def normal(self, new): self._normal = np.asarray(new) / np.linalg.norm(new) else: raise ValueError( - "normal should be a 3 element vector [nx, ny, nz], got {}".format(new) + "normal should be a 3 element vector in " + + self._basis + + " coordinates, got {}".format(new) ) @optimizable_parameter @@ -685,6 +689,54 @@ def set_coeffs(self, n, r=None): if rr is not None: self.r_n = put(self.r_n, idx, rr) + def compute( + self, + names, + grid=None, + params=None, + transforms=None, + data=None, + override_grid=True, + **kwargs, + ): + """Compute the quantity given by name on 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 + override_grid : bool + If True, override the user supplied grid if necessary and use a full + resolution grid to compute quantities and then downsample to user requested + grid. If False, uses only the user specified grid, which may lead to + inaccurate values for surface or volume averages. + + Returns + ------- + data : dict of ndarray + Computed quantity and intermediate variables. + + """ + return super().compute( + names=names, + grid=grid, + params=params, + transforms=transforms, + data=data, + override_grid=override_grid, + basis_in=self._basis, + **kwargs, + ) + class SplineXYZCurve(Curve): """Curve parameterized by spline knots in X,Y,Z. diff --git a/tests/test_curves.py b/tests/test_curves.py index ce1a04e3a9..66c4a12edf 100644 --- a/tests/test_curves.py +++ b/tests/test_curves.py @@ -554,6 +554,30 @@ def test_coords(self): np.testing.assert_allclose(y, -11) np.testing.assert_allclose(z, 1) + @pytest.mark.unit + def test_basis(self): + """Test xyz vs rpz basis.""" + cxyz = FourierPlanarCurve(center=[1, 1, 0], normal=[-1, 1, 0], basis="xyz") + crpz = FourierPlanarCurve( + center=[np.sqrt(2), np.pi / 4, 0], normal=[0, 1, 0], basis="rpz" + ) + + x_xyz = cxyz.compute("x")["x"] + x_rpz = crpz.compute("x")["x"] + np.testing.assert_allclose(x_xyz, x_rpz) + + xs_xyz = cxyz.compute("x_s")["x_s"] + xs_rpz = crpz.compute("x_s")["x_s"] + np.testing.assert_allclose(xs_xyz, xs_rpz, atol=2e-15) + + xss_xyz = cxyz.compute("x_ss")["x_ss"] + xss_rpz = crpz.compute("x_ss")["x_ss"] + np.testing.assert_allclose(xss_xyz, xss_rpz, atol=2e-15) + + xsss_xyz = cxyz.compute("x_sss")["x_sss"] + xsss_rpz = crpz.compute("x_sss")["x_sss"] + np.testing.assert_allclose(xsss_xyz, xsss_rpz, atol=2e-15) + @pytest.mark.unit def test_misc(self): """Test getting/setting misc attributes of FourierPlanarCurve.""" diff --git a/tests/test_examples.py b/tests/test_examples.py index adad7d7104..b84cd967b1 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -10,7 +10,7 @@ from qsc import Qsc from desc.backend import jnp -from desc.coils import FourierRZCoil +from desc.coils import FourierPlanarCoil, FourierRZCoil, MixedCoilSet from desc.continuation import solve_continuation_automatic from desc.equilibrium import EquilibriaFamily, Equilibrium from desc.examples import get @@ -1294,3 +1294,22 @@ def test_second_stage_optimization(): np.testing.assert_allclose(field[0].R0, 3.5) # this value was fixed np.testing.assert_allclose(field[0].B0, 1) # toroidal field (no change) np.testing.assert_allclose(field[1].B0, 0, atol=1e-12) # vertical field (vanishes) + + +# TODO: replace this with the solution to Issue #1021 +@pytest.mark.unit +def test_optimize_with_fourier_planar_coil(): + """Test optimizing a FourierPlanarCoil.""" + # single coil + c = FourierPlanarCoil() + objective = ObjectiveFunction(CoilLength(c, target=11)) + optimizer = Optimizer("fmintr") + (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0) + np.testing.assert_allclose(c.compute("length")["length"], 11, atol=1e-3) + + # in MixedCoilSet + c = MixedCoilSet(FourierRZCoil(), FourierPlanarCoil()) + objective = ObjectiveFunction(CoilLength(c, target=11)) + optimizer = Optimizer("fmintr") + (c,), _ = optimizer.optimize(c, objective=objective, maxiter=200, ftol=0, xtol=0) + np.testing.assert_allclose(c.compute("length")[1]["length"], 11, atol=1e-3)