From 7d47f01dbba6a923d52dad585d68b1221c81314b Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sat, 8 Apr 2023 00:35:06 -0400 Subject: [PATCH 01/46] initial commit of dommaschk potentials implemented in JAX --- desc/magnetic_fields.py | 179 ++++++++++++++++++++++++++++++++++ tests/test_magnetic_fields.py | 20 ++++ 2 files changed, 199 insertions(+) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 442c10e278..6044d7549b 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -795,3 +795,182 @@ def odefun(rpz, s): r = x[:, :, 0].T.reshape((len(phis), *rshape)) z = x[:, :, 2].T.reshape((len(phis), *rshape)) return r, z + + +### Dommaschk potential utility functions ### + +# based off Representations for vacuum potentials in stellarators +# https://doi.org/10.1016/0010-4655(86)90109-8 + +# written with naive for loops initially and can jaxify later + + +def gamma(n): + """Gamma function, only implemented for integers (equiv to factorial of (n-1)).""" + return jnp.prod(jnp.arange(1, n)) + + +def alpha(m, n): + """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2 ** (2 * n + m)) + + +def alphastar(m, n): + """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + return (2 * n + m) * alpha(m, n) + + +def beta(m, n): + """Beta of eq 28.""" + return gamma(m - n) / (gamma(n + 1) * 2 ** (2 * n - m + 1)) + + +def betastar(m, n): + """Betastar of eq 28.""" + return (2 * n - m) * beta(m, n) + + +def CD_m_k(R, m, k): + """Eq 25 of Dommaschk paper.""" + if m == k == 0: # the initial term, defined in eqn 8 + return jnp.ones_like(R) + + sum1 = 0 + for j in range(k + 1): + sum1 += alpha(m, j) * betastar(m, k - j) * R ** (2 * j) + sum2 = 0 + for j in range(k + 1): + sum2 += alphastar(m, k - j) * beta(m, j) * R ** (2 * j) + return -(R ** (m)) * sum1 + R ** (-m) * sum2 + + +def CN_m_k(R, m, k): + """Eq 26 of Dommaschk paper.""" + if m == k == 0: # the initial term, defined in eqn 9 + return jnp.log(R) + + sum1 = 0 + for j in range(k + 1): + sum1 += alpha(m, j) * beta(m, k - j) * R ** (2 * j) + sum2 = 0 + for j in range(k + 1): + sum2 += alpha(m, k - j) * beta(m, j) * R ** (2 * j) + return R ** (m) * sum1 - R ** (-m) * sum2 + + +def D_m_n(R, Z, m, n): + """D_m_n term in eqn 8 of Dommaschk paper.""" + # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper + + max_ind = ( + n // 2 + ) # find the top of the summation ind, the range should be up to and including this + # i.e. this is the index of k at + # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum + result = 0 + for k in range(max_ind + 1): + result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) + return result + + +def N_m_n(R, Z, m, n): + """N_m_n term in eqn 9 of Dommaschk paper.""" + # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper + + max_ind = ( + n // 2 + ) # find the top of the summation ind, the range should be up to and including this + # i.e. this is the index of k at + # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum + result = 0 + for k in range(max_ind + 1): + result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) + return result + + +# TODO: note on how stell sym affects, i.e. +# if stell sym then a=d=0 for even l and b=c=0 for odd l + + +def V_m_l(R, Z, phi, m, l, a, b, c, d): + """Eq 12 of Dommaschk paper. + + Parameters + ---------- + R,Z,phi : array-like + Cylindrical coordinates (1-D arrays of each of size num_eval_pts) + to evaluate the Dommaschk potential term at. + m : int + first index of V_m_l term + l : int + second index of V_m_l term + a : float + a_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*D_m_l + b : float + b_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*D_m_l + c : float + c_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*N_m_l-1 + d : float + d_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*N_m_l-1 + + Returns + ------- + value : array-like + Value of this V_m_l term evaluated at the given R,phi,Z points + (same size as the size of the given R,phi, or Z arrays). + + """ + return (a * jnp.cos(m * phi) + b * jnp.sin(m * phi)) * D_m_n(R, Z, m, l) + ( + c * jnp.cos(m * phi) + d * jnp.sin(m * phi) + ) * N_m_n(R, Z, m, l - 1) + + +def dommaschk_potential(R, Z, phi, ms, ls, a_arr, b_arr, c_arr, d_arr): + """Eq 1 of Dommaschk paper. + + this is the total dommaschk potential for + a given set of m,l indices and their corresponding + coefficients a_ml, b_ml, c_ml d_ml. + + Parameters + ---------- + R,Z,phi : array-like + Cylindrical coordinates (1-D arrays of each of size num_eval_pts) + to evaluate the Dommaschk potential term at. + ms : 1D array-like of int + first indices of V_m_l terms + ls : 1D array-like of int + second indices of V_m_l terms + a_arr : 1D array-like of float + a_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*D_m_l + b_arr : 1D array-like of float + b_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*D_m_l + c_arr : 1D array-like of float + c_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*N_m_l-1 + d_arr : 1D array-like of float + d_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*N_m_l-1 + + Returns + ------- + value : array-like + Value of the total dommaschk potential evaluated + at the given R,phi,Z points + (same size as the size of the given R,phi, or Z arrays). + """ + value = phi # phi term + + # make sure all are 1D arrays + ms = jnp.atleast_1d(ms) + ls = jnp.atleast_1d(ls) + a_arr = jnp.atleast_1d(a_arr) + b_arr = jnp.atleast_1d(b_arr) + c_arr = jnp.atleast_1d(c_arr) + d_arr = jnp.atleast_1d(d_arr) + + assert ( + ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size + ), "Passed in arrays must all be of the same size!" + + for m, l, a, b, c, d in zip(ms, ls, a_arr, b_arr, c_arr, d_arr): + value += V_m_l(R, Z, phi, m, l, a, b, c, d) + return value diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 23e98b97b9..aaa9252574 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -5,6 +5,8 @@ from desc.backend import jnp from desc.magnetic_fields import ( + CD_m_k, + CN_m_k, PoloidalMagneticField, ScalarPotentialField, SplineMagneticField, @@ -130,3 +132,21 @@ def test_field_line_integrate(self): r, z = field_line_integrate(r0, z0, phis, field) np.testing.assert_allclose(r[-1], 10, rtol=1e-6, atol=1e-6) np.testing.assert_allclose(z[-1], 0.001, rtol=1e-6, atol=1e-6) + + +@pytest.mark.unit +def test_dommaschk_CN_CD_m_0(): + """Test of CD_m_k and CN_m_k when k=0.""" + # based off eqn 8 and 9 of Dommaschk paper + # TODO: put paper DOI here + for m in range(1, 6): + # test of CD_m_k based off eqn 8 + R = np.linspace(0.1, 1, 100) + res1 = CD_m_k(R, m, 0) + res2 = 0.5 * (R**m + R ** (-m)) + np.testing.assert_allclose(res1, res2) + + # test of CN_m_k based off eqn 9 + res1 = CN_m_k(R, m, 0) + res2 = 0.5 * (R**m - R ** (-m)) / m + np.testing.assert_allclose(res1, res2) From 27fd9d3a2eaaf0da09a1a5348ece93a59b7ef606 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sat, 8 Apr 2023 00:40:16 -0400 Subject: [PATCH 02/46] add doi to test --- tests/test_magnetic_fields.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index aaa9252574..b90b059626 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -138,7 +138,7 @@ def test_field_line_integrate(self): def test_dommaschk_CN_CD_m_0(): """Test of CD_m_k and CN_m_k when k=0.""" # based off eqn 8 and 9 of Dommaschk paper - # TODO: put paper DOI here + # https://doi.org/10.1016/0010-4655(86)90109-8 for m in range(1, 6): # test of CD_m_k based off eqn 8 R = np.linspace(0.1, 1, 100) From 52a1f5758de3e8c1ad4a0454783f3d8851de20b9 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sat, 8 Apr 2023 00:44:37 -0400 Subject: [PATCH 03/46] add test for assertion error --- tests/test_magnetic_fields.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index b90b059626..f0bb12d15a 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -12,6 +12,7 @@ SplineMagneticField, ToroidalMagneticField, VerticalMagneticField, + dommaschk_potential, field_line_integrate, ) @@ -150,3 +151,19 @@ def test_dommaschk_CN_CD_m_0(): res1 = CN_m_k(R, m, 0) res2 = 0.5 * (R**m - R ** (-m)) / m np.testing.assert_allclose(res1, res2) + + +@pytest.mark.unit +def test_dommaschk_potential_arr_equal_error(): + """Test the assert statement of the potential function.""" + phi = 1 + R = 1 + Z = 1 + ms = [1] + ls = [1] + a_arr = [1] + b_arr = [1] + c_arr = [1] + d_arr = [1, 1] + with pytest.raises(AssertionError): + dommaschk_potential(R, Z, phi, ms, ls, a_arr, b_arr, c_arr, d_arr) From 5fbc429cbd950b8ce01421471be9bfaf989e28fd Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 9 Apr 2023 14:05:39 -0400 Subject: [PATCH 04/46] add tests and DommaschkMagneticField, need to fix vertical field test still --- desc/magnetic_fields.py | 43 ++++++++++++++++++++++++++++-- tests/test_magnetic_fields.py | 50 +++++++++++++++++++++++++++++++++-- 2 files changed, 89 insertions(+), 4 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 6044d7549b..1d4d47fc5a 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -746,6 +746,45 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz"): return B +class DommaschkPotentialField(ScalarPotentialField): + """Magnetic field due to a Dommaschk scalar magnetic potential in rpz coordinates. + + From Dommaschk 1986 paper https://doi.org/10.1016/0010-4655(86)90109-8 + + this is the field due to the dommaschk potential (eq. 1) for + a given set of m,l indices and their corresponding + coefficients a_ml, b_ml, c_ml d_ml. + + Parameters + ---------- + ms : 1D array-like of int + first indices of V_m_l terms (eq. 12 of reference) + ls : 1D array-like of int + second indices of V_m_l terms (eq. 12 of reference) + 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 + b_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*D_m_l terms + c_arr : 1D array-like of float + c_m_l coefficients of V_m_l terms, which multiply the cos(m*phi)*N_m_l-1 term + d_arr : 1D array-like of float + d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms + + + """ + + def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr): + params = {} + params["ms"] = ms + params["ls"] = ls + params["a_arr"] = a_arr + params["b_arr"] = b_arr + params["c_arr"] = c_arr + params["d_arr"] = d_arr + + super().__init__(dommaschk_potential, params) + + def field_line_integrate( r0, z0, phis, field, params={}, rtol=1e-8, atol=1e-8, maxstep=1000 ): @@ -925,7 +964,7 @@ def V_m_l(R, Z, phi, m, l, a, b, c, d): ) * N_m_n(R, Z, m, l - 1) -def dommaschk_potential(R, Z, phi, ms, ls, a_arr, b_arr, c_arr, d_arr): +def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr): """Eq 1 of Dommaschk paper. this is the total dommaschk potential for @@ -972,5 +1011,5 @@ def dommaschk_potential(R, Z, phi, ms, ls, a_arr, b_arr, c_arr, d_arr): ), "Passed in arrays must all be of the same size!" for m, l, a, b, c, d in zip(ms, ls, a_arr, b_arr, c_arr, d_arr): - value += V_m_l(R, Z, phi, m, l, a, b, c, d) + value += V_m_l(R, phi, Z, m, l, a, b, c, d) return value diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index f0bb12d15a..4f64f3372c 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -7,6 +7,7 @@ from desc.magnetic_fields import ( CD_m_k, CN_m_k, + DommaschkPotentialField, PoloidalMagneticField, ScalarPotentialField, SplineMagneticField, @@ -155,7 +156,7 @@ def test_dommaschk_CN_CD_m_0(): @pytest.mark.unit def test_dommaschk_potential_arr_equal_error(): - """Test the assert statement of the potential function.""" + """Test the assert statement of the Dommaschk potential function.""" phi = 1 R = 1 Z = 1 @@ -164,6 +165,51 @@ def test_dommaschk_potential_arr_equal_error(): a_arr = [1] b_arr = [1] c_arr = [1] - d_arr = [1, 1] + d_arr = [1, 1] # length is not equal to the rest shuld with pytest.raises(AssertionError): dommaschk_potential(R, Z, phi, ms, ls, a_arr, b_arr, c_arr, d_arr) + + +@pytest.mark.unit +def test_dommaschk_radial_field(): + """Test the Dommaschk potential for a pure toroidal (Bphi~1/R) field.""" + phi = np.linspace(0, 2 * np.pi, 10) + R = np.linspace(0.1, 1.5, 50) + Z = np.linspace(-0.05, 0.5, 50) + R, phi, Z = np.meshgrid(R, phi, Z) + coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T + + ms = [0] + ls = [0] + a_arr = [0] + b_arr = [0] + c_arr = [0] + d_arr = [0] + B = DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) + B_dom = B.compute_magnetic_field(coords) + np.testing.assert_allclose(B_dom[:, 0], 0) + np.testing.assert_array_equal(B_dom[:, 1], 1 / R.flatten()) + np.testing.assert_allclose(B_dom[:, 2], 0) + + +@pytest.mark.unit +def test_dommaschk_vertical_field(): + """Test the Dommaschk potential for a 1/R toroidal + pure vertical field.""" + phi = np.linspace(0, 2 * np.pi, 10) + R = np.linspace(0.1, 1.5, 50) + Z = np.linspace(-0.05, 0.5, 50) + R, phi, Z = np.meshgrid(R, phi, Z) + coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T + + ms = [0] + ls = [1] + a_arr = [1] + b_arr = [0] + c_arr = [0] + d_arr = [0] + B = DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) + B_dom = B.compute_magnetic_field(coords) + ones = np.ones_like(B_dom[:, 0]) + np.testing.assert_allclose(B_dom[:, 0], 0) + np.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten()) + np.testing.assert_array_equal(B_dom[:, 2], ones) From 9ce422a9f5c7590f0eaaaef9a00153fd161226fc Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 11 Apr 2023 15:25:00 -0400 Subject: [PATCH 05/46] fix typo in args of V_m_l causing test to fail --- desc/magnetic_fields.py | 2 +- tests/test_magnetic_fields.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 1d4d47fc5a..0a9008cf03 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -931,7 +931,7 @@ def N_m_n(R, Z, m, n): # if stell sym then a=d=0 for even l and b=c=0 for odd l -def V_m_l(R, Z, phi, m, l, a, b, c, d): +def V_m_l(R, phi, Z, m, l, a, b, c, d): """Eq 12 of Dommaschk paper. Parameters diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 4f64f3372c..9f29126b52 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -197,7 +197,7 @@ def test_dommaschk_vertical_field(): """Test the Dommaschk potential for a 1/R toroidal + pure vertical field.""" phi = np.linspace(0, 2 * np.pi, 10) R = np.linspace(0.1, 1.5, 50) - Z = np.linspace(-0.05, 0.5, 50) + Z = np.linspace(-0.5, 0.5, 50) R, phi, Z = np.meshgrid(R, phi, Z) coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T From c22b7020251a25c0806de4b607803e96f9a2a7d1 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 11 Apr 2023 15:25:54 -0400 Subject: [PATCH 06/46] update docstring --- desc/magnetic_fields.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 0a9008cf03..e138329164 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -936,7 +936,7 @@ def V_m_l(R, phi, Z, m, l, a, b, c, d): Parameters ---------- - R,Z,phi : array-like + R,phi,Z : array-like Cylindrical coordinates (1-D arrays of each of size num_eval_pts) to evaluate the Dommaschk potential term at. m : int @@ -973,7 +973,7 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr): Parameters ---------- - R,Z,phi : array-like + R,phi,Z : array-like Cylindrical coordinates (1-D arrays of each of size num_eval_pts) to evaluate the Dommaschk potential term at. ms : 1D array-like of int From 67d3bbcf9d46ccbe9c4fc8ae83bb0f0c6f31e35c Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 11 Apr 2023 15:26:22 -0400 Subject: [PATCH 07/46] update arg in test --- tests/test_magnetic_fields.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 9f29126b52..344c121838 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -167,7 +167,7 @@ def test_dommaschk_potential_arr_equal_error(): c_arr = [1] d_arr = [1, 1] # length is not equal to the rest shuld with pytest.raises(AssertionError): - dommaschk_potential(R, Z, phi, ms, ls, a_arr, b_arr, c_arr, d_arr) + dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr) @pytest.mark.unit From 9a186713dba08da3d39bb12de30b705cac1b64ad Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 20 Apr 2023 02:58:56 -0400 Subject: [PATCH 08/46] add function to fit a magnetic field with Dommaschk potentials --- desc/magnetic_fields.py | 95 ++++++++++++++++++++++++++++++++--- tests/test_magnetic_fields.py | 32 ++++++++++++ 2 files changed, 121 insertions(+), 6 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index e138329164..9f31f89ec1 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod import numpy as np +import scipy.linalg from netCDF4 import Dataset from desc.backend import jit, jnp, odeint @@ -769,11 +770,12 @@ class DommaschkPotentialField(ScalarPotentialField): c_m_l coefficients of V_m_l terms, which multiply the cos(m*phi)*N_m_l-1 term d_arr : 1D array-like of float d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms - + B0: float + scale strength of the magnetic field's 1/R portion """ - def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr): + def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1.0): params = {} params["ms"] = ms params["ls"] = ls @@ -781,9 +783,86 @@ def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr): params["b_arr"] = b_arr params["c_arr"] = c_arr params["d_arr"] = d_arr + params["B0"] = B0 super().__init__(dommaschk_potential, params) + @classmethod + def fit_magnetic_field(cls, field, coords, max_m, max_l): + """Fit a vacuum magnetic field with a Dommaschk Potential field. + + Parameters + ---------- + field (MagneticField or callable): magnetic field to fit + 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_l (int): maximum l to use for Dommaschk Potentials + """ + if not isinstance(field, MagneticField): + B = field(coords) + else: + B = field.compute_magnetic_field(coords) + + num_nodes = coords.shape[0] # number of coordinate nodes + # we will have the rhs be 3*num_nodes in length (bc of vector B) + + rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") + # b is made, now do A + + num_modes = 1 + for l in range(max_l + 1): + for m in range(max_m + 1): + num_modes += 4 + # FIXME: do this properly with combinations + + A = jnp.zeros((3 * num_nodes, num_modes)) + + # c will be [B0, a_00, a_10, a_01, a_11... etc] + + # B0 first, the scale magnitude of the 1/R field + Bf_temp = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) + B_temp = Bf_temp.compute_magnetic_field(coords) + for i in range(B_temp.shape[1]): + A = A.at[i * num_nodes : (i + 1) * num_nodes, 0:1].add( + B_temp[:, i].reshape(num_nodes, 1) + ) + + # mode numbers + ms = [] + ls = [] + + # order of coeffs are a_ml and + coef_ind = 1 + for l in range(max_l + 1): + for m in range(max_m + 1): + ms.append(m) + ls.append(l) + for which_coef in range(4): + a = 1 if which_coef == 0 else 0 + b = 1 if which_coef == 1 else 0 + c = 1 if which_coef == 2 else 0 + d = 1 if which_coef == 3 else 0 + + Bf_temp = DommaschkPotentialField(m, l, a, b, c, d, 0) + B_temp = Bf_temp.compute_magnetic_field(coords) + for i in range(B_temp.shape[1]): + A = A.at[ + i * num_nodes : (i + 1) * num_nodes, coef_ind : coef_ind + 1 + ].add(B_temp[:, i].reshape(num_nodes, 1)) + coef_ind += 1 + # now solve Ac=b for the coefficients c + + Ainv = scipy.linalg.pinv(A, rcond=None) + c = jnp.matmul(Ainv, rhs) + print(c.size) + + B0 = c[0] + a_arr = c[1::4] + b_arr = c[2::4] + c_arr = c[3::4] + d_arr = c[4::4] + return cls(ms, ls, a_arr, b_arr, c_arr, d_arr, B0) + def field_line_integrate( r0, z0, phis, field, params={}, rtol=1e-8, atol=1e-8, maxstep=1000 @@ -851,7 +930,7 @@ def gamma(n): def alpha(m, n): """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" - return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2 ** (2 * n + m)) + return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) def alphastar(m, n): @@ -861,7 +940,7 @@ def alphastar(m, n): def beta(m, n): """Beta of eq 28.""" - return gamma(m - n) / (gamma(n + 1) * 2 ** (2 * n - m + 1)) + return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) def betastar(m, n): @@ -964,7 +1043,7 @@ def V_m_l(R, phi, Z, m, l, a, b, c, d): ) * N_m_n(R, Z, m, l - 1) -def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr): +def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): """Eq 1 of Dommaschk paper. this is the total dommaschk potential for @@ -996,7 +1075,11 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr): at the given R,phi,Z points (same size as the size of the given R,phi, or Z arrays). """ - value = phi # phi term + # TODO: think of if this is best way to generalize + # B0 is what scales the 1/R part of the magnetic field, + # and is the magnetic field strength at R=1 + + value = B0 * phi # phi term # make sure all are 1D arrays ms = jnp.atleast_1d(ms) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 344c121838..380cd23863 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -213,3 +213,35 @@ def test_dommaschk_vertical_field(): np.testing.assert_allclose(B_dom[:, 0], 0) np.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten()) np.testing.assert_array_equal(B_dom[:, 2], ones) + + +@pytest.mark.unit +def test_dommaschkfit_field(): + """Test the Dommaschk potential fit for a 1/R toroidal scaled to 2 T.""" + phi = np.linspace(0, 2 * np.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 = 1 + max_m = 1 + 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) + + B_dom = B.compute_magnetic_field(coords) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=3e-15) + np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-15) + np.testing.assert_allclose(B_dom[:, 2], np.zeros_like(R.flatten()), atol=1e-15) + + # only nonzero coefficient of the field should be the B0 + np.testing.assert_allclose(B._params["B0"], B0, atol=1e-15) + for coef in ["a_arr", "b_arr", "c_arr", "d_arr"]: + np.testing.assert_allclose(B._params[coef], 0, atol=1e-15) From 060ac5406f8f66b554b0d68e814204247ba32876 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 25 Apr 2023 00:43:56 -0400 Subject: [PATCH 09/46] fix typo in Neumann term of Dommaschk potential. Add assertions to ensure summations are only being used where they are valid --- desc/magnetic_fields.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 9f31f89ec1..24c72fe0fa 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -948,11 +948,14 @@ def betastar(m, n): return (2 * n - m) * beta(m, n) +# TODO: generalize the below 2 terms according to +# eqns 31,32 so they are valid for all indices of m,k>=0? def CD_m_k(R, m, k): """Eq 25 of Dommaschk paper.""" if m == k == 0: # the initial term, defined in eqn 8 return jnp.ones_like(R) - + else: + assert m > k, "Sum only valid for m > k!" sum1 = 0 for j in range(k + 1): sum1 += alpha(m, j) * betastar(m, k - j) * R ** (2 * j) @@ -966,6 +969,8 @@ def CN_m_k(R, m, k): """Eq 26 of Dommaschk paper.""" if m == k == 0: # the initial term, defined in eqn 9 return jnp.log(R) + else: + assert m > k, "Sum only valid for m > k!" sum1 = 0 for j in range(k + 1): @@ -993,7 +998,7 @@ def D_m_n(R, Z, m, n): def N_m_n(R, Z, m, n): """N_m_n term in eqn 9 of Dommaschk paper.""" - # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper + # the sum comes from fact that N_mn = I_mn and the def of I_mn in eq 2 of the paper max_ind = ( n // 2 @@ -1002,7 +1007,7 @@ def N_m_n(R, Z, m, n): # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum result = 0 for k in range(max_ind + 1): - result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) + result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CN_m_k(R, m, k) return result @@ -1067,18 +1072,16 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): c_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*N_m_l-1 d_arr : 1D array-like of float d_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*N_m_l-1 + B0: float, toroidal magnetic field strength scale, this is the strength of the + 1/R part of the magnetic field and is the Bphi at R=1. Returns ------- value : array-like Value of the total dommaschk potential evaluated at the given R,phi,Z points - (same size as the size of the given R,phi, or Z arrays). + (same size as the size of the given R,phi, Z arrays). """ - # TODO: think of if this is best way to generalize - # B0 is what scales the 1/R part of the magnetic field, - # and is the magnetic field strength at R=1 - value = B0 * phi # phi term # make sure all are 1D arrays From 797797d8a777080fa5600372cbe63622941800fc Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 25 Apr 2023 00:56:08 -0400 Subject: [PATCH 10/46] add test for dommaschk fitting using MagneticField class --- tests/test_magnetic_fields.py | 37 ++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 380cd23863..3cfa856e53 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -216,7 +216,7 @@ def test_dommaschk_vertical_field(): @pytest.mark.unit -def test_dommaschkfit_field(): +def test_dommaschk_fit_toridal_field(): """Test the Dommaschk potential fit for a 1/R toroidal scaled to 2 T.""" phi = np.linspace(0, 2 * np.pi, 3) R = np.linspace(0.1, 1.5, 3) @@ -245,3 +245,38 @@ def B0_over_R(coord): np.testing.assert_allclose(B._params["B0"], B0, atol=1e-15) for coef in ["a_arr", "b_arr", "c_arr", "d_arr"]: np.testing.assert_allclose(B._params[coef], 0, atol=1e-15) + + +@pytest.mark.unit +def test_dommaschk_fit_vertical_and_toroidal_field(): + """Test the Dommaschk potential fit for a toroidal and a vertical field.""" + phi = np.linspace(0, 2 * np.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 = 1 + max_m = 1 + B0 = 2 # scale strength for to 1/R field to fit + B0_Z = 1 # scale strength for to uniform vertical field to fit + field = ToroidalMagneticField(B0=B0, R0=1) + VerticalMagneticField(B0=B0_Z) + + B = DommaschkPotentialField.fit_magnetic_field(field, coords, max_m, max_l) + + B_dom = B.compute_magnetic_field(coords) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=3e-15) + np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-15) + np.testing.assert_allclose(B_dom[:, 2], B0_Z, atol=1e-15) + + np.testing.assert_allclose(B._params["B0"], B0) + + # only nonzero coefficient of the field should be the B0 and a_ml = a_01 + np.testing.assert_allclose(B._params["B0"], B0, atol=1e-15) + for coef, m, l in zip(B._params["a_arr"], B._params["ms"], B._params["ls"]): + if m == 0 and l == 1: + np.testing.assert_allclose(coef, B0_Z) + else: + np.testing.assert_allclose(coef, 0, atol=1e-15) + for name in ["b_arr", "c_arr", "d_arr"]: + np.testing.assert_allclose(B._params[name], 0, atol=1e-15) From a3137c15c6f0923a0eb05cc86f9a839230d0fa90 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 25 Apr 2023 22:16:40 -0400 Subject: [PATCH 11/46] add documentation and update dommaschk function to more succinctly calculate number of modes --- desc/magnetic_fields.py | 47 +++++++++++++++++++++++++++++++---------- 1 file changed, 36 insertions(+), 11 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 24c72fe0fa..240d2d63b0 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1,5 +1,6 @@ """Classes for magnetic fields.""" +import warnings from abc import ABC, abstractmethod import numpy as np @@ -798,27 +799,43 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): max_m (int): maximum m to use for Dommaschk Potentials max_l (int): maximum l to use for Dommaschk Potentials """ + # We seek c in Ac = b + # A will be the BR, Bphi and BZ from each individual + # dommaschk potential basis function evaluated at each node + # c is the dommaschk potential coefficients + # c will be [B0, a_00, a_10, a_01, a_11... etc] + # b is the magnetic field at each node which we are fitting + if not isinstance(field, MagneticField): B = field(coords) else: B = field.compute_magnetic_field(coords) num_nodes = coords.shape[0] # number of coordinate nodes + # we will have the rhs be 3*num_nodes in length (bc of vector B) + ######### + # make b + ######### + rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") - # b is made, now do A - num_modes = 1 - for l in range(max_l + 1): - for m in range(max_m + 1): - num_modes += 4 - # FIXME: do this properly with combinations + ##################### + # b is made, now do A + ##################### + num_modes = 1 + (max_l + 1) * (max_m + 1) * 4 + # TODO: technically we can drop some modes + # since if max_l=0, there are only ever nonzero terms + # for a and b + # and if max_m=0, there are only ever nonzero terms + # for a and c + # but since we are only fitting in a least squares sense, + # and max_l and max_m should probably be both nonzero anyways, + # this is not an issue right now A = jnp.zeros((3 * num_nodes, num_modes)) - # c will be [B0, a_00, a_10, a_01, a_11... etc] - # B0 first, the scale magnitude of the 1/R field Bf_temp = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) B_temp = Bf_temp.compute_magnetic_field(coords) @@ -850,12 +867,18 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): i * num_nodes : (i + 1) * num_nodes, coef_ind : coef_ind + 1 ].add(B_temp[:, i].reshape(num_nodes, 1)) coef_ind += 1 + # now solve Ac=b for the coefficients c Ainv = scipy.linalg.pinv(A, rcond=None) c = jnp.matmul(Ainv, rhs) - print(c.size) + res = jnp.matmul(A, c) - rhs + print(f"Mean Residual of fit: {jnp.mean(res):1.4e} T") + print(f"Max Residual of fit: {jnp.max(res):1.4e} T") + print(f"Min Residual of fit: {jnp.min(res):1.4e} T") + + # recover the params from the c coefficient vector B0 = c[0] a_arr = c[1::4] b_arr = c[2::4] @@ -955,7 +978,8 @@ def CD_m_k(R, m, k): if m == k == 0: # the initial term, defined in eqn 8 return jnp.ones_like(R) else: - assert m > k, "Sum only valid for m > k!" + if m > k: + warnings.warn("Sum only valid for m > k!") sum1 = 0 for j in range(k + 1): sum1 += alpha(m, j) * betastar(m, k - j) * R ** (2 * j) @@ -970,7 +994,8 @@ def CN_m_k(R, m, k): if m == k == 0: # the initial term, defined in eqn 9 return jnp.log(R) else: - assert m > k, "Sum only valid for m > k!" + if m > k: + warnings.warn("Sum only valid for m > k!") sum1 = 0 for j in range(k + 1): From d0e4d13521a2ac7671196a90a80a092e6b693b11 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sat, 29 Apr 2023 10:44:41 -0400 Subject: [PATCH 12/46] implement generalized formulas from eqns 31 and 32 which are valid for any range of m,l >=0 --- desc/magnetic_fields.py | 84 +++++++++++++++++++++++------------ tests/test_magnetic_fields.py | 10 +++-- 2 files changed, 63 insertions(+), 31 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 240d2d63b0..a839387fb5 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1,6 +1,5 @@ """Classes for magnetic fields.""" -import warnings from abc import ABC, abstractmethod import numpy as np @@ -953,57 +952,83 @@ def gamma(n): def alpha(m, n): """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + # modified for eqns 31 and 32 + if n < 0: + return 0 return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) def alphastar(m, n): """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + # modified for eqns 31 and 32 + if n < 0: + return 0 return (2 * n + m) * alpha(m, n) def beta(m, n): - """Beta of eq 28.""" + """Beta of eq 28, modified for eqns 31 and 32.""" + if n < 0 or n >= m: + return 0 return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) def betastar(m, n): - """Betastar of eq 28.""" + """Betastar of eq 28, modified for eqns 31 and 32.""" + if n < 0 or n >= m: + return 0 return (2 * n - m) * beta(m, n) -# TODO: generalize the below 2 terms according to -# eqns 31,32 so they are valid for all indices of m,k>=0? +def gamma_n(m, n): + """gamma_n of eq 33.""" + if n <= 0: + return 0 + return ( + alpha(m, n) / 2 * jnp.sum(jnp.array([1 / i + 1 / (m + i) for i in range(1, n)])) + ) + + +def gamma_nstar(m, n): + """gamma_n star of eq 33.""" + if n <= 0: + return 0 + return (2 * n + m) * gamma_n(m, n) + + def CD_m_k(R, m, k): - """Eq 25 of Dommaschk paper.""" - if m == k == 0: # the initial term, defined in eqn 8 - return jnp.ones_like(R) - else: - if m > k: - warnings.warn("Sum only valid for m > k!") + """Eq 31 of Dommaschk paper.""" sum1 = 0 for j in range(k + 1): - sum1 += alpha(m, j) * betastar(m, k - j) * R ** (2 * j) - sum2 = 0 - for j in range(k + 1): - sum2 += alphastar(m, k - j) * beta(m, j) * R ** (2 * j) - return -(R ** (m)) * sum1 + R ** (-m) * sum2 + sum1 += ( + -( + alpha(m, j) + * ( + alphastar(m, k - m - j) * jnp.log(R) + + gamma_nstar(m, k - m - j) + - alpha(m, k - m - j) + ) + - gamma_n(m, j) * alphastar(m, k - m - j) + + alpha(m, j) * betastar(m, k - j) + ) + * R ** (2 * j + m) + ) + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) + return sum1 def CN_m_k(R, m, k): - """Eq 26 of Dommaschk paper.""" - if m == k == 0: # the initial term, defined in eqn 9 - return jnp.log(R) - else: - if m > k: - warnings.warn("Sum only valid for m > k!") - + """Eq 32 of Dommaschk paper.""" sum1 = 0 for j in range(k + 1): - sum1 += alpha(m, j) * beta(m, k - j) * R ** (2 * j) - sum2 = 0 - for j in range(k + 1): - sum2 += alpha(m, k - j) * beta(m, j) * R ** (2 * j) - return R ** (m) * sum1 - R ** (-m) * sum2 + sum1 += ( + ( + alpha(m, j) * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) + - gamma_n(m, j) * alpha(m, k - m - j) + + alpha(m, j) * beta(m, k - j) + ) + * R ** (2 * j + m) + ) - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) + return sum1 def D_m_n(R, Z, m, n): @@ -1120,6 +1145,9 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): assert ( ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size ), "Passed in arrays must all be of the same size!" + assert not np.any( + np.logical_or(ms < 0, ls < 0) + ), "m and l modenumbers must be >= 0!" for m, l, a, b, c, d in zip(ms, ls, a_arr, b_arr, c_arr, d_arr): value += V_m_l(R, phi, Z, m, l, a, b, c, d) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 3cfa856e53..7e81e76de3 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -155,8 +155,8 @@ def test_dommaschk_CN_CD_m_0(): @pytest.mark.unit -def test_dommaschk_potential_arr_equal_error(): - """Test the assert statement of the Dommaschk potential function.""" +def test_dommaschk_potential_errors(): + """Test the assert statements of the Dommaschk potential function.""" phi = 1 R = 1 Z = 1 @@ -165,7 +165,11 @@ def test_dommaschk_potential_arr_equal_error(): a_arr = [1] b_arr = [1] c_arr = [1] - d_arr = [1, 1] # length is not equal to the rest shuld + d_arr = [1, 1] # length is not equal to the rest + with pytest.raises(AssertionError): + dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr) + d_arr = [1] + ms = [-1] # negative modenumber with pytest.raises(AssertionError): dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr) From 46c4503044d5d19070de9a32bc28b97bed7f0b7c Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sat, 29 Apr 2023 11:10:43 -0400 Subject: [PATCH 13/46] lower test tols --- desc/magnetic_fields.py | 6 +++--- tests/test_magnetic_fields.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index a839387fb5..a20c553013 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -873,9 +873,9 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): c = jnp.matmul(Ainv, rhs) res = jnp.matmul(A, c) - rhs - print(f"Mean Residual of fit: {jnp.mean(res):1.4e} T") - print(f"Max Residual of fit: {jnp.max(res):1.4e} T") - print(f"Min Residual of fit: {jnp.min(res):1.4e} T") + print(f"Mean Residual of fit: {jnp.mean(jnp.abs(res)):1.4e} T") + print(f"Max Residual of fit: {jnp.max(jnp.abs(res)):1.4e} T") + print(f"Min Residual of fit: {jnp.min(jnp.abs(res)):1.4e} T") # recover the params from the c coefficient vector B0 = c[0] diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 7e81e76de3..a1d7e01873 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -220,7 +220,7 @@ def test_dommaschk_vertical_field(): @pytest.mark.unit -def test_dommaschk_fit_toridal_field(): +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 * np.pi, 3) R = np.linspace(0.1, 1.5, 3) @@ -269,7 +269,7 @@ def test_dommaschk_fit_vertical_and_toroidal_field(): B = DommaschkPotentialField.fit_magnetic_field(field, coords, max_m, max_l) B_dom = B.compute_magnetic_field(coords) - np.testing.assert_allclose(B_dom[:, 0], 0, atol=3e-15) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=4e-15) np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-15) np.testing.assert_allclose(B_dom[:, 2], B0_Z, atol=1e-15) From a1487597e319a9dda36e8e6cf1518da78d07fd14 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 30 Apr 2023 00:34:16 -0400 Subject: [PATCH 14/46] change test tol --- tests/test_magnetic_fields.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index a1d7e01873..eda05c60c5 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -214,8 +214,8 @@ def test_dommaschk_vertical_field(): B = DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) B_dom = B.compute_magnetic_field(coords) ones = np.ones_like(B_dom[:, 0]) - np.testing.assert_allclose(B_dom[:, 0], 0) - np.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten()) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-15) + np.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten(), atol=1e-15) np.testing.assert_array_equal(B_dom[:, 2], ones) @@ -241,7 +241,7 @@ def B0_over_R(coord): B = DommaschkPotentialField.fit_magnetic_field(B0_over_R, coords, max_m, max_l) B_dom = B.compute_magnetic_field(coords) - np.testing.assert_allclose(B_dom[:, 0], 0, atol=3e-15) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=4e-15) np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-15) np.testing.assert_allclose(B_dom[:, 2], np.zeros_like(R.flatten()), atol=1e-15) From e03fd6fa2be0f49b9852e8c06b6c0cfdfe17bb91 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 11 Jun 2023 13:02:04 -0400 Subject: [PATCH 15/46] move checks to Dommaschk class, and make some ranges into jnp ranges --- desc/magnetic_fields.py | 26 +++++++++++++++----------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index a20c553013..fe45c347fd 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -776,6 +776,15 @@ class DommaschkPotentialField(ScalarPotentialField): """ def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1.0): + + assert ( + ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size + ), "Passed in arrays must all be of the same size!" + assert not jnp.any( + jnp.logical_or(ms < 0, ls < 0) + ), "m and l modenumbers must be >= 0!" + assert jnp.isscalar(B0), "B0 should be a scalar value!" + params = {} params["ms"] = ms params["ls"] = ls @@ -985,7 +994,9 @@ def gamma_n(m, n): if n <= 0: return 0 return ( - alpha(m, n) / 2 * jnp.sum(jnp.array([1 / i + 1 / (m + i) for i in range(1, n)])) + alpha(m, n) + / 2 + * jnp.sum(jnp.array([1 / i + 1 / (m + i) for i in jnp.arange(1, n)])) ) @@ -999,7 +1010,7 @@ def gamma_nstar(m, n): def CD_m_k(R, m, k): """Eq 31 of Dommaschk paper.""" sum1 = 0 - for j in range(k + 1): + for j in jnp.arange(k + 1): sum1 += ( -( alpha(m, j) @@ -1019,7 +1030,7 @@ def CD_m_k(R, m, k): def CN_m_k(R, m, k): """Eq 32 of Dommaschk paper.""" sum1 = 0 - for j in range(k + 1): + for j in jnp.arange(k + 1): sum1 += ( ( alpha(m, j) * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) @@ -1041,7 +1052,7 @@ def D_m_n(R, Z, m, n): # i.e. this is the index of k at # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum result = 0 - for k in range(max_ind + 1): + for k in jnp.arange(max_ind + 1): result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) return result @@ -1142,13 +1153,6 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): c_arr = jnp.atleast_1d(c_arr) d_arr = jnp.atleast_1d(d_arr) - assert ( - ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size - ), "Passed in arrays must all be of the same size!" - assert not np.any( - np.logical_or(ms < 0, ls < 0) - ), "m and l modenumbers must be >= 0!" - for m, l, a, b, c, d in zip(ms, ls, a_arr, b_arr, c_arr, d_arr): value += V_m_l(R, phi, Z, m, l, a, b, c, d) return value From 3c845defcf2d733e4186f27851d9e1c067f8ad09 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 20 Jun 2023 06:28:11 -0400 Subject: [PATCH 16/46] add 1d catch and check if B0 is 1d array with size 0 to avoid jnp.isscalar(jnp.array(B0)) returning error --- desc/magnetic_fields.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index fe45c347fd..a257bfb13f 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -777,13 +777,22 @@ class DommaschkPotentialField(ScalarPotentialField): def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1.0): + ms = jnp.atleast_1d(ms) + ls = jnp.atleast_1d(ls) + a_arr = jnp.atleast_1d(a_arr) + b_arr = jnp.atleast_1d(b_arr) + c_arr = jnp.atleast_1d(c_arr) + d_arr = jnp.atleast_1d(d_arr) + assert ( ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size ), "Passed in arrays must all be of the same size!" assert not jnp.any( jnp.logical_or(ms < 0, ls < 0) ), "m and l modenumbers must be >= 0!" - assert jnp.isscalar(B0), "B0 should be a scalar value!" + assert ( + jnp.isscalar(B0) or jnp.atleast_1d(B0).size == 1 + ), "B0 should be a scalar value!" params = {} params["ms"] = ms From 73a3c58c1d80c158d6a386824d6f14557da91c69 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 23 Jun 2023 20:49:03 -0400 Subject: [PATCH 17/46] move read_mgrid to its own function --- desc/magnetic_fields.py | 91 +++++++++++++++++++++++++---------------- 1 file changed, 55 insertions(+), 36 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index a257bfb13f..7263402e57 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -606,42 +606,7 @@ def from_mgrid( period in the toroidal direction (usually 2pi/NFP) """ - mgrid = Dataset(mgrid_file, "r") - ir = int(mgrid["ir"][()]) - jz = int(mgrid["jz"][()]) - kp = int(mgrid["kp"][()]) - nfp = mgrid["nfp"][()].data - nextcur = int(mgrid["nextcur"][()]) - rMin = mgrid["rmin"][()] - rMax = mgrid["rmax"][()] - zMin = mgrid["zmin"][()] - zMax = mgrid["zmax"][()] - - br = np.zeros([kp, jz, ir]) - bp = np.zeros([kp, jz, ir]) - bz = np.zeros([kp, jz, ir]) - extcur = np.broadcast_to(extcur, nextcur) - for i in range(nextcur): - - # apply scaling by currents given in VMEC input file - scale = extcur[i] - - # sum up contributions from different coils - coil_id = "%03d" % (i + 1,) - br[:, :, :] += scale * mgrid["br_" + coil_id][()] - bp[:, :, :] += scale * mgrid["bp_" + coil_id][()] - bz[:, :, :] += scale * mgrid["bz_" + coil_id][()] - mgrid.close() - - # shift axes to correct order - br = np.moveaxis(br, (0, 1, 2), (1, 2, 0)) - bp = np.moveaxis(bp, (0, 1, 2), (1, 2, 0)) - bz = np.moveaxis(bz, (0, 1, 2), (1, 2, 0)) - - # re-compute grid knots in radial and vertical direction - Rgrid = np.linspace(rMin, rMax, ir) - Zgrid = np.linspace(zMin, zMax, jz) - pgrid = 2.0 * np.pi / (nfp * kp) * np.arange(kp) + Rgrid, pgrid, Zgrid, br, bp, bz, nfp = read_mgrid(mgrid_file, extcur) if period is None: period = 2 * np.pi / (nfp) @@ -1165,3 +1130,57 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): for m, l, a, b, c, d in zip(ms, ls, a_arr, b_arr, c_arr, d_arr): value += V_m_l(R, phi, Z, m, l, a, b, c, d) return value + + +def read_mgrid( + mgrid_file, + extcur=1, +): + """Read an "mgrid" file from MAKEGRID and return the grid and magnetic field. + + Parameters + ---------- + mgrid_file : str or path-like + path to mgrid file in netCDF format + extcur : array-like + currents for each subset of the field + + """ + mgrid = Dataset(mgrid_file, "r") + ir = int(mgrid["ir"][()]) + jz = int(mgrid["jz"][()]) + kp = int(mgrid["kp"][()]) + nfp = mgrid["nfp"][()].data + nextcur = int(mgrid["nextcur"][()]) + rMin = mgrid["rmin"][()] + rMax = mgrid["rmax"][()] + zMin = mgrid["zmin"][()] + zMax = mgrid["zmax"][()] + + br = np.zeros([kp, jz, ir]) + bp = np.zeros([kp, jz, ir]) + bz = np.zeros([kp, jz, ir]) + extcur = np.broadcast_to(extcur, nextcur) + for i in range(nextcur): + + # apply scaling by currents given in VMEC input file + scale = extcur[i] + + # sum up contributions from different coils + coil_id = "%03d" % (i + 1,) + br[:, :, :] += scale * mgrid["br_" + coil_id][()] + bp[:, :, :] += scale * mgrid["bp_" + coil_id][()] + bz[:, :, :] += scale * mgrid["bz_" + coil_id][()] + mgrid.close() + + # shift axes to correct order + br = np.moveaxis(br, (0, 1, 2), (1, 2, 0)) + bp = np.moveaxis(bp, (0, 1, 2), (1, 2, 0)) + bz = np.moveaxis(bz, (0, 1, 2), (1, 2, 0)) + + # re-compute grid knots in radial and vertical direction + Rgrid = np.linspace(rMin, rMax, ir) + Zgrid = np.linspace(zMin, zMax, jz) + pgrid = 2.0 * np.pi / (nfp * kp) * np.arange(kp) + + return Rgrid, Zgrid, pgrid, br, bp, bz, nfp From e64bea285a6b4d4ff14789d7838867d3eeddffb5 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sat, 24 Jun 2023 18:18:04 -0400 Subject: [PATCH 18/46] fix test, add check for ints since DommaschkField as written cannot handle integer coordinates --- desc/magnetic_fields.py | 2 ++ tests/test_magnetic_fields.py | 12 ++++-------- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 7263402e57..fd8871c6b3 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -693,6 +693,8 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz"): if isinstance(coords, Grid): coords = coords.nodes coords = jnp.atleast_2d(coords) + if coords.dtype == int: + coords = coords.astype(float) if basis == "xyz": coords = xyz2rpz(coords) Rq, phiq, Zq = coords.T diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index eda05c60c5..55e14b1d5d 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -13,7 +13,6 @@ SplineMagneticField, ToroidalMagneticField, VerticalMagneticField, - dommaschk_potential, field_line_integrate, ) @@ -155,11 +154,8 @@ def test_dommaschk_CN_CD_m_0(): @pytest.mark.unit -def test_dommaschk_potential_errors(): - """Test the assert statements of the Dommaschk potential function.""" - phi = 1 - R = 1 - Z = 1 +def test_dommaschk_field_errors(): + """Test the assert statements of the DommaschkField function.""" ms = [1] ls = [1] a_arr = [1] @@ -167,11 +163,11 @@ def test_dommaschk_potential_errors(): c_arr = [1] d_arr = [1, 1] # length is not equal to the rest with pytest.raises(AssertionError): - dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr) + DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) d_arr = [1] ms = [-1] # negative modenumber with pytest.raises(AssertionError): - dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr) + DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) @pytest.mark.unit From ea2d6fd32349fc406c93ca91b17e8d949978cbfb Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 27 Jun 2023 19:56:12 -0400 Subject: [PATCH 19/46] attempt to jaxify dommaschk --- desc/magnetic_fields.py | 199 +++++++++++++++++++++++++++++----------- 1 file changed, 144 insertions(+), 55 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index fd8871c6b3..b58bc83dca 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -6,7 +6,7 @@ import scipy.linalg from netCDF4 import Dataset -from desc.backend import jit, jnp, odeint +from desc.backend import cond, fori_loop, gammaln, jit, jnp, odeint from desc.derivatives import Derivative from desc.geometry.utils import rpz2xyz_vec, xyz2rpz from desc.grid import Grid @@ -697,7 +697,6 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz"): coords = coords.astype(float) if basis == "xyz": coords = xyz2rpz(coords) - Rq, phiq, Zq = coords.T if (params is None) or (len(params) == 0): params = self._params @@ -929,93 +928,176 @@ def odefun(rpz, s): # written with naive for loops initially and can jaxify later +true_fun = lambda m_n: 0.0 # used for returning 0 when conditionals evaluate to True + def gamma(n): """Gamma function, only implemented for integers (equiv to factorial of (n-1)).""" - return jnp.prod(jnp.arange(1, n)) + return jnp.exp(gammaln(n)) def alpha(m, n): """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" # modified for eqns 31 and 32 - if n < 0: - return 0 - return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) + # if n < 0: + # return 0 + + def false_fun(m_n): + m, n = m_n + return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) + + def bool_fun(n): + return jnp.any(n < 0) + + return cond( + bool_fun(n), + true_fun, + false_fun, + ( + m, + n, + ), + ) def alphastar(m, n): """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" # modified for eqns 31 and 32 - if n < 0: - return 0 - return (2 * n + m) * alpha(m, n) + # if n < 0: + # return 0 + # return (2 * n + m) * alpha(m, n) + + def false_fun(m_n): + m, n = m_n + return (2 * n + m) * alpha(m, n) + + return cond(jnp.any(n < 0), true_fun, false_fun, (m, n)) def beta(m, n): """Beta of eq 28, modified for eqns 31 and 32.""" - if n < 0 or n >= m: - return 0 - return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) + # if n < 0 or n >= m: + # return 0 + # return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) + def false_fun(m_n): + m, n = m_n + return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) + + return cond(jnp.any(jnp.logical_or(n < 0, n >= m)), true_fun, false_fun, (m, n)) def betastar(m, n): """Betastar of eq 28, modified for eqns 31 and 32.""" - if n < 0 or n >= m: - return 0 - return (2 * n - m) * beta(m, n) + # if n < 0 or n >= m: + # return 0 + # return (2 * n - m) * beta(m, n) + def false_fun(m_n): + m, n = m_n + return (2 * n - m) * beta(m, n) + + return cond(jnp.any(jnp.logical_or(n < 0, n >= m)), true_fun, false_fun, (m, n)) def gamma_n(m, n): """gamma_n of eq 33.""" - if n <= 0: - return 0 - return ( - alpha(m, n) - / 2 - * jnp.sum(jnp.array([1 / i + 1 / (m + i) for i in jnp.arange(1, n)])) - ) + # if n <= 0: + # return 0 + + def body_fun(i, val): + return val + 1 / i + 1 / (m + i) + + # temp = fori_loop(1, n, body_fun, 0) + # # jnp.sum(jnp.array([1 / i + 1 / (m + i) for i in range(1, n)])) + # return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) + def false_fun(m_n): + m, n = m_n + return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) + + return cond(jnp.any(n <= 0), true_fun, false_fun, (m, n)) def gamma_nstar(m, n): """gamma_n star of eq 33.""" - if n <= 0: - return 0 - return (2 * n + m) * gamma_n(m, n) + # if n <= 0: + # return 0 + # return (2 * n + m) * gamma_n(m, n) + def false_fun(m_n): + m, n = m_n + return (2 * n + m) * gamma_n(m, n) + + return cond(jnp.any(n <= 0), true_fun, false_fun, (m, n)) def CD_m_k(R, m, k): """Eq 31 of Dommaschk paper.""" - sum1 = 0 - for j in jnp.arange(k + 1): - sum1 += ( - -( - alpha(m, j) - * ( - alphastar(m, k - m - j) * jnp.log(R) - + gamma_nstar(m, k - m - j) - - alpha(m, k - m - j) + # sum1 = 0 + # for j in range(k + 1): + # sum1 += ( + # -( + # alpha(m, j) + # * ( + # alphastar(m, k - m - j) * jnp.log(R) + # + gamma_nstar(m, k - m - j) + # - alpha(m, k - m - j) + # ) + # - gamma_n(m, j) * alphastar(m, k - m - j) + # + alpha(m, j) * betastar(m, k - j) + # ) + # * R ** (2 * j + m) + # ) + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) + def body_fun(j, val): + result = ( + val + + ( + -( + alpha(m, j) + * ( + alphastar(m, k - m - j) * jnp.log(R) + + gamma_nstar(m, k - m - j) + - alpha(m, k - m - j) + ) + - gamma_n(m, j) * alphastar(m, k - m - j) + + alpha(m, j) * betastar(m, k - j) ) - - gamma_n(m, j) * alphastar(m, k - m - j) - + alpha(m, j) * betastar(m, k - j) + * R ** (2 * j + m) ) - * R ** (2 * j + m) - ) + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) - return sum1 + + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) + ) + return result + + return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) def CN_m_k(R, m, k): """Eq 32 of Dommaschk paper.""" - sum1 = 0 - for j in jnp.arange(k + 1): - sum1 += ( - ( - alpha(m, j) * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) - - gamma_n(m, j) * alpha(m, k - m - j) - + alpha(m, j) * beta(m, k - j) + # sum1 = 0 + # for j in range(k + 1): + # sum1 += ( + # ( + # alpha(m, j) * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) + # - gamma_n(m, j) * alpha(m, k - m - j) + # + alpha(m, j) * beta(m, k - j) + # ) + # * R ** (2 * j + m) + # ) - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) + + def body_fun(j, val): + result = ( + val + + ( + ( + alpha(m, j) + * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) + - gamma_n(m, j) * alpha(m, k - m - j) + + alpha(m, j) * beta(m, k - j) + ) + * R ** (2 * j + m) ) - * R ** (2 * j + m) - ) - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) - return sum1 + - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) + ) + return result + + return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) def D_m_n(R, Z, m, n): @@ -1028,9 +1110,13 @@ def D_m_n(R, Z, m, n): # i.e. this is the index of k at # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum result = 0 - for k in jnp.arange(max_ind + 1): - result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) - return result + + def body_fun(k, val): + return val + Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) + + # for k in range(n // 2 + 1): + # result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) + return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) def N_m_n(R, Z, m, n): @@ -1043,9 +1129,12 @@ def N_m_n(R, Z, m, n): # i.e. this is the index of k at # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum result = 0 - for k in range(max_ind + 1): - result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CN_m_k(R, m, k) - return result + # for k in range(n // 2 + 1): + # result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CN_m_k(R, m, k) + def body_fun(k, val): + return val + Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CN_m_k(R, m, k) + + return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) # TODO: note on how stell sym affects, i.e. From ed696320452069a840e8d1d6c63974d5ae342cc8 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 27 Jun 2023 20:27:58 -0400 Subject: [PATCH 20/46] tests pass but fitting is very slow now --- tests/test_magnetic_fields.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 55e14b1d5d..f4e4df2c57 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -150,7 +150,7 @@ def test_dommaschk_CN_CD_m_0(): # test of CN_m_k based off eqn 9 res1 = CN_m_k(R, m, 0) res2 = 0.5 * (R**m - R ** (-m)) / m - np.testing.assert_allclose(res1, res2) + np.testing.assert_allclose(res1, res2, atol=1e-15) @pytest.mark.unit @@ -195,11 +195,11 @@ def test_dommaschk_radial_field(): @pytest.mark.unit def test_dommaschk_vertical_field(): """Test the Dommaschk potential for a 1/R toroidal + pure vertical field.""" - phi = np.linspace(0, 2 * np.pi, 10) - R = np.linspace(0.1, 1.5, 50) - Z = np.linspace(-0.5, 0.5, 50) - R, phi, Z = np.meshgrid(R, phi, Z) - coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T + phi = jnp.linspace(0, 2 * jnp.pi, 10) + R = jnp.linspace(0.1, 1.5, 50) + Z = jnp.linspace(-0.5, 0.5, 50) + R, phi, Z = jnp.meshgrid(R, phi, Z) + coords = jnp.vstack((R.flatten(), phi.flatten(), Z.flatten())).T ms = [0] ls = [1] @@ -209,16 +209,16 @@ def test_dommaschk_vertical_field(): d_arr = [0] B = DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) B_dom = B.compute_magnetic_field(coords) - ones = np.ones_like(B_dom[:, 0]) - np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-15) - np.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten(), atol=1e-15) - np.testing.assert_array_equal(B_dom[:, 2], ones) + ones = jnp.ones_like(B_dom[:, 0]) + jnp.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) + jnp.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten(), atol=1e-14) + jnp.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 * np.pi, 3) + 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) @@ -239,7 +239,7 @@ def B0_over_R(coord): B_dom = B.compute_magnetic_field(coords) np.testing.assert_allclose(B_dom[:, 0], 0, atol=4e-15) np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-15) - np.testing.assert_allclose(B_dom[:, 2], np.zeros_like(R.flatten()), atol=1e-15) + np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=1e-15) # only nonzero coefficient of the field should be the B0 np.testing.assert_allclose(B._params["B0"], B0, atol=1e-15) From 3307c638f94e943a3eb37e43987a5d3497113f2f Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 27 Jun 2023 20:46:57 -0400 Subject: [PATCH 21/46] remove unneeded class creations in fitting, but still is very slow due to xla backend compile? --- desc/magnetic_fields.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index b58bc83dca..7b69d04c84 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -778,6 +778,8 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): Parameters ---------- field (MagneticField or callable): magnetic field to fit + if callable, must accept (num_nodes,3) ndarray as argument + and output (num_nodes,3) as the B field in cylindrical rpz basis. 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_l (int): maximum l to use for Dommaschk Potentials @@ -820,8 +822,8 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): A = jnp.zeros((3 * num_nodes, num_modes)) # B0 first, the scale magnitude of the 1/R field - Bf_temp = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) - B_temp = Bf_temp.compute_magnetic_field(coords) + domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) + B_temp = domm_field.compute_magnetic_field(coords) for i in range(B_temp.shape[1]): A = A.at[i * num_nodes : (i + 1) * num_nodes, 0:1].add( B_temp[:, i].reshape(num_nodes, 1) @@ -843,8 +845,16 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): c = 1 if which_coef == 2 else 0 d = 1 if which_coef == 3 else 0 - Bf_temp = DommaschkPotentialField(m, l, a, b, c, d, 0) - B_temp = Bf_temp.compute_magnetic_field(coords) + params = { + "ms": m, + "ls": l, + "a_arr": a, + "b_arr": b, + "c_arr": c, + "d_arr": d, + "B0": 0, + } + B_temp = domm_field.compute_magnetic_field(coords, params=params) for i in range(B_temp.shape[1]): A = A.at[ i * num_nodes : (i + 1) * num_nodes, coef_ind : coef_ind + 1 From cb23ebd41ef36934ce52310c7c31bb9d5fbb2c60 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 27 Jun 2023 23:15:43 -0400 Subject: [PATCH 22/46] jaxify all domaschk potential things further so that it can be jit'd. This does make the fitting much slower though --- desc/magnetic_fields.py | 66 +++-------------------------------------- 1 file changed, 4 insertions(+), 62 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 7b69d04c84..ed9dc2dec3 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -949,8 +949,6 @@ def gamma(n): def alpha(m, n): """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" # modified for eqns 31 and 32 - # if n < 0: - # return 0 def false_fun(m_n): m, n = m_n @@ -973,10 +971,6 @@ def bool_fun(n): def alphastar(m, n): """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" # modified for eqns 31 and 32 - # if n < 0: - # return 0 - # return (2 * n + m) * alpha(m, n) - def false_fun(m_n): m, n = m_n return (2 * n + m) * alpha(m, n) @@ -986,9 +980,7 @@ def false_fun(m_n): def beta(m, n): """Beta of eq 28, modified for eqns 31 and 32.""" - # if n < 0 or n >= m: - # return 0 - # return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) + def false_fun(m_n): m, n = m_n return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) @@ -998,9 +990,7 @@ def false_fun(m_n): def betastar(m, n): """Betastar of eq 28, modified for eqns 31 and 32.""" - # if n < 0 or n >= m: - # return 0 - # return (2 * n - m) * beta(m, n) + def false_fun(m_n): m, n = m_n return (2 * n - m) * beta(m, n) @@ -1010,15 +1000,10 @@ def false_fun(m_n): def gamma_n(m, n): """gamma_n of eq 33.""" - # if n <= 0: - # return 0 def body_fun(i, val): return val + 1 / i + 1 / (m + i) - # temp = fori_loop(1, n, body_fun, 0) - # # jnp.sum(jnp.array([1 / i + 1 / (m + i) for i in range(1, n)])) - # return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) def false_fun(m_n): m, n = m_n return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) @@ -1028,9 +1013,7 @@ def false_fun(m_n): def gamma_nstar(m, n): """gamma_n star of eq 33.""" - # if n <= 0: - # return 0 - # return (2 * n + m) * gamma_n(m, n) + def false_fun(m_n): m, n = m_n return (2 * n + m) * gamma_n(m, n) @@ -1040,21 +1023,7 @@ def false_fun(m_n): def CD_m_k(R, m, k): """Eq 31 of Dommaschk paper.""" - # sum1 = 0 - # for j in range(k + 1): - # sum1 += ( - # -( - # alpha(m, j) - # * ( - # alphastar(m, k - m - j) * jnp.log(R) - # + gamma_nstar(m, k - m - j) - # - alpha(m, k - m - j) - # ) - # - gamma_n(m, j) * alphastar(m, k - m - j) - # + alpha(m, j) * betastar(m, k - j) - # ) - # * R ** (2 * j + m) - # ) + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) + def body_fun(j, val): result = ( val @@ -1080,16 +1049,6 @@ def body_fun(j, val): def CN_m_k(R, m, k): """Eq 32 of Dommaschk paper.""" - # sum1 = 0 - # for j in range(k + 1): - # sum1 += ( - # ( - # alpha(m, j) * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) - # - gamma_n(m, j) * alpha(m, k - m - j) - # + alpha(m, j) * beta(m, k - j) - # ) - # * R ** (2 * j + m) - # ) - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) def body_fun(j, val): result = ( @@ -1114,18 +1073,9 @@ def D_m_n(R, Z, m, n): """D_m_n term in eqn 8 of Dommaschk paper.""" # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper - max_ind = ( - n // 2 - ) # find the top of the summation ind, the range should be up to and including this - # i.e. this is the index of k at - # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum - result = 0 - def body_fun(k, val): return val + Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) - # for k in range(n // 2 + 1): - # result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) @@ -1133,14 +1083,6 @@ def N_m_n(R, Z, m, n): """N_m_n term in eqn 9 of Dommaschk paper.""" # the sum comes from fact that N_mn = I_mn and the def of I_mn in eq 2 of the paper - max_ind = ( - n // 2 - ) # find the top of the summation ind, the range should be up to and including this - # i.e. this is the index of k at - # which 2k<=n, and 2(max_ind+1) would be > n and so not included in the sum - result = 0 - # for k in range(n // 2 + 1): - # result += Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CN_m_k(R, m, k) def body_fun(k, val): return val + Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CN_m_k(R, m, k) From 5c0e53e6166b88857913f55e56136562b6531a00 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 28 Jun 2023 02:09:53 -0400 Subject: [PATCH 23/46] add sym option to fitting --- desc/magnetic_fields.py | 77 +++++++++++++++++++++++++++++------ tests/test_magnetic_fields.py | 4 +- 2 files changed, 67 insertions(+), 14 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index ed9dc2dec3..1a807faacf 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -772,7 +772,7 @@ def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1.0): super().__init__(dommaschk_potential, params) @classmethod - def fit_magnetic_field(cls, field, coords, max_m, max_l): + def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): """Fit a vacuum magnetic field with a Dommaschk Potential field. Parameters @@ -783,6 +783,9 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): 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_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 """ # We seek c in Ac = b # A will be the BR, Bphi and BZ from each individual @@ -806,10 +809,17 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") + # TODO: note on how stell sym affects, i.e. + # if stell sym then a=d=0 for even l and b=c=0 for odd l + ##################### # b is made, now do A ##################### - num_modes = 1 + (max_l + 1) * (max_m + 1) * 4 + num_modes = ( + 1 + (max_l + 1) * (max_m + 1) * 4 + if not sym + else 1 + (max_l + 1) * (max_m + 1) * 2 + ) # TODO: technically we can drop some modes # since if max_l=0, there are only ever nonzero terms # for a and b @@ -821,7 +831,7 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): A = jnp.zeros((3 * num_nodes, num_modes)) - # B0 first, the scale magnitude of the 1/R field + # B0 is first coeff, the scale magnitude of the 1/R field domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) B_temp = domm_field.compute_magnetic_field(coords) for i in range(B_temp.shape[1]): @@ -833,13 +843,20 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): ms = [] ls = [] - # order of coeffs are a_ml and + # order of coeffs are a_ml, b_ml, c_ml, d_ml coef_ind = 1 + abcd_inds = [[], [], [], []] for l in range(max_l + 1): for m in range(max_m + 1): ms.append(m) ls.append(l) - for which_coef in range(4): + if not sym: + which_coefs = range(4) # no sym, use all coefs + elif l // 2 == 0: + which_coefs = [1, 2] # a=d=0 for even l with sym + elif l // 2 == 1: + which_coefs = [0, 3] # b=c=0 for odd l with sym + for which_coef in which_coefs: a = 1 if which_coef == 0 else 0 b = 1 if which_coef == 1 else 0 c = 1 if which_coef == 2 else 0 @@ -855,10 +872,12 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): "B0": 0, } B_temp = domm_field.compute_magnetic_field(coords, params=params) + for i in range(B_temp.shape[1]): A = A.at[ i * num_nodes : (i + 1) * num_nodes, coef_ind : coef_ind + 1 ].add(B_temp[:, i].reshape(num_nodes, 1)) + abcd_inds[which_coef].append(coef_ind) coef_ind += 1 # now solve Ac=b for the coefficients c @@ -873,10 +892,46 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l): # recover the params from the c coefficient vector B0 = c[0] - a_arr = c[1::4] - b_arr = c[2::4] - c_arr = c[3::4] - d_arr = c[4::4] + + if sym: + a_arr = ( + c[jnp.asarray(abcd_inds[0]).astype(int)] + if not sym + else jnp.append( + c[jnp.asarray(abcd_inds[0]).astype(int)], + jnp.zeros((len(abcd_inds[1]),)), + ) + ) + d_arr = ( + c[jnp.asarray(abcd_inds[3]).astype(int)] + if not sym + else jnp.append( + c[jnp.asarray(abcd_inds[3]).astype(int)], + jnp.zeros((len(abcd_inds[1]),)), + ) + ) + b_arr = ( + c[jnp.asarray(abcd_inds[1]).astype(int)] + if not sym + else jnp.append( + jnp.zeros((len(abcd_inds[0]),)), + c[jnp.asarray(abcd_inds[1]).astype(int)], + ) + ) + c_arr = ( + c[jnp.asarray(abcd_inds[2]).astype(int)] + if not sym + else jnp.append( + jnp.zeros((len(abcd_inds[0]),)), + c[jnp.asarray(abcd_inds[2]).astype(int)], + ) + ) + else: + a_arr = c[1::4] + b_arr = c[2::4] + c_arr = c[3::4] + d_arr = c[4::4] + return cls(ms, ls, a_arr, b_arr, c_arr, d_arr, B0) @@ -1089,10 +1144,6 @@ def body_fun(k, val): return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) -# TODO: note on how stell sym affects, i.e. -# if stell sym then a=d=0 for even l and b=c=0 for odd l - - def V_m_l(R, phi, Z, m, l, a, b, c, d): """Eq 12 of Dommaschk paper. diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index f4e4df2c57..f7dafc396c 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -234,7 +234,9 @@ def B0_over_R(coord): 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) + 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=4e-15) From 75b7cdf8b71681b012dd2e45918042b746cc0630 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 28 Jun 2023 14:11:07 -0400 Subject: [PATCH 24/46] jit compute inside of fit --- desc/magnetic_fields.py | 17 +++++++++++++++-- tests/test_magnetic_fields.py | 6 +++--- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 1a807faacf..c793b1a882 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1,6 +1,7 @@ """Classes for magnetic fields.""" from abc import ABC, abstractmethod +from functools import partial import numpy as np import scipy.linalg @@ -833,7 +834,18 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): # B0 is first coeff, the scale magnitude of the 1/R field domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) - B_temp = domm_field.compute_magnetic_field(coords) + + params = { + "ms": 0, + "ls": 0, + "a_arr": 0, + "b_arr": 0, + "c_arr": 0, + "d_arr": 0, + "B0": 1, + } + compute_B_fun = jit(domm_field.compute_magnetic_field) + B_temp = compute_B_fun(coords, params=params) for i in range(B_temp.shape[1]): A = A.at[i * num_nodes : (i + 1) * num_nodes, 0:1].add( B_temp[:, i].reshape(num_nodes, 1) @@ -871,7 +883,7 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): "d_arr": d, "B0": 0, } - B_temp = domm_field.compute_magnetic_field(coords, params=params) + B_temp = compute_B_fun(coords, params=params) for i in range(B_temp.shape[1]): A = A.at[ @@ -1144,6 +1156,7 @@ def body_fun(k, val): return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) +@partial(jnp.vectorize, signature="(k),(k),(k),(),(),(),(),(),()->(k)") def V_m_l(R, phi, Z, m, l, a, b, c, d): """Eq 12 of Dommaschk paper. diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index f7dafc396c..688c2b040a 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -210,9 +210,9 @@ def test_dommaschk_vertical_field(): B = DommaschkPotentialField(ms, ls, a_arr, b_arr, c_arr, d_arr) B_dom = B.compute_magnetic_field(coords) ones = jnp.ones_like(B_dom[:, 0]) - jnp.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) - jnp.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten(), atol=1e-14) - jnp.testing.assert_allclose(B_dom[:, 2], ones, atol=5e-15) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) + np.testing.assert_allclose(B_dom[:, 1], 1 / R.flatten(), atol=1e-14) + np.testing.assert_allclose(B_dom[:, 2], ones, atol=5e-15) @pytest.mark.unit From c022586e4fe7a73d416f883f63997d5a931ebfbc Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 28 Jun 2023 18:09:20 -0400 Subject: [PATCH 25/46] change fitting to use jacfwd instead of for loops to find A,reduce test tol --- desc/magnetic_fields.py | 139 ++++++++++++++-------------------- tests/test_magnetic_fields.py | 4 +- 2 files changed, 60 insertions(+), 83 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index c793b1a882..ddb58f0ba8 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1,10 +1,10 @@ """Classes for magnetic fields.""" from abc import ABC, abstractmethod -from functools import partial import numpy as np import scipy.linalg +from jax import jacfwd from netCDF4 import Dataset from desc.backend import cond, fori_loop, gammaln, jit, jnp, odeint @@ -810,9 +810,6 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") - # TODO: note on how stell sym affects, i.e. - # if stell sym then a=d=0 for even l and b=c=0 for odd l - ##################### # b is made, now do A ##################### @@ -832,36 +829,19 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): A = jnp.zeros((3 * num_nodes, num_modes)) - # B0 is first coeff, the scale magnitude of the 1/R field - domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) - - params = { - "ms": 0, - "ls": 0, - "a_arr": 0, - "b_arr": 0, - "c_arr": 0, - "d_arr": 0, - "B0": 1, - } - compute_B_fun = jit(domm_field.compute_magnetic_field) - B_temp = compute_B_fun(coords, params=params) - for i in range(B_temp.shape[1]): - A = A.at[i * num_nodes : (i + 1) * num_nodes, 0:1].add( - B_temp[:, i].reshape(num_nodes, 1) - ) - # mode numbers ms = [] ls = [] - # order of coeffs are a_ml, b_ml, c_ml, d_ml + # order of coeffs are B0, a_ml, b_ml, c_ml, d_ml coef_ind = 1 abcd_inds = [[], [], [], []] + a_s = [] + b_s = [] + c_s = [] + d_s = [] for l in range(max_l + 1): for m in range(max_m + 1): - ms.append(m) - ls.append(l) if not sym: which_coefs = range(4) # no sym, use all coefs elif l // 2 == 0: @@ -874,24 +854,56 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): c = 1 if which_coef == 2 else 0 d = 1 if which_coef == 3 else 0 - params = { - "ms": m, - "ls": l, - "a_arr": a, - "b_arr": b, - "c_arr": c, - "d_arr": d, - "B0": 0, - } - B_temp = compute_B_fun(coords, params=params) - - for i in range(B_temp.shape[1]): - A = A.at[ - i * num_nodes : (i + 1) * num_nodes, coef_ind : coef_ind + 1 - ].add(B_temp[:, i].reshape(num_nodes, 1)) + a_s.append(a) + b_s.append(b) + c_s.append(c) + d_s.append(d) + ms.append(m) + ls.append(l) abcd_inds[which_coef].append(coef_ind) coef_ind += 1 + params = { + "ms": ms, + "ls": ls, + "a_arr": a_s, + "b_arr": b_s, + "c_arr": c_s, + "d_arr": d_s, + "B0": 0.0, + } + n = len(ms) # how many l-m mode pairs there are, also is len(a_s) + + domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) + + def get_B_dom(coords, X, ms, ls): + """Fxn wrapper to find jacobian of dommaschk B wrt coefs a,b,c,d.""" + return domm_field.compute_magnetic_field( + coords, + params={ + "ms": jnp.asarray(ms), + "ls": jnp.asarray(ls), + "a_arr": jnp.asarray(X[1 : n + 1]), + "b_arr": jnp.asarray(X[n + 1 : 2 * n + 1]), + "c_arr": jnp.asarray(X[2 * n + 1 : 3 * n + 1]), + "d_arr": jnp.asarray(X[3 * n + 1 : 4 * n + 1]), + "B0": X[0], + }, + ) + + X = [] + for key in ["B0", "a_arr", "b_arr", "c_arr", "d_arr"]: + obj = params[key] + if type(obj) == list: + X += obj + else: + X += [obj] + X = jnp.asarray(X) + + jac = jacfwd(get_B_dom, argnums=1)(coords, X, params["ms"], params["ls"]) + + A = jac.reshape((rhs.size, len(X)), order="F") + # now solve Ac=b for the coefficients c Ainv = scipy.linalg.pinv(A, rcond=None) @@ -905,44 +917,10 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): # recover the params from the c coefficient vector B0 = c[0] - if sym: - a_arr = ( - c[jnp.asarray(abcd_inds[0]).astype(int)] - if not sym - else jnp.append( - c[jnp.asarray(abcd_inds[0]).astype(int)], - jnp.zeros((len(abcd_inds[1]),)), - ) - ) - d_arr = ( - c[jnp.asarray(abcd_inds[3]).astype(int)] - if not sym - else jnp.append( - c[jnp.asarray(abcd_inds[3]).astype(int)], - jnp.zeros((len(abcd_inds[1]),)), - ) - ) - b_arr = ( - c[jnp.asarray(abcd_inds[1]).astype(int)] - if not sym - else jnp.append( - jnp.zeros((len(abcd_inds[0]),)), - c[jnp.asarray(abcd_inds[1]).astype(int)], - ) - ) - c_arr = ( - c[jnp.asarray(abcd_inds[2]).astype(int)] - if not sym - else jnp.append( - jnp.zeros((len(abcd_inds[0]),)), - c[jnp.asarray(abcd_inds[2]).astype(int)], - ) - ) - else: - a_arr = c[1::4] - b_arr = c[2::4] - c_arr = c[3::4] - d_arr = c[4::4] + a_arr = c[1 : n + 1] + b_arr = c[n + 1 : 2 * n + 1] + c_arr = c[2 * n + 1 : 3 * n + 1] + d_arr = c[3 * n + 1 : 4 * n + 1] return cls(ms, ls, a_arr, b_arr, c_arr, d_arr, B0) @@ -1156,7 +1134,6 @@ def body_fun(k, val): return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) -@partial(jnp.vectorize, signature="(k),(k),(k),(),(),(),(),(),()->(k)") def V_m_l(R, phi, Z, m, l, a, b, c, d): """Eq 12 of Dommaschk paper. @@ -1233,9 +1210,9 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): b_arr = jnp.atleast_1d(b_arr) c_arr = jnp.atleast_1d(c_arr) d_arr = jnp.atleast_1d(d_arr) - for m, l, a, b, c, d in zip(ms, ls, a_arr, b_arr, c_arr, d_arr): value += V_m_l(R, phi, Z, m, l, a, b, c, d) + return value diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 688c2b040a..934e53bb69 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -246,7 +246,7 @@ def B0_over_R(coord): # only nonzero coefficient of the field should be the B0 np.testing.assert_allclose(B._params["B0"], B0, atol=1e-15) for coef in ["a_arr", "b_arr", "c_arr", "d_arr"]: - np.testing.assert_allclose(B._params[coef], 0, atol=1e-15) + np.testing.assert_allclose(B._params[coef], 0, atol=2e-15) @pytest.mark.unit @@ -281,4 +281,4 @@ def test_dommaschk_fit_vertical_and_toroidal_field(): else: np.testing.assert_allclose(coef, 0, atol=1e-15) for name in ["b_arr", "c_arr", "d_arr"]: - np.testing.assert_allclose(B._params[name], 0, atol=1e-15) + np.testing.assert_allclose(B._params[name], 0, atol=2e-15) From a5abda0803b8a22fec084726d36aa245620817d2 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 28 Jun 2023 23:41:40 -0400 Subject: [PATCH 26/46] add jit to fit jacfwd call --- desc/magnetic_fields.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index ddb58f0ba8..528f64a548 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -900,7 +900,7 @@ def get_B_dom(coords, X, ms, ls): X += [obj] X = jnp.asarray(X) - jac = jacfwd(get_B_dom, argnums=1)(coords, X, params["ms"], params["ls"]) + jac = jit(jacfwd(get_B_dom, argnums=1))(coords, X, params["ms"], params["ls"]) A = jac.reshape((rhs.size, len(X)), order="F") From eb29f6e8d1db16e0edebe3f0f206582b749561d7 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 8 Nov 2023 23:26:25 -0500 Subject: [PATCH 27/46] Reorganize tests to undo merge issues --- tests/test_magnetic_fields.py | 142 +++++++++++++++++----------------- 1 file changed, 71 insertions(+), 71 deletions(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 2fd1da2017..439156f4e0 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -388,6 +388,76 @@ def test_field_line_integrate(self): np.testing.assert_allclose(r[-1], 10, rtol=1e-6, atol=1e-6) np.testing.assert_allclose(z[-1], 0.001, rtol=1e-6, atol=1e-6) + @pytest.mark.unit + def test_Bnormal_calculation(self): + """Tests Bnormal calculation for simple toroidal field.""" + tfield = ToroidalMagneticField(2, 1) + surface = get("DSHAPE").surface + Bnorm, _ = tfield.compute_Bnormal(surface) + # should have 0 Bnormal because surface is axisymmetric + np.testing.assert_allclose(Bnorm, 0, atol=1e-14) + + @pytest.mark.unit + def test_Bnormal_save_and_load_DSHAPE(self, tmpdir_factory): + """Tests Bnormal save/load for simple toroidal field with DSHAPE.""" + ### test on simple field first with tokamak + tmpdir = tmpdir_factory.mktemp("BNORM_files") + path = tmpdir.join("BNORM_desc.txt") + tfield = ToroidalMagneticField(2, 1) + eq = get("DSHAPE") + grid = LinearGrid(rho=np.array(1.0), M=20, N=20, NFP=eq.NFP) + x = eq.surface.compute("x", grid=grid, basis="rpz")["x"] + Bnorm, x_from_Bnorm = tfield.compute_Bnormal( + eq.surface, eval_grid=grid, source_grid=grid, basis="rpz" + ) + + # make sure x calculation is the same + np.testing.assert_allclose(x[:, 0], x_from_Bnorm[:, 0], atol=1e-16) + np.testing.assert_allclose(x[:, 2], x_from_Bnorm[:, 2], atol=1e-16) + + np.testing.assert_allclose( + x[:, 1] % (2 * np.pi), x_from_Bnorm[:, 1] % (2 * np.pi), atol=1e-16 + ) + + # should have 0 Bnormal because surface is axisymmetric + np.testing.assert_allclose(Bnorm, 0, atol=1e-14) + + tfield.save_BNORM_file(eq, path, scale_by_curpol=False) + Bnorm_from_file = read_BNORM_file(path, eq, grid, scale_by_curpol=False) + np.testing.assert_allclose(Bnorm, Bnorm_from_file, atol=1e-14) + + # check that loading/saving with scale_by_curpol true + # but no eq passed raises error + with pytest.raises(RuntimeError): + Bnorm_from_file = read_BNORM_file(path, eq.surface, grid) + with pytest.raises(RuntimeError): + tfield.save_BNORM_file(eq.surface, path) + + @pytest.mark.unit + def test_Bnormal_save_and_load_HELIOTRON(self, tmpdir_factory): + """Tests Bnormal save/load for simple toroidal field with HELIOTRON.""" + ### test on simple field with stellarator + tmpdir = tmpdir_factory.mktemp("BNORM_files") + path = tmpdir.join("BNORM_desc_heliotron.txt") + tfield = ToroidalMagneticField(2, 1) + eq = get("HELIOTRON") + grid = LinearGrid(rho=np.array(1.0), M=20, N=20, NFP=eq.NFP) + x = eq.surface.compute("x", grid=grid, basis="xyz")["x"] + Bnorm, x_from_Bnorm = tfield.compute_Bnormal( + eq.surface, eval_grid=grid, basis="xyz" + ) + + # make sure x calculation is the same + np.testing.assert_allclose(x, x_from_Bnorm, atol=1e-16) + + tfield.save_BNORM_file(eq, path, 40, 40) + Bnorm_from_file = read_BNORM_file(path, eq, grid) + np.testing.assert_allclose(Bnorm, Bnorm_from_file, atol=1e-8) + + asym_surf = FourierRZToroidalSurface(sym=False) + with pytest.raises(AssertionError, match="sym"): + Bnorm_from_file = read_BNORM_file(path, asym_surf, grid) + @pytest.mark.unit def test_dommaschk_CN_CD_m_0(): @@ -535,74 +605,4 @@ def test_dommaschk_fit_vertical_and_toroidal_field(): else: np.testing.assert_allclose(coef, 0, atol=1e-15) for name in ["b_arr", "c_arr", "d_arr"]: - np.testing.assert_allclose(B._params[name], 0, atol=2e-15) - - @pytest.mark.unit - def test_Bnormal_calculation(self): - """Tests Bnormal calculation for simple toroidal field.""" - tfield = ToroidalMagneticField(2, 1) - surface = get("DSHAPE").surface - Bnorm, _ = tfield.compute_Bnormal(surface) - # should have 0 Bnormal because surface is axisymmetric - np.testing.assert_allclose(Bnorm, 0, atol=3e-15) - - @pytest.mark.unit - def test_Bnormal_save_and_load_DSHAPE(self, tmpdir_factory): - """Tests Bnormal save/load for simple toroidal field with DSHAPE.""" - ### test on simple field first with tokamak - tmpdir = tmpdir_factory.mktemp("BNORM_files") - path = tmpdir.join("BNORM_desc.txt") - tfield = ToroidalMagneticField(2, 1) - eq = get("DSHAPE") - grid = LinearGrid(rho=np.array(1.0), M=20, N=20, NFP=eq.NFP) - x = eq.surface.compute("x", grid=grid, basis="rpz")["x"] - Bnorm, x_from_Bnorm = tfield.compute_Bnormal( - eq.surface, eval_grid=grid, source_grid=grid, basis="rpz" - ) - - # make sure x calculation is the same - np.testing.assert_allclose(x[:, 0], x_from_Bnorm[:, 0], atol=1e-16) - np.testing.assert_allclose(x[:, 2], x_from_Bnorm[:, 2], atol=1e-16) - - np.testing.assert_allclose( - x[:, 1] % (2 * np.pi), x_from_Bnorm[:, 1] % (2 * np.pi), atol=1e-16 - ) - - # should have 0 Bnormal because surface is axisymmetric - np.testing.assert_allclose(Bnorm, 0, atol=1e-14) - - tfield.save_BNORM_file(eq, path, scale_by_curpol=False) - Bnorm_from_file = read_BNORM_file(path, eq, grid, scale_by_curpol=False) - np.testing.assert_allclose(Bnorm, Bnorm_from_file, atol=1e-14) - - # check that loading/saving with scale_by_curpol true - # but no eq passed raises error - with pytest.raises(RuntimeError): - Bnorm_from_file = read_BNORM_file(path, eq.surface, grid) - with pytest.raises(RuntimeError): - tfield.save_BNORM_file(eq.surface, path) - - @pytest.mark.unit - def test_Bnormal_save_and_load_HELIOTRON(self, tmpdir_factory): - """Tests Bnormal save/load for simple toroidal field with HELIOTRON.""" - ### test on simple field with stellarator - tmpdir = tmpdir_factory.mktemp("BNORM_files") - path = tmpdir.join("BNORM_desc_heliotron.txt") - tfield = ToroidalMagneticField(2, 1) - eq = get("HELIOTRON") - grid = LinearGrid(rho=np.array(1.0), M=20, N=20, NFP=eq.NFP) - x = eq.surface.compute("x", grid=grid, basis="xyz")["x"] - Bnorm, x_from_Bnorm = tfield.compute_Bnormal( - eq.surface, eval_grid=grid, basis="xyz" - ) - - # make sure x calculation is the same - np.testing.assert_allclose(x, x_from_Bnorm, atol=1e-16) - - tfield.save_BNORM_file(eq, path, 40, 40) - Bnorm_from_file = read_BNORM_file(path, eq, grid) - np.testing.assert_allclose(Bnorm, Bnorm_from_file, atol=1e-8) - - asym_surf = FourierRZToroidalSurface(sym=False) - with pytest.raises(AssertionError, match="sym"): - Bnorm_from_file = read_BNORM_file(path, asym_surf, grid) + np.testing.assert_allclose(B._params[name], 0, atol=1e-14) From 027c40a1727e76ffe7aa8c37d5d8601b0e81fab3 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 8 Nov 2023 23:27:05 -0500 Subject: [PATCH 28/46] Apply jit to dommaschk sub-functions --- desc/magnetic_fields.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 5aa695d5a5..fc48701f48 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1278,11 +1278,13 @@ def odefun(rpz, s): true_fun = lambda m_n: 0.0 # used for returning 0 when conditionals evaluate to True +@jit def gamma(n): """Gamma function, only implemented for integers (equiv to factorial of (n-1)).""" return jnp.exp(gammaln(n)) +@jit def alpha(m, n): """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" # modified for eqns 31 and 32 @@ -1305,6 +1307,7 @@ def bool_fun(n): ) +@jit def alphastar(m, n): """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" # modified for eqns 31 and 32 @@ -1315,6 +1318,7 @@ def false_fun(m_n): return cond(jnp.any(n < 0), true_fun, false_fun, (m, n)) +@jit def beta(m, n): """Beta of eq 28, modified for eqns 31 and 32.""" @@ -1325,6 +1329,7 @@ def false_fun(m_n): return cond(jnp.any(jnp.logical_or(n < 0, n >= m)), true_fun, false_fun, (m, n)) +@jit def betastar(m, n): """Beta* of eq 28, modified for eqns 31 and 32.""" @@ -1335,6 +1340,7 @@ def false_fun(m_n): return cond(jnp.any(jnp.logical_or(n < 0, n >= m)), true_fun, false_fun, (m, n)) +@jit def gamma_n(m, n): """gamma_n of eq 33.""" @@ -1348,6 +1354,7 @@ def false_fun(m_n): return cond(jnp.any(n <= 0), true_fun, false_fun, (m, n)) +@jit def gamma_nstar(m, n): """gamma_n star of eq 33.""" @@ -1358,6 +1365,7 @@ def false_fun(m_n): return cond(jnp.any(n <= 0), true_fun, false_fun, (m, n)) +@jit def CD_m_k(R, m, k): """Eq 31 of Dommaschk paper.""" @@ -1384,6 +1392,7 @@ def body_fun(j, val): return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) +@jit def CN_m_k(R, m, k): """Eq 32 of Dommaschk paper.""" @@ -1406,6 +1415,7 @@ def body_fun(j, val): return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) +@jit def D_m_n(R, Z, m, n): """D_m_n term in eqn 8 of Dommaschk paper.""" # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper @@ -1416,6 +1426,7 @@ def body_fun(k, val): return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) +@jit def N_m_n(R, Z, m, n): """N_m_n term in eqn 9 of Dommaschk paper.""" # the sum comes from fact that N_mn = I_mn and the def of I_mn in eq 2 of the paper @@ -1426,6 +1437,7 @@ def body_fun(k, val): return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) +@jit def V_m_l(R, phi, Z, m, l, a, b, c, d): """Eq 12 of Dommaschk paper. @@ -1459,6 +1471,7 @@ def V_m_l(R, phi, Z, m, l, a, b, c, d): ) * N_m_n(R, Z, m, l - 1) +@jit def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): """Eq 1 of Dommaschk paper. From f92cedbef13f38eb0b6c738b694490aef3c77384 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 8 Nov 2023 23:27:37 -0500 Subject: [PATCH 29/46] Fix edge case of derivative of 0**0 --- desc/magnetic_fields.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index fc48701f48..ed0f3cceff 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1421,7 +1421,11 @@ def D_m_n(R, Z, m, n): # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper def body_fun(k, val): - return val + Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CD_m_k(R, m, k) + 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 return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) @@ -1432,7 +1436,11 @@ def N_m_n(R, Z, m, n): # the sum comes from fact that N_mn = I_mn and the def of I_mn in eq 2 of the paper def body_fun(k, val): - return val + Z ** (n - 2 * k) / gamma(n - 2 * k + 1) * CN_m_k(R, m, k) + 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 return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) From 8c2a284cb1f998e2a54ae218586c8428ffc51c5d Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 8 Nov 2023 23:28:14 -0500 Subject: [PATCH 30/46] Use jax loop for dommaschk evaluation --- desc/magnetic_fields.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index ed0f3cceff..72bddab566 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1514,19 +1514,21 @@ def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): at the given R,phi,Z points (same size as the size of the given R,phi, Z arrays). """ + ms, ls, a_arr, b_arr, c_arr, d_arr = map( + jnp.atleast_1d, (ms, ls, a_arr, b_arr, c_arr, d_arr) + ) + R, phi, Z = map(jnp.atleast_1d, (R, phi, Z)) + R, phi, Z = jnp.broadcast_arrays(R, phi, Z) + ms, ls, a_arr, b_arr, c_arr, d_arr = jnp.broadcast_arrays( + ms, ls, a_arr, b_arr, c_arr, d_arr + ) value = B0 * phi # phi term - # make sure all are 1D arrays - ms = jnp.atleast_1d(ms) - ls = jnp.atleast_1d(ls) - a_arr = jnp.atleast_1d(a_arr) - b_arr = jnp.atleast_1d(b_arr) - c_arr = jnp.atleast_1d(c_arr) - d_arr = jnp.atleast_1d(d_arr) - for m, l, a, b, c, d in zip(ms, ls, a_arr, b_arr, c_arr, d_arr): - value += V_m_l(R, phi, Z, m, l, a, b, c, d) - - return value + def body(i, val): + val += V_m_l(R, phi, Z, ms[i], ls[i], a_arr[i], b_arr[i], c_arr[i], d_arr[i]) + return val + + return fori_loop(0, len(ms), body, value) def read_mgrid( From f9172ef6d9d5ede6d862a98f2aa610edcfb3057d Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 8 Nov 2023 23:29:21 -0500 Subject: [PATCH 31/46] Fix order of return values in read_mgrid --- desc/magnetic_fields.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 72bddab566..823fe75bc8 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1531,10 +1531,7 @@ def body(i, val): return fori_loop(0, len(ms), body, value) -def read_mgrid( - mgrid_file, - extcur=1, -): +def read_mgrid(mgrid_file, extcur=1): """Read an "mgrid" file from MAKEGRID and return the grid and magnetic field. Parameters @@ -1582,7 +1579,7 @@ def read_mgrid( Zgrid = np.linspace(zMin, zMax, jz) pgrid = 2.0 * np.pi / (nfp * kp) * np.arange(kp) - return Rgrid, Zgrid, pgrid, br, bp, bz, nfp + return Rgrid, pgrid, Zgrid, br, bp, bz, nfp class CurrentPotentialField(_MagneticField, FourierRZToroidalSurface): From 849711b9628169aef5bd6ff7b36676fd5a72366c Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 8 Nov 2023 23:29:39 -0500 Subject: [PATCH 32/46] Adjust test tolerances --- tests/test_magnetic_fields.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 439156f4e0..2f64683b90 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -563,14 +563,14 @@ def B0_over_R(coord): ) B_dom = B.compute_magnetic_field(coords) - np.testing.assert_allclose(B_dom[:, 0], 0, atol=4e-15) - np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-15) - np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=1e-15) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) + np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-14) + np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=1e-14) # only nonzero coefficient of the field should be the B0 - np.testing.assert_allclose(B._params["B0"], B0, atol=1e-15) + 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=2e-15) + np.testing.assert_allclose(B._params[coef], 0, atol=1e-14) @pytest.mark.unit @@ -591,18 +591,18 @@ def test_dommaschk_fit_vertical_and_toroidal_field(): B = DommaschkPotentialField.fit_magnetic_field(field, coords, max_m, max_l) B_dom = B.compute_magnetic_field(coords) - np.testing.assert_allclose(B_dom[:, 0], 0, atol=4e-15) - np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-15) - np.testing.assert_allclose(B_dom[:, 2], B0_Z, atol=1e-15) + np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) + np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-14) + np.testing.assert_allclose(B_dom[:, 2], B0_Z, atol=1e-14) np.testing.assert_allclose(B._params["B0"], B0) # only nonzero coefficient of the field should be the B0 and a_ml = a_01 - np.testing.assert_allclose(B._params["B0"], B0, atol=1e-15) + np.testing.assert_allclose(B._params["B0"], B0, atol=1e-14) for coef, m, l in zip(B._params["a_arr"], B._params["ms"], B._params["ls"]): if m == 0 and l == 1: np.testing.assert_allclose(coef, B0_Z) else: - np.testing.assert_allclose(coef, 0, atol=1e-15) + np.testing.assert_allclose(coef, 0, atol=1e-14) for name in ["b_arr", "c_arr", "d_arr"]: np.testing.assert_allclose(B._params[name], 0, atol=1e-14) From 52fb9fbb3a0812e338b2a15bbacb524a6a2243a6 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Thu, 16 Nov 2023 19:18:44 -0500 Subject: [PATCH 33/46] fix test, fix array counting so that the num_modes = sum(len(a_arr)+len(b_arr)...) + 1 as intended, fix sym with this change --- desc/magnetic_fields.py | 130 ++++++++++++++++++++-------------- tests/test_magnetic_fields.py | 10 +-- 2 files changed, 80 insertions(+), 60 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 823fe75bc8..39601fee42 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1033,7 +1033,6 @@ class DommaschkPotentialField(ScalarPotentialField): """ def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1.0): - ms = jnp.atleast_1d(ms) ls = jnp.atleast_1d(ls) a_arr = jnp.atleast_1d(a_arr) @@ -1063,7 +1062,9 @@ def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1.0): super().__init__(dommaschk_potential, params) @classmethod - def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): + def fit_magnetic_field( # noqa: C901 - FIXME - simplify + cls, field, coords, max_m, max_l, sym=False, verbose=1 + ): """Fit a vacuum magnetic field with a Dommaschk Potential field. Parameters @@ -1077,6 +1078,7 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): sym (bool): if field is stellarator symmetric or not. if True, only stellarator-symmetric modes will be included in the fitting + verbose (int): verbosity level of fitting routine, > 0 prints residuals """ # We seek c in Ac = b # A will be the BR, Bphi and BZ from each individual @@ -1090,69 +1092,78 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): else: B = field.compute_magnetic_field(coords) - num_nodes = coords.shape[0] # number of coordinate nodes - - # we will have the rhs be 3*num_nodes in length (bc of vector B) - ######### # make b ######### + # we will have the rhs be 3*num_nodes in length (bc of vector B) rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") ##################### # b is made, now do A ##################### - num_modes = ( - 1 + (max_l + 1) * (max_m + 1) * 4 - if not sym - else 1 + (max_l + 1) * (max_m + 1) * 2 - ) + num_modes = 1 + (max_l + 1) * (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 + # modes that don't fit symmetry, but in future + # should refactor code to have a 3rd index so that + # we have a = V_ml0, b = V_ml1, c = V_ml2, d = V_ml3 + # and the modes array can then be [m,l,x] where x is 0,1,2,3 + # and we dont need to keep track of a,b,c,d separately + # TODO: technically we can drop some modes - # since if max_l=0, there are only ever nonzero terms - # for a and b - # and if max_m=0, there are only ever nonzero terms - # for a and c + # since if max_l=0, there are only ever nonzero terms for a and b + # and if max_m=0, there are only ever nonzero terms for a and c # but since we are only fitting in a least squares sense, # and max_l and max_m should probably be both nonzero anyways, # this is not an issue right now - A = jnp.zeros((3 * num_nodes, num_modes)) - # mode numbers ms = [] ls = [] - # order of coeffs are B0, a_ml, b_ml, c_ml, d_ml - coef_ind = 1 - abcd_inds = [[], [], [], []] + # order of coeffs in the vector c are B0, a_ml, b_ml, c_ml, d_ml a_s = [] b_s = [] c_s = [] d_s = [] + zero_due_to_sym_inds = [] + abcd_zero_due_to_sym_inds = [ + [], + [], + [], + [], + ] # indices that should be 0 due to symmetry for l in range(max_l + 1): for m in range(max_m + 1): if not sym: - which_coefs = range(4) # no sym, use all coefs + pass # no sym, use all coefs elif l // 2 == 0: - which_coefs = [1, 2] # a=d=0 for even l with sym + zero_due_to_sym_inds = [0, 3] # a=d=0 for even l with sym elif l // 2 == 1: - which_coefs = [0, 3] # b=c=0 for odd l with sym - for which_coef in which_coefs: - a = 1 if which_coef == 0 else 0 - b = 1 if which_coef == 1 else 0 - c = 1 if which_coef == 2 else 0 - d = 1 if which_coef == 3 else 0 - - a_s.append(a) - b_s.append(b) - c_s.append(c) - d_s.append(d) - ms.append(m) - ls.append(l) - abcd_inds[which_coef].append(coef_ind) - coef_ind += 1 - + zero_due_to_sym_inds = [1, 2] # b=c=0 for odd l with sym + for which_coef in range(4): + if which_coef == 0: + a_s.append(1) + elif which_coef == 1: + b_s.append(1) + elif which_coef == 2: + c_s.append(1) + elif which_coef == 3: + d_s.append(1) + if which_coef in zero_due_to_sym_inds: + 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): + abcd_zero_due_to_sym_inds[i] = jnp.asarray(abcd_zero_due_to_sym_inds[i]) + assert ( + round(jnp.sum(len(a_s) + len(b_s) + len(c_s) + len(d_s))) == num_modes - 1 + ) params = { "ms": ms, "ls": ls, @@ -1162,21 +1173,29 @@ def fit_magnetic_field(cls, field, coords, max_m, max_l, sym=False): "d_arr": d_s, "B0": 0.0, } - n = len(ms) # how many l-m mode pairs there are, also is len(a_s) - + n = ( + round(num_modes - 1) / 4 + ) # how many l-m mode pairs there are, also is len(a_s) + n = int(n) domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) def get_B_dom(coords, X, ms, ls): """Fxn wrapper to find jacobian of dommaschk B wrt coefs a,b,c,d.""" + # zero out any terms that should be zero due to symmetry, which + # we cataloged earlier for each a_arr,b_arr,c_arr,d_arr + # that way the resulting modes after pinv don't contain them either return domm_field.compute_magnetic_field( coords, params={ "ms": jnp.asarray(ms), "ls": jnp.asarray(ls), - "a_arr": jnp.asarray(X[1 : n + 1]), - "b_arr": jnp.asarray(X[n + 1 : 2 * n + 1]), - "c_arr": jnp.asarray(X[2 * n + 1 : 3 * n + 1]), - "d_arr": jnp.asarray(X[3 * n + 1 : 4 * n + 1]), + "a_arr": jnp.asarray(X[1 : n + 1]) * abcd_zero_due_to_sym_inds[0], + "b_arr": jnp.asarray(X[n + 1 : 2 * n + 1]) + * abcd_zero_due_to_sym_inds[1], + "c_arr": jnp.asarray(X[2 * n + 1 : 3 * n + 1]) + * abcd_zero_due_to_sym_inds[2], + "d_arr": jnp.asarray(X[3 * n + 1 : 4 * n + 1]) + * abcd_zero_due_to_sym_inds[3], "B0": X[0], }, ) @@ -1196,21 +1215,23 @@ def get_B_dom(coords, X, ms, ls): # now solve Ac=b for the coefficients c - Ainv = scipy.linalg.pinv(A, rcond=None) + Ainv = scipy.linalg.pinv(A) c = jnp.matmul(Ainv, rhs) res = jnp.matmul(A, c) - rhs - print(f"Mean Residual of fit: {jnp.mean(jnp.abs(res)):1.4e} T") - print(f"Max Residual of fit: {jnp.max(jnp.abs(res)):1.4e} T") - print(f"Min Residual of fit: {jnp.min(jnp.abs(res)):1.4e} T") + if verbose > 1: + print(f"Mean Residual of fit: {jnp.mean(jnp.abs(res)):1.4e} T") + print(f"Max Residual of fit: {jnp.max(jnp.abs(res)):1.4e} T") + print(f"Min Residual of fit: {jnp.min(jnp.abs(res)):1.4e} T") # recover the params from the c coefficient vector B0 = c[0] - a_arr = c[1 : n + 1] - b_arr = c[n + 1 : 2 * n + 1] - c_arr = c[2 * n + 1 : 3 * n + 1] - d_arr = c[3 * n + 1 : 4 * n + 1] + # we zero out the terms that should be zero due to symmetry here + 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] + d_arr = c[3 * n + 1 : 4 * n + 1] * abcd_zero_due_to_sym_inds[3] return cls(ms, ls, a_arr, b_arr, c_arr, d_arr, B0) @@ -1310,8 +1331,8 @@ def bool_fun(n): @jit def alphastar(m, n): """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" - # modified for eqns 31 and 32 - def false_fun(m_n): + + def false_fun(m_n): # modified for eqns 31 and 32 m, n = m_n return (2 * n + m) * alpha(m, n) @@ -1558,7 +1579,6 @@ def read_mgrid(mgrid_file, extcur=1): bz = np.zeros([kp, jz, ir]) extcur = np.broadcast_to(extcur, nextcur) for i in range(nextcur): - # apply scaling by currents given in VMEC input file scale = extcur[i] diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 2f64683b90..d66d1d4628 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -548,8 +548,8 @@ def test_dommaschk_fit_toroidal_field(): R, phi, Z = np.meshgrid(R, phi, Z) coords = np.vstack((R.flatten(), phi.flatten(), Z.flatten())).T - max_l = 1 - max_m = 1 + max_l = 2 + max_m = 2 B0 = 2 # scale strength for to 1/R field to fit def B0_over_R(coord): @@ -563,9 +563,9 @@ def B0_over_R(coord): ) B_dom = B.compute_magnetic_field(coords) - np.testing.assert_allclose(B_dom[:, 0], 0, atol=1e-14) - np.testing.assert_allclose(B_dom[:, 1], B0 / R.flatten(), atol=1e-14) - np.testing.assert_allclose(B_dom[:, 2], jnp.zeros_like(R.flatten()), atol=1e-14) + 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) From df3740bd436178c5330d0546b85075c214d41425 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 22 Nov 2023 23:20:21 -0500 Subject: [PATCH 34/46] address some comments --- desc/magnetic_fields.py | 38 +++++++++++++++++++++++------------ tests/test_magnetic_fields.py | 15 ++++++++++++++ 2 files changed, 40 insertions(+), 13 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 39601fee42..7567215659 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -984,8 +984,7 @@ def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): if hasattr(coords, "nodes"): coords = coords.nodes coords = jnp.atleast_2d(coords) - if coords.dtype == int: - coords = coords.astype(float) + coords = coords.astype(float) # ensure coords are float if basis == "xyz": coords = xyz2rpz(coords) @@ -1032,7 +1031,16 @@ class DommaschkPotentialField(ScalarPotentialField): """ - def __init__(self, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1.0): + def __init__( + self, + ms=jnp.array([0]), + ls=jnp.array([0]), + a_arr=jnp.array([0.0]), + b_arr=jnp.array([0.0]), + c_arr=jnp.array([0.0]), + d_arr=jnp.array([0.0]), + B0=1.0, + ): ms = jnp.atleast_1d(ms) ls = jnp.atleast_1d(ls) a_arr = jnp.atleast_1d(a_arr) @@ -1069,9 +1077,11 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify Parameters ---------- - field (MagneticField or callable): magnetic field to fit + field (MagneticField or callable or ndarray): magnetic field to fit if callable, must accept (num_nodes,3) ndarray as argument - and output (num_nodes,3) as the B field in cylindrical rpz basis. + and output (num_nodes,3) as the B field in cylindrical rpz basis. + 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_l (int): maximum l to use for Dommaschk Potentials @@ -1086,11 +1096,15 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify # c is the dommaschk potential coefficients # c will be [B0, a_00, a_10, a_01, a_11... etc] # b is the magnetic field at each node which we are fitting - - if not isinstance(field, _MagneticField): - B = field(coords) - else: + if isinstance(field, _MagneticField): B = field.compute_magnetic_field(coords) + elif callable(field): + B = field(coords) + else: # it must be the field evaluated at the passed-in coords + B = field + # TODO: add basis argument for if passed-in field or callable + # evaluates rpz or xyz basis magnetic field vector, + # and what basis coords is ######### # make b @@ -1161,9 +1175,7 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify ls.append(l) for i in range(4): abcd_zero_due_to_sym_inds[i] = jnp.asarray(abcd_zero_due_to_sym_inds[i]) - assert ( - round(jnp.sum(len(a_s) + len(b_s) + len(c_s) + len(d_s))) == num_modes - 1 - ) + assert (len(a_s) + len(b_s) + len(c_s) + len(d_s)) == num_modes - 1 params = { "ms": ms, "ls": ls, @@ -1219,7 +1231,7 @@ def get_B_dom(coords, X, ms, ls): c = jnp.matmul(Ainv, rhs) res = jnp.matmul(A, c) - rhs - if verbose > 1: + if verbose > 0: print(f"Mean Residual of fit: {jnp.mean(jnp.abs(res)):1.4e} T") print(f"Max Residual of fit: {jnp.max(jnp.abs(res)):1.4e} T") print(f"Min Residual of fit: {jnp.min(jnp.abs(res)):1.4e} T") diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index d66d1d4628..5330eb988e 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -572,6 +572,21 @@ def B0_over_R(coord): 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.compute_magnetic_field(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(): From c27ac9248d45c608bb91082b18d82a7ae4ef85da Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 22 Nov 2023 23:25:20 -0500 Subject: [PATCH 35/46] change to use lstsq from jnp.linalg for fit: --- desc/magnetic_fields.py | 6 ++---- tests/test_magnetic_fields.py | 2 +- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 7567215659..efa58e8ba2 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -3,7 +3,6 @@ from abc import ABC, abstractmethod import numpy as np -import scipy.linalg from jax import jacfwd from netCDF4 import Dataset @@ -1227,10 +1226,9 @@ def get_B_dom(coords, X, ms, ls): # now solve Ac=b for the coefficients c - Ainv = scipy.linalg.pinv(A) - c = jnp.matmul(Ainv, rhs) + # TODO: use min singular value to give sense of cond number? + c, res, _, _ = jnp.linalg.lstsq(A, rhs) - res = jnp.matmul(A, c) - rhs if verbose > 0: print(f"Mean Residual of fit: {jnp.mean(jnp.abs(res)):1.4e} T") print(f"Max Residual of fit: {jnp.max(jnp.abs(res)):1.4e} T") diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 5330eb988e..791b45ce19 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -574,7 +574,7 @@ def B0_over_R(coord): # test fit from values B = DommaschkPotentialField.fit_magnetic_field( - B0_over_R.compute_magnetic_field(coords), coords, max_m, max_l, sym=True + B0_over_R(coords), coords, max_m, max_l, sym=True ) B_dom = B.compute_magnetic_field(coords) From 2382f16c30af3969f835bc5b10bca6b71a713f4b Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 24 Nov 2023 03:14:39 -0500 Subject: [PATCH 36/46] correct use of residual from lstsq --- desc/magnetic_fields.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index efa58e8ba2..cb07f0d4fb 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1230,9 +1230,7 @@ def get_B_dom(coords, X, ms, ls): c, res, _, _ = jnp.linalg.lstsq(A, rhs) if verbose > 0: - print(f"Mean Residual of fit: {jnp.mean(jnp.abs(res)):1.4e} T") - print(f"Max Residual of fit: {jnp.max(jnp.abs(res)):1.4e} T") - print(f"Min Residual of fit: {jnp.min(jnp.abs(res)):1.4e} T") + print(f"Sum of Squares Residual of fit: {res:1.4e} T") # recover the params from the c coefficient vector B0 = c[0] From 3f6fd034c29da488ce8c113ed2c7ce8492c04ea2 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Fri, 24 Nov 2023 03:16:47 -0500 Subject: [PATCH 37/46] correct use of residual from lstsq, again --- desc/magnetic_fields.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index cb07f0d4fb..9d001c549a 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1230,7 +1230,8 @@ def get_B_dom(coords, X, ms, ls): c, res, _, _ = jnp.linalg.lstsq(A, rhs) if verbose > 0: - print(f"Sum of Squares Residual of fit: {res:1.4e} T") + # res is a list of len(1) so index into it + print(f"Sum of Squares Residual of fit: {res[0]:1.4e} T") # recover the params from the c coefficient vector B0 = c[0] From 4eb1c8f24f07445c81b92f4d1b30ad019693d7a6 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sat, 20 Jan 2024 14:12:48 -0500 Subject: [PATCH 38/46] remove unneeded jnp.any statements --- desc/magnetic_fields.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 54987a6022..b5654db24e 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1521,7 +1521,7 @@ def false_fun(m_n): return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) def bool_fun(n): - return jnp.any(n < 0) + return n < 0 return cond( bool_fun(n), @@ -1542,7 +1542,7 @@ def false_fun(m_n): # modified for eqns 31 and 32 m, n = m_n return (2 * n + m) * alpha(m, n) - return cond(jnp.any(n < 0), true_fun, false_fun, (m, n)) + return cond(n < 0, true_fun, false_fun, (m, n)) @jit @@ -1553,7 +1553,7 @@ def false_fun(m_n): m, n = m_n return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) - return cond(jnp.any(jnp.logical_or(n < 0, n >= m)), true_fun, false_fun, (m, n)) + return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) @jit @@ -1564,7 +1564,7 @@ def false_fun(m_n): m, n = m_n return (2 * n - m) * beta(m, n) - return cond(jnp.any(jnp.logical_or(n < 0, n >= m)), true_fun, false_fun, (m, n)) + return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) @jit @@ -1578,7 +1578,7 @@ def false_fun(m_n): m, n = m_n return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) - return cond(jnp.any(n <= 0), true_fun, false_fun, (m, n)) + return cond(n <= 0, true_fun, false_fun, (m, n)) @jit @@ -1589,7 +1589,7 @@ def false_fun(m_n): m, n = m_n return (2 * n + m) * gamma_n(m, n) - return cond(jnp.any(n <= 0), true_fun, false_fun, (m, n)) + return cond(n <= 0, true_fun, false_fun, (m, n)) @jit From ae23a871cf5f68cffa77f1e2f83d8c0c5b3bcd80 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Sun, 4 Feb 2024 23:35:24 -0500 Subject: [PATCH 39/46] update docstring --- desc/magnetic_fields.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 942da6182a..d6fff4751b 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1312,7 +1312,7 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify Parameters ---------- field (MagneticField or callable or ndarray): magnetic field to fit - if callable, must accept (num_nodes,3) ndarray as argument + if callable, must accept (num_nodes,3) array of rpz coords as argument and output (num_nodes,3) as the B field in cylindrical rpz basis. 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) From e5b2de480f6e7bcc2a648060a4a9fe79bbba7b33 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Fri, 1 Mar 2024 13:59:28 -0500 Subject: [PATCH 40/46] Move magnetic fields to their own module --- desc/magnetic_fields/__init__.py | 16 + .../_core.py} | 1178 +---------------- desc/magnetic_fields/_current_potential.py | 678 ++++++++++ desc/magnetic_fields/_dommaschk.py | 512 +++++++ setup.cfg | 7 +- 5 files changed, 1210 insertions(+), 1181 deletions(-) create mode 100644 desc/magnetic_fields/__init__.py rename desc/{magnetic_fields.py => magnetic_fields/_core.py} (54%) create mode 100644 desc/magnetic_fields/_current_potential.py create mode 100644 desc/magnetic_fields/_dommaschk.py diff --git a/desc/magnetic_fields/__init__.py b/desc/magnetic_fields/__init__.py new file mode 100644 index 0000000000..722f64adda --- /dev/null +++ b/desc/magnetic_fields/__init__.py @@ -0,0 +1,16 @@ +"""Classes for Magnetic Fields.""" + +from ._core import ( + PoloidalMagneticField, + ScalarPotentialField, + ScaledMagneticField, + SplineMagneticField, + SumMagneticField, + ToroidalMagneticField, + VerticalMagneticField, + _MagneticField, + field_line_integrate, + read_BNORM_file, +) +from ._current_potential import CurrentPotentialField, FourierCurrentPotentialField +from ._dommaschk import DommaschkPotentialField, dommaschk_potential diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields/_core.py similarity index 54% rename from desc/magnetic_fields.py rename to desc/magnetic_fields/_core.py index 4f088dadd8..6343976982 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields/_core.py @@ -5,20 +5,18 @@ import numpy as np from interpax import approx_df, interp2d, interp3d -from jax import jacfwd from netCDF4 import Dataset, chartostring, stringtochar -from desc.backend import cond, fori_loop, gammaln, jit, jnp, odeint +from desc.backend import fori_loop, jit, jnp, odeint from desc.basis import DoubleFourierSeries -from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.compute import rpz2xyz_vec, xyz2rpz from desc.derivatives import Derivative from desc.equilibrium import EquilibriaFamily, Equilibrium -from desc.geometry import FourierRZToroidalSurface from desc.grid import LinearGrid from desc.io import IOAble from desc.optimizable import Optimizable, OptimizableCollection, optimizable_parameter from desc.transform import Transform -from desc.utils import copy_coeffs, errorif, flatten_list, setdefault, warnif +from desc.utils import copy_coeffs, flatten_list, setdefault from desc.vmec_utils import ptolemy_identity_fwd, ptolemy_identity_rev @@ -1371,249 +1369,6 @@ def compute_magnetic_field( return B -class DommaschkPotentialField(ScalarPotentialField): - """Magnetic field due to a Dommaschk scalar magnetic potential in rpz coordinates. - - From Dommaschk 1986 paper https://doi.org/10.1016/0010-4655(86)90109-8 - - this is the field due to the dommaschk potential (eq. 1) for - a given set of m,l indices and their corresponding - coefficients a_ml, b_ml, c_ml d_ml. - - Parameters - ---------- - ms : 1D array-like of int - first indices of V_m_l terms (eq. 12 of reference) - ls : 1D array-like of int - second indices of V_m_l terms (eq. 12 of reference) - 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 - b_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*D_m_l terms - c_arr : 1D array-like of float - c_m_l coefficients of V_m_l terms, which multiply the cos(m*phi)*N_m_l-1 term - d_arr : 1D array-like of float - d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms - B0: float - scale strength of the magnetic field's 1/R portion - - """ - - def __init__( - self, - ms=jnp.array([0]), - ls=jnp.array([0]), - a_arr=jnp.array([0.0]), - b_arr=jnp.array([0.0]), - c_arr=jnp.array([0.0]), - d_arr=jnp.array([0.0]), - B0=1.0, - ): - ms = jnp.atleast_1d(ms) - ls = jnp.atleast_1d(ls) - a_arr = jnp.atleast_1d(a_arr) - b_arr = jnp.atleast_1d(b_arr) - c_arr = jnp.atleast_1d(c_arr) - d_arr = jnp.atleast_1d(d_arr) - - assert ( - ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size - ), "Passed in arrays must all be of the same size!" - assert not jnp.any( - jnp.logical_or(ms < 0, ls < 0) - ), "m and l mode numbers must be >= 0!" - assert ( - jnp.isscalar(B0) or jnp.atleast_1d(B0).size == 1 - ), "B0 should be a scalar value!" - - params = {} - params["ms"] = ms - params["ls"] = ls - params["a_arr"] = a_arr - params["b_arr"] = b_arr - params["c_arr"] = c_arr - params["d_arr"] = d_arr - params["B0"] = B0 - - super().__init__(dommaschk_potential, params) - - @classmethod - def fit_magnetic_field( # noqa: C901 - FIXME - simplify - cls, field, coords, max_m, max_l, sym=False, verbose=1 - ): - """Fit a vacuum magnetic field with a Dommaschk Potential field. - - Parameters - ---------- - field (MagneticField or callable or ndarray): magnetic field to fit - if callable, must accept (num_nodes,3) array of rpz coords as argument - and output (num_nodes,3) as the B field in cylindrical rpz basis. - 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_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 - verbose (int): verbosity level of fitting routine, > 0 prints residuals - """ - # We seek c in Ac = b - # A will be the BR, Bphi and BZ from each individual - # dommaschk potential basis function evaluated at each node - # c is the dommaschk potential coefficients - # c will be [B0, a_00, a_10, a_01, a_11... etc] - # b is the magnetic field at each node which we are fitting - if isinstance(field, _MagneticField): - B = field.compute_magnetic_field(coords) - elif callable(field): - B = field(coords) - else: # it must be the field evaluated at the passed-in coords - B = field - # TODO: add basis argument for if passed-in field or callable - # evaluates rpz or xyz basis magnetic field vector, - # and what basis coords is - - ######### - # make b - ######### - # we will have the rhs be 3*num_nodes in length (bc of vector B) - - rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") - - ##################### - # b is made, now do A - ##################### - num_modes = 1 + (max_l + 1) * (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 - # modes that don't fit symmetry, but in future - # should refactor code to have a 3rd index so that - # we have a = V_ml0, b = V_ml1, c = V_ml2, d = V_ml3 - # and the modes array can then be [m,l,x] where x is 0,1,2,3 - # and we dont need to keep track of a,b,c,d separately - - # TODO: technically we can drop some modes - # since if max_l=0, there are only ever nonzero terms for a and b - # and if max_m=0, there are only ever nonzero terms for a and c - # but since we are only fitting in a least squares sense, - # and max_l and max_m should probably be both nonzero anyways, - # this is not an issue right now - - # mode numbers - ms = [] - ls = [] - - # order of coeffs in the vector c are B0, a_ml, b_ml, c_ml, d_ml - a_s = [] - b_s = [] - c_s = [] - d_s = [] - zero_due_to_sym_inds = [] - abcd_zero_due_to_sym_inds = [ - [], - [], - [], - [], - ] # indices that should be 0 due to symmetry - for l in range(max_l + 1): - for m in range(max_m + 1): - if not sym: - pass # no sym, use all coefs - elif l // 2 == 0: - zero_due_to_sym_inds = [0, 3] # a=d=0 for even l with sym - elif l // 2 == 1: - zero_due_to_sym_inds = [1, 2] # b=c=0 for odd l with sym - for which_coef in range(4): - if which_coef == 0: - a_s.append(1) - elif which_coef == 1: - b_s.append(1) - elif which_coef == 2: - c_s.append(1) - elif which_coef == 3: - d_s.append(1) - if which_coef in zero_due_to_sym_inds: - 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): - abcd_zero_due_to_sym_inds[i] = jnp.asarray(abcd_zero_due_to_sym_inds[i]) - assert (len(a_s) + len(b_s) + len(c_s) + len(d_s)) == num_modes - 1 - params = { - "ms": ms, - "ls": ls, - "a_arr": a_s, - "b_arr": b_s, - "c_arr": c_s, - "d_arr": d_s, - "B0": 0.0, - } - n = ( - round(num_modes - 1) / 4 - ) # how many l-m mode pairs there are, also is len(a_s) - n = int(n) - domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) - - def get_B_dom(coords, X, ms, ls): - """Fxn wrapper to find jacobian of dommaschk B wrt coefs a,b,c,d.""" - # zero out any terms that should be zero due to symmetry, which - # we cataloged earlier for each a_arr,b_arr,c_arr,d_arr - # that way the resulting modes after pinv don't contain them either - return domm_field.compute_magnetic_field( - coords, - params={ - "ms": jnp.asarray(ms), - "ls": jnp.asarray(ls), - "a_arr": jnp.asarray(X[1 : n + 1]) * abcd_zero_due_to_sym_inds[0], - "b_arr": jnp.asarray(X[n + 1 : 2 * n + 1]) - * abcd_zero_due_to_sym_inds[1], - "c_arr": jnp.asarray(X[2 * n + 1 : 3 * n + 1]) - * abcd_zero_due_to_sym_inds[2], - "d_arr": jnp.asarray(X[3 * n + 1 : 4 * n + 1]) - * abcd_zero_due_to_sym_inds[3], - "B0": X[0], - }, - ) - - X = [] - for key in ["B0", "a_arr", "b_arr", "c_arr", "d_arr"]: - obj = params[key] - if isinstance(obj, list): - X += obj - else: - X += [obj] - X = jnp.asarray(X) - - jac = jit(jacfwd(get_B_dom, argnums=1))(coords, X, params["ms"], params["ls"]) - - A = jac.reshape((rhs.size, len(X)), order="F") - - # now solve Ac=b for the coefficients c - - # TODO: use min singular value to give sense of cond number? - c, res, _, _ = jnp.linalg.lstsq(A, rhs) - - if verbose > 0: - # res is a list of len(1) so index into it - print(f"Sum of Squares Residual of fit: {res[0]:1.4e} T") - - # recover the params from the c coefficient vector - B0 = c[0] - - # we zero out the terms that should be zero due to symmetry here - 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] - d_arr = c[3 * n + 1 : 4 * n + 1] * abcd_zero_due_to_sym_inds[3] - - return cls(ms, ls, a_arr, b_arr, c_arr, d_arr, B0) - - def field_line_integrate( r0, z0, @@ -1675,930 +1430,3 @@ def odefun(rpz, s): r = x[:, :, 0].T.reshape((len(phis), *rshape)) z = x[:, :, 2].T.reshape((len(phis), *rshape)) return r, z - - -### Dommaschk potential utility functions ### - -# based off Representations for vacuum potentials in stellarators -# https://doi.org/10.1016/0010-4655(86)90109-8 - -# written with naive for loops initially and can jax-ify later - -true_fun = lambda m_n: 0.0 # used for returning 0 when conditionals evaluate to True - - -@jit -def gamma(n): - """Gamma function, only implemented for integers (equiv to factorial of (n-1)).""" - return jnp.exp(gammaln(n)) - - -@jit -def alpha(m, n): - """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" - # modified for eqns 31 and 32 - - def false_fun(m_n): - m, n = m_n - return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) - - def bool_fun(n): - return n < 0 - - return cond( - bool_fun(n), - true_fun, - false_fun, - ( - m, - n, - ), - ) - - -@jit -def alphastar(m, n): - """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" - - def false_fun(m_n): # modified for eqns 31 and 32 - m, n = m_n - return (2 * n + m) * alpha(m, n) - - return cond(n < 0, true_fun, false_fun, (m, n)) - - -@jit -def beta(m, n): - """Beta of eq 28, modified for eqns 31 and 32.""" - - def false_fun(m_n): - m, n = m_n - return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) - - return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) - - -@jit -def betastar(m, n): - """Beta* of eq 28, modified for eqns 31 and 32.""" - - def false_fun(m_n): - m, n = m_n - return (2 * n - m) * beta(m, n) - - return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) - - -@jit -def gamma_n(m, n): - """gamma_n of eq 33.""" - - def body_fun(i, val): - return val + 1 / i + 1 / (m + i) - - def false_fun(m_n): - m, n = m_n - return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) - - return cond(n <= 0, true_fun, false_fun, (m, n)) - - -@jit -def gamma_nstar(m, n): - """gamma_n star of eq 33.""" - - def false_fun(m_n): - m, n = m_n - return (2 * n + m) * gamma_n(m, n) - - return cond(n <= 0, true_fun, false_fun, (m, n)) - - -@jit -def CD_m_k(R, m, k): - """Eq 31 of Dommaschk paper.""" - - def body_fun(j, val): - result = ( - val - + ( - -( - alpha(m, j) - * ( - alphastar(m, k - m - j) * jnp.log(R) - + gamma_nstar(m, k - m - j) - - alpha(m, k - m - j) - ) - - gamma_n(m, j) * alphastar(m, k - m - j) - + alpha(m, j) * betastar(m, k - j) - ) - * R ** (2 * j + m) - ) - + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) - ) - return result - - return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) - - -@jit -def CN_m_k(R, m, k): - """Eq 32 of Dommaschk paper.""" - - def body_fun(j, val): - result = ( - val - + ( - ( - alpha(m, j) - * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) - - gamma_n(m, j) * alpha(m, k - m - j) - + alpha(m, j) * beta(m, k - j) - ) - * R ** (2 * j + m) - ) - - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) - ) - return result - - return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) - - -@jit -def D_m_n(R, Z, m, n): - """D_m_n term in eqn 8 of Dommaschk paper.""" - # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper - - 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 - - return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) - - -@jit -def N_m_n(R, Z, m, n): - """N_m_n term in eqn 9 of Dommaschk paper.""" - # the sum comes from fact that N_mn = I_mn and the def of I_mn in eq 2 of the paper - - 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 - - return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) - - -@jit -def V_m_l(R, phi, Z, m, l, a, b, c, d): - """Eq 12 of Dommaschk paper. - - Parameters - ---------- - R,phi,Z : array-like - Cylindrical coordinates (1-D arrays of each of size num_eval_pts) - to evaluate the Dommaschk potential term at. - m : int - first index of V_m_l term - l : int - second index of V_m_l term - a : float - a_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*D_m_l - b : float - b_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*D_m_l - c : float - c_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*N_m_l-1 - d : float - d_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*N_m_l-1 - - Returns - ------- - value : array-like - Value of this V_m_l term evaluated at the given R,phi,Z points - (same size as the size of the given R,phi, or Z arrays). - - """ - return (a * jnp.cos(m * phi) + b * jnp.sin(m * phi)) * D_m_n(R, Z, m, l) + ( - c * jnp.cos(m * phi) + d * jnp.sin(m * phi) - ) * N_m_n(R, Z, m, l - 1) - - -@jit -def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): - """Eq 1 of Dommaschk paper. - - this is the total dommaschk potential for - a given set of m,l indices and their corresponding - coefficients a_ml, b_ml, c_ml d_ml. - - Parameters - ---------- - R,phi,Z : array-like - Cylindrical coordinates (1-D arrays of each of size num_eval_pts) - to evaluate the Dommaschk potential term at. - ms : 1D array-like of int - first indices of V_m_l terms - ls : 1D array-like of int - second indices of V_m_l terms - a_arr : 1D array-like of float - a_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*D_m_l - b_arr : 1D array-like of float - b_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*D_m_l - c_arr : 1D array-like of float - c_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*N_m_l-1 - d_arr : 1D array-like of float - d_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*N_m_l-1 - B0: float, toroidal magnetic field strength scale, this is the strength of the - 1/R part of the magnetic field and is the Bphi at R=1. - - Returns - ------- - value : array-like - Value of the total dommaschk potential evaluated - at the given R,phi,Z points - (same size as the size of the given R,phi, Z arrays). - """ - ms, ls, a_arr, b_arr, c_arr, d_arr = map( - jnp.atleast_1d, (ms, ls, a_arr, b_arr, c_arr, d_arr) - ) - R, phi, Z = map(jnp.atleast_1d, (R, phi, Z)) - R, phi, Z = jnp.broadcast_arrays(R, phi, Z) - ms, ls, a_arr, b_arr, c_arr, d_arr = jnp.broadcast_arrays( - ms, ls, a_arr, b_arr, c_arr, d_arr - ) - value = B0 * phi # phi term - - def body(i, val): - val += V_m_l(R, phi, Z, ms[i], ls[i], a_arr[i], b_arr[i], c_arr[i], d_arr[i]) - return val - - return fori_loop(0, len(ms), body, value) - - -class CurrentPotentialField(_MagneticField, FourierRZToroidalSurface): - """Magnetic field due to a surface current potential on a toroidal surface. - - surface current K is assumed given by - K = n x ∇ Φ - where: - n is the winding surface unit normal. - Phi is the current potential function, - which is a function of theta and zeta. - This function then uses biot-savart to find the - B field from this current density K on the surface. - - Parameters - ---------- - potential : callable - function to compute the current potential. Should have a signature of - the form potential(theta,zeta,**params) -> ndarray. - theta,zeta are poloidal and toroidal angles on the surface - potential_dtheta: callable - function to compute the theta derivative of the current potential - potential_dzeta: callable - function to compute the zeta derivative of the current potential - params : dict, optional - default parameters to pass to potential function (and its derivatives) - R_lmn, Z_lmn : array-like, shape(k,) - Fourier coefficients for winding surface R and Z in cylindrical coordinates - modes_R : array-like, shape(k,2) - poloidal and toroidal mode numbers [m,n] for R_lmn. - modes_Z : array-like, shape(k,2) - mode numbers associated with Z_lmn, defaults to modes_R - NFP : int - number of field periods - sym : bool - whether to enforce stellarator symmetry for the surface geometry. - Default is "auto" which enforces if modes are symmetric. If True, - non-symmetric modes will be truncated. - M, N: int or None - Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R - and modes_Z. - name : str - name for this field - check_orientation : bool - ensure that this surface has a right handed orientation. Do not set to False - unless you are sure the parameterization you have given is right handed - (ie, e_theta x e_zeta points outward from the surface). - - """ - - _io_attrs_ = ( - _MagneticField._io_attrs_ - + FourierRZToroidalSurface._io_attrs_ - + [ - "_params", - ] - ) - - def __init__( - self, - potential, - potential_dtheta, - potential_dzeta, - params=None, - R_lmn=None, - Z_lmn=None, - modes_R=None, - modes_Z=None, - NFP=1, - sym="auto", - M=None, - N=None, - name="", - check_orientation=True, - ): - assert callable(potential), "Potential must be callable!" - assert callable(potential_dtheta), "Potential derivative must be callable!" - assert callable(potential_dzeta), "Potential derivative must be callable!" - - self._potential = potential - self._potential_dtheta = potential_dtheta - self._potential_dzeta = potential_dzeta - self._params = params - - super().__init__( - R_lmn=R_lmn, - Z_lmn=Z_lmn, - modes_R=modes_R, - modes_Z=modes_Z, - NFP=NFP, - sym=sym, - M=M, - N=N, - name=name, - check_orientation=check_orientation, - ) - - @property - def params(self): - """Dict of parameters to pass to potential function and its derivatives.""" - return self._params - - @params.setter - def params(self, new): - warnif( - len(new) != len(self._params), - UserWarning, - "Length of new params is different from length of current params! " - "May cause errors unless potential function is also changed.", - ) - self._params = new - - @property - def potential(self): - """Potential function, signature (theta,zeta,**params) -> potential value.""" - return self._potential - - @potential.setter - def potential(self, new): - if new != self._potential: - assert callable(new), "Potential must be callable!" - self._potential = new - - @property - def potential_dtheta(self): - """Phi poloidal deriv. function, signature (theta,zeta,**params) -> value.""" - return self._potential_dtheta - - @potential_dtheta.setter - def potential_dtheta(self, new): - if new != self._potential_dtheta: - assert callable(new), "Potential derivative must be callable!" - self._potential_dtheta = new - - @property - def potential_dzeta(self): - """Phi toroidal deriv. function, signature (theta,zeta,**params) -> value.""" - return self._potential_dzeta - - @potential_dzeta.setter - def potential_dzeta(self, new): - if new != self._potential_dzeta: - assert callable(new), "Potential derivative must be callable!" - self._potential_dzeta = new - - def save(self, file_name, file_format=None, file_mode="w"): - """Save the object. - - **Not supported for this object! - - Parameters - ---------- - file_name : str file path OR file instance - location to save object - file_format : str (Default hdf5) - format of save file. Only used if file_name is a file path - file_mode : str (Default w - overwrite) - mode for save file. Only used if file_name is a file path - - """ - raise OSError( - "Saving CurrentPotentialField is not supported," - " as the potential function cannot be serialized." - ) - - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None - ): - """Compute magnetic field at a set of points. - - Parameters - ---------- - coords : array-like shape(n,3) - Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. - params : dict or array-like of dict, optional - Dictionary of optimizable parameters, eg field.params_dict. - basis : {"rpz", "xyz"} - Basis for input coordinates and returned magnetic field. - source_grid : Grid, int or None or array-like, optional - Source grid upon which to evaluate the surface current density K. - - Returns - ------- - field : ndarray, shape(N,3) - magnetic field at specified points - - """ - source_grid = source_grid or LinearGrid( - M=30 + 2 * self.M, - N=30 + 2 * self.N, - NFP=self.NFP, - ) - return _compute_magnetic_field_from_CurrentPotentialField( - field=self, - coords=coords, - params=params, - basis=basis, - source_grid=source_grid, - ) - - @classmethod - def from_surface( - cls, - surface, - potential, - potential_dtheta, - potential_dzeta, - params=None, - ): - """Create CurrentPotentialField using geometry provided by given surface. - - Parameters - ---------- - surface: FourierRZToroidalSurface, optional, default None - Existing FourierRZToroidalSurface object to create a - CurrentPotentialField with. - potential : callable - function to compute the current potential. Should have a signature of - the form potential(theta,zeta,**params) -> ndarray. - theta,zeta are poloidal and toroidal angles on the surface - potential_dtheta: callable - function to compute the theta derivative of the current potential - potential_dzeta: callable - function to compute the zeta derivative of the current potential - params : dict, optional - default parameters to pass to potential function (and its derivatives) - - """ - errorif( - not isinstance(surface, FourierRZToroidalSurface), - TypeError, - "Expected type FourierRZToroidalSurface for argument surface, " - f"instead got type {type(surface)}", - ) - - R_lmn = surface.R_lmn - Z_lmn = surface.Z_lmn - modes_R = surface._R_basis.modes[:, 1:] - modes_Z = surface._Z_basis.modes[:, 1:] - NFP = surface.NFP - sym = surface.sym - name = surface.name - - return cls( - potential, - potential_dtheta, - potential_dzeta, - params, - R_lmn, - Z_lmn, - modes_R, - modes_Z, - NFP, - sym, - name=name, - check_orientation=False, - ) - - -class FourierCurrentPotentialField( - _MagneticField, FourierRZToroidalSurface, Optimizable -): - """Magnetic field due to a surface current potential on a toroidal surface. - - surface current K is assumed given by - - K = n x ∇ Φ - Φ(θ,ζ) = Φₛᵥ(θ,ζ) + Gζ/2π + Iθ/2π - - where: - n is the winding surface unit normal. - Phi is the current potential function, - which is a function of theta and zeta, - and is given as a secular linear term in theta/zeta - and a double Fourier series in theta/zeta. - This function then uses biot-savart to find the - B field from this current density K on the surface. - - Parameters - ---------- - Phi_mn : ndarray - Fourier coefficients of the double FourierSeries part of the current potential. - modes_Phi : array-like, shape(k,2) - Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn - coefficients. - I : float - Net current linking the plasma and the surface toroidally - Denoted I in the algorithm - G : float - Net current linking the plasma and the surface poloidally - Denoted G in the algorithm - NOTE: a negative G will tend to produce a positive toroidal magnetic field - B in DESC, as in DESC the poloidal angle is taken to be positive - and increasing when going in the clockwise direction, which with the - convention n x grad(phi) will result in a toroidal field in the negative - toroidal direction. - sym_Phi : {False,"cos","sin"} - whether to enforce a given symmetry for the DoubleFourierSeries part of the - current potential. - M_Phi, N_Phi: int or None - Maximum poloidal and toroidal mode numbers for the single valued part of the - current potential. - R_lmn, Z_lmn : array-like, shape(k,) - Fourier coefficients for winding surface R and Z in cylindrical coordinates - modes_R : array-like, shape(k,2) - poloidal and toroidal mode numbers [m,n] for R_lmn. - modes_Z : array-like, shape(k,2) - mode numbers associated with Z_lmn, defaults to modes_R - NFP : int - number of field periods - sym : bool - whether to enforce stellarator symmetry for the surface geometry. - Default is "auto" which enforces if modes are symmetric. If True, - non-symmetric modes will be truncated. - M, N: int or None - Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R - and modes_Z. - name : str - name for this field - check_orientation : bool - ensure that this surface has a right handed orientation. Do not set to False - unless you are sure the parameterization you have given is right handed - (ie, e_theta x e_zeta points outward from the surface). - - """ - - _io_attrs_ = ( - _MagneticField._io_attrs_ - + FourierRZToroidalSurface._io_attrs_ - + ["_Phi_mn", "_I", "_G", "_Phi_basis", "_M_Phi", "_N_Phi", "_sym_Phi"] - ) - - def __init__( - self, - Phi_mn=np.array([0.0]), - modes_Phi=np.array([[0, 0]]), - I=0, - G=0, - sym_Phi=False, - M_Phi=None, - N_Phi=None, - R_lmn=None, - Z_lmn=None, - modes_R=None, - modes_Z=None, - NFP=1, - sym="auto", - M=None, - N=None, - name="", - check_orientation=True, - ): - Phi_mn, modes_Phi = map(np.asarray, (Phi_mn, modes_Phi)) - assert ( - Phi_mn.size == modes_Phi.shape[0] - ), "Phi_mn size and modes_Phi.shape[0] must be the same size!" - - assert np.issubdtype(modes_Phi.dtype, np.integer) - - M_Phi = setdefault(M_Phi, np.max(abs(modes_Phi[:, 0]))) - N_Phi = setdefault(N_Phi, np.max(abs(modes_Phi[:, 1]))) - - self._M_Phi = M_Phi - self._N_Phi = N_Phi - - self._sym_Phi = sym_Phi - self._Phi_basis = DoubleFourierSeries(M=M_Phi, N=N_Phi, NFP=NFP, sym=sym_Phi) - self._Phi_mn = copy_coeffs(Phi_mn, modes_Phi, self._Phi_basis.modes[:, 1:]) - - assert np.isscalar(I) or np.asarray(I).size == 1, "I must be a scalar" - assert np.isscalar(G) or np.asarray(G).size == 1, "G must be a scalar" - self._I = float(I) - self._G = float(G) - - super().__init__( - R_lmn=R_lmn, - Z_lmn=Z_lmn, - modes_R=modes_R, - modes_Z=modes_Z, - NFP=NFP, - sym=sym, - M=M, - N=N, - name=name, - check_orientation=check_orientation, - ) - - @optimizable_parameter - @property - def I(self): # noqa: E743 - """Net current linking the plasma and the surface toroidally.""" - return self._I - - @I.setter - def I(self, new): # noqa: E743 - assert np.isscalar(new) or np.asarray(new).size == 1, "I must be a scalar" - self._I = float(new) - - @optimizable_parameter - @property - def G(self): - """Net current linking the plasma and the surface poloidally.""" - return self._G - - @G.setter - def G(self, new): - assert np.isscalar(new) or np.asarray(new).size == 1, "G must be a scalar" - self._G = float(new) - - @optimizable_parameter - @property - def Phi_mn(self): - """Fourier coefficients describing single-valued part of potential.""" - return self._Phi_mn - - @Phi_mn.setter - def Phi_mn(self, new): - if len(new) == self.Phi_basis.num_modes: - self._Phi_mn = jnp.asarray(new) - else: - raise ValueError( - f"Phi_mn should have the same size as the basis, got {len(new)} for " - + f"basis with {self.Phi_basis.num_modes} modes." - ) - - @property - def Phi_basis(self): - """DoubleFourierSeries: Spectral basis for Phi.""" - return self._Phi_basis - - @property - def sym_Phi(self): - """str: Type of symmetry of periodic part of Phi (no symmetry if False).""" - return self._sym_Phi - - @property - def M_Phi(self): - """int: Poloidal resolution of periodic part of Phi.""" - return self._M_Phi - - @property - def N_Phi(self): - """int: Toroidal resolution of periodic part of Phi.""" - return self._N_Phi - - def change_Phi_resolution(self, M=None, N=None, NFP=None, sym_Phi=None): - """Change the maximum poloidal and toroidal resolution for Phi. - - Parameters - ---------- - M : int - Poloidal resolution to change Phi basis to. - If None, defaults to current self.Phi_basis poloidal resolution - N : int - Toroidal resolution to change Phi basis to. - If None, defaults to current self.Phi_basis toroidal resolution - NFP : int - Number of field periods for surface and Phi basis. - If None, defaults to current NFP. - Note: will change the NFP of the surface geometry as well as the - Phi basis. - sym_Phi : {"auto","cos","sin",False} - whether to enforce a given symmetry for the DoubleFourierSeries part of the - current potential. Default is "auto" which enforces if modes are symmetric. - If True, non-symmetric modes will be truncated. - - """ - M = M or self._M_Phi - N = N or self._M_Phi - NFP = NFP or self.NFP - sym_Phi = sym_Phi or self.sym_Phi - - Phi_modes_old = self.Phi_basis.modes - self.Phi_basis.change_resolution(M=M, N=N, NFP=self.NFP, sym=sym_Phi) - - self._Phi_mn = copy_coeffs(self.Phi_mn, Phi_modes_old, self.Phi_basis.modes) - self._M_Phi = M - self._N_Phi = N - self._sym_Phi = sym_Phi - self.change_resolution( - NFP=NFP - ) # make sure surface and Phi basis NFP are the same - - def compute_magnetic_field( - self, coords, params=None, basis="rpz", source_grid=None - ): - """Compute magnetic field at a set of points. - - Parameters - ---------- - coords : array-like shape(n,3) - Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. - params : dict or array-like of dict, optional - Dictionary of optimizable parameters, eg field.params_dict. - basis : {"rpz", "xyz"} - Basis for input coordinates and returned magnetic field. - source_grid : Grid, int or None or array-like, optional - Source grid upon which to evaluate the surface current density K. - - Returns - ------- - field : ndarray, shape(N,3) - magnetic field at specified points - - """ - source_grid = source_grid or LinearGrid( - M=30 + 2 * max(self.M, self.M_Phi), - N=30 + 2 * max(self.N, self.N_Phi), - NFP=self.NFP, - ) - return _compute_magnetic_field_from_CurrentPotentialField( - field=self, - coords=coords, - params=params, - basis=basis, - source_grid=source_grid, - ) - - @classmethod - def from_surface( - cls, - surface, - Phi_mn=np.array([0.0]), - modes_Phi=np.array([[0, 0]]), - I=0, - G=0, - sym_Phi=False, - M_Phi=None, - N_Phi=None, - ): - """Create FourierCurrentPotentialField using geometry of given surface. - - Parameters - ---------- - surface: FourierRZToroidalSurface, optional, default None - Existing FourierRZToroidalSurface object to create a - CurrentPotentialField with. - Phi_mn : ndarray - Fourier coefficients of the double FourierSeries of the current potential. - Should correspond to the given DoubleFourierSeries basis object passed in. - modes_Phi : array-like, shape(k,2) - Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn - coefficients - I : float - Net current linking the plasma and the surface toroidally - Denoted I in the algorithm - G : float - Net current linking the plasma and the surface poloidally - Denoted G in the algorithm - NOTE: a negative G will tend to produce a positive toroidal magnetic field - B in DESC, as in DESC the poloidal angle is taken to be positive - and increasing when going in the clockwise direction, which with the - convention n x grad(phi) will result in a toroidal field in the negative - toroidal direction. - sym_Phi : {False,"cos","sin"} - whether to enforce a given symmetry for the DoubleFourierSeries part of the - current potential. - M_Phi, N_Phi: int or None - Maximum poloidal and toroidal mode numbers for the single valued part of the - current potential. - - """ - if not isinstance(surface, FourierRZToroidalSurface): - raise TypeError( - "Expected type FourierRZToroidalSurface for argument surface, " - f"instead got type {type(surface)}" - ) - R_lmn = surface.R_lmn - Z_lmn = surface.Z_lmn - modes_R = surface._R_basis.modes[:, 1:] - modes_Z = surface._Z_basis.modes[:, 1:] - NFP = surface.NFP - sym = surface.sym - name = surface.name - - return cls( - Phi_mn=Phi_mn, - modes_Phi=modes_Phi, - I=I, - G=G, - sym_Phi=sym_Phi, - M_Phi=M_Phi, - N_Phi=N_Phi, - R_lmn=R_lmn, - Z_lmn=Z_lmn, - modes_R=modes_R, - modes_Z=modes_Z, - NFP=NFP, - sym=sym, - name=name, - check_orientation=False, - ) - - -def _compute_magnetic_field_from_CurrentPotentialField( - field, - coords, - source_grid, - params=None, - basis="rpz", -): - """Compute magnetic field at a set of points. - - Parameters - ---------- - field : CurrentPotentialField or FourierCurrentPotentialField - current potential field object from which to compute magnetic field. - coords : array-like shape(N,3) - cylindrical or cartesian coordinates - source_grid : Grid, - source grid upon which to evaluate the surface current density K - params : dict, optional - parameters to pass to compute function - should include the potential - basis : {"rpz", "xyz"} - basis for input coordinates and returned magnetic field - - - Returns - ------- - field : ndarray, shape(N,3) - magnetic field at specified points - - """ - assert basis.lower() in ["rpz", "xyz"] - coords = jnp.atleast_2d(coords) - if basis == "rpz": - coords = rpz2xyz(coords) - - # compute surface current, and store grid quantities - # needed for integration in class - # TODO: does this have to be xyz, or can it be computed in rpz as well? - data = field.compute(["K", "x"], grid=source_grid, basis="xyz", params=params) - - _rs = xyz2rpz(data["x"]) - _K = xyz2rpz_vec(data["K"], phi=source_grid.nodes[:, 2]) - - # surface element, must divide by NFP to remove the NFP multiple on the - # surface grid weights, as we account for that when doing the for loop - # over NFP - _dV = source_grid.weights * data["|e_theta x e_zeta|"] / source_grid.NFP - - def nfp_loop(j, f): - # calculate (by rotating) rs, rs_t, rz_t - phi = (source_grid.nodes[:, 2] + j * 2 * jnp.pi / source_grid.NFP) % ( - 2 * jnp.pi - ) - # new coords are just old R,Z at a new phi (bc of discrete NFP symmetry) - rs = jnp.vstack((_rs[:, 0], phi, _rs[:, 2])).T - rs = rpz2xyz(rs) - K = rpz2xyz_vec(_K, phi=phi) - fj = biot_savart_general( - coords, - rs, - K, - _dV, - ) - f += fj - return f - - B = fori_loop(0, source_grid.NFP, nfp_loop, jnp.zeros_like(coords)) - if basis == "rpz": - B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) - return B diff --git a/desc/magnetic_fields/_current_potential.py b/desc/magnetic_fields/_current_potential.py new file mode 100644 index 0000000000..c414285512 --- /dev/null +++ b/desc/magnetic_fields/_current_potential.py @@ -0,0 +1,678 @@ +"""Magnetic field due to sheet current on a winding surface.""" + + +import numpy as np + +from desc.backend import fori_loop, jnp +from desc.basis import DoubleFourierSeries +from desc.compute import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec +from desc.geometry import FourierRZToroidalSurface +from desc.grid import LinearGrid +from desc.optimizable import Optimizable, optimizable_parameter +from desc.utils import copy_coeffs, errorif, setdefault, warnif + +from ._core import _MagneticField, biot_savart_general + + +class CurrentPotentialField(_MagneticField, FourierRZToroidalSurface): + """Magnetic field due to a surface current potential on a toroidal surface. + + surface current K is assumed given by + K = n x ∇ Φ + where: + n is the winding surface unit normal. + Phi is the current potential function, + which is a function of theta and zeta. + This function then uses biot-savart to find the + B field from this current density K on the surface. + + Parameters + ---------- + potential : callable + function to compute the current potential. Should have a signature of + the form potential(theta,zeta,**params) -> ndarray. + theta,zeta are poloidal and toroidal angles on the surface + potential_dtheta: callable + function to compute the theta derivative of the current potential + potential_dzeta: callable + function to compute the zeta derivative of the current potential + params : dict, optional + default parameters to pass to potential function (and its derivatives) + R_lmn, Z_lmn : array-like, shape(k,) + Fourier coefficients for winding surface R and Z in cylindrical coordinates + modes_R : array-like, shape(k,2) + poloidal and toroidal mode numbers [m,n] for R_lmn. + modes_Z : array-like, shape(k,2) + mode numbers associated with Z_lmn, defaults to modes_R + NFP : int + number of field periods + sym : bool + whether to enforce stellarator symmetry for the surface geometry. + Default is "auto" which enforces if modes are symmetric. If True, + non-symmetric modes will be truncated. + M, N: int or None + Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R + and modes_Z. + name : str + name for this field + check_orientation : bool + ensure that this surface has a right handed orientation. Do not set to False + unless you are sure the parameterization you have given is right handed + (ie, e_theta x e_zeta points outward from the surface). + + """ + + _io_attrs_ = ( + _MagneticField._io_attrs_ + + FourierRZToroidalSurface._io_attrs_ + + [ + "_params", + ] + ) + + def __init__( + self, + potential, + potential_dtheta, + potential_dzeta, + params=None, + R_lmn=None, + Z_lmn=None, + modes_R=None, + modes_Z=None, + NFP=1, + sym="auto", + M=None, + N=None, + name="", + check_orientation=True, + ): + assert callable(potential), "Potential must be callable!" + assert callable(potential_dtheta), "Potential derivative must be callable!" + assert callable(potential_dzeta), "Potential derivative must be callable!" + + self._potential = potential + self._potential_dtheta = potential_dtheta + self._potential_dzeta = potential_dzeta + self._params = params + + super().__init__( + R_lmn=R_lmn, + Z_lmn=Z_lmn, + modes_R=modes_R, + modes_Z=modes_Z, + NFP=NFP, + sym=sym, + M=M, + N=N, + name=name, + check_orientation=check_orientation, + ) + + @property + def params(self): + """Dict of parameters to pass to potential function and its derivatives.""" + return self._params + + @params.setter + def params(self, new): + warnif( + len(new) != len(self._params), + UserWarning, + "Length of new params is different from length of current params! " + "May cause errors unless potential function is also changed.", + ) + self._params = new + + @property + def potential(self): + """Potential function, signature (theta,zeta,**params) -> potential value.""" + return self._potential + + @potential.setter + def potential(self, new): + if new != self._potential: + assert callable(new), "Potential must be callable!" + self._potential = new + + @property + def potential_dtheta(self): + """Phi poloidal deriv. function, signature (theta,zeta,**params) -> value.""" + return self._potential_dtheta + + @potential_dtheta.setter + def potential_dtheta(self, new): + if new != self._potential_dtheta: + assert callable(new), "Potential derivative must be callable!" + self._potential_dtheta = new + + @property + def potential_dzeta(self): + """Phi toroidal deriv. function, signature (theta,zeta,**params) -> value.""" + return self._potential_dzeta + + @potential_dzeta.setter + def potential_dzeta(self, new): + if new != self._potential_dzeta: + assert callable(new), "Potential derivative must be callable!" + self._potential_dzeta = new + + def save(self, file_name, file_format=None, file_mode="w"): + """Save the object. + + **Not supported for this object! + + Parameters + ---------- + file_name : str file path OR file instance + location to save object + file_format : str (Default hdf5) + format of save file. Only used if file_name is a file path + file_mode : str (Default w - overwrite) + mode for save file. Only used if file_name is a file path + + """ + raise OSError( + "Saving CurrentPotentialField is not supported," + " as the potential function cannot be serialized." + ) + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + source_grid = source_grid or LinearGrid( + M=30 + 2 * self.M, + N=30 + 2 * self.N, + NFP=self.NFP, + ) + return _compute_magnetic_field_from_CurrentPotentialField( + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, + ) + + @classmethod + def from_surface( + cls, + surface, + potential, + potential_dtheta, + potential_dzeta, + params=None, + ): + """Create CurrentPotentialField using geometry provided by given surface. + + Parameters + ---------- + surface: FourierRZToroidalSurface, optional, default None + Existing FourierRZToroidalSurface object to create a + CurrentPotentialField with. + potential : callable + function to compute the current potential. Should have a signature of + the form potential(theta,zeta,**params) -> ndarray. + theta,zeta are poloidal and toroidal angles on the surface + potential_dtheta: callable + function to compute the theta derivative of the current potential + potential_dzeta: callable + function to compute the zeta derivative of the current potential + params : dict, optional + default parameters to pass to potential function (and its derivatives) + + """ + errorif( + not isinstance(surface, FourierRZToroidalSurface), + TypeError, + "Expected type FourierRZToroidalSurface for argument surface, " + f"instead got type {type(surface)}", + ) + + R_lmn = surface.R_lmn + Z_lmn = surface.Z_lmn + modes_R = surface._R_basis.modes[:, 1:] + modes_Z = surface._Z_basis.modes[:, 1:] + NFP = surface.NFP + sym = surface.sym + name = surface.name + + return cls( + potential, + potential_dtheta, + potential_dzeta, + params, + R_lmn, + Z_lmn, + modes_R, + modes_Z, + NFP, + sym, + name=name, + check_orientation=False, + ) + + +class FourierCurrentPotentialField( + _MagneticField, FourierRZToroidalSurface, Optimizable +): + """Magnetic field due to a surface current potential on a toroidal surface. + + surface current K is assumed given by + + K = n x ∇ Φ + Φ(θ,ζ) = Φₛᵥ(θ,ζ) + Gζ/2π + Iθ/2π + + where: + n is the winding surface unit normal. + Phi is the current potential function, + which is a function of theta and zeta, + and is given as a secular linear term in theta/zeta + and a double Fourier series in theta/zeta. + This function then uses biot-savart to find the + B field from this current density K on the surface. + + Parameters + ---------- + Phi_mn : ndarray + Fourier coefficients of the double FourierSeries part of the current potential. + modes_Phi : array-like, shape(k,2) + Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn + coefficients. + I : float + Net current linking the plasma and the surface toroidally + Denoted I in the algorithm + G : float + Net current linking the plasma and the surface poloidally + Denoted G in the algorithm + NOTE: a negative G will tend to produce a positive toroidal magnetic field + B in DESC, as in DESC the poloidal angle is taken to be positive + and increasing when going in the clockwise direction, which with the + convention n x grad(phi) will result in a toroidal field in the negative + toroidal direction. + sym_Phi : {False,"cos","sin"} + whether to enforce a given symmetry for the DoubleFourierSeries part of the + current potential. + M_Phi, N_Phi: int or None + Maximum poloidal and toroidal mode numbers for the single valued part of the + current potential. + R_lmn, Z_lmn : array-like, shape(k,) + Fourier coefficients for winding surface R and Z in cylindrical coordinates + modes_R : array-like, shape(k,2) + poloidal and toroidal mode numbers [m,n] for R_lmn. + modes_Z : array-like, shape(k,2) + mode numbers associated with Z_lmn, defaults to modes_R + NFP : int + number of field periods + sym : bool + whether to enforce stellarator symmetry for the surface geometry. + Default is "auto" which enforces if modes are symmetric. If True, + non-symmetric modes will be truncated. + M, N: int or None + Maximum poloidal and toroidal mode numbers. Defaults to maximum from modes_R + and modes_Z. + name : str + name for this field + check_orientation : bool + ensure that this surface has a right handed orientation. Do not set to False + unless you are sure the parameterization you have given is right handed + (ie, e_theta x e_zeta points outward from the surface). + + """ + + _io_attrs_ = ( + _MagneticField._io_attrs_ + + FourierRZToroidalSurface._io_attrs_ + + ["_Phi_mn", "_I", "_G", "_Phi_basis", "_M_Phi", "_N_Phi", "_sym_Phi"] + ) + + def __init__( + self, + Phi_mn=np.array([0.0]), + modes_Phi=np.array([[0, 0]]), + I=0, + G=0, + sym_Phi=False, + M_Phi=None, + N_Phi=None, + R_lmn=None, + Z_lmn=None, + modes_R=None, + modes_Z=None, + NFP=1, + sym="auto", + M=None, + N=None, + name="", + check_orientation=True, + ): + Phi_mn, modes_Phi = map(np.asarray, (Phi_mn, modes_Phi)) + assert ( + Phi_mn.size == modes_Phi.shape[0] + ), "Phi_mn size and modes_Phi.shape[0] must be the same size!" + + assert np.issubdtype(modes_Phi.dtype, np.integer) + + M_Phi = setdefault(M_Phi, np.max(abs(modes_Phi[:, 0]))) + N_Phi = setdefault(N_Phi, np.max(abs(modes_Phi[:, 1]))) + + self._M_Phi = M_Phi + self._N_Phi = N_Phi + + self._sym_Phi = sym_Phi + self._Phi_basis = DoubleFourierSeries(M=M_Phi, N=N_Phi, NFP=NFP, sym=sym_Phi) + self._Phi_mn = copy_coeffs(Phi_mn, modes_Phi, self._Phi_basis.modes[:, 1:]) + + assert np.isscalar(I) or np.asarray(I).size == 1, "I must be a scalar" + assert np.isscalar(G) or np.asarray(G).size == 1, "G must be a scalar" + self._I = float(I) + self._G = float(G) + + super().__init__( + R_lmn=R_lmn, + Z_lmn=Z_lmn, + modes_R=modes_R, + modes_Z=modes_Z, + NFP=NFP, + sym=sym, + M=M, + N=N, + name=name, + check_orientation=check_orientation, + ) + + @optimizable_parameter + @property + def I(self): # noqa: E743 + """Net current linking the plasma and the surface toroidally.""" + return self._I + + @I.setter + def I(self, new): # noqa: E743 + assert np.isscalar(new) or np.asarray(new).size == 1, "I must be a scalar" + self._I = float(new) + + @optimizable_parameter + @property + def G(self): + """Net current linking the plasma and the surface poloidally.""" + return self._G + + @G.setter + def G(self, new): + assert np.isscalar(new) or np.asarray(new).size == 1, "G must be a scalar" + self._G = float(new) + + @optimizable_parameter + @property + def Phi_mn(self): + """Fourier coefficients describing single-valued part of potential.""" + return self._Phi_mn + + @Phi_mn.setter + def Phi_mn(self, new): + if len(new) == self.Phi_basis.num_modes: + self._Phi_mn = jnp.asarray(new) + else: + raise ValueError( + f"Phi_mn should have the same size as the basis, got {len(new)} for " + + f"basis with {self.Phi_basis.num_modes} modes." + ) + + @property + def Phi_basis(self): + """DoubleFourierSeries: Spectral basis for Phi.""" + return self._Phi_basis + + @property + def sym_Phi(self): + """str: Type of symmetry of periodic part of Phi (no symmetry if False).""" + return self._sym_Phi + + @property + def M_Phi(self): + """int: Poloidal resolution of periodic part of Phi.""" + return self._M_Phi + + @property + def N_Phi(self): + """int: Toroidal resolution of periodic part of Phi.""" + return self._N_Phi + + def change_Phi_resolution(self, M=None, N=None, NFP=None, sym_Phi=None): + """Change the maximum poloidal and toroidal resolution for Phi. + + Parameters + ---------- + M : int + Poloidal resolution to change Phi basis to. + If None, defaults to current self.Phi_basis poloidal resolution + N : int + Toroidal resolution to change Phi basis to. + If None, defaults to current self.Phi_basis toroidal resolution + NFP : int + Number of field periods for surface and Phi basis. + If None, defaults to current NFP. + Note: will change the NFP of the surface geometry as well as the + Phi basis. + sym_Phi : {"auto","cos","sin",False} + whether to enforce a given symmetry for the DoubleFourierSeries part of the + current potential. Default is "auto" which enforces if modes are symmetric. + If True, non-symmetric modes will be truncated. + + """ + M = M or self._M_Phi + N = N or self._M_Phi + NFP = NFP or self.NFP + sym_Phi = sym_Phi or self.sym_Phi + + Phi_modes_old = self.Phi_basis.modes + self.Phi_basis.change_resolution(M=M, N=N, NFP=self.NFP, sym=sym_Phi) + + self._Phi_mn = copy_coeffs(self.Phi_mn, Phi_modes_old, self.Phi_basis.modes) + self._M_Phi = M + self._N_Phi = N + self._sym_Phi = sym_Phi + self.change_resolution( + NFP=NFP + ) # make sure surface and Phi basis NFP are the same + + def compute_magnetic_field( + self, coords, params=None, basis="rpz", source_grid=None + ): + """Compute magnetic field at a set of points. + + Parameters + ---------- + coords : array-like shape(n,3) + Nodes to evaluate field at in [R,phi,Z] or [X,Y,Z] coordinates. + params : dict or array-like of dict, optional + Dictionary of optimizable parameters, eg field.params_dict. + basis : {"rpz", "xyz"} + Basis for input coordinates and returned magnetic field. + source_grid : Grid, int or None or array-like, optional + Source grid upon which to evaluate the surface current density K. + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + source_grid = source_grid or LinearGrid( + M=30 + 2 * max(self.M, self.M_Phi), + N=30 + 2 * max(self.N, self.N_Phi), + NFP=self.NFP, + ) + return _compute_magnetic_field_from_CurrentPotentialField( + field=self, + coords=coords, + params=params, + basis=basis, + source_grid=source_grid, + ) + + @classmethod + def from_surface( + cls, + surface, + Phi_mn=np.array([0.0]), + modes_Phi=np.array([[0, 0]]), + I=0, + G=0, + sym_Phi=False, + M_Phi=None, + N_Phi=None, + ): + """Create FourierCurrentPotentialField using geometry of given surface. + + Parameters + ---------- + surface: FourierRZToroidalSurface, optional, default None + Existing FourierRZToroidalSurface object to create a + CurrentPotentialField with. + Phi_mn : ndarray + Fourier coefficients of the double FourierSeries of the current potential. + Should correspond to the given DoubleFourierSeries basis object passed in. + modes_Phi : array-like, shape(k,2) + Poloidal and Toroidal mode numbers corresponding to passed-in Phi_mn + coefficients + I : float + Net current linking the plasma and the surface toroidally + Denoted I in the algorithm + G : float + Net current linking the plasma and the surface poloidally + Denoted G in the algorithm + NOTE: a negative G will tend to produce a positive toroidal magnetic field + B in DESC, as in DESC the poloidal angle is taken to be positive + and increasing when going in the clockwise direction, which with the + convention n x grad(phi) will result in a toroidal field in the negative + toroidal direction. + sym_Phi : {False,"cos","sin"} + whether to enforce a given symmetry for the DoubleFourierSeries part of the + current potential. + M_Phi, N_Phi: int or None + Maximum poloidal and toroidal mode numbers for the single valued part of the + current potential. + + """ + if not isinstance(surface, FourierRZToroidalSurface): + raise TypeError( + "Expected type FourierRZToroidalSurface for argument surface, " + f"instead got type {type(surface)}" + ) + R_lmn = surface.R_lmn + Z_lmn = surface.Z_lmn + modes_R = surface._R_basis.modes[:, 1:] + modes_Z = surface._Z_basis.modes[:, 1:] + NFP = surface.NFP + sym = surface.sym + name = surface.name + + return cls( + Phi_mn=Phi_mn, + modes_Phi=modes_Phi, + I=I, + G=G, + sym_Phi=sym_Phi, + M_Phi=M_Phi, + N_Phi=N_Phi, + R_lmn=R_lmn, + Z_lmn=Z_lmn, + modes_R=modes_R, + modes_Z=modes_Z, + NFP=NFP, + sym=sym, + name=name, + check_orientation=False, + ) + + +def _compute_magnetic_field_from_CurrentPotentialField( + field, + coords, + source_grid, + params=None, + basis="rpz", +): + """Compute magnetic field at a set of points. + + Parameters + ---------- + field : CurrentPotentialField or FourierCurrentPotentialField + current potential field object from which to compute magnetic field. + coords : array-like shape(N,3) + cylindrical or cartesian coordinates + source_grid : Grid, + source grid upon which to evaluate the surface current density K + params : dict, optional + parameters to pass to compute function + should include the potential + basis : {"rpz", "xyz"} + basis for input coordinates and returned magnetic field + + + Returns + ------- + field : ndarray, shape(N,3) + magnetic field at specified points + + """ + assert basis.lower() in ["rpz", "xyz"] + coords = jnp.atleast_2d(coords) + if basis == "rpz": + coords = rpz2xyz(coords) + + # compute surface current, and store grid quantities + # needed for integration in class + # TODO: does this have to be xyz, or can it be computed in rpz as well? + data = field.compute(["K", "x"], grid=source_grid, basis="xyz", params=params) + + _rs = xyz2rpz(data["x"]) + _K = xyz2rpz_vec(data["K"], phi=source_grid.nodes[:, 2]) + + # surface element, must divide by NFP to remove the NFP multiple on the + # surface grid weights, as we account for that when doing the for loop + # over NFP + _dV = source_grid.weights * data["|e_theta x e_zeta|"] / source_grid.NFP + + def nfp_loop(j, f): + # calculate (by rotating) rs, rs_t, rz_t + phi = (source_grid.nodes[:, 2] + j * 2 * jnp.pi / source_grid.NFP) % ( + 2 * jnp.pi + ) + # new coords are just old R,Z at a new phi (bc of discrete NFP symmetry) + rs = jnp.vstack((_rs[:, 0], phi, _rs[:, 2])).T + rs = rpz2xyz(rs) + K = rpz2xyz_vec(_K, phi=phi) + fj = biot_savart_general( + coords, + rs, + K, + _dV, + ) + f += fj + return f + + B = fori_loop(0, source_grid.NFP, nfp_loop, jnp.zeros_like(coords)) + if basis == "rpz": + B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) + return B diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py new file mode 100644 index 0000000000..46681075c1 --- /dev/null +++ b/desc/magnetic_fields/_dommaschk.py @@ -0,0 +1,512 @@ +"""Dommaschk potential utility functions. + +based off Representations for vacuum potentials in stellarators +https://doi.org/10.1016/0010-4655(86)90109-8 + +written with naive for loops initially and can jax-ify later +""" +from desc.backend import cond, fori_loop, gammaln, jit, jnp +from desc.derivatives import Derivative + +from ._core import ScalarPotentialField, _MagneticField + + +class DommaschkPotentialField(ScalarPotentialField): + """Magnetic field due to a Dommaschk scalar magnetic potential in rpz coordinates. + + From Dommaschk 1986 paper https://doi.org/10.1016/0010-4655(86)90109-8 + + this is the field due to the dommaschk potential (eq. 1) for + a given set of m,l indices and their corresponding + coefficients a_ml, b_ml, c_ml d_ml. + + Parameters + ---------- + ms : 1D array-like of int + first indices of V_m_l terms (eq. 12 of reference) + ls : 1D array-like of int + second indices of V_m_l terms (eq. 12 of reference) + 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 + b_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*D_m_l terms + c_arr : 1D array-like of float + c_m_l coefficients of V_m_l terms, which multiply the cos(m*phi)*N_m_l-1 term + d_arr : 1D array-like of float + d_m_l coefficients of V_m_l terms, which multiply the sin(m*phi)*N_m_l-1 terms + B0: float + scale strength of the magnetic field's 1/R portion + + """ + + def __init__( + self, + ms=jnp.array([0]), + ls=jnp.array([0]), + a_arr=jnp.array([0.0]), + b_arr=jnp.array([0.0]), + c_arr=jnp.array([0.0]), + d_arr=jnp.array([0.0]), + B0=1.0, + ): + ms = jnp.atleast_1d(ms) + ls = jnp.atleast_1d(ls) + a_arr = jnp.atleast_1d(a_arr) + b_arr = jnp.atleast_1d(b_arr) + c_arr = jnp.atleast_1d(c_arr) + d_arr = jnp.atleast_1d(d_arr) + + assert ( + ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size + ), "Passed in arrays must all be of the same size!" + assert not jnp.any( + jnp.logical_or(ms < 0, ls < 0) + ), "m and l mode numbers must be >= 0!" + assert ( + jnp.isscalar(B0) or jnp.atleast_1d(B0).size == 1 + ), "B0 should be a scalar value!" + + params = {} + params["ms"] = ms + params["ls"] = ls + params["a_arr"] = a_arr + params["b_arr"] = b_arr + params["c_arr"] = c_arr + params["d_arr"] = d_arr + params["B0"] = B0 + + super().__init__(dommaschk_potential, params) + + @classmethod + def fit_magnetic_field( # noqa: C901 - FIXME - simplify + cls, field, coords, max_m, max_l, sym=False, verbose=1 + ): + """Fit a vacuum magnetic field with a Dommaschk Potential field. + + Parameters + ---------- + field (MagneticField or callable or ndarray): magnetic field to fit + if callable, must accept (num_nodes,3) array of rpz coords as argument + and output (num_nodes,3) as the B field in cylindrical rpz basis. + 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_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 + verbose (int): verbosity level of fitting routine, > 0 prints residuals + """ + # We seek c in Ac = b + # A will be the BR, Bphi and BZ from each individual + # dommaschk potential basis function evaluated at each node + # c is the dommaschk potential coefficients + # c will be [B0, a_00, a_10, a_01, a_11... etc] + # b is the magnetic field at each node which we are fitting + if isinstance(field, _MagneticField): + B = field.compute_magnetic_field(coords) + elif callable(field): + B = field(coords) + else: # it must be the field evaluated at the passed-in coords + B = field + # TODO: add basis argument for if passed-in field or callable + # evaluates rpz or xyz basis magnetic field vector, + # and what basis coords is + + ######### + # make b + ######### + # we will have the rhs be 3*num_nodes in length (bc of vector B) + + rhs = jnp.vstack((B[:, 0], B[:, 1], B[:, 2])).T.flatten(order="F") + + ##################### + # b is made, now do A + ##################### + num_modes = 1 + (max_l + 1) * (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 + # modes that don't fit symmetry, but in future + # should refactor code to have a 3rd index so that + # we have a = V_ml0, b = V_ml1, c = V_ml2, d = V_ml3 + # and the modes array can then be [m,l,x] where x is 0,1,2,3 + # and we dont need to keep track of a,b,c,d separately + + # TODO: technically we can drop some modes + # since if max_l=0, there are only ever nonzero terms for a and b + # and if max_m=0, there are only ever nonzero terms for a and c + # but since we are only fitting in a least squares sense, + # and max_l and max_m should probably be both nonzero anyways, + # this is not an issue right now + + # mode numbers + ms = [] + ls = [] + + # order of coeffs in the vector c are B0, a_ml, b_ml, c_ml, d_ml + a_s = [] + b_s = [] + c_s = [] + d_s = [] + zero_due_to_sym_inds = [] + abcd_zero_due_to_sym_inds = [ + [], + [], + [], + [], + ] # indices that should be 0 due to symmetry + for l in range(max_l + 1): + for m in range(max_m + 1): + if not sym: + pass # no sym, use all coefs + elif l // 2 == 0: + zero_due_to_sym_inds = [0, 3] # a=d=0 for even l with sym + elif l // 2 == 1: + zero_due_to_sym_inds = [1, 2] # b=c=0 for odd l with sym + for which_coef in range(4): + if which_coef == 0: + a_s.append(1) + elif which_coef == 1: + b_s.append(1) + elif which_coef == 2: + c_s.append(1) + elif which_coef == 3: + d_s.append(1) + if which_coef in zero_due_to_sym_inds: + 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): + abcd_zero_due_to_sym_inds[i] = jnp.asarray(abcd_zero_due_to_sym_inds[i]) + assert (len(a_s) + len(b_s) + len(c_s) + len(d_s)) == num_modes - 1 + params = { + "ms": ms, + "ls": ls, + "a_arr": a_s, + "b_arr": b_s, + "c_arr": c_s, + "d_arr": d_s, + "B0": 0.0, + } + n = ( + round(num_modes - 1) / 4 + ) # how many l-m mode pairs there are, also is len(a_s) + n = int(n) + domm_field = DommaschkPotentialField(0, 0, 0, 0, 0, 0, 1) + + def get_B_dom(coords, X, ms, ls): + """Fxn wrapper to find jacobian of dommaschk B wrt coefs a,b,c,d.""" + # zero out any terms that should be zero due to symmetry, which + # we cataloged earlier for each a_arr,b_arr,c_arr,d_arr + # that way the resulting modes after pinv don't contain them either + return domm_field.compute_magnetic_field( + coords, + params={ + "ms": jnp.asarray(ms), + "ls": jnp.asarray(ls), + "a_arr": jnp.asarray(X[1 : n + 1]) * abcd_zero_due_to_sym_inds[0], + "b_arr": jnp.asarray(X[n + 1 : 2 * n + 1]) + * abcd_zero_due_to_sym_inds[1], + "c_arr": jnp.asarray(X[2 * n + 1 : 3 * n + 1]) + * abcd_zero_due_to_sym_inds[2], + "d_arr": jnp.asarray(X[3 * n + 1 : 4 * n + 1]) + * abcd_zero_due_to_sym_inds[3], + "B0": X[0], + }, + ) + + X = [] + for key in ["B0", "a_arr", "b_arr", "c_arr", "d_arr"]: + obj = params[key] + if isinstance(obj, list): + X += obj + else: + X += [obj] + X = jnp.asarray(X) + + jac = jit(Derivative(get_B_dom, argnum=1))( + coords, X, params["ms"], params["ls"] + ) + + A = jac.reshape((rhs.size, len(X)), order="F") + + # now solve Ac=b for the coefficients c + + # TODO: use min singular value to give sense of cond number? + c, res, _, _ = jnp.linalg.lstsq(A, rhs) + + if verbose > 0: + # res is a list of len(1) so index into it + print(f"Sum of Squares Residual of fit: {res[0]:1.4e} T") + + # recover the params from the c coefficient vector + B0 = c[0] + + # we zero out the terms that should be zero due to symmetry here + 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] + d_arr = c[3 * n + 1 : 4 * n + 1] * abcd_zero_due_to_sym_inds[3] + + return cls(ms, ls, a_arr, b_arr, c_arr, d_arr, B0) + + +true_fun = lambda m_n: 0.0 # used for returning 0 when conditionals evaluate to True + + +@jit +def gamma(n): + """Gamma function, only implemented for integers (equiv to factorial of (n-1)).""" + return jnp.exp(gammaln(n)) + + +@jit +def alpha(m, n): + """Alpha of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + # modified for eqns 31 and 32 + + def false_fun(m_n): + m, n = m_n + return (-1) ** n / (gamma(m + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + m)) + + def bool_fun(n): + return n < 0 + + return cond( + bool_fun(n), + true_fun, + false_fun, + ( + m, + n, + ), + ) + + +@jit +def alphastar(m, n): + """Alphastar of eq 27, 1st ind comes from C_m_k, 2nd is the subscript of alpha.""" + + def false_fun(m_n): # modified for eqns 31 and 32 + m, n = m_n + return (2 * n + m) * alpha(m, n) + + return cond(n < 0, true_fun, false_fun, (m, n)) + + +@jit +def beta(m, n): + """Beta of eq 28, modified for eqns 31 and 32.""" + + def false_fun(m_n): + m, n = m_n + return gamma(m - n) / (gamma(n + 1) * 2.0 ** (2 * n - m + 1)) + + return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) + + +@jit +def betastar(m, n): + """Beta* of eq 28, modified for eqns 31 and 32.""" + + def false_fun(m_n): + m, n = m_n + return (2 * n - m) * beta(m, n) + + return cond(jnp.logical_or(n < 0, n >= m), true_fun, false_fun, (m, n)) + + +@jit +def gamma_n(m, n): + """gamma_n of eq 33.""" + + def body_fun(i, val): + return val + 1 / i + 1 / (m + i) + + def false_fun(m_n): + m, n = m_n + return alpha(m, n) / 2 * fori_loop(1, n, body_fun, 0) + + return cond(n <= 0, true_fun, false_fun, (m, n)) + + +@jit +def gamma_nstar(m, n): + """gamma_n star of eq 33.""" + + def false_fun(m_n): + m, n = m_n + return (2 * n + m) * gamma_n(m, n) + + return cond(n <= 0, true_fun, false_fun, (m, n)) + + +@jit +def CD_m_k(R, m, k): + """Eq 31 of Dommaschk paper.""" + + def body_fun(j, val): + result = ( + val + + ( + -( + alpha(m, j) + * ( + alphastar(m, k - m - j) * jnp.log(R) + + gamma_nstar(m, k - m - j) + - alpha(m, k - m - j) + ) + - gamma_n(m, j) * alphastar(m, k - m - j) + + alpha(m, j) * betastar(m, k - j) + ) + * R ** (2 * j + m) + ) + + beta(m, j) * alphastar(m, k - j) * R ** (2 * j - m) + ) + return result + + return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def CN_m_k(R, m, k): + """Eq 32 of Dommaschk paper.""" + + def body_fun(j, val): + result = ( + val + + ( + ( + alpha(m, j) + * (alpha(m, k - m - j) * jnp.log(R) + gamma_n(m, k - m - j)) + - gamma_n(m, j) * alpha(m, k - m - j) + + alpha(m, j) * beta(m, k - j) + ) + * R ** (2 * j + m) + ) + - beta(m, j) * alpha(m, k - j) * R ** (2 * j - m) + ) + return result + + return fori_loop(0, k + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def D_m_n(R, Z, m, n): + """D_m_n term in eqn 8 of Dommaschk paper.""" + # the sum comes from fact that D_mn = I_mn and the def of I_mn in eq 2 of the paper + + 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 + + return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def N_m_n(R, Z, m, n): + """N_m_n term in eqn 9 of Dommaschk paper.""" + # the sum comes from fact that N_mn = I_mn and the def of I_mn in eq 2 of the paper + + 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 + + return fori_loop(0, n // 2 + 1, body_fun, jnp.zeros_like(R)) + + +@jit +def V_m_l(R, phi, Z, m, l, a, b, c, d): + """Eq 12 of Dommaschk paper. + + Parameters + ---------- + R,phi,Z : array-like + Cylindrical coordinates (1-D arrays of each of size num_eval_pts) + to evaluate the Dommaschk potential term at. + m : int + first index of V_m_l term + l : int + second index of V_m_l term + a : float + a_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*D_m_l + b : float + b_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*D_m_l + c : float + c_m_l coefficient of V_m_l term, which multiplies cos(m*phi)*N_m_l-1 + d : float + d_m_l coefficient of V_m_l term, which multiplies sin(m*phi)*N_m_l-1 + + Returns + ------- + value : array-like + Value of this V_m_l term evaluated at the given R,phi,Z points + (same size as the size of the given R,phi, or Z arrays). + + """ + return (a * jnp.cos(m * phi) + b * jnp.sin(m * phi)) * D_m_n(R, Z, m, l) + ( + c * jnp.cos(m * phi) + d * jnp.sin(m * phi) + ) * N_m_n(R, Z, m, l - 1) + + +@jit +def dommaschk_potential(R, phi, Z, ms, ls, a_arr, b_arr, c_arr, d_arr, B0=1): + """Eq 1 of Dommaschk paper. + + this is the total dommaschk potential for + a given set of m,l indices and their corresponding + coefficients a_ml, b_ml, c_ml d_ml. + + Parameters + ---------- + R,phi,Z : array-like + Cylindrical coordinates (1-D arrays of each of size num_eval_pts) + to evaluate the Dommaschk potential term at. + ms : 1D array-like of int + first indices of V_m_l terms + ls : 1D array-like of int + second indices of V_m_l terms + a_arr : 1D array-like of float + a_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*D_m_l + b_arr : 1D array-like of float + b_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*D_m_l + c_arr : 1D array-like of float + c_m_l coefficients of V_m_l terms, which multiplies cos(m*phi)*N_m_l-1 + d_arr : 1D array-like of float + d_m_l coefficients of V_m_l terms, which multiplies sin(m*phi)*N_m_l-1 + B0: float, toroidal magnetic field strength scale, this is the strength of the + 1/R part of the magnetic field and is the Bphi at R=1. + + Returns + ------- + value : array-like + Value of the total dommaschk potential evaluated + at the given R,phi,Z points + (same size as the size of the given R,phi, Z arrays). + """ + ms, ls, a_arr, b_arr, c_arr, d_arr = map( + jnp.atleast_1d, (ms, ls, a_arr, b_arr, c_arr, d_arr) + ) + R, phi, Z = map(jnp.atleast_1d, (R, phi, Z)) + R, phi, Z = jnp.broadcast_arrays(R, phi, Z) + ms, ls, a_arr, b_arr, c_arr, d_arr = jnp.broadcast_arrays( + ms, ls, a_arr, b_arr, c_arr, d_arr + ) + value = B0 * phi # phi term + + def body(i, val): + val += V_m_l(R, phi, Z, ms[i], ls[i], a_arr[i], b_arr[i], c_arr[i], d_arr[i]) + return val + + return fori_loop(0, len(ms), body, value) diff --git a/setup.cfg b/setup.cfg index 033cbe2a0b..20e90c9dde 100644 --- a/setup.cfg +++ b/setup.cfg @@ -71,12 +71,7 @@ ignore = D105, per-file-ignores = # need to import things to top level even if they aren't used there - desc/compute/__init__.py: F401 - desc/equilibrium/__init__.py: F401 - desc/geometry/__init__.py: F401 - desc/io/__init__.py: F401 - desc/objectives/__init__.py: F401 - desc/optimize/__init__.py: F401 + desc/*/__init__.py: F401 # too many long lines to deal with now desc/compute/data_index.py: E501 # need imports in weird order for selecting device before benchmarks From 05673fae8dfcb9755414a71020bdbcd0e660c63f Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Fri, 1 Mar 2024 14:13:46 -0500 Subject: [PATCH 41/46] Fix imports --- desc/compute/_surface.py | 22 ++++++++++++++-------- desc/compute/data_index.py | 8 ++++---- desc/objectives/_free_boundary.py | 4 ++-- tests/inputs/master_compute_data.pkl | Bin 5753002 -> 5753040 bytes tests/test_compute_funs.py | 6 +++--- tests/test_magnetic_fields.py | 3 +-- 6 files changed, 24 insertions(+), 19 deletions(-) diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 2b87cab2bf..a65cbbaa0a 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -626,7 +626,9 @@ def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.FourierCurrentPotentialField", + parameterization=[ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField" + ], ) def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi"] = ( @@ -649,7 +651,9 @@ def _Phi_FourierCurrentPotentialField(params, transforms, profiles, data, **kwar profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.FourierCurrentPotentialField", + parameterization=[ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField" + ], ) def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_t"] = ( @@ -670,7 +674,9 @@ def _Phi_t_FourierCurrentPotentialField(params, transforms, profiles, data, **kw profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.FourierCurrentPotentialField", + parameterization=[ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField" + ], ) def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_z"] = ( @@ -691,7 +697,7 @@ def _Phi_z_FourierCurrentPotentialField(params, transforms, profiles, data, **kw profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.CurrentPotentialField", + parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField", ) def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi"] = transforms["potential"]( @@ -714,7 +720,7 @@ def _Phi_CurrentPotentialField(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.CurrentPotentialField", + parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField", ) def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_t"] = transforms["potential_dtheta"]( @@ -737,7 +743,7 @@ def _Phi_t_CurrentPotentialField(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="tz", data=[], - parameterization="desc.magnetic_fields.CurrentPotentialField", + parameterization="desc.magnetic_fields._current_potential.CurrentPotentialField", ) def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs): data["Phi_z"] = transforms["potential_dzeta"]( @@ -762,8 +768,8 @@ def _Phi_z_CurrentPotentialField(params, transforms, profiles, data, **kwargs): coordinates="tz", data=["Phi_t", "Phi_z", "e_theta", "e_zeta", "|e_theta x e_zeta|"], parameterization=[ - "desc.magnetic_fields.CurrentPotentialField", - "desc.magnetic_fields.FourierCurrentPotentialField", + "desc.magnetic_fields._current_potential.CurrentPotentialField", + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField", ], ) def _K_CurrentPotentialField(params, transforms, profiles, data, **kwargs): diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index cc91e326ff..285d851d77 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -211,15 +211,15 @@ def _decorator(func): "desc.geometry.curve.FourierPlanarCurve", "desc.geometry.core.Curve", ], - "desc.magnetic_fields.CurrentPotentialField": [ + "desc.magnetic_fields._current_potential.CurrentPotentialField": [ "desc.geometry.surface.FourierRZToroidalSurface", "desc.geometry.core.Surface", - "desc.magnetic_fields.MagneticField", + "desc.magnetic_fields._core.MagneticField", ], - "desc.magnetic_fields.FourierCurrentPotentialField": [ + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField": [ "desc.geometry.surface.FourierRZToroidalSurface", "desc.geometry.core.Surface", - "desc.magnetic_fields.MagneticField", + "desc.magnetic_fields._core.MagneticField", ], "desc.coils.SplineXYZCoil": [ "desc.geometry.curve.SplineXYZCurve", diff --git a/desc/objectives/_free_boundary.py b/desc/objectives/_free_boundary.py index 313db2298c..cc4c0f4feb 100644 --- a/desc/objectives/_free_boundary.py +++ b/desc/objectives/_free_boundary.py @@ -633,14 +633,14 @@ def compute(self, eq_params, sheet_params=None, constants=None): sheet_params["R_lmn"] = eq_params["Rb_lmn"] sheet_params["Z_lmn"] = eq_params["Zb_lmn"] sheet_source_data = compute_fun( - "desc.magnetic_fields.FourierCurrentPotentialField", + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField", self._sheet_data_keys, params=sheet_params, transforms=constants["sheet_source_transforms"], profiles={}, ) sheet_eval_data = compute_fun( - "desc.magnetic_fields.FourierCurrentPotentialField", + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField", self._sheet_data_keys, params=sheet_params, transforms=constants["sheet_eval_transforms"], diff --git a/tests/inputs/master_compute_data.pkl b/tests/inputs/master_compute_data.pkl index 702deaeb7797bf396afd6fc3918bdc6a5c468462..8ff60241da5298598f26336ddca264d43106098d 100644 GIT binary patch delta 2482 zcmchXYfKbZ6vw$UvpYOQUPD0_dA}AFQF(~+5L~KO7Iv3K3>LPE;DUglkbV#v(imwK zgLZ(MJ}f0|l_0TDQg6gINSdam7Q~3)6I7%PwXv;jTD57~bML?wwK46NZt}~Sd*TPQF;t6W3o!6+* z_1~dJJ2ELb4pdh)nVK3nb0O*fX4N^mmyCVIjsE?nyig9ft8r+)I2YN`AGsB)ZkKRf zN|HSdZukX)9c-V$m3=}OHcSZ-?6>2}$B3ddrl!Qj9O&%78B#-7G7bV@= zrf}>0*R2{zRa1+`Xat0@FFZP|tQ)p+Kc2)=n5sfZZQ9iw|b(sS1(m+2TWp(eE^e_0Yc(tjQQKBOM_bm{{yLawm) z;1C+7w!4yNa+5EDu&N~4dnOX}9DWIENOrdOQgRjcX9)$hEo21LLfZ~zw09U{ba0fg93Gf7-1H1rlU>-0Z&;UNb0>BsW1N?yiAP@)wf`Jeq6bJ*t0WA;# zECdz-i-AaB2@nM=1(pHJfoLEGhy~(+6+k?&5?BQ!0Es{nkPNH_Qh-!o4WI+kfaig= zz&c<(@B-O+f5ZA=v{K%Jys>YZ(L=t$i=1@>4)`*Ck-d?1!pgI^XLHT$Ek2rNj6vbd zw|%o&a?0?tvr_`i)JfjmY~JdL4gGqNq})jNx`X;iWOlYzb+WgOs*tyo2fMj=M|@k zxG~G%jo&+>cRTr-$yvm$t{L{Il|_=MM$-b36;*Oj@l>?fJnf37>&UdHJ})$b{k>a{@t0a`yl@_afaJG_Yh$>cVj3yzsB$xkL!~SK1w0q2e3j;BpZGgX9^*GYYjuZ zQUiO>YG6+-xb0%57p@34xYADz@sKO^5b#{$72%+wUwL%!Nk3Q8uDFJpA(U*l7)+BW z7We9m8g1gX!rcXtWm~qF7giM)Sc-}ZOSkQgv=mhD+ErLyWvSR%1*UjQX(Z0iHP1)3 zrp|8QKX%nN91^~sYj(uP(u~pQ@d>dtzVV^b?d2k{Jx8G5i+GSp<^P|owO4F2pbAm3{de0De^~YG}(_a9h=U54;B!oKbmIRna`ZN=iPVi zIp1^cvmPE~ngMPPSI@L?0ut3+0xqzq?Zr+MixBQvF-;KL7=>tLRAM9>Ej~e7@k1tF zEMc{xnOP(rK?@jDqFBT&MRDQ{#=szhSjGB@JCTZY&{)JsE>KM4RpKzGL3yGlN)}(_ z6T}m23`!IS*ktfEaVf1nDZPSpY^%1vTR`ot$N0u@#?;;#7j>Qf{b}tjYQ8DIOU;>m zL)09>H>jyNeUBQ$rC~7Sl`j4^SX?l~&v2;!-4)rXO?D?{UvkK<)q_8W6uaBknR{b=^lgNnT3k`nAu zu`-;PAMq#XK20*$ORc4u8xR~v_i;I;wq5TEUiJQNa#}SphRr1k=$m6FpZ zk(Atp5tdd*Fw5PfR!nCGsjc1!b}X*4v=TQJE}BL5?qsBFJJ!i>#r4AJAd1`ltAfL5 z5==#FQh%rIU9O_omXnnEd*SF?$SrzmZaRZ3&|enRPjlGr%k>D36Ovy$m!lma9=UL= z)987Q<)sxkI)Du&3m1Tvj1iHcRCY+ElR^^NJ2bsmBu!*tQ zI>~ERMH_GAM!3PNe))!qOjObKxRoD*8zFDkAKEX+5G&J>$mOU4DyRJ z;}q4Ca{r}{m-_GW6)W`$J{;rPa;aicl$!kd0Wax(`xCjJBS8wS*8NXOA=OOZS=lxO@*S(EIt*27k^dA&Y+7^(63lq{8Z#HQlvqR8^X z%@xJlON-W*looH@RE?j)qtS}!1PELzGzCw;tefv#Hs0rdF%WfsPMQL2S!Z+tm7Y5Y z&cTvEx{U097zAmPq>{9H?jT%IKilpAC8#uZs3UNA*oUh$gzJ}~7&&|#M8~-y2+5i9m diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 7d65b6b5ac..bb9c6dd167 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1239,14 +1239,14 @@ def test_compute_everything(): **elliptic_cross_section_with_torsion ), # magnetic fields - "desc.magnetic_fields.CurrentPotentialField": CurrentPotentialField( + "desc.magnetic_fields._current_potential.CurrentPotentialField": CurrentPotentialField( # noqa:E501 **elliptic_cross_section_with_torsion, potential=lambda theta, zeta, G: G * zeta / 2 / np.pi, potential_dtheta=lambda theta, zeta, G: np.zeros_like(theta), potential_dzeta=lambda theta, zeta, G: G * np.ones_like(theta) / 2 / np.pi, params={"G": 1e7}, ), - "desc.magnetic_fields.FourierCurrentPotentialField": ( + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField": ( FourierCurrentPotentialField( **elliptic_cross_section_with_torsion, I=0, G=1e7 ) @@ -1270,7 +1270,7 @@ def test_compute_everything(): ), } assert things.keys() == data_index.keys(), ( - f"Missing the parameterizations {data_index.keys() - things.keys()}" + f"Missing the parameterization {data_index.keys() - things.keys()}" f" to test against master." ) # use this low resolution grid for equilibria to reduce file size diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index 13ffec2d4d..5611eacb99 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -12,8 +12,6 @@ from desc.grid import LinearGrid from desc.io import load from desc.magnetic_fields import ( - CD_m_k, - CN_m_k, CurrentPotentialField, DommaschkPotentialField, FourierCurrentPotentialField, @@ -25,6 +23,7 @@ field_line_integrate, read_BNORM_file, ) +from desc.magnetic_fields._dommaschk import CD_m_k, CN_m_k def phi_lm(R, phi, Z, a, m): From 6bb8f48dd5d79c40a3d55d95a91312533910070c Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 6 Mar 2024 14:14:36 -0500 Subject: [PATCH 42/46] add comments --- desc/magnetic_fields/_dommaschk.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py index 46681075c1..192d622312 100644 --- a/desc/magnetic_fields/_dommaschk.py +++ b/desc/magnetic_fields/_dommaschk.py @@ -3,8 +3,8 @@ based off Representations for vacuum potentials in stellarators https://doi.org/10.1016/0010-4655(86)90109-8 -written with naive for loops initially and can jax-ify later """ + from desc.backend import cond, fori_loop, gammaln, jit, jnp from desc.derivatives import Derivative @@ -157,6 +157,9 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify [], [], ] # indices that should be 0 due to symmetry + # 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): if not sym: @@ -184,6 +187,7 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify for i in range(4): abcd_zero_due_to_sym_inds[i] = jnp.asarray(abcd_zero_due_to_sym_inds[i]) assert (len(a_s) + len(b_s) + len(c_s) + len(d_s)) == num_modes - 1 + params = { "ms": ms, "ls": ls, From b4e3d37bda79f088ec8df0dc51d251f11dd3d2f8 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 6 Mar 2024 14:37:27 -0500 Subject: [PATCH 43/46] update changelog --- CHANGELOG.md | 11 +++++++---- desc/magnetic_fields/_dommaschk.py | 4 ++++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 56a93d27b0..b69ca99d02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,11 +15,14 @@ in the MAKEGRID format for use with other codes. least squares fit is now weighted inversely with the distance from the axis to improve the accuracy for low aspect ratio. - Adds a bounding box to the `field_line_integrate` defined by `bounds_R` and `bounds_Z` -keyword arguments, which form a hollow cylindrical bounding box. If the field line -trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying -function of the distance to the box to stop the trajectory and prevent tracing the field -line out to infinity, which is both costly and unnecessary when making a Poincare plot, +keyword arguments, which form a hollow cylindrical bounding box. If the field line +trajectory exits these bounds, the RHS will be multiplied by an exponentially decaying +function of the distance to the box to stop the trajectory and prevent tracing the field +line out to infinity, which is both costly and unnecessary when making a Poincare plot, the principle purpose of the function. +- Adds a new class ``DommaschkPotentialField`` which allows creation of magnetic fields based +off of the vacuum potentials detailed in Representations for Vacuum Potentials in Stellarators +https://doi.org/10.1016/0010-4655(86)90109-8. Speed Improvements diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py index 192d622312..e002efa7fc 100644 --- a/desc/magnetic_fields/_dommaschk.py +++ b/desc/magnetic_fields/_dommaschk.py @@ -252,6 +252,10 @@ def get_B_dom(coords, X, ms, ls): B0 = c[0] # 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] 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] From 1cc8b0ac7a33b7945273ede9a984b4601d8042a1 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 6 Mar 2024 21:17:10 -0500 Subject: [PATCH 44/46] Cast lists to jax arrays --- desc/magnetic_fields/_dommaschk.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py index e002efa7fc..7ba98960e6 100644 --- a/desc/magnetic_fields/_dommaschk.py +++ b/desc/magnetic_fields/_dommaschk.py @@ -49,12 +49,12 @@ def __init__( d_arr=jnp.array([0.0]), B0=1.0, ): - ms = jnp.atleast_1d(ms) - ls = jnp.atleast_1d(ls) - a_arr = jnp.atleast_1d(a_arr) - b_arr = jnp.atleast_1d(b_arr) - c_arr = jnp.atleast_1d(c_arr) - d_arr = jnp.atleast_1d(d_arr) + ms = jnp.atleast_1d(jnp.asarray(ms)) + ls = jnp.atleast_1d(jnp.asarray(ls)) + a_arr = jnp.atleast_1d(jnp.asarray(a_arr)) + b_arr = jnp.atleast_1d(jnp.asarray(b_arr)) + c_arr = jnp.atleast_1d(jnp.asarray(c_arr)) + d_arr = jnp.atleast_1d(jnp.asarray(d_arr)) assert ( ms.size == ls.size == a_arr.size == b_arr.size == c_arr.size == d_arr.size From acfbebefaf2d44d3d6d4470ad32be60e757c2db5 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 12 Mar 2024 15:51:40 -0400 Subject: [PATCH 45/46] add dommaschk to API --- docs/api.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/api.rst b/docs/api.rst index f88c5b0a5e..2d14391b21 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -141,6 +141,7 @@ Magnetic Fields desc.magnetic_fields.PoloidalMagneticField desc.magnetic_fields.SplineMagneticField desc.magnetic_fields.ScalarPotentialField + desc.magnetic_fields.DommaschkPotentialField desc.magnetic_fields.field_line_integrate desc.magnetic_fields.read_BNORM_file From 5830208e01448a8d6327727f06dc128dfe452154 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Wed, 13 Mar 2024 13:06:10 -0400 Subject: [PATCH 46/46] remove extra indent from docstring --- desc/magnetic_fields/_dommaschk.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/desc/magnetic_fields/_dommaschk.py b/desc/magnetic_fields/_dommaschk.py index 7ba98960e6..ac5720224a 100644 --- a/desc/magnetic_fields/_dommaschk.py +++ b/desc/magnetic_fields/_dommaschk.py @@ -85,18 +85,18 @@ def fit_magnetic_field( # noqa: C901 - FIXME - simplify Parameters ---------- - field (MagneticField or callable or ndarray): magnetic field to fit - if callable, must accept (num_nodes,3) array of rpz coords as argument - and output (num_nodes,3) as the B field in cylindrical rpz basis. - 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_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 - verbose (int): verbosity level of fitting routine, > 0 prints residuals + field (MagneticField or callable or ndarray): magnetic field to fit + if callable, must accept (num_nodes,3) array of rpz coords as argument + and output (num_nodes,3) as the B field in cylindrical rpz basis. + 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_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 + verbose (int): verbosity level of fitting routine, > 0 prints residuals """ # We seek c in Ac = b # A will be the BR, Bphi and BZ from each individual