def compute(cls, precip_rate, wind_direction, criteria): # Checking input types if isinstance(precip_rate, netCDF4._netCDF4.Dataset): pr = NCdata(precip_rate) elif isinstance(precip_rate, NCdata): pr = precip_rate else: raise TypeError( 'precip_rate of class ', precip_rate.__class__, ' must be either a netCDF4 reader or a ' 'NCData object') if isinstance(wind_direction, netCDF4._netCDF4.Dataset): wind_dir = NCdata(wind_direction) elif isinstance(wind_direction, NCdata): wind_dir = wind_direction else: raise TypeError( 'wind_direction of class ', wind_direction.__class__, ' must be either a netCDF4 reader or ' 'a NCData object') # Create new monsoon object NCmonsoon = MonsoonData.__new__(cls) # Check that pr and wind_dir have properly aligning grids # TODO shouldn't need this if statement if align_grids always works properly if pr.lats != wind_dir.lats or pr.lons != wind_dir.lons: pr.variable, wind_dir.variable, new_lats, new_lons = cls.align_grids( pr, wind_dir) pr.lats = wind_dir.lats = new_lats pr.lons = wind_dir.lons = new_lons # Check that pr and wind_dir are compatible; if so assign shared variables to monsoon object for attr in [ 'lats', 'lons', 'time', 'calendar', 'time_units', 'model_id' ]: if (np.asarray( getattr(pr, attr) != getattr(wind_dir, attr))).all(): raise TypeError(attr, ' not aligned in pr and wind_dir') setattr(NCmonsoon, attr, getattr(pr, attr)) # Create boolean monsoon from pr and wind_dir passing_pr = criteria['precip'](pr.variable) passing_wd = criteria['wind_dir'](wind_dir.variable) blank_monsoon = passing_pr & passing_wd # Add newly calculated data to the monsoon object NCmonsoon.var_name = 'monsoon' NCmonsoon.var_units = 'boolean' NCmonsoon.variable = blank_monsoon.astype(float) # Created combined global attributes monsoon_globalattrs = {} all_keys = list( set( list(pr.globalattrs.keys()) + list(wind_dir.globalattrs.keys()))) for key in all_keys: if key not in pr.globalattrs.keys(): monsoon_globalattrs[key + '_wd'] = wind_dir.globalattrs[key] elif key not in wind_dir.globalattrs.keys(): monsoon_globalattrs[key + '_pr'] = pr.globalattrs[key] elif pr.globalattrs[key] == wind_dir.globalattrs[key]: monsoon_globalattrs[key] = pr.globalattrs[key] else: monsoon_globalattrs[key + '_pr'] = pr.globalattrs[key] monsoon_globalattrs[key + '_wd'] = wind_dir.globalattrs[key] # monsoon_globalattrs['monsoon_criteria'] = criteria['summary_dict'] NCmonsoon.globalattrs = monsoon_globalattrs return NCmonsoon
from wnpmonsoon.ncdata import NCdata import numpy as np uas = NCdata(file_, 'uas') vas = NCdata(file_, 'uas') wind_sp = np.sqrt(np.square(uas.variable) + np.square(vas.variable)) # TODO decide if this is a good idea / if there is a better way to do this # Inheriting from another object? variable = 'ws' ws = NCdata('notrealfile', variable, create_new=True) ws.var_name = variable ws.model = uas.model_id ws.var_units = uas.var_units ws.variable = wind_sp ws.lats = uas.lats ws.lons = uas.lons ws.time = uas.time ws.calendar = uas.calendar ws.t_units = uas.t_units ws.write(outfile)