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 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/_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/desc/objectives/getters.py b/desc/objectives/getters.py index a5218e27e2..ec2a36ecc1 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,23 +81,15 @@ 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 ---------- eq : Equilibrium - Equilibrium being constrained. + Equilibrium 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 add constraints to fix kinetic profiles or pressure - anisotropy : bool - Whether to add constraint to fix anisotropic pressure. + If True, also include constraints to fix all profiles assigned to equilibrium. normalize : bool Whether to apply constraints in normalized units. @@ -102,58 +105,24 @@ 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 ---------- eq : Equilibrium - Equilibrium to constraint. + Equilibrium 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. + If True, also include constraints to fix all profiles assigned to equilibrium. 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, @@ -227,13 +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. - 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. + If True, also include constraints to fix all profiles assigned to equilibrium. 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/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}", ) 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/docs/installation.rst b/docs/installation.rst index 462f2e8076..7fef9c279a 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: @@ -152,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 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)" ] }, { diff --git a/docs/theory_general.rst b/docs/theory_general.rst index 03e15688d5..fecbd5bc6d 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}`, @@ -96,27 +108,34 @@ The ideal magnetohydrodynamic equilibrium force balance is defined as .. math:: \mathbf{F} \equiv \mathbf{J} \times \mathbf{B} - \nabla p = \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_{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}^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}`. @@ -138,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). 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_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: diff --git a/tests/test_examples.py b/tests/test_examples.py index 9e29f47a60..37fe8a6160 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"]) @@ -547,13 +547,12 @@ 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, 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: @@ -635,8 +634,7 @@ 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, 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, )