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

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

torchrir/plotting.py

@@ -20,7 +20,15 @@ def plot_scene_static(
     title: Optional[str] = None,
     show: bool = False,
 ):
-    """Plot a static room with source and mic positions."""
+    """Plot a static room with source and mic positions.
+
+    Example:
+        >>> ax = plot_scene_static(
+        ...     room=[6.0, 4.0, 3.0],
+        ...     sources=[[1.0, 2.0, 1.5]],
+        ...     mics=[[2.0, 2.0, 1.5]],
+        ... )
+    """
     plt, ax = _setup_axes(ax, room)
 
     size = _room_size(room, ax)
@@ -46,21 +54,35 @@ def plot_scene_dynamic(
     src_traj: Tensor | Sequence,
     mic_traj: Tensor | Sequence,
     step: int = 1,
+    src_pos: Optional[Tensor | Sequence] = None,
+    mic_pos: Optional[Tensor | Sequence] = None,
     ax: Any | None = None,
     title: Optional[str] = None,
     show: bool = False,
 ):
-    """Plot source and mic trajectories within a room."""
+    """Plot source and mic trajectories within a room.
+
+    If trajectories are static, only positions are plotted.
+
+    Example:
+        >>> ax = plot_scene_dynamic(
+        ...     room=[6.0, 4.0, 3.0],
+        ...     src_traj=src_traj,
+        ...     mic_traj=mic_traj,
+        ... )
+    """
     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)
+    src_traj = _as_trajectory(src_traj)
+    mic_traj = _as_trajectory(mic_traj)
+    src_pos_t = _extract_positions(src_pos, ax) if src_pos is not None else src_traj[0]
+    mic_pos_t = _extract_positions(mic_pos, ax) if mic_pos is not None else mic_traj[0]
 
-    _plot_trajectories(ax, src_traj, step=step, label="source path")
-    _plot_trajectories(ax, mic_traj, step=step, label="mic path")
+    _plot_entity(ax, src_traj, src_pos_t, step=step, label="sources", marker="^")
+    _plot_entity(ax, mic_traj, mic_pos_t, step=step, label="mics", marker="o")
 
     if title:
         ax.set_title(title)
@@ -70,7 +92,9 @@ def plot_scene_dynamic(
     return ax
 
 
-def _setup_axes(ax: Any | None, room: Room | Sequence[float] | Tensor) -> tuple[Any, Any]:
+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
 
@@ -109,8 +133,9 @@ 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")
+    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())
@@ -164,7 +189,9 @@ def _draw_room_3d(ax: Any, size: Tensor) -> None:
     ax.set_zlabel("z")
 
 
-def _extract_positions(entity: Source | MicrophoneArray | Tensor | Sequence, ax: Any | None) -> Tensor:
+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
@@ -176,18 +203,34 @@ def _extract_positions(entity: Source | MicrophoneArray | Tensor | Sequence, ax:
     return pos
 
 
-def _scatter_positions(ax: Any, positions: Tensor, *, label: str, marker: str) -> None:
+def _scatter_positions(
+    ax: Any,
+    positions: Tensor,
+    *,
+    label: str,
+    marker: str,
+    color: Optional[str] = None,
+) -> 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)
+        ax.scatter(
+            positions[:, 0], positions[:, 1], label=label, marker=marker, color=color
+        )
     else:
-        ax.scatter(positions[:, 0], positions[:, 1], positions[:, 2], label=label, marker=marker)
+        ax.scatter(
+            positions[:, 0],
+            positions[:, 1],
+            positions[:, 2],
+            label=label,
+            marker=marker,
+            color=color,
+        )
 
 
-def _as_trajectory(traj: Tensor | Sequence, ax: Any | None) -> Tensor:
+def _as_trajectory(traj: Tensor | Sequence) -> Tensor:
     """Validate and normalize a trajectory tensor."""
     traj = as_tensor(traj)
     if traj.ndim != 3:
@@ -195,16 +238,73 @@ def _as_trajectory(traj: Tensor | Sequence, ax: Any | None) -> Tensor:
     return traj
 
 
-def _plot_trajectories(ax: Any, traj: Tensor, *, step: int, label: str) -> None:
-    """Plot trajectories for each entity."""
+def _plot_entity(
+    ax: Any,
+    traj: Tensor,
+    positions: Tensor,
+    *,
+    step: int,
+    label: str,
+    marker: str,
+) -> None:
+    """Plot trajectories and/or static positions with a unified legend entry."""
     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}")
+    import matplotlib.pyplot as plt
+
+    if positions.shape != traj.shape[1:]:
+        positions = traj[0]
+    moving = _is_moving(traj, positions)
+    colors = plt.rcParams.get("axes.prop_cycle", None)
+    if colors is not None:
+        palette = colors.by_key().get("color", [])
     else:
-        for idx in range(traj.shape[1]):
-            xyz = traj[::step, idx]
-            ax.plot(xyz[:, 0], xyz[:, 1], xyz[:, 2], label=f"{label} {idx}")
+        palette = []
+    if not palette:
+        palette = ["C0", "C1", "C2", "C3", "C4", "C5"]
+
+    dim = traj.shape[2]
+    for idx in range(traj.shape[1]):
+        color = palette[idx % len(palette)]
+        lbl = label if idx == 0 else "_nolegend_"
+        if moving:
+            if dim == 2:
+                xy = traj[::step, idx]
+                ax.plot(
+                    xy[:, 0],
+                    xy[:, 1],
+                    label=lbl,
+                    color=color,
+                    marker=marker,
+                    markevery=[0],
+                )
+            else:
+                xyz = traj[::step, idx]
+                ax.plot(
+                    xyz[:, 0],
+                    xyz[:, 1],
+                    xyz[:, 2],
+                    label=lbl,
+                    color=color,
+                    marker=marker,
+                    markevery=[0],
+                )
+        pos = positions[idx : idx + 1]
+        _scatter_positions(ax, pos, label="_nolegend_", marker=marker, color=color)
+    if not moving:
+        # ensure legend marker uses the group label
+        _scatter_positions(
+            ax,
+            positions[:1],
+            label=label,
+            marker=marker,
+            color=palette[0],
+        )
+
+
+def _is_moving(traj: Tensor, positions: Tensor, *, tol: float = 1e-6) -> bool:
+    """Return True if any trajectory deviates from the provided positions."""
+    if traj.numel() == 0:
+        return False
+    pos0 = positions.unsqueeze(0).expand_as(traj)
+    return bool(torch.any(torch.linalg.norm(traj - pos0, dim=-1) > tol).item())

torchrir/plotting_utils.py

@@ -26,6 +26,8 @@ def plot_scene_and_save(
 ) -> tuple[list[Path], list[Path]]:
     """Plot static and dynamic scenes and save images to disk.
 
+    Dynamic plots show trajectories for moving entities and points for fixed ones.
+
     Args:
         out_dir: Output directory for PNGs.
         room: Room size tensor or sequence.
@@ -41,6 +43,17 @@ def plot_scene_and_save(
 
     Returns:
         Tuple of (static_paths, dynamic_paths).
+
+    Example:
+        >>> plot_scene_and_save(
+        ...     out_dir=Path("outputs"),
+        ...     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,
+        ...     prefix="scene",
+        ... )
     """
     out_dir = Path(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
@@ -85,11 +98,12 @@ def plot_scene_and_save(
                 room=view_room,
                 src_traj=view_src_traj,
                 mic_traj=view_mic_traj,
+                src_pos=view_src,
+                mic_pos=view_mic,
                 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)
@@ -113,7 +127,10 @@ def _positions_to_cpu(entity: torch.Tensor | object) -> torch.Tensor:
     return pos
 
 
-def _traj_steps(src_traj: Optional[torch.Tensor | Sequence], mic_traj: Optional[torch.Tensor | Sequence]) -> int:
+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])
@@ -142,32 +159,3 @@ def _save_axes(ax: Any, path: Path, *, show: bool) -> None:
     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

@@ -13,7 +13,13 @@ from .scene import Scene
 
 @dataclass(frozen=True)
 class RIRResult:
-    """Container for RIRs with metadata."""
+    """Container for RIRs with metadata.
+
+    Example:
+        >>> from torchrir import ISMSimulator
+        >>> result = ISMSimulator().simulate(scene, config)
+        >>> rirs = result.rirs
+    """
 
     rirs: Tensor
     scene: Scene

torchrir/room.py

@@ -13,7 +13,11 @@ from .utils import as_tensor, ensure_dim
 
 @dataclass(frozen=True)
 class Room:
-    """Room geometry and acoustic parameters."""
+    """Room geometry and acoustic parameters.
+
+    Example:
+        >>> room = Room.shoebox(size=[6.0, 4.0, 3.0], fs=16000, beta=[0.9] * 6)
+    """
 
     size: Tensor
     fs: float
@@ -43,7 +47,11 @@ class Room:
         device: Optional[torch.device | str] = None,
         dtype: Optional[torch.dtype] = None,
     ) -> "Room":
-        """Create a rectangular (shoebox) room."""
+        """Create a rectangular (shoebox) room.
+
+        Example:
+            >>> room = Room.shoebox(size=[6.0, 4.0, 3.0], fs=16000, beta=[0.9] * 6)
+        """
         size_t = as_tensor(size, device=device, dtype=dtype)
         size_t = ensure_dim(size_t)
         beta_t = None
@@ -54,7 +62,11 @@ class Room:
 
 @dataclass(frozen=True)
 class Source:
-    """Source container with positions and optional orientation."""
+    """Source container with positions and optional orientation.
+
+    Example:
+        >>> sources = Source.from_positions([[1.0, 2.0, 1.5]])
+    """
 
     positions: Tensor
     orientation: Optional[Tensor] = None
@@ -70,20 +82,6 @@ class 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,
@@ -103,7 +101,11 @@ class Source:
 
 @dataclass(frozen=True)
 class MicrophoneArray:
-    """Microphone array container."""
+    """Microphone array container.
+
+    Example:
+        >>> mics = MicrophoneArray.from_positions([[2.0, 2.0, 1.5]])
+    """
 
     positions: Tensor
     orientation: Optional[Tensor] = None
@@ -119,20 +121,6 @@ class 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,

torchrir/scene.py

@@ -13,7 +13,12 @@ from .room import MicrophoneArray, Room, Source
 
 @dataclass(frozen=True)
 class Scene:
-    """Container for room, sources, microphones, and optional trajectories."""
+    """Container for room, sources, microphones, and optional trajectories.
+
+    Example:
+        >>> scene = Scene(room=room, sources=sources, mics=mics, src_traj=src_traj, mic_traj=mic_traj)
+        >>> scene.validate()
+    """
 
     room: Room
     sources: Source