Let the kernels take their array library from their inputs

The propagation kernels are the only part of this worth a GPU: a loop over FFTs
on a long array, where everything else is scalar arithmetic or a closed form
evaluated a few dozen times. They were written as array-to-array functions from
the beginning so this could be added without touching anything above them, and
`maiman.backend` is the whole of the addition.

The arrays decide, not a setting. A kernel handed CuPy arrays runs on CuPy and
returns CuPy arrays; handed NumPy arrays it runs on NumPy. No global mode and no
flag on the context — the kernels are pure functions, and a hidden mode would be
the one piece of state that could make the same inputs give two answers.
Dispatch is on the array's own type, so there is no registry to keep in sync.

CuPy is not exercised here and the docs say so: there is no device and no
install in CI. What is tested is the half that would break a port. A second
array library sets `__array_function__` to None — NumPy's own way for a type to
say it is not NumPy's — so `np.fft.fft` on one of its arrays raises. Universal
functions are deliberately left working, because `np.exp` on a CuPy array
dispatches and comes back a CuPy array; refusing them would be testing a rule
that is not true. What breaks a port is anything that *allocates*, and that is
what the guard covers.

Every converted kernel runs on both libraries and the results are compared, and
the names the second was asked for are recorded and asserted as an equality
rather than a lower bound: thirteen, all of which CuPy provides. A change that
reaches for something only NumPy has now fails in this repository rather than on
somebody's GPU, and one that stops needing something fails too.

Only the propagation path is converted, and a test names which functions are in
and which are out so the line is a decision rather than an oversight.

9 of 9 sabotages caught, after two got through. One was a host copy taken with
`np.asarray` instead of the array's own `get`, which works on a NumPy subclass
and raises on a real device array — the stub grew a `get` that records being
called. The other was the split-step building its frequency grid on the host:
the numbers come out right, because multiplying a host array by a device array
works by transferring the host one, once per step, over a bus. Correctness could
not see it, so the test asserts the allocation instead.

860 tests.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Author: ehsun-sh

README.md

@@ -447,6 +447,43 @@ noise bins rendered onto one grid, the way an instrument shows them. Its resolut
 not cosmetic — widening it raises the ASE trace decibel for decibel and leaves a carrier exactly
 where it is, which is the clearest demonstration of why OSNR needs a stated reference bandwidth.
 
+## The kernels do not know what they are running on
+
+The propagation kernels are the only part of this worth a GPU: a loop over FFTs on a long array,
+where everything else is scalar arithmetic or a closed form evaluated a few dozen times.
+[`maiman/kernels.py`](src/maiman/kernels.py) was written as array-to-array functions from the
+beginning so that this could be added without touching anything above it, and
+[`maiman/backend.py`](src/maiman/backend.py) is the whole of the addition.
+
+**The arrays decide, not a setting.** A kernel handed CuPy arrays runs on CuPy and returns CuPy
+arrays; handed NumPy arrays it runs on NumPy. There is no global mode and no flag on the context,
+which matters because the kernels are pure functions and a hidden mode would be the one piece of
+state that could make the same inputs give different answers. Dispatch is on the array's own type —
+`type(a).__module__` names the package it came from — so there is no registry to keep in sync.
+
+**CuPy is not exercised here**, and saying otherwise would be the kind of claim this project exists
+to avoid: there is no device and no install in CI. What *is* tested is the half that would actually
+break a port. A second array library — [`tests/hostile_backend.py`](tests/hostile_backend.py) — sets
+`__array_function__` to `None`, which is NumPy's own way for a type to say it is not NumPy's, so
+`np.fft.fft` on one of its arrays raises. Universal functions are deliberately left working, because
+`np.exp` on a CuPy array dispatches and comes back a CuPy array; refusing them would be testing a
+rule that is not true. What breaks a port is anything that *allocates* — `np.fft.fftfreq` and
+`np.zeros` build on the host, and a kernel calling one inside its loop pays a transfer every step.
+
+Every kernel is then run on both libraries and the results compared, and the names the second one
+was asked for are recorded. That set is the contract, and it is asserted as an equality rather than
+a lower bound:
+
+    abs  complex128  conj  exp  float64  max  pi
+    fft.fft  fft.fftfreq  fft.ifft  fft.irfft  fft.rfft  fft.rfftfreq
+
+Thirteen names. A change that reaches for something only NumPy has fails in this repository rather
+than on somebody's GPU, and one that stops needing something fails too. CuPy provides all thirteen.
+
+Only the propagation path is converted. The four-wave-mixing closed forms and the 2×2 Jones algebra
+are scalar work a device would slow down, and there is a test naming which functions are in and
+which are out, so the line is a decision rather than an oversight.
+
 ## 400G and 800G, and what they cost
 
 Three reference transceivers, all dual-polarization coherent, all derived from one number and
@@ -966,7 +1003,7 @@ time window, and results are reproducible.
 | **1 — MVP: linear link** *(essentially done)* | ✅ PRBS → NRZ → laser → MZM → fiber (α + CD) → PIN → filter → eye/Q/BER, validated end to end. **Python only, no GUI.** | ~2–3 months |
 | **1.5 — Nonlinear & amplified** ✅ | Adaptive-step SSFM, Kerr, EDFA with ASE, OSNR, PMD, APD, dispersion slope and its third-order term, cross-polarization Kerr coupling, inter-channel stimulated Raman scattering | ~2 months |
 | **2 — Coherent transceiver** ✅ | Gray-coded M-QAM to 256, IQ modulator with bias and quadrature error, 90° hybrid, balanced detection, blind carrier phase recovery, dual polarization with a blind butterfly equaliser, root-raised-cosine shaping and matched filtering, differential quadrant encoding, receiver-side dispersion compensation over spans to 1000 km with blind estimation of the accumulated value, EVM/MER, constellation diagram, validated against closed-form SER | ~3 months |
-| **3 — GUI & WDM** | ✅ Wavelength-selective filters, an OSA, coupled-channel propagation (XPM with walk-off, FWM accumulating coherently across spans), the session server, a schematic editor — add, wire, move and delete blocks, edit parameters, run, sweep, open and save — and 400G/800G reference designs validated against the OSNR relations · CuPy back-end | ~6 months |
+| **3 — GUI & WDM** | ✅ Wavelength-selective filters, an OSA, coupled-channel propagation (XPM with walk-off, FWM accumulating coherently across spans), the session server, a schematic editor — add, wire, move and delete blocks, edit parameters, run, sweep, open and save — and 400G/800G reference designs validated against the OSNR relations, and a back-end indirection the propagation kernels dispatch through — CuPy runs it where a device exists; it is not exercised in CI | ~6 months |
 | **4 — PIC** | Waveguides, ring resonators, MMI, MZI via integration with an existing S-matrix solver; PDK import | — |
 
 ¹ One developer, part-time. Estimates, not commitments.

docs/ARCHITECTURE.md

@@ -350,6 +350,15 @@ The original plan specified a C++20 core with pybind11 bindings. Reconsidered:
 GPU, and a native C++/Rust module later if profiling shows a specific kernel needs it. The
 architecture keeps that door open; it does not walk through it on day one.
 
+**Since built:** `maiman.backend` is the door, and it opens onto the propagation kernels only —
+the closed forms and the 2x2 Jones algebra are scalar work a device would slow down. Dispatch is on
+the array's own type rather than on a mode: a kernel handed CuPy arrays runs on CuPy, because the
+kernels are pure functions and a global mode would be the one piece of state that could make the
+same inputs give two answers. The surface a back-end must provide is thirteen names, asserted as an
+equality in `tests/test_backend.py` against a second array library that refuses NumPy's allocating
+API. CuPy itself is not exercised: there is no device in CI, and the part that is untested is
+CuPy's numerics rather than this interface.
+
 ### 6.2 FFT library and licensing — a real constraint
 
 **FFTW is GPL-2.0-or-later** (commercial licenses are sold separately). Linking it makes the

src/maiman/backend.py

@@ -0,0 +1,88 @@
+"""Which array library a kernel runs on, and how it finds out.
+
+The propagation kernels are the only part of this project a GPU would help: they
+are a loop over FFTs on a long array, and everything else is either scalar
+arithmetic or a closed form evaluated a few dozen times. :mod:`maiman.kernels`
+was written as array-to-array functions from the beginning so that this could be
+added without touching anything above it, and this module is the whole of the
+addition.
+
+**The arrays decide, not a setting.** A kernel handed CuPy arrays runs on CuPy
+and returns CuPy arrays; handed NumPy arrays it runs on NumPy. There is no global
+mode, no flag on the context, and nothing to get out of step between two calls —
+which matters because the kernels are pure functions and a hidden mode would be
+the one piece of state that could make the same inputs give different answers.
+It is also the idiom CuPy itself recommends.
+
+Dispatch is by the array's own type: ``type(a).__module__`` names the package it
+came from, and if that package is imported and looks like an array library it is
+the one to use. No registry, nothing to keep in sync, and a library nobody has
+heard of works if it presents the same surface.
+
+**What a back-end has to provide** is small, and
+``tests/test_backend.py`` holds it to exactly that list — a kernel that reaches
+past it for something only NumPy has fails there rather than on somebody's GPU.
+
+**CuPy is not exercised here.** No CUDA device or CuPy install is available in
+this project's CI, so what is tested is that the kernels never touch NumPy
+directly and that a second, deliberately hostile array library gets identical
+answers out of them. That is the part that would break a port; the part that
+remains untested is CuPy's own numerics, which are not this project's to test.
+"""
+
+from __future__ import annotations
+
+import sys
+from types import ModuleType
+
+import numpy as np
+
+#: Attributes a module must have before it is believed to be an array library.
+#: ``fft`` is the one every kernel here needs and no ordinary module has.
+_MARKERS = ("fft", "exp", "abs")
+
+
+def array_module(*arrays: object) -> ModuleType:
+    """The array library ``arrays`` belong to, defaulting to NumPy.
+
+    The first argument that names a plausible array library wins, so a kernel
+    mixing a device array with a plain Python float still runs on the device.
+    Scalars, lists and NumPy arrays all fall through to NumPy, which is what
+    makes this free for every existing caller.
+    """
+    for array in arrays:
+        root = type(array).__module__.partition(".")[0]
+        module = sys.modules.get(root)
+        if module is None or module is np:
+            continue
+        if all(hasattr(module, marker) for marker in _MARKERS):
+            return module
+    return np
+
+
+def to_numpy(array: object) -> np.ndarray:
+    """Bring a result back to the host, whatever it was computed on.
+
+    Measurement and encoding are host work — a constellation histogram is
+    kilobytes and a JSON payload is not going anywhere near a device — so this is
+    the boundary the rest of the project sees.
+    """
+    getter = getattr(array, "get", None)
+    if callable(getter):  # cupy.ndarray.get, and anything that copies it
+        return np.asarray(getter())
+    return np.asarray(array)
+
+
+def available() -> dict[str, bool]:
+    """Which back-ends this interpreter could actually use.
+
+    Reported rather than assumed: ``pip install cupy-cuda12x`` and a device is
+    what makes the second entry true, and neither is something this project can
+    check for by importing hopefully.
+    """
+    import importlib.util
+
+    return {
+        "numpy": True,
+        "cupy": importlib.util.find_spec("cupy") is not None,
+    }

src/maiman/kernels.py

@@ -19,16 +19,25 @@
 
 import numpy as np
 
+from .backend import array_module
 from .units import C_LIGHT
 
 
-def angular_frequency_grid(num_samples: int, sample_rate: float) -> np.ndarray:
+def angular_frequency_grid(
+    num_samples: int, sample_rate: float, *, like: object = None
+) -> np.ndarray:
     """Angular frequency offsets from the band centre [rad/s], in FFT order.
 
     Returned in `numpy.fft` output order (positive frequencies first, then
     negative), so it multiplies an un-shifted spectrum directly.
+
+    ``like`` names the array library to build it on — pass the field it is about
+    to multiply and the grid is created wherever that field lives, which is the
+    difference between one kernel and a device round trip per step. See
+    :mod:`maiman.backend`.
     """
-    return 2.0 * np.pi * np.fft.fftfreq(num_samples, d=1.0 / sample_rate)
+    xp = array_module(like)
+    return 2.0 * xp.pi * xp.fft.fftfreq(num_samples, d=1.0 / sample_rate)
 
 
 def dispersion_to_beta2(dispersion: float, wavelength: float) -> float:
@@ -108,12 +117,13 @@ def propagate_dispersion(
     the caller casts the result back. Correctness first; if profiling later shows
     this matters, the precision policy belongs here, in one place.
     """
+    xp = array_module(field)
     if distance == 0.0 or (beta2 == 0.0 and beta3 == 0.0):
-        return field.astype(np.complex128, copy=True)
+        return field.astype(xp.complex128, copy=True)
 
-    omega = angular_frequency_grid(field.shape[0], sample_rate)
-    transfer = np.exp(1j * (0.5 * beta2 * omega**2 - beta3 * omega**3 / 6.0) * distance)
-    return np.fft.ifft(np.fft.fft(field.astype(np.complex128)) * transfer)
+    omega = angular_frequency_grid(field.shape[0], sample_rate, like=field)
+    transfer = xp.exp(1j * (0.5 * beta2 * omega**2 - beta3 * omega**3 / 6.0) * distance)
+    return xp.fft.ifft(xp.fft.fft(field.astype(xp.complex128)) * transfer)
 
 
 def soliton_peak_power(beta2: float, gamma: float, width: float, order: int = 1) -> float:
@@ -405,7 +415,8 @@ def propagate_coupled_ssfm(
     if any(value not in (0, 1) for value in axis):
         raise ValueError(f"polarization entries must be 0 or 1, got {sorted(set(axis))}")
 
-    a = [f.astype(np.complex128, copy=True) for f in fields]
+    xp = array_module(*fields)
+    a = [f.astype(xp.complex128, copy=True) for f in fields]
     if not a:
         return a, PropagationDiagnostics(0, distance, 0.0, 0.0, 0.0)
 
@@ -421,7 +432,7 @@ def propagate_coupled_ssfm(
     if distance == 0.0:
         return a, PropagationDiagnostics(0, 0.0, 0.0, 0.0, 0.0)
 
-    omega = angular_frequency_grid(a[0].shape[0], sample_rate)
+    omega = angular_frequency_grid(a[0].shape[0], sample_rate, like=a[0])
     # One expansion of beta(omega), one operator: the group delay, the dispersion
     # and its slope are the first three terms of the same series, which is what
     # fixes their relative signs. See propagate_dispersion.
@@ -451,11 +462,11 @@ def propagate_coupled_ssfm(
             # always did.
             peak_effective = 0.0
             for field, ax in zip(a, axis, strict=True):
-                effective = 2.0 * _axis_power[ax] - np.abs(field) ** 2
+                effective = 2.0 * _axis_power[ax] - xp.abs(field) ** 2
                 other = _axis_power.get(1 - ax)
                 if other is not None:
                     effective = effective + ORTHOGONAL_KERR_WEIGHT * other
-                peak_effective = max(peak_effective, float(np.max(effective)))
+                peak_effective = max(peak_effective, float(xp.max(effective)))
             if peak_effective > 0.0:
                 step = min(step, max_nonlinear_phase / (abs(gamma) * peak_effective))
             # Only the nonlinear operator cares where the channels sit relative
@@ -467,8 +478,8 @@ def propagate_coupled_ssfm(
         # without this an extreme peak power would stall the loop.
         step = max(min(step, remaining), remaining * 1e-9)
 
-        half = [np.exp(-alpha * step / 4.0 + op * (step / 2.0)) for op in operators]
-        a = [np.fft.ifft(np.fft.fft(f) * h) for f, h in zip(a, half, strict=True)]
+        half = [xp.exp(-alpha * step / 4.0 + op * (step / 2.0)) for op in operators]
+        a = [xp.fft.ifft(xp.fft.fft(f) * h) for f, h in zip(a, half, strict=True)]
 
         if gamma != 0.0:
             # Summed per polarization, because power in the orthogonal component
@@ -478,16 +489,16 @@ def propagate_coupled_ssfm(
             per_axis = _power_per_axis(a, axis)
             rotated = []
             for field, ax in zip(a, axis, strict=True):
-                effective = 2.0 * per_axis[ax] - np.abs(field) ** 2
+                effective = 2.0 * per_axis[ax] - xp.abs(field) ** 2
                 other = per_axis.get(1 - ax)
                 if other is not None:
                     effective = effective + ORTHOGONAL_KERR_WEIGHT * other
                 phase = gamma * effective * step
-                peak_phase = max(peak_phase, float(np.max(np.abs(phase))))
-                rotated.append(field * np.exp(1j * phase))
+                peak_phase = max(peak_phase, float(xp.max(xp.abs(phase))))
+                rotated.append(field * xp.exp(1j * phase))
             a = rotated
 
-        a = [np.fft.ifft(np.fft.fft(f) * h) for f, h in zip(a, half, strict=True)]
+        a = [xp.fft.ifft(xp.fft.fft(f) * h) for f, h in zip(a, half, strict=True)]
 
         travelled += step
         steps += 1
@@ -500,7 +511,7 @@ def propagate_coupled_ssfm(
         # of the accumulated -1j*w*omega*distance, so it cannot remove anything
         # the propagation put there.
         a = [
-            np.fft.ifft(np.fft.fft(f) * np.exp(1j * w * omega * distance))
+            xp.fft.ifft(xp.fft.fft(f) * xp.exp(1j * w * omega * distance))
             for f, w in zip(a, walkoff, strict=True)
         ]
 
@@ -534,9 +545,10 @@ def _power_per_axis(fields: Sequence[np.ndarray], axis: Sequence[int]) -> dict[i
 
 def _total_power(fields: Sequence[np.ndarray]) -> np.ndarray:
     """Summed instantaneous power of co-propagating fields [W], sample by sample."""
-    total = np.abs(fields[0]) ** 2
+    xp = array_module(fields[0])
+    total = xp.abs(fields[0]) ** 2
     for field in fields[1:]:
-        total += np.abs(field) ** 2
+        total += xp.abs(field) ** 2
     return total
 
 
@@ -918,21 +930,26 @@ def apply_pmd(
     what the DGD measurement reports are the same chain.
     """
     if not sections:
-        return ex.astype(np.complex128, copy=True), ey.astype(np.complex128, copy=True)
+        module = array_module(ex, ey)
+        return (
+            ex.astype(module.complex128, copy=True),
+            ey.astype(module.complex128, copy=True),
+        )
 
-    omega = angular_frequency_grid(ex.shape[0], sample_rate)
-    spectrum_x = np.fft.fft(ex.astype(np.complex128))
-    spectrum_y = np.fft.fft(ey.astype(np.complex128))
+    xp = array_module(ex, ey)
+    omega = angular_frequency_grid(ex.shape[0], sample_rate, like=ex)
+    spectrum_x = xp.fft.fft(ex.astype(xp.complex128))
+    spectrum_y = xp.fft.fft(ey.astype(xp.complex128))
 
     for section in sections:
-        phase = np.exp(0.5j * omega * section.dgd)
+        phase = xp.exp(0.5j * omega * section.dgd)
         delayed_x = spectrum_x * phase
-        delayed_y = spectrum_y * np.conj(phase)
+        delayed_y = spectrum_y * xp.conj(phase)
         u = section.unitary
         spectrum_x = u[0, 0] * delayed_x + u[0, 1] * delayed_y
         spectrum_y = u[1, 0] * delayed_x + u[1, 1] * delayed_y
 
-    return np.fft.ifft(spectrum_x), np.fft.ifft(spectrum_y)
+    return xp.fft.ifft(spectrum_x), xp.fft.ifft(spectrum_y)
 
 
 def gaussian_lowpass_response(frequency: np.ndarray, bandwidth: float) -> np.ndarray:
@@ -943,7 +960,8 @@ def gaussian_lowpass_response(frequency: np.ndarray, bandwidth: float) -> np.nda
     """
     if bandwidth <= 0.0:
         raise ValueError(f"bandwidth must be positive, got {bandwidth}")
-    return np.exp(-0.5 * np.log(2.0) * (frequency / bandwidth) ** 2)
+    xp = array_module(frequency)
+    return xp.exp(-0.5 * math.log(2.0) * (frequency / bandwidth) ** 2)
 
 
 def super_gaussian_response(frequency: np.ndarray, bandwidth: float, order: int) -> np.ndarray:
@@ -964,7 +982,8 @@ def super_gaussian_response(frequency: np.ndarray, bandwidth: float, order: int)
         raise ValueError(f"bandwidth must be positive, got {bandwidth}")
     if order < 1:
         raise ValueError(f"order must be >= 1, got {order}")
-    return np.exp(-0.5 * np.log(2.0) * (2.0 * frequency / bandwidth) ** (2 * order))
+    xp = array_module(frequency)
+    return xp.exp(-0.5 * math.log(2.0) * (2.0 * frequency / bandwidth) ** (2 * order))
 
 
 def super_gaussian_noise_bandwidth(bandwidth: float, order: int) -> float:
@@ -1015,6 +1034,7 @@ def lowpass_filter(samples: np.ndarray, sample_rate: float, bandwidth: float) ->
     window are contaminated by the wrap; analysis blocks drop them.
     """
     n = samples.shape[0]
-    spectrum = np.fft.rfft(samples.astype(np.float64))
-    response = gaussian_lowpass_response(np.fft.rfftfreq(n, d=1.0 / sample_rate), bandwidth)
-    return np.fft.irfft(spectrum * response, n)
+    xp = array_module(samples)
+    spectrum = xp.fft.rfft(samples.astype(xp.float64))
+    frequency = xp.asarray(xp.fft.rfftfreq(n, d=1.0 / sample_rate))
+    return xp.fft.irfft(spectrum * gaussian_lowpass_response(frequency, bandwidth), n)