diff --git a/sup3r/models/linear.py b/sup3r/models/linear.py index 4a5557bd0..e049a3efd 100644 --- a/sup3r/models/linear.py +++ b/sup3r/models/linear.py @@ -2,6 +2,7 @@ """Simple models for super resolution such as linear interp models.""" import numpy as np import logging +from inspect import signature import os import json from sup3r.utilities.utilities import st_interp @@ -45,7 +46,9 @@ def load(cls, model_dir, verbose=False): Parameters ---------- model_dir : str - Directory to load LinearInterp model files from. + Directory to load LinearInterp model files from. Must + have a model_params.json file containing "meta" key with all of the + class init args. verbose : bool Flag to log information about the loaded model. @@ -59,11 +62,10 @@ def load(cls, model_dir, verbose=False): with open(fp_params, 'r') as f: params = json.load(f) - meta = params.get('meta', {'class': 'Sup3rGan'}) - model = cls(features=meta['training_features'], - s_enhance=meta['s_enhance'], - t_enhance=meta['t_enhance'], - t_centered=meta['t_centered']) + meta = params['meta'] + args = signature(cls.__init__).parameters + kwargs = {k: v for k, v in meta.items() if k in args} + model = cls(**kwargs) if verbose: logger.info('Loading LinearInterp with meta data: {}' @@ -74,7 +76,8 @@ def load(cls, model_dir, verbose=False): @property def meta(self): """Get meta data dictionary that defines the model params""" - return {'s_enhance': self._s_enhance, + return {'features': self._features, + 's_enhance': self._s_enhance, 't_enhance': self._t_enhance, 't_centered': self._t_centered, 'training_features': self.training_features, diff --git a/sup3r/models/surface.py b/sup3r/models/surface.py index 6bd5abfcf..8d8e4a508 100644 --- a/sup3r/models/surface.py +++ b/sup3r/models/surface.py @@ -1,19 +1,22 @@ # -*- coding: utf-8 -*- """Special models for surface meteorological data.""" +import os +import json import logging +from inspect import signature from fnmatch import fnmatch import numpy as np from PIL import Image from sklearn import linear_model from warnings import warn -from sup3r.models.abstract import AbstractInterface +from sup3r.models.linear import LinearInterp from sup3r.utilities.utilities import spatial_coarsening logger = logging.getLogger(__name__) -class SurfaceSpatialMetModel(AbstractInterface): +class SurfaceSpatialMetModel(LinearInterp): """Model to spatially downscale daily-average near-surface temperature, relative humidity, and pressure @@ -43,7 +46,8 @@ class SurfaceSpatialMetModel(AbstractInterface): def __init__(self, features, s_enhance, noise_adders=None, temp_lapse=None, w_delta_temp=None, w_delta_topo=None, - pres_div=None, pres_exp=None): + pres_div=None, pres_exp=None, interp_method='LANCZOS', + fix_bias=True): """ Parameters ---------- @@ -85,6 +89,15 @@ def __init__(self, features, s_enhance, noise_adders=None, pres_div : None | float Exponential factor in the pressure scale height equation. Defaults to the cls.PRES_EXP attribute. + interp_method : str + Name of the interpolation method to use from PIL.Image.Resampling + (NEAREST, BILINEAR, BICUBIC, LANCZOS) + LANCZOS is default and has been tested to work best for + SurfaceSpatialMetModel. + fix_bias : bool + Some local bias can be introduced by the bilinear interp + lapse + rate, this flag will attempt to correct that bias by using the + low-resolution deviation from the input data """ self._features = features @@ -95,6 +108,8 @@ def __init__(self, features, s_enhance, noise_adders=None, self._w_delta_topo = w_delta_topo or self.W_DELTA_TOPO self._pres_div = pres_div or self.PRES_DIV self._pres_exp = pres_exp or self.PRES_EXP + self._fix_bias = fix_bias + self._interp_method = getattr(Image.Resampling, interp_method) if isinstance(self._noise_adders, (int, float)): self._noise_adders = [self._noise_adders] * len(self._features) @@ -103,42 +118,6 @@ def __len__(self): """Get number of model steps (match interface of MultiStepGan)""" return 1 - @classmethod - def load(cls, features, s_enhance, verbose=False, **kwargs): - """Load the GAN with its sub-networks from a previously saved-to output - directory. - - Parameters - ---------- - features : list - List of feature names that this model will operate on for both - input and output. This must match the feature axis ordering in the - array input to generate(). Typically this is a list containing: - temperature_*m, relativehumidity_*m, and pressure_*m. The list can - contain multiple instances of each variable at different heights. - relativehumidity_*m entries must have corresponding temperature_*m - entires at the same hub height. - s_enhance : int - Integer factor by which the spatial axes are to be enhanced. - verbose : bool - Flag to log information about the loaded model. - kwargs : None | dict - Optional kwargs to initialize SurfaceSpatialMetModel - - Returns - ------- - out : SurfaceSpatialMetModel - Returns an initialized SurfaceSpatialMetModel - """ - - model = cls(features, s_enhance, **kwargs) - - if verbose: - logger.info('Loading SurfaceSpatialMetModel with meta data: {}' - .format(model.meta)) - - return model - @staticmethod def _get_s_enhance(topo_lr, topo_hr): """Get the spatial enhancement factor given low-res and high-res @@ -227,8 +206,39 @@ def _get_temp_rh_ind(self, idf_rh): return idf_temp + def _fix_downscaled_bias(self, single_lr, single_hr, + method=Image.Resampling.LANCZOS): + """Fix any bias introduced by the spatial downscaling with lapse rate. + + Parameters + ---------- + single_lr : np.ndarray + Single timestep raster data with shape + (lat, lon) matching the low-resolution input data. + single_hr : np.ndarray + Single timestep downscaled raster data with shape + (lat, lon) matching the high-resolution input data. + method : Image.Resampling.LANCZOS + An Image.Resampling method (NEAREST, BILINEAR, BICUBIC, LANCZOS). + NEAREST enforces zero bias but makes slightly more spatial seams. + + Returns + ------- + single_hr : np.ndarray + Single timestep downscaled raster data with shape + (lat, lon) matching the high-resolution input data. + """ + + re_coarse = spatial_coarsening(np.expand_dims(single_hr, axis=-1), + s_enhance=self._s_enhance, + obs_axis=False)[..., 0] + bias = re_coarse - single_lr + bc = self.downscale_arr(bias, s_enhance=self._s_enhance, method=method) + single_hr -= bc + return single_hr + @staticmethod - def downscale_arr(arr, s_enhance, method=Image.Resampling.BILINEAR): + def downscale_arr(arr, s_enhance, method=Image.Resampling.LANCZOS): """Downscale a 2D array of data Image.resize() method Parameters @@ -238,9 +248,9 @@ def downscale_arr(arr, s_enhance, method=Image.Resampling.BILINEAR): (lat, lon) s_enhance : int Integer factor by which the spatial axes are to be enhanced. - method : Image.Resampling.BILINEAR + method : Image.Resampling.LANCZOS An Image.Resampling method (NEAREST, BILINEAR, BICUBIC, LANCZOS). - BILINEAR is default and has been tested to work best for + LANCZOS is default and has been tested to work best for SurfaceSpatialMetModel. """ im = Image.fromarray(arr) @@ -284,9 +294,15 @@ def downscale_temp(self, single_lr_temp, topo_lr, topo_hr): assert len(topo_hr.shape) == 2, 'Bad shape for topo_hr' lower_data = single_lr_temp.copy() + topo_lr * self._temp_lapse - hi_res_temp = self.downscale_arr(lower_data, self._s_enhance) + hi_res_temp = self.downscale_arr(lower_data, self._s_enhance, + method=self._interp_method) hi_res_temp -= topo_hr * self._temp_lapse + if self._fix_bias: + hi_res_temp = self._fix_downscaled_bias(single_lr_temp, + hi_res_temp, + method=self._interp_method) + return hi_res_temp def downscale_rh(self, single_lr_rh, single_lr_temp, single_hr_temp, @@ -336,9 +352,12 @@ def downscale_rh(self, single_lr_rh, single_lr_temp, single_hr_temp, assert len(topo_lr.shape) == 2, 'Bad shape for topo_lr' assert len(topo_hr.shape) == 2, 'Bad shape for topo_hr' - interp_rh = self.downscale_arr(single_lr_rh, self._s_enhance) - interp_temp = self.downscale_arr(single_lr_temp, self._s_enhance) - interp_topo = self.downscale_arr(topo_lr, self._s_enhance) + interp_rh = self.downscale_arr(single_lr_rh, self._s_enhance, + method=self._interp_method) + interp_temp = self.downscale_arr(single_lr_temp, self._s_enhance, + method=self._interp_method) + interp_topo = self.downscale_arr(topo_lr, self._s_enhance, + method=self._interp_method) delta_temp = single_hr_temp - interp_temp delta_topo = topo_hr - interp_topo @@ -347,6 +366,10 @@ def downscale_rh(self, single_lr_rh, single_lr_temp, single_hr_temp, + self._w_delta_temp * delta_temp + self._w_delta_topo * delta_topo) + if self._fix_bias: + hi_res_rh = self._fix_downscaled_bias(single_lr_rh, hi_res_rh, + method=self._interp_method) + return hi_res_rh def downscale_pres(self, single_lr_pres, topo_lr, topo_hr): @@ -388,21 +411,28 @@ def downscale_pres(self, single_lr_pres, topo_lr, topo_hr): warn(msg) const = 101325 * (1 - (1 - topo_lr / self._pres_div)**self._pres_exp) - single_lr_pres = single_lr_pres.copy() + const + lr_pres_adj = single_lr_pres.copy() + const - if np.min(single_lr_pres) < 0.0: + if np.min(lr_pres_adj) < 0.0: msg = ('Spatial interpolation of surface pressure ' 'resulted in negative values. Incorrectly ' 'scaled/unscaled values or incorrect units are ' - 'the most likely causes.') + 'the most likely causes. All pressure data should be ' + 'in Pascals.') logger.error(msg) raise ValueError(msg) - hi_res_pres = self.downscale_arr(single_lr_pres, self._s_enhance) + hi_res_pres = self.downscale_arr(lr_pres_adj, self._s_enhance, + method=self._interp_method) const = 101325 * (1 - (1 - topo_hr / self._pres_div)**self._pres_exp) hi_res_pres -= const + if self._fix_bias: + hi_res_pres = self._fix_downscaled_bias(single_lr_pres, + hi_res_pres, + method=self._interp_method) + if np.min(hi_res_pres) < 0.0: msg = ('Spatial interpolation of surface pressure ' 'resulted in negative values. Incorrectly ' @@ -524,25 +554,11 @@ def meta(self): 'pressure_exponent': self._pres_exp, 'training_features': self.training_features, 'output_features': self.output_features, + 'interp_method': str(self._interp_method), + 'fix_bias': self._fix_bias, 'class': self.__class__.__name__, } - @property - def training_features(self): - """Get the list of input feature names that the generative model was - trained on. - - Note that topography needs to be passed into generate() as an exogenous - data input. - """ - return self._features - - @property - def output_features(self): - """Get the list of output feature names that the generative model - outputs""" - return self._features - def train(self, true_hr_temp, true_hr_rh, true_hr_topo): """This method trains the relative humidity linear model. The temperature and surface lapse rate models are parameterizations taken diff --git a/sup3r/preprocessing/data_handling.py b/sup3r/preprocessing/data_handling.py index d73e38569..363be6de8 100644 --- a/sup3r/preprocessing/data_handling.py +++ b/sup3r/preprocessing/data_handling.py @@ -1106,7 +1106,7 @@ def preflight(self): msg = (f'sample_shape[2] ({self.sample_shape[2]}) cannot be larger ' 'than the number of time steps in the raw data ' f'({len(self.raw_time_index)}).') - if len(self.raw_time_index) >= self.sample_shape[2]: + if len(self.raw_time_index) < self.sample_shape[2]: logger.warning(msg) warnings.warn(msg) diff --git a/sup3r/utilities/loss_metrics.py b/sup3r/utilities/loss_metrics.py index 5a4034a25..312397f8c 100644 --- a/sup3r/utilities/loss_metrics.py +++ b/sup3r/utilities/loss_metrics.py @@ -1,6 +1,6 @@ """Loss metrics for Sup3r""" -from tensorflow.keras.losses import MeanSquaredError +from tensorflow.keras.losses import MeanSquaredError, MeanAbsoluteError import tensorflow as tf @@ -171,3 +171,39 @@ def __call__(self, x1, x2): x1_coarse = tf.reduce_mean(x1, axis=(1, 2)) x2_coarse = tf.reduce_mean(x2, axis=(1, 2)) return self.MSE_LOSS(x1_coarse, x2_coarse) + + +class TemporalExtremesLoss(tf.keras.losses.Loss): + """Loss class that encourages accuracy of the min/max values in the + timeseries""" + + MAE_LOSS = MeanAbsoluteError() + + def __call__(self, x1, x2): + """Custom content loss that encourages temporal min/max accuracy + + Parameters + ---------- + x1 : tf.tensor + synthetic generator output + (n_observations, spatial_1, spatial_2, temporal, features) + x2 : tf.tensor + high resolution data + (n_observations, spatial_1, spatial_2, temporal, features) + + Returns + ------- + tf.tensor + 0D tensor with loss value + """ + x1_min = tf.reduce_min(x1, axis=3) + x2_min = tf.reduce_min(x2, axis=3) + + x1_max = tf.reduce_max(x1, axis=3) + x2_max = tf.reduce_max(x2, axis=3) + + mae = self.MAE_LOSS(x1, x2) + mae_min = self.MAE_LOSS(x1_min, x2_min) + mae_max = self.MAE_LOSS(x1_max, x2_max) + + return mae + mae_min + mae_max diff --git a/sup3r/version.py b/sup3r/version.py index eb9b139c7..ae5c54150 100644 --- a/sup3r/version.py +++ b/sup3r/version.py @@ -1,4 +1,4 @@ # -*- coding: utf-8 -*- """SUP3R Version""" -__version__ = '0.0.8' +__version__ = '0.0.9' diff --git a/tests/test_custom_loss.py b/tests/test_custom_loss.py index e334eaee9..124d630a9 100644 --- a/tests/test_custom_loss.py +++ b/tests/test_custom_loss.py @@ -3,7 +3,8 @@ import numpy as np import tensorflow as tf -from sup3r.utilities.loss_metrics import MmdMseLoss, CoarseMseLoss +from sup3r.utilities.loss_metrics import (MmdMseLoss, CoarseMseLoss, + TemporalExtremesLoss) def test_mmd_loss(): @@ -52,3 +53,26 @@ def test_coarse_mse_loss(): assert mse.numpy().size == 1 assert coarse_mse.numpy().size == 1 assert mse.numpy() > 10 * coarse_mse.numpy() + + +def test_tex_loss(): + """Test custom TemporalExtremesLoss function that looks at min/max values + in the timeseries.""" + loss_obj = TemporalExtremesLoss() + + x = np.zeros((1, 1, 1, 72, 1)) + y = np.zeros((1, 1, 1, 72, 1)) + + # loss should be dominated by special min/max values + x[..., 24, 0] = 20 + y[..., 25, 0] = 25 + loss = loss_obj(x, y) + assert loss.numpy() > 5 + assert loss.numpy() < 6 + + # loss should be dominated by special min/max values + x[..., 24, 0] = -10 + y[..., 25, 0] = -15 + loss = loss_obj(x, y) + assert loss.numpy() > 5 + assert loss.numpy() < 6 diff --git a/tests/test_multi_step.py b/tests/test_multi_step.py index 81caa03f6..6c9901e71 100644 --- a/tests/test_multi_step.py +++ b/tests/test_multi_step.py @@ -95,7 +95,8 @@ def test_multi_step_norm(norm_option): out1 = model1.generate(x) out2 = model2.generate(out1) out3 = model3.generate(out2) - assert np.allclose(out, out3, atol=5e-4) + err = np.abs(out - out3).mean() / np.abs(out.mean()) + assert err < 1e4 def test_spatial_then_temporal_gan(): diff --git a/tests/test_surface_model.py b/tests/test_surface_model.py index 18ded9e9c..0373394e2 100644 --- a/tests/test_surface_model.py +++ b/tests/test_surface_model.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Test the temperature and relative humidity scaling relationships of the SurfaceSpatialMetModel""" +import json import os import tempfile import pytest @@ -53,17 +54,24 @@ def test_surface_model(s_enhance=5): low_res, true_hi_res, topo_lr, topo_hr = get_inputs(s_enhance) - model = SurfaceSpatialMetModel.load(features=FEATURES, s_enhance=s_enhance) + kwargs = {'meta': {'features': FEATURES, 's_enhance': s_enhance}} + with tempfile.TemporaryDirectory() as td: + fp_params = os.path.join(td, 'model_params.json') + with open(fp_params, 'w') as f: + json.dump(kwargs, f) + + model = SurfaceSpatialMetModel.load(model_dir=td) + hi_res = model.generate(low_res, exogenous_data=[topo_lr, topo_hr]) diff = true_hi_res - hi_res # high res temperature should have very low bias and MAE < 1C - assert np.abs(diff[..., 0].mean()) < 1e-6 + assert np.abs(diff[..., 0].mean()) < 1e-4 assert np.abs(diff[..., 0]).mean() < 5 # high res relative humidity should have very low bias and MAE < 3% - assert np.abs(diff[..., 1].mean()) < 1e-6 + assert np.abs(diff[..., 1].mean()) < 1e-4 assert np.abs(diff[..., 1]).mean() < 2 # high res pressure should have very low bias and MAE < 200 Pa @@ -116,14 +124,21 @@ def test_multi_step_surface(s_enhance=2, t_enhance=2): t_enhance=t_enhance) with tempfile.TemporaryDirectory() as td: - fp = os.path.join(td, 'model') - model.save(fp) + temporal_dir = os.path.join(td, 'model') + model.save(temporal_dir) + + surface_model_kwargs = {'meta': {'features': FEATURES, + 's_enhance': s_enhance}} + + surface_dir = os.path.join(td, 'surface/') + os.makedirs(surface_dir) + fp_params = os.path.join(surface_dir, 'model_params.json') + with open(fp_params, 'w') as f: + json.dump(surface_model_kwargs, f) - surface_model_kwargs = {'features': FEATURES, 's_enhance': s_enhance} - temporal_model_kwargs = {'model_dirs': fp} ms_model = MultiStepSurfaceMetGan.load( - surface_model_kwargs=surface_model_kwargs, - temporal_model_kwargs=temporal_model_kwargs) + surface_model_kwargs={'model_dir': surface_dir}, + temporal_model_kwargs={'model_dirs': temporal_dir}) for model in ms_model.models: assert isinstance(model.s_enhance, int)