Encoding and decoding#
This chapter encodes the Kodak test image one sample at a time in plain Python: predict, pool the statistic, select a cell, emit the split symbol or extend a zero run. It then decodes the stream by mirroring every step, and verifies the inline bytes are identical to the C engine’s. The package’s own pure-Python reference encoder (GGDLPC.reference) follows the same spec; the decoder here is written inline because outside this notebook decoding exists only in C. The chapter ends with the real rate of the full image compared against the differentiable proxy.
Parameters#
N = 128 # alphabet bound: 8-bit RGB planes centered to [-128, 127]
RANK = 2 # spatial rank (image)
SEED = 20260820
FIT_IMAGES = 16 # LSDIR_8k training images used for the fit (Kodak is test-only)
FIT_CROP = 256 # center-crop side for the fitting images
KODAK_INDEX = 22 # the test image: the same Kodak image as the FRAPPE v3 notebooks
Data, fit, codec#
Identical fit to the previous chapters (LSDIR crops; Kodak is test-only); Codec(blob) generates the tables (chapter 2) and compiles them for the C engine. The sample-by-sample demonstration uses a 64×64 crop of the test image, encoded as its own call, since pure-Python encoding of the full image would take minutes. The full image is encoded by the C engine at the end.
import numpy as np
from datasets import load_dataset
def center_crop(img, side):
w, h = img.size
left, top = (w - side) // 2, (h - side) // 2
return img.crop((left, top, left + side, top + side))
lsdir = load_dataset('danjacobellis/LSDIR_8k', split='train')
fit_set = []
for sample in lsdir:
img = sample['image'].convert('RGB')
if min(img.size) < FIT_CROP:
continue
arr = np.asarray(center_crop(img, FIT_CROP), np.int64) - 128
fit_set.append(arr.transpose(2, 0, 1)) # (3, H, W) in [-128, 127]
if len(fit_set) == FIT_IMAGES:
break
kodak = load_dataset('danjacobellis/kodak', split='validation')
test_img = kodak[KODAK_INDEX]['image'].convert('RGB')
test_tensor = np.asarray(test_img, np.int64).transpose(2, 0, 1) - 128
import torch
import GGDLPC
def loader_factory():
batches = [np.stack(fit_set[i:i + 4]) for i in range(0, len(fit_set), 4)]
return (torch.from_numpy(b) for b in batches)
channels, provenance = GGDLPC.fit(loader_factory, N=N, rank=RANK)
blob = GGDLPC.new_blob(N, RANK, channels, provenance)
GGDLPC fit N=128 rank=2: predictor moments over 4 batches (0s)
GGDLPC fit N=128 rank=2: dense h(S, r) over 4 batches (0s)
GGDLPC fit N=128 rank=2: ch 0: mu=+0.003 a=0.8852 b=0.4220 u=0.692 v=+0.080 wq=[179, 200, -105, -7, -14] bq=-35
GGDLPC fit N=128 rank=2: ch 1: mu=-0.079 a=0.8896 b=0.4631 u=0.818 v=+0.047 wq=[173, 196, -97, -7, -14] bq=-55
GGDLPC fit N=128 rank=2: ch 2: mu=-0.017 a=0.8912 b=0.3859 u=0.628 v=+0.093 wq=[175, 199, -95, -9, -17] bq=-70
codec = GGDLPC.Codec(blob)
builds = codec.builds # the per-channel tables (chapter 2)
TRACE_Y, TRACE_X, TRACE_SIDE = 384, 128, 64
z = test_tensor[:, TRACE_Y:TRACE_Y + TRACE_SIDE, TRACE_X:TRACE_X + TRACE_SIDE]
C, H, W = z.shape
print(f'tracing {z.shape}, {z.size} samples')
display(test_img.crop((TRACE_X, TRACE_Y, TRACE_X + TRACE_SIDE, TRACE_Y + TRACE_SIDE))
.resize((4 * TRACE_SIDE, 4 * TRACE_SIDE), 0))
tracing (3, 64, 64), 12288 samples
A bit writer#
The stream is a plain MSB-first bit sequence, padded with zeros to a byte boundary once at the very end. Channel prefixes are literal bit prefixes: truncating the byte stream at a channel boundary leaves every earlier channel decodable.
class BitWriter:
def __init__(self):
self.bits = []
def put(self, value, nbits):
for i in range(nbits - 1, -1, -1):
self.bits.append((int(value) >> i) & 1)
def tobytes(self):
bits = self.bits + [0] * (-len(self.bits) % 8)
out = bytearray()
for i in range(0, len(bits), 8):
byte = 0
for b in bits[i:i + 8]:
byte = (byte << 1) | b
out.append(byte)
return bytes(out)
Per-sample prediction and statistic terms#
This is the scalar form of the maps from chapter 1, for one position of one rank-2 plane: the fixed-point FIR prediction over the causal template (arithmetic-mean fallback at the borders), and the template’s absolute differences \((A, \text{cnt})\) that pool with the residual window into the statistic.
def pred_terms(zp, i, j, wq, bq):
n0 = int(zp[i - 1, j]) if i >= 1 else 0 # north
n1 = int(zp[i, j - 1]) if j >= 1 else 0 # west
p01 = int(zp[i - 1, j - 1]) if (i >= 1 and j >= 1) else 0
m0 = int(zp[i - 2, j]) if i >= 2 else 0
m1 = int(zp[i, j - 2]) if j >= 2 else 0
if i >= 2 and j >= 2:
acc = bq + wq[0] * n0 + wq[1] * n1 + wq[2] * p01 + wq[3] * m0 + wq[4] * m1
pred = max(-N, min(N, (acc + 128) // 256))
else:
s = (n0 if i >= 1 else 0) + (n1 if j >= 1 else 0)
cnt = int(i >= 1) + int(j >= 1)
pred = max(-N, min(N, (2 * s + cnt) // (2 * cnt))) if cnt else 0
A, cnt2 = 0, 0
if i >= 2:
A += abs(n0 - m0)
cnt2 += 1
if j >= 2:
A += abs(n1 - m1)
cnt2 += 1
if i >= 1 and j >= 1:
A += abs(n0 - p01) + abs(n1 - p01)
cnt2 += 2
return pred, A, cnt2
The encode loop#
The encoder is one flat loop per channel, in scan order. Each sample gets a prediction, a statistic (template terms pooled with the window of the last 16 coded absolute residuals), a cell selected by the derived edges, and the fold to \(w = (r + N) \bmod (2N+1)\). The sample is then either emitted as a split symbol or, when the cell is at or below run_jmax, absorbed into a zero-run: the run length \(L\) is coded as \(\lfloor L / 2^g \rfloor\) one-bits, and a zero bit plus \(g\) remainder bits plus an exclusion-table symbol codes the sample that interrupts it. The trace arrays remain in scope.
from collections import deque
def canonical_codes(lengths):
lengths = np.asarray(lengths, np.int64)
order = sorted(range(len(lengths)), key=lambda s: (lengths[s], s))
codes = np.zeros(len(lengths), np.int64)
code, prev = 0, int(lengths[order[0]])
for sym in order:
code <<= int(lengths[sym]) - prev
prev = int(lengths[sym])
codes[sym] = code
code += 1
return codes
def emit_symbol(wtr, split, codes, ws):
hi = ws >> split.k
if split.win_lo <= hi < split.win_lo + split.n_win:
s = hi - split.win_lo
wtr.put(codes[s], int(split.lengths[s]))
else:
wtr.put(codes[split.n_win], int(split.lengths[split.n_win])) # escape
wtr.put(hi, split.hi_bits)
if split.k:
wtr.put(ws & ((1 << split.k) - 1), split.k)
nsym = 2 * N + 1
STAT_W = 16
positions = [(i, j) for i in range(H) for j in range(W)]
T = len(positions)
wtr = BitWriter()
ch_bits_inline = []
trace = {'t': [], 'pred': [], 'stat': [], 'cell': [], 'res': [], 'bits': []} # channel 0
run_events = [] # (channel, start t, length)
for c in range(C):
b = builds[c]
codes = [canonical_codes(sp.lengths) for sp in b.splits]
ecodes = [canonical_codes(sp.lengths) if sp is not None else None
for sp in b.excl]
zp, wq, bq = z[c], b.wq, int(b.bq)
win, wsum = deque(), 0
bits_start = len(wtr.bits)
def push(absres):
global wsum
win.append(absres)
wsum += absres
if len(win) > STAT_W:
wsum -= win.popleft()
t = 0
while t < T:
pred, A, cnt = pred_terms(zp, *positions[t], wq, bq)
den = cnt + len(win)
stat = min(b.stat_cap, (16 * (A + wsum)) // den) if den else 0
cell = int(np.searchsorted(b.edges, stat, side='right'))
r = int(zp[positions[t]]) - pred
ws = (r + N) % nsym
if b.run_jmax < 0 or cell > b.run_jmax: # ordinary symbol
nb0 = len(wtr.bits)
emit_symbol(wtr, b.splits[cell], codes[cell], ws)
if c == 0:
trace['t'].append(t)
trace['pred'].append(pred)
trace['stat'].append(stat)
trace['cell'].append(cell)
trace['res'].append(r)
trace['bits'].append(len(wtr.bits) - nb0)
push(abs(r))
t += 1
continue
g = int(b.run_g[cell]) # zero-run mode
m = 1 << g
u = t
while u < T:
pu, _, _ = pred_terms(zp, *positions[u], wq, bq)
if int(zp[positions[u]]) != pu:
break
u += 1
L = u - t
run_events.append((c, t, L))
for _ in range(L >> g):
wtr.put(1, 1)
for _ in range(L):
push(0)
if u == T: # run reaches the end
if L & (m - 1):
wtr.put(1, 1)
t = T
continue
wtr.put(0, 1)
if g:
wtr.put(L & (m - 1), g)
r_i = int(zp[positions[u]]) - pu # the interrupting sample
emit_symbol(wtr, b.excl[cell], ecodes[cell], (r_i + N) % nsym)
push(abs(r_i))
t = u + 1
ch_bits_inline.append(len(wtr.bits) - bits_start)
data_inline = wtr.tobytes()
print(f'{len(data_inline)} bytes, per-channel bits {ch_bits_inline}, '
f'{len(run_events)} runs')
6068 bytes, per-channel bits [16235, 16185, 16118], 7 runs
Inspecting the trace#
Each trace entry is one ordinarily coded sample of channel 0: its position, prediction, statistic, selected cell, residual, and actual emitted bits. Samples coded inside runs, and the samples that interrupt them, do not appear in these arrays; run_events records those, and they stay blank in the maps below. On photographic data the statistic rarely falls low enough to enter run mode — the low cells belong to sparse sources such as autoencoder latents — so runs are few here.
import matplotlib.pyplot as plt
for k in range(5):
print(f"t={trace['t'][k]}: pred={trace['pred'][k]:+4d} stat={trace['stat'][k]:4d} "
f"cell={trace['cell'][k]:2d} r={trace['res'][k]:+4d} "
f"bits={trace['bits'][k]}")
cell_map = np.full(T, np.nan)
bits_map = np.full(T, np.nan)
cell_map[trace['t']] = trace['cell']
bits_map[trace['t']] = trace['bits']
fig, axs = plt.subplots(1, 2, figsize=(11, 4))
im = axs[0].imshow(cell_map.reshape(H, W))
axs[0].set_title('selected cell (channel 0)')
plt.colorbar(im, ax=axs[0])
im = axs[1].imshow(bits_map.reshape(H, W))
axs[1].set_title('emitted bits per sample')
plt.colorbar(im, ax=axs[1])
plt.tight_layout()
for c in range(C):
zeros = sum(L for cc, _, L in run_events if cc == c)
print(f'channel {c}: {zeros} samples coded inside runs')
t=1: pred= -38 stat= 608 cell=36 r= +3 bits=6
t=2: pred= -35 stat= 234 cell=31 r= -2 bits=5
t=3: pred= -37 stat= 180 cell=30 r= +0 bits=5
t=4: pred= -37 stat= 137 cell=28 r= +2 bits=4
t=5: pred= -35 stat= 125 cell=28 r= -3 bits=5
channel 0: 0 samples coded inside runs
channel 1: 0 samples coded inside runs
channel 2: 12 samples coded inside runs
The decode loop#
Decoding mirrors the encoder exactly. The bit sequence a run emits does not depend on sample values, so the reader first recovers the run length from the counting bits, then fills the run positions with their own predictions (each prediction computed from already-decoded samples, exactly as the encoder saw them). Ordinary symbols decode through the canonical table — a (length, code) -> symbol lookup fed one bit at a time — and unfold back to the residual using the constraint that \(z = \text{pred} + r\) must lie in \([-N, N]\).
class BitReader:
def __init__(self, data):
self.data = data
self.pos = 0
def get(self, nbits):
v = 0
for _ in range(nbits):
byte = self.data[self.pos >> 3]
v = (v << 1) | ((byte >> (7 - (self.pos & 7))) & 1)
self.pos += 1
return v
def read_symbol(rdr, split, table):
code, length = 0, 0
while True:
code = (code << 1) | rdr.get(1)
length += 1
s = table.get((length, code))
if s is not None:
break
if split.esc and s == split.n_win: # escape: hi written raw
hi = rdr.get(split.hi_bits)
else:
hi = split.win_lo + s
lo = rdr.get(split.k) if split.k else 0
return (hi << split.k) | lo
def unfold(ws, pred):
r = ws - N
if r < -N - pred:
r += nsym
elif r > N - pred:
r -= nsym
return r
z_dec = np.zeros_like(z)
rdr = BitReader(data_inline)
for c in range(C):
b = builds[c]
tables = [{(int(sp.lengths[s]), int(cd)): s
for s, cd in enumerate(canonical_codes(sp.lengths))}
for sp in b.splits]
etables = [{(int(sp.lengths[s]), int(cd)): s
for s, cd in enumerate(canonical_codes(sp.lengths))}
if sp is not None else None for sp in b.excl]
zp, wq, bq = z_dec[c], b.wq, int(b.bq)
win, wsum = deque(), 0
def push(absres):
global wsum
win.append(absres)
wsum += absres
if len(win) > STAT_W:
wsum -= win.popleft()
t = 0
while t < T:
pred, A, cnt = pred_terms(zp, *positions[t], wq, bq)
den = cnt + len(win)
stat = min(b.stat_cap, (16 * (A + wsum)) // den) if den else 0
cell = int(np.searchsorted(b.edges, stat, side='right'))
if b.run_jmax < 0 or cell > b.run_jmax: # ordinary symbol
ws = read_symbol(rdr, b.splits[cell], tables[cell])
r = unfold(ws, pred)
zp[positions[t]] = pred + r
push(abs(r))
t += 1
continue
g = int(b.run_g[cell]) # zero-run mode
m = 1 << g
zeros, interrupted, rem = 0, False, 0
while True:
if t + zeros == T:
break # exact multiple, no tail bit
bit = rdr.get(1)
if bit == 1:
remaining = T - t - zeros
if remaining < m:
zeros += remaining # tail: rest of channel
break
zeros += m
else:
interrupted = True
rem = rdr.get(g) if g else 0
zeros += rem
break
for _ in range(zeros): # fill: z = its prediction
p_t, _, _ = pred_terms(zp, *positions[t], wq, bq)
zp[positions[t]] = p_t
push(0)
t += 1
if interrupted:
p_t, _, _ = pred_terms(zp, *positions[t], wq, bq)
ws = read_symbol(rdr, b.excl[cell], etables[cell])
r = unfold(ws, p_t)
zp[positions[t]] = p_t + r
push(abs(r))
t += 1
assert np.array_equal(z_dec, z)
print('inline decode is bit-exact')
inline decode is bit-exact
Byte identity with the C engine#
codec.encode runs the C engine over the same crop. The stream must match the inline one byte for byte, and the C decoder must reproduce the tensor.
data_c, ch_bits = codec.encode(z)
assert data_c == data_inline
assert ch_bits.tolist() == ch_bits_inline
assert np.array_equal(codec.decode(data_c, z.shape), z)
print(f'byte-identical: {len(data_c)} bytes')
print(f'crop rate: {8 * len(data_c)} bits = {8 * len(data_c) / z.size:.4f} bits/sample '
f'(vs 8 raw)')
print('first bytes:', data_c[:16].hex(' '))
byte-identical: 6068 bytes
crop rate: 48544 bits = 3.9505 bits/sample (vs 8 raw)
first bytes: 3f ff ff fd 2d 16 24 23 c9 4c 87 78 2f 17 a4 9b
The full image through the C engine#
The C engine encodes the whole Kodak image as one call, round-trip verified. Every rate below is real bits of this byte-padded stream.
data_full, ch_bits_full = codec.encode(test_tensor)
assert np.array_equal(codec.decode(data_full, test_tensor.shape), test_tensor)
real_bits = 8 * len(data_full)
n_pixels = test_tensor.shape[1] * test_tensor.shape[2]
print(f'{len(data_full)} bytes, round-trip verified')
print(f'real rate: {real_bits / test_tensor.size:.4f} bits/sample = '
f'{real_bits / n_pixels:.4f} bpp (vs 24 raw)')
print('per-channel bits:', ch_bits_full.tolist())
562525 bytes, round-trip verified
real rate: 3.8149 bits/sample = 11.4446 bpp (vs 24 raw)
per-channel bits: [1500435, 1490266, 1509494]
Truncation at a channel boundary#
Channel prefixes are literal bit prefixes of the stream. Keeping only the bytes that cover channel 0’s bits leaves channel 0 decodable; the per-channel bit counts returned by encode locate the boundary.
n0_bytes = (int(ch_bits_full[0]) + 7) // 8
z0 = codec.decode(data_full[:n0_bytes], test_tensor.shape, n_channels=1)
assert np.array_equal(z0[0], test_tensor[0])
print(f'channel 0 decoded from the first {n0_bytes} of {len(data_full)} bytes')
channel 0 decoded from the first 187555 of 562525 bytes
Real bits against the differentiable proxy#
The proxy charges one encoder call the ideal code length of the same folded GGD conditional, evaluated continuously, plus an expected 3.5 bits of byte padding. On the integers it lies near the real rate (Huffman and run-mode overheads differ from the ideal); on noise-perturbed inputs it is the training-time surrogate, and its gradient with respect to the input supplies the training signal for a transform that feeds the coder. Real bits come from the stream; the proxy numbers are the model.
z_t = torch.from_numpy(test_tensor)[None]
proxy_int = float(GGDLPC.proxy_call_bits(z_t, blob))
gen = torch.Generator().manual_seed(SEED)
z_noisy = z_t.double() + (torch.rand(z_t.shape, generator=gen, dtype=torch.float64) - 0.5)
z_noisy.requires_grad_()
proxy_bits = GGDLPC.proxy_call_bits(z_noisy, blob)
proxy_bits.backward()
print(f'real stream: {real_bits} bits ({real_bits / n_pixels:.4f} bpp)')
print(f'proxy (integers): {proxy_int:.1f} bits ({proxy_int / n_pixels:.4f} bpp)')
print(f'proxy (noisy): {proxy_bits.detach().item():.1f} bits '
f'({proxy_bits.detach().item() / n_pixels:.4f} bpp)')
plt.figure(figsize=(7, 4))
plt.imshow(z_noisy.grad[0, 0], vmin=-0.5, vmax=0.5, cmap='gray')
plt.colorbar()
plt.title('d(proxy bits) / d(input), channel 0')
plt.tight_layout()
real stream: 4500200 bits (11.4446 bpp)
proxy (integers): 4431194.6 bits (11.2691 bpp)
proxy (noisy): 4457035.0 bits (11.3348 bpp)