Scale S2 pattern map to correct S2 AFT (#196)

* set turned off pmts before scaling pattern map by aft

* rescale each space point

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* small change

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Diego Ramírez García <[email protected]>

fuse/context.py

@@ -242,18 +242,22 @@ def pattern_map(map_data, pmt_mask, method="WeightedNearestNeighbors"):
 
 
 @URLConfig.register("s2_aft_scaling")
-def modify_s2_pattern_map(s2_pattern_map, s2_mean_area_fraction_top, n_tpc_pmts, n_top_pmts):
+def modify_s2_pattern_map(
+    s2_pattern_map, s2_mean_area_fraction_top, n_tpc_pmts, n_top_pmts, turned_off_pmts
+):
     """Modify the S2 pattern map to match a given input AFT."""
     if s2_mean_area_fraction_top > 0:
         s2map = deepcopy(s2_pattern_map)
-        s2map_topeff_ = s2map.data["map"][..., 0:n_top_pmts].sum(axis=2)
-        s2map_toteff_ = s2map.data["map"].sum(axis=2)
-        orig_aft_ = np.mean((s2map_topeff_ / s2map_toteff_)[s2map_toteff_ > 0.0])
+        # First we need to set turned off pmts before scaling
+        s2map.data["map"][..., turned_off_pmts] = 0
+        s2map_topeff_ = s2map.data["map"][..., :n_top_pmts].sum(axis=2, keepdims=True)
+        s2map_toteff_ = s2map.data["map"].sum(axis=2, keepdims=True)
+        orig_aft_ = np.nanmean(s2map_topeff_ / s2map_toteff_)
         # Getting scales for top/bottom separately to preserve total efficiency
         scale_top_ = s2_mean_area_fraction_top / orig_aft_
         scale_bot_ = (1 - s2_mean_area_fraction_top) / (1 - orig_aft_)
-        s2map.data["map"][:, :, 0:n_top_pmts] *= scale_top_
-        s2map.data["map"][:, :, n_top_pmts:n_tpc_pmts] *= scale_bot_
+        s2map.data["map"][..., :n_top_pmts] *= scale_top_
+        s2map.data["map"][..., n_top_pmts:n_tpc_pmts] *= scale_bot_
         s2_pattern_map.__init__(s2map.data)
     return s2_pattern_map
 

fuse/plugins/detector_physics/s2_photon_propagation.py

@@ -211,7 +211,8 @@ class S2PhotonPropagationBase(FuseBaseDownChunkingPlugin):
         "&pmt_mask=plugin.pmt_mask"
         "&s2_mean_area_fraction_top=plugin.s2_mean_area_fraction_top"
         "&n_tpc_pmts=plugin.n_tpc_pmts"
-        "&n_top_pmts=plugin.n_top_pmts",
+        "&n_top_pmts=plugin.n_top_pmts"
+        "&turned_off_pmts=plugin.turned_off_pmts",
         cache=True,
         help="S2 pattern map",
     )
@@ -444,9 +445,6 @@ def photon_channels(self, n_electron, z_obs, positions, drift_time_mean, n_photo
                 pattern, [[0, 0], [0, len(bottom_index)]], "constant", constant_values=1
             )
 
-        # Remove turned off pmts
-        pattern[:, np.in1d(channels, self.turned_off_pmts)] = 0
-
         sum_pat = np.sum(pattern, axis=1).reshape(-1, 1)
         pattern = np.divide(pattern, sum_pat, out=np.zeros_like(pattern), where=sum_pat != 0)
 

fuse/plugins/detector_physics/secondary_scintillation.py

@@ -160,7 +160,8 @@ class SecondaryScintillation(FuseBasePlugin):
         "&pmt_mask=plugin.pmt_mask"
         "&s2_mean_area_fraction_top=plugin.s2_mean_area_fraction_top"
         "&n_tpc_pmts=plugin.n_tpc_pmts"
-        "&n_top_pmts=plugin.n_top_pmts",
+        "&n_top_pmts=plugin.n_top_pmts"
+        "&turned_off_pmts=plugin.turned_off_pmts",
         cache=True,
         help="S2 pattern map",
     )