From 703efca22212c380b763498c7a5c8fc777020056 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 26 Aug 2023 00:10:11 -0400 Subject: [PATCH 1/5] Fixes to linear objectives to work with bounds --- desc/objectives/linear_objectives.py | 243 +++++++++++++-------------- desc/objectives/objective_funs.py | 2 +- 2 files changed, 116 insertions(+), 129 deletions(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index c306d1eaf5..f726999c12 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -14,7 +14,7 @@ from desc.backend import jnp from desc.basis import zernike_radial, zernike_radial_coeffs -from desc.utils import setdefault +from desc.utils import errorif, setdefault from .normalization import compute_scaling_factors from .objective_funs import _Objective @@ -39,6 +39,25 @@ def update_target(self, eq): if self._use_jit: self.jit() + def _parse_target_from_user( + self, target_from_user, default_target, default_bounds, idx + ): + if target_from_user is None: + target = default_target + bounds = default_bounds + elif isinstance(target_from_user, tuple) and ( + len(target_from_user) == 2 + ): # treat as bounds + target = None + bounds = ( + np.broadcast_to(target_from_user[0], self._dim_f).copy()[idx], + np.broadcast_to(target_from_user[1], self._dim_f).copy()[idx], + ) + else: + target = np.broadcast_to(target_from_user, self._dim_f).copy()[idx] + bounds = None + return target, bounds + class BoundaryRSelfConsistency(_Objective): """Ensure that the boundary and interior surfaces are self-consistent. @@ -476,7 +495,7 @@ def __init__( ): self._modes = modes - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) self._surface_label = surface_label super().__init__( things=eq, @@ -506,9 +525,11 @@ def build(self, eq=None, use_jit=False, verbose=1): if self._modes is False or self._modes is None: # no modes modes = np.array([[]], dtype=int) idx = np.array([], dtype=int) + modes_idx = idx elif self._modes is True: # all modes modes = eq.surface.R_basis.modes idx = np.arange(eq.surface.R_basis.num_modes) + modes_idx = idx else: # specified modes modes = np.atleast_2d(self._modes) dtype = { @@ -535,23 +556,12 @@ def build(self, eq=None, use_jit=False, verbose=1): ) self._dim_f = idx.size - if self._target_from_user is not None: - if self._modes is True or self._modes is False: - raise RuntimeError( - "Attempting to provide target for R boundary modes without " - + "providing modes array!" - + "You must pass in the modes corresponding to the" - + "provided target" - ) - # rearrange given target to match modes order - self.target = self._target_from_user[modes_idx] - # Rb_lmn -> Rb optimization space self._A = np.eye(eq.surface.R_basis.num_modes)[idx, :] - # use surface parameters as target if needed - if self._target_from_user is None: - self.target = eq.surface.R_lmn[idx] + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.surface.R_lmn[idx], None, modes_idx + ) if self._normalize: scales = compute_scaling_factors(eq) @@ -635,7 +645,7 @@ def __init__( ): self._modes = modes - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) self._surface_label = surface_label super().__init__( things=eq, @@ -665,9 +675,11 @@ def build(self, eq=None, use_jit=False, verbose=1): if self._modes is False or self._modes is None: # no modes modes = np.array([[]], dtype=int) idx = np.array([], dtype=int) + modes_idx = idx elif self._modes is True: # all modes modes = eq.surface.Z_basis.modes idx = np.arange(eq.surface.Z_basis.num_modes) + modes_idx = idx else: # specified modes modes = np.atleast_2d(self._modes) dtype = { @@ -694,23 +706,12 @@ def build(self, eq=None, use_jit=False, verbose=1): ) self._dim_f = idx.size - if self._target_from_user is not None: - if self._modes is True or self._modes is False: - raise RuntimeError( - "Attempting to provide target for Z boundary modes without " - + "providing modes array!" - + "You must pass in the modes corresponding to the" - + "provided target" - ) - # rearrange given target to match modes order - self.target = self._target_from_user[modes_idx] - # Zb_lmn -> Zb optimization space self._A = np.eye(eq.surface.Z_basis.num_modes)[idx, :] - # use surface parameters as target if needed - if self._target_from_user is None: - self.target = eq.surface.Z_lmn[idx] + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.surface.Z_lmn[idx], None, modes_idx + ) if self._normalize: scales = compute_scaling_factors(eq) @@ -954,7 +955,7 @@ def __init__( ): self._modes = modes - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -984,9 +985,11 @@ def build(self, eq=None, use_jit=False, verbose=1): if self._modes is False or self._modes is None: # no modes modes = np.array([[]], dtype=int) idx = np.array([], dtype=int) + modes_idx = idx elif self._modes is True: # all modes modes = eq.axis.R_basis.modes idx = np.arange(eq.axis.R_basis.num_modes) + modes_idx = idx else: # specified modes modes = np.atleast_1d(self._modes) dtype = { @@ -1013,23 +1016,12 @@ def build(self, eq=None, use_jit=False, verbose=1): ) self._dim_f = idx.size - if self._target_from_user is not None: - if self._modes is True or self._modes is False: - raise RuntimeError( - "Attempting to provide target for R axis modes without " - + "providing modes array!" - + "You must pass in the modes corresponding to the" - + "provided target" - ) - # rearrange given target to match modes order - self.target = self._target_from_user[modes_idx] - # Ra_lmn -> Ra optimization space self._A = np.eye(eq.axis.R_basis.num_modes)[idx, :] - # use surface parameters as target if needed - if self._target_from_user is None: - self.target = eq.axis.R_n[idx] + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.axis.R_n[idx], None, modes_idx + ) if self._normalize: scales = compute_scaling_factors(eq) @@ -1105,7 +1097,7 @@ def __init__( ): self._modes = modes - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -1135,9 +1127,11 @@ def build(self, eq=None, use_jit=False, verbose=1): if self._modes is False or self._modes is None: # no modes modes = np.array([[]], dtype=int) idx = np.array([], dtype=int) + modes_idx = idx elif self._modes is True: # all modes modes = eq.axis.Z_basis.modes idx = np.arange(eq.axis.Z_basis.num_modes) + modes_idx = idx else: # specified modes modes = np.atleast_1d(self._modes) dtype = { @@ -1164,23 +1158,12 @@ def build(self, eq=None, use_jit=False, verbose=1): ) self._dim_f = idx.size - if self._target_from_user is not None: - if self._modes is True or self._modes is False: - raise RuntimeError( - "Attempting to provide target for Z axis modes without " - + "providing modes array!" - + "You must pass in the modes corresponding to the" - + "provided target" - ) - # rearrange given target to match modes order - self.target = self._target_from_user[modes_idx] - - # Ra_lmn -> Ra optimization space + # Za_lmn -> Za optimization space self._A = np.eye(eq.axis.Z_basis.num_modes)[idx, :] - # use surface parameters as target if needed - if self._target_from_user is None: - self.target = eq.axis.Z_n[idx] + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.axis.Z_n[idx], None, modes_idx + ) if self._normalize: scales = compute_scaling_factors(eq) @@ -1262,7 +1245,7 @@ def __init__( raise ValueError( f"modes kwarg must be specified or True with FixModeR! got {modes}" ) - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -1315,18 +1298,9 @@ def build(self, eq=None, use_jit=False, verbose=1): self._dim_f = modes_idx.size - # use current eq's coefficients as target if needed - if self._target_from_user is None: - self.target = eq.R_lmn[self._idx] - else: # rearrange given target to match modes order - if self._modes is True or self._modes is False: - raise RuntimeError( - "Attempting to provide target for R fixed modes without " - + "providing modes array!" - + "You must pass in the modes corresponding to the" - + "provided target modes" - ) - self.target = self._target_from_user[modes_idx] + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.R_lmn[idx], None, modes_idx + ) super().build(things=eq, use_jit=use_jit, verbose=verbose) @@ -1404,7 +1378,7 @@ def __init__( raise ValueError( f"modes kwarg must be specified or True with FixModeZ! got {modes}" ) - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -1457,18 +1431,9 @@ def build(self, eq=None, use_jit=False, verbose=1): self._dim_f = modes_idx.size - # use current eq's coefficients as target if needed - if self._target_from_user is None: - self.target = eq.Z_lmn[self._idx] - else: # rearrange given target to match modes order - if self._modes is True or self._modes is False: - raise RuntimeError( - "Attempting to provide target for Z fixed modes without " - + "providing modes array!" - + "You must pass in the modes corresponding to the" - + "provided target modes" - ) - self.target = self._target_from_user[modes_idx] + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.Z_lmn[idx], None, modes_idx + ) super().build(things=eq, use_jit=use_jit, verbose=verbose) @@ -1549,20 +1514,28 @@ def __init__( name="Fix Sum Modes R", ): + errorif( + modes is None or modes is False, + ValueError, + f"modes kwarg must be specified or True with FixSumModesR! got {modes}", + ) + errorif( + target is not None and np.asarray(target).size > 1, + ValueError, + "FixSumModesR only accepts 1 target value, please use multiple" + + " FixSumModesR objectives if you wish to have multiple" + + " sets of constrained mode sums!", + ) + errorif( + bounds is not None and np.asarray(bounds)[0].size > 1, + ValueError, + "FixSumModesR only accepts 1 target value, please use multiple" + + " FixSumModesR objectives if you wish to have multiple" + + " sets of constrained mode sums!", + ) self._modes = modes - if modes is None or modes is False: - raise ValueError( - f"modes kwarg must be specified or True with FixSumModesR! got {modes}" - ) self._sum_weights = sum_weights - if target is not None: - if target.size > 1: - raise ValueError( - "FixSumModesR only accepts 1 target value, please use multiple" - + " FixSumModesR objectives if you wish to have multiple" - + " sets of constrained mode sums!" - ) - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -1629,9 +1602,12 @@ def build(self, eq=None, use_jit=False, verbose=1): j = eq.R_basis.get_idx(L=l, M=m, N=n) self._A[0, j] = sum_weights[i] - # use current sum as target if needed - if self._target_from_user is None: - self.target = np.dot(sum_weights.T, eq.R_lmn[self._idx]) + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, + np.dot(sum_weights.T, eq.R_lmn[self._idx]), + None, + np.array([0]), + ) super().build(things=eq, use_jit=use_jit, verbose=verbose) @@ -1711,21 +1687,28 @@ def __init__( modes=True, name="Fix Sum Modes Z", ): - + errorif( + modes is None or modes is False, + ValueError, + f"modes kwarg must be specified or True with FixSumModesZ! got {modes}", + ) + errorif( + target is not None and np.asarray(target).size > 1, + ValueError, + "FixSumModesZ only accepts 1 target value, please use multiple" + + " FixSumModesZ objectives if you wish to have multiple" + + " sets of constrained mode sums!", + ) + errorif( + bounds is not None and np.asarray(bounds)[0].size > 1, + ValueError, + "FixSumModesZ only accepts 1 target value, please use multiple" + + " FixSumModesZ objectives if you wish to have multiple" + + " sets of constrained mode sums!", + ) self._modes = modes - if modes is None or modes is False: - raise ValueError( - f"modes kwarg must be specified or True with FixSumModesZ! got {modes}" - ) self._sum_weights = sum_weights - if target is not None: - if target.size > 1: - raise ValueError( - "FixSumModesZ only accepts 1 target value, please use multiple" - + " FixSumModesZ objectives if you wish to have multiple sets of" - + " constrained mode sums!" - ) - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -1793,11 +1776,13 @@ def build(self, eq=None, use_jit=False, verbose=1): j = eq.Z_basis.get_idx(L=l, M=m, N=n) self._A[0, j] = sum_weights[i] - # use current sum as target if needed - if self._target_from_user is None: - self.target = np.dot(sum_weights.T, eq.Z_lmn[self._idx]) - - super().build(things=eq, use_jit=use_jit, verbose=verbose) + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, + np.dot(sum_weights.T, eq.Z_lmn[self._idx]), + None, + np.array([0]), + ) + super().build(eq=eq, use_jit=use_jit, verbose=verbose) def compute(self, params, constants=None): """Compute Sum mode Z errors. @@ -1874,7 +1859,7 @@ def __init__( self._profile = profile self._indices = indices - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -1914,9 +1899,10 @@ def build(self, eq=None, profile=None, use_jit=False, verbose=1): self._idx = np.atleast_1d(self._indices) self._dim_f = self._idx.size - # use profile parameters as target if needed - if self._target_from_user is None: - self.target = self._profile.params[self._idx] + + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, self._profile.params[self._idx], None, self._idx + ) super().build(things=eq, use_jit=use_jit, verbose=verbose) @@ -2727,7 +2713,7 @@ def __init__( normalize_target=True, name="fixed-Psi", ): - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -2755,8 +2741,9 @@ def build(self, eq=None, use_jit=False, verbose=1): eq = self.things[0] self._dim_f = 1 - if self._target_from_user is None: - self.target = eq.Psi + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.Psi, None, np.array([0]) + ) if self._normalize: scales = compute_scaling_factors(eq) diff --git a/desc/objectives/objective_funs.py b/desc/objectives/objective_funs.py index 54635a876c..158594d48f 100644 --- a/desc/objectives/objective_funs.py +++ b/desc/objectives/objective_funs.py @@ -919,7 +919,7 @@ def target(self): @target.setter def target(self, target): - self._target = np.atleast_1d(target) + self._target = np.atleast_1d(target) if target is not None else target self._check_dimensions() @property From 3ce2fb0fc51fd35a07f3965eeb7bb6ce7db2660c Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 26 Aug 2023 00:12:00 -0400 Subject: [PATCH 2/5] Update constraint parsing to allow linear inequality constraints --- desc/optimize/optimizer.py | 8 ++++++-- tests/test_optimizer.py | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/desc/optimize/optimizer.py b/desc/optimize/optimizer.py index 23e418d7ac..978f9ecaaa 100644 --- a/desc/optimize/optimizer.py +++ b/desc/optimize/optimizer.py @@ -315,11 +315,15 @@ def _parse_constraints(constraints): """ if not isinstance(constraints, (tuple, list)): constraints = (constraints,) + # we treat linear bound constraints as nonlinear since they can't be easily + # factorized like linear equality constraints linear_constraints = tuple( - constraint for constraint in constraints if constraint.linear + constraint + for constraint in constraints + if (constraint.linear and (constraint.bounds is None)) ) nonlinear_constraints = tuple( - constraint for constraint in constraints if not constraint.linear + constraint for constraint in constraints if constraint not in linear_constraints ) # check for incompatible constraints if any(isinstance(lc, FixCurrent) for lc in linear_constraints) and any( diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 4e2b31e246..b10c740b77 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -861,7 +861,7 @@ def test_constrained_AL_lsq(): FixBoundaryZ(eq=eq), # fix Z shape but not R FixPressure(eq=eq), # fix pressure profile FixIota(eq=eq), # fix rotational transform profile - FixPsi(eq=eq), # fix total toroidal magnetic flux + FixPsi(eq=eq, bounds=(eq.Psi * 0.99, eq.Psi * 1.01)), # linear inequality ) # some random constraints to keep the shape from getting wacky V = eq.compute("V")["V"] From cab8752cf20559753dd0e1f1810f7898e3959325 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Fri, 3 Nov 2023 22:33:39 -0600 Subject: [PATCH 3/5] Update expected error type --- desc/objectives/linear_objectives.py | 76 ++++++++++++++-------------- tests/test_linear_objectives.py | 44 ++++++++-------- 2 files changed, 60 insertions(+), 60 deletions(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 603fd406e4..33a21db466 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -1351,7 +1351,7 @@ def __init__( self._modes = modes if modes is None or modes is False: raise ValueError( - f"modes kwarg must be specified or True with FixModeR! got {modes}" + f"modes kwarg must be specified or True with FixModeR got {modes}" ) self._target_from_user = setdefault(bounds, target) super().__init__( @@ -1480,7 +1480,7 @@ def __init__( self._modes = modes if modes is None or modes is False: raise ValueError( - f"modes kwarg must be specified or True with FixModeZ! got {modes}" + f"modes kwarg must be specified or True with FixModeZ got {modes}" ) self._target_from_user = setdefault(bounds, target) super().__init__( @@ -1610,7 +1610,7 @@ def __init__( if modes is None or modes is False: raise ValueError( "modes kwarg must be specified" - + f" or True with FixModeLambda! got {modes}" + + f" or True with FixModeLambda got {modes}" ) self._target_from_user = target super().__init__( @@ -1662,18 +1662,9 @@ def build(self, use_jit=False, verbose=1): self._dim_f = modes_idx.size - # use current eq's coefficients as target if needed - if self._target_from_user is None: - self.target = eq.L_lmn[self._idx] - else: # rearrange given target to match modes order - if self._modes is True or self._modes is False: - raise RuntimeError( - "Attempting to provide target for lambda fixed modes without " - + "providing modes array!" - + "You must pass in the modes corresponding to the" - + "provided target modes" - ) - self.target = self._target_from_user[modes_idx] + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, eq.L_lmn[idx], None, modes_idx + ) super().build(use_jit=use_jit, verbose=verbose) @@ -1754,21 +1745,21 @@ def __init__( errorif( modes is None or modes is False, ValueError, - f"modes kwarg must be specified or True with FixSumModesR! got {modes}", + f"modes kwarg must be specified or True with FixSumModesR got {modes}", ) errorif( target is not None and np.asarray(target).size > 1, ValueError, "FixSumModesR only accepts 1 target value, please use multiple" + " FixSumModesR objectives if you wish to have multiple" - + " sets of constrained mode sums!", + + " sets of constrained mode sums", ) errorif( bounds is not None and np.asarray(bounds)[0].size > 1, ValueError, "FixSumModesR only accepts 1 target value, please use multiple" + " FixSumModesR objectives if you wish to have multiple" - + " sets of constrained mode sums!", + + " sets of constrained mode sums", ) self._modes = modes self._sum_weights = sum_weights @@ -1921,21 +1912,21 @@ def __init__( errorif( modes is None or modes is False, ValueError, - f"modes kwarg must be specified or True with FixSumModesZ! got {modes}", + f"modes kwarg must be specified or True with FixSumModesZ got {modes}", ) errorif( target is not None and np.asarray(target).size > 1, ValueError, "FixSumModesZ only accepts 1 target value, please use multiple" + " FixSumModesZ objectives if you wish to have multiple" - + " sets of constrained mode sums!", + + " sets of constrained mode sums", ) errorif( bounds is not None and np.asarray(bounds)[0].size > 1, ValueError, "FixSumModesZ only accepts 1 target value, please use multiple" + " FixSumModesZ objectives if you wish to have multiple" - + " sets of constrained mode sums!", + + " sets of constrained mode sums", ) self._modes = modes self._sum_weights = sum_weights @@ -2088,22 +2079,28 @@ def __init__( modes=True, name="Fix Sum Modes lambda", ): + errorif( + modes is None or modes is False, + ValueError, + f"modes kwarg must be specified or True with FixSumModesLambda got {modes}", + ) + errorif( + target is not None and np.asarray(target).size > 1, + ValueError, + "FixSumModesLambda only accepts 1 target value, please use multiple" + + " FixSumModesLambda objectives if you wish to have multiple" + + " sets of constrained mode sums", + ) + errorif( + bounds is not None and np.asarray(bounds)[0].size > 1, + ValueError, + "FixSumModesLambda only accepts 1 target value, please use multiple" + + " FixSumModesLambda objectives if you wish to have multiple" + + " sets of constrained mode sums", + ) self._modes = modes - if modes is None or modes is False: - raise ValueError( - "modes kwarg must be specified or True with" - + f" FixSumModesLambda! got {modes}" - ) self._sum_weights = sum_weights - if target is not None: - if target.size > 1: - raise ValueError( - "FixSumModesLambda only accepts 1 target value, please use" - + " multiple FixSumModesLambda objectives if you wish" - + " to have multiple sets of" - + " constrained mode sums!" - ) - self._target_from_user = target + self._target_from_user = setdefault(bounds, target) super().__init__( things=eq, target=target, @@ -2168,9 +2165,12 @@ def build(self, use_jit=False, verbose=1): j = eq.L_basis.get_idx(L=l, M=m, N=n) self._A[0, j] = sum_weights[i] - # use current sum as target if needed - if self._target_from_user is None: - self.target = np.dot(sum_weights.T, eq.L_lmn[self._idx]) + self.target, self.bounds = self._parse_target_from_user( + self._target_from_user, + np.dot(sum_weights.T, eq.L_lmn[self._idx]), + None, + np.array([0]), + ) super().build(use_jit=use_jit, verbose=verbose) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index a1302335d7..2890f8cca8 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -713,17 +713,17 @@ def test_FixBoundary_with_single_weight(): def test_FixBoundary_passed_target_no_passed_modes_error(): """Test Fixing boundary with no passed-in modes.""" eq = Equilibrium() - FixZ = FixBoundaryZ(eq=eq, modes=True, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixZ = FixBoundaryZ(eq=eq, modes=True, target=np.array([0, 0])) + with pytest.raises(ValueError): FixZ.build() - FixZ = FixBoundaryZ(eq=eq, modes=False, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixZ = FixBoundaryZ(eq=eq, modes=False, target=np.array([0, 0])) + with pytest.raises(ValueError): FixZ.build() - FixR = FixBoundaryR(eq=eq, modes=True, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixR = FixBoundaryR(eq=eq, modes=True, target=np.array([0, 0])) + with pytest.raises(ValueError): FixR.build() - FixR = FixBoundaryR(eq=eq, modes=False, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixR = FixBoundaryR(eq=eq, modes=False, target=np.array([0, 0])) + with pytest.raises(ValueError): FixR.build() @@ -731,17 +731,17 @@ def test_FixBoundary_passed_target_no_passed_modes_error(): def test_FixAxis_passed_target_no_passed_modes_error(): """Test Fixing Axis with no passed-in modes.""" eq = Equilibrium() - FixZ = FixAxisZ(eq=eq, modes=True, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixZ = FixAxisZ(eq=eq, modes=True, target=np.array([0, 0])) + with pytest.raises(ValueError): FixZ.build() - FixZ = FixAxisZ(eq=eq, modes=False, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixZ = FixAxisZ(eq=eq, modes=False, target=np.array([0, 0])) + with pytest.raises(ValueError): FixZ.build() - FixR = FixAxisR(eq=eq, modes=True, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixR = FixAxisR(eq=eq, modes=True, target=np.array([0, 0])) + with pytest.raises(ValueError): FixR.build() - FixR = FixAxisR(eq=eq, modes=False, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixR = FixAxisR(eq=eq, modes=False, target=np.array([0, 0])) + with pytest.raises(ValueError): FixR.build() @@ -749,14 +749,14 @@ def test_FixAxis_passed_target_no_passed_modes_error(): def test_FixMode_passed_target_no_passed_modes_error(): """Test Fixing Modes with no passed-in modes.""" eq = Equilibrium() - FixZ = FixModeZ(eq=eq, modes=True, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixZ = FixModeZ(eq=eq, modes=True, target=np.array([0, 0])) + with pytest.raises(ValueError): FixZ.build() - FixR = FixModeR(eq=eq, modes=True, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixR = FixModeR(eq=eq, modes=True, target=np.array([0, 0])) + with pytest.raises(ValueError): FixR.build(eq) - FixL = FixModeLambda(eq=eq, modes=True, target=np.array([[0]])) - with pytest.raises(RuntimeError): + FixL = FixModeLambda(eq=eq, modes=True, target=np.array([0, 0])) + with pytest.raises(ValueError): FixL.build(eq) From 776dbf2ccea47fa0b62a52f79839044013a5afd4 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 9 Nov 2023 18:22:54 -0500 Subject: [PATCH 4/5] Update bounds parsing for new objectives --- desc/objectives/linear_objectives.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/desc/objectives/linear_objectives.py b/desc/objectives/linear_objectives.py index 33a21db466..9e45d4bd7f 100644 --- a/desc/objectives/linear_objectives.py +++ b/desc/objectives/linear_objectives.py @@ -59,7 +59,7 @@ def _parse_target_from_user( # TODO: make this work with above, but for multiple target args? -class FixParameter(_Objective): +class FixParameter(_FixedObjective): """Fix specific degrees of freedom associated with a given Optimizable object. Parameters @@ -167,14 +167,24 @@ def build(self, use_jit=False, verbose=1): idx = np.arange(thing.dimensions[par]) indices[par] = np.atleast_1d(idx) self._indices = indices - - target = setdefault( - self._target_from_user, - {par: thing.params_dict[par][self._indices[par]] for par in params}, - ) self._dim_f = sum(t.size for t in self._indices.values()) - self.target = jnp.concatenate([target[par] for par in params]) + default_target = { + par: thing.params_dict[par][self._indices[par]] for par in params + } + default_bounds = None + target, bounds = self._parse_target_from_user( + self._target_from_user, default_target, default_bounds, indices + ) + if target: + self.target = jnp.concatenate([target[par] for par in params]) + self.bounds = None + else: + self.target = None + self.bounds = ( + jnp.concatenate([bounds[0][par] for par in params]), + jnp.concatenate([bounds[1][par] for par in params]), + ) super().build(use_jit=use_jit, verbose=verbose) def compute(self, params, constants=None): From cfe1f97e9957a0e0783645ba9906af34a206ec26 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Tue, 14 Nov 2023 22:30:39 -0500 Subject: [PATCH 5/5] Add test for FixParameter with bounds --- tests/test_optimizer.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 36aafb4175..ba6b1421cb 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -24,6 +24,7 @@ FixBoundaryZ, FixCurrent, FixIota, + FixParameter, FixPressure, FixPsi, ForceBalance, @@ -863,8 +864,10 @@ def test_constrained_AL_lsq(): constraints = ( FixBoundaryR(eq=eq, modes=[0, 0, 0]), # fix specified major axis position - FixPressure(eq=eq), # fix pressure profile - FixIota(eq=eq), # fix rotational transform profile + FixPressure(eq=eq), + FixParameter( + eq, "i_l", bounds=(eq.i_l * 0.9, eq.i_l * 1.1) + ), # linear inequality FixPsi(eq=eq, bounds=(eq.Psi * 0.99, eq.Psi * 1.01)), # linear inequality ) # some random constraints to keep the shape from getting wacky @@ -900,6 +903,9 @@ def test_constrained_AL_lsq(): Dwell = constraints[-2].compute_scaled(*constraints[-2].xs(eq2)) assert (ARbounds[0] - ctol) < AR2 < (ARbounds[1] + ctol) assert (Vbounds[0] - ctol) < V2 < (Vbounds[1] + ctol) + assert (0.99 * eq.Psi - ctol) <= eq2.Psi <= (1.01 * eq.Psi + ctol) + np.testing.assert_array_less((0.9 * eq.i_l - ctol), eq2.i_l) + np.testing.assert_array_less(eq2.i_l, (1.1 * eq.i_l + ctol)) assert eq2.is_nested() np.testing.assert_array_less(-Dwell, ctol)