torchrir 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

torchrir/__init__.py

@@ -4,6 +4,8 @@ from .config import SimulationConfig, default_config
 from .core import simulate_dynamic_rir, simulate_rir
 from .dynamic import DynamicConvolver
 from .logging_utils import LoggingConfig, get_logger, setup_logging
+from .animation import animate_scene_gif
+from .metadata import build_metadata, save_metadata_json
 from .plotting import plot_scene_dynamic, plot_scene_static
 from .plotting_utils import plot_scene_and_save
 from .room import MicrophoneArray, Room, Source
@@ -61,6 +63,8 @@ __all__ = [
     "get_logger",
     "list_cmu_arctic_speakers",
     "LoggingConfig",
+    "animate_scene_gif",
+    "build_metadata",
     "resolve_device",
     "SentenceLike",
     "load_dataset_sources",
@@ -75,6 +79,7 @@ __all__ = [
     "plot_scene_and_save",
     "plot_scene_static",
     "save_wav",
+    "save_metadata_json",
     "Scene",
     "setup_logging",
     "SimulationConfig",

torchrir/animation.py

@@ -0,0 +1,172 @@
+from __future__ import annotations
+
+"""Animation helpers for dynamic scenes."""
+
+from pathlib import Path
+from typing import Optional, Sequence
+
+import torch
+
+from .plotting_utils import _positions_to_cpu, _to_cpu, _traj_steps, _trajectory_to_cpu
+
+
+def animate_scene_gif(
+    *,
+    out_path: 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,
+    step: int = 1,
+    fps: Optional[float] = None,
+    signal_len: Optional[int] = None,
+    fs: Optional[float] = None,
+    duration_s: Optional[float] = None,
+    plot_2d: bool = True,
+    plot_3d: bool = False,
+) -> Path:
+    """Render a GIF showing source/mic trajectories.
+
+    Args:
+        out_path: Destination GIF path.
+        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).
+        step: Subsampling step for trajectories.
+        fps: Frames per second for the GIF (auto if None).
+        signal_len: Optional signal length (samples) to infer elapsed time.
+        fs: Sample rate used with signal_len.
+        duration_s: Optional total duration in seconds (overrides signal_len/fs).
+        plot_2d: Use 2D projection if True.
+        plot_3d: Use 3D projection if True and dim == 3.
+
+    Returns:
+        The output path.
+
+    Example:
+        >>> animate_scene_gif(
+        ...     out_path=Path("outputs/scene.gif"),
+        ...     room=[6.0, 4.0, 3.0],
+        ...     sources=[[1.0, 2.0, 1.5]],
+        ...     mics=[[2.0, 2.0, 1.5]],
+        ...     src_traj=src_traj,
+        ...     mic_traj=mic_traj,
+        ...     signal_len=16000,
+        ...     fs=16000,
+        ... )
+    """
+    import matplotlib.pyplot as plt
+    from matplotlib import animation
+
+    out_path = Path(out_path)
+    out_path.parent.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())
+    view_dim = 3 if (plot_3d and dim == 3) else 2
+    view_room = room_size[:view_dim]
+    view_src = src_pos[:, :view_dim]
+    view_mic = mic_pos[:, :view_dim]
+
+    if src_traj is None and mic_traj is None:
+        raise ValueError("at least one trajectory is required for animation")
+    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]
+
+    if view_dim == 3:
+        fig = plt.figure()
+        ax = fig.add_subplot(111, projection="3d")
+        ax.set_xlim(0, view_room[0].item())
+        ax.set_ylim(0, view_room[1].item())
+        ax.set_zlim(0, view_room[2].item())
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+        ax.set_zlabel("z")
+    else:
+        fig, ax = plt.subplots()
+        ax.set_xlim(0, view_room[0].item())
+        ax.set_ylim(0, view_room[1].item())
+        ax.set_aspect("equal", adjustable="box")
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+
+    src_scatter = ax.scatter([], [], marker="^", color="tab:green", label="sources")
+    mic_scatter = ax.scatter([], [], marker="o", color="tab:orange", label="mics")
+    src_lines = []
+    mic_lines = []
+    for _ in range(view_src_traj.shape[1]):
+        if view_dim == 2:
+            line, = ax.plot([], [], color="tab:green", alpha=0.6)
+        else:
+            line, = ax.plot([], [], [], color="tab:green", alpha=0.6)
+        src_lines.append(line)
+    for _ in range(view_mic_traj.shape[1]):
+        if view_dim == 2:
+            line, = ax.plot([], [], color="tab:orange", alpha=0.6)
+        else:
+            line, = ax.plot([], [], [], color="tab:orange", alpha=0.6)
+        mic_lines.append(line)
+
+    ax.legend(loc="best")
+
+    if duration_s is None and signal_len is not None and fs is not None:
+        duration_s = float(signal_len) / float(fs)
+
+    def _frame(i: int):
+        idx = min(i * step, view_src_traj.shape[0] - 1)
+        src_frame = view_src_traj[: idx + 1]
+        mic_frame = view_mic_traj[: idx + 1]
+        src_pos_frame = view_src_traj[idx]
+        mic_pos_frame = view_mic_traj[idx]
+
+        if view_dim == 2:
+            src_scatter.set_offsets(src_pos_frame)
+            mic_scatter.set_offsets(mic_pos_frame)
+            for s_idx, line in enumerate(src_lines):
+                xy = src_frame[:, s_idx, :]
+                line.set_data(xy[:, 0], xy[:, 1])
+            for m_idx, line in enumerate(mic_lines):
+                xy = mic_frame[:, m_idx, :]
+                line.set_data(xy[:, 0], xy[:, 1])
+        else:
+            src_scatter._offsets3d = (
+                src_pos_frame[:, 0],
+                src_pos_frame[:, 1],
+                src_pos_frame[:, 2],
+            )
+            mic_scatter._offsets3d = (
+                mic_pos_frame[:, 0],
+                mic_pos_frame[:, 1],
+                mic_pos_frame[:, 2],
+            )
+            for s_idx, line in enumerate(src_lines):
+                xyz = src_frame[:, s_idx, :]
+                line.set_data(xyz[:, 0], xyz[:, 1])
+                line.set_3d_properties(xyz[:, 2])
+            for m_idx, line in enumerate(mic_lines):
+                xyz = mic_frame[:, m_idx, :]
+                line.set_data(xyz[:, 0], xyz[:, 1])
+                line.set_3d_properties(xyz[:, 2])
+        if duration_s is not None and steps > 1:
+            t = (idx / (steps - 1)) * duration_s
+            ax.set_title(f"t = {t:.2f} s")
+        return [src_scatter, mic_scatter, *src_lines, *mic_lines]
+
+    frames = max(1, (view_src_traj.shape[0] + step - 1) // step)
+    if fps is None or fps <= 0:
+        if duration_s is not None and duration_s > 0:
+            fps = frames / duration_s
+        else:
+            fps = 6.0
+    anim = animation.FuncAnimation(fig, _frame, frames=frames, interval=1000 / fps, blit=False)
+    anim.save(out_path, writer="pillow", fps=fps)
+    plt.close(fig)
+    return out_path

torchrir/config.py

@@ -10,7 +10,12 @@ import torch
 
 @dataclass(frozen=True)
 class SimulationConfig:
-    """Configuration values for RIR simulation and convolution."""
+    """Configuration values for RIR simulation and convolution.
+
+    Example:
+        >>> cfg = SimulationConfig(max_order=6, tmax=0.3, device="auto")
+        >>> cfg.validate()
+    """
 
     fs: Optional[float] = None
     max_order: Optional[int] = None
@@ -53,7 +58,11 @@ class SimulationConfig:
 
 
 def default_config() -> SimulationConfig:
-    """Return the default simulation configuration."""
+    """Return the default simulation configuration.
+
+    Example:
+        >>> cfg = default_config()
+    """
     cfg = SimulationConfig()
     cfg.validate()
     return cfg

torchrir/core.py

@@ -58,6 +58,18 @@ def simulate_rir(
 
     Returns:
         Tensor of shape (n_src, n_mic, nsample).
+
+    Example:
+        >>> room = Room.shoebox(size=[6.0, 4.0, 3.0], fs=16000, beta=[0.9] * 6)
+        >>> sources = Source.positions([[1.0, 2.0, 1.5]])
+        >>> mics = MicrophoneArray.positions([[2.0, 2.0, 1.5]])
+        >>> rir = simulate_rir(
+        ...     room=room,
+        ...     sources=sources,
+        ...     mics=mics,
+        ...     max_order=6,
+        ...     tmax=0.3,
+        ... )
     """
     cfg = config or default_config()
     cfg.validate()
@@ -208,6 +220,24 @@ def simulate_dynamic_rir(
 
     Returns:
         Tensor of shape (T, n_src, n_mic, nsample).
+
+    Example:
+        >>> room = Room.shoebox(size=[6.0, 4.0, 3.0], fs=16000, beta=[0.9] * 6)
+        >>> from torchrir import linear_trajectory
+        >>> src_traj = torch.stack(
+        ...     [linear_trajectory(torch.tensor([1.0, 2.0, 1.5]),
+        ...                        torch.tensor([4.0, 2.0, 1.5]), 8)],
+        ...     dim=1,
+        ... )
+        >>> mic_pos = torch.tensor([[2.0, 2.0, 1.5]])
+        >>> mic_traj = mic_pos.unsqueeze(0).repeat(8, 1, 1)
+        >>> rirs = simulate_dynamic_rir(
+        ...     room=room,
+        ...     src_traj=src_traj,
+        ...     mic_traj=mic_traj,
+        ...     max_order=4,
+        ...     tmax=0.3,
+        ... )
     """
     cfg = config or default_config()
     cfg.validate()

torchrir/datasets/cmu_arctic.py

@@ -49,6 +49,13 @@ class CmuArcticSentence:
 
 
 class CmuArcticDataset:
+    """CMU ARCTIC dataset loader.
+
+    Example:
+        >>> dataset = CmuArcticDataset(Path("datasets/cmu_arctic"), speaker="bdl", download=True)
+        >>> audio, fs = dataset.load_wav("arctic_a0001")
+    """
+
     def __init__(self, root: Path, speaker: str = "bdl", download: bool = False) -> None:
         """Initialize a CMU ARCTIC dataset handle.
 
@@ -182,7 +189,11 @@ def _parse_text_line(line: str) -> Tuple[str, str]:
 
 
 def load_wav_mono(path: Path) -> Tuple[torch.Tensor, int]:
-    """Load a wav file and return mono audio and sample rate."""
+    """Load a wav file and return mono audio and sample rate.
+
+    Example:
+        >>> audio, fs = load_wav_mono(Path("datasets/cmu_arctic/ARCTIC/.../wav/arctic_a0001.wav"))
+    """
     import soundfile as sf
 
     audio, sample_rate = sf.read(str(path), dtype="float32", always_2d=True)
@@ -195,7 +206,11 @@ def load_wav_mono(path: Path) -> Tuple[torch.Tensor, int]:
 
 
 def save_wav(path: Path, audio: torch.Tensor, sample_rate: int) -> None:
-    """Save a mono or multi-channel wav to disk."""
+    """Save a mono or multi-channel wav to disk.
+
+    Example:
+        >>> save_wav(Path("outputs/example.wav"), audio, sample_rate)
+    """
     import soundfile as sf
 
     audio = audio.detach().cpu().clamp(-1.0, 1.0).to(torch.float32)

torchrir/datasets/utils.py

@@ -11,7 +11,12 @@ from .base import BaseDataset, SentenceLike
 
 
 def choose_speakers(dataset: BaseDataset, num_sources: int, rng: random.Random) -> List[str]:
-    """Select unique speakers for the requested number of sources."""
+    """Select unique speakers for the requested number of sources.
+
+    Example:
+        >>> rng = random.Random(0)
+        >>> speakers = choose_speakers(dataset, num_sources=2, rng=rng)
+    """
     speakers = dataset.list_speakers()
     if not speakers:
         raise RuntimeError("no speakers available")
@@ -27,7 +32,20 @@ def load_dataset_sources(
     duration_s: float,
     rng: random.Random,
 ) -> Tuple[torch.Tensor, int, List[Tuple[str, List[str]]]]:
-    """Load and concatenate utterances for each speaker into fixed-length signals."""
+    """Load and concatenate utterances for each speaker into fixed-length signals.
+
+    Example:
+        >>> from pathlib import Path
+        >>> from torchrir import CmuArcticDataset
+        >>> rng = random.Random(0)
+        >>> root = Path("datasets/cmu_arctic")
+        >>> signals, fs, info = load_dataset_sources(
+        ...     dataset_factory=lambda spk: CmuArcticDataset(root, speaker=spk, download=True),
+        ...     num_sources=2,
+        ...     duration_s=10.0,
+        ...     rng=rng,
+        ... )
+    """
     dataset0 = dataset_factory(None)
     speakers = choose_speakers(dataset0, num_sources, rng)
     signals: List[torch.Tensor] = []

torchrir/dynamic.py

@@ -17,7 +17,12 @@ from .signal import _ensure_dynamic_rirs, _ensure_signal
 
 @dataclass(frozen=True)
 class DynamicConvolver:
-    """Convolver for time-varying RIRs."""
+    """Convolver for time-varying RIRs.
+
+    Example:
+        >>> convolver = DynamicConvolver(mode="trajectory")
+        >>> y = convolver.convolve(signal, rirs)
+    """
 
     mode: str = "trajectory"
     hop: Optional[int] = None
@@ -28,7 +33,11 @@ class DynamicConvolver:
         return self.convolve(signal, rirs)
 
     def convolve(self, signal: Tensor, rirs: Tensor) -> Tensor:
-        """Convolve signals with time-varying RIRs."""
+        """Convolve signals with time-varying RIRs.
+
+        Example:
+            >>> y = DynamicConvolver(mode="hop", hop=1024).convolve(signal, rirs)
+        """
         if self.mode not in ("trajectory", "hop"):
             raise ValueError("mode must be 'trajectory' or 'hop'")
         if self.mode == "hop":

torchrir/logging_utils.py

@@ -9,7 +9,12 @@ from typing import Optional
 
 @dataclass(frozen=True)
 class LoggingConfig:
-    """Configuration for torchrir logging."""
+    """Configuration for torchrir logging.
+
+    Example:
+        >>> config = LoggingConfig(level="INFO")
+        >>> logger = setup_logging(config)
+    """
 
     level: str | int = "INFO"
     format: str = "%(levelname)s:%(name)s:%(message)s"
@@ -33,7 +38,12 @@ class LoggingConfig:
 
 
 def setup_logging(config: LoggingConfig, *, name: str = "torchrir") -> logging.Logger:
-    """Configure and return the base torchrir logger."""
+    """Configure and return the base torchrir logger.
+
+    Example:
+        >>> logger = setup_logging(LoggingConfig(level="DEBUG"))
+        >>> logger.info("ready")
+    """
     logger = logging.getLogger(name)
     level = config.resolve_level()
     logger.setLevel(level)
@@ -47,7 +57,11 @@ def setup_logging(config: LoggingConfig, *, name: str = "torchrir") -> logging.L
 
 
 def get_logger(name: Optional[str] = None) -> logging.Logger:
-    """Return a torchrir logger, namespaced under the torchrir root."""
+    """Return a torchrir logger, namespaced under the torchrir root.
+
+    Example:
+        >>> logger = get_logger("examples.static")
+    """
     if not name:
         return logging.getLogger("torchrir")
     if name.startswith("torchrir"):

torchrir/metadata.py

@@ -0,0 +1,216 @@
+from __future__ import annotations
+
+"""Metadata helpers for simulation outputs."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import json
+import torch
+from torch import Tensor
+
+from .room import MicrophoneArray, Room, Source
+
+
+@dataclass(frozen=True)
+class ArrayAttributes:
+    """Structured description of a microphone array."""
+
+    geometry_name: str
+    positions: Tensor
+    orientation: Optional[Tensor]
+    center: Tensor
+    normal: Optional[Tensor]
+    spacing: Optional[float]
+
+
+def build_metadata(
+    *,
+    room: Room,
+    sources: Source,
+    mics: MicrophoneArray,
+    rirs: Tensor,
+    src_traj: Optional[Tensor] = None,
+    mic_traj: Optional[Tensor] = None,
+    timestamps: Optional[Tensor] = None,
+    signal_len: Optional[int] = None,
+    source_info: Optional[Any] = None,
+    extra: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Build JSON-serializable metadata for a simulation output.
+
+    Example:
+        >>> metadata = build_metadata(
+        ...     room=room,
+        ...     sources=sources,
+        ...     mics=mics,
+        ...     rirs=rirs,
+        ...     src_traj=src_traj,
+        ...     mic_traj=mic_traj,
+        ...     signal_len=signal.shape[-1],
+        ... )
+        >>> save_metadata_json(Path(\"outputs/scene_metadata.json\"), metadata)
+    """
+    nsample = int(rirs.shape[-1])
+    fs = float(room.fs)
+    time_axis = {
+        "fs": fs,
+        "nsample": nsample,
+        "t": _to_serializable(torch.arange(nsample, dtype=torch.float32) / fs),
+    }
+
+    src_pos = sources.positions
+    mic_pos = mics.positions
+    dim = int(room.size.numel())
+    src_traj_n = _normalize_traj(src_traj, src_pos, dim, "src_traj")
+    mic_traj_n = _normalize_traj(mic_traj, mic_pos, dim, "mic_traj")
+
+    t_steps = max(src_traj_n.shape[0], mic_traj_n.shape[0])
+    if src_traj_n.shape[0] == 1 and t_steps > 1:
+        src_traj_n = src_traj_n.expand(t_steps, -1, -1)
+    if mic_traj_n.shape[0] == 1 and t_steps > 1:
+        mic_traj_n = mic_traj_n.expand(t_steps, -1, -1)
+    if src_traj_n.shape[0] != mic_traj_n.shape[0]:
+        raise ValueError("src_traj and mic_traj must have matching time steps")
+
+    azimuth, elevation = _compute_doa(src_traj_n, mic_traj_n)
+    doa = {
+        "frame": "world",
+        "unit": "radians",
+        "azimuth": _to_serializable(azimuth),
+        "elevation": _to_serializable(elevation),
+    }
+
+    timestamps_out: Optional[Tensor] = None
+    if timestamps is not None:
+        timestamps_out = timestamps
+    elif t_steps > 1 and signal_len is not None:
+        duration = max(0.0, (float(signal_len) - 1.0) / fs)
+        timestamps_out = torch.linspace(0.0, duration, t_steps, dtype=torch.float32)
+
+    array_attrs = _array_attributes(mics)
+
+    metadata: Dict[str, Any] = {
+        "room": {
+            "size": _to_serializable(room.size),
+            "c": float(room.c),
+            "beta": _to_serializable(room.beta) if room.beta is not None else None,
+            "t60": float(room.t60) if room.t60 is not None else None,
+            "fs": fs,
+        },
+        "sources": {
+            "positions": _to_serializable(src_pos),
+            "orientation": _to_serializable(sources.orientation),
+        },
+        "mics": {
+            "positions": _to_serializable(mic_pos),
+            "orientation": _to_serializable(mics.orientation),
+        },
+        "trajectories": {
+            "sources": _to_serializable(src_traj_n if t_steps > 1 else None),
+            "mics": _to_serializable(mic_traj_n if t_steps > 1 else None),
+        },
+        "array": {
+            "geometry": array_attrs.geometry_name,
+            "positions": _to_serializable(array_attrs.positions),
+            "orientation": _to_serializable(array_attrs.orientation),
+            "center": _to_serializable(array_attrs.center),
+            "normal": _to_serializable(array_attrs.normal),
+            "spacing": array_attrs.spacing,
+        },
+        "time_axis": time_axis,
+        "doa": doa,
+        "timestamps": _to_serializable(timestamps_out),
+        "rirs_shape": list(rirs.shape),
+        "dynamic": bool(t_steps > 1),
+    }
+
+    if source_info is not None:
+        metadata["source_info"] = _to_serializable(source_info)
+    if extra:
+        metadata["extra"] = _to_serializable(extra)
+    return metadata
+
+
+def save_metadata_json(path: Path, metadata: Dict[str, Any]) -> None:
+    """Save metadata as JSON to the given path.
+
+    Example:
+        >>> save_metadata_json(Path(\"outputs/scene_metadata.json\"), metadata)
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2)
+
+
+def _normalize_traj(traj: Optional[Tensor], pos: Tensor, dim: int, name: str) -> Tensor:
+    if traj is None:
+        if pos.ndim != 2 or pos.shape[1] != dim:
+            raise ValueError(f"{name} default positions must have shape (N, {dim})")
+        return pos.unsqueeze(0)
+    if not torch.is_tensor(traj):
+        raise TypeError(f"{name} must be a Tensor")
+    if traj.ndim == 2:
+        if traj.shape[1] != dim:
+            raise ValueError(f"{name} must have shape (T, {dim})")
+        return traj.unsqueeze(1)
+    if traj.ndim == 3:
+        if traj.shape[2] != dim:
+            raise ValueError(f"{name} must have shape (T, N, {dim})")
+        return traj
+    raise ValueError(f"{name} must have shape (T, N, {dim})")
+
+
+def _compute_doa(src_traj: Tensor, mic_traj: Tensor) -> tuple[Tensor, Tensor]:
+    vec = src_traj[:, :, None, :] - mic_traj[:, None, :, :]
+    x = vec[..., 0]
+    y = vec[..., 1]
+    azimuth = torch.atan2(y, x)
+    if vec.shape[-1] < 3:
+        elevation = torch.zeros_like(azimuth)
+    else:
+        z = vec[..., 2]
+        r_xy = torch.sqrt(x**2 + y**2)
+        elevation = torch.atan2(z, r_xy)
+    return azimuth, elevation
+
+
+def _array_attributes(mics: MicrophoneArray) -> ArrayAttributes:
+    pos = mics.positions
+    n_mic = pos.shape[0]
+    if n_mic == 1:
+        geometry = "single"
+    elif n_mic == 2:
+        geometry = "binaural"
+    else:
+        geometry = "custom"
+    center = pos.mean(dim=0)
+    spacing = None
+    if n_mic >= 2:
+        dists = torch.cdist(pos, pos)
+        dists = dists[dists > 0]
+        if dists.numel() > 0:
+            spacing = float(dists.min().item())
+    return ArrayAttributes(
+        geometry_name=geometry,
+        positions=pos,
+        orientation=mics.orientation,
+        center=center,
+        normal=None,
+        spacing=spacing,
+    )
+
+
+def _to_serializable(value: Any) -> Any:
+    if value is None:
+        return None
+    if torch.is_tensor(value):
+        return value.detach().cpu().tolist()
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, dict):
+        return {k: _to_serializable(v) for k, v in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_to_serializable(v) for v in value]
+    return value