def get_data(name, iter=None):
if iter is None:
return mds.rdmds(name)
else:
return mds.rdmds(name, iter)
##########
firedrake_folder = "/data/2d_mitgcm_comparison/16.04.20_3_eq_param_ufricHJ99_dt600.0_dtOutput3600.0_T43200.0_ip50.0_tres86400.0_Kh0.001_Kv0.0001_dy50_dz1_closed_iterative/"
#really this is dt = 600s
mitgcm_folder = "/data/mitgcm_runs/ben_FRISP_wout_coriolis_T1.0days_original_Ks_check/run_dt30s/"
# really this is constant gamma T and time averaging. Not the final comparison!!!
firedrake_df = pd.read_csv(firedrake_folder + "matplotlib_arrays.csv")
output_folder = "/data/mitgcm_comparison_videos/"
##########
# Get firedrake coordinates
y_array_fd = firedrake_df['y_array']
z_array_fd = firedrake_df['z_array']
# Get mitgcm coordinates
YC = mds.rdmds(mitgcm_folder + 'YC')
RC = mds.rdmds(mitgcm_folder + 'RC')
# number of cells mitgcm run
nx = 1
ny = 200
nz = 102
# Reorganize RC and YC into matrices matching the shape of the T matrix
y_array_mit = np.empty((ny, nz))
z_array_mit = np.empty((ny, nz))
for j in range(nz):
y_array_mit[:, j] = YC[:, 0]
for j in range(ny):
z_array_mit[j, :] = RC[:, 0, 0]
iters = [i.split('.')[2].lstrip("0") for i in glob.glob(name + ".*.data")]
try:
iters.remove('')
except ValueError:
pass
iters = sorted([int(iter) for iter in iters])
return iters
iters = get_iters("./U")
iters_tAvg = get_iters("./uVeltave")
XC = mds.rdmds('XC')
YC = mds.rdmds('YC')
RC = mds.rdmds('RC')
nx, ny = XC.shape
CS = mds.rdmds('AngleCS')
SN = mds.rdmds('AngleSN')
U = mds.rdmds('U', np.inf)
hres = 0.5
lon = np.arange(-180, 180.01, hres) #+ hres/2
lat = np.arange(-90, 90.01, hres) #+ hres/2
lon_out, lat_out = np.meshgrid(lon, lat)
iter = int(sys.argv[1])
var_list = ['U', 'V', 'XC', 'YC', 'RC', 'W', 'CS', 'SN', 'Eta']
ds = open_mdsdataset('.',
prefix=var_list,
iters=iter,
geometry='curvilinear',
ignore_unknown_vars=True,
default_dtype=np.float)
XC = np.array(ds.XC)
YC = np.array(ds.YC)
Eta = np.array(ds.Eta)
RC = -mds.rdmds('RC').flatten()
#print(RC)
r_max = np.amax(RC)
r_min = np.amin(RC)
RC = (R / Rb - 1.0) * ((RC - r_min) / (r_max - r_min)) + 1.0
#RC = RC[::-1]
triang = tri.Triangulation(XC.flatten(), YC.flatten())
U = np.array(ds.U)[0, :, :, :] * 100
V = np.array(ds.V)[0, :, :, :] * 100
W = np.array(ds.W)[0, :, :, :] * 100
#Eta = Eta[0,:,:]
Ut = U[1, :, :]
data_dir_bump = '/Volumes/scratch/tmp/channel/taux2000_rb0110_bump/mds_orig/'
output_dir = '/Volumes/scratch/tmp/channel/taux2000_rb0110_bump/vtk'
#mod_flat = MITgcmmodel.ModelInstance(data_dir_flat, grid_dir=grid_dir)
#mod_bump = MITgcmmodel.ModelInstance(data_dir_bump, grid_dir=grid_dir)
bflat = bump_channel.BumpChannel(base_dir, 'taux2000_rb0110_flat')
bflat.load_default_fields()
bbump = bump_channel.BumpChannel(base_dir, 'taux2000_rb0110_bump')
bbump.load_default_fields()
iters = arange(0, 131520 + 1, 192)
nb = 50
nf = 145
Tf = rdmds(data_dir_flat + 'THETA', iters[nf], lev=0)
Tb = rdmds(data_dir_bump + 'THETA', iters[nb], lev=0)
clevs = arange(0, 8.1, 0.5)
X = linspace(0, 2000, bflat.Nx)
Y = X
myticks = [500, 1000, 1500]
close('all')
rc('figure.subplot',
left=0.08,
right=0.88,
bottom=0.05,
top=0.9,
wspace=0.12,
hspace=0.22)
import sys
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from MITgcmutils import mds
fig = plt.figure()
T = mds.rdmds('./res/T', itrs=np.NaN)
nt, nx, ny, nz = np.shape(T)
print "nt,nx,ny,nz", nt, nx, ny, nz
TT = np.reshape(T, [nt, nx, nz])
TT[TT**2 < 1.0e-10] = np.NaN
i = 200
#im = plt.imshow(f(i), cmap="bwr",animated=True)
im = plt.imshow(TT[i, :, :],
cmap="jet",
interpolation='bilinear',
animated=True)
#im = plt.imshow(TT[i,:,:],animated=True)
#im = plt.imshow(f(i), cmap="jet",animated=True)
#im = plt.imshow(f(i), animated=True)
#ax = Axes3D(fig)
#Z = np.arange(nz)
#X = np.arange(nx)
#Z, X = np.meshgrid(Z, X)
#print "shape of X",np.shape(X)
#print "shape of Z",np.shape(Z)
grid_cs510 = "/home/hay/Research/EuropaOceanDynamics/MITgcm_grids/cs510/"
grid_cs96 = "/home/hay/Research/EuropaOceanDynamics/MITgcm_grids/cs96/"
grid_cs510 = grid_cs96
dir_iterp_matrices = "/home/hay/Research/EuropaOceanDynamics/MITgcm_grids/"
dir_data = "./"
iteration = int(sys.argv[1])
interp_u = load_npz(dir_iterp_matrices + "cs96_to_cs96_uvel.npz")
interp_v = load_npz(dir_iterp_matrices + "cs96_to_cs96_vvel.npz")
interp_eta = load_npz(dir_iterp_matrices + "cs96_to_cs96_eta.npz")
if testing:
CS = mds.rdmds(grid_cs96 + "AngleCS")
SN = mds.rdmds(grid_cs96 + "AngleSN")
X1 = mds.rdmds(grid_cs96 + 'XC').ravel()
Y1 = mds.rdmds(grid_cs96 + 'YC').ravel()
X2 = mds.rdmds(grid_cs510 + 'XC').ravel()
Y2 = mds.rdmds(grid_cs510 + 'YC').ravel()
triang1 = tri.Triangulation(X1, Y1)
triang2 = tri.Triangulation(X2, Y2)
data_u_all = mds.rdmds(dir_data + "U", itrs=iteration)
data_v_all = mds.rdmds(dir_data + "V", itrs=iteration)
data_eta = mds.rdmds(dir_data + "Eta", itrs=iteration)
u_start = np.zeros((Nr, Nx, 6 * Nx))
def get_nearest_interp_matrix(xName, yName, zName='RC', savename="interp_mat.npz", is_2D=False, p=1.4, neighbour_num=6):
Ncs96 = 6*96*96
Ncs510 = 6*510*510
Ncs510 = Ncs96
X = mds.rdmds(grid_cs96 + xName)
Y = mds.rdmds(grid_cs96 + yName)
Z = mds.rdmds(grid_cs96 + zName)
Z = (Z + radius) / radius
if is_2D:
LATr = d2r*(Y+90)
LONr = d2r*(X+180)
R = 1.0
else:
R, LONr = np.meshgrid(Z, d2r*(X+180))
R, LATr = np.meshgrid(Z, d2r*(Y+90))
R = R.ravel()
LATr = LATr.ravel()
LONr = LONr.ravel()
x, y, z = sph2xyz(R, LATr, LONr)
xyz = np.column_stack((x, y, z))
tree = cKDTree( xyz, leafsize=100)
print("GOT TREE!")
# Load coords of higher res grid
X = mds.rdmds(grid_cs510 + xName)
Y = mds.rdmds(grid_cs510 + yName)
Z = mds.rdmds(grid_cs510 + zName)
Z = (Z + radius) / radius
if is_2D:
LATr = d2r*(Y+90)
LONr = d2r*(X+180)
R = 1.0
Nr = 1
else:
# R, LONr = np.meshgrid(Z, d2r*(X+180))
# R, LATr = np.meshgrid(Z, d2r*(Y+90))
LATr = d2r*(Y+90)
LONr = d2r*(X+180)
# R = R.ravel()
R = 1.0
Nr = 1
LATr = LATr.ravel()
LONr = LONr.ravel()
x2, y2, z2 = sph2xyz(R, LATr, LONr)
print(x2.size)
xyz2 = np.column_stack((x2, y2, z2))
print("QUERYING TREE...")
s = time()
distances, ix = tree.query( xyz2, k=neighbour_num, n_jobs=-1)
e = time()
print("GOT TREE QUERY AFTER {:1.2f} SECONDS!".format(e-s))
print(distances.shape[0])
print("COMPUTING WEIGHTS")
w = 1/distances**p
w /= np.sum(w, axis=-1)[:, np.newaxis]
print(w.shape, x.size)
cnt = 0
for i in range(w.shape[0]):
if 0.0 in distances[i]:
# print("Found zero")
at_p = np.zeros(neighbour_num, dtype=np.float32)
at_p[0] = 1.0
w[i] = at_p
# ix[i, 0] = i
print("CREATING SPARSE INTERPOLATION MATRIX")
rows = np.zeros(w.size, dtype=np.int)
cols = np.zeros(w.size, dtype=np.int)
count = 0
for ii in range(x2.size):
for j in range(neighbour_num):
rows[count] = ii
cols[count] = ix[ii, j]
count += 1
weights_mat = coo_matrix( (w.ravel(), (rows, cols)), shape=(Ncs510*Nr, Ncs96*Nr), dtype=np.float32)
weights_mat.eliminate_zeros()
weights_mat = weights_mat.tocsr()
print("SAVING MATRIX")
save_npz(savename, weights_mat)
