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()
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()
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
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()
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, x̂ 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()
op PSNR (dB) bits/voxel CR
6 13.3307 0.022124 723x
12 14.3146 0.125568 127x
18 28.3725 1.746524 9x
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()
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')
op PSNR (94) bits/voxel CR
6 18.4898 0.016149 991x
12 19.9341 0.071941 222x
18 35.0237 0.918947 17x