E2A2 — spatial audio (Aria, 7-channel)

E2A2 (E²A², Encoding Efficient Asymmetric Autoencoders, github.com/danjacobellis/E2A2) is the successor of FRAPPE v3: the same split trainer (frozen per-group encoders — one learned strided projection per latent channel — SC6 softsign companding onto the coder alphabet [−31, 31], one merged decoder per operating point) with the entropy stage replaced by GGDLPC (a learned FIR predictor per channel, two affine laws mapping a causal statistic to a folded-GGD Huffman table, plus a derived zero-run mode; the C engine builds itself on import). Every reported rate is real byte-padded stream bits, round-trip verified in this notebook by decoding the latents back from the truncated bytes.

E2A2S is the E2A2 codec for the 7-microphone Aria array at 48 kHz (danjacobellis/aria_ea_audio_preprocessed): a single-group encoder (ps 256, 64 channels, decoder k9/dim1024/depth12). The waveform is used raw — no normalization — as both codec input and distortion reference (the dataset’s peaks are not uniform, so peak-normalizing would change the reference for ~40 % of the clips). Rates are bits per frame (all 7 channels).

Stream layout all_detail: no global stream — every channel is coded per macroregion, and every macroregion is decodable from its own bytes alone (one byte-pad each). The stream for operating point n is the literal bit-prefix covering channels ≤ n.

The codec is loaded from the Hugging Face hub (danjacobellis/E2A2, spatial_audio/ subdirectory) through the compressors.e2a2 public API — there is no local checkpoint or training code in this notebook.

import json, glob
import torch, numpy as np, matplotlib.pyplot as plt
import torchaudio
from IPython.display import Audio, display
from compressors import e2a2
from compressors.e2a2.recipes import load_aria_clips

Load the codec and its entropy coder

One MergedAutoencoder per operating point, all from the hub. Each operating point has its own merged decoder (multidec); the encoders froze at merge, so every operating point’s encoder is a channel prefix of the full one. entropy.json is the GGDLPC blob (per-channel predictor taps + law scalars); the Huffman tables are generated from it when the coder loads.

NAME = 'spatial_audio'
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
config, models, coder = e2a2.load_all_codecs(NAME, device=device)
ops = list(config.cumulative_channels)
model = models[ops[-1]]
rom = e2a2.rom_facts(coder.blob)

print(f'codec         = {config.codec}/{NAME}   (hub format {config.format}, source checkpoint {config.source_checkpoint})')
print(f'modality      = {config.modality}   spatial rank = {config.dim}   input channels = {config.input_channels}')
print(f'ps ladder     = {list(config.ps)} x {list(config.group_sizes)} channels')
print(f'boundaries    = {ops}  (cumulative channels = operating points)')
for n in ops:
    d = config.decoder_configs[ops.index(n)]
    print(f'  op {n:>2}: decoder ps={d["decoder_ps"]} dim={d["decoder_dim"]} depth={d["decoder_depth"]}')
for s, (ps_s, start, end) in enumerate(model.scale_groups):
    print(f'  scale {s+1}: ps={ps_s:>4}  channels {start+1}-{end}')
print(f'\nentropy coder: GGDLPC, stream_mode={coder.stream_mode}, macroregion={coder.macroregion} frames'
      + (f', global stream = channels 1-{coder.n_coarse}' if coder.n_coarse else ', no global stream'))
print(f'  alphabet [-31, 31]; persistent model = {rom["law_scalars"]} law scalars + '
      f'{rom["predictor_ints"]} predictor ints for {rom["n_channels"]} channels')
print(f'  {rom["note"]}')

codec         = E2A2/spatial_audio   (hub format e2a2_hub1, source checkpoint e2a2s9d16.pth)
modality      = audio   spatial rank = 1   input channels = 7
ps ladder     = [256] x [64] channels
boundaries    = [64]  (cumulative channels = operating points)
  op 64: decoder ps=128 dim=1024 depth=12
  scale 1: ps= 256  channels 1-64

entropy coder: GGDLPC, stream_mode=all_detail, macroregion=16384 frames, no global stream
  alphabet [-31, 31]; persistent model = 384 law scalars + 192 predictor ints for 64 channels
  edges, folded GGD pmfs, quotient split, canonical Huffman tables and the run mode are GENERATED from the per-channel scalars at load (GGDLPC.tables)

Analysis filterbank

Each latent channel is a learned (7, ps) projection. Left: impulse response averaged over the 7 microphones; right: its magnitude response in dB (Nyquist 24 kHz); ≤ 6 channels per panel. Below: the inter-microphone weight pattern — per latent channel, the fraction of tap energy on each microphone.

def chunk_sizes(n, max_per=6):
    k = max(1, -(-n // max_per))
    return [n // k + (1 if i < n % k else 0) for i in range(k)]

fs = config.sample_rate
N = 4096
freqs_khz = np.fft.rfftfreq(N) * fs / 1e3
panels = []
for s, (ps_s, start, end) in enumerate(model.scale_groups):
    W = model.encoders[s][0].weight.data.cpu().numpy()     # (C_g, 7, ps)
    off = 0
    for sz in chunk_sizes(end - start):
        panels.append((ps_s, start + off + 1, start + off + sz, W[off:off + sz].mean(1)))
        off += sz

fig, axes = plt.subplots(len(panels), 2, figsize=(10, 1.6 * len(panels)), dpi=100, constrained_layout=True)
axes = np.atleast_2d(axes)
for (ax_imp, ax_mag), (ps_s, ch_lo, ch_hi, filt) in zip(axes, panels):
    t_ms = np.arange(filt.shape[1]) / fs * 1e3
    for c, h in enumerate(filt, start=ch_lo):
        H = 20 * np.log10(np.abs(np.fft.rfft(h, N)) + 1e-9)
        ax_imp.plot(t_ms, h, lw=0.8, label=f'ch{c}')
        ax_mag.plot(freqs_khz, H, lw=0.8, label=f'ch{c}')
    ax_imp.set_title(f'ps={ps_s}   channels {ch_lo}-{ch_hi}  (impulse response, mic-averaged)', fontsize=8)
    ax_mag.set_title(f'ps={ps_s}   channels {ch_lo}-{ch_hi}', fontsize=8)
    ax_mag.set_xlim(0, fs / 2e3); ax_mag.set_ylim(-60, 10)
    ax_imp.set_ylabel('h(t)'); ax_mag.set_ylabel('|H(f)| (dB)')
    ax_mag.legend(fontsize=6, loc='upper right', framealpha=0.6, ncol=3)
axes[-1, 0].set_xlabel('time (ms)'); axes[-1, 1].set_xlabel('frequency (kHz)')
plt.show()

mic_energy = np.concatenate([(model.encoders[s][0].weight.data.cpu().numpy() ** 2).sum(2)
                             for s in range(len(model.scale_groups))])    # (C, 7)
mic_energy = mic_energy / mic_energy.sum(1, keepdims=True)
plt.figure(figsize=(9, 2.4), dpi=110)
plt.imshow(mic_energy.T, aspect='auto', cmap='magma', vmin=0, vmax=mic_energy.max())
for n in ops[:-1]:
    plt.axvline(n - 0.5, color='w', lw=0.6, ls=':')
plt.colorbar(fraction=0.02, label='fraction of tap energy')
plt.xlabel('latent channel'); plt.ylabel('microphone'); plt.yticks(range(7))
plt.title('inter-microphone weight pattern per latent channel (dotted: operating points)', fontsize=9)
plt.tight_layout(); plt.show()

Load an example clip

Validation clip 0 of danjacobellis/aria_ea_audio_preprocessed (1215 clips of 300,000 frames; eval-only), raw stored amplitude, 48 kHz, 7 channels, read through compressors.e2a2.recipes.load_aria_clips. The 5 s window (14 macroregions = 229,376 frames = 4.78 s) with the highest RMS is taken, aligned to whole macroregions.

clips = load_aria_clips(1, verbose=False)
x_full = clips[0]                                            # (7, 300000) raw float32
clip_name = clips.seq_names[0]
unit = max(coder.macroregion, max(config.ps))
x_full = x_full[..., :unit * (x_full.shape[-1] // unit)]
n_units_win = int(5 * fs) // unit
L = n_units_win * unit
n_avail = x_full.shape[-1] // unit - n_units_win + 1
rms = np.array([x_full[:, u * unit:u * unit + L].pow(2).mean().sqrt().item() for u in range(n_avail)])
u0 = int(rms.argmax())
x = x_full[:, u0 * unit:u0 * unit + L].contiguous()
n_samples = x.shape[-1]
x_in = x.unsqueeze(0).to(device)
print(f'clip 0 ({clip_name}): {x_full.shape[-1]} frames ({x_full.shape[-1] / fs:.1f} s); window offset = {u0 * unit} frames ({u0 * unit / fs:.2f} s)')
print(f'window  = {tuple(x.shape)}  ({n_samples / fs:.2f} s = {n_samples // coder.macroregion} macroregions of '
      f'{coder.macroregion} frames),  range [{x.min():.3f}, {x.max():.3f}]')
print('peak |x| per channel: ' + '  '.join(f'{v:.3f}' for v in x.abs().amax(dim=1).tolist()))

t = np.arange(n_samples) / fs
fig, axes = plt.subplots(7, 1, figsize=(9, 6), dpi=100, sharex=True, sharey=True)
for c in range(7):
    axes[c].plot(t, x[c].numpy(), lw=0.3, color='0.3'); axes[c].set_ylabel(f'mic {c}', fontsize=8)
axes[-1].set_xlabel('time (s)'); fig.suptitle('input, 7 microphones', fontsize=9)
plt.tight_layout(); plt.show()

clip 0 (loc2_script3_seq3_rec2): 294912 frames (6.1 s); window offset = 0 frames (0.00 s)
window  = (7, 229376)  (4.78 s = 14 macroregions of 16384 frames),  range [-1.000, 1.000]
peak |x| per channel: 1.000  1.000  1.000  1.000  1.000  1.000  1.000

mel = torchaudio.transforms.MelSpectrogram(sample_rate=fs, n_fft=4096, hop_length=512)

def logmel(x_1ch):
    return (mel(x_1ch) + 1e-10).log10().numpy() * 10

def show_mels(sig, title, chans=(0, 1)):
    fig, axes = plt.subplots(1, len(chans), figsize=(4.5 * len(chans), 2.6), dpi=100)
    for a, c in zip(np.atleast_1d(axes), chans):
        a.imshow(logmel(sig[c]), aspect='auto', origin='lower', cmap='magma', vmin=-80, vmax=10,
                 extent=[0, n_samples / fs, 0, 128])
        a.set_title(f'{title} — mic {c}', fontsize=8); a.set_xlabel('time (s)'); a.set_ylabel('mel bin')
    plt.tight_layout(); plt.show()

show_mels(x, 'input (log-mel, dB)')
Audio(x[:2].numpy(), rate=fs)

Analysis transform and rounding

Encode to the per-scale latents (companding included) and round to the coder’s [-31, 31] grid.

with torch.no_grad():
    latents = model.encode(x_in)
n_latent_values = sum(z.numel() for z in latents)
print(f'{x_in.numel()} input samples -> {n_latent_values} latent values ({n_latent_values / x_in.numel():.4f}x)')
for s, z in enumerate(latents):
    print(f'  scale {s+1} (ps={model.scale_groups[s][0]:>4}): latent {tuple(z.shape)}  range [{z.min():.1f}, {z.max():.1f}]')
n_clip = sum(int((z.round().abs() > 31).sum()) for z in latents)
latents_q = e2a2.encode_to_latents(model, x_in)
latents_np = [z[0].numpy() for z in latents_q]
print(f'clipped: {n_clip}/{n_latent_values}')

plt.figure(figsize=(5, 2), dpi=110)
plt.hist(np.concatenate([z.ravel() for z in latents_np]), range=(-31.5, 31.5), bins=63, width=0.85)
plt.xlim([-15, 15]); plt.title('histogram of integer latents (zoomed)')
plt.xlabel('value'); plt.ylabel('count'); plt.tight_layout(); plt.show()

1605632 input samples -> 57344 latent values (0.0357x)
  scale 1 (ps= 256): latent (1, 64, 896)  range [-9.1, 8.7]
clipped: 0/57344

Quantized latents (prior to entropy coding)

One (channels × time) image per scale at native latent resolution — each pixel is one latent value; alphabet [-31, 31] mapped linearly to [0, 255] (0 → mid-gray). Dotted lines: macroregion boundaries.

fig, axes = plt.subplots(len(latents_np), 1, figsize=(10, 1.2 + 0.09 * sum(z.shape[0] for z in latents_np)), dpi=100,
                         gridspec_kw={'height_ratios': [z.shape[0] for z in latents_np]})
for a, z, (ps_s, start, end) in zip(np.atleast_1d(axes), latents_np, model.scale_groups):
    a.imshow(((z.astype(np.int16) + 31) * 255 // 62).astype(np.uint8), cmap='gray', vmin=0, vmax=255, aspect='auto',
             extent=[0, n_samples / fs, end - start, 0], interpolation='nearest')
    for u in range(1, n_samples // coder.macroregion):
        a.axvline(u * coder.macroregion / fs, color='c', lw=0.4, ls=':')
    a.set_ylabel(f'ps={ps_s}\nch{start+1}-{end}', fontsize=8)
    a.set_title(f'{z.shape[0]} ch x {z.shape[1]} latent frames', fontsize=8)
np.atleast_1d(axes)[-1].set_xlabel('time (s)')
plt.tight_layout(); plt.show()

GGDLPC all-detail entropy coding

encode_latents produces one stream per macroregion and nothing else — no global stream; every channel is coded per region and every region’s stream is byte-padded once. The stream for operating point n is the literal bit-prefix covering channels ≤ n. CR is against 16-bit PCM per stored sample (bits/sample = bits/frame ÷ 7).

enc = e2a2.encode_latents(coder, latents_q)
assert enc.gblob == b''
region_bytes = (enc.prefix[:, -1] + 7) // 8
print(f'{len(enc.streams)} macroregion streams, one byte-pad each; at max detail '
      f'{int(region_bytes.sum()):,} bytes total, per region min/median/max = '
      f'{int(region_bytes.min())}/{int(np.median(region_bytes))}/{int(region_bytes.max())}')
print(f'\n{"op":>4} {"stream bits":>12} {"bits/frame":>11} {"bits/sample":>12} {"kbps":>8} {"CR":>7} {"bits/latent":>12}')
for n in ops:
    bits = e2a2.op_bits(coder, enc, n)
    n_lat_n = sum(min(max(n - start, 0), end - start) * z.shape[1]
                  for z, (ps_s, start, end) in zip(latents_np, model.scale_groups))
    bpf = bits / n_samples
    print(f'{n:>4} {bits:>12,} {bpf:>11.6f} {bpf / 7:>12.6f} {bpf * fs / 1000:>8.2f} {16 / (bpf / 7):>6.0f}x {bits / n_lat_n:>12.3f}')

14 macroregion streams, one byte-pad each; at max detail 7,073 bytes total, per region min/median/max = 242/524/985

  op  stream bits  bits/frame  bits/sample     kbps      CR  bits/latent
  64       56,584    0.246687     0.035241    11.84    454x        0.987

Round-trip verification

verify_latents decodes the latents back from the truncated bytes alone at every operating point’s literal bit-prefix (the global stream from its own bytes; every macroregion from its own bytes, no cross-region input) and asserts them bit-exact against the encoder’s. The reconstructions below are decoded from those byte-decoded latents (decode_latents), not from the encoder’s copy.

e2a2.verify_latents(coder, enc, latents_q, ops)
print(f'round trip verified at ops {ops}: every op decoded from its truncated bytes alone, bit-exact')

round trip verified at ops [64]: every op decoded from its truncated bytes alone, bit-exact

Reconstructions at each operating point

One encode serves every operating point: the latents for op n are decoded from the truncated bytes and fed to that point’s own merged decoder; the rate is the same encode’s literal per-region bit-prefix. PSNR is the trainer’s: on (x+1)/2 vs (x̂+1)/2 with x̂ clamped to [-1, 1], raw amplitude, all 7 channels. Per point: a 50 ms zoom of mic 0 (original vs reconstruction), the log-mel of mic 0, per-microphone PSNR, and the mic 0–1 pair to listen to.

def psnr_of(ref, est):
    return -10 * torch.nn.functional.mse_loss(ref / 2 + 0.5, est / 2 + 0.5).log10().item()

td_len = int(0.05 * fs)
td_start = int(t[x[0].abs().argmax()] * fs) - td_len // 2
td_start = min(max(td_start, 0), n_samples - td_len)
td_sl = slice(td_start, td_start + td_len)
td_t = np.arange(td_len) / fs * 1e3 + td_start / fs * 1e3

rd_clip = []
for n in ops:
    lts = e2a2.decode_latents(coder, enc, n)
    with torch.no_grad():
        xh = models[n].decode([z.to(device) for z in lts]).clamp(-1, 1)[0].float().cpu()
    bits_n = e2a2.op_bits(coder, enc, n)
    bpf_n = bits_n / n_samples
    psnr_n = psnr_of(x, xh)
    psnr_ch = [psnr_of(x[c], xh[c]) for c in range(7)]
    rd_clip.append((bpf_n, psnr_n, n))
    print(f'{n:>2} ch:  {bpf_n:.6f} bits/frame   {bpf_n / 7:.6f} bits/sample   {bpf_n * fs / 1000:7.2f} kbps   '
          f'CR={16 / (bpf_n / 7):6.0f}x   PSNR={psnr_n:.4f} dB')
    fig, axes = plt.subplots(1, 3, figsize=(12, 2.6), dpi=100, gridspec_kw={'width_ratios': [1.3, 1.3, 0.9]})
    axes[0].plot(td_t, x[0, td_sl].numpy(), lw=0.8, c='b', alpha=0.7, label='original')
    axes[0].plot(td_t, xh[0, td_sl].numpy(), lw=0.8, c='g', alpha=0.7, label=f'{n} ch')
    axes[0].set_xlabel('time (ms)'); axes[0].set_title('mic 0, 50 ms', fontsize=8); axes[0].legend(fontsize=7)
    axes[1].imshow(logmel(xh[0]), aspect='auto', origin='lower', cmap='magma', vmin=-80, vmax=10, extent=[0, n_samples / fs, 0, 128])
    axes[1].set_title(f'mic 0 log-mel, {n} ch', fontsize=8); axes[1].set_xlabel('time (s)')
    axes[2].bar(range(7), psnr_ch, color='0.5'); axes[2].axhline(psnr_n, color='tab:red', lw=1, label=f'all: {psnr_n:.2f}')
    axes[2].set_xlabel('microphone'); axes[2].set_ylabel('PSNR (dB)'); axes[2].set_title('per-mic PSNR', fontsize=8)
    axes[2].set_ylim(min(psnr_ch) - 3, max(psnr_ch) + 3); axes[2].legend(fontsize=7)
    plt.tight_layout(); plt.show()
    display(Audio(xh[:2].numpy(), rate=fs))

64 ch:  0.246687 bits/frame   0.035241 bits/sample     11.84 kbps   CR=   454x   PSNR=22.2354 dB

Per-macroregion rate profile

Each macroregion carries its own byte-padded stream, so the rate allocation over time is directly visible.

t_edges = np.arange(len(enc.streams) + 1) * coder.macroregion / fs
fig, axes = plt.subplots(2, 1, figsize=(8, 4), dpi=110, sharex=True, gridspec_kw={'height_ratios': [1, 1.6]})
axes[0].plot(t, x[0].numpy(), lw=0.3, color='0.4')
axes[0].set_ylabel('amplitude'); axes[0].set_title('input (mic 0)', fontsize=9)
for n in ops:
    rb = 8 * ((enc.prefix[:, n - 1] + 7) // 8)
    axes[1].stairs(rb / 8, t_edges, label=f'{n} ch')
axes[1].set_xlabel('time (s)'); axes[1].set_ylabel('bytes per macroregion')
axes[1].set_title(f'per-macroregion stream size ({coder.macroregion} frames each)', fontsize=9)
axes[1].legend(fontsize=8); axes[1].grid(True, alpha=0.3)
plt.tight_layout(); plt.show()

Rate–distortion

This window’s operating points over the 1215-clip whole-validation curve from results/e2a2_spatial_audio/ (falls back to the hub rd.json; rates in bits per frame = kbps × 1000 / 48000). No in-repo baseline exists for this modality; the reference points on the same 1215 clips, quoted from the trainer ledger, are LiVeAction aria_f128c28 (33.12 dB @ 0.111 bits/frame, CR 1013) and EnCodec (27.96 dB @ CR 455 = 0.246 bits/frame).

def latest(pattern):
    paths = sorted(glob.glob(pattern))
    return paths[-1] if paths else None

rd_path = latest(f'results/e2a2_{NAME}/rate_distortion_*.json')
if rd_path is not None:
    rd = json.load(open(rd_path))
    key = 'channel_counts' if 'channel_counts' in rd else 'quality_values'
    rd_ops = list(rd[key])
    rd_label = f"E2A2 {NAME} — whole validation split ({rd.get('n_images', rd.get('n_samples'))} samples, compressors harness, verified streams)"
else:
    rd = None
    hub_rd = e2a2.hub_rd_curve(NAME)
    rd_ops = [p['n_ch'] for p in hub_rd['rd_curve']]
    rd_label = f"E2A2 {NAME} — whole validation split (hub rd.json, trainer's verified streams)"

if rd is not None:
    val_bpf = [rd['results'][str(n)]['mean']['kbps'] * 1000 / rd['sample_rate'] for n in rd_ops]
    val_psnr = [rd['results'][str(n)]['mean']['PSNR_dB'] for n in rd_ops]
else:
    val_bpf = [p['bpp'] for p in hub_rd['rd_curve']]; val_psnr = [p['psnr'] for p in hub_rd['rd_curve']]

plt.figure(figsize=(6.5, 4), dpi=110)
plt.semilogx(val_bpf, val_psnr, 'o-', color='tab:blue', label=rd_label)
plt.semilogx([p[0] for p in rd_clip], [p[1] for p in rd_clip], 's--', color='tab:blue', alpha=0.5, label='E2A2 spatial_audio — this window (clip 0)')
plt.semilogx([0.111], [33.12], '^', color='black', ms=7, label='LiVeAction aria_f128c28 (1215 clips)')
plt.semilogx([16 * 7 / 455], [27.96], 'v', color='0.5', ms=7, label='EnCodec @ CR 455 (1215 clips)')
for n, bv, psv in zip(rd_ops, val_bpf, val_psnr):
    plt.annotate(f'{n}ch', (bv, psv), fontsize=7, color='tab:blue', xytext=(3, -9), textcoords='offset points')
plt.xlabel('bits per frame (7 channels)'); plt.ylabel('PSNR (dB)')
plt.grid(True, which='both', alpha=0.3); plt.legend(fontsize=7); plt.tight_layout(); plt.show()

print(f'{"op":>4} {"PSNR (1215)":>12} {"bits/frame":>11} {"bits/sample":>12} {"kbps":>7} {"CR":>6}')
for n, bv, psv in zip(rd_ops, val_bpf, val_psnr):
    print(f'{n:>4} {psv:>12.4f} {bv:>11.6f} {bv / 7:>12.6f} {bv * fs / 1000:>7.2f} {16 / (bv / 7):>5.0f}x')

  op  PSNR (1215)  bits/frame  bits/sample    kbps     CR
  64      31.6380    0.106257     0.015180    5.10  1054x