desc.integrals.Bounce2D.batch

static Bounce2D.batch(fun, fun_data, desc_data, angle, grid, num_pitch, surf_batch_size=1, simp=True, expand_out=False)Source

Compute function fun over phase space in batches.

This is a utility method to compute some function of bounce integrals over the phase space efficiently. You may want to also JIT compile your code which calls this utility method.

Examples

  • desc/compute/_fast_ion.py::_little_gamma_c_Nemov

  • desc/compute/_neoclassical.py::_epsilon_32

Parameters:

  • fun (callable) – A function which takes a single argument fun_data and computes bounce integrals assuming fun_data holds all required quantities to construct a Bounce2D operator as well as call its methods with the flag is_fourier=True.

  • fun_data (dict[str, jnp.ndarray]) – Data to reshape, interpolate, and pass to fun. The structure of the data should match the structure returned by the registered compute functions in desc.compute. Note this dictionary will be modified.

  • desc_data (dict[str, jnp.ndarray]) – Data dictionary with the same structure as the data returned by the functions in desc.compute.

  • angle (jnp.ndarray) – Shape (num rho, X, Y). Angle returned by Bounce2D.angle.

  • grid (Grid) – Grid on which fun_data and desc_data were computed.

  • num_pitch (int) – Number of pitch angles to add to fun_data for use in the computation.

  • surf_batch_size (int) – Number of flux surfaces with which to compute simultaneously. Default is 1.

  • simp (bool) – Whether the pitch angles should be chosen for use with open Simpson rule instead of uniform weights for quadrature over velocity coordinate. Default is True.

  • expand_out (bool) – Whether to expand output to full grid so that the first dimension has size grid.num_nodes instead of grid.num_rho. Default is False.

Returns:

The output fun(fun_data).