diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py index ac5720224a..5285ef1255 100644 --- a/desc/magnetic_fields/_dommaschk.py +++ b/desc/magnetic_fields/_dommaschk.py @@ -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 @@ -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. @@ -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 @@ -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 @@ -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: @@ -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): @@ -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] @@ -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)) @@ -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)) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 5cb6d55531..635132b49a 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -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."""