From b0f7f55481b58047d7dac3ecc61a27efa1d4d0c2 Mon Sep 17 00:00:00 2001
From: bpalmeiro <brais.pp@gmail.com>
Date: Thu, 7 Mar 2024 12:43:41 +0100
Subject: [PATCH] Add selection_in_band function

---
 invisible_cities/reco/icaro_components.py | 50 ++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/invisible_cities/reco/icaro_components.py b/invisible_cities/reco/icaro_components.py
index c7e75e052..8345b1d38 100644
--- a/invisible_cities/reco/icaro_components.py
+++ b/invisible_cities/reco/icaro_components.py
@@ -1,7 +1,8 @@
 import numpy  as np
 import pandas as pd
 
-from   typing       import Tuple, Optional
+from   typing               import Tuple, Optional
+from   sklearn.linear_model import RANSACRegressor
 
 from .  corrections         import ASectorMap
 from .  corrections         import apply_geo_correction
@@ -9,9 +10,12 @@
 from .. types.symbols       import type_of_signal
 from .. types.symbols       import Strictness
 from .. types.symbols       import NormStrategy
+from .. types.symbols       import KrFitFunction
 from .. core.core_functions import check_if_values_in_interval
 from .. core.core_functions import in_range
+from .. core.fit_functions  import profileX
 from .. core.fit_functions  import fit
+from .  krmap_functions     import get_function_and_seed_lt
 
 
 def selection_nS_mask_and_checking(dst        : pd.DataFrame                         ,
@@ -134,4 +138,48 @@ def band_selector_and_check(dst         : pd.DataFrame,
     return sel_krband
 
 
+def selection_in_band(z         : np.array,
+                      e         : np.array,
+                      range_z   : Tuple[float, float],
+                      range_e   : Tuple[float, float],
+                      nsigma    : float   = 3.5) ->np.array:
+    """
+    This function returns a mask for the selection of the events that are inside the Kr E vz Z
+
+    Parameters
+    ----------
+    z: np.array
+        axial (z) values
+    e: np.array
+        energy values
+    range_z: Tuple[np.array, np.array]
+        Range in Z-axis
+    range_e: Tuple[np.array, np.array]
+        Range in Energy-axis
+    nsigma: float
+        Number of sigmas to set the band width
+    Returns
+    ----------
+        A  mask corresponding to the selection made.
+    """
+    # Reshapes and flattens are needed for RANSAC function
+
+    z_sel = z[in_range(z, *range_z)]
+    e_sel = e[in_range(e, *range_e)]
+
+    res_fit      = RANSACRegressor().fit(z_sel, np.log(e_sel).reshape(-1, 1))
+
+    in_mask      = res_fit.inlier_mask_
+    residuals_ln = e_sel[in_mask] - np.exp(res_fit.predict(z_sel[in_mask].reshape(-1, 1))).flatten()
+    resy, resx   = np.histogram(residuals_ln, 100)
+    resx         = resx[:-1] + np.diff(resx)
+    fitres       = fit(gauss, resx, resy, seed=[4e3,0,10])
+    fitsigma     = fitres.val[2]
+
+    prefict_fun  = lambda z: res_fit.predict(z.reshape(-1, 1)).flatten()
+    upper_band   = lambda z: prefict_fun(z) + nsigma * fitsigma
+    lower_band   = lambda z: prefict_fun(z) - nsigma * fitsigma
+
+    sel_inband   = in_range(np.log(e_sel), lower_band(z_sel), upper_band(z_sel))
 
+    return  sel_inband