Merge pull request #169 from NREL/bnb/exo_refactor_plus

Bnb/exo refactor plus

sup3r/models/abstract.py

@@ -21,6 +21,7 @@
 from tensorflow.keras import optimizers
 
 import sup3r.utilities.loss_metrics
+from sup3r.preprocessing.data_handling.exogenous_data_handling import ExoData
 from sup3r.utilities import VERSION_RECORD
 
 logger = logging.getLogger(__name__)
@@ -223,14 +224,11 @@ def _combine_fwp_input(self, low_res, exogenous_data=None):
  Low-resolution input data, usually a 4D or 5D array of shape:
  (n_obs, spatial_1, spatial_2, n_features)
  (n_obs, spatial_1, spatial_2, n_temporal, n_features)
- exogenous_data : dict | None
- Dictionary of exogenous feature data with entries describing
- whether features should be combined at input, a mid network layer,
- or with output. This doesn't have to include the 'model' key since
- this data is for a single step model. e.g.
- {'topography': {'steps': [
- {'combine_type': 'input', 'data': ..., 'resolution': ...},
- {'combine_type': 'layer', 'data': ..., 'resolution': ...}]}}
+ exogenous_data : dict | ExoData | None
+ Special dictionary (class:`ExoData`) of exogenous feature data with
+ entries describing whether features should be combined at input, a
+ mid network layer, or with output. This doesn't have to include
+ the 'model' key since this data is for a single step model.
 
  Returns
  -------
@@ -243,6 +241,10 @@ def _combine_fwp_input(self, low_res, exogenous_data=None):
  if exogenous_data is None:
  return low_res
 
+ if (not isinstance(exogenous_data, ExoData)
+ and exogenous_data is not None):
+ exogenous_data = ExoData(exogenous_data)
+
  training_features = ([] if self.training_features is None
  else self.training_features)
  fnum_diff = len(training_features) - low_res.shape[-1]
@@ -253,14 +255,10 @@ def _combine_fwp_input(self, low_res, exogenous_data=None):
  assert all(feature in exogenous_data for feature in exo_feats), msg
  if exogenous_data is not None and fnum_diff > 0:
  for feature in exo_feats:
- entry = exogenous_data[feature]
- combine_types = [step['combine_type']
- for step in entry['steps']]
- if 'input' in combine_types:
- idx = combine_types.index('input')
- low_res = np.concatenate((low_res,
- entry['steps'][idx]['data']),
- axis=-1)
+ exo_input = exogenous_data.get_combine_type_data(
+ feature, 'input')
+ if exo_input is not None:
+ low_res = np.concatenate((low_res, exo_input), axis=-1)
  return low_res
 
  def _combine_fwp_output(self, hi_res, exogenous_data=None):
@@ -273,14 +271,11 @@ def _combine_fwp_output(self, hi_res, exogenous_data=None):
  High-resolution output data, usually a 4D or 5D array of shape:
  (n_obs, spatial_1, spatial_2, n_features)
  (n_obs, spatial_1, spatial_2, n_temporal, n_features)
- exogenous_data : dict | None
- Dictionary of exogenous feature data with entries describing
- whether features should be combined at input, a mid network layer,
- or with output. This doesn't have to include the 'model' key since
- this data is for a single step model. e.g.
- {'topography': {'steps': [
- {'combine_type': 'input', 'data': ..., 'resolution': ...},
- {'combine_type': 'layer', 'data': ..., 'resolution': ...}]}}
+ exogenous_data : dict | ExoData | None
+ Special dictionary (class:`ExoData`) of exogenous feature data with
+ entries describing whether features should be combined at input, a
+ mid network layer, or with output. This doesn't have to include
+ the 'model' key since this data is for a single step model.
 
  Returns
  -------
@@ -293,6 +288,10 @@ def _combine_fwp_output(self, hi_res, exogenous_data=None):
  if exogenous_data is None:
  return hi_res
 
+ if (not isinstance(exogenous_data, ExoData)
+ and exogenous_data is not None):
+ exogenous_data = ExoData(exogenous_data)
+
  output_features = ([] if self.output_features is None
  else self.output_features)
  fnum_diff = len(output_features) - hi_res.shape[-1]
@@ -303,14 +302,10 @@ def _combine_fwp_output(self, hi_res, exogenous_data=None):
  assert all(feature in exogenous_data for feature in exo_feats), msg
  if exogenous_data is not None and fnum_diff > 0:
  for feature in exo_feats:
- entry = exogenous_data[feature]
- combine_types = [step['combine_type']
- for step in entry['steps']]
- if 'output' in combine_types:
- idx = combine_types.index('output')
- hi_res = np.concatenate((hi_res,
- entry['steps'][idx]['data']),
- axis=-1)
+ exo_output = exogenous_data.get_combine_type_data(
+ feature, 'output')
+ if exo_output is not None:
+ hi_res = np.concatenate((hi_res, exo_output), axis=-1)
  return hi_res
 
  def _combine_loss_input(self, high_res_true, high_res_gen):
@@ -1237,39 +1232,6 @@ def _reshape_norm_exo(self, hi_res, hi_res_exo, exo_name, norm_in=True):
 
  return hi_res_exo
 
- def _get_layer_exo_input(self, layer_name, exogenous_data):
- """Get the high-resolution exo data for the given layer name from the
- full exogenous_data dictionary.
-
- Parameters
- ----------
- layer_name : str
- Name of Sup3rAdder or Sup3rConcat layer. This should match a
- feature key in exogenous_data
- exogenous_data : dict | None
- Dictionary of exogenous feature data with entries describing
- whether features should be combined at input, a mid network layer,
- or with output. This doesn't have to include the 'model' key since
- this data is for a single step model. e.g.
- {'topography': {'steps': [
- {'combine_type': 'input', 'data': ..., 'resolution': ...},
- {'combine_type': 'layer', 'data': ..., 'resolution': ...}]}}
-
- """
- msg = (f'layer.name = {layer_name} does not match any '
- 'features in exogenous_data '
- f'({list(exogenous_data)})')
- assert layer_name in exogenous_data, msg
- steps = exogenous_data[layer_name]['steps']
- combine_types = [step['combine_type'] for step in steps]
- msg = ('Received exogenous_data without any combine_type '
- '= "layer" steps, for a model with an Adder/Concat '
- 'layer.')
- assert 'layer' in combine_types, msg
- idx = combine_types.index('layer')
- hi_res_exo = steps[idx]['data']
- return hi_res_exo
-
  def generate(self,
  low_res,
  norm_in=True,
@@ -1292,14 +1254,11 @@ def generate(self,
  un_norm_out : bool
  Flag to un-normalize synthetically generated output data to physical
  units
- exogenous_data : dict | None
- Dictionary of exogenous feature data with entries describing
- whether features should be combined at input, a mid network layer,
- or with output. This doesn't have to include the 'model' key since
- this data is for a single step model. e.g.
- {'topography': {'steps': [
- {'combine_type': 'input', 'data': ..., 'resolution': ...},
- {'combine_type': 'layer', 'data': ..., 'resolution': ...}]}}
+ exogenous_data : dict | ExoData | None
+ Special dictionary (class:`ExoData`) of exogenous feature data with
+ entries describing whether features should be combined at input, a
+ mid network layer, or with output. This doesn't have to include
+ the 'model' key since this data is for a single step model.
 
  Returns
  -------
@@ -1309,6 +1268,10 @@ def generate(self,
  (n_obs, spatial_1, spatial_2, n_features)
  (n_obs, spatial_1, spatial_2, n_temporal, n_features)
  """
+ if (not isinstance(exogenous_data, ExoData)
+ and exogenous_data is not None):
+ exogenous_data = ExoData(exogenous_data)
+
  low_res = self._combine_fwp_input(low_res, exogenous_data)
  if norm_in and self._means is not None:
  low_res = self.norm_input(low_res)
@@ -1317,8 +1280,12 @@ def generate(self,
  for i, layer in enumerate(self.generator.layers[1:]):
  try:
  if isinstance(layer, (Sup3rAdder, Sup3rConcat)):
- hi_res_exo = self._get_layer_exo_input(layer.name,
- exogenous_data)
+ msg = (f'layer.name = {layer.name} does not match any '
+ 'features in exogenous_data '
+ f'({list(exogenous_data)})')
+ assert layer.name in exogenous_data, msg
+ hi_res_exo = exogenous_data.get_combine_type_data(
+ layer.name, 'layer')
  hi_res_exo = self._reshape_norm_exo(hi_res,
  hi_res_exo,
  layer.name,