Add data preparation functions for the Kr map computation

invisible_cities/core/core_functions.py

@@ -3,13 +3,19 @@
 This module includes utility functions.
 """
 import time
+import warnings
 
 import numpy as np
 
 from typing import Sequence
 from typing import Tuple
+from typing import Optional
 
+
+from . exceptions import ValueOutOfRange
 from ..types.symbols import NormMode
+from ..types.symbols import Strictness
+
 
 def timefunc(f):
     """
@@ -77,6 +83,61 @@ def in_range(data, minval=-np.inf, maxval=np.inf, left_closed=True, right_closed
     return lower_bound & upper_bound
 
 
+def all_in_range(data         : np.ndarray                                   ,
+                 minval       : float                                        ,
+                 maxval       : float                                        ,
+                 display_name : Optional[str] = ''                           ,
+                 strictness   : Optional[Strictness] = Strictness.raise_error,
+                 **kwargs)->bool:
+    """
+    Checks whether input values are all inside the interval (minval, maxval).
+
+    Parameters
+    ----------
+    data : np.ndarray
+        Input array to check.
+    minval: float
+        Lower limit of the interval.
+    maxval: float
+        Upper limit of the interval.
+    display_name: string
+        Label to be displayed in case of warning or exception.
+    strictness: Strictness
+        Describes the behaviour when the output is `False`. If `strictness` is:
+        - `Strictness.raise_error`: an exception is raised.
+        - `Strictness.warning`: a warning is raised before returning.
+        - `Strictness.silent`: the function returns quietly.
+
+    **kwargs:
+        Optional arguments being passed to `in_range`.
+    Returns
+    -------
+    True if values are in the interval. False otherwise.
+
+    Raises
+    ------
+    ValueOutOfRange: if the output is `False` and `strictness` is `Strictness.raise_error`
+    """
+
+    values_in_interval = in_range(data, minval, maxval, **kwargs)
+    outliers   = data[~values_in_interval]
+    n_outliers = len(outliers)
+
+    if values_in_interval.all():
+        return True
+    elif strictness is Strictness.silent:
+        return False
+
+    text  = f'Variable {display_name} has {n_outliers} values out of bounds ({minval}, {maxval}\n'
+    text += f'Outliers: {outliers}'
+
+    if strictness is Strictness.warning:
+        warnings.warn(text, UserWarning)
+        return False
+    else:
+        raise ValueOutOfRange(text)
+
+
 def weighted_mean_and_var(data       : Sequence,
                           weights    : Sequence,
                           unbiased   : bool = False,

invisible_cities/core/core_functions_test.py

@@ -9,6 +9,8 @@
 from pytest import approx
 from pytest import mark
 from pytest import raises
+from pytest import warns
+
 
 from flaky                  import flaky
 from hypothesis             import given
@@ -111,6 +113,38 @@ def test_in_range_right_close_interval(data):
     assert all(output)
 
 
+@given(random_length_float_arrays(min_length = 1,
+                                  min_value  = 0,
+                                  max_value  = 100))
+def test_all_in_range_when_all_fall_inside(data):
+    assert core.all_in_range(data, 0, 100, left_closed=True,right_closed=True,
+                             strictness=core.Strictness.raise_error)
+
+
+@given(random_length_float_arrays(min_length = 1,
+                                  min_value  = 0,
+                                  max_value  = 100))
+def test_all_in_range_when_some_fall_outside_silent(data):
+    minvalue = np.min(data)
+    assert not core.all_in_range(data, minvalue+1, 100, strictness=core.Strictness.silent)
+
+@given(random_length_float_arrays(min_length = 1,
+                                  min_value  = 0,
+                                  max_value  = 100))
+def test_all_in_range_when_some_fall_outside_warns(data):
+    minvalue = np.min(data)
+    with warns(UserWarning, match="values out of bounds"):
+        core.all_in_range(data, minvalue+1, 100, strictness=core.Strictness.warning)
+
+@given(random_length_float_arrays(min_length = 1,
+                                  min_value  = 0,
+                                  max_value  = 100))
+def test_all_in_range_when_some_fall_outside_exception(data):
+    minvalue = np.min(data)
+    with raises(core.ValueOutOfRange, match="values out of bounds"):
+        core.all_in_range(data, minvalue+1, 100, strictness=core.Strictness.raise_error)
+
+
 @mark.parametrize(" first  second       norm_mode        expected".split(),
                   ((  1  ,    2  , core.NormMode.first ,   -1    ),
                    (  4  ,    2  , core.NormMode.second,    1    ),

invisible_cities/core/exceptions.py

@@ -55,3 +55,6 @@ class MCEventNotFound(ICException):
 
 class SensorIDMismatch(ICException):
     pass
+
+class ValueOutOfRange(ICException):
+    pass

invisible_cities/reco/corrections.py

@@ -243,6 +243,40 @@ def get_normalization_factor(map_e0    : ASectorMap,
 
     return norm_value
 
+def apply_geo_correction(map_e0     : ASectorMap,
+                         norm_strat : NormStrategy    = NormStrategy.max,
+                         norm_value : Optional[float] = None
+                        ) -> Callable:
+    """
+    Given a map and a correction strategy, it returns a
+    function that provides a geometric only correction factor
+    for a given hit collection when (x,y) is provided.
+
+    Parameters
+    ----------
+    map_e0 : AsectorMap
+        Correction map for geometric effects.
+    norm_strat : NormStrategy
+        Provides the desired normalization to be used.
+    norm_value : Float(optional)
+        If norm_strat is selected to be custom, user must provide the
+        desired scale.
+    Returns
+    -------
+        A function that returns geometric correction factor without passing a map.
+    """
+    normalization   = get_normalization_factor(map_e0, norm_strat, norm_value)
+    get_xy_corr_fun = maps_coefficient_getter(map_e0.mapinfo, map_e0.e0)
+
+    def geo_correction_factor(x : np.ndarray,
+                              y : np.ndarray)-> np.ndarray:
+        geo_factor = correct_geometry_(get_xy_corr_fun(x,y))
+        factor     = geo_factor * normalization
+        return factor
+
+    return geo_correction_factor
+
+
 
 def apply_all_correction_single_maps(map_e0         : ASectorMap,
                                      map_lt         : ASectorMap,

invisible_cities/reco/corrections_test.py

@@ -11,6 +11,7 @@
 from . corrections import get_normalization_factor
 from . corrections import apply_all_correction_single_maps
 from . corrections import apply_all_correction
+from . corrections import apply_geo_correction
 
 from pytest                import fixture
 from numpy.testing         import assert_allclose
@@ -268,21 +269,30 @@ def test_apply_all_correction_single_maps_raises_exception_when_invalid_map(map_
 
 @few_examples
 @given(float_arrays(size      = 1,
-                   min_value = -198,
-                   max_value = +198),
-       float_arrays(size      = 1,
-                   min_value = -198,
-                   max_value = +198),
-       float_arrays(size      = 1,
-                   min_value = 0,
-                   max_value = 5e2),
+                    min_value = -198,
+                    max_value = +198),
        float_arrays(size      = 1,
-                   min_value = 0,
-                   max_value = 1e5))
+                    min_value = -198,
+                    max_value = +198))
+def test_apply_geo_correction_properly(map_filename, x, y):
+    """
+    The input map is homogeneous, therefore the geometric
+    correction factor must be 1.
+    """
+    maps      = read_maps(map_filename)
+    load_corr = apply_geo_correction(map_e0=maps, norm_strat=NormStrategy.max)
+    corr      = load_corr(x, y)
+    assert_allclose (corr, np.ones_like(corr))
+
+@few_examples
+@given(float_arrays(size      = 1, min_value = -198, max_value =      198),
+       float_arrays(size      = 1, min_value = -198, max_value =      198),
+       float_arrays(size      = 1, min_value =    0, max_value =      500),
+       float_arrays(size      = 1, min_value =    0, max_value = 100*1000))
 def test_apply_all_correction_single_maps_properly(map_filename, x, y, z, t):
     """
-    Due the map taken as input, the geometric correction
-    factor must be 1, the temporal correction 1 and the
+    The input map is homogeneous, therefore the geometric
+    correction factor must be 1, the temporal correction 1 and the
     lifetime one: exp(Z/5000).
     """
     maps      = read_maps(map_filename)