CoxPH: accumulate the information matrix per event time instead of building per-observation (n, p, p) arrays

[derrynknife]

Split out of #329, which removed the implementation-level costs and left this one because it is a change of algorithm.

The remaining cost

After #329 the Cox fit is 3x to 10x faster and beats lifelines nearly everywhere. Two cases still do not:

case	surpyval	lifelines
n=50 000, p=10, continuous	8.3s	4.0s
n=20 000, p=5, heavy ties	0.33s	0.08s

Both have the same cause. jac_hess forms the second-moment sums by building a per-observation outer product and aggregating it:

Z2 = np.einsum("ij, ik -> ijk", Z, Z)          # (n, p, p), hoisted in #329
z2_e_beta_z = Z2 * (n * e_beta_z)[:, :, None]  # (n, p, p), every iteration
Z2Ri = gb_x.sum(z2_e_beta_z)[1]
Z2Ri = Z2Ri[::-1].cumsum(axis=0)[::-1]
Z2Ri = Z2Ri - not_yet_entered(pos, gb_tl.sum(z2_e_beta_z)[1])

At n=50 000 with ten covariates each (n, p, p) array is 40MB, and several are built, permuted, reduced and reverse-cumulated per root-finding iteration. The profile after #329 is flat across exactly those operations — no single hot spot left, just memory traffic:

ncalls  tottime  function
   150    3.036  _GroupBy.sum
   105    2.882  ndarray.cumsum
    33    2.860  c_einsum
    16    2.472  jac_hess (body)
    16    2.470  efron_hess

The flops are O(n p²) either way. What differs is that the current shape materialises that as arrays and traverses them repeatedly, where an incremental accumulation touches O(p²) of working memory and stays in cache.

The change

Walk event times in reverse, maintaining running R, ZR and Z2R accumulators, adding each observation's contribution as it enters the risk set. That is what lifelines does, and it is why it wins at large n and p while losing at small.

Complications that make this real work rather than a transcription:

• Delayed entry. The current code subtracts a not-yet-entered suffix sum (not_yet_entered, CRITICAL: Cox delayed-entry / start-stop TVC fits have corrupted gradients and Hessians #250) which is an exact gather valid for signed quantities. An incremental walk has to add and remove subjects at the right boundaries instead, under the strict (entry, exit] convention. Getting this wrong is how CRITICAL: Cox delayed-entry / start-stop TVC fits have corrupted gradients and Hessians #250 happened in the first place.
• Count weights. n is a weight and may be fractional; the Efron tie weights truncate (range(int(d))) while the divisor does not. See _efron_tie_weights.
• Stratification. _fit_stratified sums per-stratum generators; the accumulator has to reset at stratum boundaries.
• Numerical drift. An incremental accumulation over 50 000 additions is not the same summation order as a reverse cumsum, so the small differences currently seen against lifelines (1e-06 to 1e-12) may move. The existing agreement figures are the yardstick.

Guard rails already in place

#329 added tests that this work must keep passing:

• test_efron_hessian_matches_the_loop_it_replaced and test_efron_log_denominator_matches_the_loop_it_replaced — oracles running the original Python double loop over no ties, heavy ties, times with no deaths, and fractional count weights.
• test_cox_fit_is_unaffected_by_the_order_of_the_rows — both tie methods, with delayed entry, checking beta and the Hessian.
• test_efron_hessian_matches_finite_difference, test_count_weights_equivalent_to_repeated, test_cox_delayed_entry_episode_split_invariance.

Worth checking first

Whether this is worth doing at all depends on how often the p is large and n is large. At p=2 and p=5 surpyval already wins across every n measured, by up to 5x. The regression is confined to wide covariate matrices at scale, and to heavy ties. A quick look at what shapes people actually fit would settle whether this is a priority or a nicety.

Timings above are single runs and drift 10–20% between them; the lifelines column moves about as much. The 2x and 4x gaps are outside that.