torchrir 0.1.0__py3-none-any.whl

torchrir/signal.py

@@ -0,0 +1,233 @@
+from __future__ import annotations
+
+"""Signal convolution utilities.
+
+Static convolution helpers are public. Dynamic convolution is exposed via
+DynamicConvolver; internal dynamic helpers are prefixed with `_`.
+"""
+
+import math
+
+import torch
+from torch import Tensor
+
+
+def fft_convolve(signal: Tensor, rir: Tensor) -> Tensor:
+    """Convolve a 1D signal with a 1D RIR using FFT.
+
+    Args:
+        signal: 1D signal tensor.
+        rir: 1D impulse response.
+
+    Returns:
+        1D tensor of length len(signal) + len(rir) - 1.
+    """
+    if signal.ndim != 1 or rir.ndim != 1:
+        raise ValueError("fft_convolve expects 1D tensors")
+    n = signal.numel() + rir.numel() - 1
+    fft_len = 1 << (n - 1).bit_length()
+    sig_f = torch.fft.rfft(signal, n=fft_len)
+    rir_f = torch.fft.rfft(rir, n=fft_len)
+    out = torch.fft.irfft(sig_f * rir_f, n=fft_len)
+    return out[:n]
+
+
+def convolve_rir(signal: Tensor, rirs: Tensor) -> Tensor:
+    """Convolve signals with static RIRs (supports multi-source/mic).
+
+    Args:
+        signal: (n_src, n_samples) or (n_samples,) tensor.
+        rirs: (n_src, n_mic, rir_len) or compatible shape.
+
+    Returns:
+        (n_mic, n_samples + rir_len - 1) tensor or 1D for single mic.
+    """
+    signal = _ensure_signal(signal)
+    rirs = _ensure_static_rirs(rirs)
+    n_src, n_mic, rir_len = rirs.shape
+
+    if signal.shape[0] not in (1, n_src):
+        raise ValueError("signal source count does not match rirs")
+    if signal.shape[0] == 1 and n_src > 1:
+        signal = signal.expand(n_src, -1)
+
+    out_len = signal.shape[1] + rir_len - 1
+    out = torch.zeros((n_mic, out_len), dtype=signal.dtype, device=signal.device)
+
+    for s in range(n_src):
+        for m in range(n_mic):
+            out[m] += fft_convolve(signal[s], rirs[s, m])
+
+    return out.squeeze(0) if n_mic == 1 else out
+
+
+def _convolve_dynamic_rir_hop(signal: Tensor, rirs: Tensor, hop: int) -> Tensor:
+    """Dynamic convolution using fixed hop-size segments."""
+    t_steps, n_src, n_mic, rir_len = rirs.shape
+
+    frames = math.ceil(signal.shape[1] / hop)
+    frames = min(frames, t_steps)
+
+    out_len = hop * (frames - 1) + hop + rir_len - 1
+    out = torch.zeros((n_mic, out_len), dtype=signal.dtype, device=signal.device)
+
+    for t in range(frames):
+        start = t * hop
+        for s in range(n_src):
+            frame = signal[s, start : start + hop]
+            if frame.numel() == 0:
+                continue
+            for m in range(n_mic):
+                seg = fft_convolve(frame, rirs[t, s, m])
+                out[m, start : start + seg.numel()] += seg
+
+    return out.squeeze(0) if n_mic == 1 else out
+
+
+def _convolve_dynamic_rir_trajectory(
+    signal: Tensor,
+    rirs: Tensor,
+    *,
+    timestamps: Tensor | None,
+    fs: float | None,
+) -> Tensor:
+    """Dynamic convolution using variable segments like gpuRIR simulateTrajectory."""
+    n_samples = signal.shape[1]
+    t_steps, n_src, n_mic, rir_len = rirs.shape
+
+    if timestamps is not None:
+        if fs is None:
+            raise ValueError("fs must be provided when timestamps are used")
+        ts = torch.as_tensor(
+            timestamps,
+            device=signal.device,
+            dtype=torch.float32 if signal.device.type == "mps" else torch.float64,
+        )
+        if ts.ndim != 1 or ts.numel() != t_steps:
+            raise ValueError("timestamps must be 1D and match number of RIR steps")
+        if ts[0].item() != 0:
+            raise ValueError("first timestamp must be 0")
+        w_ini = (ts * fs).to(torch.long)
+    else:
+        step_fs = n_samples / t_steps
+        ts_dtype = torch.float32 if signal.device.type == "mps" else torch.float64
+        w_ini = (torch.arange(t_steps, device=signal.device, dtype=ts_dtype) * step_fs).to(
+            torch.long
+        )
+
+    w_ini = torch.cat(
+        [w_ini, torch.tensor([n_samples], device=signal.device, dtype=torch.long)]
+    )
+    w_len = w_ini[1:] - w_ini[:-1]
+
+    if signal.device.type in ("cuda", "mps"):
+        return _convolve_dynamic_rir_trajectory_batched(
+            signal, rirs, w_ini=w_ini, w_len=w_len
+        )
+
+    max_len = int(w_len.max().item())
+    segments = torch.zeros((t_steps, n_src, max_len), dtype=signal.dtype, device=signal.device)
+    for t in range(t_steps):
+        start = int(w_ini[t].item())
+        end = int(w_ini[t + 1].item())
+        if end > start:
+            segments[t, :, : end - start] = signal[:, start:end]
+
+    out = torch.zeros((n_mic, n_samples + rir_len - 1), dtype=signal.dtype, device=signal.device)
+
+    for t in range(t_steps):
+        seg_len = int(w_len[t].item())
+        if seg_len == 0:
+            continue
+        start = int(w_ini[t].item())
+        for s in range(n_src):
+            frame = segments[t, s, :seg_len]
+            for m in range(n_mic):
+                conv = fft_convolve(frame, rirs[t, s, m])
+                out[m, start : start + seg_len + rir_len - 1] += conv
+
+    return out.squeeze(0) if n_mic == 1 else out
+
+
+def _convolve_dynamic_rir_trajectory_batched(
+    signal: Tensor,
+    rirs: Tensor,
+    *,
+    w_ini: Tensor,
+    w_len: Tensor,
+    chunk_size: int = 8,
+) -> Tensor:
+    """GPU-friendly batched trajectory convolution using FFT."""
+    n_samples = signal.shape[1]
+    t_steps, n_src, n_mic, rir_len = rirs.shape
+    out = torch.zeros((n_mic, n_samples + rir_len - 1), dtype=signal.dtype, device=signal.device)
+
+    for t0 in range(0, t_steps, chunk_size):
+        t1 = min(t0 + chunk_size, t_steps)
+        lengths = w_len[t0:t1]
+        max_len = int(lengths.max().item())
+        if max_len == 0:
+            continue
+        segments = torch.zeros((t1 - t0, n_src, max_len), dtype=signal.dtype, device=signal.device)
+        for idx, t in enumerate(range(t0, t1)):
+            start = int(w_ini[t].item())
+            end = int(w_ini[t + 1].item())
+            if end > start:
+                segments[idx, :, : end - start] = signal[:, start:end]
+
+        conv_len = max_len + rir_len - 1
+        fft_len = 1 << (conv_len - 1).bit_length()
+        seg_f = torch.fft.rfft(segments, n=fft_len, dim=-1)
+        rir_f = torch.fft.rfft(rirs[t0:t1], n=fft_len, dim=-1)
+        conv_out = torch.empty(
+            (t1 - t0, n_src, n_mic, fft_len),
+            dtype=signal.dtype,
+            device=signal.device,
+        )
+        conv = torch.fft.irfft(seg_f[:, :, None, :] * rir_f, n=fft_len, dim=-1, out=conv_out)
+        conv = conv[..., :conv_len]
+        conv_sum = conv.sum(dim=1)
+
+        for idx, t in enumerate(range(t0, t1)):
+            seg_len = int(lengths[idx].item())
+            if seg_len == 0:
+                continue
+            start = int(w_ini[t].item())
+            out[:, start : start + seg_len + rir_len - 1] += conv_sum[idx, :, : seg_len + rir_len - 1]
+
+    return out.squeeze(0) if n_mic == 1 else out
+
+
+def _ensure_signal(signal: Tensor) -> Tensor:
+    """Ensure signal has shape (n_src, n_samples)."""
+    if signal.ndim == 1:
+        return signal.unsqueeze(0)
+    if signal.ndim == 2:
+        return signal
+    raise ValueError("signal must have shape (n_samples,) or (n_src, n_samples)")
+
+
+def _ensure_static_rirs(rirs: Tensor) -> Tensor:
+    """Normalize static RIR shapes to (n_src, n_mic, rir_len)."""
+    if rirs.ndim == 1:
+        return rirs.view(1, 1, -1)
+    if rirs.ndim == 2:
+        return rirs.view(1, rirs.shape[0], rirs.shape[1])
+    if rirs.ndim == 3:
+        return rirs
+    raise ValueError("rirs must have shape (rir_len,), (n_mic, rir_len), or (n_src, n_mic, rir_len)")
+
+
+def _ensure_dynamic_rirs(rirs: Tensor, signal: Tensor) -> Tensor:
+    """Normalize dynamic RIR shapes to (T, n_src, n_mic, rir_len)."""
+    if rirs.ndim == 2:
+        return rirs.view(rirs.shape[0], 1, 1, rirs.shape[1])
+    if rirs.ndim == 3:
+        if signal.ndim == 2 and rirs.shape[1] == signal.shape[0]:
+            return rirs.view(rirs.shape[0], rirs.shape[1], 1, rirs.shape[2])
+        return rirs.view(rirs.shape[0], 1, rirs.shape[1], rirs.shape[2])
+    if rirs.ndim == 4:
+        return rirs
+    raise ValueError(
+        "rirs must have shape (T, rir_len), (T, n_mic, rir_len), (T, n_src, rir_len), or (T, n_src, n_mic, rir_len)"
+    )

torchrir/simulators.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+"""Simulation strategy interfaces and implementations.
+
+Note:
+    RayTracingSimulator and FDTDSimulator are work in progress placeholders.
+"""
+
+from dataclasses import dataclass
+from typing import Protocol
+
+from .config import SimulationConfig, default_config
+from .core import simulate_dynamic_rir, simulate_rir
+from .results import RIRResult
+from .scene import Scene
+
+
+class RIRSimulator(Protocol):
+    """Strategy interface for RIR simulation backends."""
+
+    def simulate(self, scene: Scene, config: SimulationConfig | None = None) -> RIRResult:
+        """Run a simulation and return the result."""
+
+
+@dataclass(frozen=True)
+class ISMSimulator:
+    """ISM-based simulator using the current core implementation."""
+
+    def simulate(self, scene: Scene, config: SimulationConfig | None = None) -> RIRResult:
+        scene.validate()
+        cfg = config or default_config()
+        if scene.is_dynamic():
+            if scene.src_traj is None or scene.mic_traj is None:
+                raise ValueError("dynamic scene requires both src_traj and mic_traj")
+            rirs = simulate_dynamic_rir(
+                room=scene.room,
+                src_traj=scene.src_traj,
+                mic_traj=scene.mic_traj,
+                max_order=None,
+                nsample=None,
+                tmax=None,
+                directivity=None,
+                config=cfg,
+            )
+        else:
+            rirs = simulate_rir(
+                room=scene.room,
+                sources=scene.sources,
+                mics=scene.mics,
+                max_order=None,
+                nsample=None,
+                tmax=None,
+                directivity=None,
+                config=cfg,
+            )
+        return RIRResult(rirs=rirs, scene=scene, config=cfg, seed=cfg.seed)
+
+
+@dataclass(frozen=True)
+class RayTracingSimulator:
+    """Work in progress placeholder for ray tracing simulation.
+
+    Goal:
+        Provide a geometric acoustics backend that traces specular/diffuse
+        reflection paths, supports frequency-dependent absorption/scattering,
+        and returns a RIRResult compatible with the ISM path. The intent is to
+        reuse Scene/SimulationConfig for inputs and keep output shape parity.
+    """
+
+    def simulate(self, scene: Scene, config: SimulationConfig | None = None) -> RIRResult:
+        raise NotImplementedError("RayTracingSimulator is not implemented yet")
+
+
+@dataclass(frozen=True)
+class FDTDSimulator:
+    """Work in progress placeholder for FDTD simulation.
+
+    Goal:
+        Provide a wave-based solver (finite-difference time-domain) with
+        configurable grid resolution, boundary conditions, and stability
+        constraints. The solver should target CPU/GPU execution and return
+        RIRResult with the same metadata contract as ISM.
+    """
+
+    def simulate(self, scene: Scene, config: SimulationConfig | None = None) -> RIRResult:
+        raise NotImplementedError("FDTDSimulator is not implemented yet")

torchrir/utils.py

@@ -0,0 +1,281 @@
+from __future__ import annotations
+
+"""Utility functions for geometry, acoustics, and tensor handling."""
+
+import math
+import warnings
+from dataclasses import dataclass
+from typing import Iterable, Optional, Tuple
+
+import torch
+from torch import Tensor
+
+
+_DEF_SPEED_OF_SOUND = 343.0
+
+
+def as_tensor(
+    value: Tensor | Iterable[float] | float | int,
+    *,
+    device: Optional[torch.device | str] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    """Convert a value to a tensor while preserving device/dtype when possible."""
+    if isinstance(device, str):
+        device = resolve_device(device)
+    if torch.is_tensor(value):
+        out = value
+        if device is not None:
+            out = out.to(device)
+        if dtype is not None:
+            out = out.to(dtype)
+        return out
+    return torch.as_tensor(value, device=device, dtype=dtype)
+
+
+def resolve_device(
+    device: Optional[torch.device | str],
+    *,
+    prefer: Tuple[str, ...] = ("cuda", "mps", "cpu"),
+) -> torch.device:
+    """Resolve a device string (including 'auto') into a torch.device.
+
+    Falls back to CPU when the requested backend is unavailable.
+    """
+    if device is None:
+        return torch.device("cpu")
+    if isinstance(device, torch.device):
+        return device
+
+    dev = str(device).lower()
+    if dev == "auto":
+        for backend in prefer:
+            if backend == "cuda" and torch.cuda.is_available():
+                return torch.device("cuda")
+            if backend == "mps" and torch.backends.mps.is_available():
+                return torch.device("mps")
+            if backend == "cpu":
+                return torch.device("cpu")
+        return torch.device("cpu")
+
+    if dev.startswith("cuda"):
+        if torch.cuda.is_available():
+            return torch.device(device)
+        warnings.warn("CUDA not available; falling back to CPU.", RuntimeWarning)
+        return torch.device("cpu")
+    if dev == "mps":
+        if torch.backends.mps.is_available():
+            return torch.device("mps")
+        warnings.warn("MPS not available; falling back to CPU.", RuntimeWarning)
+        return torch.device("cpu")
+    if dev == "cpu":
+        return torch.device("cpu")
+
+    return torch.device(device)
+
+
+@dataclass(frozen=True)
+class DeviceSpec:
+    """Resolve device + dtype defaults consistently."""
+
+    device: Optional[torch.device | str] = None
+    dtype: Optional[torch.dtype] = None
+    prefer: Tuple[str, ...] = ("cuda", "mps", "cpu")
+
+    def resolve(self, *values) -> Tuple[torch.device, torch.dtype]:
+        """Resolve device/dtype from inputs with overrides."""
+        tensor_device: Optional[torch.device] = None
+        tensor_dtype: Optional[torch.dtype] = None
+        for value in values:
+            if torch.is_tensor(value):
+                if tensor_device is None:
+                    tensor_device = value.device
+                if tensor_dtype is None:
+                    tensor_dtype = value.dtype
+
+        if isinstance(self.device, str) and self.device.lower() == "auto":
+            device = tensor_device or resolve_device("auto", prefer=self.prefer)
+        elif self.device is None:
+            device = tensor_device or torch.device("cpu")
+        else:
+            device = resolve_device(self.device, prefer=self.prefer)
+
+        if self.dtype is None:
+            dtype = tensor_dtype or torch.float32
+        else:
+            dtype = self.dtype
+        return device, dtype
+
+
+def infer_device_dtype(
+    *values,
+    device: Optional[torch.device | str] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> Tuple[torch.device, torch.dtype]:
+    """Infer device/dtype from inputs with optional overrides."""
+    return DeviceSpec(device=device, dtype=dtype).resolve(*values)
+
+
+def ensure_dim(size: Tensor) -> Tensor:
+    """Validate room size dimensionality (2D or 3D)."""
+    if size.ndim != 1 or size.numel() not in (2, 3):
+        raise ValueError("room size must be a 1D tensor of length 2 or 3")
+    return size
+
+
+def extend_size(size: Tensor, dim: int) -> Tensor:
+    """Extend 2D room size to 3D by adding a dummy z dimension."""
+    if size.numel() == dim:
+        return size
+    if size.numel() == 2 and dim == 3:
+        pad = torch.tensor([1.0], device=size.device, dtype=size.dtype)
+        return torch.cat([size, pad])
+    raise ValueError("unsupported room dimension")
+
+
+def estimate_beta_from_t60(
+    size: Tensor,
+    t60: float,
+    *,
+    device: Optional[torch.device | str] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    """Estimate reflection coefficients from T60 using Sabine's formula."""
+    if t60 <= 0:
+        raise ValueError("t60 must be positive")
+    size = as_tensor(size, device=device, dtype=dtype)
+    size = ensure_dim(size)
+    dim = size.numel()
+    if dim == 2:
+        lx, ly = size.tolist()
+        lz = 1.0
+        volume = lx * ly * lz
+        surface = 2.0 * (lx + ly) * lz
+        alpha = 0.161 * volume / (t60 * surface)
+        alpha = max(0.0, min(alpha, 0.999))
+        beta = math.sqrt(1.0 - alpha)
+        return torch.full((4,), beta, device=size.device, dtype=size.dtype)
+    size = extend_size(size, 3)
+    lx, ly, lz = size.tolist()
+    volume = lx * ly * lz
+    surface = 2.0 * (lx * ly + ly * lz + lx * lz)
+    alpha = 0.161 * volume / (t60 * surface)
+    alpha = max(0.0, min(alpha, 0.999))
+    beta = math.sqrt(1.0 - alpha)
+    return torch.full((6,), beta, device=size.device, dtype=size.dtype)
+
+
+def estimate_t60_from_beta(
+    size: Tensor,
+    beta: Tensor,
+    *,
+    device: Optional[torch.device | str] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> float:
+    """Estimate T60 from reflection coefficients using Sabine's formula."""
+    size = as_tensor(size, device=device, dtype=dtype)
+    size = ensure_dim(size)
+    beta = as_tensor(beta, device=size.device, dtype=size.dtype)
+    dim = size.numel()
+    if dim == 2:
+        if beta.numel() != 4:
+            raise ValueError("beta must have 4 elements for 2D t60 estimation")
+        lx, ly = size.tolist()
+        lz = 1.0
+        volume = lx * ly * lz
+        surfaces = torch.tensor(
+            [ly * lz, ly * lz, lx * lz, lx * lz],
+            device=size.device,
+            dtype=size.dtype,
+        )
+        alpha = 1.0 - beta**2
+        absorption = torch.sum(surfaces * alpha).item()
+        if absorption <= 0.0:
+            return float("inf")
+        return 0.161 * volume / absorption
+    size = extend_size(size, 3)
+    if beta.numel() != 6:
+        raise ValueError("beta must have 6 elements for t60 estimation")
+    lx, ly, lz = size.tolist()
+    volume = lx * ly * lz
+    surfaces = torch.tensor(
+        [ly * lz, ly * lz, lx * lz, lx * lz, lx * ly, lx * ly],
+        device=size.device,
+        dtype=size.dtype,
+    )
+    alpha = 1.0 - beta**2
+    absorption = torch.sum(surfaces * alpha).item()
+    if absorption <= 0.0:
+        return float("inf")
+    return 0.161 * volume / absorption
+
+
+def orientation_to_unit(orientation: Tensor, dim: int) -> Tensor:
+    """Convert orientation representation to unit vectors in 2D/3D."""
+    if dim == 2:
+        if orientation.ndim == 0:
+            angle = orientation
+            vec = torch.stack([torch.cos(angle), torch.sin(angle)])
+            return normalize_orientation(vec)
+        if orientation.shape[-1] == 1:
+            angle = orientation.squeeze(-1)
+            vec = torch.stack([torch.cos(angle), torch.sin(angle)], dim=-1)
+            return normalize_orientation(vec)
+        if orientation.ndim == 1 and orientation.numel() != 2:
+            angle = orientation
+            vec = torch.stack([torch.cos(angle), torch.sin(angle)], dim=-1)
+            return normalize_orientation(vec)
+        if orientation.shape[-1] == 2:
+            return normalize_orientation(orientation)
+        raise ValueError("2D orientation must be angle or 2D vector")
+    if dim == 3:
+        if orientation.shape[-1] == 3:
+            return normalize_orientation(orientation)
+        if orientation.shape[-1] == 2:
+            azimuth = orientation[..., 0]
+            elevation = orientation[..., 1]
+            x = torch.cos(elevation) * torch.cos(azimuth)
+            y = torch.cos(elevation) * torch.sin(azimuth)
+            z = torch.sin(elevation)
+            vec = torch.stack([x, y, z], dim=-1)
+            return normalize_orientation(vec)
+        raise ValueError("3D orientation must be vector or (azimuth, elevation)")
+    raise ValueError("unsupported dimension for orientation")
+
+
+def att2t_sabine_estimation(att_db: float, t60: float) -> float:
+    """Convert attenuation (dB) to time based on T60."""
+    if t60 <= 0:
+        raise ValueError("t60 must be positive")
+    if att_db <= 0:
+        raise ValueError("att_db must be positive")
+    return (att_db / 60.0) * t60
+
+
+def att2t_SabineEstimation(att_db: float, t60: float) -> float:
+    """Legacy alias for att2t_sabine_estimation."""
+    return att2t_sabine_estimation(att_db, t60)
+
+
+def beta_SabineEstimation(room_size: Tensor, t60: float) -> Tensor:
+    """Legacy alias for estimate_beta_from_t60."""
+    return estimate_beta_from_t60(room_size, t60)
+
+
+def t2n(tmax: float, room_size: Tensor, c: float = _DEF_SPEED_OF_SOUND) -> Tensor:
+    """Estimate image counts per dimension needed to cover tmax."""
+    if tmax <= 0:
+        raise ValueError("tmax must be positive")
+    size = as_tensor(room_size)
+    size = ensure_dim(size)
+    # number of images in each dimension needed to cover the maximum distance
+    # uses the same heuristic as gpuRIR: n = ceil(tmax * c / room_size)
+    n = torch.ceil((tmax * c) / size).to(torch.int64)
+    return n
+
+
+def normalize_orientation(orientation: Tensor, *, eps: float = 1e-8) -> Tensor:
+    """Normalize orientation vectors with numerical stability."""
+    norm = torch.linalg.norm(orientation, dim=-1, keepdim=True)
+    norm = torch.clamp(norm, min=eps)
+    return orientation / norm