Allow standardizing a DenseMatrix by shifting the data

[abelianbee]

Closes #414.

standardize keeps shift = -mean/std and mult = 1/std on the StandardizedMatrix, so sandwich expands the product into four terms, one of them outer(shift, shift) * sum(d). That is a centered second moment computed from uncentered sums, so when a column's mean is large relative to its standard deviation those terms are orders of magnitude larger than their sum and the significant digits cancel away. No regrouping of the algebra avoids it. Centering the data does.

This adds an opt-in materialize_shift flag to standardize. When set, the shift and multiplier are folded into a copy of the data and the returned matrix carries a zero shift, so the expansion never happens. The default is unchanged. _materialize_standardization returns None by default, so SparseMatrix and CategoricalMatrix ignore the flag and are never densified. StandardizedMatrix holds a reference to the pre-standardization matrix, so unstandardize() stays exact and free, which glum relies on in _glm.py.

Accuracy against a longdouble oracle, and timing end-to-end including the standardize call, medians of 5, M1 Max:

n	p	dtype	rel err, expansion	rel err, materialized	sandwich per call	standardize + 10 sandwich
5,000	1,000	float64	9.1e-09	3.2e-15	1.41x	1.36x
1,000,000	100	float64	1.2e-07	4.3e-15	1.97x	1.44x
200,000	50	float64	3.9e-02	4.4e-13	1.90x	1.42x
5,000	1,000	float32	6.9e+00	2.1e-06	1.20x	1.17x

Materializing pays for a copy in standardize, so a single sandwich call is not a win on tall matrices. It breaks even after 2-3 calls and is 1.2-1.4x by 10, since each call skips three p x p outer products and an extra transpose_matvec in place of one BLAS-3 call on contiguous data. The accuracy gain does not depend on the number of calls.

Tests cover representation equivalence across all seven matrix fixtures with both flag values, operation equivalence for sandwich / matvec / transpose_matvec, and an accuracy regression against a longdouble reference. 5666 passed, 45 skipped, 4 xpassed.

I moved the unreleased changelog section from 4.2.2 to 4.3.0, since this adds functionality. Happy to put it back if you would rather keep the number.

Checklist

• Added a CHANGELOG.rst entry

[Marc-Antoine Schmidt (MarcAntoineSchmidtQC)]

Thanks for your contribution!

I tested it on a few examples and the benefits are real, especially in terms of convergence for previously ill-conditioned problems.

I see at least one problem with the state of this PR. There are a few places in the code where a method returns a StandardizedMatrix (e.g.__getitem__, astype, getcol). Each of those output will be incorrect because the concept of the cached unstandardized is not passed along.

Two other things:

1. We should make it clear in the documentation that this doubles the memory requirements (because the copy needs to be stored).
2. Can you share your benchmarking code and confirm that the speedup includes the increased time it take for the original call to standardize.

[abelianbee]

Thanks a lot for taking the time to test this, and for the careful read.

Good catch on __getitem__, getcol and astype. All three build a new StandardizedMatrix without passing the unstandardized matrix along, so unstandardize() on the result hands back the shifted copy. Fixed by threading it through in each and added a test per method.

While in astype I noticed it also drops mult, which is independent of this PR: on main, S.astype(np.float64).toarray() differs from S.toarray() whenever scale_predictors=True. Fixed it in a separate commit since it's the same line. Happy to split it out if you'd prefer.

1. Extended the note in the standardize docstring. It already mentioned the copy but not that the original is kept for unstandardize(), so roughly twice the memory of the design matrix.

2. Fair question. No, the numbers in the description timed sandwich alone. Redone end-to-end, medians of 5, M1 Max:

n	p	dtype	standardize exp / mat	one sandwich exp / mat	std + 1	breakeven	std + 10
5,000	1,000	float64	3.3 / 11.9 ms	40.0 / 28.4 ms	1.07x	1	1.36x
1,000,000	100	float64	196 / 372 ms	126 / 64 ms	0.74x	3	1.44x
200,000	50	float64	27 / 42 ms	13.2 / 6.9 ms	0.81x	3	1.42x
5,000	1,000	float32	2.8 / 6.3 ms	19.8 / 16.4 ms	0.99x	2	1.17x

Materializing costs 2-4x more in standardize because of the copy, so a single sandwich is a wash or a bit slower on the tall cases. It breaks even at 2-3 sandwich calls and is 1.2-1.4x at 10. I've replaced the table in the description with this one. Script below.

The wheel build failure is docker failing to pull the manylinux image from quay.io, unrelated to the change.

Let me know if you'd like anything else changed.

benchmark script

"""End-to-end timing for tabmat#414: standardize() + sandwich(), expansion vs materialized.

Answers the review question of whether the speedup includes the cost of the
standardize() call itself. Medians of 5 trials. Apple M1 Max.
"""
import time
import numpy as np
import tabmat as tm


def med(f, trials=5):
    ts = []
    for _ in range(trials):
        t0 = time.perf_counter()
        f()
        ts.append(time.perf_counter() - t0)
    return float(np.median(ts))


rng = np.random.default_rng(0)
# (n, p, dtype, column offset). Same shapes as the table in the PR description.
cases = [
    (5_000, 1_000, np.float64, 1e3),
    (1_000_000, 100, np.float64, 1e3),
    (200_000, 50, np.float64, 1e6),
    (5_000, 1_000, np.float32, 1e3),
]

hdr = f"{'n':>9} {'p':>5} {'dtype':>7} | {'std exp':>8} {'std mat':>8} | {'sand exp':>8} {'sand mat':>8} {'per call':>8} | {'std+1':>6} {'brkeven':>7} {'std+10':>6}"
print(hdr)
print("-" * len(hdr))
for n, p, dt, off in cases:
    X = (rng.standard_normal((n, p)) + off).astype(dt)
    w = np.full(n, 1.0 / n, dtype=dt)
    d = rng.uniform(0.5, 1.5, n).astype(dt)
    M = tm.DenseMatrix(X)

    def std_exp():
        return M.standardize(w, center_predictors=True, scale_predictors=True)[0]

    def std_mat():
        return M.standardize(
            w, center_predictors=True, scale_predictors=True, materialize_shift=True
        )[0]

    Se, Sm = std_exp(), std_mat()
    t_se, t_sm = med(std_exp), med(std_mat)
    t_ae, t_am = med(lambda: Se.sandwich(d)), med(lambda: Sm.sandwich(d))

    per_call = t_ae / t_am
    one = (t_se + t_ae) / (t_sm + t_am)
    ten = (t_se + 10 * t_ae) / (t_sm + 10 * t_am)
    extra, save = t_sm - t_se, t_ae - t_am
    be = 1 if extra <= 0 else ("inf" if save <= 0 else int(np.ceil(extra / save)))
    print(
        f"{n:>9,} {p:>5} {np.dtype(dt).name:>7} | {t_se*1e3:7.1f}ms {t_sm*1e3:7.1f}ms | "
        f"{t_ae*1e3:7.1f}ms {t_am*1e3:7.1f}ms {per_call:7.2f}x | {one:5.2f}x {str(be):>7} {ten:5.2f}x"
    )
print()
print("std = one standardize() call, sand = one sandwich() call.")
print("per call = sandwich-only speedup. std+1 / std+10 = standardize plus that many sandwich calls.")
print("brkeven = sandwich calls needed before materializing is a net win.")