Skip to content

Commit

Permalink
[evt.modules.larveto] use lambda function to pass parameters to ak tr…
Browse files Browse the repository at this point in the history
…ansform
  • Loading branch information
gipert committed May 8, 2024
1 parent cf43949 commit be05b5f
Showing 1 changed file with 10 additions and 10 deletions.
20 changes: 10 additions & 10 deletions src/pygama/evt/modules/larveto.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,8 +70,16 @@ def l200_test_stat(relative_t0, amp, ts_bkg_prob, rc_density):
n_pe_tot = np.where(n_pe_tot == 0, np.nan, n_pe_tot)

# calculate the test statistic term related to the time distribution
transform_function = transform_wrapper(ts_bkg_prob)
ts_time = -ak.sum(ak.transform(transform_function, relative_t0, amp), axis=-1)
ts_time = -ak.sum(
ak.transform(
lambda layouts, **kwargs: _ak_l200_test_stat_time_term(
layouts, ts_bkg_prob, **kwargs
),
relative_t0,
amp,
),
axis=-1,
)
# calculate the amplitude contribution
ts_amp = [l200_rc_amp_logpdf(n, rc_density) for n in n_pe_tot]

Expand All @@ -81,14 +89,6 @@ def l200_test_stat(relative_t0, amp, ts_bkg_prob, rc_density):
return t_stat


# need this to pass ts_bkg_prob parameter with ak.transform()
def transform_wrapper(ts_bkg_prob):
def _transform(layouts, **kwargs):
return _ak_l200_test_stat_time_term(layouts, ts_bkg_prob=ts_bkg_prob)

return _transform


# need to define this function and use it with ak.transform() because scipy
# routines are not NumPy universal functions
def _ak_l200_test_stat_time_term(layouts, ts_bkg_prob, **kwargs):
Expand Down

0 comments on commit be05b5f

Please sign in to comment.