PR updates

sup3r/models/abstract.py

@@ -224,7 +224,7 @@ 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 : ExoData | None
+ 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
@@ -241,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]
@@ -267,7 +271,7 @@ 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
+ 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
@@ -284,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]
@@ -1260,8 +1268,8 @@ def generate(self,
  (n_obs, spatial_1, spatial_2, n_features)
  (n_obs, spatial_1, spatial_2, n_temporal, n_features)
  """
- if (isinstance(exogenous_data, dict)
- and not isinstance(exogenous_data, ExoData)):
+ 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)

sup3r/pipeline/forward_pass.py

@@ -1128,7 +1128,7 @@ def load_exo_data(self):
  class:`ExoData` object composed of multiple
  class:`SingleExoDataStep` objects.
  """
- data = []
+ data = {}
  exo_data = None
  if self.exo_kwargs:
  self.features = [f for f in self.features
@@ -1144,7 +1144,7 @@ def load_exo_data(self):
  sig = signature(ExogenousDataHandler)
  exo_kwargs = {k: v for k, v in exo_kwargs.items()
  if k in sig.parameters}
- data += ExogenousDataHandler(**exo_kwargs).data
+ data.update(ExogenousDataHandler(**exo_kwargs).data)
  exo_data = ExoData(data)
  return exo_data
 

sup3r/preprocessing/data_handling/exogenous_data_handling.py

@@ -60,19 +60,24 @@ def __init__(self, steps):
 
  Parameters
  ----------
- steps : list | dict
- List of SingleExoDataStep objects or a feature dictionary with list
- of steps for each feature
+ steps : dict
+ Dictionary with feature keys each with entries describing whether
+ features should be combined at input, a mid network layer, or with
+ output. e.g.
+ {'topography': {'steps': [
+ {'combine_type': 'input', 'model': 0, 'data': ...,
+ 'resolution': ...},
+ {'combine_type': 'layer', 'model': 0, 'data': ...,
+ 'resolution': ...}]}}
+ Each array in in 'data' key has 3D or 4D shape:
+ (spatial_1, spatial_2, 1)
+ (spatial_1, spatial_2, n_temporal, 1)
  """
- if isinstance(steps, list):
- for step in steps:
- self.append(step.feature, step)
- elif isinstance(steps, dict):
+ if isinstance(steps, dict):
  for k, v in steps.items():
  self.__setitem__(k, v)
  else:
- msg = ('ExoData must be initialized with a dictionary of features '
- 'or list of SingleExoDataStep objects.')
+ msg = 'ExoData must be initialized with a dictionary of features.'
  logger.error(msg)
  raise ValueError(msg)
 
@@ -117,18 +122,6 @@ def split_exo_dict(self, split_step):
  spatial models and temporal models split_step should be
  len(spatial_models). If this is for a TemporalThenSpatial model
  split_step should be len(temporal_models).
- exogenous_data : dict
- Dictionary of exogenous feature data with entries describing
- whether features should be combined at input, a mid network layer,
- or with output. e.g.
- {'topography': {'steps': [
- {'combine_type': 'input', 'model': 0, 'data': ...,
- 'resolution': ...},
- {'combine_type': 'layer', 'model': 0, 'data': ...,
- 'resolution': ...}]}}
- Each array in in 'data' key has 3D or 4D shape:
- (spatial_1, spatial_2, 1)
- (spatial_1, spatial_2, n_temporal, 1)
 
  Returns
  -------
@@ -306,7 +299,7 @@ def __init__(self,
  self.input_handler = input_handler
  self.cache_data = cache_data
  self.cache_dir = cache_dir
- self.data = []
+ self.data = {feature: {'steps': []}}
 
  self.input_check()
  agg_enhance = self._get_all_agg_and_enhancement()
@@ -341,16 +334,16 @@ def __init__(self,
  t_agg_factor=t_agg_factor)
  step = SingleExoDataStep(feature, steps[i]['combine_type'],
  steps[i]['model'], data)
- self.data.append(step)
+ self.data[feature]['steps'].append(step)
  else:
  msg = (f"Can only extract {list(self.AVAILABLE_HANDLERS)}."
  f" Received {feature}.")
  raise NotImplementedError(msg)
- shapes = [None if d is None else d['data'].shape
- for d in self.data]
+ shapes = [None if step is None else step.shape
+ for step in self.data[feature]['steps']]
  logger.info(
  'Got exogenous_data of length {} with shapes: {}'.format(
- len(self.data), shapes))
+ len(self.data[feature]['steps']), shapes))
 
  def input_check(self):
  """Make sure agg factors are provided or exo_resolution and models are

tests/data_handling/test_exo_data_handling.py

@@ -43,7 +43,7 @@ def test_exo_cache(feature):
  target=TARGET, shape=SHAPE,
  input_handler='DataHandlerNCforCC',
  cache_dir=os.path.join(td, 'exo_cache'))
- for i, arr in enumerate(base.data):
+ for i, arr in enumerate(base.data[feature]['steps']):
  assert arr.shape[0] == SHAPE[0] * S_ENHANCE[i]
  assert arr.shape[1] == SHAPE[1] * S_ENHANCE[i]
 
@@ -58,5 +58,6 @@ def test_exo_cache(feature):
  cache_dir=os.path.join(td, 'exo_cache'))
  assert len(os.listdir(f'{td}/exo_cache')) == 2
 
- for arr1, arr2 in zip(base.data, cache.data):
+ for arr1, arr2 in zip(base.data[feature]['steps'],
+ cache.data[feature]['steps']):
  assert np.allclose(arr1['data'], arr2['data'])