Merge pull request #265 from NREL/gb/bc_optm

Gb/bc optm

Author: grantbuster

sup3r/bias/bias_transforms.py

@@ -575,10 +575,10 @@ def _apply_qdm(
         if no_trend
         else np.reshape(bias_fut_params, (-1, bias_fut_params.shape[-1]))
     )
+
     # The distributions at this point, after selected the respective
     # time window with `window_idx`, are 3D (space, space, N-params)
     # Collapse 3D (space, space, N) into 2D (space**2, N)
-
     QDM = QuantileDeltaMapping(
         params_oh=np.reshape(base_params, (-1, base_params.shape[-1])),
         params_mh=np.reshape(bias_params, (-1, bias_params.shape[-1])),
@@ -594,8 +594,12 @@ def _apply_qdm(
     # input 3D shape (spatial, spatial, temporal)
     # QDM expects input arr with shape (time, space)
     tmp = np.reshape(subset.data, (-1, subset.shape[-1])).T
+
     # Apply QDM correction
+    logger.info(f'Applying QDM to data with shape {tmp.shape}...')
     tmp = QDM(tmp, max_workers=max_workers)
+    logger.info(f'Finished QDM on data shape {tmp.shape}!')
+
     # Reorgnize array back from  (time, space)
     # to (spatial, spatial, temporal)
     return np.reshape(tmp.T, subset.shape)
@@ -751,9 +755,13 @@ def local_qdm_bc(
     )
 
     cfg = params['cfg']
-    base_params = params['base']
-    bias_params = params['bias']
-    bias_fut_params = params.get('bias_fut', None)
+
+    # params as dask arrays slows down QDM by several orders of magnitude
+    base_params = np.array(params['base'])
+    bias_params = np.array(params['bias'])
+    bias_fut_params = None
+    if 'bias_fut' in params:
+        bias_fut_params = np.array(params['bias_fut'])
 
     if lr_padded_slice is not None:
         spatial_slice = (lr_padded_slice[0], lr_padded_slice[1])

sup3r/bias/utilities.py

@@ -15,6 +15,7 @@
     get_date_range_kwargs,
 )
 
+
 logger = logging.getLogger(__name__)
 
 
@@ -249,8 +250,10 @@ def bias_correct_feature(
         Data corrected by the bias_correct_method ready for input to the
         forward pass through the generative model.
     """
+
     time_slice = _parse_time_slice(time_slice)
     data = input_handler[source_feature][..., time_slice]
+
     lat_lon = input_handler.lat_lon
     if bc_method is not None:
         bc_method = getattr(sup3r.bias.bias_transforms, bc_method)
@@ -260,15 +263,15 @@ def bias_correct_feature(
         if 'date_range_kwargs' in signature(bc_method).parameters:
             ti = input_handler.time_index[time_slice]
             feature_kwargs['date_range_kwargs'] = get_date_range_kwargs(ti)
-        if (
-            'lr_padded_slice' in signature(bc_method).parameters
-            and 'lr_padded_slice' not in feature_kwargs
-        ):
+
+        use_lrps = 'lr_padded_slice' in signature(bc_method).parameters
+        need_lrps = 'lr_padded_slice' not in feature_kwargs
+        if use_lrps and need_lrps:
             feature_kwargs['lr_padded_slice'] = None
-        if (
-            'temporal_avg' in signature(bc_method).parameters
-            and 'temporal_avg' not in feature_kwargs
-        ):
+
+        use_tavg = 'temporal_avg' in signature(bc_method).parameters
+        need_tavg = 'temporal_avg' not in feature_kwargs
+        if use_tavg and need_tavg:
             msg = (
                 'The kwarg "temporal_avg" was not provided in the bias '
                 'correction kwargs but is present in the bias '
@@ -306,6 +309,7 @@ def bias_correct_features(
     """
 
     time_slice = _parse_time_slice(time_slice)
+
     for feat in features:
         try:
             input_handler[feat][..., time_slice] = bias_correct_feature(
@@ -324,4 +328,5 @@ def bias_correct_features(
             )
             logger.exception(msg)
             raise RuntimeError(msg) from e
+
     return input_handler

sup3r/pipeline/forward_pass.py

@@ -24,6 +24,7 @@
 )
 from sup3r.utilities import ModuleName
 from sup3r.utilities.cli import BaseCLI
+from sup3r.utilities.utilities import Timer
 
 logger = logging.getLogger(__name__)
 
@@ -52,6 +53,7 @@ def __init__(self, strategy, node_index=0):
         node_index : int
             Index of node used to run forward pass
         """
+        self.timer = Timer()
         self.strategy = strategy
         self.model = get_model(strategy.model_class, strategy.model_kwargs)
         self.node_index = node_index
@@ -230,7 +232,8 @@ def run_generator(
         temp = cls._reshape_data_chunk(model, data_chunk, exo_data)
         data_chunk, exo_data, i_lr_t, i_lr_s = temp
         try:
-            hi_res = model.generate(data_chunk, exogenous_data=exo_data)
+            fun = Timer()(model.generate, log=True)
+            hi_res = fun(data_chunk, exogenous_data=exo_data)
         except Exception as e:
             msg = 'Forward pass failed on chunk with shape {}.'.format(
                 data_chunk.shape

sup3r/pipeline/strategy.py

@@ -504,9 +504,10 @@ def prep_chunk_data(self, chunk_index=0):
                 f'Bias correcting data for chunk_index={chunk_index}, '
                 f'with shape={input_data.shape}'
             )
-            input_data = self.timer(
-                bias_correct_features, log=True, call_id=chunk_index
-            )(
+            fun = self.timer(
+                bias_correct_features, log=True, call_id=chunk_index,
+            )
+            input_data = fun(
                 features=list(self.bias_correct_kwargs),
                 input_handler=input_data,
                 bc_method=self.bias_correct_method,