pylab inline

from grads.galab import GaLab  # for python interface for grads
def readSSTanomaly():
    #read OISST V2 anomary from 1982 to 2010

    ga = GaLab(Bin='grads',Window=False,Echo=False,Verb=0,Port=False)
    
    # read SST
    script="""
    sdfopen sst.mnmean.nc
    set time jan1982 dec2010
    deseason sst sstanom sstclim JAN1982 DEC2010
    """
    ga(script)
    sst=ga.exp("sstanom")
    script2="""
    set lat 0
    set lon 180
    define nino3=aave(sstanom,lon=210,lon=270,lat=-5,lat=5)
    """
    ga(script2)
    nino3=ga.expr("nino3")
    return sst,nino3

sst,nino3=readSSTanomaly()

from numba import autojit,jit
import numpy as np

@autojit # @jit('f8[:,:](f8[:],f8[:,:,:])')
def regress(x,Y):
    """ calculate 2d correlation
    """
    nt,ny,nx=Y.shape
    coef=np.zeros((ny,nx),dtype=np.float)
    for j in xrange(ny):
        for i in xrange(nx):
            coef[j,i]=(np.corrcoef(x,Y[:,j,i]))[0,1]
    return coef

%timeit coef=regress(nino3,sst)

coef=regress(nino3,sst)

def readLandSeaMask():
    """
read land sea mask
    """
    ga = GaLab(Bin='grads',Window=False,Echo=False,Verb=0,Port=False)
    
    # read landseamask
    script="""
    sdfopen lsmask.nc
    """
    ga(script)
    lsmask=ga.exp("MASK")
    return lsmask

lsmask=readLandSeaMask()

out=np.ma.masked_array(coef,lsmask<1)
lon=sst.grid.lon
lat=sst.grid.lat
contourf(lon,lat,out,linspace(-1,1,11))
colorbar()