Perf: speed up LIKELY ramp fitting by avoiding unnecessary allocations

[jberg5]

This PR makes determine_diffs2use much faster; previously it was a substantial proportion of likely ramp fit runtime. The prior implementation took the full diffs cube (nresultants-1, rows, columns) and then allocated a full ones array in that shape, and then threw away all but the relevant row. Doing this full-cube allocation once per row was very expensive! determine_diffs2use only needs to operate on a single row anyway (interestingly the docstring seems to have been assuming the caller was passing just that row's diffs of shape (ngroups-1, ncols), but the implementation assumed the cube).

On a synthetic 1 x 5 x 4096 x 4096 Roman WFI-sized ramp, LIKELY fitting decreased from 34s to 15s on my M3 mac. Both runs report the same mean fitted rate. See a benchmark script below.

`benchmark.py`

"""Time one LIKELY fit on synthetic data with realistic Roman WFI dimensions."""
import time
import numpy as np
from stcal.ramp_fitting.likely_fit import likely_ramp_fit
from stcal.ramp_fitting.ramp_fit_class import RampData

SIZE = 4096
RESULTANTS = 5

def main():
    """Construct one detector ramp, fit it, and report wall-clock time."""
    shape = (1, RESULTANTS, SIZE, SIZE)
    levels = np.arange(RESULTANTS, dtype=np.float32)[None, :, None, None]
    data = np.broadcast_to(100 + 20 * levels, shape).copy()
    rng = np.random.default_rng(12345)
    for resultant in range(RESULTANTS):
        noise = rng.standard_normal((SIZE, SIZE), dtype=np.float32)
        noise *= 5
        data[0, resultant] += noise
    del noise
    data[0, RESULTANTS // 2 :, ::64, ::64] += 300  # Sparse cosmic-ray steps.

    ramp = RampData()
    ramp.set_arrays(
        data,
        np.zeros(shape, dtype=np.uint8),
        np.zeros((SIZE, SIZE), dtype=np.uint32),
        np.zeros((SIZE, SIZE), dtype=np.float32),
    )
    ramp.set_meta("WFI", frame_time=3.04, group_time=15.2, groupgap=1, nframes=4)
    ramp.flags_do_not_use = 1
    ramp.flags_saturated = 2
    ramp.flags_jump_det = 4
    readnoise = np.full((SIZE, SIZE), 10, dtype=np.float32)
    gain = np.full((SIZE, SIZE), 2, dtype=np.float32)

    start = time.perf_counter()
    result = likely_ramp_fit(ramp, readnoise, gain)[0]
    elapsed = time.perf_counter() - start

    print(f"Input: 1 integration x {RESULTANTS} resultants x {SIZE} x {SIZE} pixels")
    print(f"LIKELY wall clock: {elapsed:.3f} seconds")
    print(f"Mean fitted rate: {np.nanmean(result['slope']):.9f} DN/s")

if __name__ == "__main__":
    main()

Tasks

• update or add relevant tests
• update relevant docstrings and / or docs/ page
• Does this PR change any API used downstream? (if not, label with no-changelog-entry-needed)
  • write news fragment(s) in changes/: echo "changed something" > changes/<PR#>.<changetype>.rst (see changelog readme for instructions)
  • if your change breaks existing functionality, also add a changes/<PR#>.breaking.rst news fragment
• run regression tests with this branch installed ("git+https://github.com/<fork>/stcal@<branch>")
  • jwst regression test
  • romancal regression test

[zacharyburnett]

Hello, thanks for your contribution! Did you use any generative AI tools when making this PR? That includes if the @jberg5 account itself is an automated agent. Thank you!

[jberg5]

Hi @zacharyburnett ! I'm a human, but I use AI extensively, and did so here. I'm trying out GPT6, and I instructed it to profile an end-to-end fitting workflow to look for opportunities to speed things up.

It spotted this one, but the initial proposed change diff was to leave the diffs2use signature unchanged and do the relevant row extraction inside. I saw the stale docstring reference to per-row inputs, and decided to rewrite it into basically what you see here today, where the caller is responsible for selecting the relevant row.

Of course, GPT6 then reviewed the diff before I submitted it :) and I also had it write the benchmark script you see above, and I had it run that and similar benchmarking on an x86 box just in case the allocation performance was better/worse on my ARM macbook (it's not). I was working off more realistic inputs locally but needed roman installed and the right files downloaded so I wanted something synthetic and self-contained that still had the right shape.

[schlafly]

I started a romancal regtest run here:
https://github.com/spacetelescope/RegressionTests/actions/runs/34273652803
The current code is just buggy and is doing 4000x too much work in that loop!

@t-brandt , you were looking for likelihood ramp fitting performance improvements recently; here's a low-hanging one.

[t-brandt]

I agree this is an excellent change. Depending on the opinions of others, we could get the same answer with
d2use = (gdq[1:, :] == 0) & (gdq[:-1, :] == 0)
return d2use.astype(np.uint8)
I don't know whether this is clearer or less clear than what is currently implemented (i.e. use only differences where both groups are clean). We should definitely merge a version of this PR.

[jberg5]

thanks @t-brandt ! I like your suggestion. I think it's faster as well. Totally up to you guys - the version I have here was motivated by keeping the diff minimal, happy to go with whichever you prefer.