def driver3D(args):
ocean_file = args.oceanfile
# ocean grid
D=Dataset('ocean_geometry.nc').variables['D'][:]
lonh=Dataset('ocean_geometry.nc').variables['lonh'][:]
lonq=Dataset('ocean_geometry.nc').variables['lonq'][:]
lath=Dataset('ocean_geometry.nc').variables['lath'][:]
latq=Dataset('ocean_geometry.nc').variables['latq'][:]
lonqs, latqs = np.meshgrid(lonq,latq)
lons, lats = np.meshgrid(lonh,lath)
D=np.ma.masked_where(D <= 1, D)
D.mask = np.ma.array(D); D.mask[:,:]=False
# ice shelf base
ssh = Dataset('ISOMIP_IC.nc').variables['ave_ssh'][0,:,:]
# load time
print('Time indice is:' + str(args.time))
if args.time > 0 :
time = np.array(Dataset(ocean_file).variables['time'][args.time])
tind = [args.time]
else:
time = Dataset(ocean_file).variables['time'][:]
tind = range(len(time))
if args.bergs:
rho_berg = 918.0
rho = 1030.
# mass of bergs can time dependent
mass_berg = Dataset(ocean_file).variables['mass_berg'][tind,:,:]
IS=(mass_berg/rho_berg)
else:
# ice shelf thickness, static for now
IS = Dataset('MOM_Shelf_IC.nc').variables['h_shelf'][:]
# interface and layer thickness
e=Dataset(ocean_file).variables['e'][0,:,:,:]
h=Dataset(ocean_file).variables['h'][0,:,:,:]
# correct top and bottom, for vis pourposes
h[0,:,:]=e[0,:,:]; h[-1,:,:]=e[-1,:,:]
NZ,NY,NX=h.shape
# create VTK bathymetry
VTKgen(lats,lons,D.mask,depth=D,h=h,fname=name)
if not tind>1:
# create VTK ice-shelf
VTKgen(lats,lons,D.mask,h=h,shelf_base=ssh,shelf_thick=IS,fname=name)
time_list=[] # list for time
#tm=2 # number of nc files
ind=0
# loop through time and plot surface fields
for t in range(0,len(tind),args.dt):
print 'Time is:',time[t]
# save data in the lists
#time_list.append(date)
# check if data has been saved
path_to_file = str('VTK/%s-%05d.vtk' % (name,ind))
if os.path.isfile(path_to_file):
print ' \n' + '==> ' + 'FILE EXISTS, MOVING TO THE NEXT ONE ...\n' + ''
ind=ind+1
else:
print 'Saving time indice:', t
# ocean
# structure
# layer thickness
e=Dataset(ocean_file).variables['e'][tind[t],:,:,:]
h_dum=np.abs(np.diff(e,axis=0))
h=0.5*(e[0:-1,:,:]+e[1::,:,:])
# correct top and bottom, for vis pourposes
h[0,:,:]=e[0,:,:]; h[-1,:,:]=e[-1,:,:]
#h=Dataset(ocean_file).variables['h'][t,:,:,:]
# correct top and bottom, for vis pourposes
h[0,:,:]=e[0,:,:]; h[-1,:,:]=e[-1,:,:]
temp=Dataset(ocean_file).variables['temp'][tind[t],:,:,:]
salt=Dataset(ocean_file).variables['salt'][tind[t],:,:,:]
# for isopycnal models
# temp, mark values where h_dum<10 cm with Nan
temp[h_dum<0.001]=np.nan; salt[h_dum<0.001]=np.nan
# same for salt, u and v
v=Dataset(ocean_file).variables['v'][tind[t],:,:,:]
u=Dataset(ocean_file).variables['u'][tind[t],:,:,:]
u=np.ma.masked_where(u>1000,u)
v=np.ma.masked_where(v>1000,v)
u[h_dum<0.01]=np.nan
v[h_dum<0.01]=np.nan
v.fill_value=0.0
u.fill_value=0.0
v=v.filled()
u=u.filled()
# write just bottom values
#VTKgen(lats,lons,D.mask,depth=D,h=h,temp=tt,salt=ss,rho=gamma,u=uu,v=vv,writebottom=True,fname=reg,t=ind)
print 'Saving ocean data... \n'
if args.tracer:
tr1=Dataset(ocean_file).variables['tr1'][tind[t],:,:,:]
tr1[h_dum<0.001]=np.nan
VTKgen(lats,lons,D.mask,h=h,temp=temp,salt=salt,dye=tr1,u=u,v=v,fname=name,t=tind[t])
else:
VTKgen(lats,lons,D.mask,h=h,temp=temp,salt=salt,u=u,v=v,fname=name,t=tind[t])
if args.bergs:
print 'Saving bergs data... \n'
# save ice shelf made of icebergs
VTKgen(lats,lons,D.mask,h=h,shelf_base=e[0,:,:],shelf_thick=IS[t,:,:],fname=name,t=tind[t])
print ' \n' + '==> ' + ' Done saving vtk files!\n' + ''
# the following is for displaying the time in the PNG figures generated by Visit
if args.time==0:
print ' \n' + '==> ' + ' Writting time to asn ascii file... \n' + ''
f = open('VTK/time.txt', 'wb')
for i in range(len(time)):
f.write('%8.4f \n'.format(time[i]))
f.close()
print ' \n' + '==> ' + ' DONE!\n' + ''
