AudioDataSource

class AudioDataSource(path: str | Path, *, audio_dir: str | Path | None = None, num_records: int | None = None, sample_rate: int | None = None, mono: bool = False, duration: float | None = None, pad_mode: Literal['constant', 'wrap'] = 'constant', filter_fn: Callable[[Dict], bool] | None = None, on_read_error: Literal['raise', 'skip', 'warn'] = 'raise')

A DataSource that reads audio files based on a manifest file created by AudioWriter.

This DataSource is designed to work seamlessly with the output of AudioWriter, reading audio files and restoring their associated extras from NPZ manifests.

NPZ format provides: - Efficient binary storage with 20x+ compression vs JSON for large datasets - Fast loading without text parsing - Exact numeric type preservation - Support for both compressed and uncompressed variants

Parameters:

  • path – The dataset directory an AudioWriter produced (its manifest.npz is located inside), or an explicit path to a manifest file (NPZ format)

  • audio_dir – Optional directory containing audio files. If None, uses manifest directory

  • num_records – Optional limit on number of records to load

  • sample_rate – Optional target sample rate for resampling

  • mono – Whether to convert audio to mono

  • duration –

    Optional duration to trim/pad audio to (in seconds)

    When sample_rate, mono or duration change the audio relative to what was written, the manifest’s derived fields (lufs, lufs_windows, codes, latents) are not restored – they describe the written audio, not the altered one. Annotation fields (pitch, velocity, note_duration) are restored either way.

  • pad_mode – Padding mode if duration is specified (“constant” or “wrap”)

  • filter_fn – Optional function to filter manifest entries

  • on_read_error – What to do when one entry’s audio file is missing or cannot be decoded. "raise" (default) raises AudioReadError, naming the path; "warn" drops the item (ds[i] returns None, which grain skips at iteration) and emits a UserWarning; "skip" drops it quietly. See create_audio_dataset() for the reasoning.

Example

First write some audio with AudioWriter so there is a manifest.npz to read back:

>>> import tempfile
>>> import numpy as np
>>> import jax.numpy as jnp
>>> from audiotree import AudioTree, AudioWriter
>>> out_dir = tempfile.mkdtemp()
>>> lufs = np.full((5,), -10.0, dtype=np.float32)
>>> with AudioWriter(out_dir) as writer:
...     _ = writer.write(
...         AudioTree.create(jnp.zeros((5, 1, 44100)), 44100, lufs=lufs)
...     )
>>> manifest_path = f"{out_dir}/manifest.npz"

Read straight from the NPZ manifest:

>>> source = AudioDataSource(manifest_path)
>>> len(source)
5
>>> source[0].waveform.shape
(1, 1, 44100)

Filter entries by extras while loading:

>>> source = AudioDataSource(
...     manifest_path,
...     filter_fn=lambda entry: entry.get('lufs', -float('inf')) > -20
... )

Or point straight at the dataset directory (its manifest.npz is located inside), the same way TreeDataSource is used:

>>> source = AudioDataSource(out_dir)
filter(predicate: Callable[[Dict], bool]) → AudioDataSource

Create a new AudioDataSource keeping the entries matching predicate.

Filtering narrows the current entries, so filters compose: source.filter(a).filter(b) keeps the entries matching both.

Parameters:

predicate – Function called with a manifest entry, returning whether to keep it

Returns:

New AudioDataSource with the matching entries

Raises:

ValueError – If no entry matches

filter_by_lufs(min_lufs: float | None = None, max_lufs: float | None = None) → AudioDataSource

Create a new AudioDataSource filtered by loudness range.

Filters entries based on the ‘lufs’ field in the manifest. Works with manifests created by AudioWriter in NPZ format.

Narrows the receiver’s entries, so this composes with any other filter already applied (including a constructor filter_fn and a num_records cap).

Parameters:

  • min_lufs – Minimum loudness in LUFS (inclusive)

  • max_lufs – Maximum loudness in LUFS (inclusive)

Returns:

New AudioDataSource with filtered entries

Raises:

ValueError – If no entry falls in the range

Example

Write four items with known per-item loudness so the manifest records a lufs field to filter on:

>>> import tempfile
>>> import numpy as np
>>> import jax.numpy as jnp
>>> from audiotree import AudioTree, AudioWriter
>>> out_dir = tempfile.mkdtemp()
>>> audio = AudioTree.create(
...     jnp.zeros((4, 1, 44100)), 44100,
...     lufs=np.array([-30.0, -18.0, -10.0, -25.0], dtype=np.float32),
... )
>>> with AudioWriter(out_dir) as writer:
...     _ = writer.write(audio)
>>> source = AudioDataSource(out_dir)

Keep only samples louder than -20 LUFS:

>>> loud_source = source.filter_by_lufs(min_lufs=-20.0)
>>> len(loud_source)
2

Keep samples within a specific loudness range:

>>> mid_source = source.filter_by_lufs(min_lufs=-30.0, max_lufs=-15.0)
>>> len(mid_source)
3

Filters compose, so chaining keeps only what matches both:

>>> len(loud_source.filter_by_lufs(max_lufs=-15.0))
1
filter_by_tag(tag_name: str, tag_value) → AudioDataSource

Create a new AudioDataSource filtered by a specific tag value.

Narrows the receiver’s entries, so this composes with any other filter already applied (including a constructor filter_fn and a num_records cap).

Parameters:

  • tag_name – Name of the tag to filter by

  • tag_value – Value the tag must have

Returns:

New AudioDataSource with filtered entries

Raises:

ValueError – If no entry has that tag value

get_all_entries() → List[Dict]

Get all manifest entries.

Returns:

List of all manifest entry dictionaries

get_entry(index: int) → Dict

Get the raw manifest entry for a given index.

Parameters:

index – Index of the entry

Returns:

Dictionary containing the manifest entry