Merge 80570bf into cf28f72

py/picca/delta_extraction/astronomical_objects/forest.py

@@ -8,7 +8,7 @@
 
 from picca.delta_extraction.astronomical_object import AstronomicalObject
 from picca.delta_extraction.errors import AstronomicalObjectError
-from picca.delta_extraction.utils import find_bins
+from picca.delta_extraction.utils import find_bins_lin, find_bins_log
 
 defaults = {
     "mask fields": [
@@ -58,7 +58,7 @@ def get_inner_region_slice(bincounts):
 
 @njit()
 def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
-          log_lambda_grid, log_lambda_rest_frame_grid):
+          log_lambda_grid, log_lambda_rest_frame_grid, find_bins):
     """Rebin the arrays and update control variables
     Rebinned arrays are flux, ivar, lambda_ or log_lambda, and
     transmission_correction. Control variables are mean_snr
@@ -140,29 +140,24 @@ def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
         pixel_step = log_lambda_grid[1] - log_lambda_grid[0]
         half_pixel_step = pixel_step / 2.
 
-        half_pixel_step_rest_frame = (log_lambda_rest_frame_grid[1] -
-                                      log_lambda_rest_frame_grid[0]) / 2.
-
         w1 &= log_lambda >= log_lambda_grid[0] - half_pixel_step
         w1 &= log_lambda < log_lambda_grid[-1] + half_pixel_step
         w1 &= (log_lambda - np.log10(1. + z) >=
-               log_lambda_rest_frame_grid[0] - half_pixel_step_rest_frame)
+               log_lambda_rest_frame_grid[0])
         w1 &= (log_lambda - np.log10(1. + z) <
-               log_lambda_rest_frame_grid[-1] + half_pixel_step_rest_frame)
+               log_lambda_rest_frame_grid[-1])
 
     elif wave_solution == "lin":
         pixel_step = 10**log_lambda_grid[1] - 10**log_lambda_grid[0]
         half_pixel_step = pixel_step / 2.
 
-        half_pixel_step_rest_frame = (10**log_lambda_rest_frame_grid[1] -
-                                      10**log_lambda_rest_frame_grid[0]) / 2.
         lambda_ = 10**log_lambda
         w1 &= (lambda_ >= 10**log_lambda_grid[0] - half_pixel_step)
         w1 &= (lambda_ < 10**log_lambda_grid[-1] + half_pixel_step)
         w1 &= (lambda_ / (1. + z) >=
-               10**log_lambda_rest_frame_grid[0] - half_pixel_step_rest_frame)
+               10**log_lambda_rest_frame_grid[0])
         w1 &= (lambda_ / (1. + z) <
-               10**log_lambda_rest_frame_grid[-1] + half_pixel_step_rest_frame)
+               10**log_lambda_rest_frame_grid[-1])
     else:
         raise AstronomicalObjectError("Error in Forest.rebin(). "
                                       "Class variable 'wave_solution' "
@@ -187,7 +182,8 @@ def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
     transmission_correction = transmission_correction[w1]
     # Out-of-bounds pixels are removed. IVAR=0 pixels are kept
 
-    bins = find_bins(log_lambda, log_lambda_grid, wave_solution)
+
+    bins = find_bins(log_lambda, log_lambda_grid)
     binned_arr_size = bins.max() + 1
 
     # Find non-empty bins. There will be empty bins
@@ -224,7 +220,7 @@ def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
     else:  # we have already checked that it will always be "lin" at this point
         log_lambda = np.log10(10**log_lambda_grid[0] + pixel_step *
                     np.arange(binned_arr_size)[wslice_inner])
-    rebin_bins = find_bins(log_lambda, log_lambda_grid, wave_solution)
+    rebin_bins = find_bins(log_lambda, log_lambda_grid)
 
     # finally update control variables
     snr = flux * np.sqrt(ivar)
@@ -321,6 +317,7 @@ class Forest(AstronomicalObject):
     Weights associated to the delta field. None for no information
     """
     blinding = "none"
+    find_bins = None # this will store the correct find_bins function
     flux_units = None
     log_lambda_grid = np.array([])  #None
     log_lambda_rest_frame_grid = np.array([])  #None
@@ -677,7 +674,8 @@ def rebin(self):
          self.mean_snr, bins, self.log_lambda_index, rebin_ivar, orig_ivar, w1,
          w2, wslice_inner) = rebin(self.log_lambda, self.flux, self.ivar,
                      self.transmission_correction, self.z, Forest.wave_solution,
-                     Forest.log_lambda_grid, Forest.log_lambda_rest_frame_grid)
+                     Forest.log_lambda_grid, Forest.log_lambda_rest_frame_grid,
+                     Forest.find_bins)
 
         # return weights and binning solution to be used by child classes if
         # required
@@ -720,14 +718,18 @@ def set_class_variables(cls, lambda_min, lambda_max, lambda_min_rest_frame,
                 np.log10(lambda_min),
                 np.log10(lambda_max) + pixel_step / 2, pixel_step)
             cls.log_lambda_rest_frame_grid = np.arange(
-                np.log10(lambda_min_rest_frame) + pixel_step_rest_frame / 2,
-                np.log10(lambda_max_rest_frame), pixel_step_rest_frame)
+                np.log10(lambda_min_rest_frame),
+                np.log10(lambda_max_rest_frame) + pixel_step_rest_frame / 2,
+                pixel_step_rest_frame)
+            cls.find_bins = find_bins_log
         elif wave_solution == "lin":
             cls.log_lambda_grid = np.log10(
                 np.arange(lambda_min, lambda_max + pixel_step / 2, pixel_step))
             cls.log_lambda_rest_frame_grid = np.log10(
-                np.arange(lambda_min_rest_frame + pixel_step_rest_frame / 2,
-                          lambda_max_rest_frame, pixel_step_rest_frame))
+                np.arange(lambda_min_rest_frame,
+                          lambda_max_rest_frame + pixel_step_rest_frame / 2,
+                          pixel_step_rest_frame))
+            cls.find_bins = find_bins_lin
         else:
             raise AstronomicalObjectError("Error in setting Forest class "
                                           "variables. 'wave_solution' "

py/picca/delta_extraction/expected_flux.py

@@ -11,7 +11,6 @@
 from picca.delta_extraction.astronomical_objects.forest import Forest
 from picca.delta_extraction.astronomical_objects.pk1d_forest import Pk1dForest
 from picca.delta_extraction.errors import ExpectedFluxError, AstronomicalObjectError
-from picca.delta_extraction.utils import find_bins
 
 accepted_options = [
     "iter out prefix", "num bins variance", "num processors", "out dir",
@@ -126,33 +125,21 @@ def _initialize_variance_wavelength_array(self):
         The initialized arrays are:
         - self.log_lambda_var_func_grid
         """
-        # initialize the variance-related variables (see equation 4 of
-        # du Mas des Bourboux et al. 2020 for details on these variables)
         if Forest.wave_solution == "log":
-            self.log_lambda_var_func_grid = (
-                Forest.log_lambda_grid[0] +
-                (np.arange(self.num_bins_variance) + .5) *
-                (Forest.log_lambda_grid[-1] - Forest.log_lambda_grid[0]) /
+            self.log_lambda_var_func_grid = np.linspace(
+                Forest.log_lambda_grid[0],
+                Forest.log_lambda_grid[-1],
                 self.num_bins_variance)
-        # TODO: this is related with the todo in check the effect of finding
-        # the nearest bin in log_lambda space versus lambda space infunction
-        # find_bins in utils.py. Once we understand that we can remove
-        # the dependence from Forest from here too.
         elif Forest.wave_solution == "lin":
-            self.log_lambda_var_func_grid = np.log10(
-                10**Forest.log_lambda_grid[0] +
-                (np.arange(self.num_bins_variance) + .5) *
-                (10**Forest.log_lambda_grid[-1] -
-                 10**Forest.log_lambda_grid[0]) / self.num_bins_variance)
-
-        # TODO: Replace the if/else block above by something like the commented
-        # block below. We need to check the impact of doing this on the final
-        # deltas first (eta, var_lss and fudge will be differently sampled).
-        #start of commented block
-        #resize = len(Forest.log_lambda_grid)/self.num_bins_variance
-        #print(resize)
-        #self.log_lambda_var_func_grid = Forest.log_lambda_grid[::int(resize)]
-        #end of commented block
+            self.log_lambda_var_func_grid = np.log10(np.linspace(
+                10**Forest.log_lambda_grid[0],
+                10**Forest.log_lambda_grid[-1],
+                self.num_bins_variance))
+        else:
+            raise ExpectedFluxError("Error in setting Forest class "
+                                          "variables. 'wave_solution' "
+                                          "must be either 'lin' or 'log'. "
+                                          f"Found: {Forest.wave_solution}")
 
     def __parse_config(self, config):
         """Parse the configuration options
@@ -228,8 +215,8 @@ def compute_delta_stack(self, forests, stack_from_deltas=False):
                 delta[w] = forest.flux[w] / forest.continuum[w]
                 weights = self.compute_forest_weights(forest, forest.continuum)
 
-            bins = find_bins(forest.log_lambda, Forest.log_lambda_grid,
-                             Forest.wave_solution)
+            bins = Forest.find_bins( # pylint: disable=not-callable
+                forest.log_lambda, Forest.log_lambda_grid)
             stack_delta += np.bincount(bins, weights=delta * weights, minlength=stack_delta.size)
             stack_weight += np.bincount(bins, weights=weights, minlength=stack_delta.size)
 
@@ -289,15 +276,20 @@ def _compute_mean_cont(self, forests,
         for forest in forests:
             if forest.bad_continuum_reason is not None:
                 continue
-            bins = find_bins(forest.log_lambda - np.log10(1 + forest.z),
-                             Forest.log_lambda_rest_frame_grid,
-                             Forest.wave_solution)
+            bins = Forest.find_bins( # pylint: disable=not-callable
+                forest.log_lambda - np.log10(1 + forest.z),
+                Forest.log_lambda_rest_frame_grid)
 
             weights = self.compute_forest_weights(forest, forest.continuum)
             forest_continuum = which_cont(forest)
-            mean_cont += np.bincount(bins, weights=forest_continuum * weights,
+            mean_cont += np.bincount(
+                bins,
+                weights=forest_continuum * weights,
+                minlength=mean_cont.size)
+            mean_cont_weight += np.bincount(
+                bins,
+                weights=weights,
                 minlength=mean_cont.size)
-            mean_cont_weight += np.bincount(bins, weights=weights, minlength=mean_cont.size)
 
         w = mean_cont_weight > 0
         mean_cont[w] /= mean_cont_weight[w]

py/picca/delta_extraction/expected_fluxes/true_continuum.py

@@ -12,7 +12,7 @@
 from picca.delta_extraction.astronomical_objects.pk1d_forest import Pk1dForest
 from picca.delta_extraction.errors import ExpectedFluxError
 from picca.delta_extraction.expected_flux import ExpectedFlux, defaults, accepted_options
-from picca.delta_extraction.utils import (find_bins, update_accepted_options,
+from picca.delta_extraction.utils import (update_accepted_options,
                                           update_default_options)
 
 accepted_options = update_accepted_options(accepted_options, [
@@ -514,9 +514,9 @@ def compute_var_lss(self, forests):
 
         for forest in forests:
             w = forest.ivar > 0
-            log_lambda_bins = find_bins(forest.log_lambda,
-                                        Forest.log_lambda_grid,
-                                        Forest.wave_solution)[w]
+            log_lambda_bins = Forest.find_bins( # pylint: disable=not-callable
+                forest.log_lambda,
+                Forest.log_lambda_grid)[w]
             var_pipe = 1. / forest.ivar[w] / forest.continuum[w]**2
             deltas = forest.flux[w] / forest.continuum[w] - 1
             var_lss[log_lambda_bins] += deltas**2 - var_pipe

py/picca/delta_extraction/least_squares/least_squares_var_stats.py

@@ -2,7 +2,6 @@
 import numpy as np
 
 from picca.delta_extraction.astronomical_objects.forest import Forest
-from picca.delta_extraction.utils import find_bins
 
 VAR_PIPE_MIN = np.log10(1e-5)
 VAR_PIPE_MAX = np.log10(2.)
@@ -161,14 +160,14 @@ def initialize_delta_arrays(self, forests):
                  (np.log10(var_pipe) < VAR_PIPE_MAX))
 
             # select the pipeline variance bins
-            var_pipe_bins = np.floor(
+            var_pipe_bins = (
                 (np.log10(var_pipe[w]) - VAR_PIPE_MIN) /
                 (VAR_PIPE_MAX - VAR_PIPE_MIN) * NUM_VAR_BINS).astype(int)
 
             # select the wavelength bins
-            log_lambda_bins = find_bins(forest.log_lambda[w],
-                                        self.log_lambda_var_func_grid,
-                                        Forest.wave_solution)
+            log_lambda_bins = Forest.find_bins( # pylint: disable=not-callable
+                forest.log_lambda[w],
+                self.log_lambda_var_func_grid)
 
             # compute overall bin
             bins = var_pipe_bins + NUM_VAR_BINS * log_lambda_bins

py/picca/delta_extraction/utils.py

@@ -134,11 +134,12 @@ def class_from_string(class_name, module_name):
         accepted_options = []
     return class_object, default_args, accepted_options
 
-
 @njit()
-def find_bins(original_array, grid_array, wave_solution):
+def find_bins_lin(original_array, grid_array):
     """For each element in original_array, find the corresponding bin in grid_array
 
+    This function assumes that wavelength grid that is evenly spaced on wavelength
+
     Arguments
     ---------
     original_array: array of float
@@ -147,37 +148,40 @@ def find_bins(original_array, grid_array, wave_solution):
     grid_array: array of float
     Common array, e.g. Forest.log_lambda_grid
 
-    wave_solution: "log" or "lin"
-    Specifies whether we want to construct a wavelength grid that is evenly
-    spaced on wavelength (lin) or on the logarithm of the wavelength (log)
+    Return
+    ------
+    found_bin: array of int
+    An array of size original_array.size filled with values smaller than
+    grid_array.size with the bins correspondance
+    """
+    aux = 10**grid_array[0]
+    step = 10**grid_array[1] - aux
+    found_bin = ((10**original_array - aux) / step + 0.5).astype(np.int64)
+    return found_bin
+
+@njit()
+def find_bins_log(original_array, grid_array):
+    """For each element in original_array, find the corresponding bin in grid_array
+
+    This function assumes that wavelength grid that is evenly spaced on the
+    logarithm of the wavelength
+
+    Arguments
+    ---------
+    original_array: array of float
+    Read array, e.g. forest.log_lambda
+
+    grid_array: array of float
+    Common array, e.g. Forest.log_lambda_grid
 
     Return
     ------
     found_bin: array of int
     An array of size original_array.size filled with values smaller than
     grid_array.size with the bins correspondance
     """
-    if wave_solution == "log":
-        pass
-    elif wave_solution == "lin":
-        original_array = 10**original_array
-        grid_array = 10**grid_array
-    else:  # pragma: no cover
-        raise DeltaExtractionError(
-            "Error in function find_bins from py/picca/delta_extraction/utils.py"
-            "expected wavelength solution to be either 'log' or 'lin'. ")
-    original_array_size = original_array.size
-    grid_array_size = grid_array.size
-    found_bin = np.zeros(original_array_size, dtype=np.int64)
-    for index1 in range(original_array_size):
-        min_dist = np.finfo(np.float64).max
-        for index2 in range(grid_array_size):
-            dist = np.abs(grid_array[index2] - original_array[index1])
-            if dist < min_dist:
-                min_dist = dist
-                found_bin[index1] = index2
-            else:
-                break
+    step = grid_array[1] - grid_array[0]
+    found_bin = ((original_array - grid_array[0]) / step + 0.5).astype(np.int64)
     return found_bin
 
 PROGRESS_LEVEL_NUM = 15