E2A2 — MRI k-space (fastMRI, multi-coil, 3-d coil images)#
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 kspace (E2A2K) codes raw multi-coil MRI k-space from the fastMRI collections (danjacobellis/fastmri_1e11, the frozen 16-bit companded codes of the mricompression package). The codec itself is the unchanged E2A2 trainer at spatial rank 3 with 2 input channels; what it sees is not raw k-space but a companded coil image: one item = one receiver coil’s k-space slice stack of one acquisition, (2, slices, H, W) real / imaginary, de-companded, line-cropped to the k-space rows and columns that are sampled in some slice (exact and invertible; the two masks are H + W bits of side information), inverse-FFT’d per slice, scaled by the item’s own rms (32 bits of side information) and passed through a fixed softsign compander y = a x / (1 + a |x|) with a = 1 / (κ · rms), κ = 32. That pre-transform is compressors.fastmri.coil_image (ported from the trainer’s E2A2_kspace/kspace_img.py); the decoder side inverts it (expand, FFT per slice, zero-fill the dropped lines). Three groups of 6 latent channels at cube sizes 8/4/2, decoder ps 4 / width 512 / kernel 3 / depth 12 per operating point, macroregion 8 voxels per axis.
Quality is not PSNR on the codec’s input. The reported numbers are the diagnostic-image procedure of compressors.fastmri (the same one behind the uncoded-requantization baseline in results/fastmri_requantize/): the whole volume’s k-space reassembled from every coded coil, root-sum-of-squares reconstruction, the fastMRI display crop, per-slice display normalization by the reference’s 99.5th percentile, uint8, and per-slice PSNR / SSIM / LPIPS / DISTS — plus the linear k-space SNR. Rates are bits per original complex k-space voxel (GGDLPC stream bits at the operating point’s prefix plus the side information, over the original R × A × S × C × H × W grid).
Stream layout two_stream: one global stream (the first scale group over the whole volume, one byte-pad) plus one detail stream per 8³ 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, kspace/ subdirectory) through the compressors.e2a2 public API — there is no local checkpoint or training code in this notebook. The dataset is read offline from the local Hugging Face cache.
import json, glob, time
import torch, numpy as np, matplotlib.pyplot as plt
from compressors import e2a2
from compressors.fastmri import mri_quality as mq, coil_image as ci
from compressors.e2a2.evaluate_rate_distortion import code_kspace_item, snr_db
from mricompression.quant import dequantize_kspace
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. config.input_domain records the pre-transform, including the knee κ.
NAME = 'kspace'
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)
KNEE = float(config.input_domain['knee'])
unit = config.macroregion
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} (real / imaginary)')
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"]} kernel={d["decoder_kernel_size"]} 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'\ninput domain: knee = {KNEE}, per-item rms = {config.input_domain["per_item_rms"]}, '
f'line crop = {config.input_domain["crop_sampled_lines"]}')
print(f' {config.input_domain["description"]}')
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/kspace (hub format e2a2_hub1, source checkpoint G5h_g3.pth)
modality = image spatial rank = 3 input channels = 2 (real / imaginary)
ps ladder = [8, 4, 2] x [6, 6, 6] channels
boundaries = [6, 12, 18] (cumulative channels = operating points)
op 6: decoder ps=4 dim=512 kernel=3 depth=12
op 12: decoder ps=4 dim=512 kernel=3 depth=12
op 18: decoder ps=4 dim=512 kernel=3 depth=12
scale 1: ps= 8 channels 1-6
scale 2: ps= 4 channels 7-12
scale 3: ps= 2 channels 13-18
input domain: knee = 32.0, per-item rms = True, line crop = True
coil images, not raw k-space: each receiver coil's k-space slice stack of one acquisition (2 x slices x H x W, real/imag) is de-companded from the fastmri_1e11 codes, line-cropped to the sampled k-space lines (exact, invertible), inverse-FFT'd per slice, scaled by its own rms and companded with a softsign of knee `knee` (kspace_img.py in the E2A2 repo, E2A2_kspace/); the codec reconstructs that companded coil image and the decoder side inverts the compander and FFTs back
entropy coder: GGDLPC, stream_mode=two_stream, macroregion=8 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 over the two input channels (real, imaginary). Shown per scale: the middle slice-plane of each channel’s cube for the real input (top row) and the imaginary input (bottom row) — the in-plane ps×ps impulse response.
for s, (ps_s, start, end) in enumerate(model.scale_groups):
W = model.encoders[s][0].weight.data.cpu() # (C_g, 2, ps, ps, ps): (in, slice, h, w)
n_g = W.shape[0]
fig, axes = plt.subplots(2, n_g, figsize=(1.5 * n_g, 3.2), dpi=110)
v = W[:, :, ps_s // 2].abs().max()
for j in range(n_g):
for ic, lab in enumerate(('re', 'im')):
a = axes[ic, j]
a.imshow(W[j, ic, ps_s // 2].numpy(), cmap='RdBu_r', vmin=-v, vmax=v); a.axis('off')
if ic == 0:
a.set_title(f'ch{start + j + 1}', fontsize=7)
axes[0, 0].text(-0.4, 0.5, 're', transform=axes[0, 0].transAxes, fontsize=8, ha='right', va='center')
axes[1, 0].text(-0.4, 0.5, 'im', transform=axes[1, 0].transAxes, fontsize=8, ha='right', va='center')
fig.suptitle(f'scale ps={ps_s} channels {start+1}-{end} (middle slice-plane of the {ps_s}^3 cube, per input channel)', fontsize=9)
plt.tight_layout(); plt.show()
The example volume#
The first knee volume of the validation split (dataset order, compressors.fastmri.volumes), repetition 0 / average 0, every coil. decode_volume returns the uint16 codes (2, reps, avgs, slices, coils, H, W); one coil item is the slab codes[:, r, a, :, c].
Left: the log-magnitude of one coil’s k-space at the middle slice on the original grid — the readout axis is fully sampled, the phase-encode axis carries a rectangle of exactly-zero lines (code 32768) on both sides; the row / column masks below the image are what the line crop keeps. Right: the companded coil image the codec sees for that coil and slice (real, imaginary, magnitude), after the line crop, the inverse FFT, the per-item rms scaling and the softsign compander at κ = 32.
ds = mq.load_split()
vols = mq.volumes(ds)
vol = next(v for v in vols if v['collection'] == 'knee_mc')
codes6 = mq.decode_volume(ds, vol) # (2, R, A, S, C, H, W) uint16
R, A, S, C, H, W = codes6.shape[1:]
REP, AVG, COIL = 0, 0, 0
SLICE = S // 2
print(f'{vol["name"]} ({vol["label"]}): {R} rep x {A} avg x {S} slices x {C} coils x {H} x {W} '
f'({codes6[0].size:,} complex voxels; {vol["n_slices"]} RSS slices)')
item = np.ascontiguousarray(codes6[:, REP, AVG, :, COIL]) # (2, S, H, W) uint16, one coil item
k_item = dequantize_kspace(item) # (S, H, W) complex64 on the original grid
rows, cols = ci.line_masks(item)
y, aux = ci.item_to_input(item, KNEE) # (2, max(S, 16), H', W') float32 in (-1, 1)
print(f'coil {COIL}: sampled rows {len(rows)}/{H}, sampled cols {len(cols)}/{W} -> coded grid {aux["cropped_shape"]}; '
f'rms = {aux["rms"]:.3e}; side information = {aux["side_bits"]} bits (32 + H + W)')
print(f'codec input {tuple(y.shape)} range [{y.min():.3f}, {y.max():.3f}]')
kmag = np.abs(k_item[SLICE]); kdb = 20 * np.log10(kmag / kmag.max() + 1e-7)
fig = plt.figure(figsize=(12, 4.4), dpi=110)
gs = fig.add_gridspec(2, 4, height_ratios=[20, 1], width_ratios=[1.15, 1, 1, 1], hspace=0.05, wspace=0.15)
ax = fig.add_subplot(gs[0, 0])
ax.imshow(kdb, cmap='magma', vmin=-100, vmax=0, aspect='auto'); ax.set_title(f'coil {COIL}, slice {SLICE}: |k| (dB re max), original {H}x{W} grid', fontsize=8)
ax.set_xlabel('columns (phase encode)'); ax.set_ylabel('rows (readout)')
axm = fig.add_subplot(gs[1, 0], sharex=ax)
cmask = np.zeros(W, bool); cmask[cols] = True
axm.imshow(cmask[None], cmap='gray', aspect='auto', vmin=0, vmax=1); axm.set_yticks([]); axm.set_xlabel('sampled columns (white)', fontsize=7)
yy = y[:, SLICE].numpy()
for j, (img, lab) in enumerate(((yy[0], 'companded real part'), (yy[1], 'companded imaginary part'), (np.hypot(yy[0], yy[1]), 'companded magnitude'))):
a = fig.add_subplot(gs[0, j + 1])
a.imshow(img, cmap='gray' if j == 2 else 'RdBu_r', vmin=(0 if j == 2 else -0.5), vmax=0.5, aspect='auto')
a.set_title(f'{lab} ({aux["cropped_shape"][1]}x{aux["cropped_shape"][2]})', fontsize=8); a.axis('off')
plt.show()
Using the latest cached version of the dataset since danjacobellis/fastmri_1e11 couldn't be found on the Hugging Face Hub (offline mode is enabled).
Found the latest cached dataset configuration 'default' at /home/dgj335/.cache/huggingface/datasets/danjacobellis___fastmri_1e11/default/0.0.0/757de3bbf4746c53eabe446f216ccdd7cade0b55 (last modified on Wed Sep 2 01:09:34 2026).
file1000100 (knee_mc/knee): 1 rep x 1 avg x 42 slices x 15 coils x 640 x 368 (148,377,600 complex voxels; 42 RSS slices)
coil 0: sampled rows 640/640, sampled cols 332/368 -> coded grid (42, 640, 332); rms = 9.051e-06; side information = 1040 bits (32 + H + W)
codec input (2, 42, 640, 332) range [-0.184, 0.404]
Analysis transform and rounding#
The coil item is zero-padded at the end of every spatial axis to the 8-voxel macroregion multiple, encoded to the per-scale latents (companding included) and rounded to the coder’s [-31, 31] grid.
xp, shape = ci.pad_to_multiple(y, unit)
x = xp.unsqueeze(0).to(device)
n_voxels = int(np.prod(aux['shape'])) # original complex voxels of the item (the rate denominator)
n_coded = int(np.prod(xp.shape[1:]))
print(f'coded volume {tuple(x.shape)} = {xp.shape[1] // unit} x {xp.shape[2] // unit} x {xp.shape[3] // unit} macroregions of {unit}^3 '
f'({n_coded:,} coded voxels for {n_voxels:,} original voxels)')
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()
coded volume (1, 2, 48, 640, 336) = 6 x 80 x 42 macroregions of 8^3 (10,321,920 coded voxels for 9,891,840 original voxels)
20643840 input values -> 8830080 latent values (0.4277x)
scale 1 (ps= 8): latent (1, 6, 6, 80, 42) range [-13.8, 20.8]
scale 2 (ps= 4): latent (1, 6, 12, 160, 84) range [-20.3, 20.0]
scale 3 (ps= 2): latent (1, 6, 24, 320, 168) range [-19.7, 22.7]
clipped: 0/8830080
Quantized latents (prior to entropy coding)#
One grayscale image per channel, the middle 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, 2.6), 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, aspect='auto')
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 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 8³ macroregion (one byte-pad each). The stream for operating point n is the literal bit-prefix covering channels ≤ n. The rate of this item is the stream bits plus the 32 + H + W side-information bits, per original complex voxel (the line-cropped grid is smaller than the original one); CR is against the dataset’s 32-bit complex codes.
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'side information: {aux["side_bits"]} bits (rms 32 + rows {H} + cols {W})')
print(f'\n{"op":>4} {"stream bits":>12} {"+ side":>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 + aux['side_bits']) / n_voxels
print(f'{n:>4} {bits:>12,} {bits + aux["side_bits"]:>12,} {bpv:>11.6f} {32 / bpv:>7.0f}x {bits / n_lat_n:>12.3f}')
global stream: 13,008 bytes (channels 1-6, whole volume, one byte-pad)
detail streams: 20160 macroregions, one byte-pad each; at max detail 3,326,264 bytes total, per region min/median/max = 61/183/268
side information: 1040 bits (rms 32 + rows 640 + cols 368)
op stream bits + side bits/voxel CR bits/latent
6 104,064 105,104 0.010625 3012x 0.860
12 2,514,528 2,515,568 0.254307 126x 2.310
18 26,714,176 26,715,216 2.700733 12x 3.025
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. Everything below is 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
The whole volume through every operating point#
Every coil item of the volume goes through the same path (code_kspace_item of the harness: item_to_input → pad → one top-op encode → every op decoded from its byte prefix, bit-exact → its own decoder → clamp → crop back → output_to_kspace): expand the compander, FFT per slice, zero-fill the dropped lines onto the original grid. The coded k-space of the volume is then reconstructed exactly as the reference is (mricompression.recon.rss on the full 6-d array, the fastMRI display crop) and scored with the compressors.fastmri procedure: per-slice display normalization by the reference’s 99.5th percentile (p995), uint8, per-slice PSNR / SSIM / LPIPS / DISTS through piq. Also the linear k-space SNR of the volume and the rate in bits per original complex voxel.
crop = mq.display_crop(H, W)
ref = mq.reference_rss(vol, ds) # (R, S, Hc, Wc) float32, the dataset's own RSS
k_ref = dequantize_kspace(codes6) # (R, A, S, C, H, W) complex64
k_hat = {n: np.empty_like(k_ref) for n in ops}
bits_vol = {n: 0.0 for n in ops}
t0 = time.time()
for r in range(R):
for a in range(A):
for c in range(C):
kh, b, _, _ = code_kspace_item(models, coder, ops, device, np.ascontiguousarray(codes6[:, r, a, :, c]), KNEE, unit)
for n in ops:
k_hat[n][r, a, :, c] = kh[n]; bits_vol[n] += b[n]
n_vox_vol = k_ref.size
metrics = mq.SliceMetrics(device, 32)
rd_vol, rec_rss = [], {}
print(f'{R * A * C} coil items coded in {time.time() - t0:.0f} s (volume: {n_vox_vol:,} complex voxels, {ref.shape[0] * ref.shape[1]} RSS slices, display crop {crop})\n')
print(f'{"op":>4} {"bits/voxel":>11} {"CR":>6} {"k-space SNR":>12} {"p995 PSNR":>10} {"SSIM":>7} {"LPIPS_dB":>9} {"DISTS_dB":>9}')
for n in ops:
rec_rss[n] = mq.rss_from_kspace(k_hat[n], crop)
per = mq.evaluate_volume(ref, rec_rss[n], metrics, ('p995',))['p995']
row = dict(n=n, bpv=bits_vol[n] / n_vox_vol, k_snr=snr_db(k_hat[n], k_ref),
**{q: mq.finite_mean(per[q])[0] for q in mq.QUALITY_METRICS})
rd_vol.append(row)
print(f'{n:>4} {row["bpv"]:>11.6f} {32 / row["bpv"]:>5.0f}x {row["k_snr"]:>12.2f} {row["PSNR_dB"]:>10.2f} {row["SSIM"]:>7.4f} {row["LPIPS_dB"]:>9.2f} {row["DISTS_dB"]:>9.2f}')
/home/dgj335/g/lib/python3.12/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.
warnings.warn(
/home/dgj335/g/lib/python3.12/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG16_Weights.IMAGENET1K_V1`. You can also use `weights=VGG16_Weights.DEFAULT` to get the most up-to-date weights.
warnings.warn(msg)
15 coil items coded in 104 s (volume: 148,377,600 complex voxels, 42 RSS slices, display crop 320)
op bits/voxel CR k-space SNR p995 PSNR SSIM LPIPS_dB DISTS_dB
6 0.011713 2732x 8.91 21.33 0.4096 2.51 5.34
12 0.256326 125x 9.88 24.09 0.5573 3.86 6.61
18 2.670032 12x 12.23 29.48 0.8151 7.76 10.60
Reconstructions at each operating point#
The middle slice of the volume’s RSS reconstruction, one row per operating point: reference | reconstruction | 10× |error|, all on the reference slice’s 99.5th-percentile gray scale (the p995 display normalization). The SNR in the title is that of the displayed slice.
mid = S // 2
ref_mid = ref[REP, mid].astype(np.float64)
peak = float(np.percentile(ref_mid, mq.PERCENTILE))
fig, axes = plt.subplots(len(ops), 3, figsize=(9, 3.1 * len(ops)), dpi=110, squeeze=False)
for i, n in enumerate(ops):
rec = rec_rss[n][REP, mid].astype(np.float64)
err = np.abs(rec - ref_mid)
snr = 10 * np.log10(np.sum(ref_mid ** 2) / max(np.sum(err ** 2), 1e-30))
row = rd_vol[i]
for j, (img, name) in enumerate(((ref_mid, f'slice {mid}: reference'),
(rec, f'op {n}: {row["bpv"]:.4f} bits/voxel, slice SNR {snr:.1f} dB'),
(10 * err, '10x |error|'))):
ax = axes[i, j]
ax.imshow(np.clip(img / peak, 0, 1), cmap='gray', vmin=0, vmax=1, interpolation='nearest')
ax.set_title(name, fontsize=8); ax.axis('off')
plt.tight_layout(); plt.show()
Per-macroregion rate map#
Each 8³ macroregion of a coil item carries its own byte-padded detail stream. The map is the middle slice-plane of the 3-d map of the example coil (detail-stream bytes per 8×8 in-plane region at the top operating point), next to the companded magnitude of that coil at the corresponding slice — the coil image’s energy is concentrated at the coil’s location, and so are the bytes.
nud, nuh, nuw = xp.shape[1] // unit, xp.shape[2] // unit, xp.shape[3] // unit
rmap = region_bytes.reshape(nud, nuh, nuw)
ds_mid = min(SLICE // unit, nud - 1)
fig, axes = plt.subplots(1, 2, figsize=(8, 4.4), dpi=110)
axes[0].imshow(np.hypot(yy[0], yy[1]), cmap='gray', vmin=0, vmax=0.5, aspect='auto'); axes[0].set_title(f'coil {COIL}, slice {SLICE}: companded magnitude', fontsize=8); axes[0].axis('off')
im = axes[1].imshow(rmap[ds_mid], cmap='viridis', aspect='auto')
axes[1].set_title(f'detail-stream bytes per {unit}x{unit}x{unit} macroregion, slice-plane {ds_mid} of {nud} (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 volume’s operating points over the whole-validation curve (all 161 volumes, every repetition and average, 7,707 RSS slices) from results/e2a2_kspace/ (falls back to the hub rd.json), as p995 PSNR and SSIM against bits per original complex voxel. The baseline is the uncoded scalar requantization of the 16-bit codes already evaluated in results/fastmri_requantize/ under the same procedure: the same quality drawn twice, at its uncoded rate (2k bits per voxel for k bits kept per real code) and at the rate of the same requantized planes after lossless JPEG-LS.
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))
rd_ops = list(rd['channel_counts'])
val = {n: dict(bpv=rd['results'][str(n)]['bits_per_voxel'], k_snr=rd['results'][str(n)]['k_snr'],
PSNR_dB=rd['results'][str(n)]['all_slices']['p995']['PSNR_dB']['mean'],
SSIM=rd['results'][str(n)]['all_slices']['p995']['SSIM']['mean']) for n in rd_ops}
rd_label = f"E2A2 {NAME} — whole validation split ({rd['n_volumes']} volumes, {rd['n_slices']} slices; compressors harness, verified streams)"
else:
hub_rd = e2a2.hub_rd_curve(NAME)
rd_ops = [p['n_ch'] for p in hub_rd['rd_curve']]
val = {p['n_ch']: dict(bpv=p['bpp'], k_snr=p['k_snr'], PSNR_dB=p['psnr'], SSIM=p['ssim']) for p in hub_rd['rd_curve']}
rd_label = f"E2A2 {NAME} — whole validation split (hub rd.json, trainer's verified streams)"
rq = json.load(open(latest('results/fastmri_requantize/rate_distortion_*.json')))
rq_rate = json.load(open(latest('results/fastmri_requantize/rate_jpegls_*.json')))
ks = [k for k in rq['quality_values'] if rq['results'][str(k)]['all_slices']['p995']['PSNR_dB']['mean'] is not None and str(k) in rq_rate['results']]
rq_psnr = [rq['results'][str(k)]['all_slices']['p995']['PSNR_dB']['mean'] for k in ks]
rq_ssim = [rq['results'][str(k)]['all_slices']['p995']['SSIM']['mean'] for k in ks]
rq_bpv = [rq['results'][str(k)]['bits_per_voxel'] for k in ks]
rq_bpv_jls = [rq_rate['results'][str(k)]['bits_per_voxel'] for k in ks]
fig, axes = plt.subplots(1, 2, figsize=(12, 4.2), dpi=110)
for ax, key, ylab in ((axes[0], 'PSNR_dB', 'p995 PSNR (dB), mean over slices'), (axes[1], 'SSIM', 'p995 SSIM, mean over slices')):
ax.semilogx(rq_bpv, [dict(PSNR_dB=p, SSIM=s)[key] for p, s in zip(rq_psnr, rq_ssim)], 'D-', color='tab:gray', ms=4, lw=1.0, label='requantize, uncoded (2k bits/voxel) — 161 volumes')
ax.semilogx(rq_bpv_jls, [dict(PSNR_dB=p, SSIM=s)[key] for p, s in zip(rq_psnr, rq_ssim)], 'D--', color='tab:gray', ms=4, lw=1.0, label='requantize + JPEG-LS — 161 volumes')
ax.semilogx([val[n]['bpv'] for n in rd_ops], [val[n][key] for n in rd_ops], 'o-', color='tab:blue', label=rd_label)
ax.semilogx([r['bpv'] for r in rd_vol], [r[key] for r in rd_vol], 's--', color='tab:blue', alpha=0.5, label=f'E2A2 {NAME} — this volume ({vol["name"]})')
for n in rd_ops:
ax.annotate(f'{n}ch', (val[n]['bpv'], val[n][key]), fontsize=7, color='tab:blue', xytext=(3, -9), textcoords='offset points')
ax.set_xlabel('bits per original complex voxel'); ax.set_ylabel(ylab); ax.grid(True, which='both', alpha=0.3)
axes[0].legend(fontsize=7, loc='upper left')
plt.tight_layout(); plt.show()
print(f'{"op":>4} {"bits/voxel":>11} {"CR":>6} {"k-space SNR":>12} {"p995 PSNR":>10} {"SSIM":>7} (whole split)')
for n in rd_ops:
v = val[n]
print(f'{n:>4} {v["bpv"]:>11.6f} {32 / v["bpv"]:>5.0f}x {v["k_snr"]:>12.2f} {v["PSNR_dB"]:>10.4f} {v["SSIM"]:>7.4f}')
op bits/voxel CR k-space SNR p995 PSNR SSIM (whole split)
6 0.012306 2600x 8.08 13.4450 0.1957
12 0.226387 141x 9.64 15.6352 0.3228
18 2.362440 14x 12.37 21.5019 0.6372