#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Created on Thu Jul 23 11:10:07 2020 @author: jabadgeley """ # Import packages from xarray import DataArray #import sys #sys.path.append('') from utilities import (regrid_from_nc1_to_latlonlist as regrid_list, get_lons) class model_data(): def __init__(self, name, vars_0D={}, vars_1D={}, vars_2D={}, vars_3D={}, vars_4D={}): self.vars_0D = list(vars_0D.keys()) self.vars_1D = list(vars_1D.keys()) self.vars_2D = list(vars_2D.keys()) self.vars_3D = list(vars_3D.keys()) self.vars_4D = list(vars_4D.keys()) self.vars_2D_3D_4D = [*self.vars_2D, *self.vars_3D, *self.vars_4D] self.vars_3D_4D = [*self.vars_3D, *self.vars_4D] self.vars_all = [*self.vars_0D, *self.vars_1D, *self.vars_2D, *self.vars_3D, *self.vars_4D] self.name = name for var_0D in self.vars_0D: setattr(self, var_0D, vars_0D[var_0D]) for var_1D in self.vars_1D: setattr(self, var_1D, vars_1D[var_1D]) for var_2D in self.vars_2D: setattr(self, var_2D, vars_2D[var_2D]) for var_3D in self.vars_3D: setattr(self, var_3D, vars_3D[var_3D]) for var_4D in self.vars_4D: setattr(self, var_4D, vars_4D[var_4D]) def select_lats(self, lat_tuple): for var in self.vars_2D_3D_4D: if lat_tuple[0] == lat_tuple[1]: setattr(self, var, getattr(self, var).sel(lat=lat_tuple[0], method='nearest')) else: setattr(self, var, getattr(self, var).sel(\ lat=slice(lat_tuple[0],lat_tuple[1]))) def select_lons(self, lon_tuple): for var in self.vars_2D_3D_4D: if lon_tuple[0] < lon_tuple[1]: setattr(self, var, getattr(self, var).sel(\ lon=slice(lon_tuple[0],lon_tuple[1]))) elif lon_tuple[0] > lon_tuple[1]: setattr(self, var, getattr(self, var).sel(lon=get_lons(\ getattr(self, var),lon_tuple),method='nearest')) else: setattr(self, var, getattr(self, var).sel(lon=lon_tuple[0], method='nearest')) def select_plevs(self, plev_tuple): for var in self.vars_4D: if plev_tuple[0] == plev_tuple[1]: setattr(self, var, getattr(self, var).sel(plev=plev_tuple[0], method='nearest')) else: setattr(self, var, getattr(self, var).sel(plev=slice(\ plev_tuple[0],plev_tuple[1]))) def lons_to_360(self): for var in self.vars_2D_3D_4D: setattr(self, var, getattr(self, var).assign_coords(\ lon=(getattr(self, var).lon % 360))) def lons_to_180(self): for var in self.vars_2D_3D_4D: setattr(self, var, getattr(self, var).assign_coords(\ lon=(((getattr(self, var).lon + 180) % 360) - 180))) def sortby_lons(self): for var in self.vars_2D_3D_4D: setattr(self, var, getattr(self, var).sortby('lon')) def sortby_lats(self): for var in self.vars_2D_3D_4D: setattr(self, var, getattr(self, var).sortby('lat')) def average_time(self, average_type='regular'): for var in self.vars_3D_4D: print(var) if average_type=='regular': setattr(self, var, getattr(self, var).mean('time')) elif average_type == 'monthly_weights': try: month_length = getattr(self, var).time.dt.days_in_month except TypeError as error: print(error) print('The lengths of each month will be estimated from the calendar,', 'assuming the first month is January.') if getattr(self, var).time.calendar == 'proleptic_gregorian': month_base = [31, 28.25, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] elif getattr(self, var).time.calendar == 'gregorian': month_base = [31, 28.25, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] elif getattr(self, var).time.calendar == '365_day': month_base = [31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31] num_months = len(getattr(self, var).time) if num_months < 12: month_length_temp = month_base[0:num_months] elif num_months == 12: month_length_temp = month_base else: month_base_ext = month_base.copy() for ii in range(num_months % 12 - 1): month_base_ext += month_base month_length_temp = month_base_ext[0:num_months] month_length = DataArray(data=month_length_temp, dims='time', coords={'time':getattr(self, var).time}) print('Month lengths are:', month_length.values, 'days') weights = month_length / month_length.sum() ds_weighted = (getattr(self, var) * weights).sum(dim='time', skipna=False) setattr(self, var, ds_weighted) else: print('That is not an averaging option.', 'Please use regular or monthly_weights.', 'Resporting to regular averaging.') setattr(self, var, getattr(self, var).mean('time')) def zonal_mean(self): for var in self.vars_2D_3D_4D: setattr(self, var, getattr(self, var).mean('lon')) def meridional_mean(self): for var in self.vars_2D_3D_4D: setattr(self, var, getattr(self, var).mean('lat')) def string_to_variable(self, variable_string): MDvar = getattr(self, variable_string) MDdata = getattr(MDvar, variable_string) return MDvar, MDdata def regrid_to_targetlist(self, lats, lons, include_poles=False, method='bilinear'): for var in self.vars_2D_3D_4D: nc1 = getattr(self, var) dim_names = list(nc1.dims) data_new = regrid_list(nc1, lats, lons, var, include_poles=include_poles, method=method) coords_new={} [coords_new.update({ii:([ii],nc1[ii].values)}) for ii in dim_names]; coords_new['lat']=(['lat'],lats) coords_new['lon']=(['lon'],lons) dims_new = list(coords_new.keys()) nc1_new_array = DataArray(data=data_new, dims=dims_new, coords=coords_new) nc1_new = nc1_new_array.to_dataset(name=var) setattr(self, var, nc1_new) return def interpolate_to_latlist(self, lats): for var in self.vars_2D_3D_4D: setattr(self, var, getattr(self, var).interp(lat=lats)) return