Avoid copying logk for the normal systematic type

[davidwalter2]

Summary

For systematic_type == "normal", _init_logk_scaled absorbed the param-model factor rnorm_init into logk at construction time. That allocates a second full-size copy of the [nbins, nproc, nsyst] tensor and, in sparse mode, rebuilds the CSR matrix.

rnorm_init carries no systematic index, so it factors straight out of the contraction over systematics:

sum_s (rnorm_init[b,p] * logk[b,p,s]) * theta[s]
  == rnorm_init[b,p] * sum_s logk[b,p,s] * theta[s]

This keeps logk / logk_csr aliased to the indata tensors and applies the factor to the [nbins, nproc] contraction result instead of the [nbins, nproc, nsyst] input.

Only the normal systematic type is affected; log_normal already carried the scaling multiplicatively and takes the same early return as before.

Measured effect

The scaling now runs once per loss/HVP evaluation instead of once over the full tensor at construction, so it is worth checking that this does not cost more than it saves. Measured on a 10000-bin x 2048-systematic normal-type model (logk = 819 MB), with --minimizerMaxiter pinned so both versions do matched work:

	peak RSS	wall time	n_hvp
before	2115 +- 211 MB	178 +- 36 s	1100
after	1350 +- 18 MB	190 +- 95 s	1089

• Memory: 766 +- 122 MB saved (a 6.3 sigma effect), matching the 819 MB logk tensor -- i.e. the copy is gone.
• Runtime: neutral. The wall-time difference is 12 +- 59 s, consistent with zero. Call counts are deterministic and differ by 1% (1089 vs 1100), the expected small trajectory shift from the reassociation.

The per-call cost is one elementwise multiply on [nbins, nproc], a factor nsyst smaller than the matmul it follows, and it lands on data already in cache.

The saving scales with the tensor: for a 100k-bin x 8192-systematic model the avoided copy is ~33 GB. In sparse mode the win is structural rather than just an allocation -- the old path rebuilt the CSR matrix via CSRSparseMatrix(...), which is replaced by a one-time gather_nd onto the [norm_nnz] layout the CSR contraction already returns.

rnorm_init remains a constant evaluated at xparamdefault, so the loss stays strictly linear in theta and derivatives are unchanged.

Validation

Compared against unpatched main using tests/make_tensor.py, running the fit with both and comparing fitted parameter values and variances:

case	max relative difference
dense + normal	5.7e-12
sparse + normal	4.7e-11
dense + log_normal (control)	0 (bit-identical)

Run with --expectSignal sig 2.5 so that rnorm_init != 1 and the scaling is genuinely exercised -- with the default expectSignal=1 the factor is exactly 1.0 and both orderings agree trivially. The residual differences are consistent with the change in summation order (summing then scaling instead of scaling then summing); the new ordering performs one multiply per (bin, proc) rather than nsyst of them.

🤖 Generated with Claude Code