From e2cd60b595e8f78ff28c3784597b5e2453fefc90 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 16 Jul 2024 15:47:28 -0600 Subject: [PATCH 1/7] FixSumCoilCurrent objective --- desc/objectives/__init__.py | 1 + desc/objectives/linear_objectives.py | 137 ++++++++++++++++++++++++++- docs/api.rst | 2 + tests/test_linear_objectives.py | 26 +++++ 4 files changed, 165 insertions(+), 1 deletion(-) diff --git a/desc/objectives/__init__.py b/desc/objectives/__init__.py index f14c19c96e..02a462de26 100644 --- a/desc/objectives/__init__.py +++ b/desc/objectives/__init__.py @@ -77,6 +77,7 @@ FixPressure, FixPsi, FixSheetCurrent, + FixSumCoilCurrent, FixSumModesLambda, FixSumModesR, FixSumModesZ, diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index b20317e3aa..0f9c3ed653 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -151,7 +151,7 @@ def __init__( weight=1, normalize=True, normalize_target=True, - name="Fixed parameters", + name="fixed parameters", ): self._params = params super().__init__( @@ -2726,6 +2726,141 @@ def build(self, use_jit=False, verbose=1): super().build(use_jit=use_jit, verbose=verbose) +class FixSumCoilCurrent(FixCoilCurrent): + """Fixes the sum of coil current(s) in a Coil or CoilSet. + + Parameters + ---------- + coil : Coil + Coil(s) that will be optimized to satisfy the Objective. + target : {float, ndarray}, optional + Target value(s) of the objective. Only used if bounds is None. + Must be broadcastable to Objective.dim_f. Default is FIXME. + bounds : tuple of {float, ndarray}, optional + Lower and upper bounds on the objective. Overrides target. + Both bounds must be broadcastable to to Objective.dim_f. + Default is FIXME. + weight : {float, ndarray}, optional + Weighting to apply to the Objective, relative to other Objectives. + Must be broadcastable to to Objective.dim_f + normalize : bool, optional + Whether to compute the error in physical units or non-dimensionalize. + normalize_target : bool, optional + Whether target and bounds should be normalized before comparing to computed + values. If `normalize` is `True` and the target is in physical units, + this should also be set to True. + indices : nested list of bool, optional + Pytree of bool specifying which coil currents to sum together. + See the example for how to use this on a mixed coil set. + If True/False sums all/none of the coil currents. + name : str, optional + Name of the objective function. + + Examples + -------- + .. code-block:: python + + import numpy as np + from desc.coils import ( + CoilSet, FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, MixedCoilSet + ) + from desc.objectives import FixSumCoilCurrent + + # toroidal field coil set with 4 coils + 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) + # vertical field coil set with 3 coils + 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 + ) + # another single coil + xyz_coil = FourierXYZCoil(current=2) + # full coil set with TF coils, VF coils, and other single coil + full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xyz_coil)) + + # fix the current of the 1st & 3rd TF coil + # fix none of the currents in the VF coil set + # fix the current of the other coil + obj = FixSumCoilCurrent( + full_coilset, indices=[[True, False, True, False], False, True] + ) + + """ + + _scalar = True + _linear = True + _fixed = False + _units = "(A)" + _print_value_fmt = "Summed coil current error: {:10.3e} " + + def __init__( + self, + coil, + target=None, + bounds=None, + weight=1, + normalize=True, + normalize_target=True, + indices=True, + name="summed coil current", + ): + if target is None and bounds is None: + target = 0 + super().__init__( + coil=coil, + target=target, + bounds=bounds, + weight=weight, + normalize=normalize, + normalize_target=normalize_target, + indices=indices, + name=name, + ) + + def build(self, use_jit=False, verbose=1): + """Build constant arrays. + + Parameters + ---------- + use_jit : bool, optional + Whether to just-in-time compile the objective and derivatives. + verbose : int, optional + Level of output. + + """ + super().build(use_jit=use_jit, verbose=verbose) + self._dim_f = 1 + + def compute(self, params, constants=None): + """Compute sum of coil currents. + + Parameters + ---------- + params : list of dict + List of dictionaries of degrees of freedom, eg CoilSet.params_dict + constants : dict + Dictionary of constant data, eg transforms, profiles etc. Defaults to + self.constants + + Returns + ------- + f : ndarray + Sum of coil currents. + + """ + return jnp.sum( + jnp.concatenate( + [ + jnp.atleast_1d(param[idx]) + for param, idx in zip(tree_leaves(params), self._indices) + ] + ) + ) + + class FixOmniWell(FixParameters): """Fixes OmnigenousField.B_lm coefficients. diff --git a/docs/api.rst b/docs/api.rst index 94ff742b5e..3173f8b418 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -187,6 +187,8 @@ Objective Functions desc.objectives.FixParameters desc.objectives.FixPressure desc.objectives.FixPsi + desc.objectives.FixSumCoilCurrent + desc.objectives.FixSumModesLambda desc.objectives.FixSumModesR desc.objectives.FixSumModesZ desc.objectives.FixThetaSFL diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index c6f364718a..5bf60e9ef1 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -37,6 +37,7 @@ FixParameters, FixPressure, FixPsi, + FixSumCoilCurrent, FixSumModesLambda, FixSumModesR, FixSumModesZ, @@ -1006,3 +1007,28 @@ def test_fix_coil_current(DummyMixedCoilSet): obj.build() assert obj.dim_f == 4 np.testing.assert_allclose(obj.target, [3, -1, -1, 2]) + + +@pytest.mark.unit +def test_fix_sum_coil_current(DummyMixedCoilSet): + """Tests FixSumCoilCurrent.""" + coilset = load(load_from=str(DummyMixedCoilSet["output_path"]), file_format="hdf5") + + # sum a single coil current + obj = FixSumCoilCurrent(coil=coilset.coils[1].coils[0]) + obj.build() + params = coilset.coils[1].coils[0].params_dict + np.testing.assert_allclose(obj.compute(params), -1) + + # sum all coil currents + obj = FixSumCoilCurrent(coil=coilset) + obj.build() + params = coilset.params_dict + np.testing.assert_allclose(obj.compute(params), 3) + + # only sum currents of some coils in the coil set + obj = FixSumCoilCurrent( + coil=coilset, indices=[[True], [True, False, True], True, False] + ) + obj.build() + np.testing.assert_allclose(obj.compute(params), 3) From 5d041c3a971368f1c3ca0aabc9a73f2f3aefed55 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 16 Jul 2024 16:10:00 -0600 Subject: [PATCH 2/7] add test for LinearObjectiveFromUser with OptimizableCollection --- tests/test_linear_objectives.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index 5bf60e9ef1..fb680563d9 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -6,6 +6,7 @@ from qsc import Qsc import desc.examples +from desc.backend import jnp from desc.equilibrium import Equilibrium from desc.geometry import FourierRZToroidalSurface from desc.grid import LinearGrid @@ -43,6 +44,7 @@ FixSumModesZ, FixThetaSFL, GenericObjective, + LinearObjectiveFromUser, ObjectiveFunction, get_equilibrium_objective, get_fixed_axis_constraints, @@ -1032,3 +1034,24 @@ def test_fix_sum_coil_current(DummyMixedCoilSet): ) obj.build() np.testing.assert_allclose(obj.compute(params), 3) + + +@pytest.mark.unit +def test_linear_objective_from_user_on_collection(DummyCoilSet): + """Test LinearObjectiveFromUser on an OptimizableCollection.""" + # test that LinearObjectiveFromUser can be used for the same functionality as + # FixSumCoilCurrent to sum all currents in a CoilSet + + coilset = load(load_from=str(DummyCoilSet["output_path_asym"]), file_format="hdf5") + params = coilset.params_dict + + obj1 = FixSumCoilCurrent(coil=coilset) + obj1.build() + + obj2 = LinearObjectiveFromUser( + lambda params: jnp.sum(jnp.array([param["current"] for param in params])), + coilset, + ) + obj2.build() + + np.testing.assert_allclose(obj1.compute(params), obj2.compute(params)) From cc6b88f49cd2d2da340a04e199c0637611f58070 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 16 Jul 2024 18:28:35 -0400 Subject: [PATCH 3/7] add to docstring, add default target --- desc/objectives/linear_objectives.py | 36 +++++++++++++++++++++------- tests/test_linear_objectives.py | 2 ++ 2 files changed, 30 insertions(+), 8 deletions(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 0f9c3ed653..69e2ff2bf3 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -2735,11 +2735,12 @@ class FixSumCoilCurrent(FixCoilCurrent): Coil(s) that will be optimized to satisfy the Objective. target : {float, ndarray}, optional Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Default is FIXME. + Must be broadcastable to Objective.dim_f. Defaults to the current + sum of currents in the coils indicated by indices. bounds : tuple of {float, ndarray}, optional Lower and upper bounds on the objective. Overrides target. Both bounds must be broadcastable to to Objective.dim_f. - Default is FIXME. + Default is to use target. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. Must be broadcastable to to Objective.dim_f @@ -2781,11 +2782,16 @@ class FixSumCoilCurrent(FixCoilCurrent): # full coil set with TF coils, VF coils, and other single coil full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xyz_coil)) - # fix the current of the 1st & 3rd TF coil - # fix none of the currents in the VF coil set - # fix the current of the other coil + # equilibrium G(rho=1) determines the necessary poloidal current through + # the coils (as dictated by Ampere's law) + grid_at_surf = LinearGrid(rho=1.0) + G_tot = eq.compute("G", grid=grid_at_surf)["G"][0] / mu_0 * 2 * jnp.pi + + # we want all coils that link the equilibrium poloidally to be included + # in the sum, so we include the current of the TF coilset and the xyz_coil + # and none of the vf_coilset obj = FixSumCoilCurrent( - full_coilset, indices=[[True, False, True, False], False, True] + full_coilset, indices=[[True, True, True, True], False, True] ) """ @@ -2807,8 +2813,7 @@ def __init__( indices=True, name="summed coil current", ): - if target is None and bounds is None: - target = 0 + self._target_from_user = setdefault(bounds, target) super().__init__( coil=coil, target=target, @@ -2833,6 +2838,21 @@ def build(self, use_jit=False, verbose=1): """ super().build(use_jit=use_jit, verbose=verbose) self._dim_f = 1 + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, + jnp.sum( + jnp.concatenate( + [ + jnp.atleast_1d(param[idx]) + for param, idx in zip( + tree_leaves(self.coil.params), self._indices + ) + ] + ) + ), + None, + np.array([0]), + ) def compute(self, params, constants=None): """Compute sum of coil currents. diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index fb680563d9..b2cdc75e23 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -1027,6 +1027,8 @@ def test_fix_sum_coil_current(DummyMixedCoilSet): obj.build() params = coilset.params_dict np.testing.assert_allclose(obj.compute(params), 3) + # the default target should be the original sum + np.testing.assert_allclose(obj.compute_scaled_error(params), 0) # only sum currents of some coils in the coil set obj = FixSumCoilCurrent( From c091f9613c70d7b7b85f16d60032f720f00b6d12 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 16 Jul 2024 16:33:00 -0600 Subject: [PATCH 4/7] set default target --- desc/objectives/linear_objectives.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 0f9c3ed653..595fe2fca3 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -2735,11 +2735,12 @@ class FixSumCoilCurrent(FixCoilCurrent): Coil(s) that will be optimized to satisfy the Objective. target : {float, ndarray}, optional Target value(s) of the objective. Only used if bounds is None. - Must be broadcastable to Objective.dim_f. Default is FIXME. + Must be broadcastable to Objective.dim_f. + Default is the objective value for the coil. bounds : tuple of {float, ndarray}, optional Lower and upper bounds on the objective. Overrides target. Both bounds must be broadcastable to to Objective.dim_f. - Default is FIXME. + Default is to use the target instead. weight : {float, ndarray}, optional Weighting to apply to the Objective, relative to other Objectives. Must be broadcastable to to Objective.dim_f @@ -2781,9 +2782,9 @@ class FixSumCoilCurrent(FixCoilCurrent): # full coil set with TF coils, VF coils, and other single coil full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xyz_coil)) - # fix the current of the 1st & 3rd TF coil - # fix none of the currents in the VF coil set - # fix the current of the other coil + # fix the sum of: + # the 1st & 3rd TF coil currents, none of the currents in the VF coil set, + # and the current of the other coil obj = FixSumCoilCurrent( full_coilset, indices=[[True, False, True, False], False, True] ) @@ -2807,8 +2808,9 @@ def __init__( indices=True, name="summed coil current", ): + self._default_target = False if target is None and bounds is None: - target = 0 + self._default_target = True super().__init__( coil=coil, target=target, @@ -2833,6 +2835,8 @@ def build(self, use_jit=False, verbose=1): """ super().build(use_jit=use_jit, verbose=verbose) self._dim_f = 1 + if self._default_target: + self.update_target(thing=self.things[0]) def compute(self, params, constants=None): """Compute sum of coil currents. From ebf9dc27b9cfa4a9923c51382445110162d0152d Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 16 Jul 2024 18:44:09 -0400 Subject: [PATCH 5/7] update docstring --- desc/objectives/linear_objectives.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 1e7509c001..c02ba95274 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -2729,6 +2729,11 @@ def build(self, use_jit=False, verbose=1): class FixSumCoilCurrent(FixCoilCurrent): """Fixes the sum of coil current(s) in a Coil or CoilSet. + NOTE: When using this objective, take care in knowing the signs + oof current in the coils and the rientations of the coils, it may + be possible that two coils have differing current signs but in physical + space the currents flow similarly. + Parameters ---------- coil : Coil @@ -2782,12 +2787,14 @@ class FixSumCoilCurrent(FixCoilCurrent): # full coil set with TF coils, VF coils, and other single coil full_coilset = MixedCoilSet((tf_coilset, vf_coilset, xyz_coil)) - # equilibrium G(rho=1) determines the necessary poloidal current through + # equilibrium G(rho=1) determines the necessary net poloidal current through # the coils (as dictated by Ampere's law) grid_at_surf = LinearGrid(rho=1.0, M=eq.M_grid, N=eq.N_grid) G_tot = eq.compute("G", grid=grid_at_surf)["G"][0] / (mu_0 * 2 * jnp.pi) - # only coils that link the equilibrium poloidally should be included in the sum, + # to use this objective to force the coilset to adhere to Ampere's law + # for the targeted equilibrium, we want only coils that link + # the equilibrium poloidally to be included in the sum, # which is the TF coil set and the FourierXYZ coil, but not the VF coil set obj = FixSumCoilCurrent(full_coilset, indices=[True, False, True], target=G_tot) From 1eb33e6ca509e9c8727263db1c9f0080d320259f Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Tue, 16 Jul 2024 16:54:08 -0600 Subject: [PATCH 6/7] improved grammar --- desc/objectives/linear_objectives.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index c02ba95274..6cc5007cb1 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -2729,10 +2729,10 @@ def build(self, use_jit=False, verbose=1): class FixSumCoilCurrent(FixCoilCurrent): """Fixes the sum of coil current(s) in a Coil or CoilSet. - NOTE: When using this objective, take care in knowing the signs - oof current in the coils and the rientations of the coils, it may - be possible that two coils have differing current signs but in physical - space the currents flow similarly. + NOTE: When using this objective, take care in knowing the signs of the current in + the coils and the orientations of the coils. It is possible for coils with the same + signs of their current to have currents flowing in differing directions in physical + space due to the orientation of the coils. Parameters ---------- @@ -2792,9 +2792,8 @@ class FixSumCoilCurrent(FixCoilCurrent): grid_at_surf = LinearGrid(rho=1.0, M=eq.M_grid, N=eq.N_grid) G_tot = eq.compute("G", grid=grid_at_surf)["G"][0] / (mu_0 * 2 * jnp.pi) - # to use this objective to force the coilset to adhere to Ampere's law - # for the targeted equilibrium, we want only coils that link - # the equilibrium poloidally to be included in the sum, + # to use this objective to satisfy Ampere's law for the targeted equilibrium, + # only coils that link the equilibrium poloidally should be included in the sum, # which is the TF coil set and the FourierXYZ coil, but not the VF coil set obj = FixSumCoilCurrent(full_coilset, indices=[True, False, True], target=G_tot) From d59f2b58ef833f80209d379bfcc16a7277278eb2 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 16 Jul 2024 23:06:40 -0400 Subject: [PATCH 7/7] address comments --- desc/objectives/linear_objectives.py | 4 +++- docs/api_objectives.rst | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 6cc5007cb1..43b8a20feb 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -2789,8 +2789,10 @@ class FixSumCoilCurrent(FixCoilCurrent): # equilibrium G(rho=1) determines the necessary net poloidal current through # the coils (as dictated by Ampere's law) + # the sign convention is positive poloidal current flows up through the torus + # hole grid_at_surf = LinearGrid(rho=1.0, M=eq.M_grid, N=eq.N_grid) - G_tot = eq.compute("G", grid=grid_at_surf)["G"][0] / (mu_0 * 2 * jnp.pi) + G_tot = 2*jnp.pi*eq.compute("G", grid=grid_at_surf)["G"][0] / mu_0 # to use this objective to satisfy Ampere's law for the targeted equilibrium, # only coils that link the equilibrium poloidally should be included in the sum, diff --git a/docs/api_objectives.rst b/docs/api_objectives.rst index a3cfdb7f31..273892aab9 100644 --- a/docs/api_objectives.rst +++ b/docs/api_objectives.rst @@ -146,6 +146,7 @@ Fixing degrees of freedom desc.objectives.FixSumModesZ desc.objectives.FixThetaSFL desc.objectives.FixCoilCurrent + desc.objectives.FixSumCoilCurrent desc.objectives.FixParameters