import spacegrids as sg

import numpy as np
import matplotlib.pyplot as plt
figsize(8, 4)

D = sg.info_dict()

D

ls /home/wim/PROJECTS/test_project/

P = sg.Project(D['my_project'])

P.load('O_temp')

P.load('theta')

TEMP = P['DPO']['O_temp'] 
TEMP2 = P['Lev.cdf']['theta']

TEMP # a field

TEMP.shape # show the shape of the numpy array containing the data. This is a 3D field

TEMP.grid # show the grid on which the temp data is defined

TEMP2.grid

TEMP2.grid[0].axis

P.show() # the registered experiments can be viewed here

P['DPO'].show() # the loaded fields can be viewed here

P['DPO'].axes

for c in P['DPO'].axes:
  exec c.name + ' = c'

X

sg.contourf(sg.squeeze(TEMP[Z,0]))
clb = plt.colorbar()
clb.set_label(TEMP.units)

sg.contourf(sg.squeeze(TEMP[X,50])) # a section at index 50.
clb = plt.colorbar()
clb.set_label(TEMP.units)

(TEMP[0,:,:]).shape == (TEMP[Z,0]).shape

P.load('O_sal',slices=(Z,0,X,50))

P['DPO'].show()

P['DPO']['O_sal_sliced'].grid

# Determine zonal mean mT of TEMP.

mTEMP = TEMP/X
mTEMP2 = TEMP2/X

# mask out the world ocean outside the Atlantic with NaN and take zonal mean.

mTEMP_masked=(TEMP[Y,33:].maskout(node = (X,85,Y,30) ) )/X # node is a point in the Atlantic


# --- start figure ---

lbl = ord('a')

pan = 0

height = 2
width = 1

rows = 2
cols = 1


ax = plt.subplot2grid((height, width), (int(np.floor(pan/cols)), pan%cols) )

# use the sg wrapper for the contour function.
cs= sg.contourf(mTEMP[Z,:10],linestyles='solid',showland=True); # slice only upper ocean
plt.tick_params(axis='x',labelbottom='off')
plt.xlabel('')
plt.ylabel('Depth (m)')
tle = plt.title('('+ chr(lbl) +') Global upper ocean temp ')

lbl += 1
pan += 1

ax = plt.subplot2grid((height, width), (int(np.floor(pan/cols)), pan%cols) )

cs= sg.contourf(mTEMP_masked[Z,:10],linestyles='solid',showland=True);  # slice only upper ocean

# Overiding the xlabel assigned by sg:
plt.xlabel('Latitude')
tle = plt.title('('+ chr(lbl) +') Atlantic upper ocean temp ')

mTEMP.grid  # the dimensionality of the zonal mean is less

P['DPO'].insert( ('mtemp', mTEMP ) )

P['DPO'].show()

sg.contourf(sg.squeeze(P['DPO']['O_temp'][Z,0]-P['DPC']['O_temp'][Z,0]),cmap = mpl.cm.RdYlBu_r,levels = range(-8,9,1))
clb = plt.colorbar(ticks = range(-8,9,2))
clb.set_label(TEMP.units)

TEMP2.shape

P.load(['G_kmt','A_slat'])
KMT = P['DPO']['G_kmt']
SAT = P['DPO']['A_slat']
nc = 10
sg.contourf(SAT, num_cont = nc,cmap = mpl.cm.RdYlBu_r)
cs=sg.contour(SAT,linestyles='solid', num_cont = nc, colors = 'k')
plt.clabel(cs,fmt='%1.0f')
sg.contour(sg.finer_field(KMT), colors='k',levels=[0.5,]  )

dmTEMP = mTEMP.regrid(mTEMP2.grid) - mTEMP2

dmTEMP = mTEMP(mTEMP2.grid) - mTEMP2

cont_int = np.arange(-6.,6.,1.)
cmap = mpl.cm.RdYlBu_r	# Colormap red yellow blue reversed
pl1 = sg.contourf(dmTEMP,  levels = cont_int, cmap = cmap);

clb = plt.colorbar(ticks = cont_int)
clb.set_label(TEMP.units)
plt.xlim([-80.,70.])
plt.xticks([-60,-30,0,30,60])
plt.ylabel('Depth (m)')
tle = plt.title(' Zonal avg. T difference model - observations')

dmTEMP.grid # the resulting grid is that of the Lev data

sTEMP = TEMP[Z,0]# slice to obtain surface temperature 
dTEMP = sTEMP - mTEMP[Z,0]


dTEMP = sTEMP - (sg.squeeze(mTEMP[Z,0])).regrid(sTEMP.grid)

pl1 = sg.contourf(sg.squeeze(dTEMP), cmap = mpl.cm.RdYlBu_r)
clb = plt.colorbar()
clb.set_label('C')
tle = plt.title('Sea Surface Temperature difference from zonal mean')

# Give the experiment objects convenient names E, E2:
E = P['DPO'];E2 = P['DPC']
varname = 'O_velY'

# load velocity by netcdf name.
P.load(varname)

# obtain velocity as sg field object V from project.
V = E[varname]
V2 = E2[varname]

# take the meridional stream function by taking zonal grid-integral X*V 
# and the vertical primitive of the result along Z using the pipe.
PSI = Z|(X*V)*1e-6
PSI2 = Z|(X*V2)*1e-6

lvls = np.arange(-100,100,4)
xlims = [-80,70]
ylims = [-4500., -0.]

# --- start figure ---

lbl = ord('a')

pan = 0

height = 2
width = 1

rows = 2
cols = 1

F = PSI
ax = plt.subplot2grid((height, width), (int(np.floor(pan/cols)), pan%cols) )

# use the sg wrapper for the contour function.
cs= sg.contour(F,linestyles='solid',showland=True,levels = lvls,colors='k');

plt.clabel(cs,fmt='%1.1f');

plt.xlim(xlims)
plt.ylim(ylims)
plt.tick_params(axis='x',labelbottom='off')
plt.xlabel('')
plt.ylabel('Depth (m)')
tle = plt.title('('+ chr(lbl) +') MOC Drake Passage open ')

lbl += 1
pan += 1

F = PSI2
ax = plt.subplot2grid((height, width), (int(np.floor(pan/cols)), pan%cols) )

cs= sg.contour(F,linestyles='solid',showland=True,levels = lvls,colors='k');
plt.clabel(cs,fmt='%1.1f');
plt.xlim(xlims)
plt.ylim(ylims)

# Overiding the xlabel assigned by sg:
plt.xlabel('Latitude')
tle = plt.title('('+ chr(lbl) +') MOC Drake Passage closed ')


SAT1=SAT[Y,:40];SAT2=SAT[Y,40:50];SAT3=SAT[Y,50:]

SAT_cat = sg.concatenate([SAT1,SAT2,SAT3])

(SAT_cat/X).draw()

for c in E.cstack:
  exec c.name +' = c'

depth[:] # depth coordinate points

latitude[:10] # latitudinal coordinate points (truncated)

latitude*longitude

latitude*longitude*depth

latitude**2

TEMP.grid

(TEMP[:]).shape

len(depth[:])

TEMP_shuffled = TEMP(latitude*longitude*depth)
TEMP_shuffled.grid

(TEMP_shuffled[:]).shape

F = TEMP*V

F.grid

F2 = TEMP_shuffled*V
F2.grid

V.grid

(V*TEMP).grid

P['Lev.cdf'].cstack

P['DPO'].cstack[:4] # truncated list of coord objects

depth.axis

X*(latitude*longitude)

# load heat flux by netcdf name.
P.load('F_heat')

# obtain oceanic heat flux as sg field object HF from project.
HF = P['DPO']['F_heat']
HF2 = P['DPC']['F_heat']

HF.grid

pl = sg.contourf(HF,levels = range(-200,200,20),cmap = mpl.cm.RdYlBu)
clb = plt.colorbar()
clb.set_label(HF.units)

HF_shifted = sg.roll(HF,coord = X,shift = 50)

pl = sg.contourf(HF_shifted,levels = range(-200,200,20),cmap = mpl.cm.RdYlBu)
clb = plt.colorbar()
clb.set_label(HF.units)

HF_shifted.grid

# Determine the total oceanic poleward heat transport via the y-primitive and X-intergral, and scale to Petawatt.

PHT = Y|(HF*X)*1e-15
PHT2 = Y|(HF2*X)*1e-15

pl1, = sg.plot(PHT,color = 'b');
pl2, = sg.plot(PHT2,color='r');

plt.grid()
plt.xlim([-80.,80.])
plt.xticks([-60,-30,0,30,60])
plt.ylabel('PW')
tle = plt.title(' Ocean PHT (PW) ')
lgnd = plt.legend([pl1,pl2],['DP open','DP closed'],loc=2)

PHT = latitude|(HF*longitude)*1e-15

latitude.der # the derivative on the sphere in the latitudinal direction

X.der

HF_shifted2 = longitude.coord_shift(HF,50)

HF_shifted2.grid

# obtain coord in X direction
X*(HF_shifted2.grid) 

# compare the ndarray content of the X components of the two shifted grids
X*(HF_shifted2.grid) == X*(HF_shifted.grid)

((X*(HF_shifted2.grid))[:] == (X*(HF_shifted.grid))[:]).all()

HF_sliced = HF_shifted[Y,:36]

pl = sg.contourf(HF_sliced,levels = range(-200,200,20),cmap = mpl.cm.RdYlBu)
clb = plt.colorbar()
clb.set_label(HF.units)
tle = plt.title('Heat flux in Southern Ocean')

HF3 = HF[X,:50,Y,:36]

HF3.grid

len(Y*HF3.grid)

sTEMP = longitude.sum(TEMP)
sTEMP.grid

help(longitude.vsum)

my_grid = latitude*longitude
my_grid.vsum(TEMP)[:]

my_grid.vsum(TEMP).grid

TEMP.grid.vsum(TEMP.ones())

TEMP.ones()*(X*Y*Z)

TEMP.grid.vsum(TEMP)/TEMP.grid.vsum(TEMP.ones())

# more succinctly:
TEMP.grid.mean(TEMP)

TEMP/(X*Y*Z)

# Calculate derivative in X direction:
dTEMPdX = X.der(TEMP)
dTEMPdX.grid

# compute the vertical temperature profile and plot it.
vpTEMP = TEMP/(Y*X)
pl, = sg.plot(vpTEMP)
lb = plt.xlabel('degrees C');tle = title('Vertical profile of average ocean temp')

# take the vertical derivative of the vertical temperature profile and plot it.
pl, = sg.plot(Z.der(vpTEMP))
lb = plt.xlabel('degs C per m');tle = plt.title('Vertical temperature gradient')

# test whether nparray content is the same. Excluding the first nan value.
((Z^vpTEMP).value[1:] == (Z.der(vpTEMP)).value[1:]).all()

# test whether the depth coord method vcumsum yields the same result as Z|
(depth.vcumsum(vpTEMP)[:] == (Z|vpTEMP)[:]).all() 

# load the ocean velocity in the X direction.
P.load('O_velX')

# obtain horizontal velocity fields
U = P['DPC']['O_velX']
V = P['DPO']['O_velY']
# create a 2 dimensional vector field defined on 3 dimensional space
my_vector_field = U*V
my_vector_field

U.direction

V.direction

TEMP.direction

# This gives a field
U*U

(U*V)*(U*V)

my_vector_field.direction()

# Instatiate an X-derivative operator object
ddX = sg.Der(X)
ddX

ddX*TEMP

ddX(TEMP)

ddX*ddX

# Create second order X-derivative
d2dX2 = ddX*ddX
# Apply it to the temperature field
d2dX2*TEMP

# Succesive calls of the ddX object
( (d2dX2*TEMP).value == ( ddX(ddX(TEMP)) ).value ).any()

# Create Z-derivative object
ddZ = sg.Der(Z)
# Apply the derivative operator and plot the result
pl, = sg.plot(ddZ*vpTEMP)
lb = plt.xlabel('degs C per meter')

sg.Pick(X)*my_vector_field

ddX = sg.Der(X) # derivative in X
ddY = sg.Der(Y) # etc
ddZ = sg.Der(Z)

pX = sg.Pick(X) # pick the component with direction X from vfield object (projection)
pY = sg.Pick(Y) # etc.
pZ = sg.Pick(Z)

mX = sg.Mean(X) # take zonal mean of any field argument or right-multiplicant.
mY = sg.Mean(Y) # etc
mZ = sg.Mean(Z)

sX = sg.Integ(X) # take integral
sY = sg.Integ(Y)
sZ = sg.Integ(Z)

intX = sg.Prim(X) # take primitive function of any field argument
intY = sg.Prim(Y)
intZ = sg.Prim(Z)

set2X = sg.SetDirection(X) # change the direction attribute of field argument to X
set2Y = sg.SetDirection(Y)
set2Z = sg.SetDirection(Z)
set2scalar = sg.SetDirection(sg.ID())

Avg = mZ*mX*mY  # a 3D average

curl = ddX*set2scalar*pY - ddY*set2scalar*pX      #  curl operator expressed in basic operators

div = ddX*set2scalar*pX + ddY*set2scalar*pY       # divergence 

div3 = ddX*set2scalar*pX + ddY*set2scalar*pY + ddZ*set2scalar*pZ

grad = set2X*ddX + set2Y*ddY    # gradient. To be used on single field objects.


# Take surface slices
sU = U[Z,0]
sV = V[Z,0]
surface_vectors = sU*sV
(div*surface_vectors).direction

pl = sg.contourf(sg.squeeze((div*surface_vectors)[Y,10:-10]))