database_compression.py

# Can we compress random data with Qrack?
# Robust variant: Gray code + bit-plane interleaving for equalized Hamming fidelity.
#
# Encoding modes (set via --mode argument):
#   'gray'  : Gray-coded bucket indices, sequential packing (2x fewer bit errors
#             from ±1 quantization noise vs plain binary; default)
#   'sign'  : w=1 sign-only encoding — fully equalized, most robust, lower density
#   'plane' : Gray code + bit-plane interleaving — groups same-significance bits
#             together so burst errors affect one plane, not all bits of one amplitude
#
# Hamming fidelity = 1 - (hamming_distance / total_bits)
#
# By Dan Strano and (Anthropic) Claude.

import math
import os
import random
import sys

from pyqrack import QrackSimulator


# ---------------------------------------------------------------------------
# Gray code helpers
# ---------------------------------------------------------------------------

def to_gray(n):
    return n ^ (n >> 1)

def from_gray(n):
    mask = n >> 1
    while mask:
        n ^= mask
        mask >>= 1
    return n


# ---------------------------------------------------------------------------
# Fidelity metrics
# ---------------------------------------------------------------------------

def calc_stats(ideal_ket, split_ket):
    n_pow = len(ideal_ket)
    u_u = 1.0 / n_pow
    numer = 0.0
    denom = 0.0
    l2 = 0.0
    prob_diff = 0.0

    for i in range(n_pow):
        c = ideal_ket[i]
        e = split_ket[i]
        l2 += e * c.conjugate()
        p_i = (c * c.conjugate()).real
        q_i = (e * e.conjugate()).real
        numer += (p_i - u_u) * (q_i - u_u)
        denom += (p_i - u_u) ** 2
        prob_diff += (p_i - q_i) ** 2

    xeb = numer / denom if denom > 0 else 0.0
    l2 = (l2 * l2.conjugate()).real

    return xeb, l2, prob_diff


def hamming_fidelity(orig_bits, recovered_bits):
    n = len(orig_bits)
    if n == 0:
        return 1.0
    errors = sum(a != b for a, b in zip(orig_bits, recovered_bits))
    return 1.0 - errors / n, errors


# ---------------------------------------------------------------------------
# Encoding: 'gray' mode
# Gray-coded bucket indices packed sequentially (w bits re, w bits im per amp).
# A ±1 quantization error flips exactly 1 bit (vs up to w in plain binary).
# ---------------------------------------------------------------------------

def bits_to_amps_gray(bits, w, n_amps):
    levels = 1 << w
    step   = 2.0 / levels
    amps   = []
    norm_sq = 0.0
    bit_pos = 0
    total   = len(bits)
    i_max = 0
    n_max = 0.0

    for i in range(n_amps):
        re_bucket = 0
        for b in range(w):
            bit = bits[bit_pos] if bit_pos < total else 0
            re_bucket |= bit << b
            bit_pos += 1
        im_bucket = 0
        for b in range(w):
            bit = bits[bit_pos] if bit_pos < total else 0
            im_bucket |= bit << b
            bit_pos += 1
        # Convert Gray-coded bucket index back to linear before quantizing
        re_linear = from_gray(re_bucket)
        im_linear = from_gray(im_bucket)
        re = -1.0 + (re_linear + 0.5) * step
        im = -1.0 + (im_linear + 0.5) * step
        amps.append(complex(re, im))
        nrm = re*re + im*im
        norm_sq += nrm
        if nrm > n_max:
            i_max = i
            n_max = nrm

    norm = math.sqrt(norm_sq)
    return [a / norm for a in amps], norm, i_max


def amps_to_bits_gray(amps, w, n_bits):
    levels = 1 << w
    step   = 2.0 / levels
    bits   = []

    for amp in amps:
        if len(bits) >= n_bits:
            break
        re_linear = int((amp.real + 1.0) / step)
        re_linear = max(0, min(levels - 1, re_linear))
        im_linear = int((amp.imag + 1.0) / step)
        im_linear = max(0, min(levels - 1, im_linear))
        re_gray = to_gray(re_linear)
        im_gray = to_gray(im_linear)
        for b in range(w):
            if len(bits) >= n_bits: break
            bits.append((re_gray >> b) & 1)
        for b in range(w):
            if len(bits) >= n_bits: break
            bits.append((im_gray >> b) & 1)

    return bits


# ---------------------------------------------------------------------------
# Encoding: 'radius' mode
# Data is encoded in the MAGNITUDE (radius) of each amplitude only.
# Each amplitude gets w bits of data packed into its radius via Gray code,
# and a uniformly random phase factor.
#
# This makes Z-basis probabilities |a_i|^2 the sole data carrier.
# Phases are nonphysical noise — randomized at encode, ignored at decode.
# Result: Z-basis prob. overlap survives separate() and global phase
# scrambling with near-perfect fidelity.
#
# Data density: w bits per amplitude (vs 2*w for gray/plane).
# Decode: read |a_i|, recover Gray bucket, extract w bits.
# ---------------------------------------------------------------------------

def bits_to_amps_radius(bits, w, n_amps):
    """
    Encode w bits per amplitude in the radius.
    Radius is mapped from Gray bucket to [step/2, 1] (avoiding zero).
    Phase is uniformly random — carries no data.
    """
    levels  = 1 << w
    levels_m_1 = levels - 1
    n_bits  = w * n_amps
    amps    = []
    norm_sq = 0.0
    bit_pos = 0
    total   = len(bits)
    i_max   = 0
    n_max   = 0.0

    for i in range(n_amps):
        bucket = 0
        for b in range(w):
            bit = bits[bit_pos] if bit_pos < total else 0
            bucket |= bit << b
            bit_pos += 1
        # Gray -> linear bucket index
        linear = from_gray(bucket)
        # Map to radius in (0, 1]: r = (linear + 0.5) / levels
        # This keeps all radii strictly positive (no zero-magnitude amplitudes)
        # and uniformly spaced in probability space after squaring.
        r = (linear + 0.5) / levels
        # Random phase — carries no data, makes state look Haar-random
        phase = random.uniform(0.0, 2.0 * math.pi)
        re = r * math.cos(phase)
        im = r * math.sin(phase)
        amps.append(complex(re, im))
        nrm = r * r
        norm_sq += nrm
        if nrm > n_max:
            i_max = i
            n_max = nrm

    norm = math.sqrt(norm_sq)
    return [a / norm for a in amps], norm, i_max


def amps_to_bits_radius(amps, w, n_bits, norm):
    """
    Decode w bits per amplitude from the radius |a_i|.
    Phase is ignored entirely.
    """
    levels = 1 << w
    bits   = []

    for amp in amps:
        if len(bits) >= n_bits:
            break
        r = abs(amp) * norm
        # Recover linear bucket: r = (linear + 0.5) / levels  =>
        # linear = floor(r * levels)  clamped to [0, levels-1]
        linear = int(r * levels)
        linear = max(0, min(levels - 1, linear))
        bucket = to_gray(linear)
        for b in range(w):
            if len(bits) >= n_bits: break
            bits.append((bucket >> b) & 1)

    return bits


# ---------------------------------------------------------------------------
# Encoding: 'sign' mode
# w=1: only the sign of each coordinate encodes data (1 bit re, 1 bit im).
# Most robust: a sign flip requires quantization error > amplitude magnitude
# (~1/sqrt(n_amps)), which is much larger than the typical quant step.
# Data density: 2 bits per amplitude (vs 2*w in gray mode).
# ---------------------------------------------------------------------------

def bits_to_amps_sign(bits, n_amps):
    """Encode 2 bits per amplitude via sign of re and im."""
    amps    = []
    norm_sq = 0.0
    bit_pos = 0
    total   = len(bits)
    # Fixed magnitude per coordinate; only sign carries data.
    # Use 1/sqrt(2) so each amplitude has unit magnitude before normalization.
    mag = 1.0 / math.sqrt(2.0)

    for _ in range(n_amps):
        re_bit = bits[bit_pos] if bit_pos < total else 0; bit_pos += 1
        im_bit = bits[bit_pos] if bit_pos < total else 0; bit_pos += 1
        re = mag if re_bit else -mag
        im = mag if im_bit else -mag
        amps.append(complex(re, im))
        norm_sq += re*re + im*im

    norm = math.sqrt(norm_sq)
    return [a / norm for a in amps], norm, 0


def amps_to_bits_sign(amps, n_bits, norm):
    bits = []
    for amp in amps:
        amp *= norm
        if len(bits) >= n_bits: break
        bits.append(1 if amp.real >= 0.0 else 0)
        if len(bits) >= n_bits: break
        bits.append(1 if amp.imag >= 0.0 else 0)
    return bits


# ---------------------------------------------------------------------------
# Encoding: 'plane' mode
# Gray code + bit-plane interleaving.
# Instead of packing all w bits of one amplitude together, group all amplitudes'
# bit-k together for each k in 0..w-1.  A quantization error in one amplitude
# affects one bit in one plane rather than w adjacent bits in the stream.
# Same data density as 'gray'; error rate per bit is identical, but error
# locality is improved (burst errors within one amplitude are spread across
# w different positions in the recovered bit stream).
# ---------------------------------------------------------------------------

def bits_to_amps_plane(bits, w, n_amps):
    """
    Bit-plane interleaved encoding.
    Layout: [bit0 of re for all amps] [bit0 of im for all amps]
            [bit1 of re for all amps] ... [bit(w-1) of im for all amps]
    Gray coded within each plane.
    """
    levels  = 1 << w
    step    = 2.0 / levels
    total   = len(bits)
    n_bits  = 2 * w * n_amps

    # Build (n_amps, w) arrays of re_gray_bits and im_gray_bits
    re_gray_bits = [[0]*w for _ in range(n_amps)]
    im_gray_bits = [[0]*w for _ in range(n_amps)]

    bit_pos = 0
    for plane in range(w):
        for i in range(n_amps):
            re_gray_bits[i][plane] = bits[bit_pos] if bit_pos < total else 0
            bit_pos += 1
        for i in range(n_amps):
            im_gray_bits[i][plane] = bits[bit_pos] if bit_pos < total else 0
            bit_pos += 1

    amps    = []
    norm_sq = 0.0
    i_max = 0
    n_max = 0.0
    for i in range(n_amps):
        re_gray = sum(re_gray_bits[i][b] << b for b in range(w))
        im_gray = sum(im_gray_bits[i][b] << b for b in range(w))
        re_lin  = from_gray(re_gray)
        im_lin  = from_gray(im_gray)
        re = -1.0 + (re_lin + 0.5) * step
        im = -1.0 + (im_lin + 0.5) * step
        amps.append(complex(re, im))
        nrm = re*re + im*im
        norm_sq += nrm
        if nrm > n_max:
            i_max = i
            n_max = nrm

    norm = math.sqrt(norm_sq)
    return [a / norm for a in amps], norm, i_max


def amps_to_bits_plane(amps, w, n_bits, n_amps, norm):
    levels = 1 << w
    step   = 2.0 / levels

    re_gray_bits = []
    im_gray_bits = []
    for amp in amps:
        amp *= norm
        re_lin = max(0, min(levels-1, int((amp.real + 1.0) / step)))
        im_lin = max(0, min(levels-1, int((amp.imag + 1.0) / step)))
        rg = to_gray(re_lin)
        ig = to_gray(im_lin)
        re_gray_bits.append([(rg >> b) & 1 for b in range(w)])
        im_gray_bits.append([(ig >> b) & 1 for b in range(w)])

    bits    = []
    for plane in range(w):
        for i in range(n_amps):
            if len(bits) >= n_bits: break
            bits.append(re_gray_bits[i][plane])
        for i in range(n_amps):
            if len(bits) >= n_bits: break
            bits.append(im_gray_bits[i][plane])

    return bits[:n_bits]


# ---------------------------------------------------------------------------
# Main benchmark
# ---------------------------------------------------------------------------

def bench_qrack(width, p, b, w, mode='gray', do_separate=True):
    n_amps = 1 << width

    if mode == 'sign':
        n_bits = 2 * n_amps          # 1 bit per coordinate
    elif mode == 'radius':
        n_bits = w * n_amps          # w bits per amplitude (magnitude only)
    else:
        n_bits = 2 * w * n_amps      # w bits per coordinate (re and im)

    orig_bits = [random.randint(0, 1) for _ in range(n_bits)]

    # Encode
    if mode == 'sign':
        amps, norm, key = bits_to_amps_sign(orig_bits, n_amps)
    elif mode == 'radius':
        amps, norm, key = bits_to_amps_radius(orig_bits, w, n_amps)
    elif mode == 'plane':
        amps, norm, key = bits_to_amps_plane(orig_bits, w, n_amps)
    else:  # 'gray'
        amps, norm, key = bits_to_amps_gray(orig_bits, w, n_amps)

    sim = QrackSimulator(width)
    sim.in_ket(amps)
    if do_separate:
        sim.separate(list(range(width >> 1)))
    sim.lossy_out_to_file("lda.svtq", p=p, b=b)

    szd = n_bits / 8
    szf = os.path.getsize("lda.svtq") + 8
    print(f"Mode:                     {mode}")
    print(f"Data size:                {szd:.0f} bytes")
    print(f"Compressed size (+ norm): {szf} bytes")
    print(f"Compression ratio:        {szd / szf:.4f}x")

    # Decompress
    sim2 = QrackSimulator(width)
    sim2.lossy_in_from_file("lda.svtq")
    e_amps = sim2.out_ket()
    del sim2

    xeb, l2, prob_diff = calc_stats(amps, e_amps)
    print(f"Inner-product fidelity:   {l2:.6f}")
    print(f"XEB fidelity:             {xeb:.6f}")
    print(f"Z-basis prob. diff.:      {prob_diff:.6f}")

    # Hamming fidelity — decode recovered amplitudes back to bits
    if mode == 'radius':
        rec = amps_to_bits_radius(e_amps, w, n_bits, norm)
    elif mode == 'sign':
        rec = amps_to_bits_sign(e_amps, n_bits, norm)
    elif mode == 'plane':
        rec = amps_to_bits_plane(e_amps, w, n_bits, n_amps, norm)
    else:
        rec = amps_to_bits_gray(e_amps, w, n_bits)
    hf, errors = hamming_fidelity(orig_bits, rec)
    print(f"Hamming fidelity:         {hf:.6f}  ({errors} errors / {n_bits} bits)")

    return {
        "mode":              mode,
        "compression_ratio": szd / szf,
        "l2":                l2,
        "xeb":               xeb,
        "prob_diff":         prob_diff,
        "hamming":           hf,
    }


def main():
    width = 16
    if len(sys.argv) > 1:
        width = int(sys.argv[1])
    p = 6
    if len(sys.argv) > 2:
        p = int(sys.argv[2])
    b = 4
    if len(sys.argv) > 3:
        b = int(sys.argv[3])
    w = 12
    if len(sys.argv) > 4:
        w = int(sys.argv[4])
    mode = 'plane'
    if len(sys.argv) > 5:
        mode = sys.argv[5]
    do_separate = False
    if len(sys.argv) > 6:
        do_separate = sys.argv[6].lower() not in ('0', 'false', 'no')

    result = bench_qrack(width, p, b, w, mode, do_separate)
    print(result)
    return 0


if __name__ == "__main__":
    sys.exit(main())