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
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
