Skip to content

Compression

An INR is a compressed representation of its signal: the parameters are the code. inro.compression makes this concrete for FUTON, with a codec that fits a quantized model and serializes it, PIL codecs as baselines, and one function that measures any codec's rate-distortion point.

  1. Fita FUTON, under a rate-distortion loss
  2. Quantizeevery parameter tensor
  3. Packcodes and model settings, gzipped
  4. Rebuildthe model from the bytes
  5. Evaluateit on the pixel grid

A codec in two calls

from torchvision.io import ImageReadMode, read_image
import inro

image = read_image("data/Kodak/kodim19.png", mode=ImageReadMode.RGB)  # uint8 (3, H, W)

encoded = inro.compression.encode_futon(
    image,
    basis=("cosine", {"num_components": 128}),
    combiner=("cp", {"rank": 64}),
    head="linear",
    num_steps=500,
    lr=0.1,
    alpha=0.005,
    quant_max=127,
)
reconstruction = inro.compression.decode_futon(encoded)  # uint8 (3, H, W)

encode_futon builds a FUTON from the keyword arguments (in_features and out_features default to 2 and the channel count), trains it with inro.rate_distortion_loss, quantizes every parameter tensor to quant_max levels, and packs the integer codes, their scales, the model configuration, the image ID and the spatial shape into gzip-compressed bytes. decode_futon rebuilds the model from those bytes and evaluates it on the pixel grid. The specs must be registry keys or (key, dict) pairs so that they survive the round trip. The linear head reads a whole image off the grid in one contraction.

The rate-distortion loss is \(\text{MSE} + \alpha \cdot H\), where \(H\) is the size-weighted mean of a Laplace entropy estimate of each parameter tensor, in bits; alpha=0 disables the rate term. quant_interval rounds the parameters to the quantization grid every so many steps during training, so that the model adapts to it (quantization-aware training); the final quantization always runs.

Measuring a codec

futon = inro.compression.evaluate_compression(
    image,
    inro.compression.encode_futon,
    inro.compression.decode_futon,
    basis=("cosine", {"num_components": 128}),
    combiner=("cp", {"rank": 64}),
    head="linear",
    num_steps=500,
)
jpeg = inro.compression.evaluate_compression(
    image, inro.compression.encode_pil, inro.compression.decode_pil, format="JPEG", quality=50
)
webp = inro.compression.evaluate_compression(
    image, inro.compression.encode_pil, inro.compression.decode_pil, format="WEBP", quality=50
)

Each call returns {"PSNR (dB)", "bit rate (bpp)", "encoding time (ms)", "decoding time (ms)"}, with the decoded image under "reconstructed" if return_reconstruction=True. PSNR uses data_range=255 and is infinite for a lossless codec.

The keyword arguments go to the encoder, so a rate-distortion curve is a loop over rank (or num_components, quant_max, alpha) for FUTON and over quality for the PIL codecs, and inro.LinePlot draws the curves:

records = [
    {"codec": "FUTON", "rank": rank, **inro.compression.evaluate_compression(
        image, inro.compression.encode_futon, inro.compression.decode_futon,
        basis=("cosine", {"num_components": 128}), combiner=("cp", {"rank": rank}), head="linear",
    )}
    for rank in (16, 32, 64, 128)
]
plot = inro.LinePlot(records)
fig, ax = plot.plot(x="bit rate (bpp)", y="PSNR (dB)", hue="codec")

encode_pil accepts a path, a file object, or a (channels, H, W) tensor and forwards its keyword arguments to PIL.Image.save; format is required because the image is written to memory.

Quantization in the training loop

The same machinery is available in inro.train:

result = inro.train(
    model,
    train_set,
    eval_set,
    quantize=True,
    quant_max=127,
    quant_interval=100,
    metrics={"psnr": inro.psnr},
    rate_metrics={"bpp": inro.bits_per_pixel},
)

Every evaluation then scores a quantized copy of the model, and the rate metrics see its parameters serialized with inro.encode_dict (torch.save plus gzip). result["model_dict"] holds the integer codes and scales, which inro.uniform_dequantize(model, model_dict) loads back, and checkpoint.pt keeps the quantized weights as an ordinary state dict. inro.uniform_quantize(model, quant_max) and inro.uniform_dequantize are also usable on their own.

Entropy coding

Function Purpose
inro.compression.encode_arithmetic(symbols, pmfs), decode_arithmetic Integer arithmetic coding with a separate probability mass function per position, for codecs that model the parameter distribution explicitly
inro.soft_entropy A differentiable entropy estimate from soft histogram assignments
inro.laplace_entropy The closed-form estimate that rate_distortion_loss uses

Codecs

codecs

Image codecs: PIL formats (JPEG, WebP, PNG, ...) and FUTON.

An encoder maps an image to bytes and a decoder maps the bytes back to an image tensor (channels, H, W), so evaluate_compression can benchmark any pair.

encode_pil 

encode_pil(image: str | PathLike[str] | IO[bytes] | Tensor, **kwargs: Any) -> bytes

Encode an image with a PIL codec.

Parameters:

  • image (str | PathLike[str] | IO[bytes] | Tensor) –

    Path or binary file object of an image that PIL can open, or an image tensor (channels, H, W), uint8 or floating-point in [0, 1].

  • **kwargs (Any, default: {} ) –

    Arguments of PIL.Image.Image.save. Must include format, such as "JPEG", "WEBP" or "PNG", because PIL cannot infer it for an in-memory buffer, and may include codec options such as quality.

Returns:

  • bytes –

    The encoded image file.

Raises:

  • ValueError –

    If image is not a path, a binary file object or a tensor.

Examples:

>>> image = torch.randint(0, 256, (3, 32, 32), dtype=torch.uint8)
>>> encoded = encode_pil(image, format="JPEG", quality=50)
>>> decode_pil(encoded).shape
torch.Size([3, 32, 32])
Source code in inro/compression/codecs.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def encode_pil(
    image: str | os.PathLike[str] | IO[bytes] | Tensor, **kwargs: Any
) -> bytes:
    """Encode an image with a PIL codec.

    Args:
        image: Path or binary file object of an image that PIL can open, or an image
            tensor ``(channels, H, W)``, ``uint8`` or floating-point in ``[0, 1]``.
        **kwargs: Arguments of :meth:`PIL.Image.Image.save`. Must include ``format``,
            such as ``"JPEG"``, ``"WEBP"`` or ``"PNG"``, because PIL cannot infer it for
            an in-memory buffer, and may include codec options such as ``quality``.

    Returns:
        The encoded image file.

    Raises:
        ValueError: If ``image`` is not a path, a binary file object or a tensor.

    Examples:
        >>> image = torch.randint(0, 256, (3, 32, 32), dtype=torch.uint8)
        >>> encoded = encode_pil(image, format="JPEG", quality=50)
        >>> decode_pil(encoded).shape
        torch.Size([3, 32, 32])
    """
    if isinstance(image, (str, os.PathLike, io.IOBase)):
        pil_image = Image.open(image)
    elif isinstance(image, Tensor):
        pil_image = FT.to_pil_image(image)
    else:
        raise ValueError("image must be a string, Path, IO[bytes] or Tensor")

    buffer = io.BytesIO()
    pil_image.save(buffer, **kwargs)
    return buffer.getvalue()

decode_pil 

decode_pil(encoded: bytes) -> Tensor

Decode an image file held in memory, such as the output of encode_pil.

Parameters:

  • encoded (bytes) –

    Bytes of an image file in any format PIL can read.

Returns:

  • Tensor –

    Image (channels, H, W), with the channels of the stored PIL mode and uint8 values for 8-bit images.

Source code in inro/compression/codecs.py 
64
65
66
67
68
69
70
71
72
73
74
def decode_pil(encoded: bytes) -> Tensor:
    """Decode an image file held in memory, such as the output of :func:`encode_pil`.

    Args:
        encoded: Bytes of an image file in any format PIL can read.

    Returns:
        Image ``(channels, H, W)``, with the channels of the stored PIL mode and
            ``uint8`` values for 8-bit images.
    """
    return FT.pil_to_tensor(Image.open(io.BytesIO(encoded)))

encode_futon 

encode_futon(image: str | PathLike[str] | Image | Tensor, *, subsample: float = 1.0, lr: float = 0.1, num_steps: int = 500, alpha: float = 0.005, quant_max: int = 127, quant_interval: int = -1, log_interval: int = 100, device: str | device | None = None, seed: int = 0, **kwargs: Any) -> bytes

Encode an image by fitting a FUTON and serializing its quantized parameters.

Training minimizes rate_distortion_loss.

Parameters:

  • image (str | PathLike[str] | Image | Tensor) –

    Path or PIL image, converted to RGB, or uint8 tensor (channels, H, W).

  • subsample (float, default: 1.0 ) –

    Fraction of pixels sampled at each step, in (0, 1].

  • lr (float, default: 0.1 ) –

    Initial learning rate of Adam.

  • num_steps (int, default: 500 ) –

    Number of training steps.

  • alpha (float, default: 0.005 ) –

    Weight of the parameters' Laplace entropy estimate in the loss.

  • quant_max (int, default: 127 ) –

    Largest integer level. Each parameter tensor is scaled by quant_max / max(abs(w)) and rounded.

  • quant_interval (int, default: -1 ) –

    Steps between fake quantizations of the model; non-positive values disable them. The final quantization always runs.

  • log_interval (int, default: 100 ) –

    Steps between history records; see train.

  • device (str | device | None, default: None ) –

    Device to run on; None uses CUDA when available, else the CPU.

  • seed (int, default: 0 ) –

    Seed of the global PyTorch RNG, set after the model is built, so it does not control its initialization.

  • **kwargs (Any, default: {} ) –

    Arguments of FUTON. basis, combiner and head are required, each a registry key or a (key, dict) pair, and every value must load with torch.load(weights_only=True). in_features and out_features default to 2 and the channel count.

Returns:

  • bytes –

    The integer codes, their scales, the model config, the image ID and the spatial shape, serialized with torch.save and gzip.

Raises:

  • ValueError –

    If basis, combiner or head is missing, or is neither a registry key nor a (key, dict) pair.

Examples:

>>> image = torch.randint(0, 256, (3, 16, 16), dtype=torch.uint8)
>>> encoded = encode_futon(
...     image,
...     num_steps=2,
...     basis=("cosine", {"num_components": 8}),
...     combiner=("cp", {"rank": 8}),
...     head="linear",
... )
>>> decode_futon(encoded).shape
torch.Size([3, 16, 16])
Source code in inro/compression/codecs.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def encode_futon(
    image: str | os.PathLike[str] | Image.Image | Tensor,
    *,
    subsample: float = 1.0,
    lr: float = 0.1,
    num_steps: int = 500,
    alpha: float = 0.005,
    quant_max: int = 127,
    quant_interval: int = -1,
    log_interval: int = 100,
    device: str | torch.device | None = None,
    seed: int = 0,
    **kwargs: Any,
) -> bytes:
    """Encode an image by fitting a FUTON and serializing its quantized parameters.

    Training minimizes :func:`~inro.losses.rate_distortion_loss`.

    Args:
        image: Path or PIL image, converted to RGB, or ``uint8`` tensor
            ``(channels, H, W)``.
        subsample: Fraction of pixels sampled at each step, in ``(0, 1]``.
        lr: Initial learning rate of Adam.
        num_steps: Number of training steps.
        alpha: Weight of the parameters' Laplace entropy estimate in the loss.
        quant_max: Largest integer level. Each parameter tensor is scaled by
            ``quant_max / max(abs(w))`` and rounded.
        quant_interval: Steps between fake quantizations of the model; non-positive
            values disable them. The final quantization always runs.
        log_interval: Steps between history records; see :func:`train`.
        device: Device to run on; ``None`` uses CUDA when available, else the CPU.
        seed: Seed of the global PyTorch RNG, set after the model is built, so it does
            not control its initialization.
        **kwargs: Arguments of :class:`FUTON`. ``basis``, ``combiner`` and ``head``
            are required, each a registry key or a ``(key, dict)`` pair, and every value
            must load with ``torch.load(weights_only=True)``. ``in_features`` and
            ``out_features`` default to 2 and the channel count.

    Returns:
        The integer codes, their scales, the model config, the image ID and the spatial
            shape, serialized with ``torch.save`` and gzip.

    Raises:
        ValueError: If ``basis``, ``combiner`` or ``head`` is missing, or is neither
            a registry key nor a ``(key, dict)`` pair.

    Examples:
        >>> image = torch.randint(0, 256, (3, 16, 16), dtype=torch.uint8)
        >>> encoded = encode_futon(
        ...     image,
        ...     num_steps=2,
        ...     basis=("cosine", {"num_components": 8}),
        ...     combiner=("cp", {"rank": 8}),
        ...     head="linear",
        ... )
        >>> decode_futon(encoded).shape
        torch.Size([3, 16, 16])
    """
    for name in ("basis", "combiner", "head"):
        if name not in kwargs:
            raise ValueError(f"encode_futon requires the FUTON argument `{name}`")
        _check_module_spec(name, kwargs[name])

    dataset = ImageCoordinateDataset(image=image, subsample=subsample)
    model_config = {
        "in_features": dataset.input.shape[-1],
        "out_features": dataset.target.shape[-1],
        **kwargs,
    }
    model = FUTON(**model_config)
    train_results = train(
        model,
        dataset,
        loss_fn=partial(rate_distortion_loss, alpha=alpha),
        lr=lr,
        num_steps=num_steps,
        quantize=True,
        quant_max=quant_max,
        quant_interval=quant_interval,
        log_interval=log_interval,
        eval_interval=-1,
        device=device,
        seed=seed,
    )

    metadata = {
        "image": {"id": dataset.id, "shape": tuple(dataset.grid_size)},
        "model": {"config": model_config, "state": train_results["model_dict"]},
    }
    return encode_dict(metadata)

decode_futon 

decode_futon(encoded: bytes, *, device: str | device | None = None) -> Tensor

Reconstruct an image from the bytes of encode_futon.

Rebuilds the stored FUTON, dequantizes its parameters and evaluates the whole coordinate grid in one pass. The output is clamped to [-1, 1] and mapped to uint8.

Parameters:

  • encoded (bytes) –

    Bytes returned by encode_futon.

  • device (str | device | None, default: None ) –

    Device to run on; None uses the default device.

Returns:

  • Tensor –

    Reconstructed uint8 image (channels, H, W) on device.

Source code in inro/compression/codecs.py 
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def decode_futon(encoded: bytes, *, device: str | torch.device | None = None) -> Tensor:
    """Reconstruct an image from the bytes of :func:`encode_futon`.

    Rebuilds the stored FUTON, dequantizes its parameters and evaluates the whole
    coordinate grid in one pass. The output is clamped to ``[-1, 1]`` and mapped to
    ``uint8``.

    Args:
        encoded: Bytes returned by :func:`encode_futon`.
        device: Device to run on; ``None`` uses the default device.

    Returns:
        Reconstructed ``uint8`` image ``(channels, H, W)`` on ``device``.
    """
    metadata = decode_dict(encoded)
    model_config = metadata["model"]["config"]
    model_dict = metadata["model"]["state"]
    image_shape = metadata["image"]["shape"]

    coordinates = create_coordinates(image_shape).to(device)

    model = FUTON(**model_config)
    uniform_dequantize(model, model_dict)

    model.to(device)
    model.eval()
    with torch.no_grad():
        output = model(coordinates)

    return ImageCoordinateDataset.postprocess(output)

Evaluation

evaluation

Rate-distortion evaluation of image codecs.

evaluate_compression 

evaluate_compression(image: Tensor, encode_fn: Callable[..., bytes], decode_fn: Callable[[bytes], Tensor], /, *, return_reconstruction: bool = False, **kwargs: Any) -> dict[str, float | Tensor]

Measure an image codec's PSNR, bit rate, and encoding and decoding times.

PSNR uses data_range=255 and is infinite for lossless codecs. Timings are wall-clock, from perf_counter, without CUDA synchronization.

Parameters:

  • image (Tensor) –

    Original uint8 image (channels, H, W).

  • encode_fn (Callable[..., bytes]) –

    Encoder, called as encode_fn(image, **kwargs).

  • decode_fn (Callable[[bytes], Tensor]) –

    Decoder, called as decode_fn(encoded); its result must be on the device of image.

  • return_reconstruction (bool, default: False ) –

    Whether to include the decoded image in the result.

  • **kwargs (Any, default: {} ) –

    Arguments of the encoder.

Returns:

  • dict[str, float | Tensor] –

    Dictionary with the floats "PSNR (dB)", "bit rate (bpp)", "encoding time (ms)" and "decoding time (ms)", and, if requested, the decoded tensor "reconstructed".

Examples:

>>> import torch
>>> from inro.compression import decode_pil, encode_pil
>>> image = torch.randint(0, 256, (3, 32, 32), dtype=torch.uint8)
>>> results = evaluate_compression(image, encode_pil, decode_pil, format="PNG")
>>> results["PSNR (dB)"]
inf
Source code in inro/compression/evaluation.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def evaluate_compression(
    image: Tensor,
    encode_fn: Callable[..., bytes],
    decode_fn: Callable[[bytes], Tensor],
    /,
    *,
    return_reconstruction: bool = False,
    **kwargs: Any,
) -> dict[str, float | Tensor]:
    """Measure an image codec's PSNR, bit rate, and encoding and decoding times.

    PSNR uses ``data_range=255`` and is infinite for lossless codecs. Timings are
    wall-clock, from ``perf_counter``, without CUDA synchronization.

    Args:
        image: Original ``uint8`` image ``(channels, H, W)``.
        encode_fn: Encoder, called as ``encode_fn(image, **kwargs)``.
        decode_fn: Decoder, called as ``decode_fn(encoded)``; its result must be on the
            device of ``image``.
        return_reconstruction: Whether to include the decoded image in the result.
        **kwargs: Arguments of the encoder.

    Returns:
        Dictionary with the floats ``"PSNR (dB)"``, ``"bit rate (bpp)"``,
            ``"encoding time (ms)"`` and ``"decoding time (ms)"``, and, if requested,
            the decoded tensor ``"reconstructed"``.

    Examples:
        >>> import torch
        >>> from inro.compression import decode_pil, encode_pil
        >>> image = torch.randint(0, 256, (3, 32, 32), dtype=torch.uint8)
        >>> results = evaluate_compression(image, encode_pil, decode_pil, format="PNG")
        >>> results["PSNR (dB)"]
        inf
    """
    start_time = time.perf_counter()
    encoded = encode_fn(image, **kwargs)
    encoding_time = 1000 * (time.perf_counter() - start_time)

    start_time = time.perf_counter()
    reconstruction = decode_fn(encoded)
    decoding_time = 1000 * (time.perf_counter() - start_time)

    results: dict[str, float | Tensor] = {
        "PSNR (dB)": psnr(reconstruction, image),
        "bit rate (bpp)": bits_per_pixel(encoded, image),
        "encoding time (ms)": encoding_time,
        "decoding time (ms)": decoding_time,
    }

    if return_reconstruction:
        results["reconstructed"] = reconstruction

    return results

Arithmetic coding

arithmetic_coding

Integer arithmetic coding with per-position probability mass functions.

Every position of a symbol sequence is coded under its own probability mass function (PMF), and alphabet sizes may differ between positions. PMFs are scaled to integer frequency tables and coded with a 32-bit integer arithmetic coder using bit-plus-follow (underflow) renormalization.

References

Witten et al., "Arithmetic Coding for Data Compression", Communications of the ACM 1987.

encode_arithmetic 

encode_arithmetic(symbols: Tensor | ndarray | Sequence[int], pmfs: list[_PMF] | tuple[_PMF, ...], precision: int = 1 << 16) -> bytes

Encode a symbol sequence with arithmetic coding.

Position i is coded under its own PMF pmfs[i], and alphabet sizes may differ between positions. Every symbol must have nonzero probability under its PMF: a zero-probability symbol is not rejected and silently corrupts the bitstream.

Parameters:

  • symbols (Tensor | ndarray | Sequence[int]) –

    Integer symbols, flattened to (N,). A tensor must be on the CPU.

  • pmfs (list[_PMF] | tuple[_PMF, ...]) –

    List or tuple of N PMFs, where pmfs[i][k] is the probability of symbol k at position i, not necessarily normalized. Each PMF is a tensor on any device, a NumPy array or a sequence of floats.

  • precision (int, default: 1 << 16 ) –

    Total integer frequency each PMF is scaled to, a count rather than a number of bits. Nonzero probabilities keep a frequency of at least 1, so the total can slightly exceed it. Must match decode_arithmetic.

Returns:

  • bytes –

    A 4-byte big-endian uint32 symbol count N followed by the bitstream.

Raises:

  • ValueError –

    If pmfs is not a list or tuple of length N, a PMF is empty, has a negative entry or sums to zero, or a symbol is outside the alphabet of its PMF.

Examples:

>>> pmfs = [[0.7, 0.3], [0.1, 0.2, 0.7], [0.4, 0.6]]
>>> encoded = encode_arithmetic([0, 2, 1], pmfs)
>>> decode_arithmetic(encoded, pmfs)
array([0, 2, 1])
Source code in inro/compression/arithmetic_coding.py 
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
def encode_arithmetic(
    symbols: Tensor | np.ndarray | Sequence[int],
    pmfs: list[_PMF] | tuple[_PMF, ...],
    precision: int = 1 << 16,
) -> bytes:
    """Encode a symbol sequence with arithmetic coding.

    Position ``i`` is coded under its own PMF ``pmfs[i]``, and alphabet sizes may differ
    between positions. Every symbol must have nonzero probability under its PMF: a
    zero-probability symbol is not rejected and silently corrupts the bitstream.

    Args:
        symbols: Integer symbols, flattened to ``(N,)``. A tensor must be on the CPU.
        pmfs: List or tuple of ``N`` PMFs, where ``pmfs[i][k]`` is the probability of
            symbol ``k`` at position ``i``, not necessarily normalized. Each PMF is a
            tensor on any device, a NumPy array or a sequence of floats.
        precision: Total integer frequency each PMF is scaled to, a count rather than a
            number of bits. Nonzero probabilities keep a frequency of at least ``1``, so
            the total can slightly exceed it. Must match :func:`decode_arithmetic`.

    Returns:
        A 4-byte big-endian ``uint32`` symbol count ``N`` followed by the bitstream.

    Raises:
        ValueError: If ``pmfs`` is not a list or tuple of length ``N``, a PMF is empty,
            has a negative entry or sums to zero, or a symbol is outside the alphabet of
            its PMF.

    Examples:
        >>> pmfs = [[0.7, 0.3], [0.1, 0.2, 0.7], [0.4, 0.6]]
        >>> encoded = encode_arithmetic([0, 2, 1], pmfs)
        >>> decode_arithmetic(encoded, pmfs)
        array([0, 2, 1])
    """
    symbols = np.asarray(symbols, dtype=np.int64).ravel()
    num_symbols = symbols.size
    if not isinstance(pmfs, (list, tuple)) or len(pmfs) != num_symbols:
        raise ValueError("pmfs must be a list/tuple with same length as symbols")
    tables = _frequency_tables(pmfs, precision)

    writer = _BitWriter()
    low = 0
    high = _MAX_STATE
    pending_bits = 0

    def write_bit_with_pending(bit: int) -> None:
        nonlocal pending_bits
        writer.write_bit(bit)
        while pending_bits > 0:
            writer.write_bit(1 - bit)
            pending_bits -= 1

    for index, symbol in enumerate(symbols):
        cumulative, total = tables[index]
        alphabet_size = len(cumulative) - 1
        if symbol < 0 or symbol >= alphabet_size:
            raise ValueError(
                f"symbol {symbol} out of range for PMF at position {index}"
            )
        interval_size = high - low + 1
        high = low + (interval_size * int(cumulative[symbol + 1]) // total) - 1
        low = low + (interval_size * int(cumulative[symbol]) // total)

        # Emit a shared leading bit; defer bits while the interval straddles 1/2.
        while True:
            if high < _HALF:
                write_bit_with_pending(0)
            elif low >= _HALF:
                write_bit_with_pending(1)
                low -= _HALF
                high -= _HALF
            elif low >= _QUARTER and high < _THREE_QUARTERS:
                pending_bits += 1
                low -= _QUARTER
                high -= _QUARTER
            else:
                break
            low = (low << 1) & _MAX_STATE
            high = ((high << 1) & _MAX_STATE) | 1

    # One final bit and the deferred bits identify the remaining interval.
    pending_bits += 1
    write_bit_with_pending(0 if low < _QUARTER else 1)

    return struct.pack(">I", num_symbols) + writer.finish()

decode_arithmetic 

decode_arithmetic(encoded: bytes, pmfs: list[_PMF] | tuple[_PMF, ...], precision: int = 1 << 16) -> ndarray

Decode a bitstream produced by encode_arithmetic.

Parameters:

  • encoded (bytes) –

    Bytes returned by encode_arithmetic.

  • pmfs (list[_PMF] | tuple[_PMF, ...]) –

    List or tuple of the PMFs used for encoding, one per symbol.

  • precision (int, default: 1 << 16 ) –

    The precision used for encoding.

Returns:

  • ndarray –

    The decoded symbols, an int64 array (N,).

Raises:

  • ValueError –

    If encoded is shorter than the 4-byte header, pmfs is not a list or tuple whose length equals the encoded symbol count, or a PMF is empty, has a negative entry or sums to zero.

Source code in inro/compression/arithmetic_coding.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
def decode_arithmetic(
    encoded: bytes,
    pmfs: list[_PMF] | tuple[_PMF, ...],
    precision: int = 1 << 16,
) -> np.ndarray:
    """Decode a bitstream produced by :func:`encode_arithmetic`.

    Args:
        encoded: Bytes returned by :func:`encode_arithmetic`.
        pmfs: List or tuple of the PMFs used for encoding, one per symbol.
        precision: The ``precision`` used for encoding.

    Returns:
        The decoded symbols, an ``int64`` array ``(N,)``.

    Raises:
        ValueError: If ``encoded`` is shorter than the 4-byte header, ``pmfs`` is not a
            list or tuple whose length equals the encoded symbol count, or a PMF is
            empty, has a negative entry or sums to zero.
    """
    if len(encoded) < 4:
        raise ValueError("Input too short")
    num_symbols = struct.unpack(">I", encoded[:4])[0]
    if not isinstance(pmfs, (list, tuple)) or len(pmfs) != num_symbols:
        raise ValueError(
            "pmfs must be a list/tuple with same length as encoded symbols"
        )
    tables = _frequency_tables(pmfs, precision)

    reader = _BitReader(encoded[4:])
    value = 0
    for _ in range(_STATE_BITS):
        value = (value << 1) | reader.read_bit()

    low = 0
    high = _MAX_STATE
    symbols = np.zeros(num_symbols, dtype=np.int64)

    for index, (cumulative, total) in enumerate(tables):
        interval_size = high - low + 1
        scaled = ((value - low + 1) * total - 1) // interval_size
        # Find the symbol whose cumulative-frequency interval contains scaled.
        symbol = max(bisect_right(cumulative, int(scaled)) - 1, 0)
        symbols[index] = symbol

        high = low + (interval_size * int(cumulative[symbol + 1]) // total) - 1
        low = low + (interval_size * int(cumulative[symbol]) // total)

        # Mirror the encoder's renormalization while consuming new bits.
        while True:
            if high < _HALF:
                pass
            elif low >= _HALF:
                value -= _HALF
                low -= _HALF
                high -= _HALF
            elif low >= _QUARTER and high < _THREE_QUARTERS:
                value -= _QUARTER
                low -= _QUARTER
                high -= _QUARTER
            else:
                break
            low = (low << 1) & _MAX_STATE
            high = ((high << 1) & _MAX_STATE) | 1
            value = ((value << 1) | reader.read_bit()) & _MAX_STATE

    return symbols