E2A2 — hyperspectral (AVIRIS, 3-d volume)#

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.

E2A2 hyperspectral_3d codes a 224-band AVIRIS radiance tile as a 1-channel 3-d volume (1, 224, H, W) (band axis first): three groups of 6 latent channels at cube sizes 8/4/2, decoder ps 8, macroregion 32 voxels per axis (224 = 7 × 32 keeps every band). Rates are bits per voxel.

Stream layout two_stream: one global stream (the first scale group over the whole plane, one byte-pad) plus one detail stream per macroregion (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, hyperspectral_3d/ 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 PIL.Image
from IPython.display import display
from compressors import e2a2
from compressors.e2a2.recipes import load_aviris

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 = 'hyperspectral_3d'
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} voxels per axis'
      + (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/hyperspectral_3d   (hub format e2a2_hub1, source checkpoint e2a2h3d_g3.pth)
modality      = image   spatial rank = 3   input channels = 1
ps ladder     = [8, 4, 2] x [6, 6, 6] channels
boundaries    = [6, 12, 18]  (cumulative channels = operating points)
  op  6: decoder ps=8 dim=768 depth=12
  op 12: decoder ps=8 dim=768 depth=12
  op 18: decoder ps=8 dim=768 depth=24
  scale 1: ps=   8  channels 1-6
  scale 2: ps=   4  channels 7-12
  scale 3: ps=   2  channels 13-18

entropy coder: GGDLPC, stream_mode=two_stream, macroregion=32 voxels per axis, global stream = channels 1-6
  alphabet [-31, 31]; persistent model = 108 law scalars + 180 predictor ints for 18 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 ps×ps×ps cube projection. Left: the spectral response of each channel (taps summed over the ps×ps spatial support → ps taps along the band axis). Right: the band-averaged spatial impulse response.

for s, (ps_s, start, end) in enumerate(model.scale_groups):
    W = model.encoders[s][0].weight.data.cpu()[:, 0]        # (C_g, ps, ps, ps): (band, h, w)
    n_g = end - start
    spec = W.sum(dim=(2, 3)).numpy()                       # (C_g, ps) along the band axis
    fig = plt.figure(figsize=(11, 3.0), dpi=110)
    gs = fig.add_gridspec(1, 2, width_ratios=[2, 1.6])
    ax = fig.add_subplot(gs[0])
    cols = plt.cm.turbo(np.linspace(0, 1, n_g))
    for j in range(n_g):
        ax.plot(spec[j], 'o-', lw=0.8, ms=3, color=cols[j], label=f'ch{start + j + 1}')
    ax.set_xlabel('band offset within the cube'); ax.set_ylabel('sum of spatial taps')
    ax.set_title(f'spectral response, scale ps={ps_s} (channels {start+1}-{end})', fontsize=9); ax.legend(fontsize=6, ncol=2)
    sub = gs[1].subgridspec(1, n_g, wspace=0.05)
    sp = W.mean(1).numpy()                                 # (C_g, ps, ps)
    v = np.abs(sp).max()
    for j in range(n_g):
        a = fig.add_subplot(sub[0, j])
        a.imshow(sp[j], cmap='RdBu_r', vmin=-v, vmax=v); a.axis('off')
    fig.text(0.72, 0.97, f'band-averaged {ps_s}x{ps_s} spatial impulse per channel', fontsize=8, ha='center')
    plt.tight_layout(); plt.show()
_images/d3c9531308f9b2b5cbe2f58388229b1dc1e372f7066d92d448dadfcf52010ce5.webp _images/7171cee0ab1302e6ad1184d7585badedf6f1e9a545724a07026aab0ac279149e.webp _images/fb9c5fb59747dac8069a61070f4b0edb0bfc7cf7dccf9ff7ffdcc0c7f99d32e0.webp

Load an example tile#

Validation tile 0 of danjacobellis/aviris_1k_val (94 tiles; eval-only) through the compressors.e2a2.recipes loader (an int16 .npy cache under ~/.cache/compressors/, because datasets cannot materialize these rows): x = int16 / 32768, no other scaling, as a (1, 1, 224, H, W) volume with every spatial axis cropped (never resized) to the 32-voxel macroregion multiple. False-color display uses bands 30 / 19 / 9 as R / G / B, each stretched to its 1–99 % range.

x, info = load_aviris(1, coder.macroregion, volume=True, verbose=False)[0]
x = x.to(device)
origin = info['origin']
unit = coder.macroregion
n_voxels = x[0, 0].numel()
x_01 = x / 2 + 0.5
cube = x[0, 0]                                              # (224, H, W)
print(f'tile origin {origin}   input {tuple(x.shape)}  ({n_voxels} voxels = '
      f'{x.shape[2] // unit} x {x.shape[3] // unit} x {x.shape[4] // unit} macroregions of {unit}^3),  '
      f'range [{x.min():.3f}, {x.max():.3f}]')
tile origin f110803t01p00r09_0_1   input (1, 1, 224, 736, 736)  (121339904 voxels = 7 x 23 x 23 macroregions of 32^3),  range [-1.000, 0.993]
RGB_BANDS = (30, 19, 9)

def false_color(t, bands=RGB_BANDS, lo_hi=None):
    """(224, H, W) tensor -> uint8 (H, W, 3), per-band 1-99 % stretch (shared lo/hi if given)."""
    a = t[list(bands)].float().cpu().numpy()
    if lo_hi is None:
        lo_hi = [np.percentile(b, (1, 99)) for b in a]
    out = np.stack([np.clip((b - lo) / (hi - lo + 1e-9), 0, 1) for b, (lo, hi) in zip(a, lo_hi)], -1)
    return (out * 255).astype(np.uint8), lo_hi

rgb, lo_hi = false_color(cube)
mid = 112
fig, axes = plt.subplots(1, 3, figsize=(12, 4), dpi=110)
axes[0].imshow(rgb); axes[0].set_title(f'false color, bands {RGB_BANDS}', fontsize=9); axes[0].axis('off')
g = cube[mid].cpu().numpy()
axes[1].imshow(g, cmap='gray', vmin=np.percentile(g, 1), vmax=np.percentile(g, 99))
axes[1].set_title(f'band {mid} (grayscale)', fontsize=9); axes[1].axis('off')
spec_mean = cube.mean(dim=(1, 2)).cpu().numpy(); spec_std = cube.std(dim=(1, 2)).cpu().numpy()
axes[2].plot(spec_mean, lw=1); axes[2].fill_between(np.arange(224), spec_mean - spec_std, spec_mean + spec_std, alpha=0.25)
axes[2].set_xlabel('band index'); axes[2].set_ylabel('radiance / 32768'); axes[2].set_title('mean spectrum ± 1σ', fontsize=9)
plt.tight_layout(); plt.show()
_images/59a9c16fe2640e32337139a94803730f49b7beaae2f19af2a95dd6ffbf3df8fe.webp

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)
n_latent_values = sum(z.numel() for z in latents)
print(f'{x.numel()} input values -> {n_latent_values} latent values ({n_latent_values / x.numel():.4f}x)')
for s, z in enumerate(latents):
    print(f'  scale {s+1} (ps={model.scale_groups[s][0]:>2}): 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)
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()
121339904 input values -> 103802496 latent values (0.8555x)
  scale 1 (ps= 8): latent (1, 6, 28, 92, 92)  range [-7.7, 7.3]
  scale 2 (ps= 4): latent (1, 6, 56, 184, 184)  range [-3.7, 3.8]
  scale 3 (ps= 2): latent (1, 6, 112, 368, 368)  range [-7.9, 8.1]
clipped: 0/103802496
_images/d937f34846e32bde2c04e6255364182faf5277bb34a3477b819644051a6fbd14.webp

Quantized latents (prior to entropy coding)#

One grayscale image per channel, the middle band-slice of each latent cube at native resolution, grouped by scale; alphabet [-31, 31] mapped linearly to [0, 255] (0 → mid-gray).

for s, (ps_s, start, end) in enumerate(model.scale_groups):
    z = latents_np[s]                                       # (C_g, D_s, H_s, W_s)
    n_g = z.shape[0]
    fig, axes = plt.subplots(1, n_g, figsize=(1.6 * n_g, 1.9), dpi=110)
    for j, a in enumerate(np.atleast_1d(axes)):
        a.imshow(((z[j, z.shape[1] // 2].astype(np.int16) + 31) * 255 // 62).astype(np.uint8), cmap='gray', vmin=0, vmax=255)
        a.set_title(f'ch{start + j + 1}', fontsize=7); a.axis('off')
    fig.suptitle(f'ps={ps_s}  channels {start+1}-{end}  ({z.shape[1]}x{z.shape[2]}x{z.shape[3]} native, middle band-slice)', fontsize=9)
    plt.tight_layout(); plt.show()
_images/654d5473d9f34a8e132e552b93a42d6d1f2c8a3816e86b073582063852310a19.webp _images/23ffad6c5a6c808f8238294243676751676329ac20014bc363c17c892db1746a.webp _images/5b75943fdc30eb4e608af1e480437b9f409b357c9f6a1ba8cf114dd10d91adcd.webp

GGDLPC two-stream entropy coding#

encode_latents produces one global stream (channels 1–6, the whole volume, one byte-pad) plus one detail stream per 32³ macroregion (one byte-pad each). The stream for operating point n is the literal bit-prefix covering channels ≤ n. CR is against the 16-bit samples.

enc = e2a2.encode_latents(coder, latents_q)
print(f'global stream:  {len(enc.gblob):,} bytes  (channels 1-{coder.n_coarse}, whole volume, one byte-pad)')
region_bytes = (enc.prefix[:, -1] + 7) // 8
print(f'detail streams: {len(enc.streams)} macroregions, 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/voxel":>11} {"CR":>8} {"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) * int(np.prod(z.shape[1:]))
                  for z, (ps_s, start, end) in zip(latents_np, model.scale_groups))
    bpv = bits / n_voxels
    print(f'{n:>4} {bits:>12,} {bpv:>11.6f} {16 / bpv:>7.0f}x {bits / n_lat_n:>12.3f}')
global stream:  335,561 bytes  (channels 1-6, whole volume, one byte-pad)
detail streams: 3703 macroregions, one byte-pad each; at max detail 26,154,825 bytes total, per region min/median/max = 52/7191/10750

  op  stream bits  bits/voxel       CR  bits/latent
   6    2,684,488    0.022124     723x        1.888
  12   15,236,456    0.125568     127x        1.191
  18  211,923,088    1.746524       9x        2.042

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 [6, 12, 18]: 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. PSNR is the trainer’s: on (x+1)/2 vs (x̂+1)/2, clamped to [-1, 1], over the whole volume. A fixed 256×256 spatial crop (false color, shared stretch) of the original (left) and the reconstruction (right) per point, then the per-band mean |x − x̂| over the tile.

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

ch_, cw_ = 256, 256
ct = (x.shape[3] - ch_) // 2
cl = (x.shape[4] - cw_) // 2
rd_tile, err_bands = [], {}
fig, axes = plt.subplots(len(ops), 2, figsize=(6, 3.1 * len(ops)), dpi=110)
for i, n in enumerate(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)
    bpv_n = e2a2.op_bits(coder, enc, n) / n_voxels
    psnr_n = psnr_of(xh)
    rd_tile.append((bpv_n, psnr_n, n))
    err_bands[n] = (cube - xh[0, 0]).abs().mean(dim=(1, 2)).cpu().numpy()
    axes[i, 0].imshow(rgb[ct:ct+ch_, cl:cl+cw_]); axes[i, 0].set_title('original', fontsize=8)
    axes[i, 1].imshow(false_color(xh[0, 0, :, ct:ct+ch_, cl:cl+cw_], lo_hi=lo_hi)[0])
    axes[i, 1].set_title(f'{n} ch: {bpv_n:.5f} bits/voxel, {psnr_n:.2f} dB', fontsize=8)
    for a in axes[i]: a.axis('off')
    del xh
plt.tight_layout(); plt.show()

print(f'{"op":>4} {"PSNR (dB)":>10} {"bits/voxel":>11} {"CR":>8}')
for bpv_n, psnr_n, n in rd_tile:
    print(f'{n:>4} {psnr_n:>10.4f} {bpv_n:>11.6f} {16 / bpv_n:>7.0f}x')

plt.figure(figsize=(6.5, 3), dpi=110)
for n in ops:
    plt.plot(err_bands[n], lw=0.9, label=f'{n} ch')
plt.xlabel('band index'); plt.ylabel('mean |x - x̂|  (radiance / 32768)')
plt.title('per-band reconstruction error, this tile', fontsize=9); plt.legend(fontsize=7)
plt.tight_layout(); plt.show()
_images/bf896e094da4dcd87864680bab3957dea6be98cfc354c83baf529af6e9d9c495.webp
  op  PSNR (dB)  bits/voxel       CR
   6    13.3307    0.022124     723x
  12    14.3146    0.125568     127x
  18    28.3725    1.746524       9x
_images/ab3d4543492ae6a68732ee17d100bbb1176524e15854fcd68b7a966ce79d08ca.webp

Per-macroregion rate map#

Each 32³ macroregion carries its own byte-padded detail stream. The map sums the stream bytes over the 7 band-axis macroregions of each spatial 32×32 column.

nud, nuh, nuw = x.shape[2] // unit, x.shape[3] // unit, x.shape[4] // unit
fig, axes = plt.subplots(1, 2, figsize=(9, 3.8), dpi=110)
axes[0].imshow(rgb); axes[0].set_title('input'); axes[0].axis('off')
im = axes[1].imshow(region_bytes.reshape(nud, nuh, nuw).sum(0), cmap='viridis')
axes[1].set_title(f'detail-stream bytes per {unit}x{unit} spatial column ({nud} band-axis regions summed, op {ops[-1]})', fontsize=8); axes[1].axis('off')
fig.colorbar(im, ax=axes[1], fraction=0.046)
plt.tight_layout(); plt.show()
_images/f46afe5bd0efbc63fa76402b87c33cde382a795e623208951c12913f6f0fc177.webp

Rate–distortion#

This tile’s operating points over the 94-tile whole-validation curve from results/e2a2_hyperspectral_3d/ (falls back to the hub rd.json). No in-repo baseline exists for this modality; the reference points on the same 94 tiles, quoted from the trainer ledger, are LiVeAction hyper_f8c8 (18.52 dB @ 0.0278 bits/voxel, CR 575), JPEG 2000 at the same CR (18.18 dB), and the 2-d E2A2 hyperspectral ladder (results/e2a2_hyperspectral/, bpp / 224).

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_bpv = [rd['results'][str(n)]['mean']['bpp'] for n in rd_ops]
    val_psnr = [rd['results'][str(n)]['mean']['PSNR_dB'] for n in rd_ops]
else:
    val_bpv = [p['bpp'] for p in hub_rd['rd_curve']]; val_psnr = [p['psnr'] for p in hub_rd['rd_curve']]

rd2d_path = latest('results/e2a2_hyperspectral/rate_distortion_*.json')
if rd2d_path is not None:
    rd2d = json.load(open(rd2d_path))
    h2_bpv = [rd2d['results'][str(n)]['mean']['bpp'] / 224 for n in rd2d['channel_counts']]
    h2_psnr = [rd2d['results'][str(n)]['mean']['PSNR_dB'] for n in rd2d['channel_counts']]
else:
    h2 = e2a2.hub_rd_curve('hyperspectral')
    h2_bpv = [p['bpp'] / 224 for p in h2['rd_curve']]; h2_psnr = [p['psnr'] for p in h2['rd_curve']]

plt.figure(figsize=(6.5, 4), dpi=110)
plt.semilogx(h2_bpv, h2_psnr, 'D-', color='tab:gray', ms=4, lw=1.0, label='E2A2 hyperspectral (2-d, 224 ch) — 94 tiles')
plt.semilogx(val_bpv, val_psnr, 'o-', color='tab:blue', label=rd_label)
plt.semilogx([p[0] for p in rd_tile], [p[1] for p in rd_tile], 's--', color='tab:blue', alpha=0.5, label=f'E2A2 hyperspectral_3d — this tile ({origin})')
plt.semilogx([16 / 575], [18.52], '^', color='black', ms=7, label='LiVeAction hyper_f8c8 (94 tiles)')
plt.semilogx([16 / 575], [18.18], 'v', color='0.5', ms=7, label='JPEG 2000 @ CR 575 (94 tiles)')
for n, bv, psv in zip(rd_ops, val_bpv, val_psnr):
    plt.annotate(f'{n}ch', (bv, psv), fontsize=7, color='tab:blue', xytext=(3, -9), textcoords='offset points')
plt.xlabel('bits per voxel'); 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 (94)":>10} {"bits/voxel":>11} {"CR":>7}')
for n, bv, psv in zip(rd_ops, val_bpv, val_psnr):
    print(f'{n:>4} {psv:>10.4f} {bv:>11.6f} {16 / bv:>6.0f}x')
_images/5c7ae352b3f4b8554e3d9e40056b3d769cf44370e95ef8f3f213905aa51862e8.webp
  op  PSNR (94)  bits/voxel      CR
   6    18.4898    0.016149    991x
  12    19.9341    0.071941    222x
  18    35.0237    0.918947     17x