API reference¶

def download_locata(
    split: str | Sequence[str] = "dev", *, data_dir: str | Path | None = None
) -> Path:
    """Install dev/eval from the pinned official release and return its root.

    Storage lookup is ``data_dir``, then ``LOCATA_DATA_DIR``, then the platform
    user-data directory. ``LOCATA_ROOT`` does not affect downloads. Each split
    is verified and installed independently. Existing managed data is checked
    and reused; unknown or altered data raises :class:`LocataError`.

    Requires ``locata-torch[download]``. This explicit operation performs network
    I/O; Dataset construction never downloads. Call it before starting workers.
    """
    splits = select_splits(split)
    try:
        from filelock import FileLock, Timeout
    except ImportError as exc:
        raise LocataError(
            'install "locata-torch[download]" to enable downloading'
        ) from exc
    store = data_directory(data_dir)
    try:
        _check_store(store)
        required = 0
        for value in splits:
            manifest = read_manifest(store, value)
            target = dataset_directory(store) / value
            if manifest is None:
                if target.exists() or target.is_symlink():
                    raise LocataError(
                        f"{target}: unknown existing split without a manifest"
                    )
                spec = ARCHIVES[value]
                cached = store / "archives" / RELEASE / f"{value}.zip"
                retained = (
                    spec.size
                    if cached.is_file()
                    and not cached.is_symlink()
                    and cached.stat().st_size == spec.size
                    else 0
                )
                required += spec.size - retained + spec.expanded_size
        if required:
            _space(store, required)
        for relative in (
            f"archives/{RELEASE}",
            f"datasets/{RELEASE}",
            f"state/{RELEASE}",
            "staging",
            "locks",
        ):
            (store / relative).mkdir(parents=True, exist_ok=True)
        for value in splits:
            lock_path = store / "locks" / f"{RELEASE}-{value}.lock"
            if lock_path.is_symlink():
                raise LocataError(f"{lock_path}: lock file is a symlink")
            with FileLock(lock_path, timeout=60):
                _check_store(store)
                manifest = read_manifest(store, value)
                if manifest is not None:
                    _recover(store, value, manifest)
                else:
                    target = dataset_directory(store) / value
                    if target.exists() or target.is_symlink():
                        raise LocataError(f"{target}: unknown existing split")
                    _install(store, ARCHIVES[value])
        return dataset_directory(store)
    except (OSError, ValueError, zipfile.BadZipFile, Timeout) as exc:
        raise LocataError(f"{store}: LOCATA installation failed: {exc}") from exc

def collate_locata(samples: Sequence[LocataSample]) -> LocataBatch:
    """Right-pad audio; retain annotations and optional sources per sample.

    Args:
        samples: Nonempty samples sharing array, sample rate, channels, and dtype.

    Returns:
        CPU audio `[B, C, T_max]`, int64 lengths, a boolean padding-only mask,
        and lists of unmodified clocks, annotations, sources, and metadata.

    Raises:
        ValueError: Samples are empty, heterogeneous, or have non-CPU waveforms.
    """
    if not samples:
        raise ValueError("cannot collate an empty batch")
    first = samples[0]
    waveform = first["waveform"]
    signature = (
        first["metadata"]["array"],
        first["sample_rate"],
        waveform.shape[0],
        waveform.dtype,
    )
    for sample in samples:
        audio = sample["waveform"]
        if audio.ndim != 2 or audio.device.type != "cpu":
            raise ValueError("waveform must be a CPU tensor [channels, samples]")
        if (
            sample["metadata"]["array"],
            sample["sample_rate"],
            audio.shape[0],
            audio.dtype,
        ) != signature:
            raise ValueError(
                "batch must share array, sample rate, channel count and dtype"
            )
    lengths = torch.tensor([s["waveform"].shape[1] for s in samples], dtype=torch.int64)
    maximum = int(lengths.max())
    padded = waveform.new_zeros(len(samples), waveform.shape[0], maximum)
    for row, sample in enumerate(samples):
        padded[row, :, : sample["waveform"].shape[1]] = sample["waveform"]
    return {
        "waveform": padded,
        "lengths": lengths,
        "audio_mask": torch.arange(maximum).unsqueeze(0) < lengths.unsqueeze(1),
        "sample_rate": first["sample_rate"],
        "audio_time": [s["audio_time"] for s in samples],
        "time_origin": [s["time_origin"] for s in samples],
        "required_time": [s["required_time"] for s in samples],
        "array_pose": [s["array_pose"] for s in samples],
        "sources": [s["sources"] for s in samples],
        "metadata": [s["metadata"] for s in samples],
    }

def world_to_array(
    source_position: Tensor, array_position: Tensor, array_rotation: Tensor
) -> Tensor:
    """Apply R.T @ (h - p) to simultaneous, broadcastable Cartesian positions.

    Args:
        source_position: World xyz in metres, with trailing shape `[3]`.
        array_position: Array origin in world metres, with trailing shape `[3]`.
        array_rotation: Local-to-world rotation with trailing shape `[3, 3]`.

    Returns:
        Float64 array-frame vectors with broadcast leading dimensions.

    Raises:
        ValueError: Shapes, finiteness, or proper-rotation checks fail.
    """
    if (
        source_position.shape[-1:] != (3,)
        or array_position.shape[-1:] != (3,)
        or array_rotation.shape[-2:] != (3, 3)
    ):
        raise ValueError("position shape must end in 3 and rotation shape in (3, 3)")
    h, p, rotation = (
        tensor.to(dtype=torch.float64)
        for tensor in (source_position, array_position, array_rotation)
    )
    if not all(bool(torch.isfinite(tensor).all()) for tensor in (h, p, rotation)):
        raise ValueError("geometry must be finite")
    identity = torch.eye(3, dtype=torch.float64, device=rotation.device)
    if not torch.allclose(
        rotation.transpose(-1, -2) @ rotation,
        identity.expand_as(rotation),
        atol=1e-5,
        rtol=1e-5,
    ) or not torch.allclose(
        torch.linalg.det(rotation),
        torch.ones_like(torch.linalg.det(rotation)),
        atol=1e-5,
        rtol=1e-5,
    ):
        raise ValueError("rotation must be an orthonormal matrix with determinant +1")
    try:
        return (rotation.transpose(-1, -2) @ (h - p).unsqueeze(-1)).squeeze(-1)
    except RuntimeError as exc:
        raise ValueError(
            "geometry shapes must be broadcastable on the same device"
        ) from exc

def locata_doa(array_pose: Pose, source_pose: Pose) -> DOA:
    """Derive angles only for poses sharing exactly the same origin and times.

    Args:
        array_pose: Array world position and local-to-world rotation.
        source_pose: Source world position at exactly matching timestamps.

    Returns:
        Float64 vectors in metres, azimuth in `[-pi, pi)` measured from +y,
        inclination in `[0, pi]` measured from +z, and range in metres.
        Angles use radians; no temporal interpolation is performed.

    Raises:
        ValueError: Clocks or shapes differ, geometry is invalid, or distance
            is zero.
    """
    if array_pose["time_origin"] != source_pose["time_origin"]:
        raise ValueError("pose time origins must match")
    if not torch.equal(array_pose["time"], source_pose["time"]):
        raise ValueError("pose times must match exactly; align them explicitly")
    rows = len(array_pose["time"])
    if (
        array_pose["position"].shape != (rows, 3)
        or source_pose["position"].shape != (rows, 3)
        or array_pose["rotation"].shape != (rows, 3, 3)
    ):
        raise ValueError("pose shape must match its time rows")
    vector = world_to_array(
        source_pose["position"], array_pose["position"], array_pose["rotation"]
    )
    distance = torch.linalg.vector_norm(vector, dim=-1)
    if (distance == 0).any():
        raise ValueError("DOA is undefined at zero distance")
    azimuth = torch.atan2(vector[:, 1], vector[:, 0]) - math.pi / 2
    return {
        "vector": vector,
        "azimuth": torch.remainder(azimuth + math.pi, 2 * math.pi) - math.pi,
        "inclination": torch.acos((vector[:, 2] / distance).clamp(-1, 1)),
        "range": distance,
    }