cache rate index to avoid lookup in form_matrix loop

[shimwell]

Description

A tiny performance improvement that caches the rates.index_nuc and rates.index_rx so that we can avoid looking them up within the loop

I have committed as @yrrepy as this is his idea, I am just facilitating by splitting up a otherwise large piece of work.

I think this small speed improvement makes sense on its own

Fixes # (issue)

Checklist

• I have performed a self-review of my own code
• I have followed the style guidelines for Python source files (if applicable)
• I have made corresponding changes to the documentation (if applicable)
• I have added tests that prove my fix is effective or that my feature works (if applicable)

[GuySten]

I think this degrade readability a bit. I will be surprised if this will improve performance in a non negligible way. Do you have an example where this change affect performance significantly for the better?

[shimwell]

Yes good point, on my computer the form_matrix takes 5.3s on this branch and 5.6s on develop. With a bigger chain (perhaps a tendl chain) I guess it would take longer. But perhaps 0.3s is still worth it.

import time
import urllib.request
from pathlib import Path

import numpy as np
from openmc.deplete import Chain, reaction_rates

# Download the 27MB ENDF/B-VIII.0 SFR chain
CHAIN_URL = "https://raw.githubusercontent.com/openmc-data-storage/openmc_data/refs/heads/main/src/openmc_data/depletion/chain_endf_b8.0_sfr.xml"
CHAIN_FILE = Path("chain_endf_b8.0_sfr.xml")

print(f"Downloading chain from {CHAIN_URL} ...")
urllib.request.urlretrieve(CHAIN_URL, CHAIN_FILE)
print("Download complete.")

chain = Chain.from_xml(CHAIN_FILE)

# Create reaction rates with random values
nuc_names = [nuc.name for nuc in chain.nuclides]
mats = ["1"]
rates = reaction_rates.ReactionRates(mats, nuc_names, chain.reactions)

rng = np.random.default_rng(42)
for nuc in chain.nuclides:
    for r_type, target, _, br in nuc.reactions:
        # Random rate in [0, 10) reactions/sec
        rates.set("1", nuc.name, r_type, rng.random() * 10.0)

# Time form_matrix
# Warm-up call
_ = chain.form_matrix(rates[0])

n_trials = 100
t0 = time.perf_counter()
for _ in range(n_trials):
    matrix = chain.form_matrix(rates[0])
t_total = time.perf_counter() - t0

print(f"form_matrix timing ({n_trials} calls):")
print(f"  Total:   {t_total:.3f} s")
print(f"  Average: {t_total / n_trials * 1000:.3f} ms")

[GuySten]

IMO a 5% speedup does not justify the readability cost.

[eepeterson]

Thanks for the PR @shimwell and @yrrepy - I think to address the readability concern you can move the _index_nuc and _index_rx variable assignments to just before the loop where they are used rather than where you currently have them. sorry I skimmed the code too quickly. Either way I think 5% speed up for ~ 5 relatively localized line changes seems totally fine to me. index_nuc and index_rx are unambiguous in the context of this function regardless of whether they are attached to the rates variable at the point of calling particularly because rates is documented in the function docstring as a np.ndarray when in fact it is a subclass of np.ndarray with additional attributes added to it.