From bbc27793f1a95854439864abff76d99a15ba7511 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Fri, 3 Nov 2023 20:40:10 -0600 Subject: [PATCH 01/13] Remove unneeded args from objective getters - When using stuff like ``get_fixed_boundary_constraints`` we have to specify both whether to fix profiles, and also which profiles the equilibrium has. - Now that ``eq`` is a required argument to the get_* functions, this now redundant, since we can just look and see what profiles the equilibrium has. - This removes the unneeded args from all the get_* methods, so that the only things needed are the equilibrium and a bool for profiles, the specific profiles will be determined from the equilibrium. --- desc/continuation.py | 36 ++----- desc/equilibrium/equilibrium.py | 24 +---- desc/objectives/getters.py | 142 +++++--------------------- desc/optimize/_constraint_wrappers.py | 6 +- desc/perturbations.py | 4 +- desc/vmec.py | 2 +- tests/test_bootstrap.py | 4 +- tests/test_examples.py | 10 +- tests/test_linear_objectives.py | 12 ++- tests/test_perturbations.py | 2 +- tests/test_profiles.py | 4 +- 11 files changed, 54 insertions(+), 192 deletions(-) diff --git a/desc/continuation.py b/desc/continuation.py index 76438fd836..eb5e092a53 100644 --- a/desc/continuation.py +++ b/desc/continuation.py @@ -98,11 +98,7 @@ def _solve_axisym( deltas = get_deltas({"surface": surf_i}, {"surface": surf_i2}) surf_i = surf_i2 - constraints_i = get_fixed_boundary_constraints( - eq=eqi, - iota=eqi.iota, - kinetic=eqi.electron_temperature, - ) + constraints_i = get_fixed_boundary_constraints(eq=eqi) objective_i = get_equilibrium_objective(eq=eqi, mode=objective) if verbose: @@ -227,11 +223,7 @@ def _add_pressure( deltas["p_l"] *= pres_step pres_ratio += pres_step - constraints_i = get_fixed_boundary_constraints( - eq=eqi, - iota=eqi.iota, - kinetic=eqi.electron_temperature, - ) + constraints_i = get_fixed_boundary_constraints(eq=eqi) objective_i = get_equilibrium_objective(eq=eqi, mode=objective) if verbose: @@ -360,11 +352,7 @@ def _add_shaping( deltas["Zb_lmn"] *= bdry_step bdry_ratio += bdry_step - constraints_i = get_fixed_boundary_constraints( - eq=eqi, - iota=eqi.iota, - kinetic=eqi.electron_temperature, - ) + constraints_i = get_fixed_boundary_constraints(eq=eqi) objective_i = get_equilibrium_objective(eq=eqi, mode=objective) if verbose: @@ -662,11 +650,7 @@ def solve_continuation( # noqa: C901 if not isinstance(optimizer, Optimizer): optimizer = Optimizer(optimizer) objective_i = get_equilibrium_objective(eq=eqfam[0], mode=objective) - constraints_i = get_fixed_boundary_constraints( - eq=eqfam[0], - iota=objective != "vacuum" and eqfam[0].iota is not None, - kinetic=eqfam[0].electron_temperature is not None, - ) + constraints_i = get_fixed_boundary_constraints(eq=eqfam[0]) ii = 0 nn = len(eqfam) @@ -713,11 +697,7 @@ def solve_continuation( # noqa: C901 # TODO: pass Jx if available eqp = eqfam[ii - 1].copy() objective_i = get_equilibrium_objective(eq=eqp, mode=objective) - constraints_i = get_fixed_boundary_constraints( - eq=eqp, - iota=eqp.iota, - kinetic=eqp.electron_temperature, - ) + constraints_i = get_fixed_boundary_constraints(eq=eqp) eqp.change_resolution(**eqi.resolution) eqp.perturb( objective=objective_i, @@ -737,11 +717,7 @@ def solve_continuation( # noqa: C901 if not stop: # TODO: add ability to rebind objectives objective_i = get_equilibrium_objective(eq=eqi, mode=objective) - constraints_i = get_fixed_boundary_constraints( - eq=eqi, - iota=eqi.iota, - kinetic=eqi.electron_temperature, - ) + constraints_i = get_fixed_boundary_constraints(eq=eqi) eqi.solve( optimizer=optimizer, objective=objective_i, diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index 434effbb5f..06451f98be 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -1704,11 +1704,7 @@ def solve( """ if constraints is None: - constraints = get_fixed_boundary_constraints( - eq=self, - iota=objective != "vacuum" and self.iota is not None, - kinetic=self.electron_temperature is not None, - ) + constraints = get_fixed_boundary_constraints(eq=self) if not isinstance(objective, ObjectiveFunction): objective = get_equilibrium_objective(eq=self, mode=objective) if not isinstance(optimizer, Optimizer): @@ -1809,11 +1805,7 @@ def optimize( if not isinstance(optimizer, Optimizer): optimizer = Optimizer(optimizer) if constraints is None: - constraints = get_fixed_boundary_constraints( - eq=self, - iota=self.iota is not None, - kinetic=self.electron_temperature is not None, - ) + constraints = get_fixed_boundary_constraints(eq=self) constraints = (ForceBalance(eq=self), *constraints) if not isinstance(constraints, (list, tuple)): constraints = tuple([constraints]) @@ -2047,17 +2039,9 @@ def perturb( objective = get_equilibrium_objective(eq=self) if constraints is None: if "Ra_n" in deltas or "Za_n" in deltas: - constraints = get_fixed_axis_constraints( - eq=self, - iota=self.iota is not None, - kinetic=self.electron_temperature is not None, - ) + constraints = get_fixed_axis_constraints(eq=self) else: - constraints = get_fixed_boundary_constraints( - eq=self, - iota=self.iota is not None, - kinetic=self.electron_temperature is not None, - ) + constraints = get_fixed_boundary_constraints(eq=self) eq = perturb( self, diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py index a5218e27e2..4910687ac8 100644 --- a/desc/objectives/getters.py +++ b/desc/objectives/getters.py @@ -30,6 +30,17 @@ from .nae_utils import calc_zeroth_order_lambda, make_RZ_cons_1st_order from .objective_funs import ObjectiveFunction +_PROFILE_CONSTRAINTS = { + "pressure": FixPressure, + "iota": FixIota, + "current": FixCurrent, + "electron_density": FixElectronDensity, + "electron_temperature": FixElectronTemperature, + "ion_temperature": FixIonTemperature, + "atomic_number": FixAtomicNumber, + "anisotropy": FixAnisotropy, +} + def get_equilibrium_objective(eq, mode="force", normalize=True): """Get the objective function for a typical force balance equilibrium problem. @@ -70,9 +81,7 @@ def get_equilibrium_objective(eq, mode="force", normalize=True): return ObjectiveFunction(objectives) -def get_fixed_axis_constraints( - eq, profiles=True, iota=True, kinetic=False, anisotropy=False, normalize=True -): +def get_fixed_axis_constraints(eq, profiles=True, normalize=True): """Get the constraints necessary for a fixed-axis equilibrium problem. Parameters @@ -81,12 +90,6 @@ def get_fixed_axis_constraints( Equilibrium being constrained. profiles : bool Whether to also return constraints to fix input profiles. - iota : bool - Whether to add FixIota or FixCurrent as a constraint. - kinetic : bool - Whether to add constraints to fix kinetic profiles or pressure - anisotropy : bool - Whether to add constraint to fix anisotropic pressure. normalize : bool Whether to apply constraints in normalized units. @@ -102,44 +105,16 @@ def get_fixed_axis_constraints( FixPsi(eq=eq, normalize=normalize, normalize_target=normalize), ) if profiles: - if kinetic: - constraints += ( - FixElectronDensity( - eq=eq, normalize=normalize, normalize_target=normalize - ), - FixElectronTemperature( - eq=eq, normalize=normalize, normalize_target=normalize - ), - FixIonTemperature( - eq=eq, normalize=normalize, normalize_target=normalize - ), - FixAtomicNumber(eq=eq, normalize=normalize, normalize_target=normalize), - ) - else: - constraints += ( - FixPressure(eq=eq, normalize=normalize, normalize_target=normalize), - ) - if anisotropy: + for name, con in _PROFILE_CONSTRAINTS.items(): + if getattr(eq, name) is not None: constraints += ( - FixAnisotropy( - eq=eq, normalize=normalize, normalize_target=normalize - ), + con(eq=eq, normalize=normalize, normalize_target=normalize), ) - if iota: - constraints += ( - FixIota(eq=eq, normalize=normalize, normalize_target=normalize), - ) - else: - constraints += ( - FixCurrent(eq=eq, normalize=normalize, normalize_target=normalize), - ) return constraints -def get_fixed_boundary_constraints( - eq, profiles=True, iota=True, kinetic=False, anisotropy=False, normalize=True -): +def get_fixed_boundary_constraints(eq, profiles=True, normalize=True): """Get the constraints necessary for a typical fixed-boundary equilibrium problem. Parameters @@ -148,12 +123,6 @@ def get_fixed_boundary_constraints( Equilibrium to constraint. profiles : bool Whether to also return constraints to fix input profiles. - iota : bool - Whether to add FixIota or FixCurrent as a constraint. - kinetic : bool - Whether to also fix kinetic profiles. - anisotropy : bool - Whether to add constraint to fix anisotropic pressure. normalize : bool Whether to apply constraints in normalized units. @@ -169,37 +138,12 @@ def get_fixed_boundary_constraints( FixPsi(eq=eq, normalize=normalize, normalize_target=normalize), ) if profiles: - if kinetic: - constraints += ( - FixElectronDensity( - eq=eq, normalize=normalize, normalize_target=normalize - ), - FixElectronTemperature( - eq=eq, normalize=normalize, normalize_target=normalize - ), - FixIonTemperature( - eq=eq, normalize=normalize, normalize_target=normalize - ), - FixAtomicNumber(eq=eq, normalize=normalize, normalize_target=normalize), - ) - else: - constraints += ( - FixPressure(eq=eq, normalize=normalize, normalize_target=normalize), - ) - if anisotropy: + for name, con in _PROFILE_CONSTRAINTS.items(): + if getattr(eq, name) is not None: constraints += ( - FixAnisotropy( - eq=eq, normalize=normalize, normalize_target=normalize - ), + con(eq=eq, normalize=normalize, normalize_target=normalize), ) - if iota: - constraints += ( - FixIota(eq=eq, normalize=normalize, normalize_target=normalize), - ) - else: - constraints += ( - FixCurrent(eq=eq, normalize=normalize, normalize_target=normalize), - ) + return constraints @@ -208,9 +152,6 @@ def get_NAE_constraints( qsc_eq, order=1, profiles=True, - iota=False, - kinetic=False, - anisotropy=False, normalize=True, N=None, fix_lambda=False, @@ -228,12 +169,6 @@ def get_NAE_constraints( order (in rho) of near-axis behavior to constrain profiles : bool Whether to also return constraints to fix input profiles. - iota : bool - Whether to add `FixIota` or `FixCurrent` as a constraint. - kinetic : bool - Whether to also fix kinetic profiles. - anisotropy : bool - Whether to add constraint to fix anisotropic pressure. normalize : bool Whether to apply constraints in normalized units. N : int @@ -256,43 +191,14 @@ def get_NAE_constraints( FixAxisZ(eq=desc_eq, normalize=normalize, normalize_target=normalize), FixPsi(eq=desc_eq, normalize=normalize, normalize_target=normalize), ) + if profiles: - if kinetic: - constraints += ( - FixElectronDensity( - eq=desc_eq, normalize=normalize, normalize_target=normalize - ), - FixElectronTemperature( - eq=desc_eq, normalize=normalize, normalize_target=normalize - ), - FixIonTemperature( - eq=desc_eq, normalize=normalize, normalize_target=normalize - ), - FixAtomicNumber( - eq=desc_eq, normalize=normalize, normalize_target=normalize - ), - ) - else: - constraints += ( - FixPressure( - eq=desc_eq, normalize=normalize, normalize_target=normalize - ), - ) - if anisotropy: + for name, con in _PROFILE_CONSTRAINTS.items(): + if getattr(desc_eq, name) is not None: constraints += ( - FixAnisotropy( - eq=desc_eq, normalize=normalize, normalize_target=normalize - ), + con(eq=desc_eq, normalize=normalize, normalize_target=normalize), ) - if iota: - constraints += ( - FixIota(eq=desc_eq, normalize=normalize, normalize_target=normalize), - ) - else: - constraints += ( - FixCurrent(eq=desc_eq, normalize=normalize, normalize_target=normalize), - ) if fix_lambda or (fix_lambda >= 0 and type(fix_lambda) is int): L_axis_constraints, _, _ = calc_zeroth_order_lambda( qsc=qsc_eq, desc_eq=desc_eq, N=N diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py index 1b5f98c386..518d262bb1 100644 --- a/desc/optimize/_constraint_wrappers.py +++ b/desc/optimize/_constraint_wrappers.py @@ -473,11 +473,7 @@ def build(self, eq=None, use_jit=None, verbose=1): # noqa: C901 timer.start("Proximal projection build") self._eq = eq - self._linear_constraints = get_fixed_boundary_constraints( - eq=eq, - iota=self._eq.iota, - kinetic=self._eq.electron_temperature, - ) + self._linear_constraints = get_fixed_boundary_constraints(eq=eq) self._linear_constraints = maybe_add_self_consistency( self._eq, self._linear_constraints ) diff --git a/desc/perturbations.py b/desc/perturbations.py index 3d392fcd6e..3993855649 100644 --- a/desc/perturbations.py +++ b/desc/perturbations.py @@ -543,9 +543,7 @@ def optimal_perturb( # noqa: C901 - FIXME: break this up into simpler pieces print("Number of objectives: {}".format(objective_g.dim_f)) # FIXME: generalize to other constraints - constraints = get_fixed_boundary_constraints( - eq=eq, iota=eq.iota is not None, kinetic=eq.electron_temperature is not None - ) + constraints = get_fixed_boundary_constraints(eq=eq) constraints = maybe_add_self_consistency(eq=eq, constraints=constraints) for con in constraints: if not con.built: diff --git a/desc/vmec.py b/desc/vmec.py index 816e716483..ee2d31f1a0 100644 --- a/desc/vmec.py +++ b/desc/vmec.py @@ -175,7 +175,7 @@ def load( constraints = get_fixed_axis_constraints( profiles=False, eq=eq - ) + get_fixed_boundary_constraints(iota=eq.iota, eq=eq) + ) + get_fixed_boundary_constraints(eq=eq) constraints = maybe_add_self_consistency(eq, constraints) objective = ObjectiveFunction(constraints, verbose=0) objective.build() diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index 8ed9841ec1..f2d25e230b 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -1301,7 +1301,7 @@ def test_bootstrap_consistency_iota(self, TmpDir): eq.solve( verbose=3, ftol=1e-8, - constraints=get_fixed_boundary_constraints(eq=eq, kinetic=True), + constraints=get_fixed_boundary_constraints(eq=eq), optimizer=Optimizer("lsq-exact"), objective=ObjectiveFunction(objectives=ForceBalance(eq=eq)), ) @@ -1413,7 +1413,7 @@ def test_bootstrap_consistency_current(self, TmpDir): eq.solve( verbose=3, ftol=1e-8, - constraints=get_fixed_boundary_constraints(eq=eq, kinetic=True, iota=False), + constraints=get_fixed_boundary_constraints(eq=eq), optimizer=Optimizer("lsq-exact"), objective=ObjectiveFunction(objectives=ForceBalance(eq=eq)), ) diff --git a/tests/test_examples.py b/tests/test_examples.py index 9e29f47a60..767e3271e9 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -89,7 +89,7 @@ def test_SOLOVEV_anisotropic_results(SOLOVEV): eq.anisotropy = anisotropy obj = ObjectiveFunction(ForceBalanceAnisotropic(eq=eq)) - constraints = get_fixed_boundary_constraints(eq=eq, anisotropy=True) + constraints = get_fixed_boundary_constraints(eq=eq) eq.solve(obj, constraints, verbose=3) rho_err, theta_err = area_difference_vmec(eq, SOLOVEV["vmec_nc_path"]) @@ -548,12 +548,12 @@ def test_NAE_QSC_solve(): # this has all the constraints we need, # iota=False specifies we want to fix current instead of iota - cs = get_NAE_constraints(eq, qsc, iota=False, order=1, fix_lambda=False, N=eq.N) + cs = get_NAE_constraints(eq, qsc, order=1, fix_lambda=False, N=eq.N) cs_lambda_fixed_0th_order = get_NAE_constraints( - eq_lambda_fixed_0th_order, qsc, iota=False, order=1, fix_lambda=0, N=eq.N + eq_lambda_fixed_0th_order, qsc, order=1, fix_lambda=0, N=eq.N ) cs_lambda_fixed_1st_order = get_NAE_constraints( - eq_lambda_fixed_1st_order, qsc, iota=False, order=1, fix_lambda=True, N=eq.N + eq_lambda_fixed_1st_order, qsc, order=1, fix_lambda=True, N=eq.N ) for c in cs: @@ -636,7 +636,7 @@ def test_NAE_QIC_solve(): # this has all the constraints we need, # iota=False specifies we want to fix current instead of iota - cs = get_NAE_constraints(eq, qic, iota=False, order=1) + cs = get_NAE_constraints(eq, qic, order=1) objectives = ForceBalance(eq=eq) obj = ObjectiveFunction(objectives) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index a1302335d7..561c330fb6 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -312,7 +312,7 @@ def test_fixed_axis_and_theta_SFL_solve(): def test_factorize_linear_constraints_asserts(): """Test error checking for factorize_linear_constraints.""" eq = Equilibrium() - constraints = get_fixed_boundary_constraints(eq=eq, iota=False) + constraints = get_fixed_boundary_constraints(eq=eq) for con in constraints: con.build(verbose=0) constraints[3].bounds = (0, 1) # bounds on FixPsi @@ -824,14 +824,15 @@ def _is_any_instance(things, cls): def test_FixAxis_util_correct_objectives(): """Test util for fix axis constraints.""" eq = Equilibrium() - cs = get_fixed_axis_constraints(eq, iota=False) + cs = get_fixed_axis_constraints(eq) assert _is_any_instance(cs, FixAxisR) assert _is_any_instance(cs, FixAxisZ) assert _is_any_instance(cs, FixPsi) assert _is_any_instance(cs, FixPressure) assert _is_any_instance(cs, FixCurrent) - cs = get_fixed_axis_constraints(eq, iota=True, kinetic=True) + eq = Equilibrium(electron_temperature=1, electron_density=1, iota=1) + cs = get_fixed_axis_constraints(eq) assert _is_any_instance(cs, FixAxisR) assert _is_any_instance(cs, FixAxisZ) assert _is_any_instance(cs, FixPsi) @@ -847,7 +848,7 @@ def test_FixNAE_util_correct_objectives(): """Test util for fix NAE constraints.""" eq = Equilibrium() qsc = Qsc.from_paper("precise QA") - cs = get_NAE_constraints(eq, qsc, iota=False) + cs = get_NAE_constraints(eq, qsc) assert _is_any_instance(cs, FixAxisR) assert _is_any_instance(cs, FixAxisZ) assert _is_any_instance(cs, FixPsi) @@ -856,7 +857,8 @@ def test_FixNAE_util_correct_objectives(): assert _is_any_instance(cs, FixPressure) assert _is_any_instance(cs, FixCurrent) - cs = get_NAE_constraints(eq, qsc, iota=True, kinetic=True) + eq = Equilibrium(electron_temperature=1, electron_density=1, iota=1) + cs = get_NAE_constraints(eq, qsc) assert _is_any_instance(cs, FixAxisR) assert _is_any_instance(cs, FixAxisZ) assert _is_any_instance(cs, FixPsi) diff --git a/tests/test_perturbations.py b/tests/test_perturbations.py index 2ef16cd24d..29d54ab699 100644 --- a/tests/test_perturbations.py +++ b/tests/test_perturbations.py @@ -106,7 +106,7 @@ def test_perturb_with_float_without_error(): # np.concatenate cannot concatenate 0-D arrays. This test exercises the fix. eq = Equilibrium() objective = get_equilibrium_objective(eq=eq) - constraints = get_fixed_boundary_constraints(eq=eq, iota=False) + constraints = get_fixed_boundary_constraints(eq=eq) # perturb Psi with a float deltas = {"Psi": float(eq.Psi)} diff --git a/tests/test_profiles.py b/tests/test_profiles.py index 4654b1d7f8..25b5b9b13b 100644 --- a/tests/test_profiles.py +++ b/tests/test_profiles.py @@ -396,7 +396,7 @@ def test_solve_with_combined(self): sym=True, ) eq1.solve( - constraints=get_fixed_boundary_constraints(eq=eq1, kinetic=False), + constraints=get_fixed_boundary_constraints(eq=eq1), objective=ObjectiveFunction(objectives=ForceBalance(eq=eq1)), maxiter=5, ) @@ -417,7 +417,7 @@ def test_solve_with_combined(self): sym=True, ) eq2.solve( - constraints=get_fixed_boundary_constraints(eq=eq2, kinetic=True), + constraints=get_fixed_boundary_constraints(eq=eq2), objective=ObjectiveFunction(objectives=ForceBalance(eq=eq2)), maxiter=5, ) From 7c61c3af58344d177d0783e33f148e6b959328c9 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Fri, 3 Nov 2023 22:37:48 -0600 Subject: [PATCH 02/13] Update get_* calls in notebooks --- docs/notebooks/DESC_Fixed_Axis_NAE_Constraint.ipynb | 4 ++-- docs/notebooks/Toroidal_current_constraint.ipynb | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/notebooks/DESC_Fixed_Axis_NAE_Constraint.ipynb b/docs/notebooks/DESC_Fixed_Axis_NAE_Constraint.ipynb index d7332939ad..8d6a892d89 100644 --- a/docs/notebooks/DESC_Fixed_Axis_NAE_Constraint.ipynb +++ b/docs/notebooks/DESC_Fixed_Axis_NAE_Constraint.ipynb @@ -220,7 +220,7 @@ ], "source": [ "# get the fixed-boundary constraints, which include also fixing the pressure and fixing the current profile (iota=False flag means fix current)\n", - "constraints = get_fixed_boundary_constraints(eq=desc_eq, iota=False)\n", + "constraints = get_fixed_boundary_constraints(eq=desc_eq)\n", "print(constraints)\n", "\n", "# solve the equilibrium\n", @@ -554,7 +554,7 @@ "\n", "eq_NAE = eq_fit.copy()\n", "# this has all the constraints we need, iota=False specifies we want to fix current instead of iota\n", - "constraints = get_NAE_constraints(eq_NAE, qsc_eq, iota=False, order=1)\n", + "constraints = get_NAE_constraints(eq_NAE, qsc_eq, order=1)\n", "\n", "eq_NAE.solve(\n", " verbose=3,\n", diff --git a/docs/notebooks/Toroidal_current_constraint.ipynb b/docs/notebooks/Toroidal_current_constraint.ipynb index cb879c2d0e..9b8ae161e0 100644 --- a/docs/notebooks/Toroidal_current_constraint.ipynb +++ b/docs/notebooks/Toroidal_current_constraint.ipynb @@ -203,7 +203,7 @@ }, "outputs": [], "source": [ - "constraints = get_fixed_boundary_constraints(eq=eq, profiles=True, iota=False)" + "constraints = get_fixed_boundary_constraints(eq=eq, profiles=True)" ] }, { @@ -522,7 +522,7 @@ "outputs": [], "source": [ "objective = ObjectiveFunction(ForceBalance(eq=eq))\n", - "constraints = get_fixed_boundary_constraints(eq=eq, profiles=True, iota=False)" + "constraints = get_fixed_boundary_constraints(eq=eq, profiles=True)" ] }, { From 8ce138bb924a4055ec7fa90686613134876bca41 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 4 Nov 2023 00:39:04 -0600 Subject: [PATCH 03/13] Adjust test tolerance to account for CPU differences --- tests/test_compute_funs.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 4928f8532d..d53bed263c 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1292,7 +1292,8 @@ def test_compute_everything(): np.testing.assert_allclose( actual=this_branch_data[p][name], desired=master_data[p][name], - atol=1e-12, + atol=1e-10, + rtol=1e-10, err_msg=f"Parameterization: {p}. Name: {name}.", ) except AssertionError as e: From 8334abe56dce93ddbca9e323c5cdad5844f83dd3 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 6 Nov 2023 10:58:30 -0500 Subject: [PATCH 04/13] update stellar install documentation --- docs/installation.rst | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/installation.rst b/docs/installation.rst index 462f2e8076..94918f879f 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -123,15 +123,16 @@ specific JAX GPU installation instructions, as that is the main installation dif Della and Stellar Clusters (Princeton) ++++++++++++++++++++++++++++++++++++++ -These instructions were tested and confirmed to work on the Della and Stellar clusters at Princeton as of 8-26-2023. +These instructions were tested and confirmed to work on the Della and Stellar clusters at Princeton as of 11-6-2023. First, install JAX (check `this tutorial `__ ) for the latest version of `jaxlib` available on the Princeton clusters): .. code-block:: sh module load anaconda3/2023.3 - CONDA_OVERRIDE_CUDA="11.2" conda create --name desc-env jax "jaxlib==0.4.10=cuda112*" -c conda-forge + conda create --name desc-env python=3.9 conda activate desc-env + pip install --upgrade "jax[cuda11_pip]<0.4.15" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html Then, install DESC: From d02b1dcbcd68efaf5021b2dbd6acf436217af055 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 6 Nov 2023 12:36:17 -0500 Subject: [PATCH 05/13] add traverse commit --- docs/installation.rst | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/installation.rst b/docs/installation.rst index 94918f879f..7fef9c279a 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -153,7 +153,10 @@ On Clusters with IBM Power Architecture If pre-built JAX binaries are not available, you will first need to build JAX from source. More info can be found here: https://jax.readthedocs.io/en/latest/developer.html -These instructions were tested and confirmed to work on the Traverse supercomputer at Princeton as of 7-10-2023. +These instructions were tested and confirmed to work on the Traverse supercomputer at Princeton as of 11-6-2023. + +NOTE: You must use an older version of DESC in order to use Traverse, as there are some compatibility issues with JAX and the architecture. +Commit `a2fe711ffa3f` (an older version of the `master` branch) was tested to work fine on Traverse with these instructions. .. code-block:: sh From aada5a84184715abc39da5cdc69c60d1a1ec05f9 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 7 Nov 2023 18:06:10 -0500 Subject: [PATCH 06/13] fix inconsistency with the equations on the paper --- docs/theory_general.rst | 32 +++++++++++++++++++++++++------- 1 file changed, 25 insertions(+), 7 deletions(-) diff --git a/docs/theory_general.rst b/docs/theory_general.rst index 03e15688d5..4ef8a95f0e 100644 --- a/docs/theory_general.rst +++ b/docs/theory_general.rst @@ -68,7 +68,19 @@ Magnetic Field & Current Density By assuming nested flux surfaces, :math:`\mathbf{B} \cdot \nabla \rho = 0`, and invoking Gauss's Law, :math:`\nabla \cdot \mathbf{B} = 0`, the magnetic field is written in flux coordinates as .. math:: - \mathbf{B} = B^\theta \mathbf{e}_\theta + B^\zeta \mathbf{e}_\zeta = \frac{\partial_\rho \psi}{2 \pi \sqrt{g}} \cdot ((\iota - \partial_\zeta \lambda) \mathbf{e}_\theta + (1 + \partial_\theta \lambda) \mathbf{e}_\zeta) + \begin{aligned} + \mathbf{B} &= B^\theta \mathbf{e}_\theta + B^\zeta \mathbf{e}_\zeta \\[.3cm] + &= \frac{\partial_\rho \psi}{2 \pi \sqrt{g}} \left((\iota - \cfrac{\partial \lambda}{\partial \zeta}) \mathbf{e}_\theta + (1 + \cfrac{\partial \lambda}{\partial \theta}) \mathbf{e}_\zeta \right) + \end{aligned} + +So, + +.. math:: + \begin{aligned} + B^\rho &= 0 \\[.2cm] + B^\theta &= \frac{\partial_\rho \psi}{2 \pi \sqrt{g}} \left(\iota - \cfrac{\partial \lambda}{\partial \zeta}\right)\\[.2cm] + B^\zeta &= \frac{\partial_\rho \psi}{2 \pi \sqrt{g}} \left(1 + \cfrac{\partial \lambda}{\partial \theta}\right) + \end{aligned} The current density is then calculated from Ampere's Law, :math:`\nabla \times \mathbf{B} = \mu_0 \mathbf{J}`, @@ -94,19 +106,25 @@ Equilibrium Force Balance The ideal magnetohydrodynamic equilibrium force balance is defined as .. math:: - \mathbf{F} \equiv \mathbf{J} \times \mathbf{B} - \nabla p = \mathbf{0} + \mathbf{F} \equiv \nabla p - \mathbf{J} \times \mathbf{B} = \mathbf{0} + +Using cross product in curvilinear coordinates, we can write + +.. math:: + \mathbf{J} \times \mathbf{B} = \sqrt{g} (J^\theta B^\zeta - J^\zeta B^\theta)\mathbf{e}^\rho \hspace{.2cm}-\hspace{.2cm} + \sqrt{g} J^\rho (B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta) When written in flux coordinates there are only two independent components: .. math:: \begin{aligned} - \mathbf{F} &= F_\rho \nabla \rho + F_\beta \mathbf{\beta} \\[.2cm] - F_\rho &= \sqrt{g} (B^\zeta J^\theta - B^\theta J^\zeta) - \partial_\rho p \\[.2cm] - F_\beta &= \sqrt{g} B^\zeta J^\rho \\[.2cm] - \mathbf{\beta} &= \nabla \theta - \iota \nabla \zeta + \mathbf{F} &= F_\rho \mathbf{e}^\rho + F_\beta \mathbf{\beta}_{DESC} \\[.3cm] + F_\rho &= \frac{\partial p}{\partial \rho} + \sqrt{g} (J^\zeta B^\theta - J^\theta B^\zeta)\\[.3cm] + F_\beta &= \sqrt{g} J^\rho \\[.3cm] + \mathbf{\beta_{DESC}} &= (B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta) \end{aligned} -These forces in both the radial and helical directions must vanish in equilibrium. +where :math:`\mathbf{e}^u = \nabla u` is the contravariant basis vector of parameter u. These forces in both the radial and helical directions must vanish in equilibrium. DESC solves this force balance locally by evaluating the residual errors at discrete points in real space: .. math:: From 53dff1c24f9b2129eb2930c5c70c9eea73ab41f1 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 7 Nov 2023 18:17:20 -0500 Subject: [PATCH 07/13] move Jacobian to beta vector --- docs/theory_general.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/theory_general.rst b/docs/theory_general.rst index 4ef8a95f0e..29405808af 100644 --- a/docs/theory_general.rst +++ b/docs/theory_general.rst @@ -120,8 +120,8 @@ When written in flux coordinates there are only two independent components: \begin{aligned} \mathbf{F} &= F_\rho \mathbf{e}^\rho + F_\beta \mathbf{\beta}_{DESC} \\[.3cm] F_\rho &= \frac{\partial p}{\partial \rho} + \sqrt{g} (J^\zeta B^\theta - J^\theta B^\zeta)\\[.3cm] - F_\beta &= \sqrt{g} J^\rho \\[.3cm] - \mathbf{\beta_{DESC}} &= (B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta) + F_\beta &= J^\rho \\[.3cm] + \mathbf{\beta_{DESC}} &= \sqrt{g} (B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta) \end{aligned} where :math:`\mathbf{e}^u = \nabla u` is the contravariant basis vector of parameter u. These forces in both the radial and helical directions must vanish in equilibrium. From 6e13b7955318bfd96856267b6a4210ffd26619c5 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Tue, 7 Nov 2023 22:43:27 -0500 Subject: [PATCH 08/13] change sign of force errorfor consistency with the code --- docs/theory_general.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/theory_general.rst b/docs/theory_general.rst index 29405808af..28a212adaa 100644 --- a/docs/theory_general.rst +++ b/docs/theory_general.rst @@ -106,7 +106,7 @@ Equilibrium Force Balance The ideal magnetohydrodynamic equilibrium force balance is defined as .. math:: - \mathbf{F} \equiv \nabla p - \mathbf{J} \times \mathbf{B} = \mathbf{0} + \mathbf{F} \equiv \mathbf{J} \times \mathbf{B} - \nabla p = \mathbf{0} Using cross product in curvilinear coordinates, we can write @@ -119,7 +119,7 @@ When written in flux coordinates there are only two independent components: .. math:: \begin{aligned} \mathbf{F} &= F_\rho \mathbf{e}^\rho + F_\beta \mathbf{\beta}_{DESC} \\[.3cm] - F_\rho &= \frac{\partial p}{\partial \rho} + \sqrt{g} (J^\zeta B^\theta - J^\theta B^\zeta)\\[.3cm] + F_\rho &= -\frac{\partial p}{\partial \rho} + \sqrt{g} (J^\theta B^\zeta - J^\zeta B^\theta)\\[.3cm] F_\beta &= J^\rho \\[.3cm] \mathbf{\beta_{DESC}} &= \sqrt{g} (B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta) \end{aligned} From 58a3f497fc19ee6f1bd2d6f289c50f2130167474 Mon Sep 17 00:00:00 2001 From: YigitElma Date: Wed, 8 Nov 2023 11:22:23 -0500 Subject: [PATCH 09/13] move back the jacobian to F_beta --- docs/theory_general.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/theory_general.rst b/docs/theory_general.rst index 28a212adaa..2304d6ef56 100644 --- a/docs/theory_general.rst +++ b/docs/theory_general.rst @@ -120,8 +120,8 @@ When written in flux coordinates there are only two independent components: \begin{aligned} \mathbf{F} &= F_\rho \mathbf{e}^\rho + F_\beta \mathbf{\beta}_{DESC} \\[.3cm] F_\rho &= -\frac{\partial p}{\partial \rho} + \sqrt{g} (J^\theta B^\zeta - J^\zeta B^\theta)\\[.3cm] - F_\beta &= J^\rho \\[.3cm] - \mathbf{\beta_{DESC}} &= \sqrt{g} (B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta) + F_\beta &= \sqrt{g} J^\rho \\[.3cm] + \mathbf{\beta_{DESC}} &= B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta \end{aligned} where :math:`\mathbf{e}^u = \nabla u` is the contravariant basis vector of parameter u. These forces in both the radial and helical directions must vanish in equilibrium. From 05507ad9368c198ab1fa183c25710a80350836c5 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Wed, 8 Nov 2023 10:58:17 -0700 Subject: [PATCH 10/13] edit force balance equations --- desc/objectives/_equilibrium.py | 4 ++-- docs/theory_general.rst | 19 ++++++++++--------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/desc/objectives/_equilibrium.py b/desc/objectives/_equilibrium.py index d1057da714..79c26c73a9 100644 --- a/desc/objectives/_equilibrium.py +++ b/desc/objectives/_equilibrium.py @@ -15,13 +15,13 @@ class ForceBalance(_Objective): Given force densities: - Fᵨ = √g (B^ζ J^θ - B^θ J^ζ) - ∇ p + Fᵨ = √g (J^θ B^ζ - J^ζ B^θ) - ∇ p Fₕₑₗᵢ √g J^ρ and helical basis vector: - 𝐞ʰᵉˡⁱ = −B^ζ ∇ θ + B^θ ∇ ζ + 𝐞ʰᵉˡⁱ = B^ζ ∇ θ - B^θ ∇ ζ Minimizes the magnitude of the forces: diff --git a/docs/theory_general.rst b/docs/theory_general.rst index 2304d6ef56..fecbd5bc6d 100644 --- a/docs/theory_general.rst +++ b/docs/theory_general.rst @@ -118,23 +118,24 @@ When written in flux coordinates there are only two independent components: .. math:: \begin{aligned} - \mathbf{F} &= F_\rho \mathbf{e}^\rho + F_\beta \mathbf{\beta}_{DESC} \\[.3cm] - F_\rho &= -\frac{\partial p}{\partial \rho} + \sqrt{g} (J^\theta B^\zeta - J^\zeta B^\theta)\\[.3cm] - F_\beta &= \sqrt{g} J^\rho \\[.3cm] - \mathbf{\beta_{DESC}} &= B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta + \mathbf{F} &= F_\rho \mathbf{e}^\rho + F_{helical} \mathbf{e}^{helical} \\[.3cm] + F_\rho &= \sqrt{g} (J^\theta B^\zeta - J^\zeta B^\theta) - \frac{\partial p}{\partial \rho} \\[.3cm] + F_{helical} &= \sqrt{g} J^\rho \\[.3cm] + \mathbf{e}^{helical} &= B^\zeta\mathbf{e}^\theta - B^\theta\mathbf{e}^\zeta \end{aligned} -where :math:`\mathbf{e}^u = \nabla u` is the contravariant basis vector of parameter u. These forces in both the radial and helical directions must vanish in equilibrium. +where :math:`\mathbf{e}^x = \nabla x` is the contravariant basis vector of parameter x. +These forces in both the radial and helical directions must vanish in equilibrium. DESC solves this force balance locally by evaluating the residual errors at discrete points in real space: .. math:: \begin{aligned} f_\rho &= F_\rho ||\nabla \rho|| \Delta V \\[.2cm] - f_\beta &= F_\beta ||\mathbf{\beta}|| \Delta V + f_{helical} &= F_{helical} ||\mathbf{e}^{helical}|| \Delta V \end{aligned} -These equations :math:`f_\rho` and :math:`f_\beta` represent the force errors (in Newtons) in the unit of volume :math:`\Delta V = \sqrt{g} \Delta \rho \Delta \theta \Delta \zeta` surrounding a collocation point :math:`(\rho, \theta, \zeta)`. -[Note: this definition of :math:`\mathbf{\beta}` is slightly different from that given in the original paper, but the resulting equation for :math:`f_\beta` is equivalent. +These equations :math:`f_\rho` and :math:`f_{helical}` represent the force errors (in Newtons) in the unit of volume :math:`\Delta V = \sqrt{g} \Delta \rho \Delta \theta \Delta \zeta` surrounding a collocation point :math:`(\rho, \theta, \zeta)`. +[Note: this definition of :math:`\mathbf{e}^{helical}` is slightly different from that of :math:`\mathbf{\beta}` given in the original paper, but the resulting equation for :math:`f_{helical}` is equivalent to the original :math:`f_\beta`. The publication also included an additional sign term in the equations for :math:`f_\rho` and :math:`f_\beta` that has been dropped.] In summary, the equilibrium problem is formulated as a system of nonlinear equations :math:`\mathbf{f}(\mathbf{x}, \mathbf{c}) = \mathbf{0}`. @@ -156,7 +157,7 @@ The equations :math:`\mathbf{f}` are the force error residuals at a series of co .. math:: \mathbf{f} = \begin{bmatrix} - f_\rho \\ f_\beta + f_\rho \\ f_{helical} \end{bmatrix} DESC allows flexibility in the choice of optimization algorithm used to solve this system of equations; popular approaches include Newton-Raphson methods and least-squares minimization (as the collocation grids are often oversampled, which has been found to improve convergence and robustness). From 4d8cf9933be5ca76e4077a7bf845c144506ee4f0 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 8 Nov 2023 15:00:05 -0500 Subject: [PATCH 11/13] reduce lin con tol --- desc/objectives/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py index b53b50f102..be962d9e3a 100644 --- a/desc/objectives/utils.py +++ b/desc/objectives/utils.py @@ -157,8 +157,8 @@ def recover(x_reduced): np.testing.assert_allclose( y1, y2, - atol=1e-14, - rtol=1e-14, + atol=2e-14, + rtol=5e-14, err_msg="Incompatible constraints detected, cannot satisfy " + f"constraint {con}", ) From 600e6e8ba24e87d845eba88e1e7b8ef022b96deb Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 8 Nov 2023 22:04:17 -0500 Subject: [PATCH 12/13] Fix docstrings --- desc/objectives/getters.py | 10 +++++----- tests/test_examples.py | 2 -- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/desc/objectives/getters.py b/desc/objectives/getters.py index 4910687ac8..ec2a36ecc1 100644 --- a/desc/objectives/getters.py +++ b/desc/objectives/getters.py @@ -87,9 +87,9 @@ def get_fixed_axis_constraints(eq, profiles=True, normalize=True): Parameters ---------- eq : Equilibrium - Equilibrium being constrained. + Equilibrium to constrain. profiles : bool - Whether to also return constraints to fix input profiles. + If True, also include constraints to fix all profiles assigned to equilibrium. normalize : bool Whether to apply constraints in normalized units. @@ -120,9 +120,9 @@ def get_fixed_boundary_constraints(eq, profiles=True, normalize=True): Parameters ---------- eq : Equilibrium - Equilibrium to constraint. + Equilibrium to constrain. profiles : bool - Whether to also return constraints to fix input profiles. + If True, also include constraints to fix all profiles assigned to equilibrium. normalize : bool Whether to apply constraints in normalized units. @@ -168,7 +168,7 @@ def get_NAE_constraints( order : int order (in rho) of near-axis behavior to constrain profiles : bool - Whether to also return constraints to fix input profiles. + If True, also include constraints to fix all profiles assigned to equilibrium. normalize : bool Whether to apply constraints in normalized units. N : int diff --git a/tests/test_examples.py b/tests/test_examples.py index 767e3271e9..37fe8a6160 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -547,7 +547,6 @@ def test_NAE_QSC_solve(): eq_lambda_fixed_1st_order = eq.copy() # this has all the constraints we need, - # iota=False specifies we want to fix current instead of iota cs = get_NAE_constraints(eq, qsc, order=1, fix_lambda=False, N=eq.N) cs_lambda_fixed_0th_order = get_NAE_constraints( eq_lambda_fixed_0th_order, qsc, order=1, fix_lambda=0, N=eq.N @@ -635,7 +634,6 @@ def test_NAE_QIC_solve(): eq_fit = eq.copy() # this has all the constraints we need, - # iota=False specifies we want to fix current instead of iota cs = get_NAE_constraints(eq, qic, order=1) objectives = ForceBalance(eq=eq) From 40931262e924688e5455b88d894058ec9a5e25d0 Mon Sep 17 00:00:00 2001 From: daniel-dudt Date: Sat, 11 Nov 2023 15:20:55 -0700 Subject: [PATCH 13/13] hot fix for rescale --- desc/compat.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/desc/compat.py b/desc/compat.py index 982ead88a0..7ab9fb5548 100644 --- a/desc/compat.py +++ b/desc/compat.py @@ -204,4 +204,8 @@ def rescale(eq, L=("R0", None), B=("B0", None), verbose=0): if eq.current is not None: eq.c_l *= cL * cB + # boundary & axis + eq.axis = eq.get_axis() + eq.surface = eq.get_surface_at(rho=1) + return eq