torchrir 0.1.0__py3-none-any.whl

torchrir/core.py

@@ -0,0 +1,741 @@
+from __future__ import annotations
+
+"""Core RIR simulation functions (static and dynamic)."""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from .config import SimulationConfig, default_config
+from .directivity import directivity_gain, split_directivity
+from .room import MicrophoneArray, Room, Source
+from .utils import (
+    as_tensor,
+    ensure_dim,
+    estimate_beta_from_t60,
+    estimate_t60_from_beta,
+    infer_device_dtype,
+    normalize_orientation,
+    orientation_to_unit,
+    resolve_device,
+)
+
+
+def simulate_rir(
+    *,
+    room: Room,
+    sources: Source | Tensor,
+    mics: MicrophoneArray | Tensor,
+    max_order: int | None,
+    nb_img: Optional[Tensor | Tuple[int, ...]] = None,
+    nsample: Optional[int] = None,
+    tmax: Optional[float] = None,
+    tdiff: Optional[float] = None,
+    directivity: str | tuple[str, str] | None = "omni",
+    orientation: Optional[Tensor | tuple[Tensor, Tensor]] = None,
+    config: Optional[SimulationConfig] = None,
+    device: Optional[torch.device | str] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    """Simulate a static RIR using the image source method.
+
+    Args:
+        room: Room configuration (geometry, fs, reflection coefficients).
+        sources: Source positions or a Source object.
+        mics: Microphone positions or a MicrophoneArray object.
+        max_order: Maximum reflection order (uses config if None).
+        nb_img: Optional per-dimension image counts (overrides max_order).
+        nsample: Output length in samples.
+        tmax: Output length in seconds (used if nsample is None).
+        tdiff: Optional time to start diffuse tail modeling.
+        directivity: Directivity pattern(s) for source and mic (uses config if None).
+        orientation: Orientation vectors or angles.
+        config: Optional simulation configuration overrides.
+        device: Output device.
+        dtype: Output dtype.
+
+    Returns:
+        Tensor of shape (n_src, n_mic, nsample).
+    """
+    cfg = config or default_config()
+    cfg.validate()
+
+    if device is None and cfg.device is not None:
+        device = cfg.device
+
+    if max_order is None:
+        if cfg.max_order is None:
+            raise ValueError("max_order must be provided if not set in config")
+        max_order = cfg.max_order
+
+    if tmax is None and nsample is None and cfg.tmax is not None:
+        tmax = cfg.tmax
+
+    if directivity is None:
+        directivity = cfg.directivity or "omni"
+
+    if not isinstance(room, Room):
+        raise TypeError("room must be a Room instance")
+    if nsample is None and tmax is None:
+        raise ValueError("nsample or tmax must be provided")
+    if nsample is None:
+        nsample = int(math.ceil(tmax * room.fs))
+    if nsample <= 0:
+        raise ValueError("nsample must be positive")
+    if max_order < 0:
+        raise ValueError("max_order must be non-negative")
+
+    if isinstance(device, str):
+        device = resolve_device(device)
+
+    src_pos, src_ori = _prepare_entities(
+        sources, orientation, which="source", device=device, dtype=dtype
+    )
+    mic_pos, mic_ori = _prepare_entities(
+        mics, orientation, which="mic", device=device, dtype=dtype
+    )
+
+    device, dtype = infer_device_dtype(
+        src_pos, mic_pos, room.size, device=device, dtype=dtype
+    )
+    src_pos = as_tensor(src_pos, device=device, dtype=dtype)
+    mic_pos = as_tensor(mic_pos, device=device, dtype=dtype)
+
+    if src_ori is not None:
+        src_ori = as_tensor(src_ori, device=device, dtype=dtype)
+    if mic_ori is not None:
+        mic_ori = as_tensor(mic_ori, device=device, dtype=dtype)
+
+    room_size = as_tensor(room.size, device=device, dtype=dtype)
+    room_size = ensure_dim(room_size)
+    dim = room_size.numel()
+
+    if src_pos.ndim == 1:
+        src_pos = src_pos.unsqueeze(0)
+    if mic_pos.ndim == 1:
+        mic_pos = mic_pos.unsqueeze(0)
+    if src_pos.ndim != 2 or src_pos.shape[1] != dim:
+        raise ValueError("sources must be of shape (n_src, dim)")
+    if mic_pos.ndim != 2 or mic_pos.shape[1] != dim:
+        raise ValueError("mics must be of shape (n_mic, dim)")
+
+    beta = _resolve_beta(room, room_size, device=device, dtype=dtype)
+    beta = _validate_beta(beta, dim)
+
+    n_vec = _image_source_indices(max_order, dim, device=device, nb_img=nb_img)
+    refl = _reflection_coefficients(n_vec, beta)
+
+    src_pattern, mic_pattern = split_directivity(directivity)
+    mic_dir = None
+    if mic_pattern != "omni":
+        if mic_ori is None:
+            raise ValueError("mic orientation required for non-omni directivity")
+        mic_dir = orientation_to_unit(mic_ori, dim)
+
+    n_src = src_pos.shape[0]
+    n_mic = mic_pos.shape[0]
+    rir = torch.zeros((n_src, n_mic, nsample), device=device, dtype=dtype)
+    fdl = cfg.frac_delay_length
+    fdl2 = (fdl - 1) // 2
+    img_chunk = cfg.image_chunk_size
+    if img_chunk <= 0:
+        img_chunk = n_vec.shape[0]
+
+    src_dirs = None
+    if src_pattern != "omni":
+        if src_ori is None:
+            raise ValueError("source orientation required for non-omni directivity")
+        src_dirs = orientation_to_unit(src_ori, dim)
+        if src_dirs.ndim == 1:
+            src_dirs = src_dirs.unsqueeze(0).repeat(n_src, 1)
+        if src_dirs.ndim != 2 or src_dirs.shape[0] != n_src:
+            raise ValueError("source orientation must match number of sources")
+
+    for start in range(0, n_vec.shape[0], img_chunk):
+        end = min(start + img_chunk, n_vec.shape[0])
+        n_vec_chunk = n_vec[start:end]
+        refl_chunk = refl[start:end]
+        sample_chunk, attenuation_chunk = _compute_image_contributions_batch(
+            src_pos,
+            mic_pos,
+            room_size,
+            n_vec_chunk,
+            refl_chunk,
+            room,
+            fdl2,
+            src_pattern=src_pattern,
+            mic_pattern=mic_pattern,
+            src_dirs=src_dirs,
+            mic_dir=mic_dir,
+        )
+        _accumulate_rir_batch(rir, sample_chunk, attenuation_chunk, cfg)
+
+    if tdiff is not None and tmax is not None and tdiff < tmax:
+        rir = _apply_diffuse_tail(rir, room, beta, tdiff, tmax, seed=cfg.seed)
+    return rir
+
+
+def simulate_dynamic_rir(
+    *,
+    room: Room,
+    src_traj: Tensor,
+    mic_traj: Tensor,
+    max_order: int | None,
+    nsample: Optional[int] = None,
+    tmax: Optional[float] = None,
+    directivity: str | tuple[str, str] | None = "omni",
+    orientation: Optional[Tensor | tuple[Tensor, Tensor]] = None,
+    config: Optional[SimulationConfig] = None,
+    device: Optional[torch.device | str] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    """Simulate time-varying RIRs for source/mic trajectories.
+
+    Args:
+        room: Room configuration.
+        src_traj: Source trajectory (T, n_src, dim).
+        mic_traj: Microphone trajectory (T, n_mic, dim).
+        max_order: Maximum reflection order (uses config if None).
+        nsample: Output length in samples.
+        tmax: Output length in seconds (used if nsample is None).
+        directivity: Directivity pattern(s) for source and mic (uses config if None).
+        orientation: Orientation vectors or angles.
+        config: Optional simulation configuration overrides.
+        device: Output device.
+        dtype: Output dtype.
+
+    Returns:
+        Tensor of shape (T, n_src, n_mic, nsample).
+    """
+    cfg = config or default_config()
+    cfg.validate()
+
+    if device is None and cfg.device is not None:
+        device = cfg.device
+
+    if max_order is None:
+        if cfg.max_order is None:
+            raise ValueError("max_order must be provided if not set in config")
+        max_order = cfg.max_order
+
+    if tmax is None and nsample is None and cfg.tmax is not None:
+        tmax = cfg.tmax
+
+    if directivity is None:
+        directivity = cfg.directivity or "omni"
+
+    if isinstance(device, str):
+        device = resolve_device(device)
+
+    src_traj = as_tensor(src_traj, device=device, dtype=dtype)
+    mic_traj = as_tensor(mic_traj, device=device, dtype=dtype)
+
+    if src_traj.ndim == 2:
+        src_traj = src_traj.unsqueeze(1)
+    if mic_traj.ndim == 2:
+        mic_traj = mic_traj.unsqueeze(1)
+    if src_traj.ndim != 3:
+        raise ValueError("src_traj must be of shape (T, n_src, dim)")
+    if mic_traj.ndim != 3:
+        raise ValueError("mic_traj must be of shape (T, n_mic, dim)")
+    if src_traj.shape[0] != mic_traj.shape[0]:
+        raise ValueError("src_traj and mic_traj must have the same time length")
+
+    t_steps = src_traj.shape[0]
+    rirs = []
+    for t_idx in range(t_steps):
+        rir = simulate_rir(
+            room=room,
+            sources=src_traj[t_idx],
+            mics=mic_traj[t_idx],
+            max_order=max_order,
+            nsample=nsample,
+            tmax=tmax,
+            directivity=directivity,
+            orientation=orientation,
+            config=config,
+            device=device,
+            dtype=dtype,
+        )
+        rirs.append(rir)
+    return torch.stack(rirs, dim=0)
+
+
+def _prepare_entities(
+    entities: Source | MicrophoneArray | Tensor,
+    orientation: Optional[Tensor | tuple[Tensor, Tensor]],
+    *,
+    which: str,
+    device: Optional[torch.device | str],
+    dtype: Optional[torch.dtype],
+) -> Tuple[Tensor, Optional[Tensor]]:
+    """Extract positions and orientations from entities or raw tensors.
+
+    Returns:
+        Tuple of (positions, orientation).
+    """
+    if isinstance(entities, (Source, MicrophoneArray)):
+        pos = entities.positions
+        ori = entities.orientation
+    else:
+        pos = entities
+        ori = None
+    if orientation is not None:
+        if isinstance(orientation, (list, tuple)):
+            if len(orientation) != 2:
+                raise ValueError("orientation tuple must have length 2")
+            ori = orientation[0] if which == "source" else orientation[1]
+        else:
+            ori = orientation
+    pos = as_tensor(pos, device=device, dtype=dtype)
+    if ori is not None:
+        ori = as_tensor(ori, device=device, dtype=dtype)
+    return pos, ori
+
+
+def _resolve_beta(
+    room: Room, room_size: Tensor, *, device: torch.device, dtype: torch.dtype
+) -> Tensor:
+    """Resolve reflection coefficients from beta/t60/defaults.
+
+    Returns:
+        Tensor of reflection coefficients per wall.
+    """
+    if room.beta is not None:
+        return as_tensor(room.beta, device=device, dtype=dtype)
+    if room.t60 is not None:
+        return estimate_beta_from_t60(room_size, room.t60, device=device, dtype=dtype)
+    dim = room_size.numel()
+    default_faces = 4 if dim == 2 else 6
+    return torch.ones((default_faces,), device=device, dtype=dtype)
+
+
+def _validate_beta(beta: Tensor, dim: int) -> Tensor:
+    """Validate beta size against room dimension.
+
+    Returns:
+        The validated beta tensor.
+    """
+    expected = 4 if dim == 2 else 6
+    if beta.numel() != expected:
+        raise ValueError(f"beta must have {expected} elements for {dim}D")
+    return beta
+
+
+def _image_source_indices(
+    max_order: int,
+    dim: int,
+    *,
+    device: torch.device,
+    nb_img: Optional[Tensor | Tuple[int, ...]] = None,
+) -> Tensor:
+    """Generate image source index vectors up to the given order.
+
+    Returns:
+        Tensor of shape (n_images, dim).
+    """
+    if nb_img is not None:
+        nb = as_tensor(nb_img, device=device, dtype=torch.int64)
+        if nb.numel() != dim:
+            raise ValueError("nb_img must match room dimension")
+        ranges = [torch.arange(-n, n + 1, device=device, dtype=torch.int64) for n in nb]
+        grids = torch.meshgrid(*ranges, indexing="ij")
+        return torch.stack([g.reshape(-1) for g in grids], dim=-1)
+    rng = torch.arange(-max_order, max_order + 1, device=device, dtype=torch.int64)
+    grids = torch.meshgrid(*([rng] * dim), indexing="ij")
+    n_vec = torch.stack([g.reshape(-1) for g in grids], dim=-1)
+    order = torch.sum(torch.abs(n_vec), dim=-1)
+    return n_vec[order <= max_order]
+
+
+def _image_positions(src: Tensor, room_size: Tensor, n_vec: Tensor) -> Tensor:
+    """Compute image source positions for a given source.
+
+    Returns:
+        Tensor of image positions (n_images, dim).
+    """
+    n_vec_f = n_vec.to(dtype=src.dtype)
+    sign = torch.where((n_vec % 2) == 0, 1.0, -1.0).to(dtype=src.dtype)
+    n = torch.floor_divide(n_vec + 1, 2).to(dtype=src.dtype)
+    return 2.0 * room_size * n + sign * src
+
+
+def _image_positions_batch(src_pos: Tensor, room_size: Tensor, n_vec: Tensor) -> Tensor:
+    """Compute image source positions for multiple sources."""
+    sign = torch.where((n_vec % 2) == 0, 1.0, -1.0).to(dtype=src_pos.dtype)
+    n = torch.floor_divide(n_vec + 1, 2).to(dtype=src_pos.dtype)
+    base = 2.0 * room_size * n
+    return base[None, :, :] + sign[None, :, :] * src_pos[:, None, :]
+
+
+def _reflection_coefficients(n_vec: Tensor, beta: Tensor) -> Tensor:
+    """Compute reflection coefficients for each image source.
+
+    Returns:
+        Tensor of shape (n_images,) with per-image gains.
+    """
+    dim = n_vec.shape[1]
+    beta = beta.view(dim, 2)
+    beta_lo = beta[:, 0]
+    beta_hi = beta[:, 1]
+
+    n = n_vec
+    k = torch.abs(n)
+    n_hi = torch.where(n >= 0, (n + 1) // 2, k // 2)
+    n_lo = torch.where(n >= 0, n // 2, (k + 1) // 2)
+
+    n_hi = n_hi.to(dtype=beta.dtype)
+    n_lo = n_lo.to(dtype=beta.dtype)
+
+    coeff = (beta_hi**n_hi) * (beta_lo**n_lo)
+    return torch.prod(coeff, dim=1)
+
+
+def _compute_image_contributions(
+    src: Tensor,
+    mic_pos: Tensor,
+    room_size: Tensor,
+    n_vec: Tensor,
+    refl: Tensor,
+    room: Room,
+    fdl2: int,
+    *,
+    src_pattern: str,
+    mic_pattern: str,
+    src_dir: Optional[Tensor],
+    mic_dir: Optional[Tensor],
+) -> Tuple[Tensor, Tensor]:
+    """Compute sample positions and attenuation for a source and all mics."""
+    img = _image_positions(src, room_size, n_vec)
+    vec = mic_pos[:, None, :] - img[None, :, :]
+    dist = torch.linalg.norm(vec, dim=-1)
+    dist = torch.clamp(dist, min=1e-6)
+    time = dist / room.c
+    time = time + (fdl2 / room.fs)
+    sample = time * room.fs
+
+    gain = refl[None, :]
+    if src_pattern != "omni":
+        if src_dir is None:
+            raise ValueError("source orientation required for non-omni directivity")
+        cos_theta = _cos_between(vec, src_dir)
+        gain = gain * directivity_gain(src_pattern, cos_theta)
+    if mic_pattern != "omni":
+        if mic_dir is None:
+            raise ValueError("mic orientation required for non-omni directivity")
+        cos_theta = _cos_between(-vec, mic_dir)
+        gain = gain * directivity_gain(mic_pattern, cos_theta)
+
+    attenuation = gain / dist
+    return sample, attenuation
+
+
+def _compute_image_contributions_batch(
+    src_pos: Tensor,
+    mic_pos: Tensor,
+    room_size: Tensor,
+    n_vec: Tensor,
+    refl: Tensor,
+    room: Room,
+    fdl2: int,
+    *,
+    src_pattern: str,
+    mic_pattern: str,
+    src_dirs: Optional[Tensor],
+    mic_dir: Optional[Tensor],
+) -> Tuple[Tensor, Tensor]:
+    """Compute samples/attenuation for all sources/mics/images in batch."""
+    img = _image_positions_batch(src_pos, room_size, n_vec)
+    vec = mic_pos[None, :, None, :] - img[:, None, :, :]
+    dist = torch.linalg.norm(vec, dim=-1)
+    dist = torch.clamp(dist, min=1e-6)
+    time = dist / room.c
+    time = time + (fdl2 / room.fs)
+    sample = time * room.fs
+
+    gain = refl.view(1, 1, -1)
+    if src_pattern != "omni":
+        if src_dirs is None:
+            raise ValueError("source orientation required for non-omni directivity")
+        src_dirs = src_dirs[:, None, None, :]
+        cos_theta = _cos_between(vec, src_dirs)
+        gain = gain * directivity_gain(src_pattern, cos_theta)
+    if mic_pattern != "omni":
+        if mic_dir is None:
+            raise ValueError("mic orientation required for non-omni directivity")
+        mic_dir = mic_dir[None, :, None, :] if mic_dir.ndim == 2 else mic_dir.view(1, 1, 1, -1)
+        cos_theta = _cos_between(-vec, mic_dir)
+        gain = gain * directivity_gain(mic_pattern, cos_theta)
+
+    attenuation = gain / dist
+    return sample, attenuation
+
+
+def _select_orientation(orientation: Tensor, idx: int, count: int, dim: int) -> Tensor:
+    """Pick the correct orientation vector for a given entity index."""
+    if orientation.ndim == 0:
+        return orientation_to_unit(orientation, dim)
+    if orientation.ndim == 1:
+        return orientation_to_unit(orientation, dim)
+    if orientation.ndim == 2 and orientation.shape[0] == count:
+        return orientation_to_unit(orientation[idx], dim)
+    raise ValueError("orientation must be shape (dim,), (count, dim), or angles")
+
+
+def _cos_between(vec: Tensor, orientation: Tensor) -> Tensor:
+    """Compute cosine between direction vectors and orientation."""
+    orientation = normalize_orientation(orientation)
+    unit = vec / torch.linalg.norm(vec, dim=-1, keepdim=True)
+    return torch.sum(unit * orientation, dim=-1)
+
+
+def _accumulate_rir(
+    rir: Tensor, sample: Tensor, amplitude: Tensor, cfg: SimulationConfig
+) -> None:
+    """Accumulate fractional-delay contributions into the RIR tensor."""
+    idx0 = torch.floor(sample).to(torch.int64)
+    frac = sample - idx0.to(sample.dtype)
+
+    n_mic, nsample = rir.shape
+    fdl = cfg.frac_delay_length
+    lut_gran = cfg.sinc_lut_granularity
+    use_lut = cfg.use_lut and rir.device.type != "mps"
+    fdl2 = (fdl - 1) // 2
+
+    dtype = amplitude.dtype
+    n = _get_fdl_grid(fdl, device=rir.device, dtype=dtype)
+    offsets = _get_fdl_offsets(fdl, device=rir.device)
+    window = _get_fdl_window(fdl, device=rir.device, dtype=dtype)
+
+    if use_lut:
+        sinc_lut = _get_sinc_lut(fdl, lut_gran, device=rir.device, dtype=dtype)
+
+    mic_offsets = (torch.arange(n_mic, device=rir.device, dtype=torch.int64) * nsample).view(
+        n_mic, 1, 1
+    )
+    rir_flat = rir.view(-1)
+
+    chunk_size = cfg.accumulate_chunk_size
+    n_img = idx0.shape[1]
+    for start in range(0, n_img, chunk_size):
+        end = min(start + chunk_size, n_img)
+        idx = idx0[:, start:end]
+        amp = amplitude[:, start:end]
+        frac_m = frac[:, start:end]
+
+        if use_lut:
+            x_off_frac = (1.0 - frac_m) * lut_gran
+            lut_gran_off = torch.floor(x_off_frac).to(torch.int64)
+            x_off = x_off_frac - lut_gran_off.to(dtype)
+            lut_pos = lut_gran_off[..., None] + (n[None, None, :].to(torch.int64) * lut_gran)
+
+            s0 = torch.take(sinc_lut, lut_pos)
+            s1 = torch.take(sinc_lut, lut_pos + 1)
+            interp = s0 + x_off[..., None] * (s1 - s0)
+            filt = interp * window[None, None, :]
+        else:
+            t = n[None, None, :] - fdl2 - frac_m[..., None]
+            filt = torch.sinc(t) * window[None, None, :]
+
+        contrib = amp[..., None] * filt
+        target = idx[..., None] + offsets[None, None, :]
+        valid = (target >= 0) & (target < nsample)
+        if not valid.any():
+            continue
+
+        target = target + mic_offsets
+        target_flat = target[valid].to(torch.int64)
+        values_flat = contrib[valid]
+        rir_flat.scatter_add_(0, target_flat, values_flat)
+
+
+def _accumulate_rir_batch(
+    rir: Tensor, sample: Tensor, amplitude: Tensor, cfg: SimulationConfig
+) -> None:
+    """Accumulate fractional-delay contributions for all sources/mics."""
+    fn = _get_accumulate_fn(cfg, rir.device, amplitude.dtype)
+    return fn(rir, sample, amplitude)
+
+
+def _accumulate_rir_batch_impl(
+    rir: Tensor,
+    sample: Tensor,
+    amplitude: Tensor,
+    *,
+    fdl: int,
+    lut_gran: int,
+    use_lut: bool,
+    chunk_size: int,
+) -> None:
+    """Implementation for batch accumulation (optionally compiled)."""
+    idx0 = torch.floor(sample).to(torch.int64)
+    frac = sample - idx0.to(sample.dtype)
+
+    n_src, n_mic, nsample = rir.shape
+    n_sm = n_src * n_mic
+    idx0 = idx0.view(n_sm, -1)
+    frac = frac.view(n_sm, -1)
+    amplitude = amplitude.view(n_sm, -1)
+
+    fdl2 = (fdl - 1) // 2
+
+    n = _get_fdl_grid(fdl, device=rir.device, dtype=sample.dtype)
+    offsets = _get_fdl_offsets(fdl, device=rir.device)
+    window = _get_fdl_window(fdl, device=rir.device, dtype=sample.dtype)
+
+    if use_lut:
+        sinc_lut = _get_sinc_lut(fdl, lut_gran, device=rir.device, dtype=sample.dtype)
+
+    sm_offsets = (torch.arange(n_sm, device=rir.device, dtype=torch.int64) * nsample).view(
+        n_sm, 1, 1
+    )
+    rir_flat = rir.view(-1)
+
+    n_img = idx0.shape[1]
+    for start in range(0, n_img, chunk_size):
+        end = min(start + chunk_size, n_img)
+        idx = idx0[:, start:end]
+        amp = amplitude[:, start:end]
+        frac_m = frac[:, start:end]
+
+        if use_lut:
+            x_off_frac = (1.0 - frac_m) * lut_gran
+            lut_gran_off = torch.floor(x_off_frac).to(torch.int64)
+            x_off = x_off_frac - lut_gran_off.to(sample.dtype)
+            lut_pos = lut_gran_off[..., None] + (n[None, None, :].to(torch.int64) * lut_gran)
+
+            s0 = torch.take(sinc_lut, lut_pos)
+            s1 = torch.take(sinc_lut, lut_pos + 1)
+            interp = s0 + x_off[..., None] * (s1 - s0)
+            filt = interp * window[None, None, :]
+        else:
+            t = n[None, None, :] - fdl2 - frac_m[..., None]
+            filt = torch.sinc(t) * window[None, None, :]
+
+        contrib = amp[..., None] * filt
+        target = idx[..., None] + offsets[None, None, :]
+        valid = (target >= 0) & (target < nsample)
+        if not valid.any():
+            continue
+
+        target = target + sm_offsets
+        target_flat = target[valid].to(torch.int64)
+        values_flat = contrib[valid]
+        rir_flat.scatter_add_(0, target_flat, values_flat)
+
+
+_SINC_LUT_CACHE: dict[tuple[int, int, str, torch.dtype], Tensor] = {}
+_FDL_GRID_CACHE: dict[tuple[int, str, torch.dtype], Tensor] = {}
+_FDL_OFFSETS_CACHE: dict[tuple[int, str], Tensor] = {}
+_FDL_WINDOW_CACHE: dict[tuple[int, str, torch.dtype], Tensor] = {}
+_ACCUM_BATCH_COMPILED: dict[tuple[str, torch.dtype, int, int, bool, int], callable] = {}
+
+
+def _get_accumulate_fn(
+    cfg: SimulationConfig, device: torch.device, dtype: torch.dtype
+) -> callable:
+    """Return an accumulation function with config-bound constants."""
+    use_lut = cfg.use_lut and device.type != "mps"
+    fdl = cfg.frac_delay_length
+    lut_gran = cfg.sinc_lut_granularity
+    chunk_size = cfg.accumulate_chunk_size
+
+    def _fn(rir: Tensor, sample: Tensor, amplitude: Tensor) -> None:
+        _accumulate_rir_batch_impl(
+            rir,
+            sample,
+            amplitude,
+            fdl=fdl,
+            lut_gran=lut_gran,
+            use_lut=use_lut,
+            chunk_size=chunk_size,
+        )
+
+    if device.type not in ("cuda", "mps") or not cfg.use_compile:
+        return _fn
+    key = (str(device), dtype, fdl, lut_gran, use_lut, chunk_size)
+    compiled = _ACCUM_BATCH_COMPILED.get(key)
+    if compiled is None:
+        compiled = torch.compile(_fn, dynamic=True)
+        _ACCUM_BATCH_COMPILED[key] = compiled
+    return compiled
+
+
+def _get_fdl_grid(fdl: int, *, device: torch.device, dtype: torch.dtype) -> Tensor:
+    key = (fdl, str(device), dtype)
+    cached = _FDL_GRID_CACHE.get(key)
+    if cached is None:
+        cached = torch.arange(fdl, device=device, dtype=dtype)
+        _FDL_GRID_CACHE[key] = cached
+    return cached
+
+
+def _get_fdl_offsets(fdl: int, *, device: torch.device) -> Tensor:
+    key = (fdl, str(device))
+    cached = _FDL_OFFSETS_CACHE.get(key)
+    if cached is None:
+        fdl2 = (fdl - 1) // 2
+        cached = torch.arange(fdl, device=device, dtype=torch.int64) - fdl2
+        _FDL_OFFSETS_CACHE[key] = cached
+    return cached
+
+
+def _get_fdl_window(fdl: int, *, device: torch.device, dtype: torch.dtype) -> Tensor:
+    key = (fdl, str(device), dtype)
+    cached = _FDL_WINDOW_CACHE.get(key)
+    if cached is None:
+        cached = torch.hann_window(fdl, periodic=False, device=device, dtype=dtype)
+        _FDL_WINDOW_CACHE[key] = cached
+    return cached
+
+
+def _get_sinc_lut(fdl: int, lut_gran: int, *, device: torch.device, dtype: torch.dtype) -> Tensor:
+    """Create a sinc lookup table for fractional delays."""
+    key = (fdl, lut_gran, str(device), dtype)
+    cached = _SINC_LUT_CACHE.get(key)
+    if cached is not None:
+        return cached
+    fdl2 = (fdl - 1) // 2
+    lut_size = (fdl + 1) * lut_gran + 1
+    n = torch.linspace(-fdl2 - 1, fdl2 + 1, lut_size, device=device, dtype=dtype)
+    cached = torch.sinc(n)
+    _SINC_LUT_CACHE[key] = cached
+    return cached
+
+
+def _apply_diffuse_tail(
+    rir: Tensor,
+    room: Room,
+    beta: Tensor,
+    tdiff: float,
+    tmax: float,
+    *,
+    seed: Optional[int] = None,
+) -> Tensor:
+    """Apply a diffuse reverberation tail after tdiff.
+
+    Returns:
+        RIR tensor with diffuse tail applied.
+    """
+    nsample = rir.shape[-1]
+    tdiff_idx = min(nsample - 1, int(math.floor(tdiff * room.fs)))
+    if tdiff_idx <= 0:
+        return rir
+    tail_len = nsample - tdiff_idx
+    t = torch.arange(tail_len, device=rir.device, dtype=rir.dtype) / room.fs
+
+    t60 = estimate_t60_from_beta(room.size, beta)
+    if math.isinf(t60) or t60 <= 0:
+        decay = torch.exp(-t * 3.0)
+    else:
+        tau = t60 / 6.9078
+        decay = torch.exp(-t / tau)
+
+    gen = torch.Generator(device=rir.device)
+    gen.manual_seed(0 if seed is None else seed)
+    noise = torch.randn(rir[..., tdiff_idx:].shape, device=rir.device, dtype=rir.dtype, generator=gen)
+    scale = torch.linalg.norm(rir[..., tdiff_idx - 1 : tdiff_idx], dim=-1, keepdim=True) + 1e-8
+    rir[..., tdiff_idx:] = noise * decay * scale
+    return rir