petitRADTRANS.radtrans_core.ck_mixing_ffi#
Registration and calling helpers for the native correlated-k mixing kernel.
The kernel lives in petitRADTRANS/cpp_src/prt_ck_mixing_ffi.cc and is
built as the petitRADTRANS.prt_ck_mixing_ffi extension module. It
replaces the former Fortran combine_ck_opacities pure_callback path:
the forward pass stays inside the XLA program as a typed FFI custom call,
runs in parallel on the XLA intra-op thread pool, and consumes the
mass-fraction-weighted opacities in their natural
(n_frequencies * n_layers, n_species, n_g) row-major layout (no
Fortran-order transpose, no host round-trip, no GIL).
Attributes#
Functions#
|
Register the c-k mixing FFI targets once per process. |
|
Mix correlated-k opacities across species with the native CPU kernel. |
Module Contents#
- petitRADTRANS.radtrans_core.ck_mixing_ffi._CK_MIX_TARGETS#
- petitRADTRANS.radtrans_core.ck_mixing_ffi._REGISTERED = False#
- petitRADTRANS.radtrans_core.ck_mixing_ffi._REGISTRATION_LOCK#
- petitRADTRANS.radtrans_core.ck_mixing_ffi.ensure_ck_mixing_ffi_registered() None#
Register the c-k mixing FFI targets once per process.
Raises an informative ImportError if the native extension is missing, since there is no fallback implementation on CPU.
- petitRADTRANS.radtrans_core.ck_mixing_ffi.combine_ck_mixing_ffi(weighted_opacities: jax.typing.ArrayLike, continuum_opacities_flat: jax.typing.ArrayLike, g_gauss: jax.typing.ArrayLike, weights_gauss: jax.typing.ArrayLike) jax.typing.ArrayLike#
Mix correlated-k opacities across species with the native CPU kernel.
Parameters#
- weighted_opacitiesjax.Array, shape (n_fl, n_species, n_g)
Mass-fraction-weighted line opacities, flattened over (frequency, layer) cells, species axis second, g axis last (row-major). The continuum must NOT be folded in; it is passed separately so the kernel can reproduce the Fortran threshold semantics exactly.
- continuum_opacities_flatjax.Array, shape (n_fl,)
Continuum opacities for each flattened cell, added to species 0 inside the kernel.
- g_gauss, weights_gaussjax.Array, shape (n_g,)
Gaussian quadrature nodes and weights of the g grid.
Returns#
- jax.Array, shape (n_fl, n_g)
Mixed opacities on the input g grid, in the input dtype.