Re-using fill_and_smooth for QDM

sup3r/bias/bias_calc.py

@@ -721,7 +721,91 @@ def _reduce_base_data(base_ti, base_data, base_cs_ghi, base_dset,
  return base_data, daily_ti
 
 
-class LinearCorrection(DataRetrievalBase):
+class FillAndSmoothMixin():
+ """Fill and extend parameters for calibration on missing positions"""
+ def fill_and_smooth(self,
+ out,
+ fill_extend=True,
+ smooth_extend=0,
+ smooth_interior=0):
+ """For a given set of parameters, fill and extend missing positions
+
+ Fill data extending beyond the base meta data extent by doing a
+ nearest neighbor gap fill. Smooth interior and extended region with
+ given smoothing values.
+ Interior smoothing can reduce the affect of extreme values
+ within aggregations over large number of pixels.
+ The interior is assumed to be defined by the region without nan values.
+ The extended region is assumed to be the region with nan values.
+
+ Parameters
+ ----------
+ out : dict
+ Dictionary of values defining the mean/std of the bias + base
+ data and the scalar + adder factors to correct the biased data
+ like: bias_data * scalar + adder. Each value is of shape
+ (lat, lon, time).
+ fill_extend : bool
+ Whether to fill data extending beyond the base meta data with
+ nearest neighbor values.
+ smooth_extend : float
+ Option to smooth the scalar/adder data outside of the spatial
+ domain set by the threshold input. This alleviates the weird seams
+ far from the domain of interest. This value is the standard
+ deviation for the gaussian_filter kernel
+ smooth_interior : float
+ Value to use to smooth the scalar/adder data inside of the spatial
+ domain set by the threshold input. This can reduce the effect of
+ extreme values within aggregations over large number of pixels.
+ This value is the standard deviation for the gaussian_filter
+ kernel.
+
+ Returns
+ -------
+ out : dict
+ Dictionary of values defining the mean/std of the bias + base
+ data and the scalar + adder factors to correct the biased data
+ like: bias_data * scalar + adder. Each value is of shape
+ (lat, lon, time).
+ """
+ if len(self.bad_bias_gids) > 0:
+ logger.info('Found {} bias gids that are out of bounds: {}'
+ .format(len(self.bad_bias_gids), self.bad_bias_gids))
+
+ for key, arr in out.items():
+ nan_mask = np.isnan(arr[..., 0])
+ for idt in range(arr.shape[-1]):
+
+ arr_smooth = arr[..., idt]
+
+ needs_fill = (np.isnan(arr_smooth).any()
+ and fill_extend) or smooth_interior > 0
+
+ if needs_fill:
+ logger.info('Filling NaN values outside of valid spatial '
+ 'extent for dataset "{}" for timestep {}'
+ .format(key, idt))
+ arr_smooth = nn_fill_array(arr_smooth)
+
+ arr_smooth_int = arr_smooth_ext = arr_smooth
+
+ if smooth_extend > 0:
+ arr_smooth_ext = gaussian_filter(arr_smooth_ext,
+ smooth_extend,
+ mode='nearest')
+
+ if smooth_interior > 0:
+ arr_smooth_int = gaussian_filter(arr_smooth_int,
+ smooth_interior,
+ mode='nearest')
+
+ out[key][nan_mask, idt] = arr_smooth_ext[nan_mask]
+ out[key][~nan_mask, idt] = arr_smooth_int[~nan_mask]
+
+ return out
+
+
+class LinearCorrection(FillAndSmoothMixin, DataRetrievalBase):
  """Calculate linear correction *scalar +adder factors to bias correct data
 
  This calculation operates on single bias sites for the full time series of
@@ -820,85 +904,6 @@ def _run_single(cls,
  base_dset)
  return out
 
- def fill_and_smooth(self,
- out,
- fill_extend=True,
- smooth_extend=0,
- smooth_interior=0):
- """Fill data extending beyond the base meta data extent by doing a
- nearest neighbor gap fill. Smooth interior and extended region with
- given smoothing values.
- Interior smoothing can reduce the affect of extreme values
- within aggregations over large number of pixels.
- The interior is assumed to be defined by the region without nan values.
- The extended region is assumed to be the region with nan values.
-
- Parameters
- ----------
- out : dict
- Dictionary of values defining the mean/std of the bias + base
- data and the scalar + adder factors to correct the biased data
- like: bias_data * scalar + adder. Each value is of shape
- (lat, lon, time).
- fill_extend : bool
- Whether to fill data extending beyond the base meta data with
- nearest neighbor values.
- smooth_extend : float
- Option to smooth the scalar/adder data outside of the spatial
- domain set by the threshold input. This alleviates the weird seams
- far from the domain of interest. This value is the standard
- deviation for the gaussian_filter kernel
- smooth_interior : float
- Value to use to smooth the scalar/adder data inside of the spatial
- domain set by the threshold input. This can reduce the effect of
- extreme values within aggregations over large number of pixels.
- This value is the standard deviation for the gaussian_filter
- kernel.
-
- Returns
- -------
- out : dict
- Dictionary of values defining the mean/std of the bias + base
- data and the scalar + adder factors to correct the biased data
- like: bias_data * scalar + adder. Each value is of shape
- (lat, lon, time).
- """
- if len(self.bad_bias_gids) > 0:
- logger.info('Found {} bias gids that are out of bounds: {}'
- .format(len(self.bad_bias_gids), self.bad_bias_gids))
-
- for key, arr in out.items():
- nan_mask = np.isnan(arr[..., 0])
- for idt in range(arr.shape[-1]):
-
- arr_smooth = arr[..., idt]
-
- needs_fill = (np.isnan(arr_smooth).any()
- and fill_extend) or smooth_interior > 0
-
- if needs_fill:
- logger.info('Filling NaN values outside of valid spatial '
- 'extent for dataset "{}" for timestep {}'
- .format(key, idt))
- arr_smooth = nn_fill_array(arr_smooth)
-
- arr_smooth_int = arr_smooth_ext = arr_smooth
-
- if smooth_extend > 0:
- arr_smooth_ext = gaussian_filter(arr_smooth_ext,
- smooth_extend,
- mode='nearest')
-
- if smooth_interior > 0:
- arr_smooth_int = gaussian_filter(arr_smooth_int,
- smooth_interior,
- mode='nearest')
-
- out[key][nan_mask, idt] = arr_smooth_ext[nan_mask]
- out[key][~nan_mask, idt] = arr_smooth_int[~nan_mask]
-
- return out
-
  def write_outputs(self, fp_out, out):
  """Write outputs to an .h5 file.
 
@@ -1176,7 +1181,7 @@ def get_linear_correction(bias_data, base_data, bias_feature, base_dset):
  return out
 
 
-class QuantileDeltaMappingCorrection(DataRetrievalBase):
+class QuantileDeltaMappingCorrection(FillAndSmoothMixin, DataRetrievalBase):
  """Estimate probability distributions required by Quantile Delta Mapping
 
  The main purpose of this class is to estimate the probability
@@ -1457,7 +1462,10 @@ def write_outputs(self, fp_out, out=None):
  def run(self,
  fp_out=None,
  max_workers=None,
- daily_reduction='avg'):
+ daily_reduction='avg',
+ fill_extend=True,
+ smooth_extend=0,
+ smooth_interior=0):
  """Estimate the statistical distributions for each location
 
  Parameters
@@ -1572,6 +1580,9 @@ def run(self,
 
  logger.info('Finished calculating bias correction factors.')
 
+ self.out = self.fill_and_smooth(self.out, fill_extend, smooth_extend,
+ smooth_interior)
+
  self.write_outputs(fp_out, self.out)
 
  return copy.deepcopy(self.out)