Skip to content

Commit

Permalink
Merge branch 'ripple' into Gamma_c
Browse files Browse the repository at this point in the history
  • Loading branch information
unalmis authored Nov 26, 2024
2 parents 67ba404 + 92fe128 commit 158fade
Show file tree
Hide file tree
Showing 76 changed files with 1,182,038 additions and 321 deletions.
14 changes: 13 additions & 1 deletion CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,13 +2,25 @@ Changelog
=========

New Features
- Add ``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 ``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)
* ``M_coil`` is the number of poloidal transits a coil makes before returning to itself, while ``N_coil`` is the number of toroidal transits a coil makes before returning to itself (this is sort of like the QS ``helicity``)
* if multiple values of the regularization parameter are input, will return a family of surface current fields (as a list) corresponding to the solution at each regularization value
- Adds method ``to_CoilSet`` to ``FourierCurrentPotentialField`` which implements a coil cutting algorithm to discretize the surface current into coils
* works for both modular and helical coils
- Adds a new objective ``SurfaceCurrentRegularization`` (which minimizes ``w*|K|``, the regularization term from surface current in the REGCOIL algorithm, with `w` being the objective weight which act as the regularization parameter)
* use of both this and the ``QuadraticFlux`` objective allows for REGCOIL solutions to be obtained through the optimization framework, and combined with other objectives as well.
- Changes local area weighting of Bn in QuadraticFlux objective to be the square root of the local area element (Note that any existing optimizations using this objective may need different weights to achieve the same result now.)
- Adds a new tutorial showing how to use``REGCOIL`` features.


Bug Fixes

- Fixes bug that occurs when taking the gradient of ``root`` and ``root_scalar`` with newer versions of JAX (>=0.4.34) and unpins the JAX version
- Changes ``FixLambdaGauge`` constraint to now enforce zero flux surface average for lambda, instead of enforcing lambda(rho,0,0)=0 as it was incorrectly doing before.
- Fixes bug in ``softmin/softmax`` implementation.


v0.12.3
-------
Expand Down
6 changes: 3 additions & 3 deletions desc/backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@
)
)

if use_jax: # noqa: C901 - FIXME: simplify this, define globally and then assign?
if use_jax: # noqa: C901
from jax import custom_jvp, jit, vmap

imap = jax.lax.map
Expand All @@ -76,7 +76,7 @@
from jax.numpy.fft import irfft, rfft, rfft2
from jax.scipy.fft import dct, idct
from jax.scipy.linalg import block_diag, cho_factor, cho_solve, qr, solve_triangular
from jax.scipy.special import gammaln, logsumexp
from jax.scipy.special import gammaln
from jax.tree_util import (
register_pytree_node,
tree_flatten,
Expand Down Expand Up @@ -445,7 +445,7 @@ def tangent_solve(g, y):
qr,
solve_triangular,
)
from scipy.special import gammaln, logsumexp # noqa: F401
from scipy.special import gammaln # noqa: F401
from scipy.special import softmax as softargmax # noqa: F401

trapezoid = np.trapezoid if hasattr(np, "trapezoid") else np.trapz
Expand Down
5 changes: 2 additions & 3 deletions desc/basis.py
Original file line number Diff line number Diff line change
Expand Up @@ -1098,13 +1098,13 @@ def evaluate(
if not len(modes):
return np.array([]).reshape((len(nodes), 0))

# TODO: avoid duplicate calculations when mixing derivatives
# TODO(#1243): avoid duplicate calculations when mixing derivatives
r, t, z = nodes.T
l, m, n = modes.T
lm = modes[:, :2]

if unique:
# TODO: can avoid this here by using grid.unique_idx etc
# TODO(#1243): can avoid this here by using grid.unique_idx etc
# and adding unique_modes attributes to basis
_, ridx, routidx = np.unique(
r, return_index=True, return_inverse=True, axis=0
Expand Down Expand Up @@ -1364,7 +1364,6 @@ def polyval_vec(p, x, prec=None):
def _polyval_exact(p, x, prec):
p = np.atleast_2d(p)
x = np.atleast_1d(x).flatten()
# TODO: possibly multithread this bit
mpmath.mp.dps = prec
y = np.array([np.asarray(mpmath.polyval(list(pi), x)) for pi in p])
return y.astype(float)
Expand Down
1 change: 0 additions & 1 deletion desc/batching.py
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,6 @@ def f_(carry, x):
return res_append


# TODO in_axes a la vmap?
def _scanmap(fun, scan_fun, argnums=0):
"""A helper function to wrap f with a scan_fun."""

Expand Down
37 changes: 28 additions & 9 deletions desc/coils.py
Original file line number Diff line number Diff line change
Expand Up @@ -1633,7 +1633,14 @@ def compute_magnetic_vector_potential(

@classmethod
def linspaced_angular(
cls, coil, current=None, axis=[0, 0, 1], angle=2 * np.pi, n=10, endpoint=False
cls,
coil,
current=None,
axis=[0, 0, 1],
angle=2 * np.pi,
n=10,
endpoint=False,
check_intersection=True,
):
"""Create a CoilSet by repeating a coil at equal spacing around the torus.
Expand All @@ -1651,6 +1658,8 @@ def linspaced_angular(
Number of copies of original coil.
endpoint : bool
Whether to include a coil at final rotation angle. Default = False.
check_intersection : bool
whether to check the resulting coilsets for intersecting coils.
"""
assert isinstance(coil, _Coil) and not isinstance(coil, CoilSet)
Expand All @@ -1664,11 +1673,17 @@ def linspaced_angular(
coili.rotate(axis=axis, angle=phi[i])
coili.current = currents[i]
coils.append(coili)
return cls(*coils)
return cls(*coils, check_intersection=check_intersection)

@classmethod
def linspaced_linear(
cls, coil, current=None, displacement=[2, 0, 0], n=4, endpoint=False
cls,
coil,
current=None,
displacement=[2, 0, 0],
n=4,
endpoint=False,
check_intersection=True,
):
"""Create a CoilSet by repeating a coil at equal spacing in a straight line.
Expand All @@ -1685,6 +1700,8 @@ def linspaced_linear(
Number of copies of original coil.
endpoint : bool
Whether to include a coil at final displacement location. Default = False.
check_intersection : bool
whether to check the resulting coilsets for intersecting coils.
"""
assert isinstance(coil, _Coil) and not isinstance(coil, CoilSet)
Expand All @@ -1699,10 +1716,10 @@ def linspaced_linear(
coili.translate(a[i] * displacement)
coili.current = currents[i]
coils.append(coili)
return cls(*coils)
return cls(*coils, check_intersection=check_intersection)

@classmethod
def from_symmetry(cls, coils, NFP=1, sym=False):
def from_symmetry(cls, coils, NFP=1, sym=False, check_intersection=True):
"""Create a coil group by reflection and symmetry.
Given coils over one field period, repeat coils NFP times between
Expand All @@ -1721,6 +1738,8 @@ def from_symmetry(cls, coils, NFP=1, sym=False):
sym : bool (optional)
Whether to enforce stellarator symmetry.
If True, the coils will be duplicated 2*NFP times. Default = False.
check_intersection : bool
whether to check the resulting coilsets for intersecting coils.
Returns
-------
Expand Down Expand Up @@ -1763,7 +1782,7 @@ def from_symmetry(cls, coils, NFP=1, sym=False):
rotated_coils.rotate(axis=[0, 0, 1], angle=2 * jnp.pi * k / NFP)
coilset += rotated_coils

return cls(*coilset)
return cls(*coilset, check_intersection=check_intersection)

@classmethod
def from_makegrid_coilfile(cls, coil_file, method="cubic", check_intersection=True):
Expand Down Expand Up @@ -1901,8 +1920,8 @@ def save_in_makegrid_format(self, coilsFilename, NFP=None, grid=None):
if None, will default to the coil compute functions's
default grid
"""
# TODO: name each group based off of CoilSet name?
# TODO: have CoilGroup be automatically assigned based off of
# TODO(#1376): name each group based off of CoilSet name?
# TODO(#1376): have CoilGroup be automatically assigned based off of
# CoilSet if current coilset is a collection of coilsets?

NFP = 1 if NFP is None else NFP
Expand Down Expand Up @@ -2697,7 +2716,7 @@ def insert(self, i, new_item):
self._coils.insert(i, new_item)

@classmethod
def from_makegrid_coilfile( # noqa: C901 - FIXME: simplify this
def from_makegrid_coilfile( # noqa: C901
cls, coil_file, method="cubic", ignore_groups=False, check_intersection=True
):
"""Create a MixedCoilSet of SplineXYZCoils from a MAKEGRID coil txtfile.
Expand Down
6 changes: 3 additions & 3 deletions desc/compute/_geometry.py
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,7 @@ def _A_of_z(params, transforms, profiles, data, **kwargs):
data=["Z", "n_rho", "e_theta|r,p", "rho"],
parameterization=["desc.geometry.surface.FourierRZToroidalSurface"],
resolution_requirement="rt", # just need max(rho) near 1
# FIXME: Add source grid requirement once omega is nonzero.
# TODO(#568): Add source grid requirement once omega is nonzero.
)
def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs):
# Denote any vector v = [vᴿ, v^ϕ, vᶻ] with a tuple of its contravariant components.
Expand All @@ -213,7 +213,7 @@ def _A_of_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg
line_integrals(
transforms["grid"],
data["Z"] * n[:, 2] * safenorm(data["e_theta|r,p"], axis=-1),
# FIXME: Works currently for omega = zero, but for nonzero omega
# TODO(#568): Works currently for omega = zero, but for nonzero omega
# we need to integrate over theta at constant phi.
# Should be simple once we have coordinate mapping and source grid
# logic from GitHub pull request #1024.
Expand Down Expand Up @@ -449,7 +449,7 @@ def _perimeter_of_z(params, transforms, profiles, data, **kwargs):
line_integrals(
transforms["grid"],
safenorm(data["e_theta|r,p"], axis=-1),
# FIXME: Works currently for omega = zero, but for nonzero omega
# TODO(#568): Works currently for omega = zero, but for nonzero omega
# we need to integrate over theta at constant phi.
# Should be simple once we have coordinate mapping and source grid
# logic from GitHub pull request #1024.
Expand Down
4 changes: 2 additions & 2 deletions desc/compute/_omnigenity.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,7 +47,7 @@ def fitfun(x):
return data


# TODO: do math to change definition of nu so that we can just use B_zeta_mn here
# TODO (#568): do math to change definition of nu so that we can just use B_zeta_mn here
@register_compute_fun(
name="B_phi_mn",
label="B_{\\phi, m, n}",
Expand All @@ -63,7 +63,7 @@ def fitfun(x):
data=["B_phi|r,t"],
resolution_requirement="tz",
grid_requirement={"is_meshgrid": True},
aliases="B_zeta_mn", # TODO: remove when phi != zeta
aliases="B_zeta_mn", # TODO(#568): remove when phi != zeta
M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M",
N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N",
)
Expand Down
6 changes: 3 additions & 3 deletions desc/compute/_profiles.py
Original file line number Diff line number Diff line change
Expand Up @@ -1382,7 +1382,7 @@ def _iota_num_rrr(params, transforms, profiles, data, **kwargs):
- beta * data["sqrt(g)_rrr"],
data["sqrt(g)"],
),
# Todo: axis limit of beta_rrr
# TODO(#587): axis limit of beta_rrr
# Computed with four applications of l’Hôpital’s rule.
# Requires sqrt(g)_rrrr and fourth derivatives of basis vectors.
jnp.nan,
Expand Down Expand Up @@ -1660,7 +1660,7 @@ def _iota_den_rrr(params, transforms, profiles, data, **kwargs):
- gamma * data["sqrt(g)_rrr"],
data["sqrt(g)"],
),
# Todo: axis limit
# TODO(#587): axis limit
# Computed with four applications of l’Hôpital’s rule.
# Requires sqrt(g)_rrrr and fourth derivatives of basis vectors.
jnp.nan,
Expand Down Expand Up @@ -1717,7 +1717,7 @@ def _q(params, transforms, profiles, data, **kwargs):
return data


# TODO: add K(rho,theta,zeta)*grad(rho) term
# TODO (#1381): add K(rho,theta,zeta)*grad(rho) term
@register_compute_fun(
name="I",
label="I",
Expand Down
2 changes: 1 addition & 1 deletion desc/compute/_stability.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ def _D_current(params, transforms, profiles, data, **kwargs):
/ data["|grad(psi)|"] ** 3
* dot(Xi, data["B"])
),
# Todo: implement equivalent of equation 4.3 in desc coordinates
# TODO(#671): implement equivalent of equation 4.3 in desc coordinates
jnp.nan,
)
)
Expand Down
Loading

0 comments on commit 158fade

Please sign in to comment.