Ejemplo n.º 1
0
def sunhis2initial(hisfile,icfile):
    """
    Main function
    """

    print '#####\nCreating initial condition file (%s) from history file (%s)...'%(icfile,hisfile)

    # Load the history file
    sunhis = Spatial(hisfile, tstep=-1, klayer=[-99])

    # Load the initial condition object
    timeic = datetime.strftime(sunhis.time[-1],'%Y%m%d.%H%M%S')
    sunic = InitialCond(hisfile,timeic)

    # Load the last time step into each variable
    sunic.h = sunhis.loadData(variable='eta').reshape((1,sunic.h.shape))
    sunic.T = sunhis.loadData(variable='temp').reshape((1,)+sunic.T.shape)
    sunic.S = sunhis.loadData(variable='salt').reshape((1,)+sunic.S.shape)
    sunic.uc = sunhis.loadData(variable='uc').reshape((1,)+sunic.uc.shape)
    sunic.vc = sunhis.loadData(variable='vc').reshape((1,)+sunic.vc.shape)

    print sunic.h.shape
    print sunic.T.shape

    # Load the age
    if sunhis.hasVar('agec'):
        sunic.agec = sunhis.loadData(variable='agec').reshape((1,)+sunic.agec.shape)       
        sunic.agealpha = sunhis.loadData(variable='agealpha').reshape((1,)+sunic.agealpha.shape)       

    # Write to the output file
    sunic.writeNC(icfile)
Ejemplo n.º 2
0
 def _dumpic(self):
     """
     Dump initial condition plots
     
     For each variable (uc,vc,T,S,eta):
         - surface plot
         - seabed plot
     """
     
     varnames = ['uc','vc','temp','salt','eta']
     
     sun = Spatial(self.suntanspath+'/'+self.icfile,klayer=[0])
     
     # Plot the surface variables
     for vv in varnames:
         h=plt.figure()
         sun.variable=vv
         sun.loadData()
         sun.clim = [sun.data.min(),sun.data.max()]
         sun.plot()            
         outfile = '%s/IC_%s_surface.png'%(self.plotdir,vv)
         sun.savefig(outfile)
         del h
         
     # Plot the seabed variables
     sun.klayer=[-1]
     for vv in varnames:
         h=plt.figure()
         sun.variable=vv
         sun.loadData()
         sun.clim = [sun.data.min(),sun.data.max()]
         sun.plot()
         outfile = '%s/IC_%s_seabed.png'%(self.plotdir,vv)
         sun.savefig(outfile)  
         del h
Ejemplo n.º 3
0
def sunhis2initial(hisfile, icfile):
    """
    Main function
    """

    print '#####\nCreating initial condition file (%s) from history file (%s)...' % (
        icfile, hisfile)

    # Load the history file
    sunhis = Spatial(hisfile, tstep=-1, klayer=[-99])

    # Load the initial condition object
    timeic = datetime.strftime(sunhis.time[-1], '%Y%m%d.%H%M%S')
    sunic = InitialCond(hisfile, timeic)

    # Load the last time step into each variable
    sunic.h = sunhis.loadData(variable='eta').reshape((1, sunic.h.shape))
    sunic.T = sunhis.loadData(variable='temp').reshape((1, ) + sunic.T.shape)
    sunic.S = sunhis.loadData(variable='salt').reshape((1, ) + sunic.S.shape)
    sunic.uc = sunhis.loadData(variable='uc').reshape((1, ) + sunic.uc.shape)
    sunic.vc = sunhis.loadData(variable='vc').reshape((1, ) + sunic.vc.shape)

    print sunic.h.shape
    print sunic.T.shape

    # Load the age
    if sunhis.hasVar('agec'):
        sunic.agec = sunhis.loadData(variable='agec').reshape((1, ) +
                                                              sunic.agec.shape)
        sunic.agealpha = sunhis.loadData(
            variable='agealpha').reshape((1, ) + sunic.agealpha.shape)

    # Write to the output file
    sunic.writeNC(icfile)
Ejemplo n.º 4
0
    def _dumpic(self):
        """
        Dump initial condition plots
        
        For each variable (uc,vc,T,S,eta):
            - surface plot
            - seabed plot
        """

        varnames = ['uc', 'vc', 'temp', 'salt', 'eta']

        sun = Spatial(self.suntanspath + '/' + self.icfile, klayer=[0])

        # Plot the surface variables
        for vv in varnames:
            h = plt.figure()
            sun.variable = vv
            sun.loadData()
            sun.clim = [sun.data.min(), sun.data.max()]
            sun.plot()
            outfile = '%s/IC_%s_surface.png' % (self.plotdir, vv)
            sun.savefig(outfile)
            del h

        # Plot the seabed variables
        sun.klayer = [-1]
        for vv in varnames:
            h = plt.figure()
            sun.variable = vv
            sun.loadData()
            sun.clim = [sun.data.min(), sun.data.max()]
            sun.plot()
            outfile = '%s/IC_%s_seabed.png' % (self.plotdir, vv)
            sun.savefig(outfile)
            del h
Ejemplo n.º 5
0
 def loadData(self, variable=None):
     """
     Overloaded loadData function - updates the unstructured grid object
     """
     Spatial.loadData(self, variable=variable)
     if self.is3D:
         self.data=np.ravel(self.data[self.mask3D])
     else:
         self.data=np.ravel(self.data)
         
     self.ug.cell_data.scalars = self.data
     self.ug.cell_data.scalars.name = 'suntans_scalar' 
     self.ug.modified()      
Ejemplo n.º 6
0
 def loadData(self):
     """
     Overloaded loadData function - updates the unstructured grid object
     """
     Spatial.loadData(self)
     if self.is3D:
         self.data=np.ravel(self.data[self.mask3D])
     else:
         self.data=np.ravel(self.data)
         
     self.ug.cell_data.scalars = self.data
     self.ug.cell_data.scalars.name = 'suntans_scalar' 
     self.ug.modified()      
Ejemplo n.º 7
0
def river_discharge(ncfile,shpfiles):
    """
    Calculates the river flux of all type-2 boundaries located with each polygon 
    specified with a shpfile polygon
    """

    # Load the spatial object
    sun = Spatial(ncfile,klayer=[-99])

    # Indentify the river cells
    ind = np.argwhere(sun.mark==2).ravel()
    #sun.j = ind

    # Calculate the mask region
    if type(shpfiles) != type([]):
        shpfiles = [shpfiles]

    masks = []
    polynames = []
    for shpfile in shpfiles:
        mask, maskpoly = maskShpPoly(sun.xe[ind],sun.ye[ind],shpfile)
        masks.append(mask)
        polynames.append(os.path.splitext(os.path.basename(shpfile))[0])

    # Create a dictionary with output info
    data={}
    for poly  in polynames:
        data.update({poly:{'Q_r':np.zeros((sun.Nt,)),'time':sun.time}})

    print 'Loading the data...'
    #U = sun.loadData(variable='U_F')
    sun.tstep = range(sun.Nt)
    U = np.zeros((sun.Nt,sun.Nkmax,ind.shape[0]))
    for tt in range(sun.Nt):
        sun.tstep = tt 
        tmp = sun.loadData(variable='U_F')
        U[tt,...] = tmp[:,ind]
        print '\t%d of %d...'%(tt,sun.Nt)

    for mask,poly in zip(masks,polynames):
        tmp_dA  = np.sum( np.sum(U*mask,axis=-1), axis=-1)
        data[poly]['Q_r']=tmp_dA


    return data
Ejemplo n.º 8
0
def area_integrate(ncfile,varnames,shpfiles):
    """
    Area integrate a suntans variable for all time in the domain 
    specified with a shpfile polygon
    """
    # Use numexpr to try and speed things up
    import numexpr as ne

    # Load the spatial object
    sun = Spatial(ncfile,klayer=[-99])
    Ac = sun.Ac

    # Calculate the mask region
    if type(shpfiles) != type([]):
        shpfiles = [shpfiles]

    masks = []
    polynames = []
    for shpfile in shpfiles:
        mask, maskpoly = maskShpPoly(sun.xv,sun.yv,shpfile)
        masks.append(mask)
        polynames.append(os.path.splitext(os.path.basename(shpfile))[0])

    # Create a dictionary with output info
    data={}
    for poly  in polynames:
        data.update({poly:{'V':np.zeros((sun.Nt,)),'time':sun.time}})
        for varname in varnames:
            data[poly].update({varname:np.zeros((sun.Nt,))})

    sun.tstep = range(sun.Nt)
    for varname in varnames:
        print 'Area integrating varibles: %s ...'%(varname)
        tmp = sun.loadData(variable=varname)
        for mask,poly in zip(masks,polynames):
            tmp_dA = ne.evaluate("sum(tmp*Ac*mask,axis=1)")
            data[poly][varname]=tmp_dA


    return data
Ejemplo n.º 9
0
def energy_budget(energyfile,polyfile,trange):
    """
    # Area-integrate the energy terms
    """
    varnames = ['KEz','PEz','uP','uKE','uPE','ueta','W_work','B_flux','diss']

    # Load the energy file as a suntans object
    sun = Spatial(energyfile)

    # Create the mask
    mask,maskpoly = maskShpPoly(sun.xv,sun.yv,polyfile)

    # Initialise the output dictionary
    tstep = range(0,sun.Nt)[trange[0]:trange[1]]
    nt = len(tstep)

    budget ={}
    for vv in varnames:
        budget.update({vv:np.zeros((nt,))})

    for ii,tt in enumerate(tstep):
        print 'Area-integrating step: %d of %d...'%(ii,tstep[-1])
        for vv in varnames:
            sun.tstep=[tt]
            data = sun.loadData(variable=vv)
            budget[vv][ii],areatotal = sun.areaint(data,mask)

    budget.update({'time':sun.time[tstep]})

    # Calculate the time-rate of change of KE and PE
    dt = sun.timeraw[1]-sun.timeraw[0]
    budget.update({'dKE_dt':np.zeros((nt,))})
    budget.update({'dPE_dt':np.zeros((nt,))})
    budget['dKE_dt'][1::] = (budget['KEz'][1::]-budget['KEz'][0:-1])/dt
    budget['dPE_dt'][1::] = (budget['PEz'][1::]-budget['PEz'][0:-1])/dt

    return budget
Ejemplo n.º 10
0
    def suntans2ic(self,hisfile,setUV=False,seth=False):
        """
        Uses data from another suntans file as initial conditions

        Data needs to be on the same grid

        """
        # Load the history file
        sunhis = Spatial(hisfile, tstep=-1, klayer=[-99])

        # Set the time step to grab from the history file
        #...
        tstep = sunhis.getTstep(self.time,self.time)
        sunhis.tstep = [tstep[0]]
        print 'Setting the intial condition with time step: %s\nfrom the file:%s'\
            %(datetime.strftime(sunhis.time[tstep[0]],\
            '%Y-%m-%d %H-%M-%S'),hisfile)

        # Npw grab each variable and save in the IC object
        if seth:
            self.h = sunhis.loadData(variable='eta').reshape((1,self.h.shape))

        if sunhis.hasVar('temp'):
            self.T = sunhis.loadData(variable='temp').reshape((1,)+self.T.shape)
        if sunhis.hasVar('salt'):
            self.S = sunhis.loadData(variable='salt').reshape((1,)+self.S.shape)

        if setUV:
            self.uc = sunhis.loadData(variable='uc').reshape((1,)+self.uc.shape)
            self.vc = sunhis.loadData(variable='vc').reshape((1,)+self.vc.shape)

        # Load the age
        if sunhis.hasVar('agec'):
            self.agec = sunhis.loadData(variable='agec').reshape((1,)+self.agec.shape)       
            self.agealpha = sunhis.loadData(variable='agealpha').reshape((1,)+self.agealpha.shape)       

        print 'Done setting initial condition data from file.'
Ejemplo n.º 11
0
wroms = np.load("calcs/mask_bayshelf_weights_radial.npz")["wroms"]
wromsu = wroms[:, 1:]  # don't want to interpolate this
wromsv = wroms[1:, :]

etablend = np.empty_like(gridblend.x_rho)

# Read in SUNTANS output
for File in Files:

    # klayer = 0 is surface, tstep=-99 reads in all time steps
    # Do u, v, eta
    sun = Spatial(File, klayer=2, tstep=-99, variable="eta")  # , clim=clim)
    # sun = Spatial(File, klayer=0, tstep=-99, variable='eta')#, clim=clim)
    datesbay = sun.time  # get dates available in file
    eta = sun.loadData()

    # Loop through times in bay file
    for i, date in enumerate(datesbay):

        # Find the same time from the shelf model output
        # Assumes that there is a time in the shelf model that perfectly aligns with the
        # times in the bay model at some point
        tindshelf = find(datesshelf == date)
        print date
        # pdb.set_trace()
        if not find(datesshelf == date):  # if a shelf time doesn't align, don't use this bay model output
            continue

        # have to do u and v by time step for some reason
        u = Spatial(File, klayer=["surface"], tstep=i, variable="uc").loadData()
Ejemplo n.º 12
0
def suntans2untrim(ncfile, outfile, tstart, tend, grdfile=None):
    """
    Converts a suntans averages netcdf file into untrim format
    for use in particle tracking
    """
    ####
    # Step 1: Load the suntans data object
    ####
    sun = Spatial(ncfile, klayer=[-99])

    # Calculate some other variables
    sun.de = sun.get_edgevar(sun.dv, method='min')
    sun.mark[sun.mark == 5] = 0
    sun.mark[sun.mark == 3] = 2
    sun.facemark = np.zeros((sun.Nc, ), dtype=np.int)

    # Update the grad variable from the ascii grid file if supplied
    if not grdfile == None:
        print('Updating grid with ascii values...')
        grd = Grid(grdfile)
        sun.grad = grd.grad[:, ::-1]

    ###
    # Step 2: Write the grid variables to a netcdf file
    ###
    nc = Dataset(outfile, 'w', format='NETCDF4_CLASSIC')

    # Global variable
    nc.Description = 'UnTRIM history file converted from SUNTANS output'

    # Write the dimensions
    for dd in list(untrim_griddims.keys()):
        if dd == 'time':
            nc.createDimension(untrim_griddims[dd], 0)
        elif dd == 'numsides':
            nc.createDimension(untrim_griddims[dd], sun.maxfaces)
        else:
            nc.createDimension(untrim_griddims[dd], sun[dd])

    for dd in other_dims:
        nc.createDimension(dd, other_dims[dd])

    ###
    # Step 3: Initialize all of the grid variables
    ###
    def create_nc_var(name, dimensions, attdict,data=None, \
        dtype='f8',zlib=False,complevel=0,fill_value=999999.0):

        tmp=nc.createVariable(name, dtype, dimensions,\
            zlib=zlib,complevel=complevel,fill_value=fill_value)

        for aa in list(attdict.keys()):
            tmp.setncattr(aa, attdict[aa])

        if not data == None:
            nc.variables[name][:] = data

    # Make sure the masked cells have a value of -1
    mask = sun['cells'].mask.copy()
    sun['cells'][mask] = FILLVALUE
    sun['face'][mask] = FILLVALUE

    for vv in list(untrim_gridvars.keys()):
        vname = untrim_gridvars[vv]
        print('Writing grid variable %s (%s)...' % (vname, vv))

        if vv == 'time':
            continue

        # add dz_min attribute to z_r variable
        if vv == 'z_r':
            ugrid[vname]['attributes'].update({'dz_min': 1e-5})
            #sun[vv][:]=sun[vv][::-1]
            sun[vv][:] = sun['z_w'][0:-1][::-1]

        # Reverse the order of grad(???)
        if vv == 'grad':
            sun[vv][:] = sun[vv][:, ::-1]

        ## Fix one-based indexing
        #if vv in ['cells','edges','grad']:
        #    mask = sun[vv][:]==-1
        #    tmp = sun[vv][:]+1
        #    tmp[mask]=-1
        #    #sun[vv][:]=sun[vv][:]+1
        #    create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
        #        data=tmp,dtype=ugrid[vname]['dtype'])

        create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            data=sun[vv],dtype=ugrid[vname]['dtype'])

    # Initialize the two time variables
    vname = untrim_gridvars['time']
    create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            dtype=ugrid[vname]['dtype'])
    vname = 'Mesh2_data_time_string'
    create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            dtype=ugrid[vname]['dtype'])

    ###
    # Step 4: Initialize all of the time-varying variables (but don't write)
    ###
    for vname in varnames:
        print('Creating variable %s...' % (vname))

        create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            dtype=ugrid[vname]['dtype'],zlib=True,complevel=1,fill_value=999999.)

    ###
    # Step 5: Loop through all of the time steps and write the variables
    ###
    tsteps = sun.getTstep(tstart, tend)
    tdays = othertime.DaysSince(sun.time, basetime=datetime(1899, 12, 31))
    for ii, tt in enumerate(tsteps):
        # Convert the time to the untrim formats
        timestr = datetime.strftime(sun.time[tt], '%Y-%m-%d %H:%M:%S')

        print('Writing data at time %s (%d of %d)...' %
              (timestr, tt, tsteps[-1]))

        #Write the time variables
        nc.variables['Mesh2_data_time'][ii] = tdays[ii]
        nc.variables['Mesh2_data_time_string'][:, ii] = timestr

        # Load each variable or calculate it and convert it to the untrim format
        sun.tstep = [tt]

        ###
        # Compute a few terms first
        eta = sun.loadData(variable='eta')
        U = sun.loadData(variable='U_F')
        dzz = sun.getdzz(eta)
        dzf = sun.getdzf(eta)

        vname = 'Mesh2_sea_surface_elevation'
        #print '\tVariable: %s...'%vname
        nc.variables[vname][:, ii] = eta

        vname = 'Mesh2_salinity_3d'
        #print '\tVariable: %s...'%vname
        tmp3d = sun.loadData(variable='salt')
        nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]

        vname = 'Mesh2_vertical_diffusivity_3d'
        #print '\tVariable: %s...'%vname
        tmp3d = sun.loadData(variable='nu_v')
        nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]

        vname = 'h_flow_avg'
        #print '\tVariable: %s...'%vname
        nc.variables[vname][:, :, ii] = U.swapaxes(0, 1)[:, ::-1]

        vname = 'v_flow_avg'
        #print '\tVariable: %s...'%vname
        tmp3d = sun.loadData(variable='w') * sun.Ac  # m^3/s
        nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]

        # Need to calculate a few terms for the other variables

        vname = 'Mesh2_edge_wet_area'
        #print '\tVariable: %s...'%vname
        #dzf = sun.loadData(variable='dzf')
        tmp3d = dzf * sun.df
        nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]

        vname = 'Mesh2_face_water_volume'
        #print '\tVariable: %s...'%vname
        #dzz = sun.loadData(variable='dzz')
        tmp3d = dzz * sun.Ac
        nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]

        vname = 'Mesh2_face_wet_area'
        #print '\tVariable: %s...'%vname
        tmp3d = np.repeat(sun.Ac[np.newaxis, ...], sun.Nkmax, axis=0)
        nc.variables[vname][:, :, ii] = tmp3d.swapaxes(0, 1)[:, ::-1]

        # UnTRIM references from bottom to top i.e.
        # k = 0 @ bed ; k = Nkmax-1 @ top

        vname = 'Mesh2_edge_bottom_layer'
        #print '\tVariable: %s...'%vname
        #tmp2d = sun.Nkmax-sun.Nke # zero based
        tmp2d = sun.Nkmax - sun.Nke + 1  # one based
        nc.variables[vname][:, ii] = tmp2d

        vname = 'Mesh2_edge_top_layer'
        #print '\tVariable: %s...'%vname
        etop = sun.loadData(variable='etop')
        #tmp2d = sun.Nkmax-etop-1 # zero based
        tmp2d = sun.Nkmax - etop  # one based
        nc.variables[vname][:, ii] = tmp2d

        vname = 'Mesh2_face_bottom_layer'
        #print '\tVariable: %s...'%vname
        #tmp2d = sun.Nkmax-sun.Nk + 1 # zero based
        tmp2d = sun.Nkmax - sun.Nk  # one based
        nc.variables[vname][:, ii] = tmp2d

        vname = 'Mesh2_face_top_layer'
        #print '\tVariable: %s...'%vname
        ctop = sun.loadData(variable='ctop')
        #tmp2d = sun.Nkmax-ctop-1 # zero based
        tmp2d = sun.Nkmax - ctop  # one based
        nc.variables[vname][:, ii] = tmp2d

    print(72 * '#')
    print('\t Finished SUNTANS->UnTRIM conversion')
    print(72 * '#')

    # close the file
    nc.close()
Ejemplo n.º 13
0
"""
Generates a series of animations
"""

from sunpy import Spatial
import numpy as np
import matplotlib.pyplot as plt

runname='SFBay3D'
ncfile = '%s/%s_0*.nc'%('data',runname)
k = 1 # depth layer

plt.figure()
sun = Spatial(ncfile,klayer=[k],variable='salt')
sun.tstep=np.arange(0,len(sun.time))
sun.loadData()
sun.clim = [28.0,32.0]
sun.animate(vector_overlay=False)
sun.saveanim('plots/%s_salt.mov'%runname)

sun.clim=None
plt.figure()
sun = Spatial(ncfile,klayer=[k],variable='uc')
sun.tstep=np.arange(0,len(sun.time))
sun.loadData()
sun.animate(vector_overlay=False)
sun.saveanim('plots/%s_uc.mov'%runname)

sun.clim=None
plt.figure()
sun = Spatial(ncfile,klayer=[k],variable='eta')
Ejemplo n.º 14
0
def suntans2untrim(ncfile,outfile,tstart,tend,grdfile=None):
    """
    Converts a suntans averages netcdf file into untrim format
    for use in particle tracking
    """
    ####
    # Step 1: Load the suntans data object
    ####
    sun = Spatial(ncfile,klayer=[-99])

    # Calculate some other variables
    sun.de = sun.get_edgevar(sun.dv,method='min')
    sun.mark[sun.mark==5]=0
    sun.mark[sun.mark==3]=2
    sun.facemark = np.zeros((sun.Nc,),dtype=np.int)

    # Update the grad variable from the ascii grid file if supplied
    if not grdfile == None:
        print 'Updating grid with ascii values...'
        grd = Grid(grdfile)
        sun.grad = grd.grad[:,::-1]

    ###
    # Step 2: Write the grid variables to a netcdf file
    ###
    nc = Dataset(outfile,'w',format='NETCDF4_CLASSIC')

    # Global variable
    nc.Description = 'UnTRIM history file converted from SUNTANS output'

    # Write the dimensions
    for dd in untrim_griddims.keys():
        if dd == 'time':
            nc.createDimension(untrim_griddims[dd],0)
        elif dd =='numsides':
            nc.createDimension(untrim_griddims[dd],sun.maxfaces)
        else:
            nc.createDimension(untrim_griddims[dd],sun[dd])


    for dd in other_dims:
        nc.createDimension(dd,other_dims[dd])

    ###
    # Step 3: Initialize all of the grid variables
    ###
    def create_nc_var(name, dimensions, attdict,data=None, \
        dtype='f8',zlib=False,complevel=0,fill_value=999999.0):
            
        tmp=nc.createVariable(name, dtype, dimensions,\
            zlib=zlib,complevel=complevel,fill_value=fill_value)

        for aa in attdict.keys():
            tmp.setncattr(aa,attdict[aa])
        
        if not data==None:
            nc.variables[name][:] = data
     
    # Make sure the masked cells have a value of -1
    mask = sun['cells'].mask.copy()
    sun['cells'][mask]=FILLVALUE
    sun['face'][mask]=FILLVALUE

    for vv in untrim_gridvars.keys():
        vname = untrim_gridvars[vv]
        print 'Writing grid variable %s (%s)...'%(vname,vv)

        if vv=='time':
            continue

        # add dz_min attribute to z_r variable
        if vv == 'z_r':
            ugrid[vname]['attributes'].update({'dz_min':1e-5})
            #sun[vv][:]=sun[vv][::-1]
            sun[vv][:]=sun['z_w'][0:-1][::-1]

        # Reverse the order of grad(???)
        if vv=='grad':
            sun[vv][:]=sun[vv][:,::-1]

        ## Fix one-based indexing
        #if vv in ['cells','edges','grad']:
        #    mask = sun[vv][:]==-1
        #    tmp = sun[vv][:]+1
        #    tmp[mask]=-1
        #    #sun[vv][:]=sun[vv][:]+1
        #    create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
        #        data=tmp,dtype=ugrid[vname]['dtype'])

        create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            data=sun[vv],dtype=ugrid[vname]['dtype'])

            
    # Initialize the two time variables
    vname=untrim_gridvars['time']
    create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            dtype=ugrid[vname]['dtype'])
    vname = 'Mesh2_data_time_string'
    create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            dtype=ugrid[vname]['dtype'])

    ###
    # Step 4: Initialize all of the time-varying variables (but don't write) 
    ###
    for vname  in varnames:
        print 'Creating variable %s...'%(vname)

        create_nc_var(vname,ugrid[vname]['dimensions'],ugrid[vname]['attributes'],\
            dtype=ugrid[vname]['dtype'],zlib=True,complevel=1,fill_value=999999.)

    ###
    # Step 5: Loop through all of the time steps and write the variables
    ###
    tsteps = sun.getTstep(tstart,tend)
    tdays = othertime.DaysSince(sun.time,basetime=datetime(1899,12,31))
    for ii, tt in enumerate(tsteps):
        # Convert the time to the untrim formats
        timestr = datetime.strftime(sun.time[tt],'%Y-%m-%d %H:%M:%S')

        print 'Writing data at time %s (%d of %d)...'%(timestr,tt,tsteps[-1]) 

        #Write the time variables
        nc.variables['Mesh2_data_time'][ii]=tdays[ii]
        nc.variables['Mesh2_data_time_string'][:,ii]=timestr

        # Load each variable or calculate it and convert it to the untrim format
        sun.tstep=[tt]

        ###
        # Compute a few terms first
        eta = sun.loadData(variable='eta' )  
        U = sun.loadData(variable='U_F' )  
        dzz = sun.getdzz(eta)
        dzf = sun.getdzf(eta)

        
        vname='Mesh2_sea_surface_elevation'
        #print '\tVariable: %s...'%vname
        nc.variables[vname][:,ii]=eta

        vname = 'Mesh2_salinity_3d'
        #print '\tVariable: %s...'%vname
        tmp3d = sun.loadData(variable='salt' )  
        nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]

        vname = 'Mesh2_vertical_diffusivity_3d'
        #print '\tVariable: %s...'%vname
        tmp3d = sun.loadData(variable='nu_v' )  
        nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]

        vname = 'h_flow_avg'
        #print '\tVariable: %s...'%vname
        nc.variables[vname][:,:,ii]=U.swapaxes(0,1)[:,::-1]


        vname = 'v_flow_avg'
        #print '\tVariable: %s...'%vname
        tmp3d = sun.loadData(variable='w' ) * sun.Ac # m^3/s
        nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]

        # Need to calculate a few terms for the other variables

        vname = 'Mesh2_edge_wet_area'
        #print '\tVariable: %s...'%vname
        #dzf = sun.loadData(variable='dzf')
        tmp3d = dzf*sun.df
        nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]

        vname = 'Mesh2_face_water_volume'
        #print '\tVariable: %s...'%vname
        #dzz = sun.loadData(variable='dzz')
        tmp3d = dzz*sun.Ac
        nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]

        vname = 'Mesh2_face_wet_area'
        #print '\tVariable: %s...'%vname
        tmp3d = np.repeat(sun.Ac[np.newaxis,...],sun.Nkmax,axis=0)
        nc.variables[vname][:,:,ii]=tmp3d.swapaxes(0,1)[:,::-1]

        # UnTRIM references from bottom to top i.e.
        # k = 0 @ bed ; k = Nkmax-1 @ top

        vname = 'Mesh2_edge_bottom_layer'
        #print '\tVariable: %s...'%vname
        #tmp2d = sun.Nkmax-sun.Nke # zero based
        tmp2d = sun.Nkmax-sun.Nke+1 # one based
        nc.variables[vname][:,ii]=tmp2d

        vname = 'Mesh2_edge_top_layer'
        #print '\tVariable: %s...'%vname
        etop = sun.loadData(variable='etop')
        #tmp2d = sun.Nkmax-etop-1 # zero based
        tmp2d = sun.Nkmax-etop # one based
        nc.variables[vname][:,ii]=tmp2d

        vname = 'Mesh2_face_bottom_layer'
        #print '\tVariable: %s...'%vname
        #tmp2d = sun.Nkmax-sun.Nk + 1 # zero based
        tmp2d = sun.Nkmax-sun.Nk # one based
        nc.variables[vname][:,ii]=tmp2d

        vname = 'Mesh2_face_top_layer'
        #print '\tVariable: %s...'%vname
        ctop = sun.loadData(variable='ctop')
        #tmp2d = sun.Nkmax-ctop-1 # zero based
        tmp2d = sun.Nkmax-ctop # one based
        nc.variables[vname][:,ii]=tmp2d

    print 72*'#'
    print '\t Finished SUNTANS->UnTRIM conversion'
    print 72*'#'






    # close the file
    nc.close()
Ejemplo n.º 15
0
def volume_integrate(ncfile,varnames,shpfiles,constantdzz=False):
    """
    Volume integrate a suntans variable for all time in the domain 
    specified with a shpfile polygon
    """
    # Use numexpr to try and speed things up
    import numexpr as ne

    # Load the spatial object
    sun = Spatial(ncfile,klayer=[-99])
    Ac = sun.Ac

    # Calculate the mask region
    if type(shpfiles) != type([]):
        shpfiles = [shpfiles]

    masks = []
    polynames = []
    for shpfile in shpfiles:
        mask, maskpoly = maskShpPoly(sun.xv,sun.yv,shpfile)
        masks.append(mask)
        polynames.append(os.path.splitext(os.path.basename(shpfile))[0])

    # Create a dictionary with output info
    data={}
    for poly  in polynames:
        data.update({poly:{'V':np.zeros((sun.Nt,)),'time':sun.time}})
        for varname in varnames:
            data[poly].update({varname:np.zeros((sun.Nt,))})

    # Fix dzz for testing
    sun.tstep = [0]
    dzz = sun.loadData(variable='dzz')
    h = ne.evaluate("sum(dzz,axis=0)")

    #dzz = np.repeat(sun.dz[:,np.newaxis],sun.Nc,axis=1)
    for ii in range(sun.Nt):
        sun.tstep = [ii]
        print 'Volume integrating for time step: %d of %d...'%(ii,sun.Nt)

        # Load the depth and mean age arrays
        if not constantdzz:
            dzz = sun.loadData(variable='dzz')
            # Calculate the total volume 
            #h = np.sum(dzz,axis=0)
            h = ne.evaluate("sum(dzz,axis=0)")
        
        for varname in varnames:
            tmp = sun.loadData(variable=varname)

            for mask,poly in zip(masks,polynames):
                V, A = sun.areaint(h,mask=mask) # Depth*area
                data[poly]['V'][ii] = V

                # Get the depth-integral 
                #tmp_dz = sun.depthint(tmp,dz=dzz)
                tmp_dz = ne.evaluate("sum(tmp*dzz,axis=0)")
                
                # Calculate the volume-integral
                #tmp_dV, A = sun.areaint(tmp_dz,mask=mask)
                tmp_dV = ne.evaluate("sum(tmp_dz*Ac*mask)")

                data[poly][varname][ii]=tmp_dV/V

    return data