Merge branch 'master' into rc/mirror_ratio

CHANGELOG.md

@@ -4,6 +4,7 @@ Changelog
 New Feature
 
 - Adds a new profile class ``PowerProfile`` for raising profiles to a power.
+- Add ``desc.objectives.LinkingCurrentConsistency`` for ensuring that coils in a stage 2 or single stage optimization provide the required linking current for a given equilibrium.
 - Adds an option ``scaled_termination`` (defaults to True) to all of the desc optimizers to measure the norms for ``xtol`` and ``gtol`` in the scaled norm provided by ``x_scale`` (which defaults to using an adaptive scaling based on the Jacobian or Hessian). This should make things more robust when optimizing parameters with widely different magnitudes. The old behavior can be recovered by passing ``options={"scaled_termination": False}``.
 - ``desc.objectives.Omnigenity`` is now vectorized and able to optimize multiple surfaces at the same time. Previously it was required to use a different objective for each surface.
 - Adds a new objective ``desc.objectives.MirrorRatio`` for targeting a particular mirror ratio on each flux surface, for either an ``Equilibrium`` or ``OmnigenousField``.
@@ -16,7 +17,6 @@ v0.13.0
 -------
 
 New Features
-
 - Adds ``from_input_file`` method to ``Equilibrium`` class to generate an ``Equilibrium`` object with boundary, profiles, resolution and flux specified in a given DESC or VMEC input file
 - Adds function ``solve_regularized_surface_current`` to ``desc.magnetic_fields`` module that implements the REGCOIL algorithm (Landreman, (2017)) for surface current normal field optimization
     * Can specify the tuple ``current_helicity=(M_coil, N_coil)`` to determine if resulting contours correspond to helical topology (both ``(M_coil, N_coil)`` not equal to 0), modular (``N_coil`` equal to 0 and ``M_coil`` nonzero) or windowpane/saddle (``M_coil`` and ``N_coil`` both zero)
@@ -299,7 +299,7 @@ optimization. Set to False by default.
 non-singular, non-degenerate) coordinate mappings for initial guesses. This is applied
 automatically when creating a new `Equilibrium` if the default initial guess of scaling
 the boundary surface produces self-intersecting surfaces. This can be disabled by
-passing `ensure_nested=False` when constructing the `Equilibrum`.
+passing `ensure_nested=False` when constructing the `Equilibrium`.
 - Adds `loss_function` argument to all `Objective`s for applying one of min/max/mean
 to objective function values (for targeting the average value of a profile, etc).
 - `Equilibrium.get_profile` now allows user to choose a profile type (power series, spline, etc)
@@ -451,7 +451,7 @@ Breaking changes
 - Renames ``theta_sfl`` to ``theta_PEST`` in compute functions to avoid confusion with
 other straight field line coordinate systems.
 - Makes plotting kwargs a bit more uniform. ``zeta``, ``nzeta``, ``nphi`` have all been
-superceded by ``phi`` which can be an integer for equally spaced angles or a float or
+superseded by ``phi`` which can be an integer for equally spaced angles or a float or
 array of float to specify angles manually.
 
 Bug fixes
@@ -521,7 +521,7 @@ the future all quantities should evaluate correctly at the magnetic axis. Note t
 evaluating quantities at the axis generally requires higher order derivatives and so
 can be much more expensive than evaluating at nonsingular points, so during optimization
 it is not recommended to include a grid point at the axis. Generally a small finite value
-such as ``rho = 1e-6`` will avoid the singuarlity with a negligible loss in accuracy for
+such as ``rho = 1e-6`` will avoid the singularity with a negligible loss in accuracy for
 analytic quantities.
 - Adds new optimizers ``fmin-auglag`` and ``lsq-auglag`` for performing constrained
 optimization using the augmented Lagrangian method. These generally perform much better
@@ -537,7 +537,7 @@ the existing methods ``compute_theta_coordinates`` and ``compute_flux_coordinate
 but allows mapping between arbitrary coordinates.
 - Adds calculation of $\nabla \mathbf{B}$ tensor and corresponding $L_{\nabla B}$ metric
 - Adds objective ``BScaleLength`` for penalizing strong magnetic field curvature.
-- Adds objective ``ObjectiveFromUser`` for wrapping an arbitary user defined function.
+- Adds objective ``ObjectiveFromUser`` for wrapping an arbitrary user defined function.
 - Adds utilities ``desc.grid.find_least_rational_surfaces`` and
 ``desc.grid.find_most_rational_surfaces`` for finding the least/most rational surfaces
 for a given rotational transform profile.
@@ -1073,7 +1073,7 @@ New Features:
 -   Updates `Equilibrium` to make creating them more straightforward.
     -   Instead of a dictionary of arrays and values, init method now
         takes individual arguments. These can either be objects of the
-        correct type (ie `Surface` objects for boundary condiitons,
+        correct type (ie `Surface` objects for boundary conditions,
         `Profile` for pressure and iota etc,) or ndarrays which will get
         parsed into objects of the correct type (for backwards
         compatibility)
@@ -1357,7 +1357,7 @@ Major changes:
     indexing, where only M+1 points were used instead of the correct
     2\*M+1
 -   Rotated concentric grids by 2pi/3M to avoid symmetry plane at
-    theta=0,pi. Previously, for stellarator symmetic cases, the nodes at
+    theta=0,pi. Previously, for stellarator symmetric cases, the nodes at
     theta=0 did not contribute to helical force balance.
 -   Added [L\_grid]{.title-ref} parameter to specify radial resolution
     of grid nodes directly and making the API more consistent.
@@ -1523,7 +1523,7 @@ saved, and objectives getting compiled more often than necessary
 
 Major Changes:
 
--   Changes to Equilibium/EquilibriaFamily:
+-   Changes to Equilibrium/EquilibriaFamily:
     -   general switching to using properties rather than direct
         attributes when referencing things (ie, `eq.foo`, not
         `eq._foo`). This allows getter methods to have safeguards if

desc/coils.py

@@ -1300,6 +1300,15 @@ def current(self, new):
         for coil, cur in zip(self.coils, new):
             coil.current = cur
 
+    def _all_currents(self, currents=None):
+        """Return an array of all the currents (including those in virtual coils)."""
+        if currents is None:
+            currents = self.current
+        currents = jnp.asarray(currents)
+        if self.sym:
+            currents = jnp.concatenate([currents, -1 * currents[::-1]])
+        return jnp.tile(currents, self.NFP)
+
     def _make_arraylike(self, x):
         if isinstance(x, dict):
             x = [x] * len(self)
@@ -2337,6 +2346,22 @@ def num_coils(self):
         """int: Number of coils."""
         return sum([c.num_coils for c in self])
 
+    def _all_currents(self, currents=None):
+        """Return an array of all the currents (including those in virtual coils)."""
+        if currents is None:
+            currents = jnp.array(flatten_list(self.current))
+        all_currents = []
+        i = 0
+        for coil in self.coils:
+            if isinstance(coil, CoilSet):
+                curr = currents[i : i + len(coil)]
+                all_currents += [coil._all_currents(curr)]
+                i += len(coil)
+            else:
+                all_currents += [jnp.atleast_1d(currents[i])]
+                i += 1
+        return jnp.concatenate(all_currents)
+
     def compute(
         self,
         names,

desc/compute/_basis_vectors.py

@@ -34,7 +34,7 @@ def _b(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="e^rho",  # ∇ρ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates.
+    name="e^rho",  # ∇ρ is the same in any coordinate system.
     label="\\mathbf{e}^{\\rho}",
     units="m^{-1}",
     units_long="inverse meters",
@@ -927,7 +927,7 @@ def _e_sup_theta_zz(params, transforms, profiles, data, **kwargs):
 
 
 @register_compute_fun(
-    name="e^zeta",  # ∇ζ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates.
+    name="e^zeta",  # ∇ζ is the same in any coordinate system.
     label="\\mathbf{e}^{\\zeta}",
     units="m^{-1}",
     units_long="inverse meters",
@@ -2449,10 +2449,10 @@ def _e_sub_theta_over_sqrt_g(params, transforms, profiles, data, **kwargs):
 )
 def _e_sub_vartheta_rp(params, transforms, profiles, data, **kwargs):
     # constant ρ and ϕ
-    e_vartheta = (
-        data["e_theta"].T * data["phi_z"] - data["e_zeta"].T * data["phi_t"]
-    ) / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
-    data["e_theta_PEST"] = e_vartheta.T
+    data["e_theta_PEST"] = (
+        (data["e_theta"].T * data["phi_z"] - data["e_zeta"].T * data["phi_t"])
+        / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
+    ).T
     return data
 
 
@@ -2473,11 +2473,13 @@ def _e_sub_vartheta_rp(params, transforms, profiles, data, **kwargs):
 )
 def _e_sub_phi_rv(params, transforms, profiles, data, **kwargs):
     # constant ρ and ϑ
-    e_phi = (
-        data["e_zeta"].T * data["theta_PEST_t"]
-        - data["e_theta"].T * data["theta_PEST_z"]
-    ) / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
-    data["e_phi|r,v"] = e_phi.T
+    data["e_phi|r,v"] = (
+        (
+            data["e_zeta"].T * data["theta_PEST_t"]
+            - data["e_theta"].T * data["theta_PEST_z"]
+        )
+        / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
+    ).T
     return data
 
 
@@ -2499,10 +2501,10 @@ def _e_sub_phi_rv(params, transforms, profiles, data, **kwargs):
 def _e_sub_rho_vp(params, transforms, profiles, data, **kwargs):
     # constant ϑ and ϕ
     data["e_rho|v,p"] = (
-        data["e_rho"].T
-        - data["e_vartheta"].T * data["theta_PEST_r"]
-        - data["e_phi|r,v"].T * data["phi_r"]
-    ).T
+        data["e_rho"]
+        - data["e_vartheta"] * data["theta_PEST_r"][:, jnp.newaxis]
+        - data["e_phi|r,v"] * data["phi_r"][:, jnp.newaxis]
+    )
     return data
 
 
@@ -3206,7 +3208,28 @@ def _e_sub_zeta_zz(params, transforms, profiles, data, **kwargs):
             data["Z_zzz"],
         ]
     ).T
+    return data
+
 
+@register_compute_fun(
+    name="grad(phi)",
+    label="\\nabla \\phi",
+    units="m^{-1}",
+    units_long="Inverse meters",
+    description="Gradient of cylindrical toroidal angle ϕ.",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["R", "0"],
+    parameterization=[
+        "desc.equilibrium.equilibrium.Equilibrium",
+        "desc.geometry.surface.FourierRZToroidalSurface",
+    ],
+)
+def _grad_phi(params, transforms, profiles, data, **kwargs):
+    data["grad(phi)"] = jnp.column_stack([data["0"], 1 / data["R"], data["0"]])
     return data
 
 
@@ -3456,7 +3479,7 @@ def _e_sub_theta_rp(params, transforms, profiles, data, **kwargs):
     label="\\mathbf{e}_{\\rho} |_{\\alpha, \\zeta}",
     units="m",
     units_long="meters",
-    description="Tangent vector along radial field line label",
+    description="Covariant radial basis vector in (ρ, α, ζ) Clebsch coordinates.",
     dim=3,
     params=[],
     transforms={},
@@ -3477,7 +3500,7 @@ def _e_rho_az(params, transforms, profiles, data, **kwargs):
     label="\\mathbf{e}_{\\alpha}",
     units="m",
     units_long="meters",
-    description="Tangent vector along poloidal field line label",
+    description="Covariant poloidal basis vector in (ρ, α, ζ) Clebsch coordinates.",
     dim=3,
     params=[],
     transforms={},
@@ -3496,8 +3519,8 @@ def _e_alpha(params, transforms, profiles, data, **kwargs):
     label="\\partial_{\\theta} \\mathbf{e}_{\\alpha}",
     units="m",
     units_long="meters",
-    description="Tangent vector along poloidal field line label, derivative wrt"
-    " DESC poloidal angle",
+    description="Covariant poloidal basis vector in (ρ, α, ζ) Clebsch coordinates,"
+    " derivative wrt DESC poloidal angle",
     dim=3,
     params=[],
     transforms={},
@@ -3518,8 +3541,8 @@ def _e_alpha_t(params, transforms, profiles, data, **kwargs):
     label="\\partial_{\\zeta} \\mathbf{e}_{\\alpha}",
     units="m",
     units_long="meters",
-    description="Tangent vector along poloidal field line label, "
-    "derivative wrt DESC toroidal angle",
+    description="Covariant poloidal basis vector in (ρ, α, ζ) Clebsch coordinates, "
+    "derivative wrt DESC toroidal angle at fixed ρ,θ.",
     dim=3,
     params=[],
     transforms={},
@@ -3603,7 +3626,7 @@ def _e_zeta_ra_a(params, transforms, profiles, data, **kwargs):
     units="m",
     units_long="meters",
     description="Tangent vector along (collinear to) field line, "
-    "derivative wrt DESC toroidal angle",
+    "derivative wrt DESC toroidal angle at fixed ρ,θ.",
     dim=3,
     params=[],
     transforms={},
@@ -3679,3 +3702,43 @@ def _d_ell_d_zeta_z(params, transforms, profiles, data, **kwargs):
         dot(data["(e_zeta|r,a)_z|r,a"], data["e_zeta|r,a"]) / data["|e_zeta|r,a|"]
     )
     return data
+
+
+@register_compute_fun(
+    name="e_alpha|r,p",
+    label="\\mathbf{e}_{\\alpha} |_{\\rho, \\phi}",
+    units="m",
+    units_long="meters",
+    description="Covariant poloidal basis vector in (ρ, α, ϕ) Clebsch coordinates.",
+    dim=3,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e_theta", "alpha_t", "e_zeta", "alpha_z", "phi_t", "phi_z"],
+)
+def _e_alpha_rp(params, transforms, profiles, data, **kwargs):
+    data["e_alpha|r,p"] = (
+        (data["e_theta"].T * data["phi_z"] - data["e_zeta"].T * data["phi_t"])
+        / (data["alpha_t"] * data["phi_z"] - data["alpha_z"] * data["phi_t"])
+    ).T
+    return data
+
+
+@register_compute_fun(
+    name="|e_alpha|r,p|",
+    label="|\\mathbf{e}_{\\alpha} |_{\\rho, \\phi}|",
+    units="m",
+    units_long="meters",
+    description="Norm of covariant poloidal basis vector in (ρ, α, ϕ) Clebsch "
+    "coordinates.",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["e_alpha|r,p"],
+)
+def _e_alpha_rp_norm(params, transforms, profiles, data, **kwargs):
+    data["|e_alpha|r,p|"] = jnp.linalg.norm(data["e_alpha|r,p"], axis=-1)
+    return data

desc/compute/_core.py

@@ -3153,6 +3153,25 @@ def _theta_PEST_rt(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="theta_PEST_rz",
+    label="\\partial_{\\rho \\zeta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, derivative wrt "
+    "radial and DESC toroidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_rz"],
+)
+def _theta_PEST_rz(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_rz"] = data["lambda_rz"]
+    return data
+
+
 @register_compute_fun(
     name="theta_PEST_rrt",
     label="\\partial_{\\rho \\rho \\theta} \\vartheta",