diff --git a/CHANGELOG.md b/CHANGELOG.md index f609e6ad8c..5c03516b4d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,11 @@ Changelog ========= +New Features + +- Added option for ``eq_fixed`` added to ``BoundaryError`` to remove equilibrium as a degree of freedom from boundary error objective. + + Performance Improvements - Speeds up the ``"qr"`` trust-region subproblem and Newton-step solves in the least-squares optimizers by reusing the Jacobian QR factorization across the Levenberg-Marquardt parameter sweep. On ``jax >= 0.10.0`` this uses ``qr_multiply`` to additionally avoid forming ``Q`` explicitly; on older versions a fallback preserves the same results. diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index b62603289d..e99d323c9c 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -447,6 +447,7 @@ class BoundaryError(_Objective): "_B_plasma_chunk_size", "_bs_chunk_size", "_eq_data_keys", + "_eq_fixed", "_field_fixed", "_q", "_sheet_current", @@ -478,6 +479,7 @@ def __init__( eval_grid=None, field_grid=None, field_fixed=False, + eq_fixed=False, name="Boundary error", jac_chunk_size=None, *, @@ -487,12 +489,20 @@ def __init__( ): if target is None and bounds is None: target = 0 + self._eq = eq self._source_grid = source_grid self._eval_grid = eval_grid self._st, self._sz = s if isinstance(s, (tuple, list)) else (s, s) self._q = q self._field = [field] if not isinstance(field, list) else field self._field_grid = field_grid + errorif( + field_fixed and eq_fixed, + ValueError, + "At least one of eq_fixed or field_fixed must be false.", + ) + self._field_fixed = field_fixed + self._eq_fixed = eq_fixed self._bs_chunk_size = bs_chunk_size B_plasma_chunk_size = parse_argname_change( B_plasma_chunk_size, kwargs, "loop", "B_plasma_chunk_size" @@ -501,7 +511,9 @@ def __init__( B_plasma_chunk_size = None self._B_plasma_chunk_size = B_plasma_chunk_size self._sheet_current = hasattr(eq.surface, "Phi_mn") - things = [eq] + things = [] + if not eq_fixed: + things.append(self._eq) if not field_fixed: things.append(self._field) super().__init__( @@ -530,7 +542,7 @@ def build(self, use_jit=True, verbose=1): """ from desc.magnetic_fields import SumMagneticField - eq = self.things[0] + eq = self._eq if self._source_grid is None: # for axisymmetry we still need to know about toroidal effects, so its @@ -651,6 +663,74 @@ def build(self, use_jit=True, verbose=1): ) ) + if self._eq_fixed: + source_data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", + self._eq_data_keys, + params=self._eq.params_dict, + transforms=source_transforms, + profiles=source_profiles, + ) + eval_data = ( + source_data + if self._use_same_grid + else compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", + self._eq_data_keys, + params=self._eq.params_dict, + transforms=eval_transforms, + profiles=eval_profiles, + ) + ) + + Bplasma = virtual_casing_biot_savart( + eval_data, + source_data, + interpolator, + chunk_size=self._B_plasma_chunk_size, + ) + # need extra factor of B/2 bc we're evaluating on plasma surface + Bplasma = Bplasma + eval_data["B"] / 2 + + self._constants["source_data"] = source_data + self._constants["eval_data"] = eval_data + self._constants["Bplasma"] = Bplasma + + # sheet current stuff + if self._sheet_current: + p = ( + "desc.magnetic_fields._current_potential." + "FourierCurrentPotentialField" + ) + sheet_params = { + "R_lmn": self._eq.params_dict["Rb_lmn"], + "Z_lmn": self._eq.params_dict["Zb_lmn"], + "I": self._eq.params_dict["I"], + "G": self._eq.params_dict["G"], + "Phi_mn": self._eq.params_dict["Phi_mn"], + } + sheet_source_data = compute_fun( + p, + self._sheet_data_keys, + params=sheet_params, + transforms=self._constants["sheet_source_transforms"], + profiles={}, + ) + sheet_eval_data = ( + sheet_source_data + if self._use_same_grid + else compute_fun( + p, + self._sheet_data_keys, + params=sheet_params, + transforms=self._constants["sheet_eval_transforms"], + profiles={}, + ) + ) + + self._constants["sheet_eval_data"] = sheet_eval_data + source_data["K_vc"] += sheet_source_data["K"] + timer.stop("Precomputing transforms") if verbose > 1: timer.disp("Precomputing transforms") @@ -674,15 +754,17 @@ def build(self, use_jit=True, verbose=1): super().build(use_jit=use_jit, verbose=verbose) - def compute(self, eq_params, *field_params, constants=None): + def compute(self, *params, constants=None): """Compute boundary force error. Parameters ---------- - eq_params : dict - Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict - field_params : dict - Dictionary of field parameters, if field is not fixed. + params : dict + 1 or more dictionaries of params assigned in order of self.things, + respecting which degrees of freedom are fixed (via `eq_fixed` + and `field_fixed`). If eq_fixed, `params` are field_params, if + neither are fixed, `params[0]` takes equilibrium params, `params[1:]` + are field params)./ constants : dict Dictionary of constant data, eg transforms, profiles etc. Defaults to self.constants. (Deprecated) @@ -695,64 +777,84 @@ def compute(self, eq_params, *field_params, constants=None): √g||μ₀𝐊 − 𝐧 × [𝐁]|| in T*m^2 """ - if field_params == (): # common case for field_fixed=True + if self._eq_fixed: + field_params = params + eq_params = None + elif self._field_fixed: + eq_params = params[0] field_params = None + else: + eq_params = params[0] + field_params = params[1:] constants = self._get_deprecated_constants(constants) - source_data = compute_fun( - "desc.equilibrium.equilibrium.Equilibrium", - self._eq_data_keys, - params=eq_params, - transforms=constants["source_transforms"], - profiles=constants["source_profiles"], - ) - eval_data = ( - source_data - if self._use_same_grid - else compute_fun( + + if self._eq_fixed: + source_data = constants["source_data"] + eval_data = constants["eval_data"] + Bplasma = constants["Bplasma"] + if self._sheet_current: + sheet_eval_data = constants["sheet_eval_data"] + else: + source_data = compute_fun( "desc.equilibrium.equilibrium.Equilibrium", self._eq_data_keys, params=eq_params, - transforms=constants["eval_transforms"], - profiles=constants["eval_profiles"], + transforms=constants["source_transforms"], + profiles=constants["source_profiles"], ) - ) - if self._sheet_current: - p = "desc.magnetic_fields._current_potential.FourierCurrentPotentialField" - sheet_params = { - "R_lmn": eq_params["Rb_lmn"], - "Z_lmn": eq_params["Zb_lmn"], - "I": eq_params["I"], - "G": eq_params["G"], - "Phi_mn": eq_params["Phi_mn"], - } - sheet_source_data = compute_fun( - p, - self._sheet_data_keys, - params=sheet_params, - transforms=constants["sheet_source_transforms"], - profiles={}, - ) - sheet_eval_data = ( - sheet_source_data + eval_data = ( + source_data if self._use_same_grid else compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", + self._eq_data_keys, + params=eq_params, + transforms=constants["eval_transforms"], + profiles=constants["eval_profiles"], + ) + ) + + Bplasma = virtual_casing_biot_savart( + eval_data, + source_data, + constants["interpolator"], + chunk_size=self._B_plasma_chunk_size, + ) + # need extra factor of B/2 bc we're evaluating on plasma surface + Bplasma = Bplasma + eval_data["B"] / 2 + + if self._sheet_current: + p = ( + "desc.magnetic_fields._current_potential." + "FourierCurrentPotentialField" + ) + sheet_params = { + "R_lmn": eq_params["Rb_lmn"], + "Z_lmn": eq_params["Zb_lmn"], + "I": eq_params["I"], + "G": eq_params["G"], + "Phi_mn": eq_params["Phi_mn"], + } + sheet_source_data = compute_fun( p, self._sheet_data_keys, params=sheet_params, - transforms=constants["sheet_eval_transforms"], + transforms=constants["sheet_source_transforms"], profiles={}, ) - ) - source_data["K_vc"] += sheet_source_data["K"] + sheet_eval_data = ( + sheet_source_data + if self._use_same_grid + else compute_fun( + p, + self._sheet_data_keys, + params=sheet_params, + transforms=constants["sheet_eval_transforms"], + profiles={}, + ) + ) + source_data["K_vc"] += sheet_source_data["K"] - Bplasma = virtual_casing_biot_savart( - eval_data, - source_data, - constants["interpolator"], - chunk_size=self._B_plasma_chunk_size, - ) - # need extra factor of B/2 bc we're evaluating on plasma surface - Bplasma = Bplasma + eval_data["B"] / 2 x = jnp.array([eval_data["R"], eval_data["phi"], eval_data["Z"]]).T # can always pass in field params. If they're None, it just uses the # defaults for the given field. diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 5403aef6ba..d4571bac9d 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -610,6 +610,69 @@ def test_boundary_error_biest(self): # next n should be B^2 errors np.testing.assert_allclose(f[n : 2 * n], 0, atol=5e-2) + @pytest.mark.unit + def test_boundary_error_things_fixed(self): + """Test that BoundaryError runs correctly for eq_fixed/field_fixed combos.""" + eq = Equilibrium(L=3, M=3, N=3, Psi=np.pi) + eq.solve() + coil = FourierXYZCoil(5e5) + coilset = CoilSet.linspaced_angular(coil, n=100, check_intersection=False) + field = [coilset, ToroidalMagneticField(B0=0, R0=1)] + + def test(eq_fixed=False, field_fixed=False): + obj = BoundaryError( + eq, + field, + eq_fixed=eq_fixed, + field_fixed=field_fixed, + ) + obj.build() + f = obj.compute_scaled_error(*obj.xs()) + n = len(f) // 2 + # first n should be B*n errors + np.testing.assert_allclose(f[:n], 0, atol=1e-4) + # next n should be B^2 errors + np.testing.assert_allclose(f[n : 2 * n], 0, atol=5e-2) + + test(eq_fixed=False, field_fixed=False) + test(eq_fixed=True) + test(field_fixed=True) + + with pytest.raises(ValueError, match="At least one"): + BoundaryError(eq, field, eq_fixed=True, field_fixed=True) + + @pytest.mark.unit + def test_boundary_error_things_fixed_sheet_current(self): + """Test BoundaryError for eq_fixed/field_fixed combos, with sheet currents.""" + coil = FourierXYZCoil(5e5) + coilset = CoilSet.linspaced_angular(coil, n=100, check_intersection=False) + field = [coilset, ToroidalMagneticField(B0=0, R0=1)] + eq = Equilibrium(L=3, M=3, N=3, Psi=np.pi) + eq.surface = FourierCurrentPotentialField.from_surface( + eq.surface, M_Phi=eq.M, N_Phi=eq.N + ) + eq.solve() + + def test(eq_fixed=False, field_fixed=False): + obj = BoundaryError( + eq, + field, + eq_fixed=eq_fixed, + field_fixed=field_fixed, + ) + obj.build() + f = obj.compute_scaled_error(*obj.xs()) + n = len(f) // 3 + # first n should be B*n errors + np.testing.assert_allclose(f[:n], 0, atol=1e-4) + # next n should be B^2 errors + np.testing.assert_allclose(f[n : 2 * n], 0, atol=5e-2) + # last n should be K errors + np.testing.assert_allclose(f[2 * n :], 0, atol=3e-2) + + with pytest.raises(ValueError, match="At least one"): + BoundaryError(eq, field, eq_fixed=True, field_fixed=True) + @pytest.mark.unit def test_boundary_error_vacuum(self): """Test calculation of vacuum boundary error."""