# coding=utf-8 import iris import iris.analysis import iris.coord_categorisation import iris.coords import iris.exceptions import numpy as np from bscearth.utils.date import parse_date, add_months from bscearth.utils.log import Log from iris.time import PartialDateTime from earthdiagnostics.diagnostic import * from earthdiagnostics.frequency import Frequencies from earthdiagnostics.statistics.climatologicalpercentile import ClimatologicalPercentile from earthdiagnostics.utils import Utils, TempFile class DaysOverPercentile(Diagnostic): """ Calculates the montlhy percentiles :param data_manager: data management object :type data_manager: DataManager :param variable: variable to average :type variable: str """ alias = 'daysover' "Diagnostic alias for the configuration file" def __init__(self, data_manager, domain, variable, start_year, end_year, startdate, forecast_month): Diagnostic.__init__(self, data_manager) self.variable = variable self.domain = domain self.start_year = start_year self.end_year = end_year self.forecast_month = forecast_month self.startdate = startdate def __eq__(self, other): return self.startdate == other.startdate and self.domain == other.domain and \ self.variable == other.variable and self.start_year == other.start_year and \ self.end_year == other.end_year def __str__(self): return 'Days over percentile Startdate: {0.startdate} Variable: {0.domain}:{0.variable} ' \ 'Climatology: {0.start_year}-{0.end_year}'.format(self) @classmethod def generate_jobs(cls, diags, options): """ Creates a job for each chunk to compute the diagnostic :param diags: Diagnostics manager class :type diags: Diags :param options: domain, variable, percentil number, maximum depth (level) :type options: list[str] :return: """ options_available = (DiagnosticDomainOption(), DiagnosticOption('variable'), DiagnosticIntOption('start_year'), DiagnosticIntOption('end_year'), DiagnosticListIntOption('forecast_month'),) options = cls.process_options(options, options_available) job_list = list() for startdate in diags.config.experiment.startdates: for forecast_month in options['forecast_month']: job_list.append(DaysOverPercentile(diags.data_manager, options['domain'], options['variable'], options['start_year'], options['end_year'], startdate, forecast_month)) return job_list def request_data(self): var_name = '{0.variable}prct{0.start_year}{0.forecast_month}-{0.end_year}{0.forecast_month}'.format(self) self.percentiles_file = self.request_chunk(self.domain, var_name, None, None, None, frequency=Frequencies.climatology) self.variable_file = self.request_chunk(self.domain, self.variable, self.startdate, None, None) def declare_data_generated(self): var_over = self.variable + '_daysover_q{0}' var_below = self.variable + '_daysbelow_q{0}' self.days_over_file = {} self.days_below_file = {} for perc in ClimatologicalPercentile.Percentiles: self.days_over_file[perc] = self.declare_chunk(self.domain, var_over.format(int(perc * 100)), self.startdate, None, None, frequency=Frequencies.monthly, vartype=VariableType.STATISTIC) self.days_below_file[perc] = self.declare_chunk(self.domain, var_below.format(int(perc * 100)), self.startdate, None, None, frequency=Frequencies.monthly, vartype=VariableType.STATISTIC) def compute(self): """ Runs the diagnostic """ raise Exception('Pues me enfado y no respiro!!!') iris.FUTURE.netcdf_promote = True percentiles = iris.load_cube(self.percentiles_file.local_file) handler = Utils.openCdf(self.variable_file.local_file) if 'realization' in handler.variables: handler.variables[self.variable].coordinates = 'realization' handler.close() var = iris.load_cube(self.variable_file.local_file) date = parse_date(self.startdate) lead_date = add_months(date, 1, self.data_manager.config.experiment.calendar) leadtimes = {1: PartialDateTime(lead_date.year, lead_date.month, lead_date.day)} def assign_leadtime(coord, x): # noinspection PyBroadException try: leadtime_month = 1 partial_date = leadtimes[leadtime_month] while coord.units.num2date(x) >= partial_date: leadtime_month += 1 try: partial_date = leadtimes[leadtime_month] except KeyError: new_date = add_months(date, leadtime_month, self.data_manager.config.experiment.calendar) partial_date = PartialDateTime(new_date.year, new_date.month, new_date.day) leadtimes[leadtime_month] = partial_date return leadtime_month except Exception: pass iris.coord_categorisation.add_categorised_coord(var, 'leadtime', 'time', assign_leadtime) iris.coord_categorisation.add_year(var, 'time') iris.coord_categorisation.add_day_of_year(var, 'time') try: realization_coord = var.coord('realization') except iris.exceptions.CoordinateNotFoundError: realization_coord = None self.lat_coord = var.coord('latitude') self.lon_coord = var.coord('longitude') results_over = {perc: iris.cube.CubeList() for perc in ClimatologicalPercentile.Percentiles} results_below = {perc: iris.cube.CubeList() for perc in ClimatologicalPercentile.Percentiles} var_daysover = 'daysover' var_days_below = 'daysbelow' long_name_days_over = 'Proportion of days over a given percentile for {0.start_year}-{0.end_year} ' \ 'climatology'.format(self) long_name_days_below = 'Proportion of days below a given percentile for {0.start_year}-{0.end_year} ' \ 'climatology'.format(self) for leadtime in leadtimes.keys(): Log.debug('Computing startdate {0} leadtime {1}', self.startdate, leadtime) leadtime_slice = var.extract(iris.Constraint(leadtime=leadtime)) if len(percentiles.coords('leadtime')) > 0: percentiles_leadtime = percentiles.extract(iris.Constraint(leadtime=leadtime)) else: percentiles_leadtime = percentiles time_coord = iris.coords.AuxCoord.from_coord(leadtime_slice.coord('time')) first_time = time_coord.points[0] last_time = time_coord.points[-1] timesteps = leadtime_slice.coord('time').shape[0] time_coord = time_coord.copy(first_time + (last_time - first_time) / 2, (first_time, last_time)) for percentile_slice in percentiles_leadtime.slices_over('percentile'): percentile = percentile_slice.coord('percentile').points[0] # noinspection PyTypeChecker days_over = np.sum(leadtime_slice.data > percentile_slice.data, 0) / float(timesteps) result = self.create_results_cube(days_over, percentile, realization_coord, time_coord, var_daysover, long_name_days_over) results_over[percentile].append(result) # noinspection PyTypeChecker days_below = np.sum(leadtime_slice.data < percentile_slice.data, 0) / float(timesteps) result = self.create_results_cube(days_below, percentile, realization_coord, time_coord, var_days_below, long_name_days_below) results_below[percentile].append(result) Log.debug('Saving percentiles startdate {0}', self.startdate) for perc in ClimatologicalPercentile.Percentiles: iris.FUTURE.netcdf_no_unlimited = True temp = TempFile.get() iris.save(results_over[perc].merge_cube(), temp, zlib=True, unlimited_dimensions=['time']) Utils.rename_variables(temp, {'dim2': 'ensemble', 'dim1': 'ensemble'}, must_exist=False, rename_dimension=True) self.days_over_file[perc].set_local_file(temp, rename_var='daysover') temp = TempFile.get() iris.save(results_below[perc].merge_cube(), temp, zlib=True, unlimited_dimensions=['time']) Utils.rename_variables(temp, {'dim2': 'ensemble', 'dim1': 'ensemble'}, must_exist=False, rename_dimension=True) self.days_below_file[perc].set_local_file(temp, rename_var='daysbelow') del self.days_over_file del self.days_below_file del self.lat_coord del self.lon_coord def create_results_cube(self, days_over, percentile, realization_coord, time_coord, var_name, long_name): result = iris.cube.Cube(days_over.astype(np.float32), var_name=var_name, long_name=long_name, units=1.0) if realization_coord is not None: result.add_aux_coord(realization_coord, 0) result.add_dim_coord(self.lat_coord, 1) result.add_dim_coord(self.lon_coord, 2) else: result.add_dim_coord(self.lat_coord, 0) result.add_dim_coord(self.lon_coord, 1) result.add_aux_coord(iris.coords.AuxCoord(percentile, long_name='percentile')) result.add_aux_coord(time_coord) return result