diff --git a/desc/coils.py b/desc/coils.py index 1629941789..1b5f455719 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -226,6 +226,12 @@ def compute_magnetic_field( current = self.current else: current = params.pop("current", self.current) + if source_grid is None and hasattr(self, "NFP"): + # NFP=1 to ensure we have points along whole grid + # multiply by NFP in case the coil has NFP>1 + # to ensure whole coil gets counted for the + # biot savart integration + source_grid = LinearGrid(N=2 * self.N * self.NFP + 5, NFP=1, endpoint=False) if not params or not transforms: data = self.compute( diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index b68b8f2d99..af195404df 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -646,38 +646,39 @@ def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): units_long="meters", description="Position vector along curve", dim=3, - params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={"method": []}, + params=["X", "Y", "Z", "rotmat", "shift"], + transforms={"knots": []}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.SplineXYZCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.", ) def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs): xq = data["s"] Xq = interp1d( xq, - params["knots"], + transforms["knots"], params["X"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Yq = interp1d( xq, - params["knots"], + transforms["knots"], params["Y"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Zq = interp1d( xq, - params["knots"], + transforms["knots"], params["Z"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) @@ -699,38 +700,39 @@ def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs): units_long="meters", description="Position vector along curve, first derivative", dim=3, - params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={"method": []}, + params=["X", "Y", "Z", "rotmat", "shift"], + transforms={"knots": []}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.SplineXYZCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.", ) def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs): xq = data["s"] dXq = interp1d( xq, - params["knots"], + transforms["knots"], params["X"], - method=transforms["method"], + method=kwargs["method"], derivative=1, period=2 * jnp.pi, ) dYq = interp1d( xq, - params["knots"], + transforms["knots"], params["Y"], - method=transforms["method"], + method=kwargs["method"], derivative=1, period=2 * jnp.pi, ) dZq = interp1d( xq, - params["knots"], + transforms["knots"], params["Z"], - method=transforms["method"], + method=kwargs["method"], derivative=1, period=2 * jnp.pi, ) @@ -742,25 +744,25 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs): # calculate the xy coordinates to rotate to rpz Xq = interp1d( xq, - params["knots"], + transforms["knots"], params["X"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Yq = interp1d( xq, - params["knots"], + transforms["knots"], params["Y"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Zq = interp1d( xq, - params["knots"], + transforms["knots"], params["Z"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) @@ -783,38 +785,39 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs): units_long="meters", description="Position vector along curve, second derivative", dim=3, - params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={"method": []}, + params=["X", "Y", "Z", "rotmat", "shift"], + transforms={"knots": []}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.SplineXYZCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.", ) def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs): xq = data["s"] d2Xq = interp1d( xq, - params["knots"], + transforms["knots"], params["X"], - method=transforms["method"], + method=kwargs["method"], derivative=2, period=2 * jnp.pi, ) d2Yq = interp1d( xq, - params["knots"], + transforms["knots"], params["Y"], - method=transforms["method"], + method=kwargs["method"], derivative=2, period=2 * jnp.pi, ) d2Zq = interp1d( xq, - params["knots"], + transforms["knots"], params["Z"], - method=transforms["method"], + method=kwargs["method"], derivative=2, period=2 * jnp.pi, ) @@ -826,25 +829,25 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs): # calculate the xy coordinates to rotate to rpz Xq = interp1d( xq, - params["knots"], + transforms["knots"], params["X"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Yq = interp1d( xq, - params["knots"], + transforms["knots"], params["Y"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Zq = interp1d( xq, - params["knots"], + transforms["knots"], params["Z"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) @@ -866,38 +869,39 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs): units_long="meters", description="Position vector along curve, third derivative", dim=3, - params=["X", "Y", "Z", "knots", "rotmat", "shift"], - transforms={"method": []}, + params=["X", "Y", "Z", "rotmat", "shift"], + transforms={"knots": []}, profiles=[], coordinates="s", data=["s"], parameterization="desc.geometry.curve.SplineXYZCurve", basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'", + method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.", ) def _x_sss_SplineXYZCurve(params, transforms, profiles, data, **kwargs): xq = data["s"] d3Xq = interp1d( xq, - params["knots"], + transforms["knots"], params["X"], - method=transforms["method"], + method=kwargs["method"], derivative=3, period=2 * jnp.pi, ) d3Yq = interp1d( xq, - params["knots"], + transforms["knots"], params["Y"], - method=transforms["method"], + method=kwargs["method"], derivative=3, period=2 * jnp.pi, ) d3Zq = interp1d( xq, - params["knots"], + transforms["knots"], params["Z"], - method=transforms["method"], + method=kwargs["method"], derivative=3, period=2 * jnp.pi, ) @@ -909,25 +913,25 @@ def _x_sss_SplineXYZCurve(params, transforms, profiles, data, **kwargs): # calculate the xy coordinates to rotate to rpz Xq = interp1d( xq, - params["knots"], + transforms["knots"], params["X"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Yq = interp1d( xq, - params["knots"], + transforms["knots"], params["Y"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) Zq = interp1d( xq, - params["knots"], + transforms["knots"], params["Z"], - method=transforms["method"], + method=kwargs["method"], derivative=0, period=2 * jnp.pi, ) @@ -1080,14 +1084,15 @@ def _length(params, transforms, profiles, data, **kwargs): description="Length of the curve", dim=0, params=[], - transforms={"method": []}, + transforms={}, profiles=[], coordinates="", data=["ds", "x", "x_s"], parameterization="desc.geometry.curve.SplineXYZCurve", + method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.", ) def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs): - if transforms["method"] == "nearest": # cannot use derivative method as deriv=0 + if kwargs["method"] == "nearest": # cannot use derivative method as deriv=0 coords = data["x"] if kwargs.get("basis", "rpz").lower() == "rpz": coords = rpz2xyz(coords) diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index b928818cdc..f5281e1efd 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -822,7 +822,7 @@ def __init__( knots = knots[:-1] if closed_flag else knots self._knots = knots - self.method = method + self._method = method @optimizable_parameter @property @@ -925,6 +925,46 @@ def method(self, new): + f"instead got unknown method {new} " ) + def compute( + self, + names, + grid=None, + params=None, + transforms=None, + data=None, + **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 + + Returns + ------- + data : dict of ndarray + Computed quantity and intermediate variables. + """ + return super().compute( + names=names, + grid=grid, + params=params, + transforms=transforms, + data=data, + method=self._method, + **kwargs, + ) + @classmethod def from_values(cls, coords, knots=None, method="cubic", name="", basis="xyz"): """Create SplineXYZCurve from coordinate values. diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 9892e22b78..3be3415535 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -128,7 +128,10 @@ def _prune_coilset_tree(coilset): # map grid to list of length coils if grid is None: - grid = [LinearGrid(N=2 * c.N + 5, endpoint=False) for c in coils] + grid = [] + for c in coils: + NFP = c.NFP if hasattr(c, "NFP") else 1 + grid.append(LinearGrid(N=2 * c.N + 5, NFP=NFP, endpoint=False)) if isinstance(grid, numbers.Integral): grid = LinearGrid(N=self._grid, endpoint=False) if isinstance(grid, _Grid): diff --git a/tests/conftest.py b/tests/conftest.py index 3d57ea5ab4..197fcfef39 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,6 +15,7 @@ FourierRZCoil, FourierXYZCoil, MixedCoilSet, + SplineXYZCoil, ) from desc.compute import rpz2xyz_vec from desc.equilibrium import EquilibriaFamily, Equilibrium @@ -275,13 +276,23 @@ def DummyMixedCoilSet(tmpdir_factory): output_path = output_dir.join("DummyMixedCoilSet.h5") tf_coil = FourierPlanarCoil(current=3, center=[2, 0, 0], normal=[0, 1, 0], r_n=[1]) - tf_coilset = CoilSet.linspaced_angular(tf_coil, n=4) + tf_coil.rotate(angle=np.pi / 4) + tf_coilset = CoilSet(tf_coil, NFP=2, sym=True) + vf_coil = FourierRZCoil(current=-1, R_n=3, Z_n=-1) vf_coilset = CoilSet.linspaced_linear( vf_coil, displacement=[0, 0, 2], n=3, endpoint=True ) xyz_coil = FourierXYZCoil(current=2) - full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xyz_coil)) + phi = 2 * np.pi * np.linspace(0, 1, 20, endpoint=True) ** 2 + spline_coil = SplineXYZCoil( + current=1, + X=np.cos(phi), + Y=np.sin(phi), + Z=np.zeros_like(phi), + knots=np.linspace(0, 2 * np.pi, len(phi)), + ) + full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xyz_coil, spline_coil)) full_coilset.save(output_path) DummyMixedCoilSet_out = {"output_path": output_path} diff --git a/tests/test_coils.py b/tests/test_coils.py index 4c5d38e4ff..7faf6de27c 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -130,6 +130,20 @@ def test_biot_savart_all_coils(self): B_true_rpz_phi, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" ) + # FourierRZCoil with NFP>1 + coil = FourierRZCoil(I, R_n=np.array([R]), modes_R=np.array([0]), NFP=2) + B_xyz = coil.compute_magnetic_field(grid_xyz, basis="xyz", source_grid=None) + B_rpz = coil.compute_magnetic_field(grid_rpz, basis="rpz", source_grid=None) + np.testing.assert_allclose( + B_true_xyz, B_xyz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + ) + np.testing.assert_allclose( + B_true_rpz_xy, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + ) + np.testing.assert_allclose( + B_true_rpz_phi, B_rpz, rtol=1e-3, atol=1e-10, err_msg="Using FourierRZCoil" + ) + @pytest.mark.unit def test_properties(self): """Test getting/setting attributes for Coil class.""" diff --git a/tests/test_examples.py b/tests/test_examples.py index 053e3c2bac..70ff2023c6 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -9,13 +9,14 @@ from qic import Qic from qsc import Qsc -from desc.backend import jnp +from desc.backend import jnp, tree_leaves from desc.coils import ( CoilSet, FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, MixedCoilSet, + _Coil, ) from desc.continuation import solve_continuation_automatic from desc.equilibrium import EquilibriaFamily, Equilibrium @@ -1185,44 +1186,6 @@ def test_example_get_current(self): ) -@pytest.mark.unit -def test_single_coil_optimization(): - """Test that single coil (not coilset) optimization works.""" - # testing that the objectives work and that the optimization framework - # works when a single coil is passed in. - opt = Optimizer("fmintr") - coil = FourierRZCoil() - coil.change_resolution(N=1) - target_R = 9 - target_length = 2 * np.pi * target_R - target_curvature = 1 / target_R - target_torsion = 0 - grid = LinearGrid(N=2) - - # length and curvature - obj = ObjectiveFunction( - ( - CoilLength(coil, target=target_length), - CoilCurvature(coil, target=target_curvature, grid=grid), - ), - ) - opt.optimize([coil], obj) - np.testing.assert_allclose( - coil.compute("length")["length"], target_length, rtol=3e-3 - ) - np.testing.assert_allclose( - coil.compute("curvature", grid=grid)["curvature"], target_curvature, rtol=3e-3 - ) - - # torsion - coil.Z_n = coil.Z_n.at[0].set(0.1) # initialize with some torsion - obj = ObjectiveFunction(CoilTorsion(coil, target=target_torsion, grid=grid)) - opt.optimize([coil], obj) - np.testing.assert_allclose( - coil.compute("torsion", grid=grid)["torsion"], target_torsion, atol=1e-5 - ) - - @pytest.mark.unit def test_quadratic_flux_optimization_with_analytic_field(): """Test analytic field optimization to reduce quadratic flux. @@ -1336,30 +1299,84 @@ def test_second_stage_optimization_CoilSet(): np.testing.assert_allclose(field[0].current, 0, atol=1e-12) -# TODO: replace this with the solution to Issue #1021 +@pytest.mark.slow @pytest.mark.unit -def test_optimize_with_fourier_planar_coil(): - """Test optimizing a FourierPlanarCoil.""" +def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): + """Test optimizing for every type of coil and dummy coil sets.""" + sym_coils = load(load_from=str(DummyCoilSet["output_path_sym"]), file_format="hdf5") + asym_coils = load( + load_from=str(DummyCoilSet["output_path_asym"]), file_format="hdf5" + ) + mixed_coils = load( + load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5" + ) + nested_coils = MixedCoilSet(sym_coils, mixed_coils) + eq = Equilibrium() + # not attempting to accurately calc B for this test, + # so make the grids very coarse + quad_eval_grid = LinearGrid(M=2, sym=True) + quad_field_grid = LinearGrid(N=2) + + def test(c, method): + target = 11 + rtol = 1e-3 + # first just check that quad flux works for a couple iterations + # as this is an expensive objective to compute + obj = ObjectiveFunction( + QuadraticFlux( + eq=eq, + field=c, + vacuum=True, + weight=1e-4, + eval_grid=quad_eval_grid, + field_grid=quad_field_grid, + ) + ) + optimizer = Optimizer(method) + (c,), _ = optimizer.optimize(c, obj, maxiter=2, ftol=0, xtol=1e-15) + + # now check with optimizing geometry and actually check result + objs = [ + CoilLength(c, target=target), + ] + extra_msg = "" + if isinstance(c, MixedCoilSet): + # just to check they work without error + objs.extend( + [ + CoilCurvature(c, target=0.5, weight=1e-2), + CoilTorsion(c, target=0, weight=1e-2), + ] + ) + rtol = 3e-2 + extra_msg = " with curvature and torsion obj" + + obj = ObjectiveFunction(objs) + + (c,), _ = optimizer.optimize(c, obj, maxiter=25, ftol=5e-3, xtol=1e-15) + flattened_coils = tree_leaves( + c, is_leaf=lambda x: isinstance(x, _Coil) and not isinstance(x, CoilSet) + ) + lengths = [coil.compute("length")["length"] for coil in flattened_coils] + np.testing.assert_allclose( + lengths, target, rtol=rtol, err_msg=f"lengths {c}" + extra_msg + ) + + spline_coil = mixed_coils.coils[-1].copy() + # 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) - - # in CoilSet - c = CoilSet(FourierPlanarCoil(), sym=True, NFP=2) - 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")[0]["length"], 11, atol=1e-3) + test(FourierPlanarCoil(), "fmintr") + test(FourierRZCoil(), "fmintr") + test(FourierXYZCoil(), "fmintr") + test(spline_coil, "fmintr") + + # CoilSet + test(sym_coils, "lsq-exact") + test(asym_coils, "lsq-exact") + + # MixedCoilSet + test(mixed_coils, "lsq-exact") + test(nested_coils, "lsq-exact") @pytest.mark.unit diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index bed8c6482a..c6f364718a 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -959,16 +959,14 @@ def test_fix_subset_of_params_in_collection(DummyMixedCoilSet): params = [ [ {"current": True}, - {"center": True, "normal": np.array([1])}, - {"r_n": True}, - {}, ], {"shift": True, "rotmat": True}, {"X_n": np.array([1, 2]), "Y_n": False, "Z_n": np.array([0])}, + {}, ] target = np.concatenate( ( - np.array([3, 2, 0, 0, 1, 1]), + np.array([3]), np.eye(3).flatten(), np.array([0, 0, 0]), np.eye(3).flatten(), @@ -998,13 +996,13 @@ def test_fix_coil_current(DummyMixedCoilSet): # fix all coil currents obj = FixCoilCurrent(coil=coilset) obj.build() - assert obj.dim_f == 8 - np.testing.assert_allclose(obj.target, [3, 3, 3, 3, -1, -1, -1, 2]) + assert obj.dim_f == 6 + np.testing.assert_allclose(obj.target, [3, -1, -1, -1, 2, 1]) # only fix currents of some coils in the coil set obj = FixCoilCurrent( - coil=coilset, indices=[[True, False, True, False], False, True] + coil=coilset, indices=[[True], [True, False, True], True, False] ) obj.build() - assert obj.dim_f == 3 - np.testing.assert_allclose(obj.target, [3, 3, 2]) + assert obj.dim_f == 4 + np.testing.assert_allclose(obj.target, [3, -1, -1, 2])