Vectorise matrix extraction in MetisLMSSpectralTrace.get_matrices

[astronomyk]

Adding both @oczoske and @teutoburg as it's more a mechanics issue, than an instrument specific one.
Local tests produce identical results.

TL;DR - A speed boost for SpectralTraceList

From the testing notebook, removing these nested loops results in a 10x speed boost for extracting the polynomial coefficients

Not that impressive on a single run, but when we have to run this hundreds of times for METIS_Simulations, it all adds up.

import time
from astropy.io import fits
from scopesim.effects.metis_lms_trace_list import (MetisLMSSpectralTraceList,
                                                   echelle_setting)

cmd = scopesim.UserCommands(use_instrument="METIS", set_modes=["lms"])
lms = MetisLMSSpectralTraceList(
    filename="TRACE_LMS.fits", wavelen=4.2, cmds=cmd,
    wave_colname="wavelength", s_colname="xi", slice_samples=5)
trace = lms.spectral_traces["Slice 1"]
order, angle = trace.meta["order"], trace.meta["echelle"]
print(f"order {order}, echelle angle {angle:.3f} deg")

# Reference: the removed per-element implementation
def get_matrices_reference(trace):
    poly = trace.table
    order = trace.meta["order"]; spslice = trace.meta["slice"]
    angle = trace.meta["echelle"]
    out = {}
    for matid, name in enumerate(["A", "B", "AI", "BI"]):
        sel = ((poly["Ord"] == order) & (poly["Sli"] == spslice)
               & (poly["Mat"] == matid))
        sub = poly[sel]
        mat = np.zeros((4, 4))
        for i in range(4):
            for j in range(4):
                s = (sub["Row"] == i) & (sub["Col"] == j)
                mat[i, j] = (sub["P3"][s][0]*angle**3 + sub["P2"][s][0]*angle**2
                             + sub["P1"][s][0]*angle + sub["P0"][s][0])
        out[name] = mat
    return out

t0 = time.perf_counter()
ref = {sid: get_matrices_reference(spt) for sid, spt in lms.spectral_traces.items()}
t1 = time.perf_counter()
new = {sid: spt.get_matrices() for sid, spt in lms.spectral_traces.items()}
t2 = time.perf_counter()

worst = max(np.abs(new[sid][m] - ref[sid][m]).max()
            for sid in ref for m in ref[sid])
print(f"per-element loops (28 slices): {t1-t0:.3f} s")
print(f"vectorised       (28 slices): {t2-t1:.3f} s   ({(t1-t0)/(t2-t1):.0f}x)")
print(f"max |difference| over all 28 x 4 matrices: {worst}")
assert worst == 0.0

Resulting in

per-element loops (28 slices): 0.356 s
vectorised       (28 slices): 0.037 s   (10x)
max |difference| over all 28 x 4 matrices: 0.0

Toaster description below:

---

What

The A/B/AI/BI matrices are polynomial evaluations at the echelle angle, with coefficients selected from the Polynomial coefficients table (37,632 rows in TRACE_LMS.fits). The old loop ran one boolean mask over the full table per matrix, plus one mask per matrix element on the sub-table — (4 + 64) table scans for each of the 28 slices, on every trace-list construction.

The new code selects the (order, slice) block once, evaluates the angle polynomial on whole columns, and scatters the values into the 4x4 matrices via their Row/Col index columns. The resulting matrices are exactly identical (same additions in the same order — unit-tested element-by-element with zero tolerance); the KeyError for an unknown order/slice combination is kept.

Testing

• TestGetMatrices::test_matches_elementwise_reference — exact equality against the previous per-element evaluation for slices 1, 14 and 28 of the real coefficient table.
• TestGetMatrices::test_raises_for_unknown_order — KeyError behaviour preserved.
• A notebook with the equivalence check and timing on the full 28-slice TRACE_LMS is attached below.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 76.82%. Comparing base (beb191d) to head (af6f67a).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #988      +/-   ##
==========================================
+ Coverage   76.46%   76.82%   +0.36%     
==========================================
  Files          69       69              
  Lines        9025     9024       -1     
==========================================
+ Hits         6901     6933      +32     
+ Misses       2124     2091      -33     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.