torchrir 0.1.0__py3-none-any.whl

torchrir/plotting.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+"""Matplotlib-based plotting helpers for room scenes."""
+
+from typing import Any, Iterable, Optional, Sequence, Tuple
+
+import torch
+from torch import Tensor
+
+from .room import MicrophoneArray, Room, Source
+from .utils import as_tensor, ensure_dim
+
+
+def plot_scene_static(
+    *,
+    room: Room | Sequence[float] | Tensor,
+    sources: Source | Tensor | Sequence,
+    mics: MicrophoneArray | Tensor | Sequence,
+    ax: Any | None = None,
+    title: Optional[str] = None,
+    show: bool = False,
+):
+    """Plot a static room with source and mic positions."""
+    plt, ax = _setup_axes(ax, room)
+
+    size = _room_size(room, ax)
+    _draw_room(ax, size)
+
+    src = _extract_positions(sources, ax)
+    mic = _extract_positions(mics, ax)
+
+    _scatter_positions(ax, src, label="sources", marker="^")
+    _scatter_positions(ax, mic, label="mics", marker="o")
+
+    if title:
+        ax.set_title(title)
+    ax.legend(loc="best")
+    if show:
+        plt.show()
+    return ax
+
+
+def plot_scene_dynamic(
+    *,
+    room: Room | Sequence[float] | Tensor,
+    src_traj: Tensor | Sequence,
+    mic_traj: Tensor | Sequence,
+    step: int = 1,
+    ax: Any | None = None,
+    title: Optional[str] = None,
+    show: bool = False,
+):
+    """Plot source and mic trajectories within a room."""
+    plt, ax = _setup_axes(ax, room)
+
+    size = _room_size(room, ax)
+    _draw_room(ax, size)
+
+    src_traj = _as_trajectory(src_traj, ax)
+    mic_traj = _as_trajectory(mic_traj, ax)
+
+    _plot_trajectories(ax, src_traj, step=step, label="source path")
+    _plot_trajectories(ax, mic_traj, step=step, label="mic path")
+
+    if title:
+        ax.set_title(title)
+    ax.legend(loc="best")
+    if show:
+        plt.show()
+    return ax
+
+
+def _setup_axes(ax: Any | None, room: Room | Sequence[float] | Tensor) -> tuple[Any, Any]:
+    """Create 2D/3D axes based on room dimension."""
+    import matplotlib.pyplot as plt
+
+    size = _room_size(room, ax)
+    dim = size.numel()
+    if ax is None:
+        if dim == 3:
+            fig = plt.figure()
+            ax = fig.add_subplot(111, projection="3d")
+        else:
+            _, ax = plt.subplots()
+    return plt, ax
+
+
+def _room_size(room: Room | Sequence[float] | Tensor, ax: Any | None) -> Tensor:
+    """Normalize room size input to a 1D tensor."""
+    if isinstance(room, Room):
+        size = room.size
+    else:
+        size = room
+    size = as_tensor(size)
+    size = ensure_dim(size)
+    return size
+
+
+def _draw_room(ax: Any, size: Tensor) -> None:
+    """Draw a 2D or 3D room outline."""
+    dim = size.numel()
+    if dim == 2:
+        _draw_room_2d(ax, size)
+    else:
+        _draw_room_3d(ax, size)
+
+
+def _draw_room_2d(ax: Any, size: Tensor) -> None:
+    """Draw a 2D rectangular room."""
+    import matplotlib.patches as patches
+
+    rect = patches.Rectangle((0.0, 0.0), size[0].item(), size[1].item(),
+                             fill=False, edgecolor="black")
+    ax.add_patch(rect)
+    ax.set_xlim(0, size[0].item())
+    ax.set_ylim(0, size[1].item())
+    ax.set_aspect("equal", adjustable="box")
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+
+
+def _draw_room_3d(ax: Any, size: Tensor) -> None:
+    """Draw a 3D box representing the room."""
+    x, y, z = size.tolist()
+    corners = torch.tensor(
+        [
+            [0, 0, 0],
+            [x, 0, 0],
+            [x, y, 0],
+            [0, y, 0],
+            [0, 0, z],
+            [x, 0, z],
+            [x, y, z],
+            [0, y, z],
+        ],
+        dtype=torch.float32,
+    )
+    edges = [
+        (0, 1),
+        (1, 2),
+        (2, 3),
+        (3, 0),
+        (4, 5),
+        (5, 6),
+        (6, 7),
+        (7, 4),
+        (0, 4),
+        (1, 5),
+        (2, 6),
+        (3, 7),
+    ]
+    for a, b in edges:
+        ax.plot(
+            [corners[a, 0], corners[b, 0]],
+            [corners[a, 1], corners[b, 1]],
+            [corners[a, 2], corners[b, 2]],
+            color="black",
+        )
+    ax.set_xlim(0, x)
+    ax.set_ylim(0, y)
+    ax.set_zlim(0, z)
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")
+    ax.set_zlabel("z")
+
+
+def _extract_positions(entity: Source | MicrophoneArray | Tensor | Sequence, ax: Any | None) -> Tensor:
+    """Extract positions from Source/MicrophoneArray or raw tensor."""
+    if isinstance(entity, (Source, MicrophoneArray)):
+        pos = entity.positions
+    else:
+        pos = entity
+    pos = as_tensor(pos)
+    if pos.ndim == 1:
+        pos = pos.unsqueeze(0)
+    return pos
+
+
+def _scatter_positions(ax: Any, positions: Tensor, *, label: str, marker: str) -> None:
+    """Scatter-plot positions in 2D or 3D."""
+    if positions.numel() == 0:
+        return
+    dim = positions.shape[1]
+    if dim == 2:
+        ax.scatter(positions[:, 0], positions[:, 1], label=label, marker=marker)
+    else:
+        ax.scatter(positions[:, 0], positions[:, 1], positions[:, 2], label=label, marker=marker)
+
+
+def _as_trajectory(traj: Tensor | Sequence, ax: Any | None) -> Tensor:
+    """Validate and normalize a trajectory tensor."""
+    traj = as_tensor(traj)
+    if traj.ndim != 3:
+        raise ValueError("trajectory must be of shape (T, N, dim)")
+    return traj
+
+
+def _plot_trajectories(ax: Any, traj: Tensor, *, step: int, label: str) -> None:
+    """Plot trajectories for each entity."""
+    if traj.numel() == 0:
+        return
+    dim = traj.shape[2]
+    if dim == 2:
+        for idx in range(traj.shape[1]):
+            xy = traj[::step, idx]
+            ax.plot(xy[:, 0], xy[:, 1], label=f"{label} {idx}")
+    else:
+        for idx in range(traj.shape[1]):
+            xyz = traj[::step, idx]
+            ax.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2], label=f"{label} {idx}")

torchrir/plotting_utils.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+"""Higher-level plotting utilities used by examples."""
+
+from pathlib import Path
+from typing import Any, Optional, Sequence
+
+import torch
+
+from .plotting import plot_scene_dynamic, plot_scene_static
+
+
+def plot_scene_and_save(
+    *,
+    out_dir: Path,
+    room: Sequence[float] | torch.Tensor,
+    sources: object | torch.Tensor | Sequence,
+    mics: object | torch.Tensor | Sequence,
+    src_traj: Optional[torch.Tensor | Sequence] = None,
+    mic_traj: Optional[torch.Tensor | Sequence] = None,
+    prefix: str = "scene",
+    step: int = 1,
+    show: bool = False,
+    plot_2d: bool = True,
+    plot_3d: bool = True,
+) -> tuple[list[Path], list[Path]]:
+    """Plot static and dynamic scenes and save images to disk.
+
+    Args:
+        out_dir: Output directory for PNGs.
+        room: Room size tensor or sequence.
+        sources: Source positions or Source-like object.
+        mics: Microphone positions or MicrophoneArray-like object.
+        src_traj: Optional source trajectory (T, n_src, dim).
+        mic_traj: Optional mic trajectory (T, n_mic, dim).
+        prefix: Filename prefix for saved images.
+        step: Subsampling step for trajectories.
+        show: Whether to show figures interactively.
+        plot_2d: Save 2D projections.
+        plot_3d: Save 3D projections (only if dim == 3).
+
+    Returns:
+        Tuple of (static_paths, dynamic_paths).
+    """
+    out_dir = Path(out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    room_size = _to_cpu(room)
+    src_pos = _positions_to_cpu(sources)
+    mic_pos = _positions_to_cpu(mics)
+    dim = int(room_size.numel())
+
+    static_paths: list[Path] = []
+    dynamic_paths: list[Path] = []
+
+    for view_dim, enabled in ((2, plot_2d), (3, plot_3d)):
+        if not enabled:
+            continue
+        if view_dim == 2 and dim < 2:
+            continue
+        if view_dim == 3 and dim < 3:
+            continue
+        view_room = room_size[:view_dim]
+        view_src = src_pos[:, :view_dim]
+        view_mic = mic_pos[:, :view_dim]
+
+        ax = plot_scene_static(
+            room=view_room,
+            sources=view_src,
+            mics=view_mic,
+            title=f"Room scene ({view_dim}D static)",
+            show=False,
+        )
+        static_path = out_dir / f"{prefix}_static_{view_dim}d.png"
+        _save_axes(ax, static_path, show=show)
+        static_paths.append(static_path)
+
+        if src_traj is not None or mic_traj is not None:
+            steps = _traj_steps(src_traj, mic_traj)
+            src_traj = _trajectory_to_cpu(src_traj, src_pos, steps)
+            mic_traj = _trajectory_to_cpu(mic_traj, mic_pos, steps)
+            view_src_traj = src_traj[:, :, :view_dim]
+            view_mic_traj = mic_traj[:, :, :view_dim]
+            ax = plot_scene_dynamic(
+                room=view_room,
+                src_traj=view_src_traj,
+                mic_traj=view_mic_traj,
+                step=step,
+                title=f"Room scene ({view_dim}D trajectories)",
+                show=False,
+            )
+            _overlay_positions(ax, view_src, view_mic)
+            dynamic_path = out_dir / f"{prefix}_dynamic_{view_dim}d.png"
+            _save_axes(ax, dynamic_path, show=show)
+            dynamic_paths.append(dynamic_path)
+
+    return static_paths, dynamic_paths
+
+
+def _to_cpu(value: Any) -> torch.Tensor:
+    """Move a value to CPU as a tensor."""
+    if torch.is_tensor(value):
+        return value.detach().cpu()
+    return torch.as_tensor(value).detach().cpu()
+
+
+def _positions_to_cpu(entity: torch.Tensor | object) -> torch.Tensor:
+    """Extract positions from an entity and move to CPU."""
+    pos = getattr(entity, "positions", entity)
+    pos = _to_cpu(pos)
+    if pos.ndim == 1:
+        pos = pos.unsqueeze(0)
+    return pos
+
+
+def _traj_steps(src_traj: Optional[torch.Tensor | Sequence], mic_traj: Optional[torch.Tensor | Sequence]) -> int:
+    """Infer the number of trajectory steps."""
+    if src_traj is not None:
+        return int(_to_cpu(src_traj).shape[0])
+    return int(_to_cpu(mic_traj).shape[0])
+
+
+def _trajectory_to_cpu(
+    traj: Optional[torch.Tensor | Sequence], fallback_pos: torch.Tensor, steps: int
+) -> torch.Tensor:
+    """Normalize trajectory to CPU tensor with shape (T, N, dim)."""
+    if traj is None:
+        return fallback_pos.unsqueeze(0).repeat(steps, 1, 1)
+    traj = _to_cpu(traj)
+    if traj.ndim != 3:
+        raise ValueError("trajectory must be of shape (T, N, dim)")
+    return traj
+
+
+def _save_axes(ax: Any, path: Path, *, show: bool) -> None:
+    """Save a matplotlib axis to disk."""
+    import matplotlib.pyplot as plt
+
+    fig = ax.figure
+    fig.tight_layout()
+    fig.savefig(path, dpi=150)
+    if show:
+        plt.show()
+    plt.close(fig)
+
+
+def _overlay_positions(ax: Any, sources: torch.Tensor, mics: torch.Tensor) -> None:
+    """Overlay static source and mic positions on an axis."""
+    if sources.numel() > 0:
+        if sources.shape[1] == 2:
+            ax.scatter(sources[:, 0], sources[:, 1], marker="^", label="sources", color="tab:green")
+        else:
+            ax.scatter(
+                sources[:, 0],
+                sources[:, 1],
+                sources[:, 2],
+                marker="^",
+                label="sources",
+                color="tab:green",
+            )
+    if mics.numel() > 0:
+        if mics.shape[1] == 2:
+            ax.scatter(mics[:, 0], mics[:, 1], marker="o", label="mics", color="tab:orange")
+        else:
+            ax.scatter(
+                mics[:, 0],
+                mics[:, 1],
+                mics[:, 2],
+                marker="o",
+                label="mics",
+                color="tab:orange",
+            )
+    ax.legend(loc="best")

torchrir/results.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+"""Result containers for simulation outputs."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from torch import Tensor
+
+from .config import SimulationConfig
+from .scene import Scene
+
+
+@dataclass(frozen=True)
+class RIRResult:
+    """Container for RIRs with metadata."""
+
+    rirs: Tensor
+    scene: Scene
+    config: SimulationConfig
+    timestamps: Optional[Tensor] = None
+    seed: Optional[int] = None

torchrir/room.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+"""Room, source, and microphone geometry models."""
+
+from dataclasses import dataclass, replace
+from typing import Optional, Sequence
+
+import torch
+from torch import Tensor
+
+from .utils import as_tensor, ensure_dim
+
+
+@dataclass(frozen=True)
+class Room:
+    """Room geometry and acoustic parameters."""
+
+    size: Tensor
+    fs: float
+    c: float = 343.0
+    beta: Optional[Tensor] = None
+    t60: Optional[float] = None
+
+    def __post_init__(self) -> None:
+        """Validate room size and reflection parameters."""
+        size = ensure_dim(self.size)
+        object.__setattr__(self, "size", size)
+        if self.beta is not None and self.t60 is not None:
+            raise ValueError("beta and t60 are mutually exclusive")
+
+    def replace(self, **kwargs) -> "Room":
+        """Return a new Room with updated fields."""
+        return replace(self, **kwargs)
+
+    @staticmethod
+    def shoebox(
+        size: Sequence[float] | Tensor,
+        *,
+        fs: float,
+        c: float = 343.0,
+        beta: Optional[Sequence[float] | Tensor] = None,
+        t60: Optional[float] = None,
+        device: Optional[torch.device | str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "Room":
+        """Create a rectangular (shoebox) room."""
+        size_t = as_tensor(size, device=device, dtype=dtype)
+        size_t = ensure_dim(size_t)
+        beta_t = None
+        if beta is not None:
+            beta_t = as_tensor(beta, device=device, dtype=dtype)
+        return Room(size=size_t, fs=fs, c=c, beta=beta_t, t60=t60)
+
+
+@dataclass(frozen=True)
+class Source:
+    """Source container with positions and optional orientation."""
+
+    positions: Tensor
+    orientation: Optional[Tensor] = None
+
+    def __post_init__(self) -> None:
+        pos = as_tensor(self.positions)
+        object.__setattr__(self, "positions", pos)
+        if self.orientation is not None:
+            ori = as_tensor(self.orientation)
+            object.__setattr__(self, "orientation", ori)
+
+    def replace(self, **kwargs) -> "Source":
+        """Return a new Source with updated fields."""
+        return replace(self, **kwargs)
+
+    @classmethod
+    def positions(
+        cls,
+        positions: Sequence[Sequence[float]] | Tensor,
+        *,
+        orientation: Optional[Sequence[float] | Tensor] = None,
+        device: Optional[torch.device | str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "Source":
+        """Construct a Source from positions."""
+        return cls.from_positions(
+            positions, orientation=orientation, device=device, dtype=dtype
+        )
+
+    @classmethod
+    def from_positions(
+        cls,
+        positions: Sequence[Sequence[float]] | Tensor,
+        *,
+        orientation: Optional[Sequence[float] | Tensor] = None,
+        device: Optional[torch.device | str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "Source":
+        """Convert positions/orientation to tensors and build a Source."""
+        pos = as_tensor(positions, device=device, dtype=dtype)
+        ori = None
+        if orientation is not None:
+            ori = as_tensor(orientation, device=device, dtype=dtype)
+        return cls(pos, ori)
+
+
+@dataclass(frozen=True)
+class MicrophoneArray:
+    """Microphone array container."""
+
+    positions: Tensor
+    orientation: Optional[Tensor] = None
+
+    def __post_init__(self) -> None:
+        pos = as_tensor(self.positions)
+        object.__setattr__(self, "positions", pos)
+        if self.orientation is not None:
+            ori = as_tensor(self.orientation)
+            object.__setattr__(self, "orientation", ori)
+
+    def replace(self, **kwargs) -> "MicrophoneArray":
+        """Return a new MicrophoneArray with updated fields."""
+        return replace(self, **kwargs)
+
+    @classmethod
+    def positions(
+        cls,
+        positions: Sequence[Sequence[float]] | Tensor,
+        *,
+        orientation: Optional[Sequence[float] | Tensor] = None,
+        device: Optional[torch.device | str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "MicrophoneArray":
+        """Construct a MicrophoneArray from positions."""
+        return cls.from_positions(
+            positions, orientation=orientation, device=device, dtype=dtype
+        )
+
+    @classmethod
+    def from_positions(
+        cls,
+        positions: Sequence[Sequence[float]] | Tensor,
+        *,
+        orientation: Optional[Sequence[float] | Tensor] = None,
+        device: Optional[torch.device | str] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "MicrophoneArray":
+        """Convert positions/orientation to tensors and build a MicrophoneArray."""
+        pos = as_tensor(positions, device=device, dtype=dtype)
+        ori = None
+        if orientation is not None:
+            ori = as_tensor(orientation, device=device, dtype=dtype)
+        return cls(pos, ori)

torchrir/scene.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+"""Scene container for simulation inputs."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+from .room import MicrophoneArray, Room, Source
+
+
+@dataclass(frozen=True)
+class Scene:
+    """Container for room, sources, microphones, and optional trajectories."""
+
+    room: Room
+    sources: Source
+    mics: MicrophoneArray
+    src_traj: Optional[Tensor] = None
+    mic_traj: Optional[Tensor] = None
+
+    def is_dynamic(self) -> bool:
+        """Return True if any trajectory is provided."""
+        return self.src_traj is not None or self.mic_traj is not None
+
+    def validate(self) -> None:
+        """Validate scene consistency and trajectory shapes."""
+        if not isinstance(self.room, Room):
+            raise TypeError("room must be a Room instance")
+        if not isinstance(self.sources, Source):
+            raise TypeError("sources must be a Source instance")
+        if not isinstance(self.mics, MicrophoneArray):
+            raise TypeError("mics must be a MicrophoneArray instance")
+
+        dim = self.room.size.numel()
+        n_src = self.sources.positions.shape[0]
+        n_mic = self.mics.positions.shape[0]
+
+        t_src = _validate_traj(self.src_traj, n_src, dim, "src_traj")
+        t_mic = _validate_traj(self.mic_traj, n_mic, dim, "mic_traj")
+        if t_src is not None and t_mic is not None and t_src != t_mic:
+            raise ValueError("src_traj and mic_traj must have matching time steps")
+
+
+def _validate_traj(
+    traj: Optional[Tensor],
+    count: int,
+    dim: int,
+    name: str,
+) -> Optional[int]:
+    if traj is None:
+        return None
+    if not torch.is_tensor(traj):
+        raise TypeError(f"{name} must be a Tensor")
+    if traj.ndim == 2:
+        if count != 1:
+            raise ValueError(f"{name} must have shape (T, {count}, {dim})")
+        if traj.shape[1] != dim:
+            raise ValueError(f"{name} must have shape (T, {dim}) for single entity")
+        return traj.shape[0]
+    if traj.ndim == 3:
+        if traj.shape[1] != count or traj.shape[2] != dim:
+            raise ValueError(f"{name} must have shape (T, {count}, {dim})")
+        return traj.shape[0]
+    raise ValueError(f"{name} must have shape (T, {count}, {dim})")

torchrir/scene_utils.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+"""Scene-agnostic utilities for example setups."""
+
+import random
+from typing import List
+
+import torch
+
+
+def sample_positions(
+    *,
+    num: int,
+    room_size: torch.Tensor,
+    rng: random.Random,
+    margin: float = 0.5,
+) -> torch.Tensor:
+    """Sample random positions within a room with a safety margin."""
+    dim = room_size.numel()
+    low = [margin] * dim
+    high = [float(room_size[i].item()) - margin for i in range(dim)]
+    coords: List[List[float]] = []
+    for _ in range(num):
+        point = [rng.uniform(low[i], high[i]) for i in range(dim)]
+        coords.append(point)
+    return torch.tensor(coords, dtype=torch.float32)
+
+
+def linear_trajectory(start: torch.Tensor, end: torch.Tensor, steps: int) -> torch.Tensor:
+    """Create a linear trajectory between start and end."""
+    return torch.stack(
+        [start + (end - start) * t / (steps - 1) for t in range(steps)],
+        dim=0,
+    )
+
+
+def binaural_mic_positions(center: torch.Tensor, offset: float = 0.08) -> torch.Tensor:
+    """Create a two-mic binaural layout around a center point."""
+    dim = center.numel()
+    offset_vec = torch.zeros((dim,), dtype=torch.float32)
+    offset_vec[0] = offset
+    left = center - offset_vec
+    right = center + offset_vec
+    return torch.stack([left, right], dim=0)
+
+
+def clamp_positions(positions: torch.Tensor, room_size: torch.Tensor, margin: float = 0.1) -> torch.Tensor:
+    """Clamp positions to remain inside the room with a margin."""
+    min_v = torch.full_like(room_size, margin)
+    max_v = room_size - margin
+    return torch.max(torch.min(positions, max_v), min_v)