jaxdem.analysis.engine#

JAX accumulation of pure pair kernels into analysis bins.

Precompute pair indices and bin identifiers on the host, vmap a pure kernel over pairs, and reduce into bins with jax.ops.segment_sum.

Functions

clear_jit_cache()

Clear analysis-owned compiled-function references.

evaluate_binned(kernel, arrays, binspec, *)

Run a kernel over bins and average the results in JAX.

Classes

Binned(sums, counts, mean, pairs)

Binned accumulation output.
jaxdem.analysis.engine.clear_jit_cache() → None[source]#

Clear analysis-owned compiled-function references.

This does not clear JAX’s process-wide compilation caches.

class jaxdem.analysis.engine.Binned(sums: Any, counts: Array, mean: Any, pairs: Pairs)#

Bases: object

Binned accumulation output.

Variables:

  • sums (Any) – pytree with each leaf shaped (B, …)

  • counts (jax.Array) – integer array shape (B,)

  • mean (Any) – pytree with each leaf shaped (B, …)

  • pairs (jaxdem.analysis.pairs.Pairs) – flattened pair representation used for the run (host arrays)

sums: Any#
counts: Array#
mean: Any#
pairs: Pairs#
jaxdem.analysis.engine.evaluate_binned(kernel: Any, arrays: Mapping[str, Any], binspec: BinSpec, *, kernel_kwargs: dict[str, Any] | None = None, jit: bool = True, chunk_size: int | None = None, max_pairs: int | None = None) → Binned[source]#

Run a kernel over bins and average the results in JAX.

Parameters:

  • kernel – pure function called as kernel(arrays, t0, t1, **kernel_kwargs).

  • arrays – mapping of field name -> array with leading time axis, e.g. pos: (T,N,d) or (T,S,N,d). At least one frame is required to infer kernel outputs, including for an empty pair selection.

  • binspec – bin specification built on the host. Defines which indices to use.

  • kernel_kwargs – extra keyword arguments for the kernel.

  • jit – whether to jit the core compute.

  • chunk_size – optional maximum number of pairs transferred to the device and evaluated per chunk. With the default None, one jax.vmap call processes all pairs. Set a positive integer to process pairs in chunks with repeated compiled calls. Chunking keeps peak device pair storage proportional to chunk_size. The returned pairs arrays still retain one host index triple per pair.

  • max_pairs – optional bound on the number of host pair-index triples. The default None retains all pairs without a host-memory bound. The estimate is checked before allocation and the bound is also enforced while enumerating custom bin specifications.