Merge f0abe71 into 3f55f74

py/picca/delta_extraction/astronomical_objects/desi_pk1d_forest.py

@@ -197,22 +197,33 @@ def rebin(self):
         AstronomicalObjectError if Forest.wave_solution is not 'lin' or 'log'
         """
         bins, rebin_ivar, orig_ivar, w1, w2 = super().rebin()
-        if len(rebin_ivar) == 0:
+        if len(rebin_ivar) == 0 or np.sum(w2) == 0:
             self.resolution_matrix = np.array([[]])
-            return [], [], [], [], []
+            return [], [], [], [], np.array([])
 
         # apply mask due to cuts in bin
         self.resolution_matrix = self.resolution_matrix[:, w1]
 
+        # Find non-empty bins
+        binned_arr_size = bins.max() + 1
+        bincounts = np.bincount(bins, minlength=binned_arr_size)
+        wnonempty_bins = bincounts != 0
+
+        # Do a simple average when ivar=0
+        orig_ivar_2 = orig_ivar[w1]
+        w__ = orig_ivar_2>0
+        orig_ivar_2[~w__] = 1
+        rebin_reso_ivar = np.bincount(bins, weights=orig_ivar_2, minlength=binned_arr_size)
+
         # rebin resolution_matrix
         rebin_reso_matrix_aux = np.zeros(
-            (self.resolution_matrix.shape[0], bins.max() + 1))
+            (self.resolution_matrix.shape[0], binned_arr_size))
         for index, reso_matrix_col in enumerate(self.resolution_matrix):
             rebin_reso_matrix_aux[index, :] = np.bincount(
-                bins, weights=orig_ivar[w1] * reso_matrix_col)
+                bins, weights=orig_ivar_2 * reso_matrix_col)
         # apply mask due to rebinned inverse vairane
-        self.resolution_matrix = rebin_reso_matrix_aux[:, w2] / rebin_ivar[
-            np.newaxis, w2]
+        self.resolution_matrix = rebin_reso_matrix_aux[:, wnonempty_bins] / rebin_reso_ivar[
+            np.newaxis, wnonempty_bins]
 
         # return weights and binning solution to be used by child classes if
         # required

py/picca/delta_extraction/astronomical_objects/forest.py

@@ -14,7 +14,6 @@
     "mask fields": ["flux", "ivar", "transmission_correction", "log_lambda"],
 }
 
-
 @njit()
 def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
           log_lambda_grid, log_lambda_rest_frame_grid):
@@ -94,6 +93,7 @@ def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
     pixel_step = np.nan
 
     # compute bins
+    # Remove only out-of-bounds pixels
     if wave_solution == "log":
         pixel_step = log_lambda_grid[1] - log_lambda_grid[0]
         half_pixel_step = pixel_step / 2.
@@ -107,7 +107,6 @@ def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
                log_lambda_rest_frame_grid[0] - half_pixel_step_rest_frame)
         w1 &= (log_lambda - np.log10(1. + z) <
                log_lambda_rest_frame_grid[-1] + half_pixel_step_rest_frame)
-        w1 &= (ivar > 0.)
 
     elif wave_solution == "lin":
         pixel_step = 10**log_lambda_grid[1] - 10**log_lambda_grid[0]
@@ -122,73 +121,80 @@ def rebin(log_lambda, flux, ivar, transmission_correction, z, wave_solution,
                10**log_lambda_rest_frame_grid[0] - half_pixel_step_rest_frame)
         w1 &= (lambda_ / (1. + z) <
                10**log_lambda_rest_frame_grid[-1] + half_pixel_step_rest_frame)
-        w1 &= (ivar > 0.)
     else:
         raise AstronomicalObjectError("Error in Forest.rebin(). "
                                       "Class variable 'wave_solution' "
                                       "must be either 'lin' or 'log'.")
 
+    if (w1 & (ivar > 0.)).sum() == 0:
+        log_lambda = np.zeros(0)
+        flux = np.zeros(0)
+        ivar = np.zeros(0)
+        transmission_correction = np.zeros(0)
+        mean_snr = 0.0
+        bins = np.zeros(0, dtype=np.int64)
+        rebin_ivar = np.zeros(0)
+        w1 = np.zeros(0, dtype=bool_)
+        w2 = np.zeros(0, dtype=bool_)
+        return (log_lambda, flux, ivar, transmission_correction, mean_snr, bins,
+                rebin_ivar, orig_ivar, w1, w2)
+
     log_lambda = log_lambda[w1]
     flux = flux[w1]
     ivar = ivar[w1]
     transmission_correction = transmission_correction[w1]
-    if w1.sum() == 0:
-        log_lambda = np.zeros(log_lambda.size)
-        flux = np.zeros(log_lambda.size)
-        ivar = np.zeros(log_lambda.size)
-        transmission_correction = np.zeros(log_lambda.size)
-        mean_snr = 0.0
-        bins = np.zeros(log_lambda.size, dtype=np.int64)
-        rebin_ivar = np.zeros(log_lambda.size)
-        w1 = np.zeros(log_lambda.size, dtype=bool_)
-        w2 = np.zeros(log_lambda.size, dtype=bool_)
-        return (log_lambda, flux, ivar, transmission_correction, mean_snr, bins,
-                rebin_ivar, orig_ivar, w1, w2)
+    # Out-of-bounds pixels are removed. IVAR=0 pixels are kept
 
     bins = find_bins(log_lambda, log_lambda_grid, wave_solution)
+    binned_arr_size = bins.max() + 1
     log_lambda = log_lambda_grid[0] + bins * pixel_step
 
+    # Find non-empty bins. There will be empty bins
+    # at the lower end by construction.
+    bincounts = np.bincount(bins, minlength=binned_arr_size)
+    wnonempty_bins = bincounts != 0
+    final_arr_size = np.sum(wnonempty_bins)
+
     # rebin flux, ivar and transmission_correction
-    rebin_flux = np.zeros(bins.max() + 1)
-    rebin_transmission_correction = np.zeros(bins.max() + 1)
-    rebin_ivar = np.zeros(bins.max() + 1)
-    rebin_flux_aux = np.bincount(bins, weights=ivar * flux)
-    rebin_transmission_correction_aux = np.bincount(
-        bins, weights=(ivar * transmission_correction))
-    rebin_ivar_aux = np.bincount(bins, weights=ivar)
-    rebin_flux[:len(rebin_flux_aux)] += rebin_flux_aux
-    rebin_transmission_correction[:len(rebin_transmission_correction_aux
-                                      )] += rebin_transmission_correction_aux
-    rebin_ivar[:len(rebin_ivar_aux)] += rebin_ivar_aux
+    rebin_flux = np.bincount(bins, weights=ivar * flux, minlength=binned_arr_size)
+    rebin_transmission_correction = np.bincount(
+            bins, weights=(ivar * transmission_correction), minlength=binned_arr_size)
+    rebin_ivar = np.bincount(bins, weights=ivar, minlength=binned_arr_size)
 
     # this condition should always be non-zero for at least one pixel
     # this does not mean that all rebin_ivar pixels will be non-zero,
     # as we could have a masked region of the spectra
-    w2 = (rebin_ivar > 0.)
-    flux = rebin_flux[w2] / rebin_ivar[w2]
-    transmission_correction = rebin_transmission_correction[w2] / rebin_ivar[w2]
-    ivar = rebin_ivar[w2]
+    w2_ = (rebin_ivar > 0.) & wnonempty_bins
+    w2  = w2_[wnonempty_bins]
+    flux = np.zeros(final_arr_size)
+    transmission_correction = np.zeros(final_arr_size)
+    ivar = np.zeros(final_arr_size)
+
+    # Remove the empty pixels at the lower end
+    flux[w2] = rebin_flux[w2_] / rebin_ivar[w2_]
+    transmission_correction[w2] = rebin_transmission_correction[
+        w2_] / rebin_ivar[w2_]
+    ivar[w2] = rebin_ivar[w2_]
 
     # then rebin wavelength
     if wave_solution == "log":
         rebin_log_lambda = (log_lambda_grid[0] +
-                            np.arange(bins.max() + 1) * pixel_step)
-        log_lambda = rebin_log_lambda[w2]
+                            np.arange(binned_arr_size) * pixel_step)
+        log_lambda = rebin_log_lambda[wnonempty_bins]
     else:  # we have already checked that it will always be "lin" at this point
         rebin_lambda = (10**log_lambda_grid[0] +
-                        np.arange(bins.max() + 1) * pixel_step)
-        log_lambda = np.log10(rebin_lambda[w2])
+                        np.arange(binned_arr_size) * pixel_step)
+        log_lambda = np.log10(rebin_lambda[wnonempty_bins])
 
     # finally update control variables
     snr = flux * np.sqrt(ivar)
-    mean_snr = np.sum(snr) / float(snr.size)
+    mean_snr = np.mean(snr[w2])
 
     # return weights and binning solution to be used by child classes if
     # required
     return (log_lambda, flux, ivar, transmission_correction, mean_snr, bins,
             rebin_ivar, orig_ivar, w1, w2)
 
-
 class Forest(AstronomicalObject):
     """Forest Object
 

py/picca/delta_extraction/astronomical_objects/pk1d_forest.py

@@ -238,44 +238,52 @@ def rebin(self):
         AstronomicalObjectError if Forest.wave_solution is not 'lin' or 'log'
         """
         bins, rebin_ivar, orig_ivar, w1, w2 = super().rebin()
-        if len(rebin_ivar) == 0:
+        if len(rebin_ivar) == 0 or np.sum(w2) == 0:
             self.exposures_diff = np.array([])
             self.reso = np.array([])
             self.reso_pix = np.array([])
-            return [], [], [], [], []
+            return [], [], [], [], np.array([])
 
         # apply mask due to cuts in bin
         self.exposures_diff = self.exposures_diff[w1]
         self.reso = self.reso[w1]
         self.reso_pix = self.reso_pix[w1]
 
-        # rebin exposures_diff and reso
-        rebin_exposures_diff = np.zeros(bins.max() + 1)
-        rebin_reso = np.zeros(bins.max() + 1)
-        rebin_reso_pix = np.zeros(bins.max() + 1)
-        rebin_exposures_diff_aux = np.bincount(bins,
-                                               weights=orig_ivar[w1] *
-                                               self.exposures_diff)
-        rebin_reso_aux = np.bincount(bins, weights=orig_ivar[w1] * self.reso)
-        rebin_reso_pix_aux = np.bincount(bins, weights=orig_ivar[w1] * self.reso_pix)
-        rebin_exposures_diff[:len(rebin_exposures_diff_aux
-                                 )] += rebin_exposures_diff_aux
-        rebin_reso[:len(rebin_reso_aux)] += rebin_reso_aux
-        rebin_reso_pix[:len(rebin_reso_pix_aux)] += rebin_reso_pix_aux
+        # Find non-empty bins
+        binned_arr_size = bins.max() + 1
+        bincounts = np.bincount(bins, minlength=binned_arr_size)
+        wnonempty_bins = bincounts != 0
+        final_arr_size = np.sum(wnonempty_bins)
 
+        # rebin exposures_diff and reso
+        rebin_exposures_diff = np.bincount(bins,
+            weights=orig_ivar[w1] * self.exposures_diff, minlength=binned_arr_size)
+        rebin_reso = np.bincount(bins, weights=orig_ivar[w1] * self.reso, minlength=binned_arr_size)
+        rebin_reso_pix = np.bincount(bins, weights=orig_ivar[w1] * self.reso_pix,
+                                     minlength=binned_arr_size)
+
+        # Remove empty bins but not ivar
+        w2_ = (rebin_ivar > 0.) & wnonempty_bins
+        self.exposures_diff = np.zeros(final_arr_size)
+        self.reso = np.zeros(final_arr_size)
+        self.reso_pix = np.zeros(final_arr_size)
 
         # apply mask due to rebinned inverse vairane
-        self.exposures_diff = rebin_exposures_diff[w2] / rebin_ivar[w2]
-        self.reso = rebin_reso[w2] / rebin_ivar[w2]
-        self.reso_pix = rebin_reso_pix[w2] / rebin_ivar[w2]
+        self.exposures_diff[w2] = rebin_exposures_diff[w2_] / rebin_ivar[w2_]
+        self.reso[w2] = rebin_reso[w2_] / rebin_ivar[w2_]
+        self.reso_pix[w2] = rebin_reso_pix[w2_] / rebin_ivar[w2_]
 
         # finally update control variables
-        self.mean_reso = self.reso.mean()
+        self.mean_reso = self.reso[w2].mean()
         self.mean_z = (
             (np.power(10., self.log_lambda[len(self.log_lambda) - 1]) +
              np.power(10., self.log_lambda[0])) / 2. /
             Pk1dForest.lambda_abs_igm - 1.0)
-        self.mean_reso_pix = self.reso_pix.mean()
+        self.mean_reso_pix = self.reso_pix[w2].mean()
+
+        # maybe replace empty resolution values with the mean?
+        self.reso[~w2] = self.mean_reso
+        self.reso_pix[~w2] = self.mean_reso_pix
 
         # return weights and binning solution to be used by child classes if
         # required

py/picca/delta_extraction/config.py

@@ -34,7 +34,7 @@
     "overwrite", "logging level console", "logging level file", "log",
     "out dir", "num processors"
 ]
-accepted_masks_options = ["num masks", "type {int}", "module name {int}"]
+accepted_masks_options = ["num masks", "keep pixels", "type {int}", "module name {int}"]
 
 default_config = {
     "general": {

py/picca/delta_extraction/data.py

@@ -345,7 +345,7 @@ def filter_forests(self):
 
         remove_indexs = []
         for index, forest in enumerate(self.forests):
-            if forest.flux.size < self.min_num_pix:
+            if np.sum(forest.ivar>0) < self.min_num_pix:
                 # store information for logs
                 self.add_to_rejection_log(forest.get_header(), forest.flux.size,
                                           "short_forest")

py/picca/delta_extraction/expected_flux.py

@@ -203,29 +203,28 @@ def compute_delta_stack(self, forests, stack_from_deltas=False):
                 # ignore forest if continuum could not be computed
                 if forest.continuum is None:
                     continue
-                delta = forest.flux / forest.continuum
-                variance = self.compute_forest_variance(forest,
-                                                        forest.continuum)
-                weights = 1. / variance
+
+                delta = np.zeros_like(forest.log_lambda)
+                w = forest.ivar > 0
+                delta[w] = forest.flux[w] / forest.continuum[w]
+                weights = self.compute_forest_weight(forest, forest.continuum)
 
             bins = find_bins(forest.log_lambda, Forest.log_lambda_grid,
                              Forest.wave_solution)
-            rebin = np.bincount(bins, weights=delta * weights)
-            stack_delta[:len(rebin)] += rebin
-            rebin = np.bincount(bins, weights=weights)
-            stack_weight[:len(rebin)] += rebin
+            stack_delta += np.bincount(bins, weights=delta * weights, minlength=stack_delta.size)
+            stack_weight += np.bincount(bins, weights=weights, minlength=stack_delta.size)
 
         w = stack_weight > 0
         stack_delta[w] /= stack_weight[w]
 
         self.get_stack_delta = interp1d(
-            Forest.log_lambda_grid[stack_weight > 0.],
-            stack_delta[stack_weight > 0.],
+            Forest.log_lambda_grid[w],
+            stack_delta[w],
             kind="nearest",
             fill_value="extrapolate")
         self.get_stack_delta_weights = interp1d(
-            Forest.log_lambda_grid[stack_weight > 0.],
-            stack_weight[stack_weight > 0.],
+            Forest.log_lambda_grid[w],
+            stack_weight[w],
             kind="nearest",
             fill_value=0.0,
             bounds_error=False)
@@ -268,6 +267,24 @@ def compute_forest_variance(self, forest, continuum):  # pylint: disable=no-self
         raise ExpectedFluxError("Function 'compute_forest_variance' was not "
                                 "overloaded by child class")
 
+    def compute_forest_weight(self, forest, continuum):
+        """Compute the forest weights by setting ivar=0 to 0
+
+        Arguments
+        ---------
+        forest: Forest
+        A forest instance where the variance will be computed
+
+        continuum: array of float
+        Quasar continuum associated with the forest
+        """
+        weights = np.zeros_like(forest.log_lambda)
+        w = forest.ivar > 0
+        variance = self.compute_forest_variance(forest, continuum)
+        weights[w] = 1. / variance[w]
+
+        return weights
+
     def extract_deltas(self, forest):
         """Apply the continuum to compute the delta field