diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b34ead0ff..5bbcd0085a 100644 --- a/CHANGELOG.md +++ b/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 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. + v0.10.4 ------- diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 61608785a5..cb855fda29 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -1380,6 +1380,9 @@ def field_line_integrate( rtol=1e-8, atol=1e-8, maxstep=1000, + bounds_R=(0, np.inf), + bounds_Z=(-np.inf, np.inf), + decay_accel=1e6, ): """Trace field lines by integration. @@ -1401,6 +1404,26 @@ def field_line_integrate( 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 to the bounding box, + 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 to the bounding box, + 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) + decay_accel : float, optional + An extra factor to the exponential that decays the RHS, i.e. + the RHS is multiplied by exp(-r * decay_accel), this is to + accelerate the decay of the RHS and stop the integration sooner + after exiting the bounds. Defaults to 1e6 + Returns ------- @@ -1410,6 +1433,7 @@ def field_line_integrate( """ r0, z0, phis = map(jnp.asarray, (r0, z0, phis)) assert r0.shape == z0.shape, "r0 and z0 must have the same shape" + assert decay_accel > 0, "decay_accel must be positive" rshape = r0.shape r0 = r0.flatten() z0 = z0.flatten() @@ -1419,12 +1443,45 @@ def field_line_integrate( 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 decay_accel to accelerate the decay so that the + # integration is stopped soon after the bounds are exited. + jnp.exp(-(decay_accel * (r - bounds_R[0]) ** 2)), + jnp.exp(-(decay_accel * (r - bounds_R[1]) ** 2)), + jnp.exp(-(decay_accel * (z - bounds_Z[0]) ** 2)), + jnp.exp(-(decay_accel * (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) diff --git a/tests/test_magnetic_fields.py b/tests/test_magnetic_fields.py index d6b668bca6..8d06ea9dd7 100644 --- a/tests/test_magnetic_fields.py +++ b/tests/test_magnetic_fields.py @@ -679,6 +679,26 @@ 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_bounds(self): + """Test field line integration with bounding box.""" + # q=4, field line should rotate 1/4 turn after 1 toroidal transit + # from outboard midplane to top center + field = ToroidalMagneticField(2, 10) + PoloidalMagneticField(2, 10, 0.25) + # test that bounds work correctly, and stop integration when trajectory + # hits the bounds + r0 = [10.1] + z0 = [0.0] + phis = [0, 2 * np.pi] + # 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): """Tests Bnormal calculation for simple toroidal field."""