diff --git a/acme_diags/driver/default_diags/streamflow_model_vs_model.cfg b/acme_diags/driver/default_diags/streamflow_model_vs_model.cfg new file mode 100644 index 000000000000..72c88af768c5 --- /dev/null +++ b/acme_diags/driver/default_diags/streamflow_model_vs_model.cfg @@ -0,0 +1,6 @@ +[#] +sets = ["streamflow"] +case_id = "seasonality" +variables = ["RIVER_DISCHARGE_OVER_LAND_LIQ"] +regions = ["global"] +seasons = ["ANN"] diff --git a/acme_diags/driver/default_diags/streamflow_model_vs_obs.cfg b/acme_diags/driver/default_diags/streamflow_model_vs_obs.cfg new file mode 100644 index 000000000000..f45d60e09a05 --- /dev/null +++ b/acme_diags/driver/default_diags/streamflow_model_vs_obs.cfg @@ -0,0 +1,8 @@ +[#] +sets = ["streamflow"] +case_id = "seasonality" +variables = ["RIVER_DISCHARGE_OVER_LAND_LIQ"] +ref_name = "GSIM" +reference_name = "GSIM monthly streamflow (1986-1995)" +regions = ["global"] +seasons = ["ANN"] diff --git a/acme_diags/driver/streamflow_driver.py b/acme_diags/driver/streamflow_driver.py new file mode 100644 index 000000000000..a91920d52a30 --- /dev/null +++ b/acme_diags/driver/streamflow_driver.py @@ -0,0 +1,316 @@ +from __future__ import print_function + +import csv +import numpy +import scipy.io +import time +import cdutil +from acme_diags.driver import utils +from acme_diags.plot.cartopy.streamflow_plot import plot_seasonality + + +def get_drainage_area_error(radius, resolution, lon_ref, lat_ref, area_upstream, area_ref): + k_bound = len(range(-radius, radius + 1)) + k_bound *= k_bound + area_test = numpy.zeros((k_bound, 1)) + error_test = numpy.zeros((k_bound, 1)) + lat_lon_test = numpy.zeros((k_bound, 2)) + k = 0 + for i in range(-radius, radius + 1): + for j in range(-radius, radius + 1): + x = ((lon_ref + j * resolution) - (-180 + resolution / 2)) / resolution + y = ((lat_ref + i * resolution) - (-90 + resolution / 2)) / resolution + area_test[k] = area_upstream[int(x), int(y)] / 1000000 + error_test[k] = abs(area_test[k] - area_ref) / area_ref + lat_lon_test[k, 0] = lat_ref + i * resolution + lat_lon_test[k, 1] = lon_ref + j * resolution + k += 1 + # The id of the center grid in the searching area + center_id = (max(error_test.shape) - 1) / 2 + + lat_lon_ref = [lat_ref, lon_ref] + drainage_area_error = error_test[int(center_id)] + return drainage_area_error, lat_lon_ref + + +def get_seasonality(monthly): + # See https://agupubs.onlinelibrary.wiley.com/doi/epdf/10.1029/2018MS001603 Equations 1 and 2 + num_years = monthly.shape[1] + p_k = numpy.zeros((12, 1)) + # The total streamflow for each year (sum of Q_ij in the denominator of Equation 1, for all j) + # 1 x num_years + total_streamflow = numpy.sum(monthly, axis=0) + for month in range(12): + # The streamflow for this month in each year (Q_ij in the numerator of Equation 1, for all j) + # 1 x num_years + streamflow_month_all_years = monthly[month, :] + # Proportion that this month contributes to streamflow that year. + # 1 x num_years + streamflow_proportion = streamflow_month_all_years / total_streamflow + # The sum is the sum over j in Equation 1. + # Dividing the sum of proportions by num_years gives the *average* proportion of annual streamflow during + # this month. + # Multiplying by 12 makes it so that Pk_i (`p_k[month]`) will be 1 if all months have equal streamflow and + # 12 if all streamflow occurs in one month. + # These steps produce the 12/n factor in Equation 1. + p_k[month] = numpy.nansum(streamflow_proportion) * 12 / num_years + # From Equation 2 + seasonality_index = numpy.max(p_k) + # `p_k == numpy.max(p_k)` produces a Boolean matrix, True if the value (i.e., streamflow) is the max value. + # `np.where(p_k == numpy.max(p_k))` produces the indices (i.e., months) where the max value is reached. + peak_month = numpy.where(p_k == numpy.max(p_k))[0] + # If more than one month has peak streamflow, simply define the peak month as the first one of the peak months. + peak_month = peak_month[0] + return seasonality_index, peak_month + + +def run_diag_seasonality(parameter): + # Set path to the gauge metadata + with open(parameter.gauges_path) as gauges_path: + gauges_list = list(csv.reader(gauges_path)) + # Remove headers + gauges_list.pop(0) + gauges = numpy.array(gauges_list) + if parameter.print_statements: + print('gauges.shape={}'.format(gauges.shape)) + + variables = parameter.variables + for var in variables: + + if not parameter.ref_mat_file: + ref_data = utils.dataset.Dataset(parameter, ref=True) + ref_data_ts = ref_data.get_timeseries_variable(var) + var_array = ref_data_ts(cdutil.region.domain(latitude=(-90., 90, 'ccb'))) + if parameter.print_statements: + print('ref var original dimensions={}'.format(var_array.shape)) + var_transposed = numpy.transpose(var_array, (2, 1, 0)) + if parameter.print_statements: + print('ref var transposed dimensions={}'.format(var_transposed.shape)) + ref_array = var_transposed + else: + # Load the observed streamflow dataset (GSIM) + # the data has been reorganized to a 1380 * 30961 matrix. 1380 is the month + # number from 1901.1 to 2015.12. 30961 include two columns for year and month plus + # streamflow at 30959 gauge locations reported by GSIM + ref_mat = scipy.io.loadmat(parameter.ref_mat_file) + ref_array = ref_mat['GSIM'] + if parameter.print_statements: + # GSIM: 1380 x 30961 + # wrmflow: 720 x 360 x 360 + print('ref_array.shape={}'.format(ref_array.shape)) + + # Load E3SM simulated streamflow dataset + if not parameter.test_mat_file: + # `Dataset` will take the time slice from test_start_yr to test_end_yr + test_data = utils.dataset.Dataset(parameter, test=True) + test_data_ts = test_data.get_timeseries_variable(var) + var_array = test_data_ts(cdutil.region.domain(latitude=(-90., 90, 'ccb'))) + if parameter.print_statements: + print('test var original dimensions={}'.format(var_array.shape)) + var_transposed = numpy.transpose(var_array, (2, 1, 0)) + if parameter.print_statements: + print('test var transposed dimensions={}'.format(var_transposed.shape)) + test_array = var_transposed + areatotal2 = test_data.get_static_variable('areatotal2', var) + area_upstream = numpy.transpose(areatotal2, (1, 0)) + if parameter.print_statements: + print('area_upstream dimensions={}'.format(area_upstream.shape)) + else: + data_mat = scipy.io.loadmat(parameter.test_mat_file) + if 'E3SMflow' in data_mat.keys(): + # `edison` file uses this block + e3sm_flow = data_mat['E3SMflow'] + if parameter.print_statements: + print('e3sm_flow = data_mat["E3SMflow"]') + else: + # `test` file uses this block + e3sm_flow = data_mat + if parameter.print_statements: + print('e3sm_flow = data_mat') + try: + if e3sm_flow['uparea'].shape == (1, 1): + # `edison` file uses this block + area_upstream = e3sm_flow['uparea'][0][0] + if parameter.print_statements: + print('e3sm_flow["uparea"] was indexed into for later use') + else: + area_upstream = e3sm_flow['uparea'] + if parameter.print_statements: + print('e3sm_flow["uparea"] will be used') + except KeyError: + # `test` file uses this block + area_upstream = None + if parameter.print_statements: + print('WARNING: uparea not found and will thus not be used') + if e3sm_flow['wrmflow'].shape == (1, 1): + # `edison` file uses this block + test_array = e3sm_flow['wrmflow'][0][0] + if parameter.print_statements: + print('e3sm_flow["wrmflow"] was indexed into for later use') + else: + # `test` file uses this block + test_array = e3sm_flow['wrmflow'] + if parameter.print_statements: + print('e3sm_flow["wrmflow"] will be used') + if parameter.print_statements: + print('test_array.shape={}'.format(test_array.shape)) + + # Resolution of MOSART output + resolution = 0.5 + # Search radius (number of grids around the center point) + radius = 1 + bins = numpy.floor(gauges[:, 7:9].astype(numpy.float64) / resolution) + # Move the ref lat lon to grid center + lat_lon = (bins + 0.5) * resolution + if parameter.print_statements: + print('lat_lon.shape={}'.format(lat_lon.shape)) + + # Define the export matrix + # Annual mean of test, annual mean of ref, error for area, lat, lon + export = numpy.zeros((lat_lon.shape[0], 9)) + if parameter.print_statements: + print('export.shape={}'.format(export.shape)) + t0 = time.time() + for i in range(lat_lon.shape[0]): + if parameter.print_statements and (i % 1000 == 0): + print('On gauge #{}'.format(i)) + if parameter.max_num_gauges and i > parameter.max_num_gauges: + break + lat_ref = lat_lon[i, 1] + lon_ref = lat_lon[i, 0] + # Estimated drainage area (km^2) from ref + area_ref = gauges[i, 13].astype(numpy.float64) + + if area_upstream is not None: + drainage_area_error, lat_lon_ref = get_drainage_area_error( + radius, resolution, lon_ref, lat_ref, area_upstream, area_ref) + else: + # Use the center location + lat_lon_ref = [lat_ref, lon_ref] + if parameter.ref_mat_file: + origin_id = gauges[i, 1].astype(numpy.int64) + # Column 0 -- year + # Column 1 -- month + # Column origin_id + 1 -- the ref streamflow from gauge with the corresponding origin_id + extracted = ref_array[:, [0, 1, origin_id + 1]] + monthly_mean = numpy.zeros((12,1)) + numpy.nan + for month in range(12): + # Add 1 to month to account for the months being 1-indexed + if sum(extracted[:, 1] == month + 1) > 0: + # `extracted[:,1]`: for all x, examine `extracted[x,1]` + # `extracted[:,1] == m`: Boolean array where 0 means the item in position [x,1] is NOT m, + # and 1 means it is m + # Example: + # a = [[1,2,3,4], + # [5,6,7,8], + # [1,2,3,4]] + # a[:,1]: [[2], + # [6], + # [2]] + # a[:,1] == 2: [[1], # False is 0, True is 1 + # [0], + # [1]] + # a[a[:,1] == 2, 2]: [[3], + # [3]] + monthly_mean[month] = numpy.nanmean(extracted[extracted[:, 1] == month + 1, 2]) + # This is ref annual mean streamflow + annual_mean_ref = numpy.mean(monthly_mean) + if parameter.ref_mat_file and numpy.isnan(annual_mean_ref): + # All elements of row i will be nan + export[i,:] = numpy.nan + else: + if parameter.ref_mat_file: + # Reshape extracted[:,2] into a 12 x ? matrix; -1 means to calculate the size of the missing dimension. + mmat = numpy.reshape(extracted[:, 2], (12, -1)) + mmat_id = numpy.sum(mmat, axis=0).transpose() + if numpy.sum(~numpy.isnan(mmat_id), axis=0) > 0: + # There's at least one year of record + monthly = mmat[:, ~numpy.isnan(mmat_id)] + else: + monthly = monthly_mean + seasonality_index_ref, peak_month_ref = get_seasonality(monthly) + else: + ref_lon = int(1 + (lat_lon_ref[1] - (-180 + resolution / 2)) / resolution) + ref_lat = int(1 + (lat_lon_ref[0] - (-90 + resolution / 2)) / resolution) + ref = numpy.squeeze(ref_array[ref_lon, ref_lat, :]) + mmat = numpy.reshape(ref, (12, -1)) + monthly_mean_ref = numpy.nanmean(mmat, axis=1) + annual_mean_ref = numpy.mean(monthly_mean_ref) + if numpy.isnan(annual_mean_ref) == 1: + # The identified grid is in the ocean + monthly = numpy.ones((12, 1)) + else: + monthly = mmat + seasonality_index_ref, peak_month_ref = get_seasonality(monthly) + + test_lon = int(1 + (lat_lon_ref[1] - (-180 + resolution / 2)) / resolution) + test_lat = int(1 + (lat_lon_ref[0] - (-90 + resolution / 2)) / resolution) + test = numpy.squeeze(test_array[test_lon, test_lat, :]) + mmat = numpy.reshape(test, (12, -1)) + monthly_mean_test = numpy.nanmean(mmat, axis=1) + annual_mean_test = numpy.mean(monthly_mean_test) + if numpy.isnan(annual_mean_test) == 1: + # The identified grid is in the ocean + monthly = numpy.ones((12, 1)) + else: + monthly = mmat + seasonality_index_test, peak_month_test = get_seasonality(monthly) + + export[i, 0] = annual_mean_ref + export[i, 1] = annual_mean_test + if area_upstream is not None: + export[i, 2] = drainage_area_error * 100 # From fraction to percentage of the drainage area bias + export[i, 3] = seasonality_index_ref # Seasonality index of ref + export[i, 4] = peak_month_ref # Max flow month of ref + export[i, 5] = seasonality_index_test # Seasonality index of test + export[i, 6] = peak_month_test # Max flow month of test + export[i, 7:9] = lat_lon_ref # latlon of ref + t1 = time.time() + if parameter.print_statements: + print('Loop time={}s={}m'.format(t1 - t0, (t1 - t0) / 60)) + # Remove the gauges with nan flow + # `export[:,0]` => get first column of export + # `numpy.isnan(export[:,0])` => Boolean column, True if value in export[x,0] is nan + # `export[numpy.isnan(export[:,0]),:]` => rows of `export` where the Boolean column was True + # Gauges will thus only be plotted if they have a non-nan value for both test and ref. + if parameter.print_statements: + print('export.shape before removing ref nan means={}'.format(export.shape)) + export = export[~numpy.isnan(export[:, 0]), :] + if parameter.print_statements: + print('export.shape before removing test nan means={}'.format(export.shape)) + export = export[~numpy.isnan(export[:, 1]), :] + if parameter.print_statements: + print('export.shape after both nan removals={}'.format(export.shape)) + + if area_upstream is not None: + # Set the max area error (percent) for all plots + max_area_error = 20 + # `export[:,2]` gives the third column of `export` + # `export[:,2]<=max_area_error` gives a Boolean array, + # `True` if the value in the third column of `export` is `<= max_area_error` + # `export[export[:,2]<=max_area_error,:]` is `export` with only the rows where the above is `True`. + export = export[export[:, 2] <= max_area_error, :] + if parameter.print_statements: + print('export.shape after max_area_error cut={}'.format(export.shape)) + + if parameter.print_statements: + print('Variable: {}'.format(var)) + parameter.var_id = '{}-seasonality'.format(var) + + # This will be the title of the plot. + parameter.main_title = 'Seasonality' + parameter.viewer_descr[var] = parameter.main_title + parameter.output_file = 'seasonality' + + # Plot + # Plot original ref and test, not regridded versions. + plot_seasonality(export, parameter) + return parameter + + +def run_diag(parameter): + # TODO: Add bias plot and scatterplot in future pull requests. + if parameter.plot_type == 'seasonality': + return run_diag_seasonality(parameter) + else: + raise Exception('Invalid plot_type={}'.format(parameter.plot_type)) diff --git a/acme_diags/driver/utils/dataset.py b/acme_diags/driver/utils/dataset.py index 43f4ae15b89a..011a01df9b31 100644 --- a/acme_diags/driver/utils/dataset.py +++ b/acme_diags/driver/utils/dataset.py @@ -156,6 +156,18 @@ def get_climo_variable(self, var, season, extra_vars=[], *args, **kwargs): return variables[0] if len(variables) == 1 else variables + def get_static_variable(self, static_var, primary_var): + if self.ref: + # Get the reference variable from timeseries files. + data_path = self.parameters.reference_data_path + elif self.test: + # Get the test variable from timeseries files. + data_path = self.parameters.test_data_path + file_path = self._get_timeseries_file_path(primary_var, data_path) + with cdms2.open(file_path) as f: + return f(static_var) + + def is_timeseries(self): """ Return True if this dataset is for timeseries data. diff --git a/acme_diags/parameter/__init__.py b/acme_diags/parameter/__init__.py index be8efd0ca37e..7324c5008f78 100644 --- a/acme_diags/parameter/__init__.py +++ b/acme_diags/parameter/__init__.py @@ -4,6 +4,7 @@ from .area_mean_time_series_parameter import AreaMeanTimeSeriesParameter from .enso_diags_parameter import EnsoDiagsParameter from .qbo_parameter import QboParameter +from .streamflow_parameter import StreamflowParameter SET_TO_PARAMETERS = { 'zonal_mean_xy': CoreParameter, @@ -14,5 +15,6 @@ 'cosp_histogram': CoreParameter, 'area_mean_time_series': AreaMeanTimeSeriesParameter, 'enso_diags': EnsoDiagsParameter, - 'qbo': QboParameter + 'qbo': QboParameter, + 'streamflow': StreamflowParameter } diff --git a/acme_diags/parameter/core_parameter.py b/acme_diags/parameter/core_parameter.py index faa84e5d92a4..684216b8d39d 100644 --- a/acme_diags/parameter/core_parameter.py +++ b/acme_diags/parameter/core_parameter.py @@ -16,7 +16,7 @@ def __init__(self): self.sets = ['zonal_mean_xy', 'zonal_mean_2d', 'meridional_mean_2d', 'lat_lon', 'polar', 'area_mean_time_series', 'cosp_histogram', - 'enso_diags', 'qbo'] + 'enso_diags', 'qbo', 'streamflow'] self.dataset = '' self.run_type = 'model_vs_obs' self.variables = [] diff --git a/acme_diags/parameter/streamflow_parameter.py b/acme_diags/parameter/streamflow_parameter.py new file mode 100644 index 000000000000..8b449c058d30 --- /dev/null +++ b/acme_diags/parameter/streamflow_parameter.py @@ -0,0 +1,62 @@ +from .core_parameter import CoreParameter + + +class StreamflowParameter(CoreParameter): + def __init__(self): + super(StreamflowParameter, self).__init__() + self.max_num_gauges = None + self.test_mat_file = None + self.ref_mat_file = None + self.plot_type = 'seasonality' + self.print_statements = False + self.ref_timeseries_input = True + self.test_timeseries_input = True + self.granulate.remove('seasons') + + def check_values(self): + if not hasattr(self, 'gauges_path'): + msg = 'gauges_path must be specified' + raise RuntimeError(msg) + + # TODO: In future pull requests add 'scatter', 'bias' + valid_plot_types = ['seasonality'] + if self.plot_type not in valid_plot_types: + msg = 'plot_type={} not in {}'.format( + self.plot_type, valid_plot_types) + raise RuntimeError(msg) + + test_ref_start_yr_both_set = hasattr(self, 'test_start_yr') and hasattr(self, 'ref_start_yr') + if hasattr(self, 'start_yr'): + # Use `start_yr` as a default value for other parameters. + if not hasattr(self, 'test_start_yr'): + self.test_start_yr = self.start_yr + if not hasattr(self, 'ref_start_yr'): + self.ref_start_yr = self.start_yr + elif test_ref_start_yr_both_set and self.test_start_yr == self.ref_start_yr: + # Derive the value of self.start_yr + self.start_yr = self.test_start_yr + + test_ref_end_yr_both_set = hasattr(self, 'test_end_yr') and hasattr(self, 'ref_end_yr') + if hasattr(self, 'end_yr'): + # Use `end_yr` as a default value for other parameters. + if not hasattr(self, 'test_end_yr'): + self.test_end_yr = self.end_yr + if not hasattr(self, 'ref_end_yr'): + self.ref_end_yr = self.end_yr + elif test_ref_end_yr_both_set and self.test_end_yr == self.ref_end_yr: + # Derive the value of self.end_yr + self.end_yr = self.test_end_yr + + if hasattr(self, 'start_yr'): + # We need to re-evaluate this variable, since these attributes could have been set. + test_ref_end_yr_both_set = hasattr(self, 'test_end_yr') and hasattr(self, 'ref_end_yr') + if not (hasattr(self, 'end_yr') or test_ref_end_yr_both_set): + msg = "To use 'start_yr' you need to also define 'end_yr' or both 'test_end_yr' and 'ref_end_yr'." + raise RuntimeError(msg) + + if hasattr(self, 'end_yr'): + # We need to re-evaluate this variable, since these attributes could have been set. + test_ref_start_yr_both_set = hasattr(self, 'test_start_yr') and hasattr(self, 'ref_start_yr') + if not (hasattr(self, 'start_yr') or test_ref_start_yr_both_set): + msg = "To use 'end_yr' you need to also define 'start_yr' or both 'test_start_yr' and 'ref_start_yr'." + raise RuntimeError(msg) diff --git a/acme_diags/parser/__init__.py b/acme_diags/parser/__init__.py index 3d4999edfa61..528256a29881 100644 --- a/acme_diags/parser/__init__.py +++ b/acme_diags/parser/__init__.py @@ -4,6 +4,7 @@ from .area_mean_time_series_parser import AreaMeanTimeSeriesParser from .enso_diags_parser import EnsoDiagsParser from .qbo_parser import QboParser +from .streamflow_parser import StreamflowParser SET_TO_PARSER = { 'zonal_mean_xy': CoreParser, @@ -14,6 +15,7 @@ 'cosp_histogram': CoreParser, 'area_mean_time_series': AreaMeanTimeSeriesParser, 'enso_diags': EnsoDiagsParser, - 'qbo': QboParser + 'qbo': QboParser, + 'streamflow': StreamflowParser } diff --git a/acme_diags/parser/streamflow_parser.py b/acme_diags/parser/streamflow_parser.py new file mode 100644 index 000000000000..da5f85fd374f --- /dev/null +++ b/acme_diags/parser/streamflow_parser.py @@ -0,0 +1,67 @@ +from .core_parser import CoreParser +from acme_diags.parameter.streamflow_parameter import StreamflowParameter + + +class StreamflowParser(CoreParser): + def __init__(self, *args, **kwargs): + if 'parameter_cls' in kwargs: + super().__init__(*args, **kwargs) + else: + super().__init__(parameter_cls=StreamflowParameter, *args, **kwargs) + + + def load_default_args(self, files=[]): + # This has '-p' and '--parameter' reserved. + super().load_default_args(files) + + self.add_argument( + '--gauges_path', + dest='gauges_path', + help='The file containing the gauge data.', + action='store_const', + const=True, + required=False) + + self.add_argument( + '--max_num_gauges', + dest='max_num_gauges', + help='The maximum number of gauges that should be processed.', + action='store_const', + const=True, + required=False) + + self.add_argument( + '--print_statements', + dest='print_statements', + help='Print information useful for debugging.', + action='store_const', + const=True, + required=False) + + self.add_argument( + '--ref_timeseries_input', + dest='ref_timeseries_input', + help='The input reference data are timeseries files.', + action='store_const', + const=True, + required=False) + + self.add_argument( + '--test_timeseries_input', + dest='test_timeseries_input', + help='The input test data are timeseries files.', + action='store_const', + const=True, + required=False) + + self.add_argument( + '--start_yr', + dest='start_yr', + help="Start year for the timeseries files.", + required=False) + + self.add_argument( + '--end_yr', + dest='end_yr', + help="End year for the timeseries files.", + required=False) diff --git a/acme_diags/plot/cartopy/streamflow_plot.py b/acme_diags/plot/cartopy/streamflow_plot.py new file mode 100644 index 000000000000..3d8f8748d9b1 --- /dev/null +++ b/acme_diags/plot/cartopy/streamflow_plot.py @@ -0,0 +1,283 @@ +from __future__ import print_function + +import matplotlib +import matplotlib.lines as lines +import numpy as np +import os + +matplotlib.use('Agg') +import matplotlib.pyplot as plt +import matplotlib.colors as colors +import cartopy.crs as ccrs +import cartopy.feature as cfeature +from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter +import cdutil +from acme_diags.derivations.default_regions import regions_specs +from acme_diags.driver.utils.general import get_output_dir + +plotTitle = {'fontsize': 11.5} +plotSideTitle = {'fontsize': 9.5} + +# Position and sizes of subplot axes in page coordinates (0 to 1) +panel = [(0.1200, 0.5500, 0.7200, 0.3000), + (0.1200, 0.1300, 0.7200, 0.3000), + ] + +# Border padding relative to subplot axes for saving individual panels +# (left, bottom, right, top) in page coordinates +border = (-0.14, -0.06, 0.04, 0.08) + + +def add_cyclic(var): + lon = var.getLongitude() + return var(longitude=(lon[0], lon[0] + 360.0, 'coe')) + + +def get_ax_size(fig, ax): + bbox = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted()) + width, height = bbox.width, bbox.height + width *= fig.dpi + height *= fig.dpi + return width, height + + +def determine_tick_step(degrees_covered): + if degrees_covered > 180: + return 60 + if degrees_covered > 60: + return 30 + elif degrees_covered > 20: + return 10 + else: + return 1 + + +def plot_panel_seasonality(plot_type, fig, proj, export, color_list, parameter): + if plot_type == 'test': + panel_index = 0 + seasonality_index_export_index = 5 + peak_month_export_index = 6 + title = (None, parameter.test_title, None) + elif plot_type == 'ref': + panel_index = 1 + seasonality_index_export_index = 3 + peak_month_export_index = 4 + title = (None, parameter.reference_title, None) + else: + raise Exception('Invalid plot_type={}'.format(plot_type)) + + # Plot of streamflow gauges. Color -> peak month, marker size -> seasonality index. + + # Position and sizes of subplot axes in page coordinates (0 to 1) + ax = fig.add_axes(panel[panel_index], projection=proj) + region_str = parameter.regions[0] + region = regions_specs[region_str] + if 'domain' in region.keys(): + # Get domain to plot + domain = region['domain'] + else: + # Assume global domain + domain = cdutil.region.domain(latitude=(-90., 90, 'ccb')) + kargs = domain.components()[0].kargs + # lon_west, lon_east, lat_south, lat_north = (0, 360, -90, 90) + lon_west, lon_east, lat_south, lat_north = (-180, 180, -90, 90) + if 'longitude' in kargs: + lon_west, lon_east, _ = kargs['longitude'] + if 'latitude' in kargs: + lat_south, lat_north, _ = kargs['latitude'] + lon_covered = lon_east - lon_west + lon_step = determine_tick_step(lon_covered) + xticks = np.arange(lon_west, lon_east, lon_step) + # Subtract 0.50 to get 0 W to show up on the right side of the plot. + # If less than 0.50 is subtracted, then 0 W will overlap 0 E on the left side of the plot. + # If a number is added, then the value won't show up at all. + xticks = np.append(xticks, lon_east - 0.50) + lat_covered = lat_north - lat_south + lat_step = determine_tick_step(lat_covered) + yticks = np.arange(lat_south, lat_north, lat_step) + yticks = np.append(yticks, lat_north) + ax.set_extent([lon_west, lon_east, lat_south, lat_north], crs=proj) + proj_function = ccrs.PlateCarree + + # Stream gauges + si_2 = 2 + si_4 = 3 + si_6 = 4 + si_large = 5 + # `export` is the array of gauges. Each gauge has multiple fields -- e.g., lat is index 7 + for gauge in export: + lat = gauge[7] + lon = gauge[8] + seasonality_index = gauge[seasonality_index_export_index] + if seasonality_index < 2: + markersize = si_2 + elif seasonality_index < 4: + markersize = si_4 + elif seasonality_index < 6: + markersize = si_6 + elif seasonality_index <= 12: + markersize = si_large + else: + raise Exception('Invalid seasonality index={}'.format(seasonality_index)) + if seasonality_index == 1: + color = 'black' + else: + peak_month = int(gauge[peak_month_export_index]) + color = color_list[peak_month] + # https://scitools.org.uk/iris/docs/v1.9.2/examples/General/projections_and_annotations.html + # Place a single marker point for each gauge. + plt.plot(lon, lat, marker='o', color=color, markersize=markersize, transform=proj_function()) + # NOTE: the "plt.annotate call" does not have a "transform=" keyword, + # so for this one we transform the coordinates with a Cartopy call. + at_x, at_y = ax.projection.transform_point(lon, lat, src_crs=proj_function()) + # https://matplotlib.org/3.1.1/gallery/text_labels_and_annotations/custom_legends.html + legend_elements = [lines.Line2D([0], [0], marker='o', color='w', label='1 <= SI < 2', + markerfacecolor='black', markersize=si_2), + lines.Line2D([0], [0], marker='o', color='w', label='2 <= SI < 4', + markerfacecolor='black', markersize=si_4), + lines.Line2D([0], [0], marker='o', color='w', label='4 <= SI < 6', + markerfacecolor='black', markersize=si_6), + lines.Line2D([0], [0], marker='o', color='w', label='6 <= SI <= 12', + markerfacecolor='black', markersize=si_large) + ] + seasonality_legend_title = 'Seasonality (SI)' + plt.legend(handles=legend_elements, title=seasonality_legend_title, prop={'size': 8}) + + # Full world would be aspect 360/(2*180) = 1 + ax.set_aspect((lon_east - lon_west) / (2 * (lat_north - lat_south))) + ax.coastlines(lw=0.3) + ax.add_feature(cfeature.RIVERS) + if title[0] is not None: + ax.set_title(title[0], loc='left', fontdict=plotSideTitle) + if title[1] is not None: + ax.set_title(title[1], fontdict=plotTitle) + if title[2] is not None: + ax.set_title(title[2], loc='right', fontdict=plotSideTitle) + ax.set_xticks(xticks, crs=proj_function()) + ax.set_yticks(yticks, crs=proj_function()) + lon_formatter = LongitudeFormatter( + zero_direction_label=True, number_format='.0f') + lat_formatter = LatitudeFormatter() + ax.xaxis.set_major_formatter(lon_formatter) + ax.yaxis.set_major_formatter(lat_formatter) + ax.tick_params(labelsize=8.0, direction='out', width=1) + ax.xaxis.set_ticks_position('bottom') + ax.yaxis.set_ticks_position('left') + + # Color bar + cbax = fig.add_axes( + (panel[panel_index][0] + 0.7535, panel[panel_index][1] + 0.0515, 0.0326, 0.1792)) + # https://matplotlib.org/tutorials/colors/colorbar_only.html + num_colors = len(color_list) + if parameter.print_statements: + print('num_colors={}'.format(num_colors)) + cmap = colors.ListedColormap(color_list) + cbar_label = 'Peak month' + + bounds = list(range(num_colors)) + # Set ticks to be in between the bounds + ticks = list(map(lambda bound: bound + 0.5, bounds)) + # Add one more bound at the bottom of the colorbar. + # `bounds` should be one longer than `ticks`. + bounds += [bounds[-1] + 1] + if parameter.print_statements: + print('bounds={}'.format(bounds)) + norm = colors.BoundaryNorm(bounds, cmap.N) + cbar = fig.colorbar( + matplotlib.cm.ScalarMappable(cmap=cmap, norm=norm), + cax=cbax, + boundaries=bounds, + ticks=ticks, + spacing='uniform', + orientation='vertical', + label=cbar_label, + ) + # https://matplotlib.org/3.1.1/gallery/ticks_and_spines/colorbar_tick_labelling_demo.html + months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec'] + cbar.ax.set_yticklabels(months) + cbar.ax.invert_yaxis() + + w, h = get_ax_size(fig, cbax) + + cbar.ax.tick_params(labelsize=9.0, length=0) + + +def plot_seasonality(export, parameter): + if parameter.backend not in ['cartopy', 'mpl', 'matplotlib']: + return + + # Create figure, projection + fig = plt.figure(figsize=parameter.figsize, dpi=parameter.dpi) + proj = ccrs.PlateCarree(central_longitude=0) + + if parameter.test_title == '': + parameter.test_title = "Test" + if parameter.reference_title == '': + parameter.reference_title = "Reference" + # test and ref color lists + color_list = [(0.12,0.96,0.87), (0.16,1.00,0.47), (0.22,1.00,0.19), + (0.70,1.00,0.20), (1.00,0.77,0.18), (1.00,0.20,0.11), + (1.00,0.05,0.30), (1.00,0.13,0.84), (0.61,0.15,1.00), + (0.07,0.14,0.99), (0.09,0.56,1.00), (0.16,0.91,0.99)] + + # First panel + plot_panel_seasonality('test', fig, proj, export, color_list, parameter) + + # Second panel + plot_panel_seasonality('ref', fig, proj, export, color_list, parameter) + + # Figure title + fig.suptitle(parameter.main_title, x=0.5, y=0.97, fontsize=15) + + # Prepare to save figure + # get_output_dir => {parameter.results_dir}/{set_name}/{parameter.case_id} + # => {parameter.results_dir}/streamflow/{parameter.case_id} + output_dir = get_output_dir(parameter.current_set, parameter) + if parameter.print_statements: + print('Output dir: {}'.format(output_dir)) + # get_output_dir => {parameter.orig_results_dir}/{set_name}/{parameter.case_id} + # => {parameter.orig_results_dir}/streamflow/{parameter.case_id} + original_output_dir = get_output_dir(parameter.current_set, parameter, ignore_container=True) + if parameter.print_statements: + print('Original output dir: {}'.format(original_output_dir)) + # parameter.output_file is defined in acme_diags/driver/streamflow_driver.py + # {parameter.results_dir}/streamflow/{parameter.case_id}/{parameter.output_file} + file_path = os.path.join(output_dir, parameter.output_file) + # {parameter.orig_results_dir}/streamflow/{parameter.case_id}/{parameter.output_file} + original_file_path = os.path.join(original_output_dir, parameter.output_file) + + # Save figure + for f in parameter.output_format: + f = f.lower().split('.')[-1] + plot_suffix = '.' + f + plot_file_path = file_path + plot_suffix + plt.savefig(plot_file_path) + # Get the filename that the user has passed in and display that. + # When running in a container, the paths are modified. + original_plot_file_path = original_file_path + plot_suffix + # Always print, even without `parameter.print_statements` + print('Plot saved in: ' + original_plot_file_path) + + # Save individual subplots + for f in parameter.output_format_subplot: + page = fig.get_size_inches() + i = 0 + for p in panel: + # Extent of subplot + subpage = np.array(p).reshape(2, 2) + subpage[1, :] = subpage[0, :] + subpage[1, :] + subpage = subpage + np.array(border).reshape(2, 2) + subpage = list((subpage * page).flatten()) + extent = matplotlib.transforms.Bbox.from_extents(*subpage) + # Save subplot + subplot_suffix = '.%i.' % i + f + subplot_file_path = file_path + subplot_suffix + plt.savefig(subplot_file_path, bbox_inches=extent) + # Get the filename that the user has passed in and display that. + # When running in a container, the paths are modified. + original_subplot_file_path = original_file_path + subplot_suffix + # Always print, even without `parameter.print_statements` + print('Sub-plot saved in: ' + original_subplot_file_path) + i += 1 + + plt.close() diff --git a/acme_diags/viewer/main.py b/acme_diags/viewer/main.py index 5544b73bd299..7c0bed9ae96c 100644 --- a/acme_diags/viewer/main.py +++ b/acme_diags/viewer/main.py @@ -1,7 +1,7 @@ import os import collections from bs4 import BeautifulSoup -from . import default_viewer, utils, area_mean_time_series_viewer, mean_2d_viewer, enso_diags_viewer, qbo_viewer +from . import default_viewer, utils, area_mean_time_series_viewer, mean_2d_viewer, enso_diags_viewer, qbo_viewer, streamflow_viewer import acme_diags # A mapping of each diagnostics set to the viewer @@ -15,7 +15,8 @@ 'cosp_histogram': default_viewer.create_viewer, 'area_mean_time_series': area_mean_time_series_viewer.create_viewer, 'enso_diags': enso_diags_viewer.create_viewer, - 'qbo': qbo_viewer.create_viewer + 'qbo': qbo_viewer.create_viewer, + 'streamflow': streamflow_viewer.create_viewer } def create_index(root_dir, title_and_url_list): diff --git a/acme_diags/viewer/streamflow_viewer.py b/acme_diags/viewer/streamflow_viewer.py new file mode 100644 index 000000000000..5bf45d282e3a --- /dev/null +++ b/acme_diags/viewer/streamflow_viewer.py @@ -0,0 +1,74 @@ +import os +from cdp.cdp_viewer import OutputViewer +from .default_viewer import create_metadata +from .utils import add_header, h1_to_h3 + + +def create_viewer(root_dir, parameters): + """ + Given a set of parameters for the streamflow set, + create a single webpage. + + Return the title and url for this page. + """ + viewer = OutputViewer(path=root_dir) + + # The name that's displayed on the viewer. + display_name = 'Streamflow' + set_name = 'streamflow' + # The title of the columns on the webpage. + # Appears in the second and third columns of the bolded rows. + cols = ['Description', 'Plot'] + viewer.add_page(display_name, short_name=set_name, columns=cols) + param_dict = {} + for param in parameters: + key = param.plot_type + if key not in param_dict.keys(): + param_dict[key] = [] + param_dict[key].append(param) + for plot_type in param_dict.keys(): + # Appears in the first column of the bolded rows. + viewer.add_group(plot_type.capitalize()) + # Only iterate through parameters with this plot type. + valid_parameters = param_dict[plot_type] + for param in valid_parameters: + # We need to make sure we have relative paths, and not absolute ones. + # This is why we don't use get_output_dir() as in the plotting script + # to get the file name. + ext = param.output_format[0] + # param.output_file is defined in acme_diags/driver/streamflow_driver.py + # This must be use param.case_id and param.output_file + # to match the file_path determined in + # acme_diags/plot/cartopy/streamflow_plot.py. + # Otherwise, the plot will not be properly linked from the viewer. + relative_path = os.path.join( + '..', set_name, param.case_id, + param.output_file) + image_relative_path = '{}.{}'.format(relative_path, ext) + if param.print_statements: + print('image_relative_path: {}'.format(image_relative_path)) + formatted_files = [] + # TODO: will param.variables ever be longer than one variable? + # If so, we'll need to create unique image_relative_paths + for var in param.variables: + # Appears in the first column of the non-bolded rows. + # This should use param.case_id to match the output_dir determined by + # get_output_dir in acme_diags/plot/cartopy/streamflow_plot.py. + # Otherwise, the plot image and the plot HTML file will have URLs + # differing in the final directory name. + viewer.add_row(param.case_id) + # Adding the description for this var to the current row. + # This was obtained and stored in the driver for this plotset. + # Appears in the second column of the non-bolded rows. + viewer.add_col(param.viewer_descr[var]) + # Link to an html version of the plot png file. + # Appears in the third column of the non-bolded rows. + viewer.add_col(image_relative_path, is_file=True, title='Plot', + other_files=formatted_files, + meta=create_metadata(param)) + + url = viewer.generate_page() + add_header(root_dir, os.path.join(root_dir, url), parameters) + h1_to_h3(os.path.join(root_dir, url)) + + return display_name, url diff --git a/setup.py b/setup.py index aa81949846bd..cf2393c2b997 100644 --- a/setup.py +++ b/setup.py @@ -29,6 +29,7 @@ def get_all_files_in_dir(directory, pattern): area_mean_time_series = get_all_files_in_dir('acme_diags/driver/default_diags', 'area_mean_time_series*') qbo = get_all_files_in_dir('acme_diags/driver/default_diags', 'qbo*') +streamflow = get_all_files_in_dir('acme_diags/driver/default_diags', 'streamflow*') enso_diags_files = get_all_files_in_dir('acme_diags/driver/default_diags', 'enso_*') @@ -68,6 +69,9 @@ def get_all_files_in_dir(directory, pattern): (os.path.join(INSTALL_PATH, 'qbo'), qbo ), + (os.path.join(INSTALL_PATH, 'streamflow'), + streamflow + ), (INSTALL_PATH, ['acme_diags/driver/acme_ne30_ocean_land_mask.nc', 'misc/e3sm_logo.png' diff --git a/tests/system/all_sets.cfg b/tests/system/all_sets.cfg index 6563d05efcbd..820c44c3df33 100644 --- a/tests/system/all_sets.cfg +++ b/tests/system/all_sets.cfg @@ -271,3 +271,20 @@ ref_file = 'U_201301_201312.nc' test_name = "system tests" variables = ["U"] output_file = 'qbo_diags' + +[streamflow] +sets = ["streamflow"] +case_id = "seasonality" +variables = ["RIVER_DISCHARGE_OVER_LAND_LIQ"] +regions = ["global"] +seasons = ["ANN"] +max_num_gauges = 12 +# TODO: Replace these paths with paths that automated testing tools can access (in future pull request). +#gauges_path = +#test_data_path = +#ref_mat_file = +start_yr = '1959' +end_yr = '1961' +results_dir = "all_sets_results_test" +test_name = "system tests" +print_statements = True diff --git a/tests/system/test_diags.py b/tests/system/test_diags.py index f98802a7c899..ed85b582ad2f 100644 --- a/tests/system/test_diags.py +++ b/tests/system/test_diags.py @@ -159,6 +159,23 @@ def check_enso_scatter_plots(self, case_id): TestAllSets.results_dir, set_name, case_id_lower) self.check_html_image(html_path, png_path) + def check_streamflow_plots(self): + case_id = 'seasonality' + case_id_lower = case_id.lower() + set_name = 'streamflow' + variables = ['RIVER_DISCHARGE_OVER_LAND_LIQ'] + for variable in variables: + variable_lower = variable.lower() + png_path = '{}/{}/{}.png'.format( + set_name, case_id_lower, case_id_lower) + full_png_path = '{}{}'.format(TestAllSets.results_dir, png_path) + # Expected: streamflow_diags_model_to_model_local/streamflow/seasonality/seasonality.png + self.assertTrue(os.path.exists(full_png_path)) + html_path = '{}viewer/{}/{}/{}/plot.html'.format( + TestAllSets.results_dir, set_name, case_id_lower, case_id_lower) + # Expected: streamflow_diags_model_to_model_local/viewer/streamflow/seasonality/seasonality/plot.html + self.check_html_image(html_path, png_path) + # Test results_dir def test_results_dir(self): self.assertTrue(TestAllSets.results_dir.endswith('all_sets_results_test/')) @@ -232,6 +249,9 @@ def test_qbo(self): TestAllSets.results_dir, set_name, case_id_lower) self.check_html_image(html_path, png_path) + def test_streamflow(self): + self.check_streamflow_plots() + def test_zonal_mean_2d(self): self.check_plots_2d('zonal_mean_2d') diff --git a/tests/test_run.py b/tests/test_run.py index e9ff6725a16e..a877adecbd8c 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -4,6 +4,7 @@ from acme_diags.parameter.area_mean_time_series_parameter import AreaMeanTimeSeriesParameter from acme_diags.parameter.core_parameter import CoreParameter from acme_diags.parameter.enso_diags_parameter import EnsoDiagsParameter +from acme_diags.parameter.streamflow_parameter import StreamflowParameter from acme_diags.parameter.zonal_mean_2d_parameter import ZonalMean2dParameter @@ -45,8 +46,13 @@ def test_all_sets_and_all_seasons(self): enso_param = EnsoDiagsParameter() enso_param.start_yr = '2000' enso_param.end_yr = '2004' - - parameters = self.runner.get_final_parameters([self.core_param, ts_param, enso_param]) + + streamflow_param = StreamflowParameter() + # TODO: Replace these paths with paths that automated testing tools can access (in future pull request). + #streamflow_param.gauges_path = + #streamflow_param.ref_mat_file = + + parameters = self.runner.get_final_parameters([self.core_param, ts_param, enso_param, streamflow_param]) # Counts the number of each set and each seasons to run the diags on. set_counter, season_counter = collections.Counter(), collections.Counter() for param in parameters: @@ -62,9 +68,10 @@ def test_all_sets_and_all_seasons(self): self.fail(msg.format(set_name, count)) all_season_counts = list(season_counter.values()) - # enso_diags only runs ANN, no seasons - # So, reduce the ANN count by the number of times enso_diags appears + # enso_diags and streamflow only run ANN, no seasons + # So, reduce the ANN count by the number of times these appear all_season_counts[0] -= set_counter['enso_diags'] + all_season_counts[0] -= set_counter['streamflow'] if not all(all_season_counts[0] == count for count in all_season_counts): self.fail('Counts for the seasons don\'t match: {}'.format(all_season_counts)) @@ -100,5 +107,6 @@ def test_zonal_mean_2d(self): msg = 'The lengths are either 0 or not equal: {} {} {}' self.fail(msg.format(len1, len2, len3)) + if __name__ == '__main__': unittest.main()