added temporal then spatial model with test

sup3r/models/__init__.py

@@ -4,7 +4,8 @@
 from .wind import WindGan
 from .solar_cc import SolarCC
 from .data_centric import Sup3rGanDC
-from .multi_step import (MultiStepGan, SpatialThenTemporalGan,
+from .multi_step import (MultiStepGan,
+ SpatialThenTemporalGan, TemporalThenSpatialGan,
  MultiStepSurfaceMetGan, SolarMultiStepGan)
 from .surface import SurfaceSpatialMetModel
 from .linear import LinearInterp

sup3r/models/multi_step.py

@@ -246,16 +246,8 @@ def model_params(self):
 
 
 class SpatialThenTemporalBase(MultiStepGan):
- """A two-step model where the first step is a spatial-only enhancement on a
- 4D tensor and the second step is (spatio)temporal enhancement on a 5D
- tensor.
-
- NOTE: The low res input to the spatial enhancement should be a 4D tensor of
- the shape (temporal, spatial_1, spatial_2, features) where temporal
- (usually the observation index) is a series of sequential timesteps that
- will be transposed to a 5D tensor of shape
- (1, spatial_1, spatial_2, temporal, features) tensor and then fed to the
- 2nd-step (spatio)temporal model.
+ """A base class for spatial-then-temporal or temporal-then-spatial multi
+ step GANs
  """
 
  def __init__(self, spatial_models, temporal_models):
@@ -272,22 +264,6 @@ def __init__(self, spatial_models, temporal_models):
  self._spatial_models = spatial_models
  self._temporal_models = temporal_models
 
- @property
- def models(self):
- """Get an ordered tuple of the Sup3rGan models that are part of this
- MultiStepGan
- """
- if isinstance(self.spatial_models, MultiStepGan):
- spatial_models = self.spatial_models.models
- else:
- spatial_models = [self.spatial_models]
- if isinstance(self.temporal_models, MultiStepGan):
- temporal_models = self.temporal_models.models
- else:
- temporal_models = [self.temporal_models]
-
- return (*spatial_models, *temporal_models)
-
  @property
  def spatial_models(self):
  """Get the MultiStepGan object for the spatial-only model(s)
@@ -308,6 +284,72 @@ def temporal_models(self):
  """
  return self._temporal_models
 
+ @classmethod
+ def load(cls, spatial_model_dirs, temporal_model_dirs, verbose=True):
+ """Load the GANs with its sub-networks from a previously saved-to
+ output directory.
+
+ Parameters
+ ----------
+ spatial_model_dirs : str | list | tuple
+ An ordered list/tuple of one or more directories containing trained
+ + saved Sup3rGan models created using the Sup3rGan.save() method.
+ This must contain only spatial models that input/output 4D
+ tensors.
+ temporal_model_dirs : str | list | tuple
+ An ordered list/tuple of one or more directories containing trained
+ + saved Sup3rGan models created using the Sup3rGan.save() method.
+ This must contain only (spatio)temporal models that input/output 5D
+ tensors.
+ verbose : bool
+ Flag to log information about the loaded model.
+
+ Returns
+ -------
+ out : MultiStepGan
+ Returns a pretrained gan model that was previously saved to
+ model_dirs
+ """
+ if isinstance(spatial_model_dirs, str):
+ spatial_model_dirs = [spatial_model_dirs]
+ if isinstance(temporal_model_dirs, str):
+ temporal_model_dirs = [temporal_model_dirs]
+
+ s_models = MultiStepGan.load(spatial_model_dirs, verbose=verbose)
+ t_models = MultiStepGan.load(temporal_model_dirs, verbose=verbose)
+
+ return cls(s_models, t_models)
+
+
+class SpatialThenTemporalGan(SpatialThenTemporalBase):
+ """A two-step GAN where the first step is a spatial-only enhancement on a
+ 4D tensor and the second step is a (spatio)temporal enhancement on a 5D
+ tensor.
+
+ NOTE: The low res input to the spatial enhancement should be a 4D tensor of
+ the shape (temporal, spatial_1, spatial_2, features) where temporal
+ (usually the observation index) is a series of sequential timesteps that
+ will be transposed to a 5D tensor of shape
+ (1, spatial_1, spatial_2, temporal, features) tensor and then fed to the
+ 2nd-step (spatio)temporal GAN.
+ """
+
+ @property
+ def models(self):
+ """Get an ordered tuple of the Sup3rGan models that are part of this
+ MultiStepGan
+ """
+ if isinstance(self.spatial_models, MultiStepGan):
+ spatial_models = self.spatial_models.models
+ else:
+ spatial_models = [self.spatial_models]
+ if isinstance(self.temporal_models, MultiStepGan):
+ temporal_models = self.temporal_models.models
+ else:
+ temporal_models = [self.temporal_models]
+
+ return (*spatial_models, *temporal_models)
+
  @property
  def meta(self):
  """Get a tuple of meta data dictionaries for all models
@@ -329,14 +371,14 @@ def meta(self):
  @property
  def training_features(self):
  """Get the list of input feature names that the first spatial
- generative model in this SpatialThenTemporalBase model requires as
+ generative model in this SpatialThenTemporalGan model requires as
  input."""
  return self.spatial_models.training_features
 
  @property
  def output_features(self):
  """Get the list of output feature names that the last spatiotemporal
- interpolation model in this SpatialThenTemporalBase model outputs."""
+ interpolation model in this SpatialThenTemporalGan model outputs."""
  return self.temporal_models.output_features
 
  def generate(self, low_res, norm_in=True, un_norm_out=True,
@@ -412,58 +454,139 @@ def generate(self, low_res, norm_in=True, un_norm_out=True,
 
  return hi_res
 
- @classmethod
- def load(cls, spatial_model_dirs, temporal_model_dirs, verbose=True):
- """Load the GANs with its sub-networks from a previously saved-to
- output directory.
+
+class TemporalThenSpatialGan(SpatialThenTemporalBase):
+ """A two-step GAN where the first step is a spatiotemporal enhancement on a
+ 5D tensor and the second step is a spatial enhancement on a 4D tensor.
+ """
+
+ @property
+ def models(self):
+ """Get an ordered tuple of the Sup3rGan models that are part of this
+ MultiStepGan
+ """
+ if isinstance(self.spatial_models, MultiStepGan):
+ spatial_models = self.spatial_models.models
+ else:
+ spatial_models = [self.spatial_models]
+ if isinstance(self.temporal_models, MultiStepGan):
+ temporal_models = self.temporal_models.models
+ else:
+ temporal_models = [self.temporal_models]
+
+ return (*temporal_models, *spatial_models)
+
+ @property
+ def meta(self):
+ """Get a tuple of meta data dictionaries for all models
+
+ Returns
+ -------
+ tuple
+ """
+ if isinstance(self.spatial_models, MultiStepGan):
+ spatial_models = self.spatial_models.meta
+ else:
+ spatial_models = [self.spatial_models.meta]
+ if isinstance(self.temporal_models, MultiStepGan):
+ temporal_models = self.temporal_models.meta
+ else:
+ temporal_models = [self.temporal_models.meta]
+
+ return (*temporal_models, *spatial_models)
+
+ @property
+ def training_features(self):
+ """Get the list of input feature names that the first temporal
+ generative model in this TemporalThenSpatialGan model requires as
+ input."""
+ return self.temporal_models.training_features
+
+ @property
+ def output_features(self):
+ """Get the list of output feature names that the last spatial
+ interpolation model in this TemporalThenSpatialGan model outputs."""
+ return self.spatial_models.output_features
+
+ def generate(self, low_res, norm_in=True, un_norm_out=True,
+ exogenous_data=None):
+ """Use the generator model to generate high res data from low res
+ input. This is the public generate function.
 
  Parameters
  ----------
- spatial_model_dirs : str | list | tuple
- An ordered list/tuple of one or more directories containing trained
- + saved Sup3rGan models created using the Sup3rGan.save() method.
- This must contain only spatial models that input/output 4D
- tensors.
- temporal_model_dirs : str | list | tuple
- An ordered list/tuple of one or more directories containing trained
- + saved Sup3rGan models created using the Sup3rGan.save() method.
- This must contain only (spatio)temporal models that input/output 5D
- tensors.
- verbose : bool
- Flag to log information about the loaded model.
+ low_res : np.ndarray
+ Low-resolution input data, a 5D array of shape:
+ (1, spatial_1, spatial_2, n_temporal, n_features)
+ norm_in : bool
+ Flag to normalize low_res input data if the self.means,
+ self.stdevs attributes are available. The generator should always
+ received normalized data with mean=0 stdev=1.
+ un_norm_out : bool
+ Flag to un-normalize synthetically generated output data to physical
+ units
+ exogenous_data : list
+ List of arrays of exogenous_data with length equal to the
+ number of model steps. e.g. If we want to include topography as
+ an exogenous feature in a temporal + spatial multistep model then
+ we need to provide a list of length=2 with topography at the low
+ spatial resolution and at the high resolution. If we include more
+ than one exogenous feature the ordering must be consistent.
+ Each array in the list has 3D or 4D shape:
+ (spatial_1, spatial_2, n_features)
+ (temporal, spatial_1, spatial_2, n_features)
 
  Returns
  -------
- out : MultiStepGan
- Returns a pretrained gan model that was previously saved to
- model_dirs
+ hi_res : ndarray
+ Synthetically generated high-resolution data output from the 2nd
+ step (spatio)temporal GAN with a 5D array shape:
+ (1, spatial_1, spatial_2, n_temporal, n_features)
  """
- if isinstance(spatial_model_dirs, str):
- spatial_model_dirs = [spatial_model_dirs]
- if isinstance(temporal_model_dirs, str):
- temporal_model_dirs = [temporal_model_dirs]
+ logger.debug('Data input to the 1st step (spatio)temporal '
+ 'enhancement has shape {}'.format(low_res.shape))
+ s_exogenous = None
+ if exogenous_data is not None:
+ s_exogenous = exogenous_data[len(self.temporal_models):]
 
- s_models = MultiStepGan.load(spatial_model_dirs, verbose=verbose)
- t_models = MultiStepGan.load(temporal_model_dirs, verbose=verbose)
+ assert low_res.shape[0] == 1, 'Low res input can only have 1 obs!'
 
- return cls(s_models, t_models)
+ try:
+ hi_res = self.temporal_models.generate(
+ low_res, norm_in=norm_in, un_norm_out=True,
+ exogenous_data=exogenous_data)
+ except Exception as e:
+ msg = ('Could not run the 1st step (spatio)temporal GAN on input '
+ 'shape {}'.format(low_res.shape))
+ logger.exception(msg)
+ raise RuntimeError(msg) from e
 
+ logger.debug('Data output from the 1st step (spatio)temporal '
+ 'enhancement has shape {}'.format(hi_res.shape))
+ hi_res = np.transpose(hi_res[0], axes=(2, 0, 1, 3))
+ logger.debug('Data from the 1st step (spatio)temporal enhancement has '
+ 'been reshaped to {}'.format(hi_res.shape))
 
-class SpatialThenTemporalGan(SpatialThenTemporalBase):
- """A two-step GAN where the first step is a spatial-only enhancement on a
- 4D tensor and the second step is a (spatio)temporal enhancement on a 5D
- tensor.
+ try:
+ hi_res = self.spatial_models.generate(
+ hi_res, norm_in=True, un_norm_out=un_norm_out,
+ exogenous_data=s_exogenous)
+ except Exception as e:
+ msg = ('Could not run the 2nd step spatial GAN on input '
+ 'shape {}'.format(low_res.shape))
+ logger.exception(msg)
+ raise RuntimeError(msg) from e
 
- NOTE: The low res input to the spatial enhancement should be a 4D tensor of
- the shape (temporal, spatial_1, spatial_2, features) where temporal
- (usually the observation index) is a series of sequential timesteps that
- will be transposed to a 5D tensor of shape
- (1, spatial_1, spatial_2, temporal, features) tensor and then fed to the
- 2nd-step (spatio)temporal GAN.
- """
+  hi_res = np.transpose(hi_res, axes=(1, 2, 0, 3))
+  hi_res = np.expand_dims(hi_res, axis=0)
+
+  logger.debug('Final multistep GAN output has shape: {}'
+  .format(hi_res.shape))
+
+  return hi_res
 
 
-class MultiStepSurfaceMetGan(SpatialThenTemporalBase):
+class MultiStepSurfaceMetGan(SpatialThenTemporalGan):
  """A two-step GAN where the first step is a spatial-only enhancement on a
  4D tensor of near-surface temperature and relative humidity data, and the
  second step is a (spatio)temporal enhancement on a 5D tensor.
@@ -612,7 +735,7 @@ def load(cls, surface_model_class='SurfaceSpatialMetModel',
  return cls(s_models, t_models)
 
 
-class SolarMultiStepGan(SpatialThenTemporalBase):
+class SolarMultiStepGan(SpatialThenTemporalGan):
  """Special multi step model for solar clearsky ratio super resolution.
 
  This model takes in two parallel models for wind-only and solar-only

sup3r/models/wind.py

@@ -50,11 +50,11 @@ def init_weights(self, lr_shape, hr_shape, device=None):
  device = self.default_device
 
  logger.info('Initializing model weights on device "{}"'.format(device))
- low_res = np.random.uniform(0, 1, lr_shape).astype(np.float32)
- hi_res = np.random.uniform(0, 1, hr_shape).astype(np.float32)
+ low_res = np.ones(lr_shape).astype(np.float32)
+ hi_res = np.ones(hr_shape).astype(np.float32)
 
  hr_topo_shape = hr_shape[:-1] + (1,)
- hr_topo = np.random.uniform(0, 1, hr_topo_shape).astype(np.float32)
+ hr_topo = np.ones(hr_topo_shape).astype(np.float32)
 
  with tf.device(device):
  _ = self._tf_generate(low_res, hr_topo)

tests/test_multi_step.py

@@ -6,7 +6,8 @@
 import tempfile
 
 from sup3r import CONFIG_DIR
-from sup3r.models import (Sup3rGan, MultiStepGan, SpatialThenTemporalGan,
+from sup3r.models import (Sup3rGan, MultiStepGan,
+ SpatialThenTemporalGan, TemporalThenSpatialGan,
  SolarMultiStepGan, LinearInterp)
 
 FEATURES = ['U_100m', 'V_100m']
@@ -129,6 +130,38 @@ def test_spatial_then_temporal_gan():
  assert out.shape == (1, 60, 60, 16, 2)
 
 
+def test_temporal_then_spatial_gan():
+ """Test the 2-step temporal-then-spatial GAN"""
+ fp_gen = os.path.join(CONFIG_DIR, 'spatial/gen_2x_2f.json')
+ fp_disc = os.path.join(CONFIG_DIR, 'spatial/disc.json')
+ model1 = Sup3rGan(fp_gen, fp_disc)
+ _ = model1.generate(np.ones((4, 10, 10, len(FEATURES))))
+
+ fp_gen = os.path.join(CONFIG_DIR, 'spatiotemporal/gen_3x_4x_2f.json')
+ fp_disc = os.path.join(CONFIG_DIR, 'spatiotemporal/disc.json')
+ model2 = Sup3rGan(fp_gen, fp_disc)
+ _ = model2.generate(np.ones((4, 10, 10, 6, len(FEATURES))))
+
+ model1.set_norm_stats([0.1, 0.2], [0.04, 0.02])
+ model2.set_norm_stats([0.3, 0.9], [0.02, 0.07])
+ model1.set_model_params(training_features=FEATURES,
+ output_features=FEATURES)
+ model2.set_model_params(training_features=FEATURES,
+ output_features=FEATURES)
+
+ with tempfile.TemporaryDirectory() as td:
+ fp1 = os.path.join(td, 'model1')
+ fp2 = os.path.join(td, 'model2')
+ model1.save(fp1)
+ model2.save(fp2)
+
+ ms_model = TemporalThenSpatialGan.load(fp1, fp2)
+
+ x = np.ones((1, 10, 10, 4, len(FEATURES)))
+ out = ms_model.generate(x)
+ assert out.shape == (1, 60, 60, 16, 2)
+
+
 def test_spatial_gan_then_linear_interp():
  """Test the 2-step spatial GAN then linear spatiotemporal interpolation"""
  fp_gen = os.path.join(CONFIG_DIR, 'spatial/gen_2x_2f.json')