Starting with hrtfpykit.datasets

hrtfpykit.datasets builds map-style datasets for HRTF workflows. A dataset object exposes len(dataset) and dataset[index], but it also keeps the HRTF dataset rules in one place: subject identifiers, file layouts, official resource variants, excluded subjects, deterministic splits, cache behavior, SOFA loading controls, and resource summaries.

This tutorial uses SONICOM because it provides a reproducible public HRTF slice. The explanations are general for the dataset layer: the same ideas apply to ARI, HUTUBS, SONICOM, and future dataset classes. The SONICOM-specific parts are the concrete subject IDs, HRTF variant, and file layout used by the downloader.

By the end, you will have downloaded a small resource slice, built acoustic datasets with specs, inspected samples, controlled SOFA loading, separated download filters from dataset filters, reused local files, managed the HRTF cache, created train/validation/test splits, batched samples for PyTorch, computed HRTF training losses with hrtf_loss, and run a small autoencoder-shaped training loop where the input HRTF representation is also the reconstruction target.

Download a small SONICOM resource slice

Dataset downloading prepares official files under the local dataset root. Dataset construction scans local files and builds the indexed samples returned by dataset[index]. hrtfpykit keeps these steps separate so a workflow can download resources once, reuse files copied manually, and build different dataset views from the same root.

download_server selects the official source used for the transfer. For SONICOM, "imperial" provides direct configured metadata, HRTF, and mesh files, while "sonicom-ecosystem" uses SONICOM ecosystem database catalogs for HRTF and mesh files. download_resources selects the resource families to fetch, such as "hrtf", "mesh", or "metadata". Variant arguments such as download_hrtf_variant select the official files to download. Later, dataset arguments such as inputs, target, dataset_hrtf_variant, subject_ids, split, and split_seed decide how local files become indexed samples.

This example downloads only measured SONICOM HRTFs for subjects P0001 through P0010, using the 44.1 kHz FreeFieldComp variant from the "imperial" server. Keeping the slice small makes the tutorial quick while still exercising the normal downloader and dataset construction path.

from pathlib import Path

from hrtfpykit.datasets import SONICOM

# Define a local dataset root for this tutorial.
root = Path("datasets/sonicom")

# Keep only P0001 through P0010 for a small reproducible tutorial slice.
selected_subject_ids = tuple(f"P{i:04d}" for i in range(1, 11))

# Select the measured 44.1 kHz FreeFieldComp HRTF files.
hrtf_variant = {
    "type": "measured",
    "sample_rate": 44100,
    "version": "FreeFieldComp",
}

# Download only the selected HRTF resource slice.
SONICOM(
    root=root,
    download=True,
    download_resources="hrtf",
    download_hrtf_variant=hrtf_variant,
    download_server="imperial",
    dataset_hrtf_variant=hrtf_variant,
    download_subject_ids=selected_subject_ids,
    subject_ids=selected_subject_ids,
    verify_checksum=True,
    verbose=True,
)

print("Selected subjects:", selected_subject_ids)

Separate download filters from dataset filters

Download filters decide which official files are transferred into root. Dataset filters decide which local subjects are scanned and used for samples. Keeping the two scopes separate is useful when a machine stores a larger local archive than a single experiment needs.

download_subject_ids and download_exclude_subject_ids affect only the download request. They do not force those subjects into the constructed dataset. subject_ids and exclude_subject_ids affect dataset construction after local files are present. In the download cell above, both scopes use the same ten subjects so the tutorial stays easy to follow.

# Download scope: which official files should be fetched.
download_subject_ids = selected_subject_ids
download_exclude_subject_ids = ()

# Dataset scope: which local subjects should be used for this experiment.
dataset_subject_ids = selected_subject_ids
dataset_exclude_subject_ids = ()

print("Download subjects:", download_subject_ids)
print("Download exclusions:", download_exclude_subject_ids)
print("Dataset subjects:", dataset_subject_ids)
print("Dataset exclusions:", dataset_exclude_subject_ids)

Build a map-style acoustic dataset

Samples are declared with specs. A spec is a small configuration object that says which value should appear in sample["inputs"] or sample["target"]. Acoustic specs request values derived from the loaded HRTF object, such as HRIR arrays, HRTF magnitudes, ITD, ILD, or spherical harmonic coefficients. Resource specs align subject resources such as meshes, metadata, anthropometry, images, or videos.

This first tutorial stays shallow on specs because specs deserve their own dedicated tutorial. The important idea here is that specs are declarative: they describe the sample values, while the dataset object handles subjects, local files, variants, splits, and indexing.

The example below builds one subject-level sample per SONICOM subject. The input and target both use the same HRTF magnitude representation. That is useful for an autoencoder-style example later, where the model tries to reconstruct the same HRTF representation and the latent vector can be treated as a compact feature representation.

from hrtfpykit.datasets import HRTFSpec

input_spec = HRTFSpec(
    domain="frequency",
    signal="tf_magnitude_db",
    ears="both",
    index_by=("subject",),
    name="input_hrtf",
)

target_spec = HRTFSpec(
    domain="frequency",
    signal="tf_magnitude_db",
    ears="both",
    index_by=("subject",),
    name="target_hrtf",
)

dataset = SONICOM(
    root=root,
    download=False,
    dataset_hrtf_variant=hrtf_variant,
    subject_ids=selected_subject_ids,
    inputs=input_spec,
    target=target_spec,
    split="all",
    verbose=True,
)

print("Number of dataset rows:", len(dataset))

Build from existing local files

After files are already under root, dataset construction does not need the downloader. Set download=False and choose the local variant with dataset_hrtf_variant. This is the normal path for repeated experiments, offline runs, shared compute nodes, and manually prepared dataset folders.

For SONICOM measured HRTFs, hrtfpykit accepts the official layout used above, such as P0001/HRTF/HRTF/44kHz/P0001_FreeFieldComp_44kHz.sofa. It also accepts semantic layouts such as P0001/hrtf/measured/44kHz/P0001_FreeFieldComp_44kHz.sofa and P0001/hrtf/measured/FreeFieldComp/44kHz/P0001_FreeFieldComp_44kHz.sofa.

local_file_dataset = SONICOM(
    root=root,
    download=False,
    dataset_hrtf_variant=hrtf_variant,
    subject_ids=("P0001",),
    inputs=input_spec,
    target=target_spec,
    verbose=False,
)

local_file_sample = local_file_dataset[0]
print("Rows from existing local files:", len(local_file_dataset))
print("Local-file subject:", local_file_sample["meta"]["subject_id"])
print("Local-file input shape:", local_file_sample["inputs"]["input_hrtf"].shape)

Control HRTF loading

Dataset acoustic specs load subject files through load_hrtf. Two constructor parameters control how that happens.

check_sofa_against_conventions=False keeps repeated dataset construction quiet for known files with accepted custom SOFA fields. Set it to True when auditing a new dataset, validating incoming files, or debugging file metadata.

sofa_open=False closes the backing netCDF4 handle after IR, TF, and source positions are loaded into the HRTF object. This is the normal training setting because it avoids holding many open file handles. Use sofa_open=True when code needs direct SOFA metadata access from cached HRTF objects, such as inspecting dimensions, variables, or global attributes after get_subject_hrtf(...).

closed_sofa_dataset = SONICOM(
    root=root,
    download=False,
    dataset_hrtf_variant=hrtf_variant,
    subject_ids=("P0001",),
    inputs=input_spec,
    target=target_spec,
    check_sofa_against_conventions=False,
    sofa_open=False,
    verbose=False,
)

closed_subject_hrtf = closed_sofa_dataset.get_subject_hrtf("P0001")
print("Closed-handle HRTF shape:", closed_subject_hrtf.IR.values.shape)
print("SOFA handle open:", closed_subject_hrtf.Sofa.is_open())
closed_sofa_dataset.clear_cache()

open_sofa_dataset = SONICOM(
    root=root,
    download=False,
    dataset_hrtf_variant=hrtf_variant,
    subject_ids=("P0001",),
    inputs=input_spec,
    target=target_spec,
    check_sofa_against_conventions=False,
    sofa_open=True,
    verbose=False,
)

open_subject_hrtf = open_sofa_dataset.get_subject_hrtf("P0001")
print("Open-handle HRTF shape:", open_subject_hrtf.IR.values.shape)
print("SOFA handle open:", open_subject_hrtf.Sofa.is_open())
print("SOFA convention:", open_subject_hrtf.Sofa.GlobalAttributes.get("SOFAConventions").value)
open_sofa_dataset.clear_cache()

Filter local subjects

The same downloaded root can support several experiments. The example below asks the local scanner to consider subjects P0001 through P0005, then removes P0005 from the constructed dataset. No files are downloaded in this step; only local files under root are scanned.

filtered_dataset = SONICOM(
    root=root,
    download=False,
    dataset_hrtf_variant=hrtf_variant,
    subject_ids=selected_subject_ids[:5],
    exclude_subject_ids=("P0005",),
    inputs=input_spec,
    target=target_spec,
    verbose=False,
)

filtered_subjects = [
    filtered_dataset[index]["meta"]["subject_id"]
    for index in range(len(filtered_dataset))
]

print("Filtered rows:", len(filtered_dataset))
print("Filtered subjects:", filtered_subjects)

Inspect dataset samples

A hrtfpykit dataset is map-style: integer indexing returns one sample dictionary. Each sample has three top-level entries.

sample["inputs"] contains the values requested by input specs. sample["target"] contains the values requested by target specs. sample["meta"] records where the row came from, including the dataset name, subject ID, and any active row context such as position, ear, frequency, or sample index.

sample["meta"] is not model input by default; it is provenance. It is useful when debugging sample construction, checking that a split contains the expected subjects, tracing a suspicious value back to a source position or ear, and analyzing predictions after a model has produced outputs. It also matters when datasets are combined: the metadata keeps the dataset name and row context attached to each sample, so a script can still know where a sample came from after concatenation or batching.

With index_by=("subject",), each row represents one subject and the HRTF arrays keep their position, ear, and frequency axes. If a later spec uses index_by=("subject", "position"), the dataset rows expand so each row represents one subject and one source position.

sample = dataset[0]

print("Sample keys:", sample.keys())
print("Input keys:", sample["inputs"].keys())
print("Target keys:", sample["target"].keys())
print("Meta:", sample["meta"])

input_hrtf = sample["inputs"]["input_hrtf"]
target_hrtf = sample["target"]["target_hrtf"]

print("Input type:", type(input_hrtf).__name__)
print("Input shape:", input_hrtf.shape)
print("Input dtype:", input_hrtf.dtype)
print("Target shape:", target_hrtf.shape)
print("Input and target shapes match:", input_hrtf.shape == target_hrtf.shape)

for index in range(3):
    sample = dataset[index]
    print(index, sample["meta"]["subject_id"], sample["inputs"]["input_hrtf"].shape)

Add ITD and ILD specs

Acoustic specs can request derived binaural cues instead of raw IR or TF arrays. ITDSpec forwards its absolute and output settings to the ITD metric. With output="time", ITD values are returned in microseconds. ILDSpec returns ILD in dB and does not expose an output parameter.

The next dataset returns one subject-level row per subject with an absolute ITD vector and a signed broad-band ILD vector over the same source positions.

from hrtfpykit.datasets import ILDSpec, ITDSpec

cue_dataset = SONICOM(
    root=root,
    download=False,
    dataset_hrtf_variant=hrtf_variant,
    subject_ids=selected_subject_ids[:3],
    inputs=[
        ITDSpec(
            absolute=True,
            output="time",
            index_by=("subject",),
            name="absolute_itd",
        ),
        ILDSpec(
            mode="broad-band",
            index_by=("subject",),
            name="signed_ild",
        ),
    ],
    verbose=False,
)

cue_sample = cue_dataset[0]
print("Cue subject:", cue_sample["meta"]["subject_id"])
print("Absolute ITD shape:", cue_sample["inputs"]["absolute_itd"].shape)
print("Signed ILD shape:", cue_sample["inputs"]["signed_ild"].shape)
print("First absolute ITD value in us:", round(float(cue_sample["inputs"]["absolute_itd"][0]), 3))
print("First signed ILD value in dB:", round(float(cue_sample["inputs"]["signed_ild"][0]), 3))

Create deterministic subject splits

Splits are part of dataset construction. hrtfpykit first determines which subjects have all resources required by the selected specs, then applies the requested split. This means train, validation, and test objects can be built from the same root and the same specs without duplicating resource matching logic in user code.

The split is deterministic for a fixed split_seed. The example below keeps the same ten SONICOM subjects and divides them into train, validation, and test datasets.

dataset_kwargs = {
    "root": root,
    "download": False,
    "dataset_hrtf_variant": hrtf_variant,
    "subject_ids": selected_subject_ids,
    "inputs": input_spec,
    "target": target_spec,
    "split_ratio": (0.7, 0.2, 0.1),
    "split_seed": 42,
}

train_dataset = SONICOM(**dataset_kwargs, split="train")
validation_dataset = SONICOM(**dataset_kwargs, split="validation")
test_dataset = SONICOM(**dataset_kwargs, split="test")

print("Train samples:", len(train_dataset))
print("Validation samples:", len(validation_dataset))
print("Test samples:", len(test_dataset))
print("First train subject:", train_dataset[0]["meta"]["subject_id"])

Apply dataset-level HRTF preprocessing

dataset_hrtf_transform applies one HRTF transform to every loaded subject before any acoustic spec extracts values. This is useful when a whole dataset should use the same preprocessing: selecting a source subset, changing the time window, padding HRIRs, changing FFT length, converting to another HRTF representation, or applying another HRTF operation.

Spec-level transforms also exist, but they are better covered in the Transform HRTFs before extraction section of the specs tutorial. For a first dataset workflow, the dataset-level transform is easier to reason about: every acoustic spec sees the same transformed HRTF object.

The next dataset keeps only three named source directions before extracting the HRTF magnitude. This makes the batched tensor small enough for a lightweight autoencoder example.

from hrtfpykit.datasets import HRTFTransform

compact_transform = HRTFTransform.select(positions=["front", "left", "right"])

compact_dataset_kwargs = {
    "root": root,
    "download": False,
    "dataset_hrtf_variant": hrtf_variant,
    "dataset_hrtf_transform": compact_transform,
    "subject_ids": selected_subject_ids,
    "inputs": input_spec,
    "target": target_spec,
    "split_ratio": (0.7, 0.2, 0.1),
    "split_seed": 42,
}

compact_train_dataset = SONICOM(**compact_dataset_kwargs, split="train")
compact_validation_dataset = SONICOM(**compact_dataset_kwargs, split="validation")

compact_sample = compact_train_dataset[0]
print("Compact input shape:", compact_sample["inputs"]["input_hrtf"].shape)
print("Compact target shape:", compact_sample["target"]["target_hrtf"].shape)

Concatenate compatible datasets

Generated hrtfpykit datasets can be concatenated with PyTorch’s ConcatDataset when they expose the same sample structure. This can combine several dataset objects from the same public dataset, such as different SONICOM splits, or different public dataset classes, such as SONICOM, HUTUBS, and ARI, when they have been prepared to return compatible samples.

The main constraint is compatibility. If a training loop expects batch["inputs"]["input_hrtf"] and batch["target"]["target_hrtf"], every dataset in the concatenation should return those keys. The returned values should also have compatible shapes and numeric types if they are going to be stacked by collate_samples. For acoustic HRTF values, that usually means matching domains, selected positions, ears, frequency bins or sample lengths, and preprocessing transforms. When combining different public HRTF datasets, pay special attention to source grids, coordinate conventions, sample rates, FFT lengths, available ears, HRTF variants, subject resources, and any custom transforms. If the datasets do not naturally agree, select, transform, resample, interpolate, or otherwise map them into a common representation before using them as one training dataset.

The example below concatenates the compact train and validation objects only to show the mechanics. In a real experiment, keep validation and test datasets separate when measuring model performance.

from torch.utils.data import ConcatDataset

combined_dataset = ConcatDataset([
    compact_train_dataset,
    compact_validation_dataset,
])

combined_sample = combined_dataset[0]

print("Combined rows:", len(combined_dataset))
print("Combined sample keys:", combined_sample.keys())
print("Combined input shape:", combined_sample["inputs"]["input_hrtf"].shape)
print("Combined target shape:", combined_sample["target"]["target_hrtf"].shape)
print("Combined sample metadata:", combined_sample["meta"])

Batch samples for PyTorch

PyTorch training code normally consumes batches, not individual samples. torch.utils.data.DataLoader is the standard PyTorch object that repeatedly calls a map-style dataset with integer indices, collects several returned samples, and yields one batch at a time.

A hrtfpykit dataset already provides the map-style part: dataset[index] returns one dictionary with inputs, target, and meta. The missing step is collation. A data loader receives a list of individual hrtfpykit samples, but it does not automatically know how every nested sample value should be stacked, preserved, or converted. That is why hrtfpykit provides collate_samples.

collate_samples lives in hrtfpykit.datasets.torch and is the bridge between hrtfpykit samples and PyTorch batches. It recursively collates the sample dictionaries, stacks compatible NumPy arrays and tensors along a new batch axis, converts homogeneous numeric values to PyTorch tensors, and keeps non numeric or ragged values as Python lists. Floating point arrays are converted to torch.float32, so normal training loops can move values to the selected device directly.

Dataset indexing stays flexible and framework neutral. Tensor conversion happens at collation time, which is the point where PyTorch needs a batch.

from torch.utils.data import DataLoader

from hrtfpykit.datasets.torch import collate_samples, hrtf_loss

train_loader = DataLoader(
    compact_train_dataset,
    batch_size=4,
    shuffle=True,
    collate_fn=collate_samples,
)

batch = next(iter(train_loader))

print("Batch input shape:", batch["inputs"]["input_hrtf"].shape)
print("Batch input dtype:", batch["inputs"]["input_hrtf"].dtype)
print("Batch target shape:", batch["target"]["target_hrtf"].shape)
print("Batch meta keys:", batch["meta"].keys())

Compute HRTF training losses

hrtf_loss is the scalar PyTorch loss used when a model output and target are HRTF or HRIR tensors. It reads ordinary tensors from the batch, reduces the final acoustic axis, then reduces the remaining batch, position, and ear axes into one value that can be used with loss.backward().

The current dataset returns tf_magnitude_db, so each tensor is arranged as batch, position, ear, and frequency. For this representation, metric="lsd" with input_scale="db" measures spectral distance along the frequency axis. metric="rmse" and metric="mae" are also available when the training objective is direct tensor error in the representation already stored in the batch.

The first comparison below uses the input and target tensors from the same spec, so the loss should be near zero. The second comparison adds a controlled +1 dB offset to the prediction; the printed losses show the scalar values that a training loop would optimize.

hrtf_prediction = batch["inputs"]["input_hrtf"]
hrtf_target = batch["target"]["target_hrtf"]

zero_rmse = hrtf_loss(hrtf_prediction, hrtf_target, metric="rmse")

biased_prediction = hrtf_prediction + 1.0
biased_rmse = hrtf_loss(biased_prediction, hrtf_target, metric="rmse")
biased_mae = hrtf_loss(biased_prediction, hrtf_target, metric="mae")
biased_lsd = hrtf_loss(
    biased_prediction,
    hrtf_target,
    metric="lsd",
    input_scale="db",
)

print("Zero-check RMSE:", float(zero_rmse.detach().cpu()))
print("+1 dB RMSE:", float(biased_rmse.detach().cpu()))
print("+1 dB MAE:", float(biased_mae.detach().cpu()))
print("+1 dB LSD:", float(biased_lsd.detach().cpu()))

Dataset performance and cache workflow

Dataset performance is controlled by three practical choices: whether SOFA convention checks run during HRTF loading, whether loaded HRTFs keep their backing SOFA handle open, and when subject HRTFs are cached.

check_sofa_against_conventions=True runs convention validation while loading subject files. That is useful when auditing a new local root or checking files from a new source. For repeated training runs on already inspected files, check_sofa_against_conventions=False avoids repeating that validation work during dataset construction and sample access.

sofa_open=False loads the IR, TF, source positions, and metadata needed by the HRTF object, then closes the backing netCDF4 handle. This is the normal setting for training because it avoids keeping many file handles open. Use sofa_open=True when cached HRTFs must expose direct hrtf.Sofa access for dimensions, variables, attributes, summaries, or other SOFA-level inspection.

By default, preload_hrtfs=False, so HRTFs are loaded lazily. The first sample request for a subject loads the HRTF, applies the dataset-level HRTF transform, stores the result in the dataset cache, and later samples from that subject reuse it. The same cache also stores spec-level transformed HRTFs and derived acoustic values such as ITD, ILD, and SH outputs when those specs are used.

Set preload_hrtfs=True when construction should load every selected subject immediately. This moves I/O and preprocessing cost to startup, uses more memory earlier, and avoids first access stalls during a long training loop. You can also call dataset.preload_hrtfs() after construction. dataset.clear_cache() clears cached HRTFs, transformed HRTFs, and derived values; if a cached HRTF owns an open SOFA dataset, clear_cache() closes that handle before dropping the cache.

training_dataset = SONICOM(
    **compact_dataset_kwargs,
    split="train",
    check_sofa_against_conventions=False,
    sofa_open=False,
    preload_hrtfs=False,
    verbose=False,
)

# The first sample request loads and caches the subject HRTF lazily.
lazy_sample = training_dataset[0]
lazy_subject_id = lazy_sample["meta"]["subject_id"]
cached_hrtf = training_dataset.get_subject_hrtf(lazy_subject_id)
print("Lazy-loaded subject:", lazy_subject_id)
print("Cached SOFA handle open:", cached_hrtf.Sofa.is_open())

# Clear cached HRTF objects, transformed HRTFs, and derived metric values.
training_dataset.clear_cache()

# Preload can be requested at construction time.
preloaded_train_dataset = SONICOM(
    **compact_dataset_kwargs,
    split="train",
    check_sofa_against_conventions=False,
    sofa_open=False,
    preload_hrtfs=True,
    verbose=False,
)

print("Preloaded rows:", len(preloaded_train_dataset))
preloaded_train_dataset.clear_cache()

# Or preload explicitly after construction.
preloaded_train_dataset.preload_hrtfs()
reloaded_sample = preloaded_train_dataset[0]
print("Sample after explicit preload:", reloaded_sample["meta"]["subject_id"])
preloaded_train_dataset.clear_cache()

Train a small HRTF autoencoder shape example

This is not intended to be a serious model. It shows how a batched hrtfpykit dataset can feed a normal PyTorch training loop.

The model receives the compact HRTF magnitude tensor and tries to reconstruct the same tensor. The loop uses hrtf_loss with metric="lsd" because the current specs return HRTF magnitudes in dB and the final tensor axis is frequency. In real experiments, the encoder output could be used as a compact HRTF representation, a feature extractor output, or a starting point for dimensionality reduction experiments. The exact architecture, loss metric, validation strategy, and preprocessing choices depend on the research question.

import numpy as np
import torch
from torch import nn


class HRTFAutoencoder(nn.Module):
    def __init__(self, input_shape, latent_dim=16):
        super().__init__()
        self.input_shape = tuple(input_shape)
        num_features = int(np.prod(self.input_shape))
        self.encoder = nn.Sequential(
            nn.Flatten(),
            nn.Linear(num_features, 128),
            nn.ReLU(),
            nn.Linear(128, latent_dim),
        )
        self.decoder = nn.Sequential(
            nn.ReLU(),
            nn.Linear(latent_dim, 128),
            nn.ReLU(),
            nn.Linear(128, num_features),
        )

    def forward(self, hrtf_batch):
        latent = self.encoder(hrtf_batch)
        reconstruction = self.decoder(latent).reshape(
            hrtf_batch.shape[0],
            *self.input_shape,
        )
        return reconstruction, latent


device = "cuda" if torch.cuda.is_available() else "cpu"
input_shape = batch["inputs"]["input_hrtf"].shape[1:]
model = HRTFAutoencoder(input_shape=input_shape, latent_dim=16).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0.0

    for batch in train_loader:
        hrtf_input = batch["inputs"]["input_hrtf"].to(device)
        hrtf_target = batch["target"]["target_hrtf"].to(device)

        reconstruction, latent = model(hrtf_input)
        loss = hrtf_loss(
            reconstruction,
            hrtf_target,
            metric="lsd",
            input_scale="db",
        )

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        epoch_loss += float(loss.detach().cpu())

    epoch_loss /= len(train_loader)
    print(f"Epoch {epoch + 1}/{num_epochs} - lsd: {epoch_loss:.6f}")