tod_core

Core numerical routines for sample-based TOD generation. All public functions are stateless and accept only arrays as arguments.

Core numerical routines for sample-based TOD generation.

All functions are stateless and take only arrays as arguments.

Rotation kernels (in tod_rotations.py)

_rodrigues_jit — fused double Rodrigues rotation (recenter + pol. roll),

materialises a (B, S, 3) buffer. Used by the numpy-fallback path and by tests.

_rodrigues_apply_one_jit — scalar per-(b, s) fused Rodrigues; inlined into

the production gather kernels so the (B, S, 3) intermediate is never materialised.

_rotation_params — per-sample scalars needed by the Rodrigues kernels. _recenter_and_rotate — fused recenter + pol-roll wrapper. precompute_rotation_vector_batch — Rodrigues vectors and pol. angle offsets for a batch.

Gather/accumulate kernels

_gather_accum_jit — scalar bilinear accumulation from pre-computed pixels/weights

(in tod_bilinear.py; used by tests).

_gather_accum_fused_jit — fully fused Rodrigues + bilinear gather + per-b

direct-mapped spin-2 cache + accumulation (in tod_bilinear.py).

_gather_accum_nearest_jit — fused Rodrigues + nearest-pixel gather + accumulation

(in tod_nearest.py).

HEALPix RING helpers (in numba_healpy.py)

_ring_above_jit, _ring_info_jit, _ring_z_jit, _get_interp_weights_jit, get_interp_weights_numba

tod_core.beam_tod_batch(nside, mp, data, rot_vecs, phi_b, theta_b, psis_b, interp_mode='bilinear', z_skip_threshold=-1.0)

Accumulate the TOD contribution of one beam entry for a batch of samples.

Uses Numba JIT kernels that fuse the Rodrigues rotation into the gather + accumulation step: no (B, S, 3) intermediate is materialised on the production mp_stacked path, and there is no S-tile loop — one kernel call per beam entry per batch.

Parameters:

• nside (int) – HEALPix nside of the sky map.

• mp (list[numpy.ndarray]) – Sky map components [I, Q, U]. Each element is a 1-D float32 array of length 12 * nside**2. Used only on the numpy-fallback path (when mp_stacked is not provided).

• data (dict) – Beam data entry as returned by prepare_beam_data(). Required keys: 'vec_orig', 'beam_vals', 'comp_indices'. Production path additionally requires 'mp_stacked'.

• rot_vecs (numpy.ndarray) – Rodrigues rotation vectors from precompute_rotation_vector_batch(), shape (B, 3).

• phi_b (numpy.ndarray) – Boresight longitude [rad], shape (B,).

• theta_b (numpy.ndarray) – Boresight colatitude [rad], shape (B,).

• psis_b (numpy.ndarray) – Combined rotation angle psi_b - beta [rad], shape (B,).

• interp_mode (str) –

  Sky-map interpolation strategy. One of:

  • 'bilinear' (default) — 4-pixel bilinear HEALPix interpolation with spin-2 Q/U frame correction.

  • 'nearest' — single nearest-pixel lookup; fastest, no pixel mixing.

  ('gaussian' and 'bicubic' are available on their respective branches.)

• z_skip_threshold (float) – Per-b spin-2 skip cutoff on |cos θ_pts|. Boresight samples with |bz| > z_skip_threshold apply the full Q/U frame correction; samples in the equatorial band (|bz| <= z_skip_threshold) bypass it. -1.0 (default) disables the optimisation — spin-2 is always applied, bit-identical to the un-optimised path. Pass the value returned by tod_bilinear.compute_spin2_skip_z_threshold() to enable.

Returns:

Mapping from Stokes component index to a

(B,) float32 array containing the beam-weighted sky-map accumulation for that component over the batch.

Return type:

dict[int, numpy.ndarray]

Execution paths in beam_tod_batch()

Three paths are selected automatically based on the contents of beam_data:

Path	Required cache keys	Description
Flat-sky	vec_rolled, psi_grid, dtheta, dphi, mp_stacked	Both Rodrigues rotations and vec2ang are skipped. Valid for narrow beams (≲ 5°). Fastest.
Single-Rodrigues	vec_rolled, psi_grid, mp_stacked	psi-roll is precomputed; only the recentering rotation is applied at runtime. ~2× faster than the full path.
Full double-Rodrigues	(no cache)	Both rotations applied per (B, S) element at runtime. Most general.