added montly bias correction data transformation method and integrate…

…d into fwp and qa

sup3r/bias/__init__.py

@@ -1,3 +1,5 @@
 # -*- coding: utf-8 -*-
 """Bias calculation and correction modules."""
-from .bias_transforms import global_linear_bc, local_linear_bc
+from .bias_transforms import (global_linear_bc, local_linear_bc,
+ monthly_local_linear_bc)
+from .bias_calc import LinearCorrection, MonthlyLinearCorrection

sup3r/bias/bias_transforms.py

@@ -1,16 +1,22 @@
 # -*- coding: utf-8 -*-
 """Bias correction transformation functions."""
 import numpy as np
+import logging
+from warnings import warn
 from rex import Resource
 
 
+logger = logging.getLogger(__name__)
+
+
 def global_linear_bc(input, scalar, adder, out_range=None):
  """Bias correct data using a simple global *scalar +adder method.
 
  Parameters
  ----------
  input : np.ndarray
- Any data to be bias corrected
+ Sup3r input data to be bias corrected, assumed to be 3D with shape
+ (spatial, spatial, temporal) for a single feature.
  scalar : float
  Scalar (multiplicative) value to apply to input data.
  adder : float
@@ -37,7 +43,8 @@ def local_linear_bc(input, feature_name, bias_fp, lr_padded_slice,
  Parameters
  ----------
  input : np.ndarray
- Any data to be bias corrected
+ Sup3r input data to be bias corrected, assumed to be 3D with shape
+ (spatial, spatial, temporal) for a single feature.
  feature_name : str
  Name of feature that is being corrected. Datasets with names
  "{feature_name}_scalar" and "{feature_name}_adder" will be retrieved
@@ -69,6 +76,11 @@ def local_linear_bc(input, feature_name, bias_fp, lr_padded_slice,
  scalar = res[scalar]
  adder = res[adder]
 
+ # 3D bias correction factors have seasonal/monthly correction in last axis
+ if len(scalar.shape) == 3 and len(adder.shape) == 3:
+ scalar = scalar.mean(axis=-1)
+ adder = adder.mean(axis=-1)
+
  if lr_padded_slice is not None:
  spatial_slice = (lr_padded_slice[0], lr_padded_slice[1])
  scalar = scalar[spatial_slice]
@@ -86,3 +98,86 @@ def local_linear_bc(input, feature_name, bias_fp, lr_padded_slice,
  out = np.minimum(out, np.max(out_range))
 
  return out
+
+
+def monthly_local_linear_bc(input, feature_name, bias_fp, lr_padded_slice,
+ time_index, temporal_avg=True, out_range=None):
+ """Bias correct data using a simple global *scalar +adder method.
+
+ Parameters
+ ----------
+ input : np.ndarray
+ Sup3r input data to be bias corrected, assumed to be 3D with shape
+ (spatial, spatial, temporal) for a single feature.
+ feature_name : str
+ Name of feature that is being corrected. Datasets with names
+ "{feature_name}_scalar" and "{feature_name}_adder" will be retrieved
+ from bias_fp.
+ bias_fp : str
+ Filepath to bias correction file from the bias calc module. Must have
+ datasets "{feature_name}_scalar" and "{feature_name}_adder" that are
+ the full low-resolution shape of the forward pass input that will be
+ sliced using lr_padded_slice for the current chunk.
+ lr_padded_slice : tuple | None
+ Tuple of length four that slices (spatial_1, spatial_2, temporal,
+ features) where each tuple entry is a slice object for that axes.
+ Note that if this method is called as part of a sup3r forward pass, the
+ lr_padded_slice will be included automatically in the kwargs for the
+ active chunk. If this is None, no slicing will be done and the full
+ bias correction source shape will be used.
+ time_index : pd.DatetimeIndex
+ DatetimeIndex object associated with the input data temporal axis
+ (assumed 3rd axis e.g. axis=2). Note that if this method is called as
+ part of a sup3r resolution forward pass, the time_index will be
+ included automatically for the current chunk.
+ temporal_avg : bool
+ Take the average scalars and adders for the chunk's time index, this
+ will smooth the transition of scalars/adders from month to month if
+ processing small chunks. If processing the full annual time index, set
+ this to False.
+ out_range : None | tuple
+ Option to set floor/ceiling values on the output data.
+
+ Returns
+ -------
+ out : np.ndarray
+ out = input * scalar + adder
+ """
+
+ scalar = f'{feature_name}_scalar'
+ adder = f'{feature_name}_adder'
+ with Resource(bias_fp) as res:
+ scalar = res[scalar]
+ adder = res[adder]
+
+ assert len(scalar.shape) == 3, 'Monthly bias correct needs 3D scalars'
+ assert len(adder.shape) == 3, 'Monthly bias correct needs 3D adders'
+
+ if lr_padded_slice is not None:
+ spatial_slice = (lr_padded_slice[0], lr_padded_slice[1])
+ scalar = scalar[spatial_slice]
+ adder = adder[spatial_slice]
+
+ imonths = time_index.month.values - 1
+ scalar = scalar[..., imonths]
+ adder = adder[..., imonths]
+
+ if temporal_avg:
+ scalar = scalar.mean(axis=-1)
+ adder = adder.mean(axis=-1)
+ scalar = np.expand_dims(scalar, axis=-1)
+ adder = np.expand_dims(adder, axis=-1)
+ scalar = np.repeat(scalar, input.shape[-1], axis=-1)
+ adder = np.repeat(adder, input.shape[-1], axis=-1)
+ if len(time_index.month.unique()) > 2:
+ msg = ('Bias correction method "monthly_local_linear_bc" was used '
+ 'with temporal averaging over a time index with >2 months.')
+ warn(msg)
+ logger.warning(msg)
+
+ out = input * scalar + adder
+ if out_range is not None:
+ out = np.maximum(out, np.min(out_range))
+ out = np.minimum(out, np.max(out_range))
+
+ return out

sup3r/pipeline/forward_pass.py

@@ -18,7 +18,7 @@
 from rex.utilities.execution import SpawnProcessPool
 
 import sup3r.models
-import sup3r.bias
+import sup3r.bias.bias_transforms
 from sup3r.preprocessing.data_handling import InputMixIn
 from sup3r.preprocessing.exogenous_data_handling import ExogenousDataHandler
 from sup3r.postprocessing.file_handling import (OutputHandlerH5,
@@ -694,10 +694,11 @@ def __init__(self, file_paths, model_kwargs, fwp_chunk_shape,
  pass
  bias_correct_method : str | None
  Optional bias correction function name that can be imported from
- the sup3r.bias module. This will transform the source data
- according to some predefined bias correction transformation along
- with the bias_correct_kwargs. As the first argument, this method
- must receive a generic numpy array of data to be bias corrected
+ the sup3r.bias.bias_transforms module. This will transform the
+ source data according to some predefined bias correction
+ transformation along with the bias_correct_kwargs. As the first
+ argument, this method must receive a generic numpy array of data to
+ be bias corrected
  bias_correct_kwargs : dict | None
  Optional namespace of kwargs to provide to bias_correct_method.
  If this is provided, it must be a dictionary where each key is a
@@ -1394,13 +1395,15 @@ def bias_correct_source_data(self, data):
  method = self.strategy.bias_correct_method
  kwargs = self.strategy.bias_correct_kwargs
  if method is not None:
- method = getattr(sup3r.bias, method)
+ method = getattr(sup3r.bias.bias_transforms, method)
  logger.info('Running bias correction with: {}'.format(method))
  for feature, feature_kwargs in kwargs.items():
  idf = self.data_handler.features.index(feature)
 
  if 'lr_padded_slice' in signature(method).parameters:
  feature_kwargs['lr_padded_slice'] = self.lr_padded_slice
+ if 'time_index' in signature(method).parameters:
+ feature_kwargs['time_index'] = self.data_handler.time_index
 
  logger.debug('Bias correcting feature "{}" at axis index {} '
  'using function: {} with kwargs: {}'

sup3r/qa/qa.py

@@ -6,12 +6,12 @@
 import numpy as np
 import xarray as xr
 import logging
+from inspect import signature
 from rex import Resource
 from rex.utilities.fun_utils import get_fun_call_str
-import sup3r.bias
+import sup3r.bias.bias_transforms
 from sup3r.preprocessing.feature_handling import Feature
 from sup3r.postprocessing.file_handling import RexOutputs, H5_ATTRS
-from sup3r.preprocessing.feature_handling import Feature
 from sup3r.utilities import ModuleName
 from sup3r.utilities.utilities import (get_input_handler_class,
  get_source_type,
@@ -102,10 +102,11 @@ def __init__(self, source_file_paths, out_file_path, s_enhance, t_enhance,
  (only .h5 is supported)
  bias_correct_method : str | None
  Optional bias correction function name that can be imported from
- the sup3r.bias module. This will transform the source data
- according to some predefined bias correction transformation along
- with the bias_correct_kwargs. As the first argument, this method
- must receive a generic numpy array of data to be bias corrected
+ the sup3r.bias.bias_transforms module. This will transform the
+ source data according to some predefined bias correction
+ transformation along with the bias_correct_kwargs. As the first
+ argument, this method must receive a generic numpy array of data to
+ be bias corrected
  bias_correct_kwargs : dict | None
  Optional namespace of kwargs to provide to bias_correct_method.
  If this is provided, it must be a dictionary where each key is a
@@ -320,9 +321,11 @@ def bias_correct_source_data(self, data, feature):
  method = self.bias_correct_method
  kwargs = self.bias_correct_kwargs
  if method is not None:
- method = getattr(sup3r.bias, method)
+ method = getattr(sup3r.bias.bias_transforms, method)
  logger.info('Running bias correction with: {}'.format(method))
  feature_kwargs = kwargs[feature]
+ if 'time_index' in signature(method).parameters:
+ feature_kwargs['time_index'] = self.time_index
  logger.debug('Bias correcting feature "{}" using function: {} '
  'with kwargs: {}'
  .format(feature, method, feature_kwargs))