Writing Datasets

AudioTree provides two writers for different use cases:

	AudioWriter	TreeWriter
Storage	Individual WAV files + NPZ manifest	Memory-mapped binary files + JSON manifest
Read speed	Decodes audio on each access	Zero-copy memmap slice
Human-readable	Yes (playable audio files)	No (raw binary)
Best for	Saving a trained model’s audio outputs to listen to or share	Pre-computing a dataset to train a model on

Use AudioWriter to save the outputs of a trained model — the generated or reconstructed audio you want to listen to, share with collaborators, or feed into non-Python tools. It writes standard WAV files and tracks per-sample data (loudness, latents, etc.) in an NPZ manifest that supports filtering.

Use TreeWriter to pre-compute data for training a model — render an augmented or feature-extracted dataset once, then read it back with no per-item decoding cost. It stores the entire pytree (AudioTree, dicts of AudioTrees, nested structures) as memory-mapped arrays — one .bin file per leaf, read as a zero-copy slice via TreeDataSource (a Grain RandomAccessDataSource).

from audiotree import TreeWriter
from audiotree.sources import TreeDataSource

# Write a dataset
with TreeWriter("dataset/", expected_samples=10000) as w:
    for batch in dataloader:
        w.write(batch)

# Read it back (Grain-compatible)
ds = TreeDataSource("dataset/")
sample = ds[0]  # reconstructed AudioTree

---

TreeWriter

TreeWriter stores each pytree leaf as its own memory-mapped .bin file and preserves the leaf’s dtype. That makes it a natural home for pre-computed features — a spectrogram, an embedding, a codec’s tokens — cached next to (or instead of) the waveform and read back later as a zero-copy memmap slice via TreeDataSource.

Note

Three named concepts, one word “metadata”. An AudioTree carries two per-item dicts that TreeWriter serializes as separate children, and the writer itself takes a third, unrelated metadata= parameter:

• extras — your per-item payload: arrays that batch and train with the audio (labels, embeddings, features). TreeWriter stores each entry as its own leaf under extras.<key>, and the namespace is entirely yours — the library plants no keys of its own there.

• AudioTree._metadata — the library’s per-item provenance: the encoded filepath and source arrays plus the offset array, recording where each item came from. The container is private (the underscore says so) — pass filepath= / source= / offset= to create() and read the .filepath, .source and .offset properties. TreeWriter stores it as a metadata node (the on-disk spelling) holding exactly those leaves; the schema is closed, so a metadata node containing any other key is rejected by name at read.

• TreeWriter(metadata=...) — true dataset-level metadata: free-form facts about the whole render (a description, a git commit, a source corpus), stored once under the manifest’s top-level "metadata" key and read back with get_metadata().

The split is what makes each name honest: extras is payload that trains with the audio, the per-item provenance container is data about the audio (where it came from), and the writer’s metadata= is data about the dataset.

Note

TreeDataSource validates the dataset when you construct it, not when you first read from it. Alongside the manifest’s declared shapes, dtypes and field names, every leaf’s .bin is measured against what the manifest claims, so a truncated render or a half-finished copy is refused by name:

ValueError: Invalid manifest dataset/manifest.json: leaf 'waveform'
declares 1000 samples of shape (1, 44100) and dtype float32 (176400000
bytes), but 'waveform.bin' is only 512 bytes. The dataset is truncated
or the manifest does not describe it.

The alternative is a bare np.memmap error at the first read — naming neither the leaf nor the manifest, and raised inside a Grain worker under the default lazy mode. A dataset that is still being written is not refused: a mid-write .bin is legitimately longer than num_samples implies, and only the prefix the manifest promises is checked.

expected_samples is a hint, not a cap

expected_samples says how many rows to pre-allocate. Under-shooting it is always safe: close() truncates every .bin file down to what was actually written. Over-shooting it is safe too — what happens is decided by the keyword-only on_overflow:

on_overflow	Behavior when a batch does not fit
"grow" (default)	Reallocate every leaf file and carry on. write() always returns the full batch size and never drops samples.
"error"	Raise ValueError naming how many samples were written, allocated and offered.
"trim"	Write as much of the batch as fits, drop the rest, and warn (UserWarning) with the same counts. This is the only policy under which write() returns less than the batch size — possibly 0, once the allocation is full.

import jax.numpy as jnp
from audiotree import AudioTree, TreeWriter

batch = AudioTree.create(jnp.zeros((8, 1, 16_000)), 16_000)

# A deliberate under-estimate: the default policy grows past it.
with TreeWriter("grown", expected_samples=2) as writer:
    written = writer.write(batch)
    stats = writer.get_stats()
    print(written, stats["expected_samples"], stats["allocated_samples"])

8 2 8

get_stats()["expected_samples"] keeps reporting the constructor argument verbatim; allocated_samples is the live allocation. The on-disk manifest is unchanged either way — it still records expected_samples as you passed it. Pass on_overflow="error" when the count is meant to be exact and a mismatch is a bug worth surfacing.

Quantizing float features to int16

Because TreeWriter keeps each leaf’s dtype, you can store a floating-point feature as int16 to halve its bytes on disk, then dequantize back to float32 in the data loader. The convention below maps the feature’s value range to [-1, 1] and then to the full int16 range [-32767, 32767]:

import numpy as np
import jax.numpy as jnp
import grain
from audiotree import AudioTree, TreeWriter
from audiotree.sources import TreeDataSource

# Stand in for a feature extractor: a batch of magnitude spectrograms of
# shape (batch, freq, frames) with values in [MAG_MIN, MAG_MAX].
MAG_MIN = 0.0
MAG_MAX = 4.0
n_items = 6
rng = np.random.default_rng(0)
spec = rng.uniform(MAG_MIN, MAG_MAX, size=(n_items, 128, 44)).astype(np.float32)

# 1. Remap the feature range [MAG_MIN, MAG_MAX] to [-1, 1].
spec_unit = (spec - MAG_MIN) / (MAG_MAX - MAG_MIN) * 2.0 - 1.0

# 2. Quantize [-1, 1] to the full int16 range for compact storage.
spec_i16 = np.round(spec_unit * 32767.0).astype(np.int16)

# Carry the int16 feature in the AudioTree's extras (a pytree node), so it
# batches and indexes alongside the waveform. TreeWriter keeps each leaf's
# dtype, so the spectrogram is written to disk as int16.
record = AudioTree.create(
    jnp.zeros((n_items, 1, 16_000)),
    16_000,
    extras={"spectrogram": spec_i16},
)
with TreeWriter("features", expected_samples=n_items) as writer:
    writer.write(record)


# In the data loader, dequantize extras["spectrogram"] back to float32.
def dequantize(audio):
    spec = audio.extras["spectrogram"].astype(np.float32) / 32767.0
    return audio.replace_extras(spectrogram=spec)


ds = grain.MapDataset.source(TreeDataSource("features")).map(dequantize)

item = ds[0]
print(item.extras["spectrogram"].dtype)
print(item.extras["spectrogram"].shape)  # the batch axis of 1 is added back
print(bool(np.all(np.abs(item.extras["spectrogram"]) <= 1.0)))

float32
(1, 128, 44)
True

The round-trip is lossy only to int16 precision (about 1 / 32767), which is negligible for most spectrogram and embedding features while cutting storage in half versus float32.

Feature-only trees and selective loading

On a large pre-rendered corpus the raw waveform often dwarfs the features you actually train on. Two knobs keep such datasets cheap:

• Drop the waveform at write time. AudioTree.replace(waveform=None) yields a feature-only tree — codes, latents, or an extras feature with no audio. TreeWriter simply omits the missing leaf, and it reads back as None.

• Skip leaves at read time. TreeDataSource accepts exclude_prefixes (dot-separated leaf-name prefixes) to avoid ever reading a memmap you don’t need — e.g., loading only mels for one training run and only audio for another, from the same directory.

import jax.numpy as jnp
import numpy as np
from audiotree import AudioTree, TreeWriter
from audiotree.sources import TreeDataSource

n = 8
# A {dry, wet} dataset: dry keeps its audio; wet keeps only a mel feature and
# drops its waveform to save disk.
dry = AudioTree.create(jnp.zeros((n, 1, 16_000)), 16_000)
wet = AudioTree.create(
    jnp.zeros((n, 1, 16_000)),
    16_000,
    extras={"mel": np.zeros((n, 80, 32), np.float32)},
).replace(waveform=None)  # feature-only: no audio is written

with TreeWriter("prerendered", expected_samples=n) as writer:
    writer.write({"dry": dry, "wet": wet})

# Full read: wet.waveform was never written, so it comes back None.
item = TreeDataSource("prerendered")[0]
print(item["wet"].waveform is None, item["wet"].extras["mel"].shape)

# Lean read: also skip the dry audio memmap (huge on a real corpus).
lean = TreeDataSource("prerendered", exclude_prefixes=["dry.waveform"])[0]
print(lean["dry"].waveform is None)

True (1, 80, 32)
True

Tip

Pass load_into_memory=True to TreeDataSource to read every (non-excluded) leaf into RAM once at construction. With fork-based multiprocessing (the default on Linux) Grain workers then inherit that data via copy-on-write instead of each re-opening the memmaps — trading memory for zero per-worker I/O. Combine it with exclude_prefixes so only the leaves you train on are held in memory. Each ds[i] still hands back its own copy of the arrays it reads, exactly as the memmap path does, so nothing downstream aliases the shared store.

Note

TreeDataSource keeps one memmap per leaf open for the life of the process — that is what makes random reads fast — and a mapped file cannot be deleted, moved, or replaced on Windows while it is open. For a training run that never matters, but if you render a dataset, read it, and then clean it up in the same process, release the handles first:

with TreeDataSource("dataset/") as ds:
    tree = ds[0]
# handles released here, so the directory can now be removed

close() does the same without the with. It is a release rather than a teardown: reading again reopens transparently, and calling it twice is fine.

Example: Creating a Training Dataset

Here’s a complete example of creating a training dataset with TreeWriter:

from tqdm import tqdm
from audiotree import TreeWriter
from audiotree.sources import create_audio_dataset, TreeDataSource
from audiotree.transforms import volume_change, shift_phase


def precompute_training_dataset(
    source_directory, num_samples, directory="precomputed_data"
):
    """Render an augmented dataset once so training never recomputes it."""

    # An infinite, shuffled stream of 3-second mono excerpts.
    ds = create_audio_dataset(
        sources=source_directory,
        sample_rate=16_000,
        duration=3.0,
        mono=True,
        shuffle=True,
        num_epochs=None,
    )

    # Seed once; each random_map derives its own distinct seed so every
    # augmentation differs.
    ds = ds.seed(42)
    ds = ds.random_map(volume_change(min_db=-6, max_db=6))
    ds = ds.random_map(shift_phase())

    # Take num_samples items from the infinite stream and write each one.
    # TreeWriter pre-allocates expected_samples rows up front, and grows
    # them if the stream turns out to be longer.
    it = iter(ds.to_iter_dataset())
    pbar = tqdm(total=num_samples, desc="Precomputing")
    with TreeWriter(
        output_dir, expected_samples=num_samples, pbar=pbar, close_pbar=True
    ) as writer:
        for _ in range(num_samples):
            writer.write(next(it))

    return output_dir


# Build the dataset once...
precompute_training_dataset("/data/audio", num_samples=10_000)

# ...then train from it with no augmentation or decoding cost.
train_ds = TreeDataSource("precomputed_data")

To pre-compute input/target pairs — say an augmented "input" beside the clean "target" — augment a {"input": ..., "target": ...} dict and restrict each transform to one key with the scope parameter; see Working with Dict[str, AudioTree] Batches.

---

AudioWriter

AudioWriter writes each item of an AudioTree batch as its own audio file and records the per-item columns (loudness, provenance, tags, extras) in a manifest alongside them.

Basic Usage

Batches written in sequence get sequential file names, and the manifest is saved when the writer closes:

from audiotree import AudioTree, AudioWriter
import numpy as np

# Create an AudioTree with 3 samples
audio = AudioTree.create(
    np.random.randn(3, 2, 44_100),  # 3 batches, stereo, 1 second
    sample_rate=44_100,
)

# Write to disk with automatic manifest
with AudioWriter("output", pattern="audio_{index:04d}.wav") as writer:
    paths = writer.write(audio)

print(len(paths))  # One file per batch item
print(paths[0].name)

3
audio_0000.wav

The Manifest

The manifest is a single NPZ file: arrays load back directly, numeric types survive exactly, and it stays compact even for datasets with thousands of files.

By default AudioWriter refuses to write into a directory that already holds a manifest, so a finished dataset is never silently overwritten (and two writers aimed at one directory are caught). The examples below reuse "output" and so pass exist_ok=True to opt in.

with AudioWriter("output", exist_ok=True) as writer:
    writer.write(audio)
# Creates output/manifest.npz

The NPZ is compressed by default; compress_manifest=False trades larger files for slightly faster writes, but the savings from compression (often 5–10x) are usually worth the small cost:

# Compressed NPZ (default) - smaller files, slightly slower writing
writer = AudioWriter("output", compress_manifest=True, exist_ok=True)

# Uncompressed NPZ - faster writing, larger files
writer = AudioWriter("output", compress_manifest=False, exist_ok=True)

Manifest Tracking

AudioWriter records every filled per-item field in the manifest automatically:

# Create AudioTree with several per-item fields filled in
meta_tree = AudioTree.create(
    np.random.randn(2, 1, 44_100),
    sample_rate=44_100,
    lufs=np.array([-20.0, -15.0]),
    pitch=np.array([60.0, 62.0]),
    velocity=np.array([64, 80]),
    note_duration=np.array([1.0, 0.5]),
    filepath=["original1.wav", "original2.wav"],
)

# Write with custom tags
with AudioWriter("output_meta") as writer:
    writer.write(meta_tree, tags={"dataset": "train", "version": 2})

The manifest will contain:

• Core info: index, filename, sample_rate, channels, samples, duration_seconds, files_written, and subtype — the soundfile subtype each item was actually encoded with (absent for a write_audio=False run, where nothing was encoded)

• AudioTree fields: lufs, lufs_windows, pitch, velocity, note_duration, codes, and latents

• Provenance columns: filepath, source and offset, dedicated bookkeeping columns recording where each item came from, mirroring the .filepath / .source / .offset properties

• Custom tags: Any additional metadata passed via the tags parameter, stored as tags_* columns

• Extras arrays: every non-nested extras entry, as an extras_* column — this namespace is purely user payload; the library plants no columns of its own there

A column that some rows lack is stored with a presence mask, so a missing value is genuinely absent on read-back rather than standing in as a sentinel like -1, NaN or "". That matters for a filter_fn: use entry.get(...) with your own default rather than entry[...] for any column that may be absent ("tags" included — it appears only on rows that have at least one tag).

The column set is closed: readers validate every column of a manifest against this schema and reject an unrecognized column by name, so a manifest that has drifted from the format fails loudly at read rather than being partially and silently ignored.

Choosing an encoding subtype

AudioWriter(subtype=...) is passed straight to soundfile. Left at its default of None, AudioWriter picks the widest subtype the container supports — FLOAT for WAV/AIFF/CAF/W64/RF64, PCM_24 for FLAC, the container’s own default otherwise — rather than libsndfile’s default of PCM_16 for WAV. Model output routinely exceeds [-1, 1], and a 16-bit default silently hard-clipped it.

# Default: lossless for a model's float output, at the cost of file size.
writer = AudioWriter("output_float", exist_ok=True)

# Ask for the old behavior explicitly when you want small, portable files.
writer = AudioWriter("output_pcm16", subtype="PCM_16", exist_ok=True)

If the effective subtype is fixed-point and an item peaks above 1.0, write() emits a RuntimeWarning naming the worst offender and its peak — one warning per write() call — because those samples do not survive the encode. Either pass subtype="FLOAT" or scale the audio down first (rescale_audio() does exactly that).

Timestamps

include_timestamp=True records when each file was written. That is worth having when you need to reconstruct how a dataset was produced, at the cost of a manifest that differs between otherwise identical runs. It is off by default:

# Without timestamps (default) - smaller, run-to-run identical manifests
writer = AudioWriter("output", include_timestamp=False, exist_ok=True)

# With timestamps - track when files were written
writer = AudioWriter("output", include_timestamp=True, exist_ok=True)

Sequential Writing

AudioWriter maintains state for sequential writing across multiple batches:

writer = AudioWriter("output_seq", pattern="sample_{index:05d}.wav")

# Write first batch
tree1 = AudioTree.create(np.random.randn(2, 1, 8000), 8000)
paths1 = writer.write(tree1)
print([p.name for p in paths1])

# Write second batch - indexing continues
tree2 = AudioTree.create(np.random.randn(3, 1, 8000), 8000)
paths2 = writer.write(tree2)
print([p.name for p in paths2])

# Get statistics
stats = writer.get_stats()
print(stats["total_files"])
print(stats["current_index"])

# Save manifest when done
manifest_path = writer.save_manifest()

['sample_00000.wav', 'sample_00001.wav']
['sample_00002.wav', 'sample_00003.wav', 'sample_00004.wav']
5
5

Progress Bars

For long renders, show_progress=True gives the writer its own tqdm bar (tqdm ships with the progress extra):

with AudioWriter(
    "output", show_progress=True, progress_desc="Writing audio", exist_ok=True
) as writer:
    for audio in audio_batches:
        writer.write(audio)
# Progress bar automatically closed

Or hand it a bar you configured yourself via pbar=, say one whose total you computed up front. Pass close_pbar=True if the writer should close it on exit:

from tqdm import tqdm

pbar = tqdm(total=1000, desc="Generating dataset")

with AudioWriter("output", pbar=pbar, close_pbar=True, exist_ok=True) as writer:
    for audio in full_audio.split(batch_size):
        if audio.lufs > -30:  # Only write loud samples
            writer.write(audio)  # advances pbar by the batch size

Either way the bar advances by the batch size of each write, so it stays accurate when batches vary in size.

Pattern Formatting

The pattern parameter supports Python string formatting:

# Zero-padded indices
pattern = "audio_{index:04d}.wav"  # audio_0000.wav, audio_0001.wav, ...

# Different padding
pattern = "sample_{index:06d}.wav"  # sample_000000.wav, sample_000001.wav, ...

# Custom prefixes
pattern = "train_{index:05d}.wav"  # train_00000.wav, train_00001.wav, ...

Reading Written Data

Use AudioDataSource to read AudioWriter output:

from audiotree.sources import AudioDataSource

# Write some data (with per-item loudness so it round-trips into the manifest)
loudness_tree = AudioTree.create(
    np.random.randn(3, 2, 44_100),
    sample_rate=44_100,
    lufs=np.array([-20.0, -15.0, -18.0]),
)
with AudioWriter("output_read") as writer:
    writer.write(loudness_tree, tags={"split": "train"})

# Read it back
source = AudioDataSource("output_read")

# Access individual items
audio = source[0]
print(audio.sample_rate)
print(audio.lufs)  # Manifest fields are restored

# Filter by manifest columns
loud_source = source.filter_by_lufs(min_lufs=-18.0)

# Filter by tags
train_source = source.filter_by_tag("split", "train")

44100
[-20.]

The restored lufs, lufs_windows, codes and latents describe the audio as it was written, so AudioDataSource restores them only when its sample_rate/mono/duration arguments leave the audio unchanged. Reading with, say, mono=True or a different sample_rate leaves those fields unset — call replace_lufs() to measure the audio actually returned. Annotation fields (pitch, velocity, note_duration) are restored either way.

Per-item extras make the same round trip. They survive transformations on the way in and land in the manifest as extras_* columns:

from audiotree import AudioTree, AudioWriter
import numpy as np

# Load with a few extras attached
audio = AudioTree.from_file(
    "input.wav",
    sample_rate=44_100,
    extras={"instrument": "guitar", "style": "rock", "bpm": 120, "key": "A minor"},
)

# Extras are preserved through transformations
processed = audio.resample(16_000)
processed = processed.replace_lufs()

# Check the extras are still there
print(processed.extras["instrument"])
print(processed.extras["bpm"])

# Write with additional tags
with AudioWriter("output_flow") as writer:
    writer.write(processed, tags={"processed": True, "version": 2})

guitar
120

When read back via AudioDataSource, the per-item extras are restored as batched arrays (so a single-item read gives array(['guitar']) for a string field):

from audiotree.sources import AudioDataSource

source = AudioDataSource("output_flow")
loaded = source[0]
print(loaded.extras["instrument"])

['guitar']

Manifest-Only: Saving Embeddings (No Audio)

Sometimes the payload you want to save is not audio but a per-item array a model produced — an embedding, a projection, a set of predicted parameters. Pass write_audio=False to write only the NPZ manifest, with your arrays carried in extras. No WAV files are written, so a large evaluation set of embeddings costs almost nothing on disk:

import os
import numpy as np
from audiotree import AudioTree, AudioWriter

# A batch of clips, each with an embedding a model produced. Carry the
# embeddings (and any ids you need) in extras — an active pytree node.
rng = np.random.default_rng(0)
n_items = 100
batch = AudioTree.create(
    rng.standard_normal((n_items, 1, 16_000)).astype(np.float32),
    16_000,
    extras={
        "embedding": rng.standard_normal((n_items, 128)).astype(np.float32),
        "label": np.arange(n_items),
    },
)

# write_audio=False writes manifest.npz only (no per-item WAVs).
with AudioWriter("embeddings", write_audio=False) as writer:
    writer.write(batch)

print(sorted(os.listdir("embeddings")))

['manifest.npz']

To read the set back for analysis, from_manifest() loads the entire manifest into a single batched AudioTree — so every item’s embedding lands in one array rather than a stream. An optional filter_fn predicate (evaluated per manifest entry) selects a subset at load time:

# The whole manifest as one tree; extras arrays round-trip exactly.
audio = AudioTree.from_manifest("embeddings/manifest.npz")
print(audio.extras["embedding"].shape)

# filter_fn sees each entry's columns as ``extras_<key>``; keep labels < 10.
subset = AudioTree.from_manifest(
    "embeddings/manifest.npz",
    filter_fn=lambda entry: entry["extras_label"] < 10,
)
print(subset.extras["embedding"].shape)

(100, 128)
(10, 128)

Note

Two readers, two shapes. from_manifest() returns one batched AudioTree with the whole manifest stacked along the batch axis — ideal for a one-shot analysis pass over saved embeddings. AudioDataSource (above) is instead a Grain RandomAccessDataSource that yields one item at a time in manifest order, for feeding a pipeline. from_manifest restores the extras_* arrays, the label fields (lufs, pitch, codes, …), and the filepath / source provenance columns — so loaded.filepath and loaded.source match what you wrote.

On-Disk Format Versioning

Pre-rendered datasets outlive the code that wrote them, so all three on-disk formats — a TreeWriter directory, an AudioWriter NPZ manifest, and the windowed-LUFS cache — carry a header (format, format_version, min_reader_version, producer), and every reader validates it:

• format names the artifact (audiotree-tree, audiotree-manifest, audiotree-lufs-windows), so pointing a reader at the wrong kind of directory fails by name instead of as a KeyError.

• A format_version major mismatch is refused in both directions; a minor difference is accepted. Minors are additive — new fields a reader may ignore — so a dataset written by any 1.x audiotree reads on any other 1.x audiotree.

• min_reader_version is what a writer raises instead of a major bump when it adds a field readers must honor: an older reader refuses rather than silently ignoring it.

• An artifact with no header was written by a pre-1.0 development build and is refused with a “re-render the dataset” message.

tests/assets/golden/ holds fixtures written once and committed, read back in tests/test_golden_formats.py against hardcoded values — so a change that would repack every dataset on disk goes red even though the ordinary tests, which write their inputs with the code under test, would stay green.

Next

That completes the Getting-started path — you can load audio, augment it, and write prepared datasets back to disk. From here, the Going further guides dig into Balanced Datasets, Windowed (Length-Aware) Datasets, Working with Dict[str, AudioTree] Batches, command-line configuration with Using ArgBind with Transforms, and Multiprocessing and Multithreading.