diff --git a/.claude/skills/comnumpy-tutorial/SKILL.md b/.claude/skills/comnumpy-tutorial/SKILL.md index 1d3c763..e9b6121 100644 --- a/.claude/skills/comnumpy-tutorial/SKILL.md +++ b/.claude/skills/comnumpy-tutorial/SKILL.md @@ -55,7 +55,13 @@ then learns the wrong lesson -- they came to see how the library is used. `compute_metric_rayleigh_theo` for the closed forms, `constellation_capacity` and `bicm_capacity` for the rates. - Monte-Carlo: `monte_carlo(chain, param, values, metrics, stimulus, seed=)` - when one chain and standard metrics cover the study. When the study + when one chain and standard metrics cover the study. Repeated trials + of **one** chain are not a Python loop: grow the stimulus a leading + batch axis -- `(n_trials, N)` -- and every block broadcasts over it, + stochastic blocks draw independently per trial, and the pooled metric + is exactly the mean of the per-trial rates (D51; the exception is a + channel whose realization is frozen at construction -- a new draw per + trial is a new instance, and that loop stays visible). When the study is *around* a chain -- several detectors on one frame, several receivers on one propagation -- write the loop out, on **one storage convention**: declare the methods and the metrics first; pre-allocate diff --git a/CHANGELOG.md b/CHANGELOG.md index 84f544f..4cdfa83 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,157 @@ one release; there is no compatibility layer. | `core.metrics.calculate_acpr` | `compute_acpr` — it was the only `calculate_*` in the library, against 17 `compute_*` | | `core.metrics.compute_effective_SNR`, `ofdm.metrics.compute_PAPR` | `compute_effective_snr`, `compute_papr` — the two capitalized outliers among functions otherwise all lowercase (`compute_ser`, `compute_ber`, `compute_evm`, `compute_ccdf`, `compute_mi`) | +### Added — the batch promise is a ratchet over the catalogue + +A convention over 91 blocks is only worth the sweep that checks it. +`tests/test_batch_contract.py` discovers every `Processor` subclass in +the library and fails unless each one is declared and verified: +`BROADCAST` (row `i` of a batch equals the block on row `i` alone -- +36 blocks checked by equality), `INDEPENDENT` (trials do not share a +realization), `REFUSES` (the ambiguous batch raises), or `EXEMPT` with +the reason written next to it. Like the pyright ratchet (D37), it makes +the contract self-enforcing: a new block cannot land without saying -- +and proving -- what a batch means for it. + +Its first sweep caught two more silent couplings. `DCCorrector`'s +default `axis=0` averaged *across* the batch (the default is now `-1`, +the canonical serial layout: one offset per converter). `HardClipper` +computed its threshold from the mean power of the whole array, so the +weakest trial of a batch was clipped against the strongest; the power +is now per row. + +### Added — the MIMO Monte-Carlo is a stacked channel (D51b) + +The MIMO question the batch contract had to answer: the channel is +redrawn every frame, and every detector needs that frame's matrix. +The answer is a **stack**. `rayleigh_channel(size=K)` returns +`(K, N_r, N_t)` -- drawn sequentially from one generator, so draw k +equals the k-th draw of the equivalent Python loop; `FlatMIMOChannel` +applies channel k to frame k (numpy's `matmul` batches over leading +axes); and a detector built with the same stack decides frame k against +channel k. `LinearDetector` (ZF and MMSE) batches through numpy's +*stacked linear algebra* -- `mmse_estimator` now transposes with +`swapaxes(-1, -2)` so its normal equations batch too -- while the +search detectors (ML, sphere, OSIC) loop per draw internally: an +enumeration, a tree and a cancellation order each depend on the one +matrix in front of them, so their batch is convenience and correctness, +not a speedup, and their docstrings say so. A mismatched stack (K +channels against K+1 frames) is refused. Locked by +`tests/mimo/test_stacked_channel.py`. + +Both MIMO studies now run on the stack. `one_shot_mimo` replaces its +zipped `monte_carlo(("channel.H", "detector.H"), ...)` sweep -- 200 +chain runs per detector per SNR -- with one batched pass (10.9 s to +9.0 s: the sphere decoder's per-sample search dominates and does not +vectorize). `monte_carlo_simulation_1` drops its 500-draw Python loop +for one stacked pass per SNR point and computes the BLER directly as +`np.mean(np.any(S_est != S_ref, axis=(-2, -1)))` -- 4.7 s to 2.5 s. + +### Added — batch axes are a contract, not an accident (D51) + +Leading axes ahead of a block's event axes are batch axes: independent +trials, run in one call. The contract has three families, now stated in +CONVENTIONS.md and locked by `tests/core/test_batch_axes.py`: +deterministic blocks broadcast their one configuration (`FIRChannel` +reshapes its kernel so scipy convolves along the last axis; +`LinearEqualizer` applies its matrix as a column product -- both +previously crashed on a batch); stochastic blocks draw independently +per event (`PhaseNoise` declares its event with `per=`: `"pair"` -- +default, one laser per polarization pair, independent across trials -- +`"row"` or `"signal"`; the ambiguous shape is refused with the +resolutions named); adaptive blocks carry independent state per event +(`BlindDualMIMOCompensator` on `(..., 2, N)` adapts one butterfly per +pair and exposes `H_` with the batch axes in front). + +A sweep of the block catalogue then held every block to the contract. +Four more deterministic blocks crashed on a batch and now broadcast: +`BWFilter` (its 1-D guard was already unnecessary -- the FFT mask acts +on the last axis), `CFO` (it read `len(x)`, i.e. the *batch* size, as +the signal length), `Delay`, and both chromatic-dispersion FIR +compensators (the same scipy kernel-rank fix as `FIRChannel`). +`BlindCFOCompensator` keeps its documented event -- one oscillator per +signal or per polarization pair `(2, N)`, estimated jointly +(`test_estimand_scope`) -- and now *refuses* a wider batch instead of +silently smearing one scalar estimate over independent trials. + +What batching buys was measured, not asserted: x4.1 on the +single-carrier chain whose block ZF equalizer rebuilds its +pseudo-inverse per call (4 trials amortize it to one build per sweep +point), x2.2 on the OFDM chain (16 trials), and **x0.5 -- slower --** +on a plain AWGN chain, where Python overhead was never the cost and +the batch only buys bigger temporaries. Batch for correctness and for +amortizing per-call operator builds; not as a blanket speed knob. + +Two silent traps died on the way. `compute_ser`/`compute_ber` with +`axis=None` used to ravel *before* truncating to the common length, so +a batch whose rows carried a tail (an OFDM frame after a `full` +convolution) compared misaligned rows -- SER 0.47 where the truth was +0.02; both now truncate along the last axis first, and the pooled rate +is exactly the mean of the per-row rates. And the OFDM tutorial's +repeat-loop became one batched call: +`monte_carlo(chain, param, values, metrics, (n_trials, N), seed=…)`. + +### Added — `print(obj)` renders `info()` + +Any object that defines `info()` -- a channel, a constellation -- now +prints it: `Processor.__str__` (and `Constellation.__str__`) render the +dictionary as one `key: value` line per entry, and fall back to `repr` +when there is no `info()`. The tutorials' `for key, value in +channel.info().items()` loops become `print(channel)`. + +### Added — the blind coherent receiver: `PMDEmulator`, BPS, and its page + +The expert review named the two absences that dated the optical layer: +no polarization impairment, and no time-varying carrier recovery -- the +`Laser` block carried a Wiener linewidth that nothing in the library +could track. Both close here, with the tutorial that needs them. + +`PMDEmulator` (optical.channels) is the standard section emulator: K +Haar-uniform Jones rotations, each followed by a DGD of tau/sqrt(K) +applied exactly in the frequency domain -- randomly oriented sections +add in quadrature, so the declared tau is the **RMS** DGD of the +ensemble, whose distribution is Maxwellian with mean 0.921 tau (Poole & +Wagner). A 300-seed test measures the DGD from the eigenvalues of the +group-delay operator and pins both moments within 10 %. The emulator is +unitary -- energy conserved to machine precision, tested -- seeded, and +refuses a single polarization. +`PhaseNoise` now draws **one** walk along the last axis and shares it +across leading axes: the phase comes from one laser, and a polarization +pair sees the same phi[n]. + +`BlindPhaseSearchCompensator` (core.compensators) is Pfau's feedforward +blind phase search: B test phases over [0, pi/2), windowed decision +distance, per-symbol argmin, unwrapped modulo pi/2. Each row of a pair +gets its own trajectory, exposed as `phase_` (D23). The quadrant +ambiguity is documented as unresolved, and tested as such. + +The page (`optical_pdm`) runs PDM-QPSK through one laser, eight PMD +sections and an amplifier, then undoes it blindly: CMA butterfly, BPS, +and the three ambiguities every blind receiver leaves -- permutation, +quadrant, equalizer delay -- resolved explicitly by known data, printed +rather than hidden. Zero errors over 123 072 symbols, and the page says +what that measures: transparency of the chain, not an error rate. + +### Added — the Gray labelling is locked by tests + +`compute_ber` expands indices in natural binary and its docstring +admitted the count is only meaningful if the constellation is labelled +accordingly. It is -- and now five tests prove it rather than the +tutorials' tail agreement suggesting it: geometric nearest neighbours +of every Gray QAM/PSK alphabet differ by exactly one index bit, natural +binary labelling fails that property (so the test measures the +labelling, not the geometry), a nearest-neighbour symbol error costs +exactly one bit through `compute_ber`, and a seeded 16-QAM run at +Eb/N0 = 11 dB lands within 20 % of the Gray closed form -- a labelling +mismatch would double it. + +### Changed — the AWGN page drops its interpolation coda + +The Eb/N0-at-BER-1e-3 table was five lines of `np.interp` on a fine +grid with an underflow guard -- machinery out of proportion with what +it added to the figure the section already shows. The spacing of the +waterfalls is now read off the figure in one sentence. + ### Changed — the simulation sections read identically across the pages Every sweep in the tutorials now shows the same three markers, in the diff --git a/CONVENTIONS.md b/CONVENTIONS.md index 27a1da7..fb1ef1f 100644 --- a/CONVENTIONS.md +++ b/CONVENTIONS.md @@ -51,6 +51,51 @@ Every block belongs to exactly one category, stated in its docstring The former `is_mimo` attribute — a shape-guessing mechanism — is removed; the declared category and `prepare()` validation replace it. +## Batch axes and event axes (D51) + +A block's category names its **event axes** — the trailing axes that +make up *one* realization of what the block models: nothing for an +element-wise block, `(N,)` for a filter, `(2, N)` for a polarization +pair, `(ant, N)` for a MIMO frame. Every axis to the left of the event +is a **batch axis**: independent trials, run in one call. The contract +has three families: + +1. **Deterministic** blocks broadcast their one configuration over the + batch — every trial goes through the same taps, the same equalizer + matrix (`FIRChannel`, `SRRCFilter`, `LinearEqualizer`, the OFDM + chains). Row `i` of a batched call equals the block applied to row + `i` alone. +2. **Stochastic** blocks draw **independently per event** — trials must + not share a realization. `AWGN` draws element-wise; `PhaseNoise` + declares its event with `per=`: `"pair"` (default — one laser per + `(2, N)` block, independent across the batch), `"row"`, or + `"signal"`. A shape the declaration cannot read one way is refused + with the resolutions named. + The exception is the block whose realization is **frozen at + construction** (`PMDEmulator`, `TappedDelayLineChannel`, + `FlatMIMOChannel`): one seeded draw *is* the configuration, so the + whole batch sees the same channel — ergodic noise over a fixed + channel. A new draw per trial is a new instance (or `set_params`), + and that loop stays visible. +3. **Adaptive** blocks carry **independent state per event** — one + equalizer per pair, never one state smeared across trials. + `BlindDualMIMOCompensator` on `(..., 2, N)` adapts one butterfly per + pair and exposes `H_` with the batch axes in front; + `BlindPhaseSearchCompensator` tracks one `phase_` trajectory per + row. + +The payoff is Monte-Carlo without a Python loop over trials: +`monte_carlo(chain, param, values, metrics, (n_trials, N), seed=…)` +pools each metric over `n_trials` independent frames per sweep point — +for equal-size trials, exactly the mean of the per-trial values. +`tests/core/test_batch_axes.py` locks one block of each family, and +`tests/test_batch_contract.py` is the **ratchet over the whole +catalogue**: it discovers every `Processor` subclass and fails unless +each is verified batched (`BROADCAST` by row equality, `INDEPENDENT` +by realization inequality, `REFUSES` by the raised error) or exempted +with a written reason. A new block cannot merge without declaring what +a batch means for it. + ## Observing signals: taps, not blocks A chain describes the communication system. Nothing else goes into diff --git a/docs/getting_started/first_simulation.rst b/docs/getting_started/first_simulation.rst index 3e63e2b..fb3fa26 100644 --- a/docs/getting_started/first_simulation.rst +++ b/docs/getting_started/first_simulation.rst @@ -146,18 +146,17 @@ itself: .. literalinclude:: ../../examples/simple/one_shot_awgn.py :language: python - :lines: 84-87 + :lines: 84-88 .. code:: - constellation.info() - family QAM - order 4 - labelling gray - bits_per_symbol 2 - energy 1.0 - min_distance 1.414213562373095 - papr_dB 0.0 + family: QAM + order: 4 + labelling: gray + bits_per_symbol: 2 + energy: 1.0 + min_distance: 1.414213562373095 + papr_dB: 0.0 Every one of those was implicit in the by-hand version. ``energy`` is the normalization the noise variance depends on; ``bits_per_symbol`` is the @@ -170,7 +169,7 @@ It also draws itself, next to what the channel did to it: .. literalinclude:: ../../examples/simple/one_shot_awgn.py :language: python - :lines: 89-97 + :lines: 90-98 .. image:: img/first_simulation_fig1.png :width: 100% diff --git a/docs/tutorials/awgn.rst b/docs/tutorials/awgn.rst index a9bad99..6d416c2 100644 --- a/docs/tutorials/awgn.rst +++ b/docs/tutorials/awgn.rst @@ -110,7 +110,7 @@ outline marks a tapped block: .. literalinclude:: ../../examples/simple/monte_carlo_awgn.py :language: python - :lines: 144-150 + :lines: 129-135 Monte Carlo Simulation """""""""""""""""""""" @@ -272,7 +272,7 @@ not guess (decision D41), so it happens at the call site: .. literalinclude:: ../../examples/simple/monte_carlo_awgn.py :language: python - :lines: 99-128 + :lines: 99-120 The metric changes with the axis. Against :math:`E_b/N_0` the natural quantity is the bit error rate, and ``compute_ber`` needs the symbol width, which @@ -289,27 +289,17 @@ below :math:`10^{-2}` the two agree to within a few percent for every order. This is an approximation behaving as documented, not a simulation disagreeing with theory. -.. code:: - - QAM order bits per symbol Eb/N0 at BER=1e-3 [dB] - -------------------------------------------------- - 4 2 6.79 - 16 4 10.52 - 64 6 14.77 - 256 8 19.38 - .. literalinclude:: ../../examples/simple/monte_carlo_awgn.py :language: python - :lines: 130-142 + :lines: 120-127 .. image:: img/monte_carlo_awgn_orders.png :width: 100% :align: center -Going from 2 to 8 bits per symbol costs **12.6 dB** of energy per bit at -:math:`10^{-3}`, spent in three roughly equal steps -- 3.7, 4.3 then 4.6 dB -per doubling of :math:`k`, the price rising slowly with the order. That is -what the rest of +The waterfalls are almost evenly spaced: each doubling of the bits per +symbol shifts the curve by **4 to 5 dB** of energy per bit, read directly +off the figure. That price is what the rest of a communication system is built to negotiate: coding buys some of it back (:doc:`coding`), shaping buys a little more (:doc:`shaping`), and a channel that is not simply AWGN changes the terms entirely (:doc:`ofdm`). diff --git a/docs/tutorials/img/monte_carlo_mimo_fig3.png b/docs/tutorials/img/monte_carlo_mimo_fig3.png index f85f97a..643dd82 100644 Binary files a/docs/tutorials/img/monte_carlo_mimo_fig3.png and b/docs/tutorials/img/monte_carlo_mimo_fig3.png differ diff --git a/docs/tutorials/img/monte_carlo_mimo_fig4.png b/docs/tutorials/img/monte_carlo_mimo_fig4.png index 9f41e21..aa506e0 100644 Binary files a/docs/tutorials/img/monte_carlo_mimo_fig4.png and b/docs/tutorials/img/monte_carlo_mimo_fig4.png differ diff --git a/docs/tutorials/img/one_shot_ofdm_fig4.png b/docs/tutorials/img/one_shot_ofdm_fig4.png index 0635ac7..9cc15a1 100644 Binary files a/docs/tutorials/img/one_shot_ofdm_fig4.png and b/docs/tutorials/img/one_shot_ofdm_fig4.png differ diff --git a/docs/tutorials/img/one_shot_ofdm_fig5.png b/docs/tutorials/img/one_shot_ofdm_fig5.png index 7d44490..6309ce4 100644 Binary files a/docs/tutorials/img/one_shot_ofdm_fig5.png and b/docs/tutorials/img/one_shot_ofdm_fig5.png differ diff --git a/docs/tutorials/img/one_shot_pdm_fig1.png b/docs/tutorials/img/one_shot_pdm_fig1.png new file mode 100644 index 0000000..2585b5a Binary files /dev/null and b/docs/tutorials/img/one_shot_pdm_fig1.png differ diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index 7115774..8e929ac 100644 --- a/docs/tutorials/index.rst +++ b/docs/tutorials/index.rst @@ -51,6 +51,11 @@ and nothing is explained twice. - A nonlinear channel: how much noise the fibre's nonlinearity makes, in closed form, and the launch power that minimizes it -- checked against a split-step simulation. + * - 9 + - :doc:`optical_pdm` + - The blind coherent receiver: a fibre that mixes the two + polarizations (``PMDEmulator``), the CMA butterfly that + unmixes them, and blind phase search against the laser. * - 9 - :doc:`optical_fiber_nonlinearity` - The performance the previous tutorial predicted, improved: digital @@ -74,4 +79,5 @@ The scripts behind these pages are also available on GitHub: `https://github.com coding shaping gn_model + optical_pdm optical_fiber_nonlinearity diff --git a/docs/tutorials/mimo.rst b/docs/tutorials/mimo.rst index 06e6cbe..d9c13ed 100644 --- a/docs/tutorials/mimo.rst +++ b/docs/tutorials/mimo.rst @@ -69,7 +69,7 @@ block: .. literalinclude:: ../../examples/mimo/one_shot_mimo.py :language: python - :lines: 182-185 + :lines: 185-188 Each chain is given the same seed before running, so the five numbers below differ by the detector alone -- same symbols, same noise, same channel: @@ -226,27 +226,30 @@ Monte Carlo Evaluation A single channel realization proves nothing: over fading, the error rate is an *average*, dominated by the rare draws where the matrix is nearly -singular. Averaging means running the chain once per realization, which is a -sweep whose parameter is the channel. :func:`~comnumpy.monte_carlo.sweep` takes -several dotted parameter names at once and zips them, so one sweep point sets -the matrix the signal goes through **and** the one the detector inverts: +singular. Averaging means one independent draw per frame, and the draws are +a **stack**: ``rayleigh_channel(size=n_channels)`` returns ``(200, 3, 2)``, +the channel block propagates frame :math:`k` through matrix :math:`k`, and +the detector -- handed the *same* stack -- decides frame :math:`k` against +matrix :math:`k`. The chain runs once per SNR point instead of once per +draw, and the pooled SER over equal-size frames is exactly the mean of the +per-draw rates: .. literalinclude:: ../../examples/mimo/one_shot_mimo.py :language: python - :lines: 83-129 + :lines: 83-132 .. code:: SER snr_dB ZF MMSE OSIC ML SD --------------------------------------------------------- - 0 0.33129 0.28710 2.826e-01 2.807e-01 2.807e-01 - 3 0.21064 0.17623 1.610e-01 1.560e-01 1.560e-01 - 6 0.10650 0.08759 6.679e-02 6.196e-02 6.196e-02 - 9 0.04285 0.03451 1.920e-02 1.588e-02 1.588e-02 - 12 0.01485 0.01159 3.738e-03 2.737e-03 2.737e-03 - 15 0.00474 0.00363 7.125e-04 5.500e-04 5.500e-04 - 18 0.00203 0.00142 3.375e-04 2.625e-04 2.625e-04 + 0 0.33086 0.28510 2.822e-01 2.791e-01 2.791e-01 + 3 0.20976 0.17386 1.595e-01 1.545e-01 1.545e-01 + 6 0.10669 0.08601 6.488e-02 5.980e-02 5.980e-02 + 9 0.04389 0.03476 1.890e-02 1.584e-02 1.584e-02 + 12 0.01510 0.01174 4.163e-03 2.488e-03 2.488e-03 + 15 0.00507 0.00394 8.000e-04 3.875e-04 3.875e-04 + 18 0.00201 0.00144 1.625e-04 1.625e-04 1.625e-04 .. image:: img/monte_carlo_mimo_fig3.png :width: 100% @@ -281,16 +284,16 @@ where that product no longer fits. We therefore time both detectors on .. literalinclude:: ../../examples/mimo/one_shot_mimo.py :language: python - :lines: 131-180 + :lines: 134-183 .. code:: - ML 0 dB 160.3 ms 65536 nodes SER 0.7900 - ML 9 dB 120.5 ms 65536 nodes SER 0.5044 - ML 18 dB 109.9 ms 65536 nodes SER 0.0112 - SD 0 dB 203.0 ms 91.8 nodes SER 0.7900 - SD 9 dB 42.0 ms 19.2 nodes SER 0.5044 - SD 18 dB 14.3 ms 5.8 nodes SER 0.0112 + ML 0 dB 142.0 ms 65536 nodes SER 0.7900 + ML 9 dB 123.5 ms 65536 nodes SER 0.5044 + ML 18 dB 140.8 ms 65536 nodes SER 0.0112 + SD 0 dB 197.5 ms 91.8 nodes SER 0.7900 + SD 9 dB 36.3 ms 19.2 nodes SER 0.5044 + SD 18 dB 11.8 ms 5.8 nodes SER 0.0112 .. image:: img/monte_carlo_mimo_fig4.png :width: 100% diff --git a/docs/tutorials/ofdm.rst b/docs/tutorials/ofdm.rst index a499bed..410a5e8 100644 --- a/docs/tutorials/ofdm.rst +++ b/docs/tutorials/ofdm.rst @@ -79,7 +79,7 @@ of it, and ``info()`` returns what the channel is: .. literalinclude:: ../../examples/ofdm/one_shot_ofdm.py :language: python - :lines: 31-54 + :lines: 31-53 .. code:: @@ -90,9 +90,12 @@ of it, and ``info()`` returns what the channel is: powers_dB: [0.0, -1.0, -2.0, -3.0, -8.0, -17.2, -20.8] rms_delay_spread_ns: 43.12922598416199 coherence_bandwidth_Hz: 4637226.739785324 + rice_k_dB: None fs_Hz: 7680000.0 + f_doppler_Hz: 0.0 n_taps: 4 resolvable_delays_samples: [0, 1, 3] + resolvable_powers: [0.5765216771059146, 0.420805847284595, 0.00267247560949035] The profile has seven arrivals within 410 ns, but at :math:`f_s = 7.68` MHz the receiver can only distinguish **4 taps**: two echoes separated by less @@ -136,14 +139,14 @@ Implementation .. literalinclude:: ../../examples/ofdm/one_shot_ofdm.py :language: python - :lines: 55-77 + :lines: 54-76 Results """"""" .. code:: - single carrier: SER 0.0117, 1829 ms + single carrier: SER 0.0117, 3669 ms .. image:: img/one_shot_ofdm_fig2.png :width: 100% @@ -153,7 +156,7 @@ Results On the left, the constellation has disappeared: each sample is a mixture of four consecutive symbols, and no decision region survives this. On the right, the equalizer has restored the sixteen clusters. Note the second number: -nearly two seconds for 1280 symbols. +more than three seconds for 1280 symbols. OFDM Communication Chain @@ -187,7 +190,7 @@ before the channel and a receiver block after it: .. literalinclude:: ../../examples/ofdm/one_shot_ofdm.py :language: python - :lines: 78-101 + :lines: 77-100 ``OFDMTransmitter`` and ``OFDMReceiver`` are themselves chains: serial-to-parallel conversion, subcarrier allocation, IFFT and cyclic prefix @@ -199,7 +202,7 @@ Results .. code:: - OFDM : SER 0.0406, 1.20 ms (1525 times faster) + OFDM : SER 0.0406, 1.00 ms (3670 times faster) .. image:: img/one_shot_ofdm_fig3.png :width: 100% @@ -218,21 +221,21 @@ range of SNR values: .. literalinclude:: ../../examples/ofdm/one_shot_ofdm.py :language: python - :lines: 102-132 + :lines: 101-128 .. code:: SER SNR [dB] single carrier OFDM -------------------------------- - 6 0.62734 0.5234 - 8 0.52695 0.3953 - 10 0.42969 0.2973 - 12 0.27070 0.1980 - 14 0.15117 0.1426 - 16 0.06562 0.0672 - 18 0.01523 0.0383 - 20 0.00195 0.0195 + 6 0.62578 0.5117 + 8 0.53750 0.3980 + 10 0.40664 0.3172 + 12 0.28164 0.1996 + 14 0.16875 0.1336 + 16 0.06172 0.0750 + 18 0.01641 0.0410 + 20 0.00352 0.0223 .. image:: img/one_shot_ofdm_fig4.png :width: 100% @@ -265,17 +268,17 @@ block length grows: .. literalinclude:: ../../examples/ofdm/one_shot_ofdm.py :language: python - :lines: 133-169 + :lines: 129-165 .. code:: receiver runtime [ms], and their ratio block length N single carrier OFDM ratio --------------------------------------------- - 128 8.74 0.613 14.2 - 256 32.64 0.644 50.7 - 512 209.35 0.999 209.5 - 1024 1685.07 1.123 1500.4 + 128 7.34 0.568 12.9 + 256 34.91 0.567 61.6 + 512 188.09 0.793 237.1 + 1024 1377.50 0.862 1598.3 .. image:: img/one_shot_ofdm_fig5.png :width: 100% diff --git a/docs/tutorials/optical_pdm.rst b/docs/tutorials/optical_pdm.rst new file mode 100644 index 0000000..5dff760 --- /dev/null +++ b/docs/tutorials/optical_pdm.rst @@ -0,0 +1,133 @@ +Polarization Demultiplexing: CMA and Blind Phase Search +======================================================= + +A coherent optical link transmits on the two polarizations of the fibre +at once. The fibre repays that by mixing them: its random birefringence +rotates the state of polarization by an unknown, wavelength-flat Jones +matrix, and polarization-mode dispersion (PMD) delays the two principal +states against each other, so each received polarization is a filtered +mix of both transmitted ones. On top of it, the transmit laser's phase +noise spins the whole constellation. None of these is known to the +receiver, and none is an obstacle: undoing all three blindly is the +standard coherent DSP chain, and this page builds it. + +**What you'll learn:** + +- How to emulate a fibre's polarization rotation and first-order PMD + with ``PMDEmulator``, and a laser's Wiener phase noise with + ``PhaseNoise``. +- How the blind 2x2 butterfly equalizer (CMA) separates the + polarizations without training. +- How blind phase search (BPS) tracks the laser phase, and which + ambiguities every blind receiver leaves behind. + +The problem +^^^^^^^^^^^ + +Two QPSK streams are shaped and sent through one laser, one fibre and +one amplifier. The phase walk is common to the two rows -- one laser -- +while the PMD emulator, a cascade of eight random unitary sections with +an RMS DGD of 10 ps (a third of the symbol at 32 GBd -- the section +delays add in quadrature, and the ensemble DGD is Maxwellian around +that value), mixes the rows into each other: + +.. literalinclude:: ../../examples/optical/one_shot_pdm.py + :language: python + :lines: 1-55 + +Sampling the received field naively at the symbol rate gives the left +panel of the figure below: a disc. The two polarizations sit on top of +each other with a relative delay, and the laser spins the sum -- there +is no constellation to see, and no error rate worth measuring. + +The method +^^^^^^^^^^ + +Two blind estimators, in the order a receiver runs them. + +The **constant modulus algorithm** (CMA) adapts a 2x2 butterfly of FIR +filters by gradient descent on the modulus error +:math:`\left(|y[n]|^2 - R^2\right)^2` (Savory, 2008). QPSK has constant +modulus, so any residual modulus variation is the channel's doing -- +the criterion needs no phase reference, which is what lets it converge +while the constellation still spins. + +**Blind phase search** (BPS, Pfau et al., 2009) then rotates each +output by a grid of test phases, decides each against the alphabet, and +retains per symbol the test phase of least windowed decision distance. +It is feedforward -- no loop to stabilize -- and estimates the phase +modulo :math:`\pi/2`, the symmetry of the constellation. + +The implementation and the measurement +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. literalinclude:: ../../examples/optical/one_shot_pdm.py + :language: python + :lines: 57-72 + +A blind receiver leaves exactly three things unresolved, and it is +worth seeing them listed rather than hidden: which output is which +polarization, a quadrant rotation per output, and the small group delay +of the equalizer. A deployed system resolves them with framing +(:doc:`../documentation/core/frames`); here the known data plays that +role, explicitly: + +.. literalinclude:: ../../examples/optical/one_shot_pdm.py + :language: python + :lines: 74-118 + +.. code:: + + output 0: polarization 0, rotated 270 deg, delayed -4 symbols + output 1: polarization 1, rotated 90 deg, delayed -4 symbols + SER after CMA + BPS: 0.00e+00 over 123072 symbols (0 errors) + +.. image:: img/one_shot_pdm_fig1.png + :width: 100% + :align: center + :alt: Received disc, ring after CMA, four QPSK clouds after BPS + +The figure is the whole argument, read left to right. The received +polarization is a disc -- two delayed streams plus a spinning phase. +After the CMA it is a ring: the modulus is restored, so the butterfly +has inverted the rotation and the DGD, but the criterion is blind to +phase and the laser still spins the ring. After the BPS the four clouds +sit on the reference crosses. + +Zero errors over 123 072 symbols does not measure an error rate -- at +this SNR the closed form predicts :math:`\sim 10^{-8}`, far below what +the run can resolve. What it measures is the claim this page makes: +the blind chain is **transparent**, from a disc to a decided +constellation, at the cost of 4 000 discarded convergence symbols and +three ambiguities resolved by framing. + +Conclusion +^^^^^^^^^^ + +You have learned how to: + +- Emulate polarization rotation, first-order PMD and laser phase noise + with three seeded blocks. +- Separate two polarizations blindly with the CMA butterfly, and read + its convergence cost. +- Track the carrier phase with BPS, and resolve the permutation, + quadrant and delay ambiguities a blind receiver always leaves. + +The natural next step is :doc:`optical_fiber_nonlinearity`, where the +fibre stops being unitary: attenuation, amplifier noise and the Kerr +effect set the limits this page's channel did not have. + +References +^^^^^^^^^^ + +S. J. Savory, "Digital filters for coherent optical receivers," +Optics Express, vol. 16, no. 2, pp. 804-817, 2008. + +T. Pfau, S. Hoffmann and R. Noe, "Hardware-Efficient Coherent Digital +Receiver Concept With Feedforward Carrier Recovery for M-QAM +Constellations," Journal of Lightwave Technology, vol. 27, no. 8, +pp. 989-999, 2009. + +C. D. Poole and R. E. Wagner, "Phenomenological approach to +polarisation dispersion in long single-mode fibres," Electronics +Letters, vol. 22, no. 19, pp. 1029-1030, 1986. diff --git a/examples/mimo/monte_carlo_simulation_1.py b/examples/mimo/monte_carlo_simulation_1.py index c318607..ec8d415 100644 --- a/examples/mimo/monte_carlo_simulation_1.py +++ b/examples/mimo/monte_carlo_simulation_1.py @@ -15,7 +15,6 @@ from comnumpy.core import Sequential from comnumpy.core.generators import SymbolGenerator from comnumpy.core.mappers import SymbolMapper -from comnumpy.core.metrics import compute_ber from comnumpy.core.utils import Constellation from comnumpy.mimo.channels import AWGN, FlatMIMOChannel from comnumpy.mimo.utils import rayleigh_channel @@ -50,34 +49,33 @@ bler[name] = np.zeros(len(snr_dB_list)) # --- simulation loop ------------------------------------------------- -# One child seed per SNR point (D6/D35): the whole figure is reproduced -# by the master seed alone. +# One child seed per SNR point (D6/D35), and one *stack* of channel +# draws per point (D51): channel k propagates frame k, and each +# detector, built with the same stack, decides frame k against channel +# k -- the comparison is over identical frames, and the chain runs once +# per point instead of once per draw. point_seeds = np.random.SeedSequence(seed).spawn(len(snr_dB_list)) for index, snr_dB in enumerate(snr_dB_list): rng = np.random.default_rng(int(point_seeds[index].generate_state(1)[0])) sigma2 = N_t * 10 ** (-snr_dB / 10) - chain.set_params(noise__sigma2=sigma2) + H = rayleigh_channel(N_r=N_r, N_t=N_t, rng=rng, size=N_test) + chain.set_params(noise__sigma2=sigma2, channel__H=H) + chain.seed(int(rng.integers(2 ** 31))) + Y = chain((N_test, N_t, N)) + S_ref = chain.tap("data_tx") - for _ in range(N_test): - # new channel realization, decided by every detector: the - # comparison is over identical frames - H = rayleigh_channel(N_r=N_r, N_t=N_t, rng=rng) - chain.seed(int(rng.integers(2 ** 31))) - chain.set_params(channel__H=H) - Y = chain((N_t, N)) - S_ref = chain.tap("data_tx") - - detectors = { - "ML": MaximumLikelihoodDetector(alphabet=constellation, H=H), - "ZF": LinearDetector(alphabet=constellation, H=H, method="zf"), - "OSIC": OrderedSuccessiveInterferenceCancellationDetector( - constellation, osic_type="sinr", H=H, sigma2=sigma2), - } - for name, detector in detectors.items(): - ber = compute_ber(S_ref, detector(Y), - width=constellation.bits_per_symbol) - # a frame is in error when at least one of its bits is - bler[name][index] += float(ber > 0) / N_test + detectors = { + "ML": MaximumLikelihoodDetector(alphabet=constellation, H=H), + "ZF": LinearDetector(alphabet=constellation, H=H, method="zf"), + "OSIC": OrderedSuccessiveInterferenceCancellationDetector( + constellation, osic_type="sinr", H=H, sigma2=sigma2), + } + for name, detector in detectors.items(): + S_est = detector(Y) + # a frame is in error when at least one of its symbols is -- + # with Gray labelling, exactly "at least one of its bits" + bler[name][index] = float(np.mean( + np.any(S_est != S_ref, axis=(-2, -1)))) # --- results: table and figure --------------------------------------- data = {"x": snr_dB_list, "curves": bler} diff --git a/examples/mimo/one_shot_mimo.py b/examples/mimo/one_shot_mimo.py index daf0dec..2bd1b81 100644 --- a/examples/mimo/one_shot_mimo.py +++ b/examples/mimo/one_shot_mimo.py @@ -3,7 +3,7 @@ import numpy as np import matplotlib.pyplot as plt -from comnumpy import monte_carlo, print_data, style +from comnumpy import print_data, style from comnumpy.core import Sequential from comnumpy.core.generators import SymbolGenerator from comnumpy.core.mappers import SymbolMapper @@ -89,21 +89,24 @@ def get_chain(detector): def average_ser(name, chain, snr_dB, seed=0): - """Average one chain over independent Rayleigh draws at one SNR.""" - rng = np.random.default_rng(seed) - draws = [] - for _ in range(n_channels): - matrix = rayleigh_channel(N_r, N_t, rng=rng) - draws.append((matrix, matrix)) + """Average one chain over independent Rayleigh draws at one SNR. + + The draws are a *stack* (D51): channel k propagates frame k, the + detector holds the same stack and decides frame k against channel + k, and the chain runs once instead of once per draw. Pooled over + equal-size frames, the SER is exactly the mean of the per-draw + rates. + """ + stacked = rayleigh_channel(N_r, N_t, seed=seed, size=n_channels) noise_variance = N_t * 10 ** (-snr_dB / 10) - params = {"noise.sigma2": noise_variance} + params = {"noise.sigma2": noise_variance, + "channel.H": stacked, "detector.H": stacked} if name in NOISE_AWARE: params["detector.sigma2"] = noise_variance chain.set_params(**params) - results = monte_carlo(chain, ("channel.H", "detector.H"), draws, - {"ser": compute_ser}, stimulus=(N_t, n_symbols), - reference="tx", seed=seed) - return float(np.mean(results["ser"])) + chain.seed(seed) + detected = chain((n_channels, N_t, n_symbols)) + return float(compute_ser(chain.tap("tx"), detected)) # --- metrics, pre-allocated ------------------------------------------ diff --git a/examples/ofdm/one_shot_ofdm.py b/examples/ofdm/one_shot_ofdm.py index 8ee4c4c..81c0c58 100644 --- a/examples/ofdm/one_shot_ofdm.py +++ b/examples/ofdm/one_shot_ofdm.py @@ -33,8 +33,7 @@ # realization of it. The channel says what it is rather than being # described from outside. channel = TappedDelayLineChannel(get_delay_profile("EPA"), fs=fs, seed=8) -for key, value in channel.info().items(): - print(f"{key}: {value}") +print(channel) # Sounding the model with an impulse gives the realization as a tap # vector -- which is what both receivers below are given. @@ -100,25 +99,22 @@ plt.savefig(f"{img_dir}/one_shot_ofdm_fig3.png") # --- error rate ------------------------------------------------------ -# One operating point is not a conclusion. Each sweep is repeated over -# independent noise seeds because 1280 symbols only resolve down to -# 8e-4; the channel is the same throughout, by construction. +# One operating point is not a conclusion, and 1280 symbols only +# resolve down to 8e-4. Instead of a Python loop over repeated runs, +# the stimulus grows a leading batch axis (D51): every block broadcasts +# over it, the noise is drawn independently per trial, and compute_ser +# pools the n_trials frames of each sweep point -- exactly the mean of +# the per-trial rates. The channel stays the same throughout, by +# construction. snr_list = np.arange(6, 22, 2) n_trials = 2 -# --- metrics, pre-allocated ------------------------------------------ -measured = {} -for name in ("single carrier", "OFDM"): - measured[name] = np.zeros(len(snr_list)) - # --- simulation loop ------------------------------------------------- +measured = {} for name, chain in (("single carrier", sc_chain), ("OFDM", ofdm_chain)): - runs = np.zeros((n_trials, len(snr_list))) - for trial in range(n_trials): - runs[trial] = monte_carlo(chain, "data_rx.snr_dB", snr_list, - {"ser": compute_ser}, N, - reference="data_tx", seed=trial + 1)["ser"] - measured[name] = np.mean(runs, axis=0) + measured[name] = monte_carlo(chain, "data_rx.snr_dB", snr_list, + {"ser": compute_ser}, (n_trials, N), + reference="data_tx", seed=1)["ser"] # --- results: table and figure --------------------------------------- diff --git a/examples/optical/one_shot_pdm.py b/examples/optical/one_shot_pdm.py new file mode 100644 index 0000000..50901d0 --- /dev/null +++ b/examples/optical/one_shot_pdm.py @@ -0,0 +1,118 @@ +"""Polarization demultiplexing of a PDM-QPSK link: CMA, then BPS. + +Run from this directory: it writes the tutorial's figures into +../../docs/tutorials/img/. +""" +import matplotlib.pyplot as plt +import numpy as np + +from comnumpy import style +from comnumpy.core import Sequential +from comnumpy.core.channels import AWGN +from comnumpy.core.compensators import BlindPhaseSearchCompensator +from comnumpy.core.filters import SRRCFilter +from comnumpy.core.generators import SymbolGenerator +from comnumpy.core.mappers import SymbolMapper +from comnumpy.core.metrics import compute_ser +from comnumpy.core.processors import Upsampler +from comnumpy.core.utils import Constellation, hard_projector +from comnumpy.core.visualizers import plot_iq +from comnumpy.mimo.compensators import BlindDualMIMOCompensator +from comnumpy.optical.channels import PhaseNoise, PMDEmulator + +style.use() + +img_dir = "../../docs/tutorials/img/" + +# parameters +constellation = Constellation("PSK", 4) +N = 2 ** 16 # symbols per polarization +oversampling = 2 +rolloff = 0.1 +R_s = 32e9 # baud rate +fs = R_s * oversampling +linewidth = 100e3 # laser linewidth, Hz +dgd = 10e-12 # total mean DGD, s -- a third of the symbol +snr_dB = 15 +n_discard = 4000 # symbols spent converging the equalizer + +# The channel: two polarizations shaped and sent through one laser's +# phase noise, the fibre's random rotation and first-order PMD, and the +# amplifier noise. The phase walk is common to the two rows -- one +# laser -- while the PMD mixes them into each other. +chain = Sequential([ + SymbolGenerator(constellation.order, name="data_tx"), + SymbolMapper(constellation, name="signal_tx"), + Upsampler(oversampling, scale=np.sqrt(oversampling)), + SRRCFilter(rolloff, oversampling, method="fft"), + PhaseNoise(2 * np.pi * linewidth / fs, name="laser"), + PMDEmulator(dgd, n_sections=8, fs=fs, seed=11, name="fibre"), + AWGN(snr_dB=snr_dB, name="noise"), + ], taps=["data_tx", "signal_tx", "noise"]) +chain.seed(0) +received = chain((2, N)) +sent = chain.tap("signal_tx") +reference = chain.tap("data_tx") + +# What the receiver faces, sampled naively at the symbol rate: the two +# polarizations are mixed by an unknown rotation, delayed against each +# other, and the constellation spins with the laser. +naive = received[:, ::oversampling] + +# The 2x2 butterfly equalizer, blind (CMA): it inverts the rotation and +# the DGD without knowing either, from the constant modulus of QPSK +# alone. The first symbols are its convergence and are discarded. +cma = BlindDualMIMOCompensator(L=9, alphabet=constellation, mu=1e-3, + oversampling=oversampling, name="cma") +equalized = cma(received)[:, n_discard:] + +# Feedforward carrier recovery (blind phase search): each polarization +# gets its own phase trajectory, estimated modulo pi/2. +bps = BlindPhaseSearchCompensator(constellation, name="bps") +recovered = bps(equalized) + +# --- results: the ambiguities, then the error rate --------------------- +# A blind receiver leaves three things unresolved: which output is +# which polarization, a quadrant rotation per output, and the small +# group delay of the equalizer. A system resolves them with framing; +# here the known data plays that role, explicitly, over a short probe. +probe = 2000 +detected_rows = [] +truth_rows = [] +for row in range(2): + best = None + for source in range(2): + for quadrant in range(4): + for lag in range(-8, 9): + start = n_discard + lag + if start < 0: + continue + candidate = recovered[row, :probe] * (1j) ** quadrant + target = sent[source, start:start + probe] + error = float(np.mean(np.abs(candidate - target) ** 2)) + if best is None or error < best[0]: + best = (error, source, quadrant, lag) + assert best is not None + _, source, quadrant, lag = best + print(f"output {row}: polarization {source}, rotated " + f"{90 * quadrant} deg, delayed {lag} symbols") + aligned = recovered[row] * (1j) ** quadrant + decisions, _ = hard_projector(aligned, constellation) + detected_rows.append(decisions) + length = decisions.shape[-1] + truth_rows.append(reference[source, n_discard + lag: + n_discard + lag + length]) + +detected = np.stack(detected_rows) +truth = np.stack(truth_rows) +ser = compute_ser(truth, detected) +print(f"SER after CMA + BPS: {ser:.2e} over {truth.size} symbols " + f"({int(round(float(ser) * truth.size))} errors)") + +fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(12, 4)) +plot_iq(naive[0, n_discard:], title="received, polarization X", ax=axes[0]) +plot_iq(equalized[0], title="after CMA", ax=axes[1]) +plot_iq(recovered[0], reference=constellation, + title="after BPS", ax=axes[2]) +plt.tight_layout() +plt.savefig(f"{img_dir}/one_shot_pdm_fig1.png") diff --git a/examples/simple/monte_carlo_awgn.py b/examples/simple/monte_carlo_awgn.py index 5660ff2..13445c2 100644 --- a/examples/simple/monte_carlo_awgn.py +++ b/examples/simple/monte_carlo_awgn.py @@ -102,12 +102,10 @@ def qfunc(x): # it needs k, which is chain-level knowledge the channel does not have. ebn0_dB_list = np.arange(0, 25, 2) N_compare = 100000 -fine_dB_list = np.arange(0, 25, 0.05) orders = np.array([4, 16, 64, 256]) bits_per_symbol = np.log2(orders).astype(int) # vectorized, not a loop measured_ber = {} theory_ber = {} -required_dB = np.zeros(len(orders)) for index, order in enumerate(orders): other = Constellation("QAM", int(order)) bits = int(bits_per_symbol[index]) @@ -119,19 +117,6 @@ def qfunc(x): order_label = f"{order}-QAM" measured_ber[order_label] = collected["ber"] theory_ber[order_label] = other.metrics(ebn0_dB_list, per="bit")["ber"] - # What the density costs, read off the closed form rather than the - # measurement: the sweep grid is 2 dB coarse and noisy at 1e-3. - fine_ber = other.metrics(fine_dB_list, per="bit")["ber"] - above_underflow = fine_ber > 0 # the tail reaches 0 in float64 - required_dB[index] = np.interp( - -3.0, np.log10(fine_ber[above_underflow])[::-1], - fine_dB_list[above_underflow][::-1]) - -print() -print_data({"x": orders, - "curves": {"bits per symbol": bits_per_symbol, - "Eb/N0 at BER=1e-3 [dB]": required_dB}}, - xlabel="QAM order") ax = plot_error_rate(ebn0_dB_list, measured_ber, theory=theory_ber, xlabel="$E_b/N_0$ [dB]", ylabel="BER", diff --git a/examples/simple/one_shot_awgn.py b/examples/simple/one_shot_awgn.py index 1a1352c..99d5e67 100644 --- a/examples/simple/one_shot_awgn.py +++ b/examples/simple/one_shot_awgn.py @@ -82,9 +82,10 @@ print(f"chain : SER = {ser:.4f}, theory = {ser_theo:.4f}") # --- 3. what the object knows about itself ---------------------------- -print("\nconstellation.info()") -for key, value in constellation.info().items(): - print(f" {key:16s} {value}") +# print(constellation) renders info() -- any object with an info() +# method prints itself this way +print() +print(constellation) # --- 4. the figures --------------------------------------------------- fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(9, 4.2)) diff --git a/src/comnumpy/core/channels.py b/src/comnumpy/core/channels.py index 1d853f5..b31c4a0 100644 --- a/src/comnumpy/core/channels.py +++ b/src/comnumpy/core/channels.py @@ -127,9 +127,9 @@ class FIRChannel(Processor): y[n] = \sum_{l=0}^{L-1} h[l] \, x[n-l] - Axes: *declared axis* -- operates on a 1D serial signal ``(N,)`` + Axes: *axis -1* -- the convolution runs along the last axis (``scipy.signal.convolve`` semantics; output length depends on - ``mode``). + ``mode``); leading axes are batch, every row through the same taps. Parameters ---------- @@ -159,7 +159,10 @@ class FIRChannel(Processor): def forward(self, x: np.ndarray) -> np.ndarray: from scipy import signal # local import (D36) - y = signal.convolve(x, self.h, mode=self.mode) + # scipy wants kernel and signal of equal rank: size-1 leading + # axes keep those dimensions untouched in every mode + h = np.reshape(self.h, (1,) * (np.ndim(x) - 1) + (-1,)) + y = signal.convolve(x, h, mode=self.mode) return y def info(self) -> dict[str, Any]: diff --git a/src/comnumpy/core/compensators.py b/src/comnumpy/core/compensators.py index d7e7692..f7f9527 100644 --- a/src/comnumpy/core/compensators.py +++ b/src/comnumpy/core/compensators.py @@ -1,7 +1,7 @@ import numpy as np import numpy.linalg as LA from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Literal, Optional, Union +from typing import Any, Literal, Optional, TYPE_CHECKING, Union from scipy import signal from scipy.linalg import toeplitz from scipy.optimize import least_squares @@ -16,7 +16,8 @@ __all__ = [ "DataAidedMixin", "DCCorrector", "Normalizer", "BlindIQCompensator", - "BlindCFOCompensator", "BlindPhaseCompensation", "LinearEqualizer", + "BlindCFOCompensator", "BlindPhaseCompensation", + "BlindPhaseSearchCompensator", "LinearEqualizer", "DataAidedFIRCompensator", "DataAidedPhaseCompensator", "DataAidedComplexGainCompensator", "DataAidedSimpleSynchronizer", "DataAidedFineSynchronizer", @@ -148,10 +149,11 @@ class DCCorrector(Processor): local oscillator adds a constant term to the baseband signal. Axes: *declared axis* -- the mean :math:`\mu_x` is computed along - ``axis`` (default 0) and broadcast over the remaining axes; use - ``axis=-1`` for the canonical serial layout ``(..., N)``. The - estimand is **per path** (D49): a DC offset belongs to a converter, - so ``axis=-1`` on a ``(..., P, N)`` signal gives each path its own. + ``axis`` (default -1, the canonical serial layout ``(..., N)``) and + broadcast over the remaining axes. The estimand is **per path** + (D49): a DC offset belongs to a converter, so the default gives each + row of a batch its own offset -- the old default of 0 averaged + *across* the batch, which is the silent coupling D51 forbids. Parameters ---------- @@ -159,7 +161,7 @@ class DCCorrector(Processor): Target mean value :math:`\alpha` of the output. Default is 0.0, i.e. a zero-mean output. axis : int, optional, keyword-only - Axis along which :math:`\mu_x` is computed. Default is 0. + Axis along which :math:`\mu_x` is computed. Default is -1. name : str, optional, keyword-only Name of the corrector instance. Default is ``"mean_corrector"``. @@ -178,7 +180,7 @@ class DCCorrector(Processor): [0. 1. 2. 5.] """ value: float = 0.0 - axis: int = field(default=0, kw_only=True) + axis: int = field(default=-1, kw_only=True) name: str = field(default="mean_corrector", kw_only=True) def forward(self, x: np.ndarray) -> np.ndarray: @@ -622,6 +624,19 @@ def __post_init__(self): self.grid_search_array = np.arange(self.grid_search_tuple[0], self.grid_search_tuple[1], self.grid_search_tuple[2]) self.history = [] + def prepare(self, x: np.ndarray) -> None: + from comnumpy.exceptions import ShapeError # local import (D36) + # the event is the polarization pair (see test_estimand_scope): + # one oscillator, ONE w0_ estimated jointly over (2, N). A wider + # leading shape is a batch of independent trials, which a single + # scalar estimate would smear silently (D51) -- refused, with + # the loop named. + if np.ndim(x) > 2 or (np.ndim(x) == 2 and np.shape(x)[-2] != 2): + raise ShapeError( + f"BlindCFOCompensator estimates one CFO per signal or " + f"per polarization pair (2, N), got shape {np.shape(x)} " + f"-- loop over the batch, one compensator per trial.") + def loss(self, x: np.ndarray, w: float) -> float: N = x.shape[-1] x4 = x**4 @@ -858,8 +873,9 @@ class LinearEqualizer(Processor): :math:`\mathbb{E}[|s[n]|^2] = 1`). The two coincide when :math:`\sigma^2 = 0`. - Axes: *declared axis* -- expects a 1D serial signal ``(N,)`` and - returns ``(N-L+1,)``. + Axes: *axis -1* -- expects a serial signal ``(..., N)`` and returns + ``(..., N-L+1)``; leading axes are batch (independent frames, one + equalizer matrix). Parameters ---------- @@ -908,12 +924,15 @@ def forward(self, X: np.ndarray) -> np.ndarray: X = np.asarray(X).astype(complex) H = self.get_H(X.shape[-1]) + # the estimator matrix acts on the last axis; the trailing + # length-1 axis makes matmul treat it as a column, so leading + # axes are batch (independent frames equalized in one call) match self.method: case "zf": - Y = zf_estimator(X, H) + Y = zf_estimator(X[..., None], H) case "mmse": - Y = mmse_estimator(X, H, self.sigma2) - return Y + Y = mmse_estimator(X[..., None], H, self.sigma2) + return Y[..., 0] @dataclass(slots=True) @@ -1670,3 +1689,100 @@ def forward(self, x: np.ndarray) -> np.ndarray: y = y[:self.signal_len] return y + + +@dataclass(slots=True) +class BlindPhaseSearchCompensator(Processor): + r"""Feedforward carrier phase recovery by blind phase search (BPS). + + Signal Model + ------------ + The input carries a slowly varying phase :math:`\phi[n]` -- laser + phase noise, typically -- on top of the constellation: + + .. math:: + + x[n] = s[n] \, e^{j \phi[n]} + b[n] + + For each of :math:`B` test phases + :math:`\varphi_i = i \, \frac{\pi/2}{B}`, the input is rotated, + decided against the alphabet, and the squared decision distance is + averaged over a sliding window of :math:`2W + 1` symbols. The + estimate at each symbol is the test phase of smallest averaged + distance, then unwrapped modulo :math:`\pi/2` so the estimate can + follow the walk beyond one quadrant. + + The :math:`\pi/2` search range is the symmetry of a square QAM (and + of QPSK): the recovered constellation is aligned up to a quadrant. + That residual **quadrant ambiguity is not resolved here** -- it is + the same ambiguity every blind phase estimator leaves, and a system + resolves it with differential encoding or a known field. The + estimated phase trajectory is exposed as ``phase_`` (D23). + + Axes: *axis -1* -- the window slides along the last axis; leading + axes (the two rows of a polarization pair) are processed + independently, each with its own trajectory. + + Parameters + ---------- + alphabet : np.ndarray or Constellation + The constellation the decision is made against. + n_test : int, optional, keyword-only + Number of test phases :math:`B` over :math:`[0, \pi/2)`. + Default is 32. + half_window : int, optional, keyword-only + Half-width :math:`W` of the averaging window. Default is 16. + The window trades noise rejection against tracking speed: it + must be short against the phase drift and long against the + additive noise. + name : str, optional, keyword-only + Block name. Default ``"bps"``. + + References + ---------- + T. Pfau, S. Hoffmann, R. Noe, "Hardware-Efficient Coherent Digital + Receiver Concept With Feedforward Carrier Recovery for M-QAM + Constellations," Journal of Lightwave Technology, vol. 27, no. 8, + pp. 989-999, 2009. + + Examples + -------- + >>> from comnumpy.core.utils import get_alphabet + >>> alphabet = get_alphabet("PSK", 4) + >>> s = alphabet[np.arange(64) % 4] + >>> rotated = s * np.exp(1j * 0.3) + >>> y = BlindPhaseSearchCompensator(alphabet, half_window=8)(rotated) + >>> bool(np.max(np.abs(y - s)) < np.pi / 2 / 32) # within one grid step + True + """ + alphabet: Any + n_test: int = field(default=32, kw_only=True) + half_window: int = field(default=16, kw_only=True) + name: str = field(default="bps", kw_only=True) + phase_: Optional[np.ndarray] = field(init=False, repr=False, default=None) + + def _track(self, x: np.ndarray) -> np.ndarray: + """The unwrapped phase trajectory of one 1D signal.""" + points = np.ravel(np.asarray(self.alphabet, dtype=complex)) + candidates = (np.pi / 2) * np.arange(self.n_test) / self.n_test + window = np.ones(2 * self.half_window + 1) + cost = np.empty((self.n_test, x.shape[-1])) + for index, phi in enumerate(candidates): + rotated = x * np.exp(-1j * phi) + distance = np.min( + np.abs(rotated[:, None] - points[None, :]) ** 2, axis=-1) + cost[index] = np.convolve(distance, window, mode="same") + raw = candidates[np.argmin(cost, axis=0)] + # unwrap modulo pi/2: a jump larger than pi/4 between neighbours + # is the estimate crossing the edge of the search range + steps = np.diff(raw) + steps -= (np.pi / 2) * np.round(steps / (np.pi / 2)) + return raw[0] + np.concatenate(([0.0], np.cumsum(steps))) + + def forward(self, x: np.ndarray) -> np.ndarray: + flat = np.reshape(x, (-1, x.shape[-1])) + trajectories = np.empty(flat.shape) + for row in range(flat.shape[0]): + trajectories[row] = self._track(flat[row]) + self.phase_ = np.reshape(trajectories, x.shape) + return x * np.exp(-1j * self.phase_) diff --git a/src/comnumpy/core/filters.py b/src/comnumpy/core/filters.py index e4f3550..428f542 100644 --- a/src/comnumpy/core/filters.py +++ b/src/comnumpy/core/filters.py @@ -244,8 +244,8 @@ class BWFilter(Processor): operation is a **circular** convolution with a periodic sinc of length :math:`N`, not a linear one. - Axes: *axis -1* -- expects a 1D signal ``(N,)``; multi-dimensional - inputs are rejected. + Axes: *axis -1* -- the FFT mask acts on the last axis; leading + axes are batch. Parameters ---------- @@ -289,12 +289,10 @@ def __post_init__(self) -> None: def forward(self, x: np.ndarray) -> np.ndarray: - if x.ndim > 1: - raise NotImplementedError("BW Filter: only 1D signals are supported.") - from comnumpy._backend import fft, ifft, fftfreq # local import (D36), cupy-compatible (D3) - NFFT = len(x) + # the FFT and the mask act on axis -1; leading axes are batch + NFFT = x.shape[-1] w = fftfreq(NFFT, d=1, like=x) # wn is normalized to Nyquist (= 1/2 cycle/sample), hence the /2 H = (abs(w) <= self.wn / 2).astype(float) diff --git a/src/comnumpy/core/generics.py b/src/comnumpy/core/generics.py index bcd08e1..a6728c6 100644 --- a/src/comnumpy/core/generics.py +++ b/src/comnumpy/core/generics.py @@ -12,6 +12,27 @@ logger = logging.getLogger(__name__) +def render_info(obj: object) -> Optional[str]: + """The ``info()`` of *obj* as ``key: value`` lines, or None. + + Shared by the ``__str__`` of :class:`Processor` and of + ``Constellation``: ``print(obj)`` shows what the object *is* when it + can say so, and falls back to ``repr`` when it cannot. + """ + info = getattr(obj, "info", None) + if not callable(info): + return None + rendered = info() + if isinstance(rendered, str): + return rendered + if not isinstance(rendered, dict): + return None + lines = [] + for key, value in rendered.items(): + lines.append(f"{key}: {value}") + return "\n".join(lines) + + @dataclass(slots=True) class Processor(): r""" @@ -66,6 +87,22 @@ def prepare(self, X: np.ndarray, /) -> None: """ pass + def __str__(self) -> str: + """``print(processor)`` shows its ``info()`` when it has one. + + A block that defines ``info()`` prints one ``key: value`` line + per entry; a block that does not falls back to its ``repr``. + + Examples + -------- + >>> from comnumpy.core.channels import FIRChannel + >>> print(FIRChannel(np.array([1.0, 0.0, 0.5]))) + kind: FIR + n_taps: 3 + ... + """ + return render_info(self) or repr(self) + def __call__(self, X: np.ndarray) -> np.ndarray: """ Process the input data by calling the forward method. diff --git a/src/comnumpy/core/impairments.py b/src/comnumpy/core/impairments.py index a80b4db..94afa37 100644 --- a/src/comnumpy/core/impairments.py +++ b/src/comnumpy/core/impairments.py @@ -67,8 +67,8 @@ class CFO(Processor): where :math:`\omega_0` is the normalized carrier frequency offset in rad/sample. - Axes: *axis -1* -- expects (N,); the time index :math:`n` runs along a - one-dimensional input. + Axes: *axis -1* -- the time index :math:`n` runs along the last + axis; leading axes are batch, every row sees the same offset. Parameters ---------- @@ -94,8 +94,7 @@ class CFO(Processor): name: str = field(default="cfo_impairment", kw_only=True) def forward(self, x: np.ndarray) -> np.ndarray: - N = len(x) - n_vect = np.arange(N) + n_vect = np.arange(x.shape[-1]) y = x * np.exp(1j*self.cfo*n_vect) return y @@ -117,8 +116,8 @@ class Delay(Processor): keep the input length :math:`N`; otherwise the output has length :math:`N - \tau`. - Axes: *axis -1* -- expects (N,); samples are dropped along a - one-dimensional input. + Axes: *axis -1* -- samples are dropped along the last axis; + leading axes are batch. Parameters ---------- @@ -156,10 +155,10 @@ def __post_init__(self): def forward(self, x: np.ndarray) -> np.ndarray: - x_delayed = x[self.tau:] + x_delayed = x[..., self.tau:] if self.pad_zeros: - y = np.zeros(len(x)) - y[:len(x_delayed)] = x_delayed + y = np.zeros(x.shape, dtype=x.dtype) + y[..., :x_delayed.shape[-1]] = x_delayed else: y = x_delayed diff --git a/src/comnumpy/core/metrics.py b/src/comnumpy/core/metrics.py index 7447890..9455a05 100644 --- a/src/comnumpy/core/metrics.py +++ b/src/comnumpy/core/metrics.py @@ -602,8 +602,11 @@ def compute_ser(X_target: np.ndarray, X_detected: np.ndarray, Axes: *declared axis* -- errors are counted along ``axis``, which must then hold the symbols; with ``axis=None`` (default) both arrays are - raveled and truncated to their common length - :math:`N = \min(N_{\text{target}}, N_{\text{detected}})`. + truncated to their common length + :math:`N = \min(N_{\text{target}}, N_{\text{detected}})` **along the + last axis**, then every element is pooled -- so a batch of frames + stays aligned row by row, and the pooled rate is the mean of the + per-row rates. Parameters ---------- @@ -638,11 +641,14 @@ def compute_ser(X_target: np.ndarray, X_detected: np.ndarray, [0.25 0.25] """ if axis is None: - x_target = np.ravel(X_target) - x_detected = np.ravel(X_detected) - N = min(len(x_detected), len(x_target)) - nb_errors = np.count_nonzero(x_target[:N] - x_detected[:N]) - return nb_errors / N + x_target = np.asarray(X_target) + x_detected = np.asarray(X_detected) + # a length mismatch is truncated along the *last* axis before + # anything is pooled: raveling first would misalign every row + # of a batch after the first -- a silent offset, not an error + N = min(x_target.shape[-1], x_detected.shape[-1]) + difference = x_target[..., :N] - x_detected[..., :N] + return np.count_nonzero(difference) / difference.size else: N = X_target.shape[axis] nb_errors = np.count_nonzero(X_target - X_detected, axis=axis) @@ -708,8 +714,13 @@ def compute_ber(X_target: np.ndarray, X_detected: np.ndarray, width: int, 0.25 """ if axis is None: - s_target = sym_2_bin(np.ravel(X_target), width) - s_detected = sym_2_bin(np.ravel(X_detected), width) + x_target = np.asarray(X_target) + x_detected = np.asarray(X_detected) + # truncated along the last axis before raveling, like + # compute_ser: a batch stays aligned row by row + N = min(x_target.shape[-1], x_detected.shape[-1]) + s_target = sym_2_bin(np.ravel(x_target[..., :N]), width) + s_detected = sym_2_bin(np.ravel(x_detected[..., :N]), width) nb_errors = np.count_nonzero(s_target - s_detected) return nb_errors / len(s_detected) else: diff --git a/src/comnumpy/core/utils.py b/src/comnumpy/core/utils.py index 84b539d..7d4d7b5 100644 --- a/src/comnumpy/core/utils.py +++ b/src/comnumpy/core/utils.py @@ -327,6 +327,19 @@ def __array__(self, dtype: Optional[npt.DTypeLike] = None, def __len__(self) -> int: return self.order + def __str__(self) -> str: + """``print(constellation)`` renders :meth:`info`, one line per key. + + Examples + -------- + >>> print(Constellation("PSK", 4)) + family: PSK + order: 4 + ... + """ + from comnumpy.core.generics import render_info + return render_info(self) or repr(self) + def info(self) -> dict[str, Any]: """What the constellation is, as a dictionary ready to print. @@ -1047,8 +1060,10 @@ def mmse_estimator(Y: np.ndarray, H: np.ndarray, sigma2: float) -> np.ndarray: [[ 0.9151 -0.8931] [ 0.9868 0.9647]] """ - _, N_t = H.shape - H_H = np.conjugate(np.transpose(H)) + N_t = H.shape[-1] + # swapaxes, not transpose: a stacked H (K, N_r, N_t) keeps its batch + # axes in front, and solve/matmul batch over them (D51) + H_H = np.conjugate(np.swapaxes(H, -1, -2)) A = np.matmul(H_H, H) + sigma2 * np.eye(N_t) Z_est = LA.solve(A, np.matmul(H_H, Y)) return Z_est diff --git a/src/comnumpy/mimo/channels.py b/src/comnumpy/mimo/channels.py index 7390546..9a5d462 100644 --- a/src/comnumpy/mimo/channels.py +++ b/src/comnumpy/mimo/channels.py @@ -103,7 +103,11 @@ class FlatMIMOChannel(BaseMIMOChannel): the noiseless product). Axes: *declared axis* -- expects ``(..., ant, N)`` with antennas on - axis -2, validated against :math:`N_t`. + axis -2, validated against :math:`N_t`. ``H`` may be a **stack** of + channels ``(K, N_r, N_t)``: against ``X`` of shape ``(K, N_t, N)`` + channel :math:`k` is applied to frame :math:`k` -- one independent + draw per trial, the batch contract of D51 for a channel whose + realization is configured rather than drawn per call. Parameters ---------- @@ -135,7 +139,9 @@ class FlatMIMOChannel(BaseMIMOChannel): def forward(self, X: np.ndarray) -> np.ndarray: assert self.H is not None # validate_input rejects a None channel - validate_input(X, self.H.shape[1]) + validate_input(X, self.H.shape[-1]) + # matmul batches over leading axes: a stacked H (K, N_r, N_t) + # against X (K, N_t, N) applies channel k to frame k (D51) return np.matmul(self.H, X) diff --git a/src/comnumpy/mimo/compensators.py b/src/comnumpy/mimo/compensators.py index e830872..ce7de2a 100644 --- a/src/comnumpy/mimo/compensators.py +++ b/src/comnumpy/mimo/compensators.py @@ -67,9 +67,14 @@ class BlindDualMIMOCompensator(Processor): takes it to within a tenth of the floor, and DD closes what is left. - Axes: *declared axis* -- expects the dual-polarization layout - ``(2, N)``, produces ``(2, N // oversampling)`` (one output sample - per ``oversampling`` input samples: fractionally spaced equalizer). + Axes: *polarization pair* -- expects ``(..., 2, N)``, produces + ``(..., 2, N // oversampling)`` (one output sample per + ``oversampling`` input samples: fractionally spaced equalizer). + Leading axes are batch, and the butterfly is a per-pair algorithm: + each pair gets its **own** equalizer state, adapted independently. + After a batched pass ``H_`` carries the batch axes in front, + ``(..., 2, 2(2L+1))``, so a further ``partial_fit`` resumes each + pair from its own state. **Delay.** The centre tap of the initial equalizer sits :math:`L` input samples back, so output sample :math:`k` is aligned with @@ -267,10 +272,45 @@ def _apply_frozen(self, X: np.ndarray, Y: np.ndarray) -> np.ndarray: return Y def forward(self, X: np.ndarray) -> np.ndarray: - - if X.shape[0] != 2: - raise ValueError(f"Blind Dual MIMO Compensator only works for dual polarization signals (X shape={X.shape})") - + from comnumpy.exceptions import ShapeError # local import (D36) + if X.ndim < 2 or X.shape[-2] != 2: + raise ShapeError( + f"BlindDualMIMOCompensator expects a polarization pair " + f"(..., 2, N), got shape {X.shape}") + if X.ndim == 2: + return self._forward_pair(X) + + # batched: one independent equalizer per leading index -- the + # butterfly is a per-pair algorithm, so its state cannot be + # shared across trials. After the pass, H_ carries the batch + # axes in front, and a further partial_fit resumes each pair + # from its own state. + assert self.H_ is not None # initialize_H ran in __post_init__ + batch_shape = X.shape[:-2] + if self.H_.ndim == 2: + initial = np.broadcast_to( + self.H_, batch_shape + self.H_.shape).copy() + elif self.H_.shape[:-2] == batch_shape: + initial = self.H_ + else: + raise ShapeError( + f"the equalizer state H_ carries batch shape " + f"{self.H_.shape[:-2]} from a previous pass, but the " + f"input carries {batch_shape} -- call initialize_H() to " + f"restart from the identity equalizer") + states = np.empty_like(initial) + outputs = np.empty( + batch_shape + (2, X.shape[-1] // self.oversampling), + dtype=complex) + for index in np.ndindex(*batch_shape): + self.H_ = initial[index].copy() + outputs[index] = self._forward_pair(X[index]) + states[index] = self.H_ + self.H_ = states + return outputs + + def _forward_pair(self, X: np.ndarray) -> np.ndarray: + """One adaptation pass over a single ``(2, N)`` pair.""" L = self.L os = self.oversampling M, N = X.shape diff --git a/src/comnumpy/mimo/detectors.py b/src/comnumpy/mimo/detectors.py index 39fd795..98349a4 100644 --- a/src/comnumpy/mimo/detectors.py +++ b/src/comnumpy/mimo/detectors.py @@ -1,7 +1,7 @@ import numpy as np import numpy.linalg as LA from dataclasses import dataclass, field -from typing import Literal, Optional +from typing import Callable, Literal, Optional from comnumpy.core.generics import Processor from comnumpy.core.utils import hard_projector, soft_projector, zf_estimator, mmse_estimator @@ -36,6 +36,28 @@ def _required_channel(H: Optional[np.ndarray], block: str) -> np.ndarray: return np.asarray(H) +def _detect_stacked(detect: "Callable[[np.ndarray, np.ndarray], np.ndarray]", + H: np.ndarray, Y: np.ndarray) -> np.ndarray: + """One detection per leading index of a stacked channel (D51). + + The search algorithms are per-draw by nature -- an enumeration, a + tree, a cancellation order all depend on the one matrix in front of + them -- so the batch is an explicit loop: channel ``k`` decides + frame ``k``, and no state leaks between draws. Convenience and + correctness, not a speedup. + """ + from comnumpy.exceptions import ShapeError + if Y.ndim < 2 or Y.shape[:-2] != H.shape[:-2]: + raise ShapeError( + f"a stacked channel with leading shape {H.shape[:-2]} " + f"decides one frame per draw and needs observations " + f"(..., N_r, N) with the same leading shape, got {Y.shape}") + S = np.empty(H.shape[:-2] + (H.shape[-1], Y.shape[-1]), dtype=int) + for index in np.ndindex(*H.shape[:-2]): + S[index] = detect(H[index], Y[index]) + return S + + def _required_noise_variance(sigma2: Optional[float], block: str) -> float: """Same, for the noise variance the MMSE-flavoured detectors need.""" if sigma2 is None: @@ -87,7 +109,9 @@ class MaximumLikelihoodDetector(Processor): antennas or a high-order constellation. Axes: *declared axis* -- expects ``(ant, N)`` with antennas on - axis -2. + axis -2. ``H`` may be a stack ``(K, N_r, N_t)`` against ``Y`` of + shape ``(K, N_r, N)``: channel k decides frame k, through an + internal per-draw loop -- convenience, not a speedup (D51). Parameters ---------- @@ -133,8 +157,7 @@ def __post_init__(self): def get_nb_candidates(self) -> int: H = _required_channel(self.H, type(self).__name__) - _, N_t = H.shape - return len(self.alphabet) ** N_t + return len(self.alphabet) ** H.shape[-1] def get_candidates(self, alphabet: np.ndarray, N_t: int) -> tuple[np.ndarray, np.ndarray]: @@ -155,6 +178,16 @@ def get_candidates(self, alphabet: np.ndarray, def forward(self, Y: np.ndarray) -> np.ndarray: H = _required_channel(self.H, type(self).__name__) + if H.ndim > 2: + S = _detect_stacked(self._detect, H, Y) + self.S = S + return S + S = self._detect(H, Y) + self.S = S + return S + + def _detect(self, H: np.ndarray, Y: np.ndarray) -> np.ndarray: + """The exhaustive search against one channel matrix.""" _, N_t = H.shape _, N = Y.shape alphabet = self.alphabet @@ -187,7 +220,6 @@ def forward(self, Y: np.ndarray) -> np.ndarray: best[better] = values[better] S[:, start:start + n_block] = S_candidates[:, winner] - self.S = S return S @@ -253,8 +285,10 @@ class SphereDecoder(Processor): 16.7 million. Axes: *declared axis* -- expects ``(ant, N)`` with antennas on - axis -2, and returns the alphabet indices, like every other detector - here. + axis -2. ``H`` may be a stack ``(K, N_r, N_t)`` against ``Y`` of + shape ``(K, N_r, N)``: channel k decides frame k, through an + internal per-draw loop -- convenience, not a speedup (D51) -- and returns the + alphabet indices, like every other detector here. Parameters ---------- @@ -329,14 +363,16 @@ class SphereDecoder(Processor): S: Optional[np.ndarray] = field(init=False, repr=False, default_factory=lambda: None) nodes_: float = field(init=False, repr=False, default=0.0) + # running counters behind nodes_, declared for slots (D40a) + _visited: int = field(init=False, repr=False, default=0) + _samples: int = field(init=False, repr=False, default=0) def __post_init__(self): # accept anything array-like, a Constellation included self.alphabet = np.asarray(self.alphabet) - def _triangularize(self) -> tuple[np.ndarray, np.ndarray]: - """The thin QR of the channel, with the guards it needs.""" - H = _required_channel(self.H, type(self).__name__) + def _triangularize(self, H: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """The thin QR of one channel matrix, with the guards it needs.""" n_rx, n_tx = H.shape if n_rx < n_tx: raise ValueError( @@ -391,21 +427,35 @@ def search(level: int, metric: float) -> None: return best, nodes def forward(self, Y: np.ndarray) -> np.ndarray: - Q, R = self._triangularize() + H = _required_channel(self.H, type(self).__name__) + if H.ndim > 2: + # nodes_ averages over every sample of every draw: the + # per-draw counts accumulate through _detect below + self._visited, self._samples = 0, 0 + S = _detect_stacked(self._detect, H, Y) + self.nodes_ = self._visited / max(1, self._samples) + self.S = S + return S + self._visited, self._samples = 0, 0 + S = self._detect(H, Y) + self.nodes_ = self._visited / max(1, self._samples) + self.S = S + return S + + def _detect(self, H: np.ndarray, Y: np.ndarray) -> np.ndarray: + """The tree search against one channel matrix.""" + Q, R = self._triangularize(H) Z = np.matmul(np.conjugate(Q).T, Y) # (N_t, N) alphabet = np.asarray(self.alphabet).ravel() gains = np.abs(np.diag(R)) ** 2 n_samples = Y.shape[-1] S = np.empty((R.shape[0], n_samples), dtype=int) - visited = 0 for sample in range(n_samples): best, nodes = self._closest(R, gains, Z[:, sample], alphabet) S[:, sample] = best - visited += nodes - - self.nodes_ = visited / max(1, n_samples) - self.S = S + self._visited += nodes + self._samples += n_samples return S @@ -439,7 +489,9 @@ class LinearDetector(Processor): the integer indices of the decisions in the alphabet. Axes: *declared axis* -- expects ``(ant, N)`` with antennas on - axis -2. + axis -2. ``H`` may be a stack ``(K, N_r, N_t)`` against + ``Y`` of shape ``(K, N_r, N)`` -- one draw per trial, decided + through numpy's stacked linear algebra in one call (D51). Parameters ---------- @@ -540,7 +592,9 @@ class OrderedSuccessiveInterferenceCancellationDetector(Processor): of the decisions in the alphabet. Axes: *declared axis* -- expects ``(ant, N)`` with antennas on - axis -2. + axis -2. ``H`` may be a stack ``(K, N_r, N_t)`` against ``Y`` of + shape ``(K, N_r, N)``: channel k decides frame k, through an + internal per-draw loop -- convenience, not a speedup (D51). Parameters ---------- @@ -629,8 +683,14 @@ def ordering(self, H: np.ndarray) -> int: raise ValueError("osic_type must be 'sinr', 'colnorm', or 'snr'.") def forward(self, Y: np.ndarray) -> np.ndarray: + H = _required_channel(self.H, type(self).__name__) + if H.ndim > 2: + return _detect_stacked(self._detect, H, Y) + return self._detect(H, Y) + + def _detect(self, H_full: np.ndarray, Y: np.ndarray) -> np.ndarray: + """The ordered cancellation against one channel matrix.""" block = type(self).__name__ - H_full = _required_channel(self.H, block) Y_temp = Y.copy() NT = H_full.shape[1] diff --git a/src/comnumpy/mimo/utils.py b/src/comnumpy/mimo/utils.py index 5880da1..b00f130 100644 --- a/src/comnumpy/mimo/utils.py +++ b/src/comnumpy/mimo/utils.py @@ -11,7 +11,8 @@ def rayleigh_channel(N_r: int, N_t: int, L: int = 1, scale_per_tap: Optional[Sequence[float]] = None, seed: Optional[int] = None, - rng: Optional[np.random.Generator] = None) -> np.ndarray: + rng: Optional[np.random.Generator] = None, + size: Optional[int] = None) -> np.ndarray: r""" Draw an i.i.d. Rayleigh MIMO channel. @@ -70,11 +71,21 @@ def rayleigh_channel(N_r: int, N_t: int, Random seed used if `rng` is not provided. rng : numpy.random.Generator, optional Pre-initialized random number generator. If provided, `seed` is ignored. + size : int, optional + Number of independent channel draws. When given, the returned + array grows a leading axis of length ``size`` -- a **stack** of + channels, one per Monte-Carlo trial, drawn sequentially from the + same generator so draw :math:`k` equals the :math:`k`-th draw of + the equivalent Python loop. The stack is what + :class:`~comnumpy.mimo.channels.FlatMIMOChannel` and the + detectors accept as a batched ``H`` (D51). Returns ------- numpy.ndarray - Array of shape (L, N_r, N_t) containing the Rayleigh channel taps. + Array of shape (L, N_r, N_t) containing the Rayleigh channel + taps -- squeezed to (N_r, N_t) when ``L == 1``, with a leading + ``size`` axis when ``size`` is given. Raises ------ @@ -102,10 +113,23 @@ def rayleigh_channel(N_r: int, N_t: int, >>> H2 = rayleigh_channel(N_r=32, N_t=32, seed=1) >>> print(round(float(np.mean(np.abs(H2)**2)), 2)) # unit average gain 1.01 + >>> stack = rayleigh_channel(N_r=2, N_t=2, seed=0, size=200) + >>> stack.shape # one draw per Monte-Carlo trial + (200, 2, 2) """ if not rng: rng = np.random.default_rng(seed) + if size is not None: + # a stack of independent draws, sequential from the same rng: + # draw k equals the k-th draw of the equivalent Python loop, so + # a study moving to the batched form keeps its channels + draws = [] + for _ in range(size): + draws.append(rayleigh_channel(N_r, N_t, L, scale_per_tap, + rng=rng)) + return np.stack(draws) + if scale_per_tap is None: scales = np.ones(L, dtype=float) else: diff --git a/src/comnumpy/ofdm/predistorders.py b/src/comnumpy/ofdm/predistorders.py index 7074c20..6fc2123 100644 --- a/src/comnumpy/ofdm/predistorders.py +++ b/src/comnumpy/ofdm/predistorders.py @@ -36,8 +36,9 @@ class HardClipper(Processor): polar form. Axes: *element-wise* -- the clipping is applied pointwise; the - threshold :math:`T_m` is computed from the mean power of the whole - array. + threshold :math:`T_m` is computed from the mean power **along the + last axis**, so each row of a batch is clipped against its own + power -- one transmitter per trial (D51). Parameters ---------- @@ -70,7 +71,9 @@ def __post_init__(self): self.cr = 10 ** (self.cr_dB / 20) def forward(self, x: np.ndarray) -> np.ndarray: - Pmoy = np.mean(np.abs(x) ** 2) + # per-row power: each trial of a batch clips against its own + # RMS, not the RMS of the whole stack (D51) + Pmoy = np.mean(np.abs(x) ** 2, axis=-1, keepdims=True) Tm = self.cr * np.sqrt(Pmoy) if np.iscomplexobj(x): y = np.where(np.abs(x) > Tm, Tm * np.exp(1j * np.angle(x)), x) diff --git a/src/comnumpy/optical/channels.py b/src/comnumpy/optical/channels.py index 8360b9e..9e8f81f 100644 --- a/src/comnumpy/optical/channels.py +++ b/src/comnumpy/optical/channels.py @@ -1,11 +1,12 @@ import numpy as np from dataclasses import dataclass, field -from typing import Optional +from typing import Literal, Optional from comnumpy.core import Processor from .utils import apply_chromatic_dispersion, apply_kerr_nonlinearity, compute_beta2 from .constants import CD_COEFFICIENT, SPEED_OF_LIGHT, WAVELENGTH, KERR_COEFFICIENT -__all__ = ["PhaseNoise", "ChromaticDispersion", "KerrNonLinearity"] +__all__ = ["PhaseNoise", "ChromaticDispersion", "KerrNonLinearity", + "PMDEmulator"] @dataclass(slots=True) @@ -24,14 +25,28 @@ class PhaseNoise(Processor): increments :math:`\delta\phi_k` (a Wiener process accumulated along the samples). The magnitude :math:`|y[n]| = |x[n]|` is preserved. - Axes: *axis -1* -- expects a 1D signal (N,); the phase random walk - accumulates along the samples. + Axes: *axis -1* -- the phase random walk accumulates along the + samples. What shares one walk is declared by ``per``, which names + the *event* -- the slice of the array that sees one laser: + + - ``"pair"`` (default): one walk per trailing ``(2, N)`` block -- + the two rows of a polarization pair come from one laser, and + every leading axis is batch: independent draws per trial. A 2-D + or deeper input whose axis -2 is not 2 is refused, because it + cannot be read both ways: use ``"row"`` for independent rows or + ``"signal"`` for one shared walk. + - ``"row"``: one independent walk per 1-D row. + - ``"signal"``: one walk broadcast over the whole array. + + A 1-D input sees one walk under every setting. Parameters ---------- sigma2 : float Variance :math:`\sigma^2` of the per-sample phase increments, in rad^2. + per : {"pair", "row", "signal"}, optional, keyword-only + The event that shares one laser. Default is ``"pair"``. seed : int, optional, keyword-only Local RNG seed. name : str, optional, keyword-only @@ -52,6 +67,8 @@ class PhaseNoise(Processor): [1. 1. 1.] """ sigma2: float + per: Literal["pair", "row", "signal"] = field(default="pair", + kw_only=True) seed: Optional[int] = field(default=None, kw_only=True) name: str = field(default="phase noise", kw_only=True) # internal state (declared for slots, D40a) @@ -61,11 +78,30 @@ class PhaseNoise(Processor): def __post_init__(self) -> None: self.rng = np.random.default_rng(self.seed) + def prepare(self, x: np.ndarray) -> None: + from comnumpy.exceptions import ShapeError # local import (D36) + if self.per == "pair" and np.ndim(x) >= 2 and np.shape(x)[-2] != 2: + raise ShapeError( + f"PhaseNoise(per='pair') expects a polarization pair " + f"(..., 2, N), got shape {np.shape(x)} -- an axis -2 of " + f"size {np.shape(x)[-2]} cannot be read as a pair. Use " + f"per='row' for one independent laser per row, or " + f"per='signal' for one walk shared by the whole array.") + def noise_rvs(self, X: np.ndarray) -> np.ndarray: - """Draw the Wiener phase increment for a signal of the length of X.""" + """Draw one Wiener walk per event declared by ``per``.""" assert self.rng is not None # set in __post_init__ - noise = self.rng.normal(loc=0, scale=np.sqrt(self.sigma2), size=len(X)) - self._b = np.cumsum(noise) + n = np.shape(X)[-1] + if self.per == "signal" or np.ndim(X) <= 1: + shape: tuple[int, ...] = (n,) + elif self.per == "pair": + # one walk per trailing (2, N) block, broadcast over the pair + shape = np.shape(X)[:-2] + (1, n) + else: # "row" + shape = np.shape(X) + noise = self.rng.normal(loc=0, scale=np.sqrt(self.sigma2), + size=shape) + self._b = np.cumsum(noise, axis=-1) return self._b def forward(self, x: np.ndarray) -> np.ndarray: @@ -238,3 +274,118 @@ def forward(self, x: np.ndarray) -> np.ndarray: return y + + +@dataclass(slots=True) +class PMDEmulator(Processor): + r"""Polarization rotation and first-order PMD, as a section emulator. + + Signal Model + ------------ + The fibre's random birefringence is emulated by :math:`K` sections, + each a random unitary Jones rotation followed by a differential + group delay (DGD) of :math:`\tau / \sqrt{K}` between the two + principal states, applied exactly in the frequency domain: + + .. math:: + + \mathbf{Y}(\omega) = \prod_{k=1}^{K} + \mathbf{D}_k(\omega) \, \mathbf{R}_k \; \mathbf{X}(\omega), + \qquad + \mathbf{D}_k(\omega) = \mathrm{diag}\!\left( + e^{+j \omega \tau / 2\sqrt{K}}, \; + e^{-j \omega \tau / 2\sqrt{K}} \right) + + Each :math:`\mathbf{R}_k` is drawn Haar-uniform from the seeded + local generator, so the emulator is unitary -- it conserves energy + exactly -- and reproducible. With ``n_sections=1`` and ``dgd=0`` it + reduces to one random rotation of the state of polarization, which + is the part every polarization demultiplexer must undo even before + PMD enters. + + Randomly oriented sections add their DGDs in quadrature, hence the + :math:`1/\sqrt{K}` per section: the declared ``dgd`` is the **RMS** + DGD of the ensemble. Over section draws the DGD is Maxwellian, with + mean :math:`\sqrt{8/3\pi}\,\tau \approx 0.921\,\tau` (Poole & + Wagner). With ``n_sections=1`` the DGD is deterministic and equals + :math:`\tau` exactly. The concatenation of sections also makes the + DGD frequency-dependent, which is what distinguishes PMD from a + wavelength-flat rotation. + + Axes: *polarization pair* -- expects ``(..., 2, N)``; the Jones + matrices act on the polarization axis, the delays on the last one. + + Parameters + ---------- + dgd : float + RMS differential group delay :math:`\tau` in seconds. + n_sections : int, optional, keyword-only + Number of concatenated sections :math:`K`. Default is 8. + fs : float, optional, keyword-only + Sampling frequency in Hz. Default is 1.0. + seed : int, optional, keyword-only + Local RNG seed for the section rotations. Same seed, same fibre. + name : str, optional, keyword-only + Block name. Default ``"pmd"``. + + Raises + ------ + ShapeError + If the input does not carry a polarization pair ``(..., 2, N)``. + + References + ---------- + C. D. Poole and R. E. Wagner, "Phenomenological approach to + polarisation dispersion in long single-mode fibres," Electronics + Letters, vol. 22, no. 19, pp. 1029-1030, 1986. + S. J. Savory, "Digital filters for coherent optical receivers," + Optics Express, vol. 16, no. 2, pp. 804-817, 2008. + + Examples + -------- + >>> x = np.zeros((2, 8), dtype=complex); x[0] = 1.0 + >>> y = PMDEmulator(0.0, n_sections=1, seed=1)(x) + >>> bool(np.abs(np.sum(np.abs(y) ** 2) - 8.0) < 1e-9) # unitary + True + >>> bool(np.abs(y[1, 0]) > 0) # the rotation mixes the two rows + True + """ + dgd: float + n_sections: int = field(default=8, kw_only=True) + fs: float = field(default=1.0, kw_only=True) + seed: Optional[int] = field(default=None, kw_only=True) + name: str = field(default="pmd", kw_only=True) + rotations_: Optional[np.ndarray] = field(init=False, repr=False, + default=None) + + def __post_init__(self) -> None: + rng = np.random.default_rng(self.seed) + sections = [] + for _ in range(self.n_sections): + # Haar-uniform 2x2 unitary: QR of a complex Gaussian matrix, + # with the phase convention fixed by the R diagonal + gaussian = (rng.standard_normal((2, 2)) + + 1j * rng.standard_normal((2, 2))) + q, r = np.linalg.qr(gaussian) + sections.append(q * (np.diag(r) / np.abs(np.diag(r)))) + self.rotations_ = np.stack(sections) + + def prepare(self, x: np.ndarray) -> None: + from comnumpy.exceptions import ShapeError # local import (D36) + from .utils import is_polarization_pair # local import (D36) + if not is_polarization_pair(x, "PMDEmulator"): + raise ShapeError( + f"PMDEmulator expects a polarization pair (..., 2, N), " + f"got shape {x.shape} -- PMD is a two-polarization " + f"effect, there is nothing to delay against on one.") + + def forward(self, x: np.ndarray) -> np.ndarray: + assert self.rotations_ is not None # set in __post_init__ + omega = 2 * np.pi * self.fs * np.fft.fftfreq(x.shape[-1]) + delay = np.exp(1j * omega * self.dgd + / (2 * np.sqrt(self.n_sections))) + spectrum = np.fft.fft(x, axis=-1) + for rotation in self.rotations_: + spectrum = np.einsum("ij,...jn->...in", rotation, spectrum) + spectrum = spectrum * np.stack([delay, np.conj(delay)]) + return np.fft.ifft(spectrum, axis=-1) diff --git a/src/comnumpy/optical/compensators.py b/src/comnumpy/optical/compensators.py index 7b849d3..79b236b 100644 --- a/src/comnumpy/optical/compensators.py +++ b/src/comnumpy/optical/compensators.py @@ -64,8 +64,9 @@ class ChromaticDispersionFIRCompensator(Processor): to s^2/km (factor :math:`10^{-24}`), :math:`z` is in km and :math:`f_s` in Hz, which makes :math:`K` dimensionless. - Axes: *declared axis* -- operates on a 1D serial signal ``(N_x,)`` - (``scipy.signal.convolve`` in ``"full"`` mode). + Axes: *axis -1* -- the convolution runs along the last axis + (``scipy.signal.convolve`` in ``"full"`` mode); leading axes are + batch, every row through the same taps. Parameters ---------- @@ -134,7 +135,11 @@ def __post_init__(self) -> None: self.K = K def forward(self, x: np.ndarray) -> np.ndarray: - y = signal.convolve(x, self.h, mode='full') + assert self.h is not None # set in __post_init__ + # scipy wants kernel and signal of equal rank: size-1 leading + # axes make the convolution run along axis -1, batch in front + h = np.reshape(self.h, (1,) * (np.ndim(x) - 1) + (-1,)) + y = signal.convolve(x, h, mode='full') return y @@ -321,8 +326,9 @@ class ChromaticDispersionLSFIRCompensator(Processor): s^2/km, :math:`z` in km and :math:`f_s` in Hz, so that :math:`K` is dimensionless. - Axes: *declared axis* -- operates on a 1D serial signal ``(N_x,)`` - (``scipy.signal.convolve`` in ``"full"`` mode). + Axes: *axis -1* -- the convolution runs along the last axis + (``scipy.signal.convolve`` in ``"full"`` mode); leading axes are + batch, every row through the same taps. Parameters ---------- @@ -433,4 +439,6 @@ def __post_init__(self) -> None: self.K = K def forward(self, x: np.ndarray) -> np.ndarray: - return signal.convolve(x, self.h) + assert self.h is not None # set in __post_init__ + h = np.reshape(self.h, (1,) * (np.ndim(x) - 1) + (-1,)) + return signal.convolve(x, h) diff --git a/tests/core/test_batch_axes.py b/tests/core/test_batch_axes.py new file mode 100644 index 0000000..91935bc --- /dev/null +++ b/tests/core/test_batch_axes.py @@ -0,0 +1,199 @@ +"""The batch contract, block by block. + +Leading axes ahead of a block's event axes are batch axes: independent +trials. The contract has three families, and each family has its own +failure mode, so each gets its own lock: + +* a **deterministic** block broadcasts its one configuration over the + batch -- row ``i`` of a batched call must equal the block applied to + row ``i`` alone; +* a **stochastic** block draws independently per event -- trials must + not share a realization, and what one event *is* must be declared; +* an **adaptive** block carries independent state per event -- one + equalizer per pair, never one equalizer smeared over the batch. +""" +import unittest + +import numpy as np + +from comnumpy.core.channels import AWGN, FIRChannel +from comnumpy.core.compensators import LinearEqualizer +from comnumpy.core.mappers import SymbolDemapper, SymbolMapper +from comnumpy.core.utils import Constellation, get_alphabet +from comnumpy.exceptions import ShapeError +from comnumpy.mimo.compensators import BlindDualMIMOCompensator +from comnumpy.ofdm.chains import OFDMReceiver, OFDMTransmitter +from comnumpy.optical.channels import PhaseNoise + + +class TestDeterministicBlocksBroadcast(unittest.TestCase): + """Row i of a batched call equals the block applied to row i.""" + + def setUp(self): + rng = np.random.default_rng(0) + self.x = (rng.standard_normal((3, 64)) + + 1j * rng.standard_normal((3, 64))) + self.h = np.array([1.0, 0.4 + 0.2j, 0.1]) + + def test_fir_channel_all_modes(self): + for mode in ("full", "same", "valid"): + with self.subTest(mode=mode): + y = FIRChannel(self.h, mode=mode)(self.x) + row = FIRChannel(self.h, mode=mode)(self.x[1]) + self.assertEqual(y.shape[:-1], (3,)) + np.testing.assert_allclose(y[1], row) + + def test_linear_equalizer_both_methods(self): + x = FIRChannel(self.h)(self.x) + for method in ("zf", "mmse"): + with self.subTest(method=method): + y = LinearEqualizer(self.h, method=method, sigma2=0.1)(x) + row = LinearEqualizer(self.h, method=method, + sigma2=0.1)(x[1]) + np.testing.assert_allclose(y[1], row) + + def test_filters_impairments_and_cd_compensators(self): + from comnumpy.core.filters import BWFilter, SRRCFilter + from comnumpy.core.impairments import CFO, Delay, IQImbalance + from comnumpy.core.processors import Downsampler, Upsampler + from comnumpy.optical.channels import ChromaticDispersion + from comnumpy.optical.compensators import ( + ChromaticDispersionFIRCompensator, + ChromaticDispersionLSFIRCompensator) + blocks = ( + ("Upsampler", lambda: Upsampler(2)), + ("Downsampler", lambda: Downsampler(2)), + ("SRRCFilter", lambda: SRRCFilter(0.25, 2, N_h=16)), + ("BWFilter", lambda: BWFilter(0.3)), + ("CFO", lambda: CFO(1e-3)), + ("Delay", lambda: Delay(3)), + ("IQImbalance", lambda: IQImbalance(1.0, 0.1)), + ("ChromaticDispersion", + lambda: ChromaticDispersion(1000, fs=32e9)), + ("CD FIR", + lambda: ChromaticDispersionFIRCompensator(1000, fs=32e9)), + ("CD LS FIR", + lambda: ChromaticDispersionLSFIRCompensator( + 1000, 15, fs=32e9, w_vect=[-np.pi, np.pi])), + ) + for name, factory in blocks: + with self.subTest(block=name): + batched = factory()(self.x) + row = factory()(self.x[1]) + np.testing.assert_allclose(batched[1], row) + + def test_mapper_demapper_and_ofdm_chains(self): + constellation = Constellation("QAM", 16) + rng = np.random.default_rng(1) + indices = rng.integers(0, 16, (3, 64)) + symbols = SymbolMapper(constellation)(indices) + waveform = OFDMTransmitter(16, 4)(symbols) + received = OFDMReceiver(16, 4, h=self.h)(waveform) + detected = SymbolDemapper(constellation)(received) + for name, batched, row in ( + ("mapper", symbols, SymbolMapper(constellation)(indices[1])), + ("tx", waveform, OFDMTransmitter(16, 4)(symbols[1])), + ("rx", received, OFDMReceiver(16, 4, h=self.h)(waveform[1])), + ("demapper", detected, + SymbolDemapper(constellation)(received[1]))): + with self.subTest(block=name): + np.testing.assert_allclose(batched[1], row) + + +class TestStochasticBlocksDrawIndependently(unittest.TestCase): + + def test_awgn_noise_differs_across_the_batch(self): + x = np.zeros((2, 4096), dtype=complex) + y = AWGN(sigma2=1.0, seed=0)(x) + self.assertFalse(np.allclose(y[0], y[1])) + + def test_one_laser_per_pair_independent_across_trials(self): + """per='pair': the two rows of a pair share the walk, and two + trials of a batch do not.""" + x = np.ones((3, 2, 512), dtype=complex) + y = PhaseNoise(1e-3, seed=0)(x) + phase = np.unwrap(np.angle(y), axis=-1) + np.testing.assert_allclose(phase[0, 0], phase[0, 1]) + self.assertFalse(np.allclose(phase[0, 0], phase[1, 0])) + + def test_a_bare_pair_still_shares_one_walk(self): + x = np.ones((2, 512), dtype=complex) + y = PhaseNoise(1e-3, seed=0)(x) + np.testing.assert_allclose(np.angle(y[0]), np.angle(y[1])) + + def test_the_ambiguous_shape_is_refused(self): + """(n, N) with n != 2 is not a pair and not declared batch: + the block refuses it and names both resolutions.""" + with self.assertRaises(ShapeError): + PhaseNoise(1e-3, seed=0)(np.ones((3, 64), dtype=complex)) + + def test_per_row_and_per_signal_resolve_it(self): + x = np.ones((3, 512), dtype=complex) + rows = PhaseNoise(1e-3, per="row", seed=0)(x) + self.assertFalse(np.allclose(np.angle(rows[0]), np.angle(rows[1]))) + shared = PhaseNoise(1e-3, per="signal", seed=0)(x) + np.testing.assert_allclose(np.angle(shared[0]), np.angle(shared[1])) + + +class TestAdaptiveBlocksCarryStatePerEvent(unittest.TestCase): + + def frames(self, n_trials=3, n=400, seed=0): + alphabet = np.asarray(get_alphabet("PSK", 4)) + rng = np.random.default_rng(seed) + s = alphabet[rng.integers(0, 4, (n_trials, 2, n))] + # a fixed rotation mixes the two polarizations of every trial + theta = 0.4 + rotation = np.array([[np.cos(theta), -np.sin(theta)], + [np.sin(theta), np.cos(theta)]]) + return alphabet, s, np.einsum("ij,kjn->kin", rotation, s) + + def test_each_pair_gets_its_own_equalizer(self): + """Pair i of a batched pass equals a fresh compensator run on + pair i alone -- the state is per event, not smeared.""" + alphabet, _, x = self.frames() + compensator = BlindDualMIMOCompensator(L=2, alphabet=alphabet, + mu=1e-3) + y = compensator(x) + assert compensator.H_ is not None + self.assertEqual(compensator.H_.shape, (3, 2, 2 * (2 * 2 + 1))) + alone = BlindDualMIMOCompensator(L=2, alphabet=alphabet, + mu=1e-3)(x[1]) + np.testing.assert_allclose(y[1], alone) + + def test_partial_fit_resumes_each_pair_from_its_own_state(self): + alphabet, _, x = self.frames() + batched = BlindDualMIMOCompensator(L=2, alphabet=alphabet, mu=1e-3) + batched(x) + batched(x) + alone = BlindDualMIMOCompensator(L=2, alphabet=alphabet, mu=1e-3) + alone(x[1]) + alone(x[1]) + assert batched.H_ is not None and alone.H_ is not None + np.testing.assert_allclose(batched.H_[1], alone.H_) + + def test_a_stale_batch_state_is_refused(self): + alphabet, _, x = self.frames(n_trials=3) + compensator = BlindDualMIMOCompensator(L=2, alphabet=alphabet, + mu=1e-3) + compensator(x) + _, _, other = self.frames(n_trials=4, seed=1) + with self.assertRaises(ShapeError): + compensator(other) + + def test_one_polarization_is_refused(self): + alphabet = np.asarray(get_alphabet("PSK", 4)) + compensator = BlindDualMIMOCompensator(L=2, alphabet=alphabet) + with self.assertRaises(ShapeError): + compensator(np.ones(64, dtype=complex)) + + def test_a_scalar_estimator_refuses_the_batch(self): + """BlindCFOCompensator optimizes ONE w0_ over the whole array: + on a batch that would smear a single estimate over independent + trials, so the batch is refused rather than misread.""" + from comnumpy.core.compensators import BlindCFOCompensator + with self.assertRaises(ShapeError): + BlindCFOCompensator()(np.ones((3, 64), dtype=complex)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/core/test_gray_labelling.py b/tests/core/test_gray_labelling.py new file mode 100644 index 0000000..41a7b75 --- /dev/null +++ b/tests/core/test_gray_labelling.py @@ -0,0 +1,110 @@ +"""The Gray promise, locked end to end. + +`compute_ber` expands symbol indices with `sym_2_bin` (natural binary), +and its docstring admits the count is only meaningful if the +constellation is labelled the same way. Nothing verified that promise: +the tutorials' BER curves matched theory in the tail, which is evidence +but not a lock. These tests are the lock. + +Two deterministic properties close the question without Monte-Carlo: +geometric neighbours must differ by exactly one bit of their *index* +(so index == label), and a nearest-neighbour symbol error must count +exactly one bit error through `compute_ber`. +""" +import unittest + +import numpy as np + +from comnumpy.core.metrics import compute_ber +from comnumpy.core.utils import Constellation, get_alphabet + + +def hamming(a, b): + return bin(int(a) ^ int(b)).count("1") + + +class TestGrayAdjacency(unittest.TestCase): + """Nearest neighbours of a Gray alphabet differ by one index bit.""" + + def nearest_pairs(self, alphabet): + """Every (i, j) pair at the minimum geometric distance.""" + points = np.asarray(alphabet) + distance = np.abs(points[:, None] - points[None, :]) + np.fill_diagonal(distance, np.inf) + minimum = distance.min() + pairs = np.argwhere(np.isclose(distance, minimum)) + return [(int(i), int(j)) for i, j in pairs] + + def test_square_qam_neighbours_differ_by_one_bit(self): + for order in (4, 16, 64, 256): + with self.subTest(order=order): + alphabet = get_alphabet("QAM", order) + for i, j in self.nearest_pairs(alphabet): + self.assertEqual( + hamming(i, j), 1, + f"QAM{order}: indices {i} and {j} are geometric " + f"neighbours but differ by {hamming(i, j)} bits -- " + f"compute_ber would overcount every symbol error") + + def test_psk_ring_neighbours_differ_by_one_bit(self): + for order in (4, 8, 16): + with self.subTest(order=order): + alphabet = get_alphabet("PSK", order) + for i, j in self.nearest_pairs(alphabet): + self.assertEqual(hamming(i, j), 1) + + def test_natural_binary_mapping_is_not_gray(self): + """The property is the labelling, not an accident of geometry.""" + alphabet = get_alphabet("QAM", 16, type="bin") + distances = [] + for i, j in self.nearest_pairs(alphabet): + distances.append(hamming(i, j)) + self.assertGreater(max(distances), 1) + + +class TestComputeBerCountsGrayBits(unittest.TestCase): + + def test_a_nearest_neighbour_error_costs_exactly_one_bit(self): + """The bridge between the labelling and the metric.""" + constellation = Constellation("QAM", 16) + alphabet = np.asarray(constellation.alphabet) + distance = np.abs(alphabet[:, None] - alphabet[None, :]) + np.fill_diagonal(distance, np.inf) + reference = np.arange(16) + detected = distance.argmin(axis=1) # each symbol -> a neighbour + ber = compute_ber(reference, detected, width=4) + self.assertAlmostEqual(float(ber), 1.0 / 4.0) + + def test_tail_ber_matches_the_gray_closed_form(self): + """One seeded run: measured BER within 20 % of Ps/k in the tail. + + At Eb/N0 = 11 dB a 16-QAM symbol error is almost surely a + nearest-neighbour event, so BER = SER/4 holds tightly; a + labelling mismatch would multiply the measured BER by ~2 and + fail this comfortably. Seeded, hence deterministic. + """ + from comnumpy.core import Sequential + from comnumpy.core.channels import AWGN + from comnumpy.core.generators import SymbolGenerator + from comnumpy.core.mappers import SymbolDemapper, SymbolMapper + from comnumpy.core.utils import ebn0_to_snr_dB + + constellation = Constellation("QAM", 16) + chain = Sequential([ + SymbolGenerator(16, name="tx"), + SymbolMapper(constellation), + AWGN(snr_dB=float(ebn0_to_snr_dB(11.0, bits_per_symbol=4))), + SymbolDemapper(constellation), + ], taps=["tx"]) + chain.seed(3) + detected = chain(400000) + measured = float(compute_ber(chain.tap("tx"), detected, width=4)) + theory = float(constellation.metrics(11.0, per="bit")["ber"]) + self.assertGreater(measured, 0.0, "no errors: raise N or the SNR") + self.assertLess(abs(measured - theory) / theory, 0.20, + f"measured {measured:.3e} vs Gray closed form " + f"{theory:.3e} -- a labelling mismatch doubles this") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/mimo/test_blind_equalizer.py b/tests/mimo/test_blind_equalizer.py index d52af20..ac09a73 100644 --- a/tests/mimo/test_blind_equalizer.py +++ b/tests/mimo/test_blind_equalizer.py @@ -155,7 +155,7 @@ def test_a_single_polarization_is_refused(self): with self.assertRaises(ValueError) as ctx: BlindDualMIMOCompensator(alphabet=alphabet())(np.zeros((1, 100), dtype=complex)) - self.assertIn("dual polarization", str(ctx.exception)) + self.assertIn("polarization pair", str(ctx.exception)) def test_an_unknown_mode_says_which_ones_exist(self): equalizer = BlindDualMIMOCompensator(L=2, alphabet=alphabet()) diff --git a/tests/mimo/test_stacked_channel.py b/tests/mimo/test_stacked_channel.py new file mode 100644 index 0000000..d9da437 --- /dev/null +++ b/tests/mimo/test_stacked_channel.py @@ -0,0 +1,95 @@ +"""The stacked MIMO channel: one draw per trial, detectors aligned. + +The MIMO Monte-Carlo redraws the channel every frame, and every +detector needs that frame's matrix. The batch contract (D51) makes the +draw a *stack*: ``rayleigh_channel(size=K)`` returns ``(K, N_r, N_t)``, +``FlatMIMOChannel`` applies channel k to frame k, and a detector built +with the same stack decides frame k against channel k. The linear +detectors batch through numpy's stacked linalg; the search detectors +(ML, OSIC, sphere) loop per draw internally -- convenience and +correctness, not a speedup, and the docstrings say so. +""" +import unittest + +import numpy as np + +from comnumpy.core.utils import Constellation +from comnumpy.exceptions import ShapeError +from comnumpy.mimo.channels import FlatMIMOChannel +from comnumpy.mimo.detectors import ( + LinearDetector, MaximumLikelihoodDetector, + OrderedSuccessiveInterferenceCancellationDetector, SphereDecoder) +from comnumpy.mimo.utils import rayleigh_channel + +CONSTELLATION = Constellation("PSK", 4) +K, N_R, N_T, N = 4, 3, 2, 60 + + +def factories(): + return { + "ZF": lambda H: LinearDetector(CONSTELLATION, H=H, method="zf"), + "MMSE": lambda H: LinearDetector(CONSTELLATION, H=H, + method="mmse", sigma2=0.02), + "ML": lambda H: MaximumLikelihoodDetector(CONSTELLATION, H=H), + "OSIC": lambda H: OrderedSuccessiveInterferenceCancellationDetector( + CONSTELLATION, osic_type="sinr", H=H, sigma2=0.02), + "SD": lambda H: SphereDecoder(CONSTELLATION, H=H), + } + + +class TestStackedDraws(unittest.TestCase): + + def test_size_grows_a_leading_axis(self): + self.assertEqual(rayleigh_channel(N_R, N_T, seed=0, size=K).shape, + (K, N_R, N_T)) + + def test_size_equals_the_python_loop_draw_for_draw(self): + """Moving a study to the batched form keeps its channels.""" + first = np.random.default_rng(7) + second = np.random.default_rng(7) + looped = [] + for _ in range(5): + looped.append(rayleigh_channel(N_R, N_T, rng=first)) + stacked = rayleigh_channel(N_R, N_T, rng=second, size=5) + np.testing.assert_allclose(np.stack(looped), stacked) + + +class TestStackedChannelAndDetectors(unittest.TestCase): + + def setUp(self): + rng = np.random.default_rng(0) + alphabet = np.asarray(CONSTELLATION.alphabet) + self.H = rayleigh_channel(N_R, N_T, seed=1, size=K) + self.sent = rng.integers(0, 4, (K, N_T, N)) + noise = 0.1 * (rng.standard_normal((K, N_R, N)) + + 1j * rng.standard_normal((K, N_R, N))) + self.Y = np.matmul(self.H, alphabet[self.sent]) + noise + + def test_channel_k_is_applied_to_frame_k(self): + x = np.asarray(CONSTELLATION.alphabet)[self.sent] + batched = FlatMIMOChannel(self.H)(x) + single = FlatMIMOChannel(self.H[2])(x[2]) + np.testing.assert_allclose(batched[2], single) + + def test_every_detector_decides_frame_k_against_channel_k(self): + for name, make in factories().items(): + with self.subTest(detector=name): + stacked = make(self.H)(self.Y) + single = make(self.H[2])(self.Y[2]) + self.assertEqual(stacked.shape, (K, N_T, N)) + np.testing.assert_array_equal(stacked[2], single) + + def test_a_mismatched_batch_is_refused(self): + """K channels cannot decide K+1 frames: no silent broadcast.""" + detector = MaximumLikelihoodDetector(CONSTELLATION, H=self.H) + with self.assertRaises(ShapeError): + detector(np.ones((K + 1, N_R, N), dtype=complex)) + + def test_the_sphere_decoder_still_counts_its_nodes(self): + decoder = SphereDecoder(CONSTELLATION, H=self.H) + decoder(self.Y) + self.assertGreater(decoder.nodes_, 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/optical/test_pmd_bps.py b/tests/optical/test_pmd_bps.py new file mode 100644 index 0000000..cf8d77c --- /dev/null +++ b/tests/optical/test_pmd_bps.py @@ -0,0 +1,161 @@ +"""The two blocks the PDM page stands on: the emulator and the tracker. + +`PMDEmulator` claims to be unitary, seeded and frequency-selective; +`BlindPhaseSearchCompensator` claims to follow a Wiener phase within +its grid and its window. Each claim is a test, because the tutorial +built on them quotes their numbers as physics. +""" +import unittest + +import numpy as np + +from comnumpy.core.compensators import BlindPhaseSearchCompensator +from comnumpy.core.utils import get_alphabet +from comnumpy.exceptions import ShapeError +from comnumpy.optical.channels import PMDEmulator + + +class TestPMDEmulator(unittest.TestCase): + + def field_pair(self, n=256, seed=0): + rng = np.random.default_rng(seed) + return (rng.standard_normal((2, n)) + + 1j * rng.standard_normal((2, n))) + + def test_energy_is_conserved(self): + """The emulator is a product of unitaries: lossless, exactly.""" + x = self.field_pair() + y = PMDEmulator(5e-12, n_sections=8, fs=64e9, seed=1)(x) + self.assertAlmostEqual(float(np.sum(np.abs(y) ** 2)), + float(np.sum(np.abs(x) ** 2)), places=8) + + def test_same_seed_same_fibre(self): + x = self.field_pair() + first = PMDEmulator(5e-12, fs=64e9, seed=3)(x) + second = PMDEmulator(5e-12, fs=64e9, seed=3)(x) + np.testing.assert_allclose(first, second) + third = PMDEmulator(5e-12, fs=64e9, seed=4)(x) + self.assertFalse(np.allclose(first, third)) + + def test_a_rotation_mixes_the_polarizations(self): + x = np.zeros((2, 64), dtype=complex) + x[0] = 1.0 + y = PMDEmulator(0.0, n_sections=1, seed=1)(x) + self.assertGreater(float(np.mean(np.abs(y[1]) ** 2)), 1e-3) + + def test_zero_dgd_is_frequency_flat(self): + """Without DGD the emulator is one wavelength-flat Jones matrix: + applying it twice with the same seed equals composing the + matrices, and every frequency sees the same rotation -- a pure + tone in must be a pure tone out.""" + n = 128 + tone = np.exp(2j * np.pi * 13 * np.arange(n) / n) + x = np.stack([tone, np.zeros(n, dtype=complex)]) + y = PMDEmulator(0.0, n_sections=4, seed=2)(x) + spectrum = np.abs(np.fft.fft(y, axis=-1)) + occupied = np.sum(spectrum > 1e-6 * spectrum.max(), axis=-1) + self.assertTrue(np.all(occupied <= 1)) + + def test_dgd_delays_the_principal_states_against_each_other(self): + """One section, aligned input: the group delay difference between + the two rows is the declared tau.""" + fs, tau = 64e9, 20e-12 + n = 4096 + rng = np.random.default_rng(5) + pulse = rng.standard_normal(n) + 1j * rng.standard_normal(n) + emulator = PMDEmulator(tau, n_sections=1, fs=fs, seed=7) + assert emulator.rotations_ is not None + rotation = emulator.rotations_[0] + # drive the two principal states one at a time, through the + # inverse rotation, so each row sees a pure delay of +-tau/2 + for row, sign in ((0, +1), (1, -1)): + x = np.zeros((2, n), dtype=complex) + x[row] = pulse + y = emulator(np.einsum("ij,jn->in", rotation.conj().T, x)) + cross = np.fft.ifft(np.fft.fft(y[row]) * + np.conj(np.fft.fft(pulse))) + lag = int(np.argmax(np.abs(cross))) + lag = lag - n if lag > n // 2 else lag + # multiplying the spectrum by exp(+j w tau/2) advances the + # signal (fft convention e^{-jwn}), hence the minus sign + self.assertEqual(lag, -round(sign * tau / 2 * fs)) + + def test_rms_dgd_matches_the_declared_value(self): + """The theory the emulator must obey: section DGDs add in + quadrature, and the ensemble DGD is Maxwellian with RMS equal + to the declared tau (Poole & Wagner). Measured at one frequency + over many section draws, from the eigenvalues of the + group-delay operator j dT/dw T^H.""" + fs, tau, n = 64e9, 10e-12, 512 + dgds = [] + for seed in range(300): + emulator = PMDEmulator(tau, n_sections=8, fs=fs, seed=seed) + responses = [] + for k in (1, 2): # T at two neighbouring frequencies + tone = np.exp(2j * np.pi * k * np.arange(n) / n) + columns = [] + for row in range(2): + x = np.zeros((2, n), dtype=complex) + x[row] = tone + y = emulator(x) + columns.append(np.fft.fft(y, axis=-1)[:, k] / n) + responses.append(np.stack(columns, axis=1)) + delta_omega = 2 * np.pi * fs / n + u = responses[1] @ responses[0].conj().T + phases = np.angle(np.linalg.eigvals(u)) + dgds.append(abs(phases[0] - phases[1]) / delta_omega) + rms = float(np.sqrt(np.mean(np.square(dgds)))) + self.assertLess(abs(rms - tau) / tau, 0.10) + mean = float(np.mean(dgds)) + self.assertLess(abs(mean - np.sqrt(8 / (3 * np.pi)) * tau) / tau, + 0.10) + + def test_one_polarization_is_refused(self): + with self.assertRaises(ShapeError): + PMDEmulator(1e-12, fs=64e9)(np.ones(64, dtype=complex)) + + +class TestBlindPhaseSearch(unittest.TestCase): + + def qpsk_frame(self, n=4096, seed=0): + alphabet = get_alphabet("PSK", 4) + rng = np.random.default_rng(seed) + return alphabet, alphabet[rng.integers(0, 4, n)] + + def test_tracks_a_wiener_walk_within_the_grid(self): + alphabet, s = self.qpsk_frame() + rng = np.random.default_rng(1) + walk = np.cumsum(rng.normal(0.0, np.sqrt(1e-4), s.size)) + compensator = BlindPhaseSearchCompensator(alphabet) + y = compensator(s * np.exp(1j * walk)) + assert compensator.phase_ is not None + interior = slice(32, -32) # the window edges are biased + self.assertLess( + float(np.max(np.abs(compensator.phase_ - walk)[interior])), + 2 * (np.pi / 2) / 32, + "the estimate left the walk by more than two grid steps") + self.assertLess(float(np.max(np.abs(y - s)[interior])), 0.2) + + def test_rows_of_a_pair_get_independent_trajectories(self): + alphabet, s = self.qpsk_frame(n=1024) + x = np.stack([s * np.exp(1j * 0.3), s * np.exp(-1j * 0.2)]) + compensator = BlindPhaseSearchCompensator(alphabet, half_window=8) + compensator(x) + assert compensator.phase_ is not None + # the estimate lives modulo pi/2 (quadrant symmetry): compare + # each trajectory to its truth in that quotient + for row, truth in ((0, 0.3), (1, -0.2)): + residual = np.median(compensator.phase_[row]) - truth + residual = (residual + np.pi / 4) % (np.pi / 2) - np.pi / 4 + self.assertAlmostEqual(float(residual), 0.0, delta=0.03) + + def test_the_quadrant_ambiguity_is_documented_not_resolved(self): + """A pi/2 offset is invisible to the search, by symmetry.""" + alphabet, s = self.qpsk_frame(n=1024) + y = BlindPhaseSearchCompensator(alphabet, + half_window=8)(s * 1j) + self.assertLess(float(np.max(np.abs(y - s * 1j))), 0.06) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_batch_contract.py b/tests/test_batch_contract.py new file mode 100644 index 0000000..9778f92 --- /dev/null +++ b/tests/test_batch_contract.py @@ -0,0 +1,365 @@ +"""The batch promise, verified over the whole catalogue. + +CONVENTIONS.md (D51) promises that the axes ahead of a block's event +axes are batch axes: independent trials in one call. A promise over a +catalogue is only worth the sweep that checks it, so this file is a +*ratchet*, like the pyright one (D37): every ``Processor`` subclass in +the library must appear in exactly one register below, and a discovery +test fails on any block that has not declared its batch behaviour. +Adding a block therefore *forces* the author to say -- and prove -- +what a batch means for it. + +Three verified families, one honest bucket: + +* ``BROADCAST`` -- deterministic: row ``i`` of a batched call equals + the block applied to row ``i`` alone. The strongest property, checked + by equality. +* ``INDEPENDENT`` -- stochastic: the batched call runs, keeps the + leading axis, and two rows do not share a realization. +* ``REFUSES`` -- the block cannot read a batch one way and must say so: + a ``ShapeError`` (or documented ``ValueError``), never a silent + misread. +* ``EXEMPT`` -- blocks the sweep cannot exercise generically, each with + the reason written next to it. An exemption is visible debt, not a + pass. +""" +import importlib +import inspect +import pkgutil +import unittest + +import numpy as np + +from comnumpy.core.generics import Processor +from comnumpy.core.utils import Constellation + +RNG = np.random.default_rng(0) +N = 64 +SERIAL = RNG.standard_normal((3, N)) + 1j * RNG.standard_normal((3, N)) +INDICES = RNG.integers(0, 4, (3, N)) +BITS = RNG.integers(0, 2, (3, N)) +QPSK = Constellation("PSK", 4) +PAIR = RNG.standard_normal((3, 2, N)) + 1j * RNG.standard_normal((3, 2, N)) + + +def _fs(): + return {"fs": 32e9} + + +# --- the registers --------------------------------------------------- +# name: (factory, batched input). The factory builds a FRESH block each +# call -- the equality check runs it twice. + +BROADCAST = { + "core.channels.FIRChannel": + (lambda: _cls("core.channels", "FIRChannel")(np.array([1.0, 0.4])), + SERIAL), + "core.compensators.LinearEqualizer": + (lambda: _cls("core.compensators", "LinearEqualizer")( + np.array([1.0, 0.4])), SERIAL), + "core.compensators.BlindIQCompensator": + (lambda: _cls("core.compensators", "BlindIQCompensator")(), SERIAL), + "core.compensators.Normalizer": + (lambda: _cls("core.compensators", "Normalizer")("max"), SERIAL), + "core.compensators.DCCorrector": + (lambda: _cls("core.compensators", "DCCorrector")(), SERIAL), + "core.devices.RappAmplifier": + (lambda: _cls("core.devices", "RappAmplifier")(1.0), SERIAL), + "core.devices.SalehAmplifier": + (lambda: _cls("core.devices", "SalehAmplifier")(), SERIAL), + "core.filters.BWFilter": + (lambda: _cls("core.filters", "BWFilter")(0.3), SERIAL), + "core.filters.SRRCFilter": + (lambda: _cls("core.filters", "SRRCFilter")(0.25, 2, N_h=8), SERIAL), + "core.impairments.CFO": + (lambda: _cls("core.impairments", "CFO")(1e-3), SERIAL), + "core.impairments.Delay": + (lambda: _cls("core.impairments", "Delay")(3), SERIAL), + "core.impairments.IQImbalance": + (lambda: _cls("core.impairments", "IQImbalance")(1.0, 0.1), SERIAL), + "core.mappers.SymbolMapper": + (lambda: _cls("core.mappers", "SymbolMapper")(QPSK), INDICES), + "core.mappers.SymbolDemapper": + (lambda: _cls("core.mappers", "SymbolDemapper")(QPSK), SERIAL), + "core.processors.Amplifier": + (lambda: _cls("core.processors", "Amplifier")(2.0), SERIAL), + "core.processors.Clipper": + (lambda: _cls("core.processors", "Clipper")(0.8), SERIAL), + "core.processors.Complex2Real": + (lambda: _cls("core.processors", "Complex2Real")(), SERIAL), + "core.processors.Downsampler": + (lambda: _cls("core.processors", "Downsampler")(2), SERIAL), + "core.processors.Upsampler": + (lambda: _cls("core.processors", "Upsampler")(2), SERIAL), + "core.processors.Serial2Parallel": + (lambda: _cls("core.processors", "Serial2Parallel")(8), SERIAL), + "core.processors.Parallel2Serial": + (lambda: _cls("core.processors", "Parallel2Serial")(), + SERIAL.reshape(3, 8, 8)), + "fec.convolutional.ConvolutionalEncoder": + (lambda: _cls("fec.convolutional", "ConvolutionalEncoder")(), BITS), + "mimo.channels.FlatMIMOChannel": + (lambda: _cls("mimo.channels", "FlatMIMOChannel")( + _rayleigh(2, 2, seed=1)), PAIR), + "mimo.detectors.LinearDetector": + (lambda: _cls("mimo.detectors", "LinearDetector")( + QPSK, H=_rayleigh(2, 2, seed=1), method="zf"), PAIR), + "ofdm.chains.OFDMTransmitter": + (lambda: _cls("ofdm.chains", "OFDMTransmitter")(16, 4), SERIAL), + "ofdm.chains.OFDMReceiver": + (lambda: _cls("ofdm.chains", "OFDMReceiver")(16, 4), + SERIAL[:, :60].repeat(2, axis=-1)[:, :80].reshape(3, 80)), + "ofdm.compensators.FrequencyDomainEqualizer": + (lambda: _cls("ofdm.compensators", "FrequencyDomainEqualizer")( + h=np.array([1.0, 0.3])), SERIAL.reshape(3, 4, 16)), + "ofdm.processors.FFTProcessor": + (lambda: _cls("ofdm.processors", "FFTProcessor")(), + SERIAL.reshape(3, 4, 16)), + "ofdm.processors.IFFTProcessor": + (lambda: _cls("ofdm.processors", "IFFTProcessor")(), + SERIAL.reshape(3, 4, 16)), + "ofdm.processors.CyclicPrefixer": + (lambda: _cls("ofdm.processors", "CyclicPrefixer")(4), + SERIAL.reshape(3, 4, 16)), + "ofdm.processors.CyclicPrefixRemover": + (lambda: _cls("ofdm.processors", "CyclicPrefixRemover")(4), + SERIAL.reshape(3, 4, 16)), + "ofdm.predistorders.HardClipper": + (lambda: _cls("ofdm.predistorders", "HardClipper")(0.8), SERIAL), + "optical.channels.ChromaticDispersion": + (lambda: _cls("optical.channels", "ChromaticDispersion")( + 100, **_fs()), SERIAL), + "optical.channels.KerrNonLinearity": + (lambda: _cls("optical.channels", "KerrNonLinearity")(1e-3), SERIAL), + "optical.compensators.ChromaticDispersionFIRCompensator": + (lambda: _cls("optical.compensators", + "ChromaticDispersionFIRCompensator")(100, **_fs()), + SERIAL), + "optical.compensators.ChromaticDispersionLSFIRCompensator": + (lambda: _cls("optical.compensators", + "ChromaticDispersionLSFIRCompensator")( + 100, 15, w_vect=[-np.pi, np.pi], **_fs()), SERIAL), +} + +INDEPENDENT = { + "core.channels.AWGN": + (lambda: _cls("core.channels", "AWGN")(sigma2=1.0), + np.zeros((3, 4096), dtype=complex)), + "core.generators.SymbolGenerator": + (lambda: _cls("core.generators", "SymbolGenerator")(4), (3, 256)), + "core.generators.GaussianGenerator": + (lambda: _cls("core.generators", "GaussianGenerator")(), (3, 256)), + "optical.channels.PhaseNoise": + (lambda: _cls("optical.channels", "PhaseNoise")(1e-3, per="row"), + np.ones((3, 512), dtype=complex)), +} + +REFUSES = { + "core.compensators.BlindCFOCompensator": + (lambda: _cls("core.compensators", "BlindCFOCompensator")(), + np.ones((3, N), dtype=complex)), + "core.compensators.DataAidedPhaseCompensator": + (lambda: _cls("core.compensators", "DataAidedPhaseCompensator")( + reference=SERIAL[0]), SERIAL), + "core.compensators.DataAidedComplexGainCompensator": + (lambda: _cls("core.compensators", + "DataAidedComplexGainCompensator")( + reference=SERIAL[0]), SERIAL), + "core.compensators.DataAidedFIRCompensator": + (lambda: _cls("core.compensators", "DataAidedFIRCompensator")( + reference=SERIAL[0]), SERIAL), +} + +EXEMPT = { + # abstract bases and containers + "mimo.channels.BaseMIMOChannel": "abstract base, forward raises", + # generic event = the whole 1-D stream by declared design; a batch + # has no meaning the docstring does not already refuse + "core.compensators.DataAidedFineSynchronizer": + "estimates one delay against one 1-D reference", + "core.compensators.DataAidedSimpleSynchronizer": + "estimates one delay against one 1-D reference", + "core.frames.Framer": "frames one 1-D stream against a FrameStructure", + "core.frames.Deframer": "deframes one 1-D stream", + "core.processors.AutoConcatenator": "concatenates along the last axis " + "with internal history: one stream per instance", + "core.processors.BlindPhaseTracker": "sequential per-sample tracker, " + "one trajectory per instance (superseded by " + "BlindPhaseSearchCompensator for pairs)", + "core.processors.DataAdder": "adds a stored companion signal of the " + "un-batched shape", + "core.processors.DataExtractor": "extracts against stored 1-D markers", + "core.processors.DelayRemover": "drops a scalar delay from one stream", + "core.processors.SampleRemover": "drops N samples of one stream", + "core.processors.Resampler": "scipy.signal.resample along a declared " + "axis; not swept here because its output length depends on the " + "rate argument", + "core.processors.WeightAmplifier": "per-branch gains tied to the " + "un-batched branch axis", + "core.compensators.BlindPhaseCompensation": "legacy single-signal " + "moment estimator (superseded by BlindPhaseSearchCompensator)", + "core.shaping.AmplitudeMapper": "shaping operates on one 1-D " + "amplitude stream", + "core.shaping.AmplitudeDemapper": "shaping operates on one 1-D " + "amplitude stream", + "core.shaping.DistributionMatcher": "arithmetic-coding state over one " + "1-D bit stream", + "core.shaping.DistributionDematcher": "arithmetic-coding state over " + "one 1-D bit stream", + "core.channels.TappedDelayLineChannel": "frozen-realization channel: " + "one seeded draw is the configuration (D51), sounded via " + "impulse_response", + "fec.convolutional.ViterbiDecoder": "decodes (n_cw, n) codeword " + "blocks -- its batch axis is the codeword axis, locked in " + "tests/fec", + "fec.ldpc.LDPCEncoder": "operates on (k, n_cw) codeword blocks, " + "locked in tests/fec", + "fec.ldpc.LDPCDecoder": "operates on (n, n_cw) LLR blocks, locked " + "in tests/fec", + "mimo.channels.SelectiveMIMOChannel": "builds an (L N_t, N) stacked " + "convolution matrix for one frame", + "mimo.coding.SpaceTimeEncoder": "Alamouti layout (N_t, N) per code " + "block, one frame per call", + "mimo.coding.SpaceTimeDecoder": "Alamouti layout per code block, one " + "frame per call", + "mimo.compensators.BlindDualMIMOCompensator": "adaptive per pair; its " + "batch (one butterfly per pair) is locked in " + "tests/core/test_batch_axes.py", + "core.compensators.BlindPhaseSearchCompensator": "adaptive per row; " + "its per-row trajectories are locked in " + "tests/optical/test_pmd_bps.py", + "mimo.detectors.MaximumLikelihoodDetector": "stacked-H batch locked " + "in tests/mimo/test_stacked_channel.py", + "mimo.detectors.SphereDecoder": "stacked-H batch locked in " + "tests/mimo/test_stacked_channel.py", + "mimo.detectors.OrderedSuccessiveInterferenceCancellationDetector": + "stacked-H batch locked in tests/mimo/test_stacked_channel.py", + "mimo.detectors.ApproximateMessagePassingDetector": "iterative " + "per-frame message passing against one H", + "mimo.detectors.OrthogonalApproximateMessagePassingDetector": + "iterative per-frame message passing against one H", + "ofdm.processors.CarrierAllocator": "validated (..., T, F_data) " + "layout, locked in tests/ofdm", + "ofdm.processors.CarrierExtractor": "validated (..., T, F) layout, " + "locked in tests/ofdm", + "ofdm.processors.HermitianPrefixer": "builds the Hermitian spectrum " + "of one block set", + "ofdm.predistorders.IctPaprReductor": "iterative clipping on one " + "block set", + "ofdm.predistorders.PtsPaprReductor": "phase search over one block " + "set", + "optical.channels.PMDEmulator": "frozen-realization channel (D51): " + "polarization pair is the event, locked in " + "tests/optical/test_pmd_bps.py", + "optical.dbp.DBP": "SSFM split-step over one field, span state along " + "the last axis; validated in validation/", + "optical.devices.ErbiumDopedFiberAmplifier": "gain plus ASE draw " + "sized on the whole input; the optical event is the field", + "optical.devices.Laser": "one laser per instance by physics; its " + "sharing convention lives in PhaseNoise per=", + "optical.devices.MachZehnderModulator": "drives one modulator with " + "one electrical signal", + "optical.devices.Optical90HybridCircuit": "mixes one signal pair by " + "construction", + "optical.devices.PowerControl": "normalizes one field's power", + "optical.links.FiberLink": "SSFM propagation of one field; EDFA " + "noise inside; validated against closed forms in validation/", + "optical.wdm.WDMMultiplexer": "combines a channel *set* -- its " + "leading axis is the WDM channel, not a batch", + "optical.wdm.WDMDemultiplexer": "splits into a channel set", +} + + +def _cls(module: str, name: str): + return getattr(importlib.import_module(f"comnumpy.{module}"), name) + + +def _rayleigh(n_r: int, n_t: int, seed: int) -> np.ndarray: + from comnumpy.mimo.utils import rayleigh_channel + return rayleigh_channel(n_r, n_t, seed=seed) + + +def discover() -> dict[str, type]: + """Every Processor subclass defined in the library, by short name.""" + found = {} + for pkg in ("core", "optical", "mimo", "ofdm", "fec"): + package = importlib.import_module(f"comnumpy.{pkg}") + for info in pkgutil.iter_modules(package.__path__): + if info.name.startswith("_"): + continue + module = importlib.import_module(f"comnumpy.{pkg}.{info.name}") + for name, cls in inspect.getmembers(module, inspect.isclass): + if (issubclass(cls, Processor) and cls is not Processor + and cls.__module__ == module.__name__): + short = f"{pkg}.{info.name}.{name}" + found[short] = cls + return found + + +class TestTheCatalogueIsCovered(unittest.TestCase): + """The ratchet: a new block must declare its batch behaviour.""" + + def test_every_block_is_declared_exactly_once(self): + declared = (set(BROADCAST) | set(INDEPENDENT) | set(REFUSES) + | set(EXEMPT)) + discovered = set(discover()) + missing = sorted(discovered - declared) + self.assertFalse( + missing, + f"blocks with no declared batch behaviour: {missing} -- add " + f"each to BROADCAST, INDEPENDENT or REFUSES in " + f"tests/test_batch_contract.py with a probe input, or to " + f"EXEMPT with the reason written out (D51)") + stale = sorted(declared - discovered) + self.assertFalse(stale, f"declared but no longer found: {stale}") + for name in sorted(declared): + registers = [name in r for r in + (BROADCAST, INDEPENDENT, REFUSES, EXEMPT)] + self.assertEqual(sum(registers), 1, + f"{name} appears in {sum(registers)} registers") + + +class TestBroadcastFamily(unittest.TestCase): + """Row i of a batched call equals the block applied to row i.""" + + def test_row_equality(self): + for name, (factory, x) in BROADCAST.items(): + with self.subTest(block=name): + batched = factory()(np.asarray(x)) + single = factory()(np.asarray(x)[1]) + self.assertEqual(batched.shape[0], 3, + f"{name} lost the batch axis") + np.testing.assert_allclose( + np.asarray(batched)[1], np.asarray(single), + err_msg=f"{name}: batched row 1 differs from the " + f"block applied to row 1 alone") + + +class TestIndependentFamily(unittest.TestCase): + """Stochastic blocks draw independently per trial.""" + + def test_rows_do_not_share_a_realization(self): + for name, (factory, x) in INDEPENDENT.items(): + with self.subTest(block=name): + y = np.asarray(factory()(x)) + self.assertEqual(y.shape[0], 3, + f"{name} lost the batch axis") + self.assertFalse( + np.array_equal(y[0], y[1]), + f"{name}: two trials of a batch share a realization") + + +class TestRefusesFamily(unittest.TestCase): + """What cannot be read one way is refused, never misread.""" + + def test_the_batch_is_refused(self): + for name, (factory, x) in REFUSES.items(): + with self.subTest(block=name): + with self.assertRaises(ValueError, + msg=f"{name} accepted a batch it " + f"cannot interpret"): + factory()(x) + + +if __name__ == "__main__": + unittest.main()