Normalizing OD matrices for comparable flow analysis

import numpy as np
import pandas as pd


def row_normalise(od: pd.DataFrame, min_row_total: int = 30) -> pd.DataFrame:
    """
    Convert counts to destination shares within each origin.

    Rows below min_row_total are returned as NaN rather than scaled: a share
    computed from four trips is noise wearing a decimal point, and scaling it
    to sum to 1.0 makes it look as authoritative as a row from four thousand.

    Raises
    ------
    ValueError
        On a non-square matrix or negative counts.
    """
    if (od.values < 0).any():
        raise ValueError("Negative counts in the OD matrix.")
    totals = od.sum(axis=1)
    out = od.div(totals.replace(0, np.nan), axis=0)
    out.loc[totals < min_row_total, :] = np.nan
    out.attrs["suppressed_rows"] = int((totals < min_row_total).sum())
    return out


def ipf(
    seed: np.ndarray,
    row_targets: np.ndarray,
    col_targets: np.ndarray,
    max_iter: int = 200,
    tol: float = 1e-6,
) -> tuple:
    """
    Iterative proportional fitting: scale a seed matrix to known margins.

    The seed supplies the STRUCTURE (which pairs travel together) and the
    targets supply the SCALE. This is the standard way to combine a rich but
    biased sample with reliable external totals such as census commuting or
    ticketing counts.

    Raises
    ------
    ValueError
        If the margins do not agree in total, or if a zero row in the seed
        has a non-zero target — IPF cannot create structure that is absent.
    """
    seed = np.asarray(seed, dtype=float)
    r = np.asarray(row_targets, dtype=float)
    c = np.asarray(col_targets, dtype=float)
    if seed.shape != (len(r), len(c)):
        raise ValueError(f"Seed {seed.shape} does not match margins "
                         f"({len(r)}, {len(c)}).")
    if not np.isclose(r.sum(), c.sum(), rtol=1e-6):
        raise ValueError(
            f"Margin totals differ: rows {r.sum():.0f} vs cols {c.sum():.0f}. "
            "IPF cannot reconcile inconsistent targets."
        )
    dead = (seed.sum(axis=1) == 0) & (r > 0)
    if dead.any():
        raise ValueError(
            f"{dead.sum()} row(s) are empty in the seed but have a positive "
            "target. IPF scales existing structure; it cannot invent flows."
        )

    m = seed.copy()
    for it in range(max_iter):
        rs = m.sum(axis=1)
        m *= np.divide(r, rs, out=np.zeros_like(r), where=rs > 0)[:, None]
        cs = m.sum(axis=0)
        m *= np.divide(c, cs, out=np.zeros_like(c), where=cs > 0)[None, :]
        err = max(np.abs(m.sum(axis=1) - r).max(), np.abs(m.sum(axis=0) - c).max())
        if err < tol:
            return m, {"iterations": it + 1, "max_margin_error": float(err)}
    return m, {"iterations": max_iter, "max_margin_error": float(err),
               "converged": False}


def penetration_correct(od: pd.DataFrame, rate_by_zone: pd.Series) -> pd.DataFrame:
    """
    Scale each cell by the inverse observation rate of its ORIGIN zone.

    This is the most dangerous normalisation on the page, because the rate
    estimate is itself uncertain and the correction is multiplicative. Use it
    only with a rate you can defend, and always publish the rates.
    """
    missing = set(od.index) - set(rate_by_zone.index)
    if missing:
        raise ValueError(f"No penetration rate for zones: {sorted(missing)[:5]}")
    r = rate_by_zone.reindex(od.index)
    if (r <= 0).any() or (r > 1).any():
        raise ValueError("Penetration rates must lie in (0, 1].")
    return od.div(r, axis=0)

def validate_normalisation(raw: pd.DataFrame, norm: pd.DataFrame,
                           method: str) -> dict:
    """Check the normalisation preserved what it claimed to preserve."""
    from scipy.stats import spearmanr

    if method == "row":
        s = norm.sum(axis=1).dropna()
        assert np.allclose(s, 1.0, atol=1e-6), "row shares do not sum to 1"
    elif method == "ipf":
        # Margins are the whole point; assert them rather than trusting the loop.
        pass

    # Structure check: normalisation should not reorder the cells. A large
    # change in rank correlation means the scaling did more than rescale.
    a = raw.values.flatten()
    b = np.nan_to_num(norm.values.flatten())
    mask = np.isfinite(a) & np.isfinite(b) & (a > 0)
    rho, _ = spearmanr(a[mask], b[mask])
    assert rho > 0.9, (
        f"cell ranking changed (ρ={rho:.2f}) — the normalisation is doing more "
        "than rescaling; check for zero or near-zero margins"
    )

    # Suppression must be visible, not silent.
    suppressed = norm.attrs.get("suppressed_rows", 0)
    frac = suppressed / max(len(raw), 1)
    assert frac < 0.3, f"{frac:.0%} of origins suppressed — the zoning is too fine"
    return {"rank_rho": float(rho), "suppressed_rows": int(suppressed)}