From e0dae7ffeefacc7abfda925e26861ee58c26c99a Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 14:52:46 -0400 Subject: [PATCH 01/29] copy over dummy coilsets and test --- tests/conftest.py | 13 +++++++- tests/test_examples.py | 68 ++++++++++++++++++++++++++++-------------- 2 files changed, 57 insertions(+), 24 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 3d57ea5ab4..e011a59379 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 @@ -276,12 +277,22 @@ def DummyMixedCoilSet(tmpdir_factory): 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_coilset = CoilSet(tf_coilset.coils, NFP=tf_coilset.NFP, 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_examples.py b/tests/test_examples.py index 053e3c2bac..2acb15dad1 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -9,7 +9,7 @@ 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, @@ -1336,30 +1336,52 @@ 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.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_coilset = load( + load_from=str(DummyCoilSet["output_path_sym"]), file_format="hdf5" + ) + asym_coilset = load( + load_from=str(DummyCoilSet["output_path_asym"]), file_format="hdf5" + ) + mixed_coilset = load( + load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5" + ) + + def test(c, maxiter=200, constraints=()): + obj = ObjectiveFunction(CoilLength(c, target=11)) + optimizer = Optimizer("fmintr") + (c,), _ = optimizer.optimize( + c, obj, constraints, maxiter=maxiter, ftol=1e-15, xtol=1e-15 + ) + flattened_coils = tree_leaves(c, is_leaf=lambda x: not hasattr(x, "__len__")) + lengths = [coil.compute("length")["length"] for coil in flattened_coils] + + assert obj.dim_f == len(flattened_coils) + np.testing.assert_allclose(lengths, 11, atol=1e-3) + + spline_coil = mixed_coilset.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()) + test(FourierRZCoil()) + test(FourierXYZCoil()) + test(spline_coil, constraints=(FixParameters(spline_coil, {"knots": True}),)) + + # CoilSet + test(sym_coilset, maxiter=500) + test(asym_coilset, maxiter=500) + + # MixedCoilSet + params = [ + {"center": True, "normal": True}, + {"R_n": True, "Z_n": True}, + {}, + {"knots": True}, + ] + constraints = (FixParameters(mixed_coilset, params),) + test(mixed_coilset, constraints=constraints, maxiter=400) @pytest.mark.unit From 89feb048f2eec10df6c51f26a418b13d2696068f Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 15:43:46 -0400 Subject: [PATCH 02/29] update _curve.py for spline optimization - add knots to transforms - add method to kwargs - replace names of these as necessary --- desc/compute/_curve.py | 109 +++++++++++++++++++++-------------------- 1 file changed, 57 insertions(+), 52 deletions(-) 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) From 1ced35927380508e28a65823e3d19cd999918ca7 Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 15:50:40 -0400 Subject: [PATCH 03/29] update curve.py for spline optimization - add compute func for overriding so _curve.py knows about method --- desc/geometry/curve.py | 41 ++++++++++++++++++++++++++++++++++++++++- 1 file changed, 40 insertions(+), 1 deletion(-) diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index b928818cdc..a4616b6642 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,45 @@ 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. From bea29ba671877ab3f6217247c2d29f2801b684f9 Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 15:51:51 -0400 Subject: [PATCH 04/29] update test - save and load coilsets - use different methods for different optimizations --- tests/test_examples.py | 37 ++++++++++++++++--------------------- 1 file changed, 16 insertions(+), 21 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 2acb15dad1..5766501c2b 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -16,6 +16,7 @@ FourierRZCoil, FourierXYZCoil, MixedCoilSet, + _Coil, ) from desc.continuation import solve_continuation_automatic from desc.equilibrium import EquilibriaFamily, Equilibrium @@ -1349,39 +1350,33 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5" ) - def test(c, maxiter=200, constraints=()): - obj = ObjectiveFunction(CoilLength(c, target=11)) - optimizer = Optimizer("fmintr") - (c,), _ = optimizer.optimize( - c, obj, constraints, maxiter=maxiter, ftol=1e-15, xtol=1e-15 + def test(c, method): + target = 11 + obj = ObjectiveFunction(CoilLength(c, target=target)) + optimizer = Optimizer(method) + (c,), _ = optimizer.optimize(c, obj, maxiter=200, ftol=0, xtol=1e-15) + flattened_coils = tree_leaves( + c, is_leaf=lambda x: isinstance(x, _Coil) and not isinstance(x, CoilSet) ) - flattened_coils = tree_leaves(c, is_leaf=lambda x: not hasattr(x, "__len__")) lengths = [coil.compute("length")["length"] for coil in flattened_coils] assert obj.dim_f == len(flattened_coils) - np.testing.assert_allclose(lengths, 11, atol=1e-3) + np.testing.assert_allclose(lengths, target, atol=1e-3) spline_coil = mixed_coilset.coils[-1].copy() # single coil - test(FourierPlanarCoil()) - test(FourierRZCoil()) - test(FourierXYZCoil()) - test(spline_coil, constraints=(FixParameters(spline_coil, {"knots": True}),)) + test(FourierPlanarCoil(), "fmintr") + test(FourierRZCoil(), "fmintr") + test(FourierXYZCoil(), "fmintr") + test(spline_coil, "fmintr") # CoilSet - test(sym_coilset, maxiter=500) - test(asym_coilset, maxiter=500) + test(sym_coilset, "lsq-exact") + test(asym_coilset, "lsq-exact") # MixedCoilSet - params = [ - {"center": True, "normal": True}, - {"R_n": True, "Z_n": True}, - {}, - {"knots": True}, - ] - constraints = (FixParameters(mixed_coilset, params),) - test(mixed_coilset, constraints=constraints, maxiter=400) + test(mixed_coilset, "lsq-exact") @pytest.mark.unit From 7ba2a7335a9273e52ba9ac7ea193b8631af3736a Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 16:11:30 -0400 Subject: [PATCH 05/29] add nested coilset --- tests/test_examples.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 5766501c2b..bf55c6058a 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1340,15 +1340,14 @@ def test_second_stage_optimization_CoilSet(): @pytest.mark.unit def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): """Test optimizing for every type of coil and dummy coil sets.""" - sym_coilset = load( - load_from=str(DummyCoilSet["output_path_sym"]), file_format="hdf5" - ) - asym_coilset = load( + 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_coilset = load( + mixed_coils = load( load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5" ) + nested_coils = MixedCoilSet(sym_coils, mixed_coils) def test(c, method): target = 11 @@ -1363,7 +1362,7 @@ def test(c, method): assert obj.dim_f == len(flattened_coils) np.testing.assert_allclose(lengths, target, atol=1e-3) - spline_coil = mixed_coilset.coils[-1].copy() + spline_coil = mixed_coils.coils[-1].copy() # single coil test(FourierPlanarCoil(), "fmintr") @@ -1372,11 +1371,12 @@ def test(c, method): test(spline_coil, "fmintr") # CoilSet - test(sym_coilset, "lsq-exact") - test(asym_coilset, "lsq-exact") + test(sym_coils, "lsq-exact") + test(asym_coils, "lsq-exact") # MixedCoilSet - test(mixed_coilset, "lsq-exact") + test(mixed_coils, "lsq-exact") + test(nested_coils, "lsq-exact") @pytest.mark.unit From 680848b9285d3bbd57177db974e615ad1733860e Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 16:26:16 -0400 Subject: [PATCH 06/29] add QuadraticFlux obj - doesn't pass due to jax bug, maybe PR #1002 will fix it? --- tests/test_examples.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index bf55c6058a..4639b72e62 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1351,7 +1351,10 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): def test(c, method): target = 11 - obj = ObjectiveFunction(CoilLength(c, target=target)) + obj = ObjectiveFunction( + CoilLength(c, target=target), + QuadraticFlux(eq=Equilibrium(), field=c, vacuum=True), + ) optimizer = Optimizer(method) (c,), _ = optimizer.optimize(c, obj, maxiter=200, ftol=0, xtol=1e-15) flattened_coils = tree_leaves( From 4fbc0c6bd41385500009dd037fe7d4c92dc295c4 Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 17:07:25 -0400 Subject: [PATCH 07/29] fix obj in test by making funcs a tuple --- tests/test_examples.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 4639b72e62..76f76f54b5 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1352,8 +1352,10 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): def test(c, method): target = 11 obj = ObjectiveFunction( - CoilLength(c, target=target), - QuadraticFlux(eq=Equilibrium(), field=c, vacuum=True), + ( + CoilLength(c, target=target), + QuadraticFlux(eq=Equilibrium(), field=c, vacuum=True), + ), ) optimizer = Optimizer(method) (c,), _ = optimizer.optimize(c, obj, maxiter=200, ftol=0, xtol=1e-15) @@ -1361,8 +1363,6 @@ def test(c, method): c, is_leaf=lambda x: isinstance(x, _Coil) and not isinstance(x, CoilSet) ) lengths = [coil.compute("length")["length"] for coil in flattened_coils] - - assert obj.dim_f == len(flattened_coils) np.testing.assert_allclose(lengths, target, atol=1e-3) spline_coil = mixed_coils.coils[-1].copy() From b7b03e075310d2230dc3171fba71b11aa90f7d4d Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 17:32:57 -0400 Subject: [PATCH 08/29] add {} to params for spline coil in other test --- tests/test_linear_objectives.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index bed8c6482a..50107423f9 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -965,6 +965,7 @@ def test_fix_subset_of_params_in_collection(DummyMixedCoilSet): ], {"shift": True, "rotmat": True}, {"X_n": np.array([1, 2]), "Y_n": False, "Z_n": np.array([0])}, + {}, ] target = np.concatenate( ( From 6652543b99c95808bc0c64d93da00901620a18d7 Mon Sep 17 00:00:00 2001 From: kianorr Date: Thu, 4 Jul 2024 17:37:05 -0400 Subject: [PATCH 09/29] fix dim_f assertion to accomodate for spline coil addition --- tests/test_linear_objectives.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index 50107423f9..7e6bf17093 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -999,13 +999,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 == 9 + np.testing.assert_allclose(obj.target, [3, 3, 3, 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, False, True, False], False, True, True] ) 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, 3, 2, 1]) From e4e692a8d955254ff114ee137a705ea1b20eafde Mon Sep 17 00:00:00 2001 From: kianorr Date: Fri, 5 Jul 2024 14:43:00 -0400 Subject: [PATCH 10/29] add curvature and torsion to replace single coil test --- tests/test_examples.py | 65 +++++++++++------------------------------- 1 file changed, 17 insertions(+), 48 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 76f76f54b5..ebdf7b0064 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1186,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. @@ -1348,22 +1310,29 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5" ) nested_coils = MixedCoilSet(sym_coils, mixed_coils) + eq = Equilibrium() - def test(c, method): + def test(c, method, add_objs=False): target = 11 - obj = ObjectiveFunction( - ( - CoilLength(c, target=target), - QuadraticFlux(eq=Equilibrium(), field=c, vacuum=True), - ), - ) + rtol = 1e-3 + objs = [ + CoilLength(c, target=target), + QuadraticFlux(eq=eq, field=c, vacuum=True), + ] + + if add_objs: + objs.extend([CoilCurvature(c), CoilTorsion(c, target=0)]) + rtol = 1e-2 + + obj = ObjectiveFunction(objs) + optimizer = Optimizer(method) - (c,), _ = optimizer.optimize(c, obj, maxiter=200, ftol=0, xtol=1e-15) + (c,), _ = optimizer.optimize(c, obj, maxiter=1000, ftol=0, 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, atol=1e-3) + np.testing.assert_allclose(lengths, target, rtol=rtol) spline_coil = mixed_coils.coils[-1].copy() @@ -1378,7 +1347,7 @@ def test(c, method): test(asym_coils, "lsq-exact") # MixedCoilSet - test(mixed_coils, "lsq-exact") + test(mixed_coils, "lsq-exact", add_objs=True) test(nested_coils, "lsq-exact") From 978dcba1b5b78ddc1a3c7641aab7a5d8102d5997 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 5 Jul 2024 17:00:40 -0400 Subject: [PATCH 11/29] fix coilset, reduce resolution of quadratix flux objective in test --- tests/conftest.py | 4 ++-- tests/test_examples.py | 14 +++++++++++++- 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index e011a59379..197fcfef39 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -276,8 +276,8 @@ 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_coilset = CoilSet(tf_coilset.coils, NFP=tf_coilset.NFP, sym=True) + 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( diff --git a/tests/test_examples.py b/tests/test_examples.py index ebdf7b0064..a8f9347256 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1299,6 +1299,7 @@ def test_second_stage_optimization_CoilSet(): np.testing.assert_allclose(field[0].current, 0, atol=1e-12) +@pytest.mark.slow @pytest.mark.unit def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): """Test optimizing for every type of coil and dummy coil sets.""" @@ -1311,13 +1312,24 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): ) 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, add_objs=False): target = 11 rtol = 1e-3 objs = [ CoilLength(c, target=target), - QuadraticFlux(eq=eq, field=c, vacuum=True), + QuadraticFlux( + eq=eq, + field=c, + vacuum=True, + weight=1e-4, + eval_grid=quad_eval_grid, + field_grid=quad_field_grid, + ), ] if add_objs: From ff38a5084d1bf85c0c3dec0d95e3d7e7322810b7 Mon Sep 17 00:00:00 2001 From: kianorr Date: Fri, 5 Jul 2024 18:11:27 -0400 Subject: [PATCH 12/29] fix coilset with symmetry in DummyMixedCoilSet --- tests/conftest.py | 4 ++-- tests/test_examples.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index e011a59379..197fcfef39 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -276,8 +276,8 @@ 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_coilset = CoilSet(tf_coilset.coils, NFP=tf_coilset.NFP, sym=True) + 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( diff --git a/tests/test_examples.py b/tests/test_examples.py index ebdf7b0064..8da68e49cb 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1309,10 +1309,10 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): mixed_coils = load( load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5" ) - nested_coils = MixedCoilSet(sym_coils, mixed_coils) + nested_coils = MixedCoilSet(sym_coils, asym_coils) eq = Equilibrium() - def test(c, method, add_objs=False): + def test(c, method): target = 11 rtol = 1e-3 objs = [ @@ -1320,7 +1320,7 @@ def test(c, method, add_objs=False): QuadraticFlux(eq=eq, field=c, vacuum=True), ] - if add_objs: + if isinstance(c, MixedCoilSet): objs.extend([CoilCurvature(c), CoilTorsion(c, target=0)]) rtol = 1e-2 @@ -1347,7 +1347,7 @@ def test(c, method, add_objs=False): test(asym_coils, "lsq-exact") # MixedCoilSet - test(mixed_coils, "lsq-exact", add_objs=True) + test(mixed_coils, "lsq-exact") test(nested_coils, "lsq-exact") From 94fa649cd752154aeeccf7f6ef0e7789bd5c1953 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 5 Jul 2024 18:22:38 -0400 Subject: [PATCH 13/29] modify test --- tests/test_examples.py | 34 +++++++++++++++++++++++++--------- 1 file changed, 25 insertions(+), 9 deletions(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index a8f9347256..06780dc80c 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1320,8 +1320,9 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): def test(c, method, add_objs=False): target = 11 rtol = 1e-3 - objs = [ - CoilLength(c, target=target), + # first just check that it quad flux works for a couple iterations + # as this is an expensive objective to compute + obj = ObjectiveFunction( QuadraticFlux( eq=eq, field=c, @@ -1329,22 +1330,37 @@ def test(c, method, add_objs=False): 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 add_objs: - objs.extend([CoilCurvature(c), CoilTorsion(c, target=0)]) - rtol = 1e-2 + # 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) - optimizer = Optimizer(method) - (c,), _ = optimizer.optimize(c, obj, maxiter=1000, ftol=0, xtol=1e-15) + (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) + np.testing.assert_allclose( + lengths, target, rtol=rtol, err_msg=f"lengths {c}" + extra_msg + ) spline_coil = mixed_coils.coils[-1].copy() From ab92b04abc01438726cde72bf5f1a7f03e6ebc2e Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 5 Jul 2024 18:28:58 -0400 Subject: [PATCH 14/29] fix comment --- tests/test_examples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index a24ed87821..0488aae450 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1320,7 +1320,7 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): def test(c, method): target = 11 rtol = 1e-3 - # first just check that it quad flux works for a couple iterations + # first just check that quad flux works for a couple iterations # as this is an expensive objective to compute obj = ObjectiveFunction( QuadraticFlux( From dbc6dc1531e58f5b2370f515c95d592f98a6fa0f Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 5 Jul 2024 22:37:36 -0400 Subject: [PATCH 15/29] fix docstring --- desc/geometry/curve.py | 1 + 1 file changed, 1 insertion(+) diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index a4616b6642..f5281e1efd 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -949,6 +949,7 @@ def compute( 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 From b2e17b6ada552bce1856f5478780c9da64c48ea8 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 5 Jul 2024 22:48:16 -0400 Subject: [PATCH 16/29] fix test --- tests/test_linear_objectives.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index 7e6bf17093..ec24523759 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -959,9 +959,6 @@ 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])}, @@ -969,7 +966,7 @@ def test_fix_subset_of_params_in_collection(DummyMixedCoilSet): ] 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(), From 4adfa2235929974ac602d990016168d9a9dabed7 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 5 Jul 2024 23:19:10 -0400 Subject: [PATCH 17/29] fix last test --- tests/test_linear_objectives.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index ec24523759..f04c854639 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -996,13 +996,11 @@ def test_fix_coil_current(DummyMixedCoilSet): # fix all coil currents obj = FixCoilCurrent(coil=coilset) obj.build() - assert obj.dim_f == 9 - np.testing.assert_allclose(obj.target, [3, 3, 3, 3, -1, -1, -1, 2, 1]) + 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, True] - ) + obj = FixCoilCurrent(coil=coilset, indices=[[True], False, True, True]) obj.build() - assert obj.dim_f == 4 - np.testing.assert_allclose(obj.target, [3, 3, 2, 1]) + assert obj.dim_f == 3 + np.testing.assert_allclose(obj.target, [3, 2, 1]) From 1319e955e5f5156fc177c089faa2d05ca0c5c820 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 7 Jul 2024 01:24:19 -0400 Subject: [PATCH 18/29] allow MixedCoilSet.from_symmetry to work when the given coils are not CoilSet-compatible --- desc/coils.py | 12 +++++++++--- tests/test_coils.py | 13 +++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index 4ac96fc5cb..1629941789 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -1153,9 +1153,15 @@ def from_symmetry(cls, coils, NFP=1, sym=False): """ if not isinstance(coils, CoilSet): - coils = CoilSet(coils) - - [_check_type(coil, coils[0]) for coil in coils] + try: + coils = CoilSet(coils) + except (TypeError, ValueError): + # likely there are multiple coil types, + # so make a MixedCoilSet + coils = MixedCoilSet(coils) + if not isinstance(coils, MixedCoilSet): + # only need to check this for a CoilSet, not MixedCoilSet + [_check_type(coil, coils[0]) for coil in coils] # check toroidal extent of coils to be repeated maxphi = 2 * np.pi / NFP / (sym + 1) diff --git a/tests/test_coils.py b/tests/test_coils.py index fb501ba13a..4c5d38e4ff 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -294,6 +294,19 @@ def test_from_symmetry(self): )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) + # With a MixedCoilSet as the base coils and only rotation + coil = FourierPlanarCoil(I) + coils = [coil] + [FourierXYZCoil(I) for i in range(N // 4 - 1)] + for i, c in enumerate(coils[1:]): + c.rotate(angle=2 * np.pi / N * (i + 1)) + coils = MixedCoilSet.from_symmetry(coils, NFP=4) + grid = LinearGrid(N=32, endpoint=False) + transforms = get_transforms(["x", "x_s", "ds"], coil, grid=grid) + B_approx = coils.compute_magnetic_field( + [10, 0, 0], basis="rpz", source_grid=grid + )[0] + np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) + @pytest.mark.unit def test_properties(self): """Test getting/setting of CoilSet attributes.""" From 94dc6a5bfc20d9e597a271a457397143c92ad830 Mon Sep 17 00:00:00 2001 From: kianorr Date: Mon, 8 Jul 2024 11:59:19 -0400 Subject: [PATCH 19/29] update test_fix_coil_current --- tests/test_linear_objectives.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index f04c854639..c6f364718a 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -1000,7 +1000,9 @@ def test_fix_coil_current(DummyMixedCoilSet): 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, True]) + obj = FixCoilCurrent( + coil=coilset, indices=[[True], [True, False, True], True, False] + ) obj.build() - assert obj.dim_f == 3 - np.testing.assert_allclose(obj.target, [3, 2, 1]) + assert obj.dim_f == 4 + np.testing.assert_allclose(obj.target, [3, -1, -1, 2]) From da819dc762921a12e9849cf5f216fb7e1121d61f Mon Sep 17 00:00:00 2001 From: kianorr Date: Mon, 8 Jul 2024 12:12:09 -0400 Subject: [PATCH 20/29] add mixed_coils to nested_coils --- tests/test_examples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_examples.py b/tests/test_examples.py index 0488aae450..70ff2023c6 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -1310,7 +1310,7 @@ def test_optimize_with_all_coil_types(DummyCoilSet, DummyMixedCoilSet): mixed_coils = load( load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5" ) - nested_coils = MixedCoilSet(sym_coils, asym_coils) + nested_coils = MixedCoilSet(sym_coils, mixed_coils) eq = Equilibrium() # not attempting to accurately calc B for this test, # so make the grids very coarse From 615555fb5a0df3de668410244f872bd392266859 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 8 Jul 2024 17:32:36 -0400 Subject: [PATCH 21/29] change default grid for coil objective to use NFP if coil has it --- desc/objectives/_coils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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): From a014d923765711f04adaa9ea808c6b59887cf5f9 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 8 Jul 2024 18:37:59 -0400 Subject: [PATCH 22/29] Revert "change default grid for coil objective to use NFP if coil has it" This reverts commit 615555fb5a0df3de668410244f872bd392266859. --- desc/objectives/_coils.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 3be3415535..9892e22b78 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -128,10 +128,7 @@ def _prune_coilset_tree(coilset): # map grid to list of length coils if grid is None: - 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)) + grid = [LinearGrid(N=2 * c.N + 5, endpoint=False) for c in coils] if isinstance(grid, numbers.Integral): grid = LinearGrid(N=self._grid, endpoint=False) if isinstance(grid, _Grid): From 6718a00ac9e272dde31d5168f5f822bd05fcb8f6 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 8 Jul 2024 18:53:26 -0400 Subject: [PATCH 23/29] change default grid for NFP>1 coils again --- desc/objectives/_coils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 9892e22b78..885be00a0c 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -128,7 +128,12 @@ 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] + for c in coils: + NFP = c.NFP if hasattr(c, "NFP") else 1 + # NFP=1 to ensure we have points along whole grid + # multiply by NFP in case the coil has nonzero NFP + # to make ensure sufficient resolution across periods + grid.append(LinearGrid(N=2 * c.N * NFP + 5, NFP=1, endpoint=False)) if isinstance(grid, numbers.Integral): grid = LinearGrid(N=self._grid, endpoint=False) if isinstance(grid, _Grid): From c28d4d45d835cc11712b4f8757aec3dd4e0a111d Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 8 Jul 2024 19:21:43 -0400 Subject: [PATCH 24/29] fix default grid --- desc/objectives/_coils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 885be00a0c..07f1185006 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -128,6 +128,7 @@ def _prune_coilset_tree(coilset): # map grid to list of length coils if grid is None: + grid = [] for c in coils: NFP = c.NFP if hasattr(c, "NFP") else 1 # NFP=1 to ensure we have points along whole grid From 5279e6340f00bb4885019e8160992eab34383879 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 9 Jul 2024 00:08:30 -0400 Subject: [PATCH 25/29] add logic for coils with NFP to coil compute_magnetic_field --- desc/coils.py | 6 ++++++ tests/test_coils.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/desc/coils.py b/desc/coils.py index 1629941789..a34fe6a591 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 nonzero NFP + # to make sure 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/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.""" From 17a3f0dd7a12787638fac92e1883e7865d281fb5 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 9 Jul 2024 10:41:53 -0600 Subject: [PATCH 26/29] add NFP correction to Curve.compute default grid --- desc/coils.py | 6 ++---- desc/geometry/core.py | 4 +++- desc/objectives/_coils.py | 5 ++--- 3 files changed, 7 insertions(+), 8 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index a34fe6a591..afcf9c7e4e 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -227,10 +227,8 @@ def compute_magnetic_field( 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 nonzero NFP - # to make sure whole coil gets counted for the - # biot savart integration + # NFP=1 and multiply N by NFP in case the coil has NFP>1 to ensure the whole + # coil is included in the biot savart integration source_grid = LinearGrid(N=2 * self.N * self.NFP + 5, NFP=1, endpoint=False) if not params or not transforms: diff --git a/desc/geometry/core.py b/desc/geometry/core.py index c087a71c6d..03db702a99 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -115,7 +115,9 @@ def compute( 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=False) + # NFP=1 and multiply N by NFP in case the coil has NFP>1 to ensure + # sufficient resolution across the whole coil + grid = LinearGrid(N=2 * self.N * NFP + 5, NFP=1, endpoint=False) elif isinstance(grid, numbers.Integral): NFP = self.NFP if hasattr(self, "NFP") else 1 grid = LinearGrid(N=grid, NFP=NFP, endpoint=False) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 07f1185006..1782c0a496 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -131,9 +131,8 @@ def _prune_coilset_tree(coilset): grid = [] for c in coils: NFP = c.NFP if hasattr(c, "NFP") else 1 - # NFP=1 to ensure we have points along whole grid - # multiply by NFP in case the coil has nonzero NFP - # to make ensure sufficient resolution across periods + # NFP=1 and multiply N by NFP in case the coil has NFP>1 to ensure + # sufficient resolution across the whole coil grid.append(LinearGrid(N=2 * c.N * NFP + 5, NFP=1, endpoint=False)) if isinstance(grid, numbers.Integral): grid = LinearGrid(N=self._grid, endpoint=False) From 28f5e33a3e7793cea0108c1be1f0f98405d31227 Mon Sep 17 00:00:00 2001 From: kianorr Date: Tue, 9 Jul 2024 14:07:00 -0400 Subject: [PATCH 27/29] Revert "add NFP correction to Curve.compute default grid" This reverts commit 17a3f0dd7a12787638fac92e1883e7865d281fb5. --- desc/coils.py | 6 ++++-- desc/geometry/core.py | 4 +--- desc/objectives/_coils.py | 5 +++-- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index afcf9c7e4e..a34fe6a591 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -227,8 +227,10 @@ def compute_magnetic_field( else: current = params.pop("current", self.current) if source_grid is None and hasattr(self, "NFP"): - # NFP=1 and multiply N by NFP in case the coil has NFP>1 to ensure the whole - # coil is included in the biot savart integration + # NFP=1 to ensure we have points along whole grid + # multiply by NFP in case the coil has nonzero NFP + # to make sure 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: diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 03db702a99..c087a71c6d 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -115,9 +115,7 @@ def compute( names = [names] if grid is None: NFP = self.NFP if hasattr(self, "NFP") else 1 - # NFP=1 and multiply N by NFP in case the coil has NFP>1 to ensure - # sufficient resolution across the whole coil - grid = LinearGrid(N=2 * self.N * NFP + 5, NFP=1, endpoint=False) + grid = LinearGrid(N=2 * self.N + 5, NFP=NFP, endpoint=False) elif isinstance(grid, numbers.Integral): NFP = self.NFP if hasattr(self, "NFP") else 1 grid = LinearGrid(N=grid, NFP=NFP, endpoint=False) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 1782c0a496..07f1185006 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -131,8 +131,9 @@ def _prune_coilset_tree(coilset): grid = [] for c in coils: NFP = c.NFP if hasattr(c, "NFP") else 1 - # NFP=1 and multiply N by NFP in case the coil has NFP>1 to ensure - # sufficient resolution across the whole coil + # NFP=1 to ensure we have points along whole grid + # multiply by NFP in case the coil has nonzero NFP + # to make ensure sufficient resolution across periods grid.append(LinearGrid(N=2 * c.N * NFP + 5, NFP=1, endpoint=False)) if isinstance(grid, numbers.Integral): grid = LinearGrid(N=self._grid, endpoint=False) From b45dff2f89ad8c6b83cfd24ae4fe74c74ce6c362 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 9 Jul 2024 14:39:02 -0400 Subject: [PATCH 28/29] undo the NFP change in coil obj build, as we can have a grid with NFP there since the coil obj are purely geometric --- desc/coils.py | 4 ++-- desc/objectives/_coils.py | 5 +---- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index a34fe6a591..1b5f455719 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -228,8 +228,8 @@ def compute_magnetic_field( 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 nonzero NFP - # to make sure whole coil gets counted for the + # 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) diff --git a/desc/objectives/_coils.py b/desc/objectives/_coils.py index 07f1185006..3be3415535 100644 --- a/desc/objectives/_coils.py +++ b/desc/objectives/_coils.py @@ -131,10 +131,7 @@ def _prune_coilset_tree(coilset): grid = [] for c in coils: NFP = c.NFP if hasattr(c, "NFP") else 1 - # NFP=1 to ensure we have points along whole grid - # multiply by NFP in case the coil has nonzero NFP - # to make ensure sufficient resolution across periods - grid.append(LinearGrid(N=2 * c.N * NFP + 5, NFP=1, endpoint=False)) + 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): From 73c6eaa41def3feec638a1580590a2170204e2ac Mon Sep 17 00:00:00 2001 From: unalmis Date: Wed, 10 Jul 2024 12:59:51 -0400 Subject: [PATCH 29/29] Resolve issue discussed here: https://github.com/PlasmaControl/DESC/pull/1024#discussion_r1667220843 --- desc/grid.py | 34 ++++++++++++++++------------------ tests/test_grid.py | 3 +++ 2 files changed, 19 insertions(+), 18 deletions(-) diff --git a/desc/grid.py b/desc/grid.py index c386c17c80..aad2c86f33 100644 --- a/desc/grid.py +++ b/desc/grid.py @@ -418,7 +418,7 @@ def spacing(self): """ errorif( - not hasattr(self, "_spacing"), + self._spacing is None, AttributeError, "Custom grids must have spacing specified by user.", ) @@ -428,7 +428,7 @@ def spacing(self): def weights(self): """ndarray: Weight for each node, either exact quadrature or volume based.""" errorif( - not hasattr(self, "_weights"), + self._weights is None, AttributeError, "Custom grids must have weights specified by user.", ) @@ -661,24 +661,21 @@ def __init__( self._node_pattern = "custom" self._coordinates = coordinates self._period = period - if source_grid is not None: - self._source_grid = source_grid + self._source_grid = source_grid self._is_meshgrid = bool(is_meshgrid) self._nodes = self._create_nodes(nodes) - if spacing is not None: - spacing = ( - jnp.atleast_2d(jnp.asarray(spacing)) - .reshape(self.nodes.shape) - .astype(float) - ) - self._spacing = spacing - if weights is not None: - weights = ( - jnp.atleast_1d(jnp.asarray(weights)) - .reshape(self.nodes.shape[0]) - .astype(float) - ) - self._weights = weights + self._spacing = ( + jnp.atleast_2d(jnp.asarray(spacing)).reshape(self.nodes.shape).astype(float) + if spacing is not None + else None + ) + self._weights = ( + jnp.atleast_1d(jnp.asarray(weights)) + .reshape(self.nodes.shape[0]) + .astype(float) + if weights is not None + else None + ) if sort: self._sort_nodes() setable_attr = [ @@ -847,6 +844,7 @@ def _create_nodes(self, nodes): @property def source_grid(self): """Coordinates from which this grid was mapped from.""" + errorif(self._source_grid is None, AttributeError) return self._source_grid diff --git a/tests/test_grid.py b/tests/test_grid.py index 0784bf03b5..9d298b8c85 100644 --- a/tests/test_grid.py +++ b/tests/test_grid.py @@ -819,6 +819,9 @@ def test_custom_jitable_grid_indexing(): with pytest.raises(AttributeError): _ = grid2.inverse_zeta_idx + assert not hasattr(grid1, "weights") + assert not hasattr(grid1, "spacing") + assert not hasattr(grid1, "source_grid") assert not hasattr(grid2, "num_rho") assert not hasattr(grid2, "num_theta") assert not hasattr(grid2, "num_zeta")