Writing a block in Python

A block of a chain is any callable f(frame, rng) that reads and writes the samples of a comnum.Frame, or any subclass of comnum.Block. Python blocks prototype an impairment or an algorithm with NumPy, run inside the same chains as the native blocks, and are checked against them trial for trial. This tutorial writes a transmitter IQ imbalance and its blind compensation, the Gram–Schmidt orthogonalisation procedure (GSOP) of Fatadin, Savory & Ives, IEEE Photon. Technol. Lett. 20(20), 2008, and compares both with the native blocks.

import matplotlib.pyplot as plt
import numpy as np

import comnum
from comnum import blocks as B

The frame. A Python block sees the current samples (frame.samples, a complex array, or frame.streams with several antennas or polarisations) and the samples per symbol (frame.sps), and may read what the source recorded: the transmitted bits, symbols and constellation. frame.record(name, value) reports a value that the runners estimate as a mean. The simplest block is a function:

def invert(frame, rng):
    """A sign flip: a deterministic block that draws nothing."""
    frame.samples = -frame.samples

An impairment as a block. A gain and phase mismatch between the I and Q paths of a transmitter turns \(x\) into \(y = x + \beta x^*\) (the image of the signal leaks into it, at \(-20\log_{10}|\beta|\) dB). A class carries the parameter; the frame keeps everything else.

class TxIqImbalance(comnum.Block):
    """y = x + β x*."""

    def __init__(self, beta):
        self.beta = complex(beta)

    def __call__(self, frame, rng):
        x = frame.samples
        frame.samples = x + self.beta * np.conj(x)

Checked against the native block. TxIq is the same model in Rust. The two chains below draw the same symbols (same seed), so the samples after the Python block and after the native block must agree to machine precision.

BETA = 0.1 + 0.15j
source = B.Source(constellation="16-QAM", n_symbols=4096)
python_chain = comnum.Chain(tx=source, iq=TxIqImbalance(BETA))
native_chain = comnum.Chain(tx=source, iq=B.TxIq(beta=[BETA.real, BETA.imag]))
a = python_chain({}, comnum.Rng(3)).samples
b = native_chain({}, comnum.Rng(3)).samples
print("Python and native IQ imbalance agree:", np.allclose(a, b))

Python and native IQ imbalance agree: True

An algorithm as a block. GSOP makes the I and Q components orthogonal and of equal power: \(I' = I/\sqrt{P_I}\), \(Q'' = Q - \rho I / P_I\) with \(\rho = \mathbb{E}[IQ]\), then \(Q' = Q''/\sqrt{P_{Q''}}\). It needs no training symbols, only the circularity of the transmitted constellation. The block records the normalised correlation it measured, which the campaign reports.

class Gsop(comnum.Block):
    """Gram–Schmidt orthogonalisation of I and Q (Fatadin, Savory & Ives 2008)."""

    def __call__(self, frame, rng):
        x = frame.samples
        i, q = x.real, x.imag
        p_i = np.mean(i**2)
        rho = np.mean(i * q)
        q = q - rho / p_i * i
        i = i / np.sqrt(p_i)
        q = q / np.sqrt(np.mean(q**2))
        frame.record("iq_correlation", rho / p_i)
        frame.samples = (i + 1j * q) / np.sqrt(2)

The constellations after the impairment and after the compensation, at 25 dB.

link = comnum.Chain(
    tx=source,
    iq=TxIqImbalance(BETA),
    channel=B.Awgn(snr_db=comnum.var("snr_db")),
    gsop=Gsop(),
)
frames = link.trace({"snr_db": 25.0}, comnum.Rng(1))
fig, axes = plt.subplots(1, 3, figsize=(10, 3.6))
for ax, (name, title) in zip(
    axes,
    {"tx": "transmitted", "channel": "with IQ imbalance", "gsop": "after GSOP"}.items(),
    strict=True,
):
    comnum.plot.constellation(frames[name], ax=ax, max_points=4000)
    ax.set_title(title)
fig.tight_layout()

The campaign. A chain with a Python block runs on the Python runner (comnum.sim.run(), sequential), with the same grid, metrics, stopping rule and random streams as the Rust runner. comnum.compare() gathers the variants; the recorded value appears as a metric.

campaign = {
    "grid": {"snr_db": [10.0, 14.0, 18.0, 22.0, 26.0]},
    "metrics": ["ber"],
    "stop": {"min_errors": {"ber": 200}, "max_trials": 100},
    "seed": 1,
}
variants = {
    "no imbalance": comnum.Chain(tx=source, channel=B.Awgn(snr_db=comnum.var("snr_db"))),
    "imbalance": link.replace(gsop=invert).replace(gsop=lambda f, r: None),
    "imbalance + GSOP": link,
}
res = comnum.compare(variants, **campaign)
rho = res.where(variant="imbalance + GSOP", snr_db=26.0)[0]["metrics"]["iq_correlation"]
expected = 2 * BETA.imag / ((1 + BETA.real) ** 2 + BETA.imag**2)
print(f"measured I/Q correlation at 26 dB: {rho['value']:.3f}, expected {expected:.3f}")

fig, ax = plt.subplots()
res.plot("snr_db", "ber", by="variant", ax=ax,
         theory=lambda s, _: comnum.theory_ber(16, s - comnum.lin2db(4), "qam"))  # fmt: skip
ax.set(xlabel="SNR per symbol (dB)", ylabel="bit error rate", ylim=(1e-5, 0.5),
       title="16-QAM, β = 0.1 + 0.15j; dotted: closed form")  # fmt: skip
fig.tight_layout()

measured I/Q correlation at 26 dB: 0.245, expected 0.243

Parity with the native runner. The two runners draw trial t of condition i from the same stream, so a Python block that leaves the frame unchanged gives the native campaign back, error for error: the way to check a Python prototype against the native block that replaces it.

native = comnum.Chain(tx=source, channel=B.Awgn(snr_db=comnum.var("snr_db")))
python = native.insert_after("tx", noop=lambda frame, rng: None)
counts = {
    "Rust runner": [r["metrics"]["ber"]["count"] for r in native.run(**campaign)],
    "Python runner": [r["metrics"]["ber"]["count"] for r in python.run(**campaign)],
}
print(counts)
assert counts["Rust runner"] == counts["Python runner"]

{'Rust runner': [31358, 4921, 222, 0, 0], 'Python runner': [31358, 4921, 222, 0, 0]}

Beyond the prototype. The Python runner is sequential: once a block is settled, port it to Rust (a Block of comnum-spec, or an extension registered with comnum_spec::blocks::register, see docs/ARCHITECTURE.md), keep the Python version as its test oracle, and the chain runs on every core.