python scripts to analysis LE model output

calculate_statistics.py

import xarray as xr
import numpy as np
from processing import processing as pr
from skill_metrics import skill_metrics as sk

def calculate_statistics(model='CESM', member='002'):
    ###======================================
    ### Load data
    ###======================================
    ds_somffn = read_peter(model=model, member=member)
    ds_cesm = read_model(model=model, member=member)
    #ds_sst = read_sst(model=model, member=member)

    ###======================================
    ### Put data into xarray dataset
    ###======================================
    ds = xr.Dataset(
        {
        'pco2_somffn':(['time','lat','lon'], ds_somffn['pco2']),
        'pco2_cesm':(['time','lat','lon'], ds_cesm['pco2']),
        #'sst':(['time','lat','lon'], ds_sst['sst']),
        },

        coords={
        'lat': (['lat'], ds_somffn['lat']),
        'lon': (['lon'], ds_somffn['lon']),
        'time': (['time'], ds_somffn['time'])
        })
    
    ###======================================
    ### Breakdown the signal and 
    ### Decompose pCO2 into T and nonT
    ###======================================
    ### SOMFFN breakdown signal
    pco2_somffn = ds['pco2_somffn'].copy()
    pco2_somffn_iav = pr(pco2_somffn).detrend().rolling_mean(window=12).values
    pco2_somffn_iav_low = pr(pco2_somffn_iav).rolling_mean(window=60).values
    #pco2_somffn_iav_low = pr(pco2_somffn).detrend().rolling_mean(window=60).values
    pco2_somffn_iav_high = pco2_somffn_iav - pco2_somffn_iav_low
    pco2_somffn_av = pr(pco2_somffn).detrend().values - pco2_somffn_iav
    
    ### Decompose into T and nonT
    #pco2_somffn_T = taro(ds['pco2_somffn'], ds['sst']).pco2_T()
    #pco2_somffn_nonT = taro(ds['pco2_somffn'], ds['sst']).pco2_nonT()
    
    ### CESM breakdown
    pco2_cesm = ds['pco2_cesm'].copy()
    pco2_cesm_iav = pr(pco2_cesm).detrend().rolling_mean(window=12).values
    pco2_cesm_iav_low = pr(pco2_cesm_iav).rolling_mean(window=60).values
    #pco2_cesm_iav_low = pr(pco2_cesm).detrend().rolling_mean(window=60).values
    pco2_cesm_iav_high = pco2_cesm_iav - pco2_cesm_iav_low
    pco2_cesm_av = pr(pco2_cesm).detrend().values - pco2_cesm_iav
    
    ### Decompose into T and nonT
    #pco2_cesm_T = taro(ds['pco2_cesm'], ds['sst']).pco2_T()
    #pco2_cesm_nonT = taro(ds['pco2_cesm'], ds['sst']).pco2_nonT()
    
    ###======================================
    ### Return dataset
    ###======================================
    ds_out = xr.Dataset(
        {
        'bias': (['lat', 'lon'], sk.avg_error(pco2_cesm, pco2_somffn) ),
        'bias_detrend':(['lat','lon'], sk.avg_error( pr(pco2_cesm).detrend().values, pr(pco2_somffn).detrend().values)),
        'bias_iav': (['lat', 'lon'], sk.avg_error(pco2_cesm_iav, pco2_somffn_iav) ),
        'bias_iav_low': (['lat', 'lon'], sk.avg_error(pco2_cesm_iav_low, pco2_somffn_iav_low) ),
        'bias_iav_high': (['lat', 'lon'], sk.avg_error(pco2_cesm_iav_high, pco2_somffn_iav_high) ),
        'bias_av': (['lat', 'lon'], sk.avg_error(pco2_cesm_av, pco2_somffn_av) ),
        #'bias_T': (['lat', 'lon'], sk.avg_error(pco2_cesm_T, pco2_somffn_T) ),
        #'bias_nonT': (['lat', 'lon'], sk.avg_error(pco2_cesm_nonT, pco2_somffn_nonT) ),
            
        'aae':(['lat','lon'], sk.avg_abs_error(pco2_cesm, pco2_somffn) ),
        'aae_detrend':(['lat','lon'], sk.avg_abs_error( pr(pco2_cesm).detrend().values, pr(pco2_somffn).detrend().values)),
        'aae_iav':(['lat','lon'], sk.avg_abs_error(pco2_cesm_iav, pco2_somffn_iav) ),
        'aae_iav_low':(['lat','lon'], sk.avg_abs_error(pco2_cesm_iav_low, pco2_somffn_iav_low) ),
        'aae_iav_high':(['lat','lon'], sk.avg_abs_error(pco2_cesm_iav_high, pco2_somffn_iav_high) ),
        'aae_av':(['lat','lon'], sk.avg_abs_error(pco2_cesm_av, pco2_somffn_av) ),
        #'aae_T':(['lat','lon'], sk.avg_abs_error(pco2_cesm_T, pco2_somffn_T) ),
        #'aae_nonT':(['lat','lon'], sk.avg_abs_error(pco2_cesm_nonT, pco2_somffn_nonT) ),
            
        'corr':(['lat','lon'], sk.correlation(pco2_cesm, pco2_somffn)),
        'corr_detrend':(['lat','lon'], sk.correlation( pr(pco2_cesm).detrend().values, pr(pco2_somffn).detrend().values)),
        'corr_iav':(['lat','lon'], sk.correlation(pco2_cesm_iav, pco2_somffn_iav) ),
        'corr_iav_low':(['lat','lon'], sk.correlation(pco2_cesm_iav_low, pco2_somffn_iav_low) ),
        'corr_iav_high':(['lat','lon'], sk.correlation(pco2_cesm_iav_high, pco2_somffn_iav_high) ),
        'corr_av':(['lat','lon'], sk.correlation(pco2_cesm_av, pco2_somffn_av) ),
        #'corr_T':(['lat','lon'], sk.correlation(pco2_cesm_T, pco2_somffn_T)),
        #'corr_nonT':(['lat','lon'], sk.correlation(pco2_cesm_nonT, pco2_somffn_nonT)),
            
        'std_star':(['lat','lon'], sk.std_star(pco2_cesm, pco2_somffn)),
        'std_star_detrend':(['lat','lon'], sk.std_star( pr(pco2_cesm).detrend().values, pr(pco2_somffn).detrend().values)),
        'std_star_iav':(['lat','lon'], sk.std_star(pco2_cesm_iav, pco2_somffn_iav) ),
        'std_star_iav_low':(['lat','lon'], sk.std_star(pco2_cesm_iav_low, pco2_somffn_iav_low) ),
        'std_star_iav_high':(['lat','lon'], sk.std_star(pco2_cesm_iav_high, pco2_somffn_iav_high) ),
        'std_star_av':(['lat','lon'], sk.std_star(pco2_cesm_av, pco2_somffn_av) ),
        #'std_star_T':(['lat','lon'], sk.std_star(pco2_cesm_T, pco2_somffn_T)),
        #'std_star_nonT':(['lat','lon'], sk.std_star(pco2_cesm_nonT, pco2_somffn_nonT)),
        },

        coords={
        'lat': (['lat'], ds_somffn['lat']),
        'lon': (['lon'], ds_somffn['lon']),
        })
        
    return ds_out

processing.py

%%writefile processing.py 

import numpy as np
import xarray as xr
import scipy.signal

class processing():
    def __init__(self, data):
        #self._data_raw = data.copy()
        self._data = data.copy()
        
    @property
    def values(self):
        return self._data
    
    def rolling_mean(self, window=12):
        """
        rolling_mean(self, window=12)
        running mean centered in the window
        """
        self._data = self._data.rolling(time=window, center=True).mean()
        return self
    
    def detrend_ufunc(self, X, axis=0):
        ### mask nan points
        mask = ~np.isnan(X)
        ## define output matrix
        out = X*np.nan
        ### detrend along axis
        out[mask] = scipy.signal.detrend(X[mask], axis=axis, type='linear')
        return out

    def detrend(self,axis=0):
        self._data = xr.apply_ufunc(self.detrend_ufunc, self._data)
        return self

    def long_term_mean(self, dim='time'):
        '''long term mean alont dimension dim'''
        self._data = self._data.mean(dim)
        return self

    def global_avg(self, dim=['lat','lon']):
        '''long term mean alont dimension dim'''
        self._data = self._data.mean(dim)
        return self

    def global_mean(self, dim=['lat','lon']):
        '''long term mean alont dimension dim'''
        self._data = self._data.mean(dim)
        return self
    
    def global_median(self, dim=['lat','lon']):
        '''long term mean alont dimension dim'''
        self._data = self._data.median(dim)
        return self

    def zonal_mean(self,dim='lon'):
        '''long term mean alont dimension dim'''
        self._data = self._data.mean(dim)
        return self
    
    def zonal_median(self,dim='lon'):
        '''long term mean alont dimension dim'''
        self._data = self._data.median(dim)
        return self

    def ensemble_mean(self, dim='ensemble'):
        '''long term mean alont dimension dim'''
        self._data = self._data.mean(dim)
        return self

    def annual_mean(self, dim='time', nyears=35):
        self._data = self._data.groupby_bins(dim, nyears).mean(dim=dim)
        return self

    def remove_mean(self, dim='time'):
        ''' 
        remove_mean(X, dim='time')
        * use with .groupby_bins().apply() to remove annual mean
        '''
        self._data =  self._data - self._data.mean(dim)
        return self

    def annual_mean_repeating(self, dim='time', nyears=35, axis=0):
        tmp = self._data.groupby_bins(dim, nyears).mean(dim=dim)
        #tmp = ds.groupby('time.year').mean(dim='time')
        self._data = xr.DataArray(np.repeat(tmp.values, 12, axis=axis), dims=['time','lat','lon'])
        return self

run_analysis.py

### Model / Members / number of members
LE_model = 'MPI'

###======================================
### Get members from model
###======================================
if LE_model=='CESM':
    members=[1,2,9,10,11,12,13,14,15,16,17,18,20,21,23,24,25,26,30,31,34,35,101,102,103,104]
    n = len(members)
    print('Loading {0} members from {1}'.format(n,LE_model))

if LE_model=='GFDL':
    members=[1,2,3,4,5,6,7,8,9,10,11,12,13,14,16,17,18,19,20,22,23,26,27,28,29,30]
    n = len(members)
    print('Loading {0} members from {1}'.format(n,LE_model))

if LE_model=='MPI':
    members=[6,9,14,20,22,24,25,27,38,43,45,46,57,60,64,70,75,77,78,80,81,83,91,95,98]
    n = len(members)
    print('Loading {0} members from {1}'.format(n,LE_model))
    
###======================================
### Load data
###======================================
fut = client.map(calculate_statistics, tuple(np.repeat(LE_model, n)), members)

###======================================
### Create dataset
###======================================
print('Concatenate dataset')
ds = xr.concat(client.gather(fut), dim='ensemble')
ds_biomes = read_biomes()
ds_data = xr.merge([ds, ds_biomes['mean_biomes']])

###======================================
### write file
###======================================
print('write to netcdf')
ds_data.to_netcdf('/local/data/artemis/workspace/gloege/SOCAT-LE/data/clean/statistics_somffn_MPI.nc')
print('Complete!!')

skill_metrics.py

%%writefile skill_metrics.py 

import numpy as np
import xarray as xr

class skill_metrics():        
    def std(x, dim='time'):
        return x.std(dim=dim)

    def covariance(x, y, dim='time'):
        return ((x - x.mean(dim=dim)) * (y - y.mean(dim=dim))).mean(dim=dim)

    def correlation(x, y,dim='time'):
        cov = ((x - x.mean(dim=dim)) * (y - y.mean(dim=dim))).mean(dim=dim)
        return cov / (x.std(dim=dim) * y.std(dim=dim))

    def avg_abs_error(obs, prd, dim='time'):
        return xr.ufuncs.fabs(prd-obs).mean(dim=dim) 

    def avg_error(obs ,prd, dim='time'):
        return (prd-obs).mean(dim=dim) 

    def std_star(obs, prd, dim='time'):
        return prd.std(dim=dim) / obs.std(dim=dim)

    def rmse(m, r):
        return xr.ufuncs.sqrt(xr.ufuncs.square((m-r)).mean(dim='time'))

    def urmse(m, r):
        return xr.ufuncs.sqrt(xr.ufuncs.square( (m - m.mean(dim='time')) - (r - r.mean(dim='time')) ).mean(dim='time'))
    
    def ri(m, r,dim='time'):
        return xr.ufuncs.exp(xr.ufuncs.sqrt( xr.ufuncs.square( xr.ufuncs.log(m/r) ).mean(dim=dim)))
    
    def nse(m, r, dim='time'):
        numer = xr.ufuncs.square(m - m.mean(dim=dim)).mean(dim=dim) - xr.ufuncs.square(r - m).mean(dim=dim)
        return numer / xr.ufuncs.square(m - m.mean(dim=dim)).mean(dim=dim)