Skip to content

Commit

Permalink
Merge pull request #169 from NREL/bnb/exo_refactor_plus
Browse files Browse the repository at this point in the history 
Bnb/exo refactor plus
  • Loading branch information
@bnb32
bnb32 authored Oct 13, 2023
2 parents 4330268 + 53d8bc7 commit 91fab51
Show file tree
Hide file tree
Showing 8 changed files with 273 additions and 325 deletions.
117 changes: 42 additions & 75 deletions sup3r/models/abstract.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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
-------
Expand All @@ -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]
Expand All @@ -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):
Expand All @@ -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
-------
Expand All @@ -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]
Expand All @@ -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):
Expand Down Expand Up @@ -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,
Expand All @@ -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
-------
Expand All @@ -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)
Expand All @@ -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,
Expand Down
Loading

0 comments on commit 91fab51

Please sign in to comment.