Merge remote-tracking branch 'origin/master' into yge/deps

Author: YigitElma

CHANGELOG.md

@@ -20,6 +20,7 @@ Bug Fixes
 - Fixes a bug in `OmnigenousField.change_resolution` when changing `L_B`.
 - Scaling a `ScaledProfile` or taking power of a `PowerProfile` now only updates the `scale`/`power` attributes instead of nesting the `ScaledProfile`/`PowerProfile`s.
 - `jax.Array`s in `_static_attrs` will be automatically converted to `np.ndarray` to prevent stalling code. In general, jax arrays should be omitted in `_static_attrs`.
+- Fixes a bug in `desc.magnetic_fields.field_integrate` when calling with an integer `bs_chunk_size`.
 
 Performance Improvements
 

desc/magnetic_fields/_core.py

@@ -19,7 +19,7 @@
 from netCDF4 import Dataset, chartostring, stringtochar
 from scipy.constants import mu_0
 
-from desc.backend import jit, jnp, sign
+from desc.backend import jnp, sign
 from desc.basis import (
     ChebyshevDoubleFourierBasis,
     ChebyshevPolynomial,
@@ -2816,8 +2816,8 @@ def _intfun_wrapper(
     )
 
 
-@jit
 def _odefun(s, rpz, args):
+    # Note: this function is already jitted by diffeqsolve
     field, params, scale, source_grid, bs_chunk_size = args
     r = rpz[0]
     br, bp, bz = (

tests/test_magnetic_fields.py

@@ -15,6 +15,7 @@
 
 from desc.backend import jax, jit, jnp
 from desc.basis import DoubleFourierSeries
+from desc.coils import CoilSet, FourierPlanarCoil
 from desc.compute.utils import get_params, get_transforms
 from desc.examples import get
 from desc.geometry import FourierRZToroidalSurface, FourierXYZCurve
@@ -1212,6 +1213,31 @@ 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_field_line_integrate_coil_bs_chunk(self):
+        """Test field line integration for coils with bs_chunk_size."""
+        # Related to issue #2214
+        # simple toroidal field
+        B0 = 1
+        R0 = 4
+        I = 2 * np.pi * B0 * R0 / mu_0
+        field = CoilSet(
+            FourierPlanarCoil(
+                current=I,
+                center=[R0, 0, 0],
+                normal=[0, 1, 0],
+                r_n=R0 / 4,
+            ),
+            NFP=40,
+            check_intersection=False,
+        )
+        r0 = [10.0]
+        z0 = [0.0]
+        phis = [0, np.pi]
+        r, z = field_line_integrate(r0, z0, phis, field, bs_chunk_size=1000)
+        np.testing.assert_allclose(r[-1], r0[0], rtol=1e-6, atol=1e-6)
+        np.testing.assert_allclose(z[-1], z0[0], rtol=1e-6, atol=1e-6)
+
     @pytest.mark.unit
     def test_field_line_integrate_jax_transforms(self, capsys):
         """Test field line integration is JAX transformable."""