Skip to content

Commit

Permalink
Dommaschk Potential Miscellaneous Fixes and Improvements (#961)
Browse files Browse the repository at this point in the history 
- Adds `NFP` to fitting function, so that only dommaschk potential modes
with the corresponding discrete symmetry are included in the fit
-Changes `0**0` derivative trick to use double where, as it seems that
the prior fix was causing some issues with the evaluation at Z=0

<img width="459" alt="image"
src="https://github.com/PlasmaControl/DESC/assets/37969854/ac3ab268-2403-4ada-b7a0-cb34e0c28f3f">

Fitted toroidal B from an elliptical vacuum equilibrium at rho=0.5 along
a curve of constant theta (Z=0 at the first point)
The first point B is way off for some reason when the check and replace
for `0**0` was present, but when removed this error vanishes (really am
unsure of why this fix works but the proof is in the result)

<img width="459" alt="image"
src="https://github.com/PlasmaControl/DESC/assets/37969854/dc17a687-64c7-45a2-91d0-47a54f1e4b4f">

same plot but after removing those lines where we check for 0**0 and
replace the exponent with 1.

Resolves #960 
Resolves #946
  • Loading branch information
@dpanici
dpanici authored Mar 29, 2024
2 parents a26926f + 7ee3b98 commit 22f6280
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 65 deletions.
37 changes: 21 additions & 16 deletions desc/magnetic_fields/_dommaschk.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,9 +23,11 @@ class DommaschkPotentialField(ScalarPotentialField):
Parameters
----------
ms : 1D array-like of int
first indices of V_m_l terms (eq. 12 of reference)
first indices of V_m_l terms (eq. 12 of reference),
corresponds to the toroidal periodicity of the mode.
ls : 1D array-like of int
second indices of V_m_l terms (eq. 12 of reference)
second indices of V_m_l terms (eq. 12 of reference),
corresponds to the poloidal periodicity of the mode.
a_arr : 1D array-like of float
a_m_l coefficients of V_m_l terms, which multiply the cos(m*phi)*D_m_l terms
b_arr : 1D array-like of float
Expand Down Expand Up @@ -79,7 +81,7 @@ def __init__(

@classmethod
def fit_magnetic_field( # noqa: C901 - FIXME - simplify
cls, field, coords, max_m, max_l, sym=False, verbose=1
cls, field, coords, max_m, max_l, sym=False, verbose=1, NFP=1
):
"""Fit a vacuum magnetic field with a Dommaschk Potential field.
Expand All @@ -91,11 +93,16 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify
if ndarray, must be an ndarray of the magnetic field in rpz,
of shape (num_nodes,3) with the columns being (B_R, B_phi, B_Z)
coords (ndarray): shape (num_nodes,3) of R,phi,Z points to fit field at
max_m (int): maximum m to use for Dommaschk Potentials
max_m (int): maximum m to use for Dommaschk Potentials, within one field period
i.e. if NFP= 2 and max_m = 3, then modes with arguments up to 3*2*phi will
be included
max_l (int): maximum l to use for Dommaschk Potentials
sym (bool): if field is stellarator symmetric or not.
if True, only stellarator-symmetric modes will
be included in the fitting
NFP (int): if the field being fit has a discrete toroidal symmetry
with field period NFP. This will only allow Dommaschk m modes
that are integer multiples of NFP.
verbose (int): verbosity level of fitting routine, > 0 prints residuals
"""
# We seek c in Ac = b
Expand Down Expand Up @@ -124,7 +131,7 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify
#####################
# b is made, now do A
#####################
num_modes = 1 + (max_l + 1) * (max_m + 1) * 4
num_modes = 1 + (max_l) * (max_m + 1) * 4
# TODO: if symmetric, technically only need half the modes
# however, the field and functions are setup to accept equal
# length arrays for a,b,c,d, so we will just zero out the
Expand Down Expand Up @@ -160,8 +167,8 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify
# if sym is True, when l is even then we need a=d=0
# and if l is odd then b=c=0

for l in range(max_l + 1):
for m in range(max_m + 1):
for l in range(1, max_l + 1):
for m in range(0, max_m * NFP + 1, NFP):
if not sym:
pass # no sym, use all coefs
elif l // 2 == 0:
Expand All @@ -181,7 +188,6 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify
abcd_zero_due_to_sym_inds[which_coef].append(0)
else:
abcd_zero_due_to_sym_inds[which_coef].append(1)

ms.append(m)
ls.append(l)
for i in range(4):
Expand Down Expand Up @@ -254,8 +260,7 @@ def get_B_dom(coords, X, ms, ls):
# we zero out the terms that should be zero due to symmetry here
# TODO: should also just not return any zeroed-out modes, but
# the way the modes are cataloged here with the ls and ms arrays,
# it is not straightforward to do that. By that I mean
# say ls = [1,2] ms = [1,1]
# it is not straightforward to do that
a_arr = c[1 : n + 1] * abcd_zero_due_to_sym_inds[0]
b_arr = c[n + 1 : 2 * n + 1] * abcd_zero_due_to_sym_inds[1]
c_arr = c[2 * n + 1 : 3 * n + 1] * abcd_zero_due_to_sym_inds[2]
Expand Down Expand Up @@ -412,9 +417,9 @@ def D_m_n(R, Z, m, n):
def body_fun(k, val):
coef = CD_m_k(R, m, k) / gamma(n - 2 * k + 1)
exp = n - 2 * k
# derivative of 0**0 is ill defined, so we do this to enforce it being 0
exp = jnp.where((Z == 0) & (exp == 0), 1, exp)
return val + coef * Z**exp
mask = (Z == 0) & (exp == 0)
exp = jnp.where(mask, 1, exp)
return val + coef * jnp.where(mask, 0, Z**exp)

return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R))

Expand All @@ -427,9 +432,9 @@ def N_m_n(R, Z, m, n):
def body_fun(k, val):
coef = CN_m_k(R, m, k) / gamma(n - 2 * k + 1)
exp = n - 2 * k
# derivative of 0**0 is ill defined, so we do this to enforce it being 0
exp = jnp.where((Z == 0) & (exp == 0), 1, exp)
return val + coef * Z**exp
mask = (Z == 0) & (exp == 0)
exp = jnp.where(mask, 1, exp)
return val + coef * jnp.where(mask, 0, Z**exp)

return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R))

Expand Down
49 changes: 0 additions & 49 deletions tests/test_magnetic_fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -854,55 +854,6 @@ def test_dommaschk_vertical_field():
np.testing.assert_allclose(B_dom[:, 2], ones, atol=5e-15)


@pytest.mark.unit
def test_dommaschk_fit_toroidal_field():
"""Test the Dommaschk potential fit for a 1/R toroidal scaled to 2 T."""
phi = np.linspace(0, 2 * jnp.pi, 3)
R = np.linspace(0.1, 1.5, 3)
Z = np.linspace(-0.5, 0.5, 3)
R, phi, Z = np.meshgrid(R, phi, Z)
coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T

max_l = 2
max_m = 2
B0 = 2 # scale strength for to 1/R field to fit

def B0_over_R(coord):
B_R = np.zeros_like(coord[:, 0])
B_phi = B0 / coord[:, 0]
B_Z = np.zeros_like(coord[:, 0])
return np.vstack((B_R, B_phi, B_Z)).T

B = DommaschkPotentialField.fit_magnetic_field(
B0_over_R, coords, max_m, max_l, sym=True
)

B_dom = B.compute_magnetic_field(coords)
np.testing.assert_allclose(B_dom[:, 0], 0, atol=2e-14)
np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=2e-14)
np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=2e-14)

# only nonzero coefficient of the field should be the B0
np.testing.assert_allclose(B._params["B0"], B0, atol=1e-14)
for coef in ["a_arr", "b_arr", "c_arr", "d_arr"]:
np.testing.assert_allclose(B._params[coef], 0, atol=1e-14)

# test fit from values
B = DommaschkPotentialField.fit_magnetic_field(
B0_over_R(coords), coords, max_m, max_l, sym=True
)

B_dom = B.compute_magnetic_field(coords)
np.testing.assert_allclose(B_dom[:, 0], 0, atol=2e-14)
np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=2e-14)
np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=2e-14)

# only nonzero coefficient of the field should be the B0
np.testing.assert_allclose(B._params["B0"], B0, atol=1e-14)
for coef in ["a_arr", "b_arr", "c_arr", "d_arr"]:
np.testing.assert_allclose(B._params[coef], 0, atol=1e-14)


@pytest.mark.unit
def test_dommaschk_fit_vertical_and_toroidal_field():
"""Test the Dommaschk potential fit for a toroidal and a vertical field."""
Expand Down

0 comments on commit 22f6280

Please sign in to comment.