audiotree¶

A flax.struct.dataclass for holding audio information including a waveform, sample rate, and extras.

The AudioTree class is inspired by Descript AudioTools’s AudioSignal.

The constructor stores its arguments verbatim – it neither reshapes the waveform nor encodes provenance. Use create() (which accepts a (Samples,) or (Channels, Samples) waveform and filepath / source strings) unless you already hold batched arrays.

Parameters:

waveform (np.ndarray or jax.Array) – Audio waveform data shaped (Batch, Channels, Samples), or None for token-only trees (codes / latents without audio).
sample_rate (int) – Sample rate of waveform, such as 44100 Hz.
lufs (np.ndarray or jax.Array, optional) – Integrated loudness of the audio waveform in LUFS, shaped (Batch,). You may not need to set this when initializing. Instead, use replace_lufs() to create a new AudioTree with lufs (and lufs_windows) calculated.
lufs_windows (np.ndarray or jax.Array, optional) – Per-window integrated loudness in LUFS, shaped (Batch, Windows) — one value per non-overlapping analysis window. Populated alongside lufs by replace_lufs().
pitch (np.ndarray or jax.Array, optional) – The MIDI pitch where 60 is middle C. The shape is (Batch,).
velocity (np.ndarray or jax.Array, optional) – The MIDI velocity between 0 and 127. The shape is (Batch,).
note_duration (np.ndarray or jax.Array, optional) – A note duration in units of your choice. The value is not necessarily the same as the duration of the audio data. The shape is (Batch,).
codes (np.ndarray or jax.Array, optional) – The neural audio codec tokens for the audio.
latents (np.ndarray or jax.Array, optional) – The latent representations of the audio.
extras (dict) – Any extra per-item data can be placed here. This dict is entirely yours: the library never plants keys of its own in it. Leaves should be arrays with the batch as their leading axis, or lists of strings (one per item) – create() validates this; a list of numbers would be flattened into scalar leaves by jax.tree_util and silently stop tracking the batch.
_metadata (dict) – Library-managed provenance container – private, as the underscore says; not for user data (put your own data in extras). It holds the encoded "filepath" and "source" arrays plus the per-item "offset" array: pass filepath= / source= / offset= to create() to fill it, and read the values back via the filepath, source and offset properties. On disk it keeps the spelling metadata.

Example

>>> audio = AudioTree.create(jnp.zeros((2, 44100)), 44100)  # stereo, 1 s
>>> audio.waveform.shape
(1, 2, 44100)
>>> audio.sample_rate
44100

Note

Every consumer that enumerates these fields derives its list from dataclasses.fields (see PYTREE_FIELDS below), so adding a field here is picked up automatically.

property backend: str¶

Which array library this tree’s arrays belong to.

One of "numpy", "jax", or "mixed". A tree is not required to be homogeneous and quietly stops being so more often than you would expect: applying a NumPy-namespace transform to a JAX tree converts the fields it touches, so audiotree.transforms.trim on a JAX tree hands back a NumPy waveform. Nothing is wrong with that until something downstream assumes otherwise – jax.jit on a NumPy leaf silently re-uploads it every call – which is what this property is for.

"mixed" is reported rather than raised, so it is safe to log. Convert with jax.device_put() / jax.device_get().

Examples

>>> import numpy as np, jax.numpy as jnp
>>> from audiotree import AudioTree
>>> AudioTree.create(np.zeros((1, 1, 8)), 16000).backend
'numpy'
>>> AudioTree.create(jnp.zeros((1, 1, 8)), 16000).backend
'jax'

static batch(items: Sequence[Any]) → Any¶

Batch function for use with grain’s IterDataset.batch().

Concatenates AudioTree objects along the batch axis (axis 0). Use this instead of grain’s default batching, which would add an extra dimension since AudioTree already has shape (batch, channels, samples).

Supports arbitrary nested structures containing AudioTrees. All arrays (including AudioTrees) are concatenated along axis 0, so data should have a leading batch dimension.

Concatenation dispatches on the leaves’ array library, so JAX in gives JAX out: batching a tree of jax.Array leaves stays on device instead of forcing a blocking host sync and silently returning NumPy.

Parameters:: items – Sequence of AudioTree objects, or structures (dicts, lists, etc.) containing AudioTree objects. Must be non-empty — there is no array library, sample rate or structure to infer from nothing.
Returns:: Batched structure with the same shape as the input items.
Raises:: ValueError – If items is empty.

Example

>>> a = AudioTree.create(jnp.zeros((1, 1, 16000)), 16000)
>>> batched = AudioTree.batch([a, a, a])  # concatenate along the batch axis
>>> batched.waveform.shape
(3, 1, 16000)
>>> isinstance(batched.waveform, jax.Array)  # JAX in, JAX out
True

With Grain, pass it as the batch_fn (each item already has a leading batch axis):

ds.to_iter_dataset().batch(32, batch_fn=AudioTree.batch)

property batch_size: int¶

Return the size of the leading (batch) axis.

Derived from waveform, falling back to codes / latents for audio-less trees (e.g. token-only training examples).

This is the leading axis, not the item count: on a mini-batched tree (rank 4, from reshape_mini_batches()) it is the number of mini-batches, which is also what len(), iteration and indexing walk over. flatten_mini_batches() first if you want items.

clear_lufs() → Self¶

Return a copy with the cached lufs and lufs_windows dropped.

Shorthand for replace(lufs=None, lufs_windows=None). Call it after changing the waveform through a bare replace(), so the next replace_lufs() measures the new audio instead of trusting a cached value that describes the old one. The built-in transforms and AudioTree methods invalidate for you; this is for hand-rolled edits.

Example

>>> audio = AudioTree.create(jnp.zeros((44100,)), 44100).replace_lufs()
>>> quieter = audio.replace(waveform=audio.waveform * 0.5).clear_lufs()
>>> quieter.lufs is None
True

codes: ndarray | None = None¶

Create an AudioTree, normalizing the waveform to (Batch, Channels, Samples).

A bare (Samples,) or (Channels, Samples) waveform gains the missing leading axes, so you don’t have to reshape by hand. filepath, source and offset go into the library-managed _metadata container; extras is left to the caller.

Parameters:

waveform – Audio of shape (Samples), (Channels, Samples), or (Batch, Channels, Samples), or None for token-only trees (e.g. codes / latents without audio).
sample_rate – Sample rate of waveform in Hz (e.g. 44100).
lufs – Optional precomputed integrated loudness (LUFS); usually left None and filled by replace_lufs().
lufs_windows – Optional precomputed per-window loudness (Batch, Windows); usually left None and filled by replace_lufs().
pitch – Optional MIDI pitch (Batch,) (60 = middle C).
velocity – Optional MIDI velocity (Batch,) in [0, 127].
note_duration – Optional per-note duration (Batch,) (not the audio duration).
codes – Optional neural-codec tokens.
latents – Optional latent representations.
extras – Optional extras dict of additional per-item leaves (copied, not mutated). Leaves must be arrays with the batch as their leading axis, lists of strings (one per item), or nested dicts of those; a Python list/tuple of non-strings is rejected – jax.tree_util would treat each element as its own scalar leaf, silently misaligning it with the batch. Convert such lists with np.asarray.
filepath – Optional path(s) for the batch; encoded into _metadata["filepath"] and read back via the filepath property. Pass a single path to tag the whole batch (it is repeated for every item), or a list with one path per batch item.
source – Optional source-group name(s) (e.g. "music"), encoded into _metadata["source"] and read back via the source property. Pass a single string to tag the whole batch, or a list with one name per batch item (unlike from_file(), which only accepts a single string).
offset – Optional source-file offset(s) in seconds, stored in _metadata["offset"] and read back via the offset property. Pass a single number to tag the whole batch, or a (Batch,) array with one offset per batch item. from_file() and the dataset loaders record this automatically.

Returns:

A new AudioTree whose waveform is (Batch, Channels, Samples).

Return type:

AudioTree

Example

>>> audio = AudioTree.create(jnp.zeros((44100,)), 44100)  # 1 s mono
>>> audio.waveform.shape
(1, 1, 44100)
>>> audio.sample_rate
44100

property device: Device | None¶

The device this tree’s arrays live on, or None on NumPy.

jax.device_put and jax.device_get already move a tree; this is the missing half – asking where it currently is without reaching into a leaf and hoping the rest agree.

Returns:: The jax.Device shared by every array leaf, or None if the tree is NumPy-backed (host memory, no JAX device).
Raises:: ValueError – If the leaves do not agree on one device – a mixed NumPy/JAX tree, or JAX arrays committed to different devices. Unlike backend this raises, because there is no honest single answer and returning one of them would be a guess.

classmethod excerpt(audio_path: str | Path, rng: Generator, duration: float, offset: float = 0.0, excerpt: ExcerptConfig | None = None, **kwargs) → Self | None¶

Create an AudioTree from one section of an audio file.

Which section is up to excerpt.strategy: "start" takes the audio at offset, "random" (the default) draws a single offset, and "loudest" defers to loudest_excerpt().

Parameters:

audio_path (str or Path) – Path to audio file.
rng (np.random.Generator) – Random number generator such as np.random.default_rng(42).
duration (float) – Duration in seconds of audio data; must be positive. The audio data will be trimmed or lengthened as necessary.
offset (float, optional) – Earliest offset in seconds the excerpt may start at.
excerpt (ExcerptConfig, optional) – How to choose the offset; defaults to a uniformly random one. See ExcerptConfig.
**kwargs – Keyword arguments passed to AudioTree.from_file.

Returns:

AudioTree, or None when excerpt searched for the loudest section, found nothing above the cutoff, and says on_failure="skip".

extras: dict¶

property filepath: List[str] | List[list]¶

Return the decoded filepaths stored in _metadata['filepath'].

One string per batch item. A mini-batched tree (rank 4, from reshape_mini_batches()) has two leading axes, so it returns one list per mini-batch — the nesting always matches the tree’s leading axes. Call flatten_mini_batches() first for a flat list.

An empty list is returned if the AudioTree does not carry filepath provenance.

filter(predicate: Callable[[Self], bool]) → Self¶

Keep only the batch items for which predicate is true.

The batch is split into one tree per item, predicate is called on each, and the survivors are concatenated back into a single tree.

Parameters:: predicate – Called with a batch-of-1 AudioTree; return True to keep that item.
Returns:: A tree holding the kept items. When nothing is kept, the result has batch_size == 0 (every array field is empty along the batch axis, every string extras leaf is []) rather than being None. Filtering an already-empty tree returns such an empty tree without calling predicate, so chained filters compose.
Return type:: AudioTree
Raises:: ValueError – If the tree is mini-batched (rank 4). Dropping items independently within each mini-batch would leave the mini-batches ragged, so there is no rank-4 answer; flatten, filter, and reshape again.

Example

>>> waveform = jnp.stack([jnp.zeros((1, 8)), jnp.ones((1, 8))])
>>> audio = AudioTree.create(waveform, 16000)
>>> loud = audio.filter(lambda item: bool(item.waveform.max() > 0.5))
>>> loud.batch_size
1

flatten_mini_batches() → Self¶

Flatten mini-batches back into a single batch dimension.

Undoes the operation performed by reshape_mini_batches(), transforming audio data from shape (num_mini_batches, mini_batch_size, C, T) back to (B, C, T). String extras leaves lose their per-mini-batch nesting the same way, back to one string per item.

Returns:: AudioTree with the mini-batch dimension flattened into the batch dimension.
Raises:: ValueError – If the waveform has fewer than 4 dimensions, i.e. it was never mini-batched.

Example

>>> x = AudioTree(np.zeros((12, 1, 44100)), 44100)
>>> x_batched = x.reshape_mini_batches(3)
>>> x_batched.waveform.shape  # 4 mini-batches of size 3
(4, 3, 1, 44100)
>>> x_unbatched = x_batched.flatten_mini_batches()
>>> x_unbatched.waveform.shape  # back to original shape
(12, 1, 44100)

Create an AudioTree from an audio file path.

Parameters:

audio_path (str) – Path to audio file.
sample_rate (int, optional) – Sample rate of audio data, such as 44100 Hz. If left as None, the file’s original sample rate will be used.
offset (float, optional) – Offset in seconds to audio data.
duration (float, optional) – Duration in seconds of audio data. The audio data will be trimmed or extended as necessary.
mono (bool, optional) – Whether to force the audio data to be single-channel.
pad_mode (Literal) – If duration is not None, and duration is less than the length of the audio, then pad_mode controls how the audio is right-padded (numpy.pad modes). Options: “constant” (zeros, default), “edge” (repeat edge), “reflect” (mirror), “symmetric” (mirror with edge), “wrap” (circular/loop), or None (no padding).
filepath (Union[str, Path, List[str | Path]], optional) – The path to store as the returned AudioTree’s provenance (_metadata["filepath"], read back via the filepath property) – a single path or a one-item list, since from_file returns a batch of 1. If None (default) the provided audio_path will be used.
source (str, optional) – The source group name for this audio file (e.g., “music”, “speech”). Stored as provenance (_metadata["source"]) and accessible via the source property.
extras (dict, optional) – Additional extras to include in the AudioTree. The dict is purely user payload; the read offset is recorded as provenance (_metadata["offset"], read back via the offset property).
lufs (np.ndarray or jax.Array, optional) – Integrated loudness (LUFS) values to assign to the AudioTree.
lufs_windows (np.ndarray or jax.Array, optional) – Per-window loudness (LUFS) values to assign to the AudioTree.
pitch (np.ndarray or jax.Array, optional) – Pitch values to assign to the AudioTree.
velocity (np.ndarray or jax.Array, optional) – Velocity values to assign to the AudioTree.
note_duration (np.ndarray or jax.Array, optional) – Note note_duration values to assign to the AudioTree.
codes (np.ndarray or jax.Array, optional) – The neural audio codec tokens for the audio.
latents (np.ndarray or jax.Array, optional) – The latent representations of the audio.

Returns:

An instance of AudioTree.

Return type:

AudioTree

classmethod from_manifest(manifest_path: str | Path, *, audio_dir: str | Path | None = None, filter_fn: Callable[[Dict[str, Any]], bool] | None = None) → Self¶

Create an AudioTree by loading all items from a manifest file.

This loads all entries from a manifest file created by AudioWriter and creates a single AudioTree with all items in the batch dimension.

Parsing the manifest – header check, presence masks, string decoding – is audiotree._manifest.read_entries(), the one reader of that format; nothing here re-derives the file’s rules.

Parameters:

manifest_path – Path to the manifest file (NPZ format)
audio_dir – Optional directory containing audio files. If None, uses manifest directory
filter_fn – Optional predicate called with one manifest entry, a dict keyed by column name ("filename", "sample_rate", the AudioTree label fields, "extras_*", plus "tags"); return True to load that entry. These are the entries of audiotree._manifest.read_entries(), the same ones AudioDataSource passes its filter_fn, so one predicate serves both (that source additionally demotes bookkeeping numbers such as sample_rate to plain Python scalars). A column with no value for an entry is absent from that entry’s dict.

Returns:

AudioTree with all manifest entries concatenated along batch dimension

Raises:

ValueError – If the manifest is unreadable, holds no entries, names an audio file outside audio_dir, or filter_fn matches nothing.

Example

First, write a small manifest with AudioWriter (here, 100 one-second stereo items; the first 60 are loud, the rest quiet, so the lufs field is recorded for filtering):

>>> import tempfile
>>> from audiotree import AudioWriter
>>> out_dir = tempfile.mkdtemp()
>>> lufs = np.where(np.arange(100) < 60, -10.0, -30.0).astype(np.float32)
>>> batch = AudioTree.create(jnp.zeros((100, 2, 44100)), 44100, lufs=lufs)
>>> with AudioWriter(out_dir) as writer:
...     _ = writer.write(batch)
>>> manifest_path = f"{out_dir}/manifest.npz"

Load every item into one batched AudioTree:

>>> audio = AudioTree.from_manifest(manifest_path)
>>> audio.waveform.shape  # 100 items, stereo, 1 second each
(100, 2, 44100)

Load only entries that pass a filter on the manifest (the 60 loud items):

>>> audio = AudioTree.from_manifest(
...     manifest_path,
...     filter_fn=lambda entry: entry.get('lufs', -float('inf')) > -20
... )
>>> audio.waveform.shape
(60, 2, 44100)

latents: ndarray | None = None¶

classmethod loudest_excerpt(audio_path: str | Path, rng: Generator, excerpt: ExcerptConfig, **kwargs) → Self | None¶

Create an AudioTree from the loudest of several candidate excerpts.

Draws up to excerpt.num_tries offsets and keeps the loudest, stopping early once one exceeds excerpt.lufs_cutoff. This is a best-of-k search rather than a filter, so on a file where nothing clears the cutoff it still has to return something – excerpt.on_failure decides what.

Parameters:

audio_path (str) – Path to audio file.
rng (np.random.Generator) – Random number generator such as np.random.default_rng(42).
excerpt (ExcerptConfig) – How to search. strategy must be "loudest".
**kwargs – Keyword arguments passed to AudioTree.from_file.

Returns:

AudioTree, or None when nothing cleared the cutoff and on_failure="skip".

lufs: ndarray | None = None¶

lufs_windows: ndarray | None = None¶

normalize_lufs(target_lufs: float, *, max_gain_db: float | None = None, device: Literal['cpu', 'gpu', 'tpu'] | Device | None = None, engine: Literal['numpy', 'jax'] | None = None) → Self¶

Normalize audio to a target LUFS level.

Computes the current loudness (if not already set), then scales the audio to achieve the target LUFS. The returned AudioTree has updated waveform, lufs, and lufs_windows fields (a constant gain shifts every window’s LUFS by the same amount). Changing the level invalidates codes and latents, which describe the audio at its previous level.

Items whose loudness is not finite are passed through unscaled. Digital silence, and any excerpt below the BS.1770 absolute gate, measure -inf LUFS, for which no gain reaches the target; scaling by the implied +inf would produce an all-NaN waveform. Their lufs stays -inf, so a silent item is still identifiable afterwards.

Parameters:

target_lufs – Target loudness in LUFS (e.g., -18.0 for broadcast standard).
max_gain_db – Optional ceiling on the applied gain, so a very quiet (but still measurable) item is not amplified without bound. None (the default) applies whatever gain the target implies; a capped item lands at lufs + max_gain_db rather than at target_lufs.
device – Where to compute the loudness when it is not already set, forwarded to replace_lufs() (see there). Ignored when lufs is already populated.
engine – Which loudness kernel to use when it is not already set, forwarded to replace_lufs() (see there). Ignored when lufs is already populated.

Returns:

AudioTree with audio scaled to target LUFS and loudness updated.

Example

>>> t = jnp.arange(44100) / 44100  # 1 s at 44.1 kHz
>>> audio = AudioTree.create(0.5 * jnp.sin(2 * jnp.pi * 1000 * t), 44100)
>>> normalized = audio.normalize_lufs(-18.0)
>>> float(normalized.lufs[0])  # now at the target LUFS
-18.0

Silence is left alone instead of becoming NaN:

>>> silent = AudioTree.create(jnp.zeros((1, 1, 44100)), 44100)
>>> out = silent.normalize_lufs(-18.0)
>>> bool(jnp.all(out.waveform == 0.0)), float(out.lufs[0])
(True, -inf)

note_duration: ndarray | None = None¶

property num_channels: int¶: Return the number of audio channels (waveform.shape[-2]).

property offset: ndarray | None¶

Return the per-item source-file offsets stored in _metadata['offset'].

Seconds from the start of each item’s source file to its sample 0, recorded by from_file() and the dataset loaders. The shape follows the tree’s leading axes: (Batch,) normally, (Mini, Batch) on a mini-batched tree. Returns None when the tree carries no offset provenance, or when a time-shifting transform such as roll() invalidated it (under prob < 1 the whole batch’s offset is dropped, since “invalidated for some items only” cannot be represented in one array). An item with no source file (e.g. a synthetic pre-built dataset group) records NaN.

pitch: ndarray | None = None¶

replace(**updates: Unpack[_AudioTreeFields]) → Self¶

Return a new AudioTree with the given fields replaced.

Parameters:: **updates – Any subset of the fields above; everything else is carried over from self, which is never mutated.
Returns:: A copy of self with updates applied.
Return type:: AudioTree

Example

>>> audio = AudioTree.create(jnp.zeros((44100,)), 44100)
>>> quiet = audio.replace(waveform=audio.waveform * 0.5)
>>> quiet.sample_rate
44100

Note

flax.struct.dataclass installs its own replace over this one at class-creation time. That implementation is this one verbatim (dataclasses.replace) minus the typing, which is the whole reason to spell it out here.

replace_extras(**kwargs) → Self¶

Return a new AudioTree with kwargs merged into extras.

Syntactic sugar for self.replace(extras={**self.extras, **kwargs}). Keys in kwargs overwrite existing extras keys with the same name; all other keys are kept. Neither the original tree nor its extras dict is mutated. For a key that isn’t a valid Python identifier, use the self.replace(extras=...) form directly.

Parameters:: **kwargs – Entries to merge into extras. Values should be arrays (or pytrees of arrays) so the result stays batchable and jittable.
Returns:: A new AudioTree with the merged extras.
Return type:: AudioTree

Example

>>> audio = AudioTree.create(jnp.zeros((44100,)), 44100)
>>> audio = audio.replace_extras(tempo=np.array([120.0]))
>>> audio.extras["tempo"]
array([120.])

replace_lufs(lufs_window_sec: float = 0.4, lufs_hop_sec: float | None = None, *, device: Literal['cpu', 'gpu', 'tpu'] | Device | None = None, engine: Literal['numpy', 'jax'] | None = None) → Self¶

Compute and set the integrated and per-window loudness (LUFS).

Returns a new AudioTree with both lufs and lufs_windows populated:

lufs — the gated integrated loudness of each batch item, per the ITU-R BS.1770-4 standard (the standard “program loudness”). Measured on the CPU with the loudness library for NumPy waveforms, or a vmapped jaxloudnorm meter for JAX waveforms.
lufs_windows — the ungated K-weighted loudness of each window (a loudness-over-time curve). The signal is K-weighted once and each window reports the K-weighted mean-square of its samples in LUFS, so windows are directly comparable and a fully silent window is -inf. NumPy uses exact IIR biquads (scipy); JAX uses jaxloudnorm’s FIR-approximated filters on the accelerator.

The two engines are not bit-identical.

Where the measurement runs and which kernel runs are separate choices, so device= and engine= are separate arguments: asking for the exact IIR meter does not commit you to a device, and moving the work to an accelerator does not silently swap in the FIR approximation.

Parameters:

lufs_window_sec – Length in seconds of each lufs_windows window. Must be at least 0.4s (the EBU momentary integration time).
lufs_hop_sec – Step in seconds between window starts. Defaults to lufs_window_sec (non-overlapping windows); a smaller value overlaps them, but it must still span at least one sample at the tree’s sample rate. The trailing partial window is dropped, so an excerpt shorter than one window yields an empty lufs_windows.
device – Where to compute — an XLA platform name ("cpu" / "gpu" / "tpu", mirroring jax.jit’s backend) or a jax.Device. None (default) leaves the waveform where it is. device="gpu" is much faster for a large batch, since the NumPy lufs path measures one item at a time.
engine –
Which kernel to run. "numpy" is the exact ITU-R BS.1770 IIR meter (the loudness C++ library + scipy biquads); it is CPU-only, so it cannot be combined with a non-CPU device. "jax" is the vmapped jaxloudnorm kernel with FIR-approximated K-weighting, and runs on whatever device says. None (default) follows the waveform’s own array library — NumPy waveform to "numpy", JAX waveform to "jax" — except that a non-CPU device implies "jax", the only engine that can run there.

The returned lufs / lufs_windows always match the waveform’s array type (a NumPy waveform yields NumPy loudness whatever the engine), so you never need a manual jax.device_put / jax.device_get round-trip.

Returns:

AudioTree with lufs shaped (*batch,) and lufs_windows shaped (*batch, num_windows).

Note

Channel Limitations: Supports up to 5 channels:

Mono (1 channel): Single channel
Stereo (2 channels): [Left, Right]
5.0/5.1 Surround (5 channels): [Left, Right, Center, Left Surround, Right Surround]

Will raise ValueError if audio has more than 5 channels.

resample(sample_rate: int, *, zeros: int = 24, rolloff: float = 0.945, output_length: int | None = None, full: bool = False) → Self¶

Resample the AudioTree’s waveform to a new sample rate. NumPy-backed waveforms are resampled on CPU with librosa (soxr); JAX-backed waveforms use a JAX port of ResampleFrac from the PyTorch library Julius. The two backends are not bit-identical.

Parameters:

sample_rate (int) – The new sample rate of audio data, such as 44100 Hz.
zeros (int, optional) – number of zero crossing to keep in the sinc filter. JAX backend only.
rolloff (float) – use a lowpass filter that is rolloff * sample_rate / 2, to ensure sufficient margin due to the imperfection of the FIR filter used. Lowering this value will reduce antialiasing, but will reduce some of the highest frequencies. JAX backend only.
output_length (None or int) – This can be set to the desired output length (last dimension). Allowed values are between 0 and ceil(length * sample_rate / old_sr). When None (default) is specified, the floored output length will be used. In order to select the largest possible size, use the full argument.
full (bool) – return the longest possible output from the input. This can be useful if you chain resampling operations, and want to give the output_length only for the last one, while passing full=True to all the other ones. JAX backend only.

Changing the sample rate changes the audio’s length and its samples, so every derived field (lufs, lufs_windows, codes, latents) is invalidated. Resampling to the rate the tree already has returns it unchanged, derived fields and all.

Returns:: A new AudioTree resampled to sample_rate (the original is unchanged).
Return type:: AudioTree

Example

>>> audio = AudioTree.create(jnp.zeros((44100,)), 44100)  # 1 s at 44.1 kHz
>>> resampled = audio.resample(22050)
>>> resampled.waveform.shape
(1, 1, 22050)
>>> resampled.sample_rate
22050

reshape_mini_batches(mini_batch_size: int) → Self¶

Reshape batch dimension into mini-batches by adding a new leading axis.

Transforms audio data from shape (B, C, T) to (num_mini_batches, mini_batch_size, C, T), where B must be evenly divisible by mini_batch_size. String extras leaves nest the same way: a list of B strings becomes num_mini_batches lists of mini_batch_size strings, so indexing a mini-batch keeps them aligned with the arrays.

Parameters:: mini_batch_size – Number of samples per mini-batch. The total batch size must be evenly divisible by this value.
Returns:: AudioTree with an additional mini-batch dimension as the first axis.
Raises:: ValueError – If the batch size is not divisible by mini_batch_size, or if the tree is already mini-batched (rank 4) — nothing in this library reads the rank-5 tree that would produce.

Example

>>> x = AudioTree(np.zeros((12, 1, 44100)), 44100)
>>> x_batched = x.reshape_mini_batches(3)
>>> x_batched.waveform.shape  # 4 mini-batches of size 3
(4, 3, 1, 44100)

sample_rate: int¶

property samples: int¶: Return the number of samples in the waveform (its last dimension).

property source: List[str] | List[list]¶

Return the decoded source names stored in _metadata['source'].

Source names indicate which data source group each item in the batch came from. For example, if create_balanced_audio_dataset() was given sources={"music": [...], "speech": [...]}, this property might return ["music", "music", "speech", "music"] for a batch of 4 items. Like filepath, a mini-batched (rank-4) tree nests one list per mini-batch.

An empty list is returned if the AudioTree does not carry source provenance.

split(n_splits: int) → List[Self]¶

Split batch dimension into a list of smaller AudioTree objects.

Divides the batch dimension evenly into n_splits separate AudioTree objects, each containing a portion of the original batch. Like indexing, this works on the leading axis, which on a mini-batched (rank-4) tree is the mini-batch axis rather than the item axis.

Parameters:: n_splits – Number of AudioTree objects to create. The batch size must be evenly divisible by this value.
Returns:: List of AudioTree objects, each with batch_size = original_batch_size / n_splits.
Raises:: ValueError – If n_splits is not positive, or if the batch size is not divisible by n_splits.

Example

>>> audio = AudioTree(np.zeros((12, 1, 44100)), 44100)
>>> audio.waveform.shape
(12, 1, 44100)
>>> halves = audio.split(2)
>>> len(halves)
2
>>> halves[0].waveform.shape  # each half has half the original batch size
(6, 1, 44100)

to_mono(strategy: Literal['average', 'left', 'right'] = 'average') → Self¶

Reduce the waveform to mono.

Changing the channel count changes the audio, so every derived field (lufs, lufs_windows, codes, latents) is invalidated. A waveform that is already mono is returned unchanged, derived fields and all.

Parameters:: strategy – "average" mixes all channels down (default); "left" / "right" select the corresponding channel of a stereo waveform.
Returns:: An instance of AudioTree.
Return type:: AudioTree

to_stereo() → Self¶

Make the waveform stereo.

Changing the channel count changes the audio, so every derived field (lufs, lufs_windows, codes, latents) is invalidated. A waveform that is already stereo is returned unchanged, derived fields and all.

Returns:: An instance of AudioTree.
Return type:: AudioTree

velocity: ndarray | None = None¶

waveform: ndarray | None¶

write(filepath: str | Path, *, subtype: str | None = None, format: str | None = None, endian: str | None = None) → Path¶

Write the waveform to an audio file using soundfile.

This is the inverse of from_file(). The AudioTree must contain a single item (batch_size == 1); index or iterate the batch first (e.g. tree[0] or for item in tree) to write each item. The sample rate is taken from self.sample_rate — call resample() beforehand if you want a different one.

Parameters:

filepath – Output path. The file format is inferred from the extension (e.g. .wav, .flac, .ogg) unless overridden by format.
subtype – soundfile subtype, e.g., "PCM_16", "PCM_24", "FLOAT". When None (default) soundfile picks the format default (PCM_16 for WAV).
format – Major format override (e.g. "WAV", "FLAC"). When None it is inferred from the filepath extension.
endian – Endianness override (e.g. "FILE", "LITTLE", "BIG").

Returns:

The path that was written.

Return type:

Path

Raises:

ValueError – If the tree has no waveform, if it is mini-batched (rank 4 — batch_size then counts mini-batches, not items), or if batch_size != 1.

class ExcerptConfig(strategy: Literal['start', 'random', 'loudest'] = 'random', num_tries: int = 8, lufs_cutoff: float = -40.0, search: str | Callable = 'uniform', on_failure: Literal['keep', 'skip', 'raise'] = 'keep')¶

How to choose which part of a file an excerpt comes from.

One named strategy rather than a set of interacting flags:

"random": A random offset drawn by search – uniform by default (the default strategy). No audio is measured, so this costs one read per item.
"start": Always offset 0. Deterministic, and the right choice for validation sets where every epoch should see identical audio.
"loudest": Draw up to num_tries candidate offsets and keep the loudest, stopping early once one exceeds lufs_cutoff. Use it on corpora with long quiet stretches; it costs up to num_tries reads and loudness measurements per item.

Note that "loudest" is a best-of-num_tries search, not a filter: on a file where nothing clears the cutoff it still has to return something, and on_failure decides what.

Variables:

strategy (Literal['start', 'random', 'loudest']) – Which of the three above.
num_tries (int) – Maximum candidate offsets to try ("loudest" only).
lufs_cutoff (float) – Integrated loudness (LUFS) that ends the search early ("loudest" only).
search (str | Callable) – How each offset is drawn ("random" and "loudest"; not meaningful under "start") — a registered name ("uniform", "bias_early"), a dotted path to an importable function, or a callable.
on_failure (Literal['keep', 'skip', 'raise']) – What to do when no candidate clears lufs_cutoff ("loudest" only). "keep" returns the loudest excerpt found, "skip" returns None (grain drops it at to_iter_dataset()), "raise" raises naming the file.

Example

>>> from audiotree import ExcerptConfig
>>> ExcerptConfig().strategy                       # random offset
'random'
>>> ExcerptConfig(strategy="start").strategy       # deterministic
'start'
>>> loud = ExcerptConfig(
...     strategy="loudest", lufs_cutoff=-30, on_failure="skip"
... )
>>> loud.num_tries, loud.on_failure
(8, 'skip')

lufs_cutoff: float = -40.0¶

num_tries: int = 8¶

on_failure: Literal['keep', 'skip', 'raise'] = 'keep'¶

property resolved_search: Callable¶: The search spec as a callable.

search: str | Callable = 'uniform'¶

strategy: Literal['start', 'random', 'loudest'] = 'random'¶