Correlation and Detection¶

import numpy as np

from doppler.spectral import Corr, Corr2D, CorrDetector, CorrDetector2D

N, LAG = 64, 17  # 1-D frame length and injected lag
NY, NX = 8, 8  # 2-D frame dimensions
ROW, COL = 3, 5  # 2-D shift (row, col)
SIGMA = 2.0  # noise amplitude; power = SIGMA² = 4  →  SNR ≈ −6 dB
DWELL = 8  # coherent integration depth
THRESHOLD = 5.0  # detection gate

rng = np.random.default_rng(42)

# Unit-magnitude BPSK PN references — fully deterministic.
ref1d = rng.choice(np.array([-1.0, 1.0], dtype=np.float32), size=N).astype(
    np.complex64
)
ref2d = rng.choice(
    np.array([-1.0, 1.0], dtype=np.float32), size=(NY, NX)
).astype(np.complex64)

noise_scale = np.float32(SIGMA / np.sqrt(2))


def noisy_frame() -> np.ndarray:
    """Return one N-sample CF32 frame: ref1d shifted by LAG, plus noise."""
    signal = np.roll(ref1d, LAG)
    noise = (rng.standard_normal(N) + 1j * rng.standard_normal(N)).astype(
        np.complex64
    ) * noise_scale
    return signal + noise


def noise_block(n_frames: int) -> np.ndarray:
    """Return n_frames*N samples of CF32 noise (no signal)."""
    total = N * n_frames
    return (
        rng.standard_normal(total) + 1j * rng.standard_normal(total)
    ).astype(np.complex64) * noise_scale



# Corr: dwell=1 vs dwell=8 — execute() accumulates frames and returns
# output only on the dwell-th call; all other calls return None.
with Corr(ref1d, dwell=1) as c:
    mag_d1 = np.abs(c.execute(noisy_frame()))

with Corr(ref1d, dwell=DWELL) as c:
    for _ in range(DWELL - 1):
        c.execute(noisy_frame())
    mag_d8 = np.abs(c.execute(noisy_frame()))

snr_d1 = mag_d1[LAG] / np.mean(np.delete(mag_d1, LAG))
snr_d8 = mag_d8[LAG] / np.mean(np.delete(mag_d8, LAG))
print(
    f"[Corr]      dwell=1  peak/mean={snr_d1:.1f}"
    f"  |  dwell={DWELL}  peak/mean={snr_d8:.1f}"
)

# The dwell=8 coherent sum must place the global peak exactly at the
# injected lag, and integrating 8 frames must sharpen the peak/mean
# ratio over a single frame (coherent integration gain).
assert int(np.argmax(mag_d8)) == LAG, "dwell=8 peak not at injected lag"
assert snr_d8 > snr_d1, "no coherent integration gain at dwell=8"

# One coherent-dwell block of the same shifted-PN + noise frames.
block = np.concatenate([noisy_frame() for _ in range(DWELL)])

det = CorrDetector(ref1d, dwell=DWELL, noise_lo=LAG + 4, noise_hi=N - 1,
                   threshold=THRESHOLD)
for lag, peak_mag, noise_est, test_stat in det.push(block):
    print(f"detection  lag={lag}  stat={test_stat:.2f}")
# detection  lag=17  stat=7.34   (stat varies with noise; lag is
#                                 deterministic)

# Corr2D: recover a (row, col) circular shift in one FFT2 call.
x2d = np.roll(np.roll(ref2d, ROW, axis=0), COL, axis=1)
with Corr2D(ref2d, dwell=1) as c:
    surf2d = np.abs(c.execute(x2d)).reshape(NY, NX)

peak_row, peak_col = np.unravel_index(surf2d.argmax(), (NY, NX))
print(
    f"[Corr2D]    peak at (row={peak_row}, col={peak_col})"
    f"  (expected ({ROW}, {COL}))"
)

# 2-D template match: the correlation surface must peak exactly at the
# injected (row, col) circular shift.
assert (peak_row, peak_col) == (ROW, COL), "2-D peak not at true shift"