In [1]:

    

%matplotlib inline
%pylab inline
import numpy as np
import pandas as pd
import xarray as xr
import seaborn as sns
import os
import matplotlib.pyplot as plt
import timeit

pylab.rcParams['figure.figsize'] = (20, 12)


    

    




Populating the interactive namespace from numpy and matplotlib

In [2]:

    

var1 = xr.DataArray(np.arange(777600).reshape(12,360,180), name='var1')
var2 = xr.DataArray(np.random.randn(777600).reshape(12,360,180), name='var2')


    

In [3]:

    

def single_pixel_deviation(x):
    return (x - x.mean()) / x.std()

def single_correlation(x): 
    x_cor = x.groupby('dim_1').apply(single_pixel_deviation)
    return x_cor


    

In [4]:

    

def mycorr(x,y,dim=None):
    return (((x-x.mean(dim=dim))*(y-y.mean(dim=dim))).sum(dim=dim)/x[dim].size)/(x.std(dim=dim)*y.std(dim=dim))


    

In [31]:

    

%%timeit
var1_cor = var1.groupby('dim_2').apply(single_correlation)
var2_cor = var2.groupby('dim_2').apply(single_correlation)


    

    




1 loop, best of 3: 5min 41s per loop

In [15]:

    

#%%timeit
cov_test = mycorr(var2,var2,dim="dim_0")


    

In [17]:

    

cov_test.min()


    

    
Out[17]:





<xarray.DataArray 'var2' ()>
array(0.9999999999999998)

In [18]:

    

cov_test


    

    
Out[18]:





<xarray.DataArray 'var2' (dim_1: 360, dim_2: 180)>
array([[ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       ..., 
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.],
       [ 1.,  1.,  1., ...,  1.,  1.,  1.]])
Coordinates:
  * dim_1    (dim_1) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
  * dim_2    (dim_2) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...

In [36]:

    

var1.groupby('dim_2')


    

    
Out[36]:





<xarray.core.groupby.DataArrayGroupBy at 0x1f5f3eaa668>

In [33]:

    

var2_cor


    

    
Out[33]:





<xarray.DataArray 'var2' (dim_0: 12, dim_1: 5, dim_2: 2)>
array([[[ 0.18710249, -0.4510858 ],
        [-1.56929836,  1.17057218],
        [-0.37523565, -0.14987089],
        [ 0.08307831,  1.19393739],
        [ 1.6640759 , -1.19414631]],

       [[-1.5680252 ,  0.60751634],
        [ 1.36221602,  0.79183935],
        [ 0.50493692,  0.89706249],
        [ 1.46465714,  0.82596639],
        [ 0.23891788, -1.00256663]],

       [[ 0.30273316,  0.73489702],
        [-0.52721015, -1.35096231],
        [-0.29362169,  1.09570489],
        [-0.94746568,  0.53487384],
        [-0.8932819 , -0.4859217 ]],

       [[-1.33299442,  0.68930313],
        [ 0.95383022, -0.20947329],
        [-0.46765436,  1.03223905],
        [ 0.90740149, -0.37320876],
        [-1.20069748, -0.49889519]],

       [[-1.96669622, -0.07890089],
        [-0.0596491 , -1.1946685 ],
        [-0.62856913,  0.1186544 ],
        [ 1.22790715, -0.5273026 ],
        [-0.02830446,  1.00831689]],

       [[ 0.10867812, -1.3062592 ],
        [ 0.65040773,  0.65190501],
        [ 0.37718872, -0.13156251],
        [-0.83084255, -1.52680948],
        [ 1.42488702, -1.35324581]],

       [[ 0.4690665 ,  2.303552  ],
        [ 0.14358582,  1.62732487],
        [ 0.05792761, -1.03666755],
        [ 0.77666824, -0.96339638],
        [ 1.00315805,  1.56614762]],

       [[ 0.88410417, -1.60267624],
        [ 0.62520406, -0.84468287],
        [-2.42663453, -0.13139241],
        [ 0.07740241,  0.85891184],
        [ 0.53353555,  0.18889505]],

       [[ 0.58224752,  0.4426039 ],
        [-1.12244262, -1.13745509],
        [ 0.51867566, -2.20516368],
        [ 0.35272099, -1.65068832],
        [-1.71926581,  1.27860907]],

       [[ 1.39123195, -0.30607486],
        [-1.84225633,  0.94182099],
        [ 0.42040044, -1.20078669],
        [-2.17904593, -0.39586585],
        [-0.74435111, -0.15761743]],

       [[ 0.60849668, -0.48663983],
        [ 0.54772842, -0.78857347],
        [ 2.09177508,  0.77323564],
        [-0.54413357,  1.37437698],
        [-0.31722618, -0.64333597]],

       [[ 0.33405526, -0.54623557],
        [ 0.83788427,  0.34235314],
        [ 0.22081092,  0.93854724],
        [-0.388348  ,  0.64920496],
        [ 0.03855255,  1.2937604 ]]])
Coordinates:
  * dim_0    (dim_0) int64 0 1 2 3 4 5 6 7 8 9 10 11
  * dim_1    (dim_1) int64 0 1 2 3 4
  * dim_2    (dim_2) int64 0 1

In [15]:

    

time, lat, lon = var1.shape
correlation = var1_cor * var2_cor / time


    

In [18]:

    

cor = correlation.groupby('dim_0').sum()


    

In [19]:

    

cor.plot()


    

    
Out[19]:





[<matplotlib.lines.Line2D at 0x1f5f3ee1630>]






    

In [ ]: