### Obtain summer (Jun-Sep, JJAS) average rainfall over India
#### and average sea surface temperature (SST) over the tropics 25N - 25S



### Read in JJAS SST average and JJAS AISMR

nrows=72
ncols=36
ntime = 59    #JJAS 1951 - JJAS 2009
ntimep = 59   # JJAS 1951 - JJAS 2009
N = nrows*ncols


### Lat - Long grid..
ygrid=seq(-87.5,87.5,by=5)
ny=length(ygrid)

xgrid=seq(27.5,382.5,by=5)
#xgrid[xgrid > 180]=xgrid[xgrid > 180]-360	#longitude on 0-360 grid if needed
xgrid[xgrid > 180]=xgrid[xgrid > 180]
nx=length(xgrid)

xygrid=matrix(0,nrow=nx*ny,ncol=2)

i=0
for(iy in 1:ny){
for(ix in 1:nx){
i=i+1
xygrid[i,1]=ygrid[iy]
xygrid[i,2]=xgrid[ix]
}

}

# REad Kaplan SST data.. for 1951 - 2009

data=readBin("Kaplan-SST-JJAS1951-JJAS2009.r4",what="numeric", n=( nrows * ncols * ntime), size=4,endian="swap")


data <- array(data = data, dim=c( nrows, ncols, ntime ) )


data1=data[,,1]
# Missing value is NaN, put it to a large number..
data1[data1 == "NaN"]=1e+30	

# the lat -long data grid..

index=1:(nx*ny)

index1=index[data1 < 20]	# only non-missing data.
xygrid1=xygrid[index1,]
x1=xygrid1[,2]

#x1[x1 < 0]= x1[x1 < 0] + 360
#xygrid1[,2]=x1

nsites=length(index1)	# locations with data -i.e. global locations
data2=data1[index1]

### SSTdata matrix - rows are years and columns are locations on the globe with data
sstdata=matrix(NA,nrow=ntimep, ncol=nsites)


for(i in 1:ntimep){
data1=data[,,i]
data1[data1 == "NaN"]=1e+30	
index1=index[data1 < 20]
data2=data1[index1]
sstdata[i,]=data2
}

## Index of locations corresponding to Global Tropics
### grids with non-missing data and between 25N-25S
xlongs = xygrid[,2]
ylats = xygrid[,1]
indextrop = index[data1 < 20 & ylats >= -25 & ylats <= 25]

rm("data")	#remove the object data to clear up space

### global Tropics
xlongs = xygrid1[,2]
ylats = xygrid1[,1]
xlong = xlongs[ylats >= -25 & ylats <= 25]
ylat = ylats[ylats >= -25 & ylats <= 25]
index1=1:length(xlongs)
index = index1[ylats >= -25 & ylats <= 25]

### Tropical seasonal average.. - the data already is seasonal average
sstanavgtrop = sstdata[,index]

xygridgp=cbind(ylat,xlong)

## write out the grid locations.. if needed
#write(t(xygridp),file="kaplan-sst-pac-locs.txt",ncol=2)

### set the data grid and the index
xygridtemp = xygrid
indextemp = indextrop   ##replace with index if you want global
###################### REAd in RAjeevan JJAS AISMR

### Read in Rajeevan Data (summer JJAS average) 1951 - 2009 and perform PCA

nrows=35
ncols=33
ntime = 59    #JJAS 1951 - JJAS 2009
ntimep = 59   # JJAS 1951 - JJAS 2000
N = nrows*ncols


### Lat - Long grid..

ygrid=seq(6.5,38.5,by=1)
ny=length(ygrid)

xgrid=seq(66.5,100.5,by=1)
nx=length(xgrid)

xygrid=matrix(0,nrow=nx*ny,ncol=2)

i=0
for(iy in 1:ny){
for(ix in 1:nx){
i=i+1
xygrid[i,1]=ygrid[iy]
xygrid[i,2]=xgrid[ix]
}

}

# REad Rajeevan seasonal average data..

data=readBin("Rajeevan-AISMR-JJAS1951-JJAS2009.r4",what="numeric", n=( nrows * ncols * ntime), size=4,endian="swap")


data <- array(data = data, dim=c( nrows, ncols, ntime ) )


data1=data[,,1]
# Missing value is NaN, change it to a negative number.
data1[data1 == "NaN"]=-99999.
# the lat -long data grid..

index=1:(nx*ny)

index1=index[data1 >= 0]	# only non-missing data.
xygrid1=xygrid[index1,]
x1=xygrid1[,2]

nsites=length(index1)	# locations with data -i.e. global locations
data2=data1[index1]

### Rain data matrix - rows are years and columns are locations on the grid over India
raindata=matrix(NA,nrow=ntimep, ncol=nsites)


for(i in 1:ntimep){
data1=data[,,i]
data1[data1 == "NaN"]=-99999.
index1=index[data1 >= 0]
data2=data1[index1]
raindata[i,]=data2
}

indexgrid = index1

rm("data")	#remove the object data to clear up space

## set the data grid and index
raingrid = xygrid1
xygridrain = xygrid
indexrain = index1
##########################################################

#PCA of SSTs
zs=var(sstanavgtrop)

#do an Eigen decomposition..
zsvd=svd(zs)

ssteof = zsvd$u

#Principal Components...
pcs= sstanavgtrop %*% zsvd$u

#Eigen Values.. - fraction variance 
lambdas=(zsvd$d/sum(zsvd$d))

npc = 10		# number of PCs to use for CCA
sstPC = pcs[,1:10]
sstpcfull = pcs

########## 
################ PCA on raindata

rainMean = apply(raindata, 2, mean)
rainSd = apply(raindata, 2, sd)


raindata1=scale(raindata)
zs=var(raindata1)

#do an Eigen decomposition..
zsvd=svd(zs)

raineof = zsvd$u

#Principal Components...
pcs=raindata1 %*% zsvd$u
rainpcfull=pcs
rainPC = pcs[,1:npc]


M=dim(sstPC)[2]
J=dim(rainPC)[2]
J=min(M,J)
N = length(rainPC[,1])
Qx1 = qr.Q(qr(sstPC))
Qy1 = qr.Q(qr(rainPC))
T11 = qr.R(qr(sstPC))
T22 = qr.R(qr(rainPC))
VV=t(Qx1) %*% Qy1
BB = svd(VV)$v
AA = svd(VV)$u
BB = solve(T22) %*% svd(VV)$v * sqrt(N-1)
wm1 = rainPC %*% BB
AA = solve(T11) %*% svd(VV)$u * sqrt(N-1)
vm1 = sstPC %*% AA
cancorln = svd(VV)$d[1:J] #canonical correlation
Fyy = var(rainPC) %*% BB
#........................................ Prediction.

### Predict the Rain PCS
betahat = solve(t(AA) %*% t(sstPC)%*% sstPC %*% AA) %*% t(AA) %*% t(sstPC) %*% rainPC
ypred=sstPC %*% AA %*% betahat


### first npc PCs from the PC forecast above and the remaining PCs are set to
## their means - i.e., 0

N1 = dim(raindata1)[2]-(npc)

rainPCpred = cbind(ypred,matrix(rep(0,N),ncol=N1,nrow=N))

## back transform to get the rainfall field
#rainpred = rainPCpred %*% raineof
#rainPCpred = cbind(ypred,rainpcfull[,npc1:dim(raindata1)[2]])


###  Keep only the first npc Eigen Vectors and set rest to zero
E = matrix(0,nrow=dim(raindata1)[2],ncol=dim(raindata1)[2])
E[,1:npc]=raineof[,1:npc]

### first npc PCs from the PC forecast above and the remaining PCs are set to
## their means - i.e., 0

N1 = dim(raindata1)[2]-(npc)
rainPCpred = cbind(ypred,matrix(rep(0,N),ncol=N1,nrow=N))

## back transform to get the rainfall field anomalies
rainpred =  rainPCpred %*% t(E)

### rescale the rainfall
rainpred=t(t(rainpred)*rainSd + rainMean)


### correlate the forecasted PCs with the historical PCs
pccor = diag(cor(ypred,rainPC[,1:npc]))

## correlate the forecasted rainfall with the historical rainfall
raincor = diag(cor(rainpred,raindata))


## map the correlations
xlong = sort(unique(xygridrain[,2]))
ylat = sort(unique(xygridrain[,1]))
nrows=35
ncols=33
nglobe = nrows*ncols   # also equal to 35*33

zfull = rep(NaN,nglobe)
zfull[indexrain]=raincor      
zmat = matrix(zfull,nrow=nrows,ncol=ncols)

image.plot(xlong,ylat,zmat,ylim=range(2,40),zlim=range(0,0.7))
contour(xlong,ylat,(zmat),ylim=range(2,40),add=TRUE,nlev=6,lwd=2)
map('world2',add=TRUE)
grid(col="black",lty=1)

###############################################

### Forecast 2014-2015

nrows=72
ncols=36
ntime = 2    #JJAS 2014 - 2015 
ntimep = 2   # 
N = nrows*ncols


### Lat - Long grid..
ygrid=seq(-87.5,87.5,by=5)
ny=length(ygrid)

xgrid=seq(27.5,382.5,by=5)
#xgrid[xgrid > 180]=xgrid[xgrid > 180]-360	#longitude on 0-360 grid if needed
xgrid[xgrid > 180]=xgrid[xgrid > 180]
nx=length(xgrid)

xygrid=matrix(0,nrow=nx*ny,ncol=2)

i=0
for(iy in 1:ny){
for(ix in 1:nx){
i=i+1
xygrid[i,1]=ygrid[iy]
xygrid[i,2]=xgrid[ix]
}

}

# REad Kaplan SST data.. for 2014 - 2015

data=readBin("Kaplan-SST-JJAS2014-2015.r4",what="numeric", n=( nrows * ncols * ntime), size=4,endian="swap")

data <- array(data = data, dim=c( nrows, ncols, ntime ) )


data1=data[,,1]
# Missing value is NaN, put it to a large number..
data1[data1 == "NaN"]=1e+30	

# the lat -long data grid..

index=1:(nx*ny)

index1=index[data1 < 20]	# only non-missing data.
xygrid1=xygrid[index1,]
x1=xygrid1[,2]

#x1[x1 < 0]= x1[x1 < 0] + 360
#xygrid1[,2]=x1

nsites=length(index1)	# locations with data -i.e. global locations
data2=data1[index1]

### SSTdata matrix - rows are years and columns are locations on the globe with data
sstdata=matrix(NA,nrow=ntimep, ncol=nsites)


for(i in 1:ntimep){
data1=data[,,i]
data1[data1 == "NaN"]=1e+30	
index1=index[data1 < 20]
data2=data1[index1]
sstdata[i,]=data2
}

## Index of locations corresponding to Global Tropics
### grids with non-missing data and between 25N-25S
xlongs = xygrid[,2]
ylats = xygrid[,1]
indextrop = index[data1 < 20 & ylats >= -25 & ylats <= 25]

rm("data")	#remove the object data to clear up space

### global Tropics
xlongs = xygrid1[,2]
ylats = xygrid1[,1]
xlong = xlongs[ylats >= -25 & ylats <= 25]
ylat = ylats[ylats >= -25 & ylats <= 25]
index1=1:length(xlongs)
index = index1[ylats >= -25 & ylats <= 25]

### Tropical seasonal average.. - the data already is seasonal average
sstpc1415 = sstdata[,index] %*% ssteof
sstpc1415 = sstpc1415[,1:npc]

## predict the PCs
ypred1415=sstpc1415 %*% AA %*% betahat

###  Keep only the first npc Eigen Vectors and set rest to zero
E = matrix(0,nrow=dim(raindata1)[2],ncol=dim(raindata1)[2])
E[,1:npc]=raineof[,1:npc]

### first npc PCs from the PC forecast above and the remaining PCs are set to
## their means - i.e., 0

N1 = dim(raindata1)[2]-(npc)
rainPCpred = cbind(ypred1415,matrix(rep(0,2),ncol=N1,nrow=2))

## back transform to get the rainfall field anomalies
rainpred =  rainPCpred %*% t(E)

### rescale the rainfall
rainpred=t(t(rainpred)*rainSd + rainMean)


## map rainfall for 2014 and 2015

xlong = sort(unique(xygridrain[,2]))
ylat = sort(unique(xygridrain[,1]))
nrows=35
ncols=33
nglobe = nrows*ncols   # also equal to 35*33

zfull = rep(NaN,nglobe)
zfull[indexrain]=rainpred[2,]    	## rainpred[2,] for 2015  
zmat = matrix(zfull,nrow=nrows,ncol=ncols)

image.plot(xlong,ylat,zmat,ylim=range(2,40),zlim=range(rainpred))
contour(xlong,ylat,(zmat),ylim=range(2,40),add=TRUE,nlev=6,lwd=2)
map('world2',add=TRUE)
grid(col="black",lty=1)