From 8c4da76b691bc5f48c285c3621aaa0a0dbe80bda Mon Sep 17 00:00:00 2001 From: Maximilian Linhoff Date: Mon, 6 Jan 2025 17:32:22 +0100 Subject: [PATCH 1/2] Cache fourier sin/cos computation --- src/ctapipe_io_lst/calibration.py | 50 +++++++++++++++++++++++++------ 1 file changed, 41 insertions(+), 9 deletions(-) diff --git a/src/ctapipe_io_lst/calibration.py b/src/ctapipe_io_lst/calibration.py index 16fe3c04..66cab445 100644 --- a/src/ctapipe_io_lst/calibration.py +++ b/src/ctapipe_io_lst/calibration.py @@ -3,6 +3,8 @@ import numpy as np import astropy.units as u from numba import njit +from numba.typed import Dict +from numba.core import types import tables from ctapipe.core import TelescopeComponent @@ -1102,6 +1104,20 @@ def ped_time(timediff): + + +_fourier_cache_key_type = types.UniTuple(types.int16, 2) +_fourier_cache_value_type = types.UniTuple(types.float64[:], 2) + + +@njit(cache=True) +def _create_fourier_cache(): + return Dict.empty( + key_type=_fourier_cache_key_type, + value_type=_fourier_cache_value_type, + ) + + @njit(cache=True) def calc_drs4_time_correction_gain_selected( first_capacitors, selected_gain_channel, fan, fbn @@ -1109,11 +1125,15 @@ def calc_drs4_time_correction_gain_selected( _n_gains, n_pixels, n_harmonics = fan.shape time = np.zeros(n_pixels) + cache = _create_fourier_cache() for pixel in range(n_pixels): gain = selected_gain_channel[pixel] first_capacitor = first_capacitors[gain, pixel] time[pixel] = calc_fourier_time_correction( - first_capacitor, fan[gain, pixel], fbn[gain, pixel] + first_capacitor, + fan[gain, pixel], + fbn[gain, pixel], + cache, ) return time @@ -1124,28 +1144,40 @@ def calc_drs4_time_correction_both_gains( ): time = np.zeros((N_GAINS, N_PIXELS)) + cache = _create_fourier_cache() for gain in range(N_GAINS): for pixel in range(N_PIXELS): first_capacitor = first_capacitors[gain, pixel] time[gain, pixel] = calc_fourier_time_correction( - first_capacitor, fan[gain, pixel], fbn[gain, pixel] + first_capacitor, + fan[gain, pixel], + fbn[gain, pixel], + cache, ) return time + @njit(cache=True) -def calc_fourier_time_correction(first_capacitor, fan, fbn): +def calc_fourier_time_correction(first_capacitor, fan, fbn, cache): n_harmonics = len(fan) time = 0 first_capacitor = first_capacitor % N_CAPACITORS_CHANNEL - for harmonic in range(1, n_harmonics): - a = fan[harmonic] - b = fbn[harmonic] - omega = harmonic * (2 * np.pi / N_CAPACITORS_CHANNEL) + cache_key = (types.int16(n_harmonics), types.int16(first_capacitor)) + cache_val = cache.get(cache_key) + if cache_val is not None: + sin_terms, cos_terms = cache_val + else: + n = np.arange(n_harmonics) + omega = 2 * np.pi / N_CAPACITORS_CHANNEL + sin_terms = np.sin(n * omega * first_capacitor) + cos_terms = np.cos(n * omega * first_capacitor) + cache[cache_key] = sin_terms, cos_terms - time += a * np.cos(omega * first_capacitor) - time += b * np.sin(omega * first_capacitor) + for harmonic in range(1, n_harmonics): + time += fan[harmonic] * cos_terms[harmonic] + time += fbn[harmonic] * sin_terms[harmonic] return time From 4a24199ab8d196045f79d130de160b15f0d48e0d Mon Sep 17 00:00:00 2001 From: Maximilian Linhoff Date: Wed, 8 Jan 2025 10:21:39 +0100 Subject: [PATCH 2/2] Simplify cache key --- src/ctapipe_io_lst/calibration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ctapipe_io_lst/calibration.py b/src/ctapipe_io_lst/calibration.py index 66cab445..107c6c91 100644 --- a/src/ctapipe_io_lst/calibration.py +++ b/src/ctapipe_io_lst/calibration.py @@ -1106,7 +1106,7 @@ def ped_time(timediff): -_fourier_cache_key_type = types.UniTuple(types.int16, 2) +_fourier_cache_key_type = types.int16 _fourier_cache_value_type = types.UniTuple(types.float64[:], 2) @@ -1165,7 +1165,7 @@ def calc_fourier_time_correction(first_capacitor, fan, fbn, cache): time = 0 first_capacitor = first_capacitor % N_CAPACITORS_CHANNEL - cache_key = (types.int16(n_harmonics), types.int16(first_capacitor)) + cache_key = types.int16(first_capacitor) cache_val = cache.get(cache_key) if cache_val is not None: sin_terms, cos_terms = cache_val