Add Bounding Box to Field Line Integration

CHANGELOG.md

@@ -1,6 +1,8 @@
 Changelog
 =========
 
+- 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 `exp(-(R**2+Z**2))` to stop the trajectory and prevent tracing the field line out to infinity, which is both costly and unnecessary when making a Poincare plot, which is the principle purpose of the function.
+
 v0.10.4
 -------
 

desc/magnetic_fields.py

@@ -1380,6 +1380,8 @@
     rtol=1e-8,
     atol=1e-8,
     maxstep=1000,
+    bounds_R=(0, np.inf),
+    bounds_Z=(-np.inf, np.inf),
 ):
     """Trace field lines by integration.
 
@@ -1401,6 +1403,21 @@
         relative and absolute tolerances for ode integration
     maxstep : int
         maximum number of steps between different phis
+    bounds_R : tuple of (float,float), optional
+        R bounds for field line integration bounding box.
+        If supplied, the RHS of the field line equations will be
+        multipled by exp(-r) where r is the distance from the origin,
+        this is meant to prevent the field lines which escape to infinity from
+        slowing the integration down by being traced to infinity.
+        defaults to (0,np.inf)
+    bounds_Z : tuple of (float,float), optional
+        Z bounds for field line integration bounding box.
+        If supplied, the RHS of the field line equations will be
+        multipled by exp(-r) where r is the distance from the origin,
+        this is meant to prevent the field lines which escape to infinity from
+        slowing the integration down by being traced to infinity.
+        Defaults to (-np.inf,np.inf)
+
 
     Returns
     -------
@@ -1419,12 +1436,45 @@
     def odefun(rpz, s):
         rpz = rpz.reshape((3, -1)).T
         r = rpz[:, 0]
+        z = rpz[:, 2]
+        # if bounds are given, will decay the magnetic field line eqn
+        # RHS if the trajectory is outside of bounds to avoid
+        # integrating the field line to infinity, which is costly
+        # and not useful in most cases
+        decay_factor = jnp.where(
+            jnp.array(
+                [
+                    jnp.less(r, bounds_R[0]),
+                    jnp.greater(r, bounds_R[1]),
+                    jnp.less(z, bounds_Z[0]),
+                    jnp.greater(z, bounds_Z[1]),
+                ]
+            ),
+            jnp.array(
+                [
+                    # we mult by 1e6 to accelerate the decay so that the integration
+                    # is stopped soon after the bounds are exited.
+                    jnp.exp(-(1e6 * (r - bounds_R[0]) ** 2)),
+                    jnp.exp(-(1e6 * (r - bounds_R[1]) ** 2)),
+                    jnp.exp(-(1e6 * (z - bounds_Z[0]) ** 2)),
+                    jnp.exp(-(1e6 * (z - bounds_Z[1]) ** 2)),
+                ]
+            ),
+            1.0,
+        )
+        # multiply all together, the conditions that are not violated
+        # are just one while the violated ones are continous decaying exponentials
+        decay_factor = jnp.prod(decay_factor, axis=0)
+
         br, bp, bz = field.compute_magnetic_field(
             rpz, params, basis="rpz", source_grid=source_grid
         ).T
-        return jnp.array(
-            [r * br / bp * jnp.sign(bp), jnp.sign(bp), r * bz / bp * jnp.sign(bp)]
-        ).squeeze()
+        return (
+            decay_factor
+            * jnp.array(
+                [r * br / bp * jnp.sign(bp), jnp.sign(bp), r * bz / bp * jnp.sign(bp)]
+            ).squeeze()
+        )
 
     intfun = lambda x: odeint(odefun, x, phis, rtol=rtol, atol=atol, mxstep=maxstep)
     x = jnp.vectorize(intfun, signature="(k)->(n,k)")(x0)

tests/test_magnetic_fields.py

@@ -678,6 +678,18 @@ 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)
+        # test that bounds work correctly, and stop integration when trajectory
+        # hits the bounds
+        r0 = [10.1]
+        z0 = [0.0]
+        # this will hit the R bound
+        # (there is no Z bound, and R would go to 10.0 if not bounded)
+        r, z = field_line_integrate(r0, z0, phis, field, bounds_R=(10.05, np.inf))
+        np.testing.assert_allclose(r[-1], 10.05, rtol=3e-4)
+        # this will hit the Z bound
+        # (there is no R bound, and Z would go to 0.1 if not bounded)
+        r, z = field_line_integrate(r0, z0, phis, field, bounds_Z=(-np.inf, 0.05))
+        np.testing.assert_allclose(z[-1], 0.05, atol=3e-3)
 
     @pytest.mark.unit
     def test_Bnormal_calculation(self):