def dist2WTP_branch5(self, Ttimes):
"""
calculate the distance to the WTP gate for branch 5
"""
## branch 5
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=5)
dist_tem5 = WB.X[1:-1][::-1] * 3.28084 ## unit: ft
ind5 = next((i for i, x in enumerate(Ttimes[1]) if x), None)
dist_tem5[:ind5] = 0
## branch 1
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=1)
dist_tem1 = WB.X[1:-1][::-1] * 3.28084 ## unit: ft
dx = dist_tem1[-3] - dist_tem1[-2]
ind1 = next((i for i, x in enumerate(Ttimes[0]) if x), None)
dist_tem1[:ind1] = 0
dist_tem5[ind5:] += dx + dist_tem1[ind1]
return [dist_tem1, dist_tem5]
def Percentage_branch5(self, Np, Nt, starttime, location_x, x_combined,
Ttimes_avg):
"""
At branch 5
calculate the particle percentage for each transect at the corresponding travel time
"""
Ttime_avg1 = Ttimes_avg[0]
Ttime_avg5 = Ttimes_avg[1]
#### read segment information for branch 5
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=5)
x_branch5 = WB.X #### segment x coordinates for branch 5
#### read segment information for branch 1
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=1)
x_branch1 = WB.X
percentage5 = np.zeros_like(x_branch5)
percentage1 = np.zeros_like(x_branch1)
percentage_combined = np.zeros_like(x_combined)
Ttime_avg_combined = Ttime_avg5.tolist() + Ttime_avg1.tolist(
)[self.DHS5 - 1:]
for ii, tt in enumerate(Ttime_avg_combined):
tt = int(tt)
if tt != 0:
tt += starttime
seg_id = ii + 2
#print ('Calculate particle percentage for time step = %s, segment = %d\n'%(str(tt), seg_id))
icount = 0
for i in range(location_x.shape[0]):
if location_x[i, tt - starttime] >= x_combined[seg_id - 1]:
icount += 1
percentage_combined[seg_id - 1] = float(icount) / Np
percentage5[1:-1] = percentage_combined[1:len(x_branch5) - 1]
percentage1[self.DHS5:] = percentage_combined[len(x_branch5) - 1:]
#pdb.set_trace()
return [percentage1, percentage5]
def Percentage_branch1(self, Np, Nt, starttime, location_x, Ttime_avg):
"""
At branch 1
calculate the particle percentage for each transect at the corresponding travel time
"""
#### read bathymetry information
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=1)
x_branch1 = WB.X
## from Ttime, find the segment index and travel time (time step) info for each
percentage = np.zeros_like(
x_branch1) ## seg ID from 1 to 46 for branch 1
for ii, tt in enumerate(Ttime_avg):
tt = int(tt)
if tt != 0:
tt += starttime
seg_id = ii + 2
#print ('Calculate particle percentage for time step = %s, segment = %d\n'%(str(tt), seg_id))
icount = 0
for i in range(location_x.shape[0]):
if location_x[i, tt - starttime] >= x_branch1[seg_id - 1]:
icount += 1
percentage[seg_id - 1] = float(icount) / Np
#pdb.set_trace()
return percentage
def dist2WTP_branch1(self, Ttime):
"""
calculate the distance to the WTP gate for branch 1
"""
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=1)
dist_tem = WB.X[1:-1][::-1] * 3.28084 ## unit: ft
ind = next((i for i, x in enumerate(Ttime) if x), None)
dist_tem[:ind] = 0
return dist_tem
def readBathymetry(self):
"""
read segment and layer info from the bathymetry file
"""
WB = W2_Bathymetry(self.bathfile)
WB.readVar()
WB.readLyr()
self.seg_length = WB.seg_length #segment length
self.seg_ori = WB.seg_ori # segment orientation
self.lyr_height = WB.lyr_height # layer height or each segment
self.lyrs_branchID = np.asarray(
[int(float(ID)) for ID in WB.lyrs_branchID])
self.lyrs_segID = np.asarray([int(float(ID)) for ID in WB.lyrs_segID])
self.lyrs = WB.lyrs
def AnimateContour(self,
varname='Tracer',
timestep=0,
branchID=1,
days=100,
PlotGrid=False):
"""
create animation
"""
import matplotlib.animation as animation
from mpl_toolkits.axes_grid1 import make_axes_locatable
Writer = animation.writers['ffmpeg']
writer = Writer(fps=5, metadata=dict(artist='Me'), bitrate=1800)
self.Readcpl()
plt.rcParams.update({'font.size': 18})
fig = plt.figure(figsize=(12.5, 8))
ax = fig.add_subplot(111)
if PlotGrid == True:
from bathymetry import W2_Bathymetry
filename = '%s\\%s' % (self.workdir, 'Bth_WB1.npt')
WB = W2_Bathymetry(filename)
pat = WB.VisBranch2(branchID)
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="3%", pad=0.05)
def animate(ii):
ii += timestep
ax.clear()
## search for index for each branch
## algorithm find the distance between two elements in self.X_flow that are large,
## this is where the two branches separate
dist = np.diff(self.X_flow[ii])
inds = np.where(dist > 1200)[0]
## this way is not necessary but clear
if branchID == 1:
ind0 = 0
ind1 = inds[0]
elif branchID == 2:
ind0 = inds[0] + 1
ind1 = inds[1]
elif branchID == 3:
ind0 = inds[1] + 1
ind1 = inds[2]
elif branchID == 4:
ind0 = inds[2] + 1
ind1 = inds[3]
elif branchID == 5:
ind0 = inds[3] + 1
ind1 = len(self.X_flow[ii])
X_flow = self.X_flow[ii][ind0:ind1 + 1]
Z_flow = self.Z_flow[ii][ind0:ind1 + 1]
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
if PlotGrid == True:
for sq in pat:
ax.add_patch(sq)
ax.autoscale_view()
## align each branch with grid
dx = WB.X.max() - X_flow.max()
X_flow += dx
U = self.U[ii][ind0:ind1 + 1]
W = self.W[ii][ind0:ind1 + 1]
U = np.asarray(U)
W = np.asarray(W)
mask = np.logical_or(U != self.mask_value, W != self.mask_value)
scale = 1.
scale = 100. / scale
Q = ax.quiver(X_flow[mask],
Z_flow[mask],
np.asarray(U[mask]) * 100.,
np.asarray(W[mask]) * 100.,
zorder=5,
width=0.001,
headwidth=4,
headlength=4.5,
scale=scale,
color='r')
qk = ax.quiverkey(Q,
0.15,
0.15,
1,
r'$1 \frac{cm}{s}$',
labelpos='W',
fontproperties={
'weight': 'bold',
'size': 20
})
var = self.var_output[varname]['value'][ii][ind0:ind1 + 1]
var = np.asarray(var)
#### quality control remove some very small values (spikes) ####
var[(var == self.mask_value) & (var < 1e-15)] = 0
#var = np.ma.masked_array(var,mask=var==self.mask_value)
triang = tri.Triangulation(X_flow, Z_flow)
isbad = np.equal(var, self.mask_value)
mask = np.any(np.where(isbad[triang.triangles], True, False),
axis=1)
triang.set_mask(mask)
levels = np.linspace(var.min(), var.max(), 100)
cmap = plt.set_cmap('bone_r')
#cs = ax.tricontourf(X_flow, Z_flow, var, cmap=cmap, levels=levels)
cs = ax.tricontourf(triang, var, cmap=cmap, levels=levels)
cb = fig.colorbar(cs, cax=cax, orientation='vertical')
cb.ax.tick_params(labelsize=12)
cb.ax.yaxis.offsetText.set_fontsize(12)
#cb.set_label('%s'%self.var_output[varname]['long_name'], fontsize=14)
cb.set_label('Concentration (mg/L)', fontsize=14)
timestr = datetime.strftime(self.runtimes[ii], '%Y-%m-%d')
ax.title.set_text('Time: %s' % timestr)
ax.set_ylim([135, 160])
ax.set_xlabel('Distance from upstream (m)')
ax.set_ylabel('Water Depth (m)')
return cs, cax
anim = animation.FuncAnimation(fig,
animate,
frames=days,
interval=600,
blit=False)
anim.save('%s\\tracer.mp4' % self.workdir, writer=writer)
def VisContour(self,
varname='Tracer',
timestep=-1,
branchID=1,
Plotuv=False,
PlotGrid=False):
"""
Create the contour plot of a variable given the variable name
variable names are provided in cpl.opt
e.g. 'U', 'W', 'Tracer'
The python dictionary for variables can be found in vardict.py
"""
self.Readcpl()
## search for index for each branch
## algorithm find the distance between two elements in self.X_flow that are large,
## this is where the two branches separate
#dist = np.diff(self.X_flow[0])
dist = np.diff(self.X_flow[timestep])
inds = np.where(dist > 1200)[0]
## this way is not necessary but clear
if branchID == 1:
ind0 = 0
ind1 = inds[0]
elif branchID == 2:
ind0 = inds[0] + 1
ind1 = inds[1]
elif branchID == 3:
ind0 = inds[1] + 1
ind1 = inds[2]
elif branchID == 4:
ind0 = inds[2] + 1
ind1 = inds[3]
elif branchID == 5:
ind0 = inds[3] + 1
#ind1 = len(self.X_flow[0])
ind1 = len(self.X_flow[timestep]) - 1
X_flow = self.X_flow[timestep][ind0:ind1 + 1]
Z_flow = self.Z_flow[timestep][ind0:ind1 + 1]
var = self.var_output[varname]['value'][timestep][ind0:ind1 + 1]
pdb.set_trace()
############### nested list ###############
def has_list(inlist):
return any(isinstance(el, list) for el in inlist)
import collections
def list_flatten(x):
if isinstance(x, collections.abc.Iterable):
return [a for i in x for a in list_flatten(i)]
else:
return [x]
if has_list(X_flow):
X_flow = list_flatten(X_flow)
if has_list(Z_flow):
Z_flow = list_flatten(Z_flow)
if has_list(var):
var = list_flatten(var)
##############################################
pdb.set_trace()
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
var = np.asarray(var)
#### quality control remove some very small values (spikes) ####
#var[(var.mask==False)&(var<1e-15)]=0
var[(var == self.mask_value) & (var < 1e-15)] = 0
plt.rcParams.update({'font.size': 18})
fig = plt.figure(figsize=(11.5, 8))
ax = fig.add_subplot(111)
if PlotGrid == True:
from bathymetry import W2_Bathymetry
filename = '%s\\%s' % (self.workdir, 'Bth_WB1.npt')
WB = W2_Bathymetry(filename)
pat = WB.VisBranch2(branchID)
for sq in pat:
ax.add_patch(sq)
ax.autoscale_view()
## align each branch with grid
## not sure how X is defined in the model
## algorithm starting from the end of the branch
dx = WB.X.max() - X_flow.max()
X_flow += dx
if Plotuv:
#X_flow = np.asarray(X_flow)
#Z_flow = np.asarray(Z_flow)
U = self.U[timestep][ind0:ind1 + 1]
W = self.W[timestep][ind0:ind1 + 1]
U = np.asarray(U)
W = np.asarray(W)
mask = np.logical_or(U != self.mask_value, W != self.mask_value)
#U = np.ma.masked_array(U,mask=U==self.mask_value)
#W = np.ma.masked_array(W,mask=W==self.mask_value)
scale = 1.
scale = 100. / scale
Q = ax.quiver(X_flow[mask],
Z_flow[mask],
np.asarray(U[mask]) * 100.,
np.asarray(W[mask]) * 100.,
zorder=5,
width=0.001,
headwidth=4,
headlength=4.5,
scale=scale,
color='r')
qk = ax.quiverkey(Q,
0.15,
0.15,
1,
r'$1 \frac{cm}{s}$',
labelpos='W',
fontproperties={
'weight': 'bold',
'size': 20
})
#var = np.ma.masked_array(var,mask=var==self.mask_value)
#### python3 does not allow for masked z values in tricontourf, a way to work around:
triang = tri.Triangulation(X_flow, Z_flow)
isbad = np.equal(var, self.mask_value)
mask = np.any(np.where(isbad[triang.triangles], True, False), axis=1)
triang.set_mask(mask)
#pdb.set_trace()
if var.min() == 0 and var.max() == 0:
levels = np.linspace(self.var_output[varname]['limits'][0],
self.var_output[varname]['limits'][1], 100)
elif varname == 'T':
#levels = np.linspace(23, 27, 100) ## low WSE
#levels = np.linspace(16, 21, 100) ## high WSE
levels = np.linspace(12, 17, 100) ## medium WSE
else:
levels = np.linspace(var.min(), var.max(), 100)
cmap = plt.set_cmap('bone_r')
#cs = ax.tricontourf(X_flow, Z_flow, var, cmap=cmap, levels=levels)
cs = ax.tricontourf(triang, var, cmap=cmap, levels=levels)
from mpl_toolkits.axes_grid1 import make_axes_locatable
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="3%", pad=0.05)
cb = fig.colorbar(cs, cax=cax, orientation='vertical')
cb.ax.tick_params(labelsize=12)
cb.ax.yaxis.offsetText.set_fontsize(12)
#cb.set_label('%s'%self.var_output[varname]['long_name'], fontsize=14)
cb.set_label('Concentration (mg/L)', fontsize=14)
timestr = datetime.strftime(self.runtimes[timestep], '%Y-%m-%d')
ax.title.set_text('Time: %s' % timestr)
#ax.set_xlim([0, 60000])
ax.set_ylim([135, 160])
ax.set_xlabel('Distance from upstream (m)')
ax.set_ylabel('Water Depth (m)')
#ax.yaxis.grid(True)
#ax.xaxis.grid(True)
plt.show()
def particle_animation(self,
Nt,
particle_location,
branchID=1,
verbose='surface'):
"""
create animation for 1D particle tracking
verbose: 'surface' or 'bottom'
"""
import matplotlib.animation as animation
Writer = animation.writers['ffmpeg']
writer = Writer(fps=5, metadata=dict(artist='Me'), bitrate=1800)
xx = np.arange(particle_location.shape[0]) + 1
if branchID == 1:
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
x_branch = WB.X
elif branchID == 5:
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
x_branch5 = WB.X #### segment x coordinates for branch 5
#### read segment information for branch 1
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=1)
x_branch1 = WB.X
#### combine the two branch cells
x_branch = x_branch5.tolist()[0:] + \
(x_branch1[self.DHS5-1:] - x_branch1[self.DHS5-1] + x_branch5[-2]).tolist()
x_branch = np.asarray(x_branch)
plt.rcParams.update({'font.size': 18})
fig = plt.figure(figsize=(8, 12.5))
ax = fig.add_subplot(111)
def animate(ii):
ax.clear()
### grid segments
for yc in x_branch:
ax.axhline(y=yc, color='gray', linestyle='-', linewidth=1)
#### particle positions
#for i in range(particle_location.shape[0]):
cs = ax.plot(xx, particle_location[:, ii], 'ok',
markersize=3.5) ## at 3rd time step
ax.title.set_text('%s \n Time step = %d' % (verbose, ii))
ax.set_ylim([-1500, 28500])
ax.set_ylim(ax.get_ylim()[::-1])
ax.set_xlabel('Particle ID')
ax.set_ylabel('Distance from upstream (m)')
return cs
anim = animation.FuncAnimation(fig,
animate,
frames=Nt,
interval=600,
blit=False)
anim.save(r'videos\particle\%s.mp4' % verbose, writer=writer)
def read_velocity(self, Nt, branchID=1):
"""
read surface and bottom velocities, specific for each branch
"""
#### read bathymetry information
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
X_surface = [] ## [time period, x grid points]
Z_surface = []
U_surface = []
X_bottom = []
Z_bottom = []
U_bottom = []
for tstep in range(self.starttime, self.starttime + Nt):
print('Time step = %s \n' % str(tstep))
## search for index for each branch
## algorithm find the distance between two elements in self.X_flow that are large,
## this is where the two branches separate
dist = np.diff(self.X_flow[tstep])
inds = np.where(dist > 1200)[0]
if branchID == 1:
ind0 = 0
ind1 = inds[0]
elif branchID == 2:
ind0 = inds[0] + 1
ind1 = inds[1]
elif branchID == 3:
ind0 = inds[1] + 1
ind1 = inds[2]
elif branchID == 4:
ind0 = inds[2] + 1
ind1 = inds[3]
elif branchID == 5:
ind0 = inds[3] + 1
ind1 = len(self.X_flow[tstep]
) - 1 ## -1 remove the array size mismatch issue
## tracer locations
X_flow = self.X_flow[tstep][ind0:ind1 + 1]
Z_flow = self.Z_flow[tstep][ind0:ind1 + 1]
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
## align coordinates with the grid
dx = WB.X.max() - X_flow.max()
X_flow += dx
## read velocity
U = self.U[tstep][ind0:ind1 + 1]
W = self.W[tstep][ind0:ind1 + 1]
U = np.asarray(U)
W = np.asarray(W)
#### find surface and bottom velocity at tstep
X_surface_tem, Z_surface_tem, U_surface_tem, \
X_bottom_tem, Z_bottom_tem, U_bottom_tem = \
self.surface_bottom_velocity(X_flow, Z_flow, U, W)
X_surface.append(X_surface_tem)
Z_surface.append(Z_surface_tem)
U_surface.append(U_surface_tem)
X_bottom.append(X_bottom_tem)
Z_bottom.append(Z_bottom_tem)
U_bottom.append(U_bottom_tem)
#pdb.set_trace()
return X_surface, Z_surface, U_surface, X_bottom, Z_bottom, U_bottom
def particle_tracking_model_1D(self,
Np,
Nt,
InitialSeg,
starttime,
branchID,
flow_condition='high',
dt=1,
transportSurface=True,
transportBottom=True,
travelTime=True):
"""
particle tracking with velocity
Np -- number of particles
Nt -- particle tracking period (unit: day)
InitialSeg -- initial spill release segment ID
dt -- time interval, default 1 day, unit: day
"""
dt *= 24 * 3600. #### conversion from day to seconds
self.starttime = starttime
self.flow_condition = flow_condition
#### read surface and bottom velocities
if branchID == 1:
self.X_surface, self.Z_surface, self.U_surface, \
self.X_bottom, self.Z_bottom, self.U_bottom = self.read_velocity(Nt, branchID=1)
## contour plot of velocity
self.plot_velocity(
self.X_surface,
self.U_surface,
figname=r'figures\flow_rate\velocity\surface_branch%d_%s.png' %
(branchID, flow_condition)) ## surface
self.plot_velocity(
self.X_bottom,
self.U_bottom,
figname=r'figures\flow_rate\velocity\bottom_branch%d_%s.png' %
(branchID, flow_condition)) ## surface
elif branchID == 5:
X_surface1, Z_surface1, U_surface1, \
X_bottom1, Z_bottom1, U_bottom1 = self.read_velocity(Nt, branchID=1)
X_surface5, Z_surface5, U_surface5, \
X_bottom5, Z_bottom5, U_bottom5 = self.read_velocity(Nt, branchID=5)
## contour plot of velocity
self.plot_velocity(
X_surface5,
U_surface5,
figname=r'figures\flow_rate\velocity\surface_branch%d_%s.png' %
(branchID, flow_condition)) ## surface
self.plot_velocity(
X_bottom5,
U_bottom5,
figname=r'figures\flow_rate\velocity\bottom_branch%d_%s.png' %
(branchID, flow_condition)) ## surface
#### read bathymetry information
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=1)
#### adding branch 5 to main branch
self.X_surface = []
self.Z_surface = []
self.U_surface = []
self.X_bottom = []
self.Z_bottom = []
self.U_bottom = []
for t in range(Nt):
## surface
xind_surface = self.findNearest(WB.X[self.DHS5 - 1],
X_surface1[t][:])
xtem_surface_branch1 = np.asarray(X_surface1[t][xind_surface:]) - X_surface1[t][xind_surface-1] \
+ X_surface5[t][-1]
self.X_surface.append(X_surface5[t] +
xtem_surface_branch1.tolist())
self.Z_surface.append(Z_surface5[t] +
Z_surface1[t][xind_surface:])
self.U_surface.append(U_surface5[t] +
U_surface1[t][xind_surface:])
## bottom
xind_bottom = self.findNearest(WB.X[self.DHS5 - 1],
X_bottom1[t][:])
xtem_bottom_branch1 = np.asarray(X_bottom1[t][xind_bottom:]) - X_bottom1[t][xind_bottom-1] \
+ X_bottom5[t][-1]
self.X_bottom.append(X_bottom5[t] +
xtem_bottom_branch1.tolist())
self.Z_bottom.append(Z_bottom5[t] + Z_bottom1[t][xind_bottom:])
self.U_bottom.append(U_bottom5[t] + U_bottom1[t][xind_bottom:])
#### read bathymetry information
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
#### particle tracking calculation
if transportSurface:
#### particle location array
self.location_x_surface = np.zeros(
[Np, Nt]) ####[Number of particles, time period]
self.grid_x_surface = np.zeros(
[Nt]) #### surface water level at each x grid
#### initial particle location
self.location_x_surface[:, 0] = WB.X[InitialSeg - 1]
#### first order Euler algorithm: x(t+1) = x(t) + U*dt + R*sqrt(6 * Dx *dt)
for i in range(Np):
for t in range(Nt - 1):
xtem = np.abs(self.X_surface[t] -
self.location_x_surface[i, t])
#### check if
if xtem.min() < 1000:
#### query index
ind = np.argwhere(xtem == xtem.min())[0][0]
utem = self.U_surface[t][ind]
R = random.uniform(
0, 2) - 1 ## random number between [-1,1]
self.location_x_surface[
i, t + 1] = self.location_x_surface[
i,
t] + utem * dt + R * np.sqrt(6 * self.Dx * dt)
elif xtem.min(
) > 1000: ## there is no close grid point, water dries at this location
utem = 0
self.location_x_surface[
i,
t + 1] = self.location_x_surface[i, t] + utem * dt
#if t in range(236, 238):
## at these steps, water at the first several cells dries, X_surface starts at 9659, while location_x_surface is 8440.
## so particles do not move at these time steps
#pdb.set_trace()
for t in range(Nt):
self.grid_x_surface[t] = self.Z_surface[t][0]
if transportBottom:
#### particle location array
self.location_x_bottom = np.zeros([Np, Nt])
self.grid_x_bottom = np.zeros(
[Nt]) #### bottom water level at each x grid
#### initial particle location
self.location_x_bottom[:, 0] = WB.X[InitialSeg - 1]
#### first order Euler algorithm
for i in range(Np):
for t in range(Nt - 1):
xtem = np.abs(self.X_bottom[t] -
self.location_x_bottom[i, t])
#### check if
if xtem.min() < 1000:
#### query index
ind = np.argwhere(xtem == xtem.min())[0][0]
utem = self.U_bottom[t][ind]
R = random.uniform(
0, 2) - 1 ## random number between [-1,1]
self.location_x_bottom[
i, t + 1] = self.location_x_bottom[
i,
t] + utem * dt + R * np.sqrt(6 * self.Dx * dt)
elif xtem.min(
) > 1000: ## there is no close grid point, water dries at this location
utem = 0
self.location_x_bottom[
i,
t + 1] = self.location_x_bottom[i, t] + utem * dt
for t in range(Nt):
self.grid_x_bottom[t] = self.Z_bottom[t][0]
## first entry: Nt or self.period or self-defined depending on how long we need the video to be
self.particle_animation(self.period,
self.location_x_surface,
branchID=branchID,
verbose='surface_branch%d_%s' %
(branchID, flow_condition))
self.particle_animation(self.period,
self.location_x_bottom,
branchID=branchID,
verbose='bottom_branch%d_%s' %
(branchID, flow_condition))
# #### For testing only: visualize particle locations
# iy = 0
# plt.rcParams.update({'font.size': 16})
# fig = plt.figure(figsize=(14,10))
# ax = fig.add_subplot(211)
# for i in range(Np):
# ax.plot(self.location_x_surface[i], self.grid_x_surface+iy, 'o')
# iy+=5
#
# ax2 = fig.add_subplot(212)
# for i in range(Np):
# ax2.plot(self.location_x_bottom[i], self.grid_x_bottom-iy, 'o')
# iy-=5
# plt.show()
if travelTime and transportSurface:
self.travel_time(
Np,
Nt,
InitialSeg,
starttime,
branchID,
self.location_x_surface,
write2shp=False,
density=0,
excelfile=r'excel\particle_surface_branch%s_%s.xlsx' %
(str(branchID), flow_condition))
if travelTime and transportBottom:
self.travel_time(
Np,
Nt,
InitialSeg,
starttime,
branchID,
self.location_x_bottom,
write2shp=False,
density=1,
excelfile=r'excel\particle_bottom_branch%s_%s.xlsx' %
(str(branchID), flow_condition))
def Concentrate(self,
starttime,
endtime,
Ttime,
branchID,
water_level=True):
"""
calculate the concentrate at each segment
if water_level is True, save the water level data too
"""
#### read bathymetry information
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
## from Ttime, find the segment index and travel time (time step) info for each
concentrate = np.zeros_like(WB.X) ## seg ID from 1 to 46 for branch 1
elevation = np.zeros_like(WB.X)
for ii, tt in enumerate(Ttime):
tt = int(tt)
if tt != 0:
tt += starttime
seg_id = ii + 2
print(
'Calculate concentration for time step = %s, segment = %d\n'
% (str(tt), seg_id))
## read grid info
dist = np.diff(self.X_flow[tt])
inds = np.where(dist > 1200)[0]
if branchID == 1:
ind0 = 0
ind1 = inds[0]
elif branchID == 5:
ind0 = inds[3] + 1
ind1 = len(
self.X_flow[tt]
) - 1 ## -1 remove the array size mismatch issue
X_flow = self.X_flow[tt][ind0:ind1 + 1]
Z_flow = self.Z_flow[tt][ind0:ind1 + 1]
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
## align coordinates with the grid
dx = WB.X.max() - X_flow.max()
X_flow += dx
## no concentrate data for particle tracking, concentration is zero instead
# ## read tracer data
# vartem = np.asarray( self.var_output['Tracer']['value'][tt][ind0:ind1+1] )
#
# #### quality control if X_flow, vartem not in the same shape, resize
# if X_flow.shape != vartem.shape:
# #pdb.set_trace()
# Lmin = np.min([X_flow.shape[0], vartem.shape[0]])
# if X_flow.shape[0] > vartem.shape[0]:
# X_flow = np.delete(X_flow, np.arange(Lmin, X_flow.shape[0]))
# elif X_flow.shape[0] < vartem.shape[0]:
# vartem = np.delete(vartem, np.arange(Lmin, vartem.shape[0]))
#
# ## segment location : WB.X[seg_id-1]
#
# ## find index
# ## There are two options
# ## Option 1: the concentrate at the exact segment
# inds = self.find_seg_index_exact(WB.X[seg_id-1], X_flow, vartem)
# ## Option 2: the concentrate beyond the segment
# #inds_beyond = self.find_seg_index_beyond(WB.X[seg_id-1], X_flow, vartem)
#
# concentrate[seg_id-1] = vartem[inds[0]]
if water_level:
eta = np.asarray(
self.var_output['Elevation']['value'][tt][ind0:ind1 +
1])
if X_flow.shape != eta.shape:
Lmin = np.min([X_flow.shape[0], eta.shape[0]])
if X_flow.shape[0] > eta.shape[0]:
X_flow = np.delete(
X_flow, np.arange(Lmin, X_flow.shape[0]))
elif X_flow.shape[0] < eta.shape[0]:
eta = np.delete(eta, np.arange(Lmin, eta.shape[0]))
inds_eta = self.find_seg_index_exact(
WB.X[seg_id - 1], X_flow, eta)
elevation[seg_id - 1] = eta[inds_eta[0]]
if water_level:
return concentrate[1:-1], elevation[1:-1] / 0.3048
def travel_time(self, Np, Nt, InitialSeg, starttime, branchID, location_x,
write2shp, density, excelfile):
"""
calculate travel time based on the particle tracking
"""
if branchID == 1:
#### read segment information
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
#### create empty array for travel time
Ttime = np.zeros([Np, WB.X.shape[0]])
#### calculate travel time
for i in range(Np):
for tstep in range(Nt):
location_x_tem = location_x[i, tstep]
ind = self.Xsearch(location_x_tem, WB.X)
if tstep == 0:
Ttime[i, InitialSeg - 1:ind + 1] = tstep + 1
ind_tem = ind
else:
ind_nonzero = np.nonzero(Ttime[i, :])[0].max(
) ## only add travel time to zero elements
if ind > max(ind_tem, ind_nonzero):
Ttime[i,
max(ind_tem + 1, ind_nonzero + 1):ind +
1] = tstep + 1
print(Ttime[i, :])
ind_tem = ind
#### calculate the average among particles
#### be careful about this average among particles
#Ttime = np.mean(Ttime, axis=0, dtype=np.int)
#### simply average may yield smaller travel time at the downstream end
#### because some particle may not travel to the last segment, which has zero travel time there
#### so only average among non-zero values.
Ttime_avg = np.zeros([WB.X.shape[0]])
for i in range(WB.X.shape[0]):
if i >= InitialSeg - 1:
Ttime_avg[i] = np.median(Ttime[:, i][np.nonzero(Ttime[:,
i])])
#pdb.set_trace()
Ttime_avg = Ttime_avg[1:-1]
## calculate concentrate/water_level based on the travel time
concentrate, water_level = self.Concentrate_branch1(
Nt, starttime, Ttime_avg)
## calculate distance to WTP gate
dist = self.dist2WTP_branch1(Ttime_avg)
## calculate particle percentage
percentage = self.Percentage_branch1(Np, Nt, starttime, location_x,
Ttime_avg)
Ttime_avg[
Ttime_avg != 0] = Ttime_avg[-1] - Ttime_avg[Ttime_avg != 0]
#### call segment class for segment information
WS = W2_Segmentation(self.Bthfile)
WS.VisSeg2()
#### save travel time data to txt file ####
#self.savetxt_Traveltime_branch1(WS, Ttime_avg, density, self.flows[self.flow_condition], txtfile)
save_excel_Traveltime_branch1(WS, Ttime_avg, density, self.solubility, self.flows[self.flow_condition], \
concentrate, water_level, dist, excelfile)
if write2shp:
from myshapefile import writeShpLines_one_branch
writeShpLines_one_branch(
WS,
Ttime_avg,
shpname='particle_surface_traveltime_branch1')
if branchID == 5:
"""
Under development
"""
#### read segment information for branch 5
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
x_branch5 = WB.X #### segment x coordinates for branch 5
#### read segment information for branch 1
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID=1)
x_branch1 = WB.X
#### create empty array for travel time
#### should not include the inactive cell at the end of branch5 when combining
# x_combined = x_branch5.tolist()[0:-1] + \
# (x_branch1[self.DHS5-1:] - x_branch1[self.DHS5-1] + x_branch5[-2]).tolist()
x_combined = x_branch5.tolist()[0:-1] + \
(x_branch1[self.DHS5-1:] - x_branch1[self.DHS5-1] + x_branch5[-2]).tolist()[1:]
x_combined = np.asarray(x_combined)
Ttime = np.zeros([Np, x_combined.shape[0]])
#### calculate travel time
for i in range(Np):
for tstep in range(Nt):
location_x_tem = location_x[i, tstep]
ind = self.Xsearch(location_x_tem, x_combined)
if tstep == 0:
Ttime[i, InitialSeg - 1:ind + 1] = tstep + 1
ind_tem = ind
else:
ind_nonzero = np.nonzero(Ttime[i, :])[0].max(
) ## only add travel time to zero elements
if ind > max(ind_tem, ind_nonzero):
Ttime[i,
max(ind_tem + 1, ind_nonzero + 1):ind +
1] = tstep + 1
print(Ttime[i, :])
ind_tem = ind
#### separate the travel time to branch 1 segments and branch 5 segments
Ttime5 = Ttime[:, 0:len(x_branch5[0:-1]) + 1]
Ttime1 = np.zeros([Np, len(x_branch1)])
#Ttime1[:, self.DHS5:] = Ttime[:, len(x_branch5[0:-1])+1:]
Ttime1[:, self.DHS5:-1] = Ttime[:, len(x_branch5[0:-1]) + 1:]
#pdb.set_trace()
#### calculate the average among particles
Ttime_avg1 = np.zeros([Ttime1.shape[1]])
for i in range(Ttime1.shape[1]):
if i >= self.DHS5 and len(Ttime1[:, i].nonzero()[0]) != 0:
Ttime_avg1[i] = np.median(Ttime1[:,
i][np.nonzero(Ttime1[:,
i])])
Ttime_avg5 = np.zeros([Ttime5.shape[1]])
for i in range(Ttime5.shape[1]):
if i >= InitialSeg - 1:
Ttime_avg5[i] = np.median(Ttime5[:,
i][np.nonzero(Ttime5[:,
i])])
Ttime_avg1 = Ttime_avg1[1:-1]
Ttime_avg5 = Ttime_avg5[1:-1]
Ttimes_avg = [Ttime_avg1, Ttime_avg5]
## calculate concentrate/water level based on the travel time
concentrates, water_levels = self.Concentrate_branch5(
Nt, starttime, Ttimes_avg)
## calculate distance to WTP gate
dists = self.dist2WTP_branch5(Ttimes_avg)
## calculate particle percentage
percentages = self.Percentage_branch5(Np, Nt, starttime,
location_x, x_combined,
Ttimes_avg)
MaxTime = Ttimes_avg[0][-1]
for Ttime in Ttimes_avg:
Ttime[Ttime != 0] = MaxTime - Ttime[Ttime != 0]
#### call segment class for segment information
WS = W2_Segmentation(self.Bthfile)
WS.VisSeg2()
#### save travel time data to txt file ####
#self.savetxt_Traveltime_branch5(WS, Ttimes_avg, density, self.flows[self.flow_condition], txtfile)
save_excel_Traveltime_branch5(WS, Ttimes_avg, density, self.solubility, self.flows[self.flow_condition], \
concentrates, water_levels, dists, excelfile)
if write2shp:
from myshapefile import writeShpLines
writeShpLines(WS,
Ttimes_avg,
shpname='particle_bottom_traveltime_branch5')
def travel_time(self, starttime=0, endtime=1460, branchID=1):
"""
calculate the travel time (or median arrival time) for conservative tracers
timestep - when to calculate the travel time
branchID - the travel time of which branch
"""
print ("Calculate travel time for branch %s ... \n"%str(branchID))
#### read bathymetry information
WB = W2_Bathymetry(self.Bthfile)
pat = WB.VisBranch2(branchID)
#### create empty array for travel time
Ttime = np.zeros_like(WB.X)
#### calculate travel time ####
#### loop over all time steps ####
#### starttime to endtime, during which period to search for tracer travel time
for tstep in range(starttime, endtime):
print ('Time step = %s \n'%str(tstep))
## search for index for each branch
## algorithm find the distance between two elements in self.X_flow that are large,
## this is where the two branches separate
dist = np.diff(self.X_flow[tstep])
inds = np.where(dist>1200)[0]
if branchID == 1:
ind0 = 0
ind1 = inds[0]
elif branchID == 2:
ind0 = inds[0]+1
ind1 = inds[1]
elif branchID == 3:
ind0 = inds[1]+1
ind1 = inds[2]
elif branchID == 4:
ind0 = inds[2]+1
ind1 = inds[3]
elif branchID == 5:
ind0 = inds[3]+1
ind1 = len(self.X_flow[tstep]) - 1 ## -1 remove the array size mismatch issue
## tracer locations
X_flow = self.X_flow[tstep][ind0:ind1+1]
Z_flow = self.Z_flow[tstep][ind0:ind1+1]
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
## align coordinates with the grid
dx = WB.X.max() - X_flow.max()
X_flow += dx
vartem = np.asarray( self.var_output['Tracer']['value'][tstep][ind0:ind1+1] )
#### quality control if X_flow, vartem not in the same shape, resize
if X_flow.shape != vartem.shape:
#pdb.set_trace()
Lmin = np.min([X_flow.shape[0], vartem.shape[0]])
if X_flow.shape[0] > vartem.shape[0]:
X_flow = np.delete(X_flow, np.arange(Lmin, X_flow.shape[0]))
elif X_flow.shape[0] < vartem.shape[0]:
vartem = np.delete(vartem, np.arange(Lmin, vartem.shape[0]))
## search nonzero X coordinates
X_nonzero = self.Xsearch(X_flow, Z_flow, vartem)
## search nonzero X index
ind_nonzero = self.GridSearch(WB.X, X_nonzero, branchID) ## output is segID
#### if tracer is released at day 385.9 at segment 18, the tracer would travel to segment 22 at day 386
#### so it is important to record the initial segment in the travel time array
if tstep == starttime and starttime >= 1:
#Ttime[np.min(ind_nonzero)-1] = tstep - 1
if self.initialBranch == 1 and branchID == 1:
Ttime[self.initialSeg-1] = tstep - 1
elif self.initialBranch == 5 and branchID == 5:
Ttime[self.initialSeg-85] = tstep - 1
if ind_nonzero.size != 0: #### there is nonzero tracer in the branch
if branchID in [1,5]:
ind_max = np.max(ind_nonzero)
## add the travel time to the pre-defined 1D array
if Ttime[ind_max] == 0:
if len(Ttime.nonzero()[0]) != 0:
ind_tem = np.where(Ttime==Ttime.max())[0][-1]
Ttime[ind_tem+1: ind_max] = tstep
else:
Ttime[ind_max] = tstep
elif branchID in [2,3,4]:
if len(Ttime.nonzero()[0])==0 and len(ind_nonzero) == 1:
Ttime[ind_nonzero[0]] = tstep
else:
ind_min = np.min(ind_nonzero)
## add the travel time to the pre-defined 1D array
if Ttime[ind_min] == 0:
if len(Ttime.nonzero()[0]) != 0:
ind_tem = np.where(Ttime==Ttime.max())[0][0]
Ttime[ind_min:ind_tem] = tstep
else:
Ttime[ind_min] = tstep
print (Ttime)
#pdb.set_trace()
if branchID in [1, 5]:
Ttime[-2] = Ttime[-3] + 1
return Ttime
def travel_time(self, endtime=1460, branchID=1):
"""
calculate the travel time (or earliest arrival time) for conservative tracers
timestep - when to calculate the travel time
branchID - the travel time of which branch
"""
print "Calculate travel time for branch %s ... \n" % str(branchID)
#### read bathymetry information
Bthfile = '%s\\%s' % (self.workdir, 'Bth_WB1.npt')
WB = W2_Bathymetry(Bthfile)
pat = WB.VisBranch2(branchID)
#### create empty array for travel time
Ttime = np.zeros_like(WB.X)
#### calculate travel time ####
#### loop over all time steps ####
for tstep in range(endtime):
print 'Time step = %s \n' % str(tstep)
## search for index for each branch
## algorithm find the distance between two elements in self.X_flow that are large,
## this is where the two branches separate
dist = np.diff(self.X_flow[tstep])
inds = np.where(dist > 1200)[0]
if branchID == 1:
ind0 = 0
ind1 = inds[0]
elif branchID == 2:
ind0 = inds[0] + 1
ind1 = inds[1]
elif branchID == 3:
ind0 = inds[1] + 1
ind1 = inds[2]
elif branchID == 4:
ind0 = inds[2] + 1
ind1 = inds[3]
elif branchID == 5:
ind0 = inds[3] + 1
ind1 = len(self.X_flow[tstep]
) - 1 ## -1 remove the array size mismatch issue
## tracer locations
X_flow = self.X_flow[tstep][ind0:ind1 + 1]
Z_flow = self.Z_flow[tstep][ind0:ind1 + 1]
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
## align coordinates with the grid
dx = WB.X.max() - X_flow.max()
X_flow += dx
vartem = np.asarray(
self.var_output['Tracer']['value'][tstep][ind0:ind1 + 1])
#### quality control if X_flow, vartem not in the same shape, resize
if X_flow.shape != vartem.shape:
#pdb.set_trace()
Lmin = np.min([X_flow.shape[0], vartem.shape[0]])
if X_flow.shape[0] > vartem.shape[0]:
X_flow = np.delete(X_flow,
np.arange(Lmin, X_flow.shape[0]))
elif X_flow.shape[0] < vartem.shape[0]:
vartem = np.delete(vartem,
np.arange(Lmin, vartem.shape[0]))
## search nonzero X coordinates
X_nonzero = self.Xsearch(X_flow, Z_flow, vartem)
## search nonzero X index
ind_nonzero = self.GridSearch(WB.X, X_nonzero, branchID)
if ind_nonzero.size != 0: #### there is nonzero tracer in the branch
if branchID == 1:
ind_max = np.max(ind_nonzero)
## add the travel time to the pre-defined 1D array
if Ttime[ind_max] == 0:
if len(Ttime.nonzero()[0]) != 0:
ind_tem = np.where(Ttime == Ttime.max())[0][-1]
Ttime[ind_tem + 1:ind_max] = tstep
else:
Ttime[ind_max] = tstep
elif branchID in [2, 3, 4, 5]:
if len(Ttime.nonzero()[0]) == 0 and len(ind_nonzero) == 1:
Ttime[ind_nonzero[0]] = tstep
else:
ind_min = np.min(ind_nonzero)
## add the travel time to the pre-defined 1D array
if Ttime[ind_min] == 0:
if len(Ttime.nonzero()[0]) != 0:
ind_tem = np.where(Ttime == Ttime.max())[0][0]
Ttime[ind_min:ind_tem] = tstep
else:
Ttime[ind_min] = tstep
print Ttime
#pdb.set_trace()
if branchID == 1:
Ttime[-2] = Ttime[-3] + 1
return Ttime
def VisParticles(self,
timestep=-1,
branchID=1,
PlotFlow=True,
PlotTemp=True,
PlotGrid=True):
"""
visualize the particle trajectories at a given time step
set PlotTemp == False for now, need to figure out how to contour plot 1D array
"""
self.ReadParticles(branchID)
X = self.X[timestep]
Z = self.Z[timestep]
if PlotFlow == True or PlotTemp == True:
self.ReadBackground(branchID)
X_flow = self.X_flow[timestep]
Z_flow = self.Z_flow[timestep]
plt.rcParams.update({'font.size': 18})
fig = plt.figure(figsize=(11.5, 8))
ax = fig.add_subplot(111)
ax.plot(X, Z, '.', color='r')
if PlotFlow == True:
U = self.U[timestep]
W = self.W[timestep]
## mask data
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
U = np.asarray(U)
W = np.asarray(W)
maskuv = np.logical_or(U != 0, W != 0)
scale = 1.
scale = 100. / scale
Q = ax.quiver(X_flow[maskuv],
Z_flow[maskuv],
U[maskuv] * 100.,
W[maskuv] * 100.,
zorder=5,
width=0.001,
headwidth=4,
headlength=4.5,
scale=scale,
color='b')
qk = ax.quiverkey(Q,
0.15,
0.15,
1,
r'$1 \frac{cm}{s}$',
labelpos='W',
fontproperties={
'weight': 'bold',
'size': 20
})
if PlotTemp == True:
import matplotlib.tri as tri
#from scipy.interpolate import griddata
T = self.T[timestep]
U = np.asarray(self.U[timestep])
W = np.asarray(self.W[timestep])
## mask data
X_flow = np.asarray(X_flow)
Z_flow = np.asarray(Z_flow)
T = np.asarray(T)
triang = tri.Triangulation(X_flow, Z_flow)
isbad = np.less_equal(np.asarray(self.T[0]), 0)
#isbad = np.equal(U, 0) & np.equal(W, 0)
mask = np.any(np.where(isbad[triang.triangles], True, False),
axis=1)
triang.set_mask(mask)
#T[T==0] = self.mask_value
#T = np.ma.masked_array(T,mask=T==self.mask_value)
#T_limits = [0,35]
T_limits = [T.min(), T.max()]
levels = np.linspace(T_limits[0], T_limits[1], 100)
cs = ax.tricontourf(triang, T, cmap=plt.cm.bone, levels=levels)
#cs = ax.tricontourf(X_flow, Z_flow, T, cmap=plt.cm.bone, levels=levels)
from mpl_toolkits.axes_grid1 import make_axes_locatable
divider = make_axes_locatable(ax)
cax = divider.append_axes("right", size="3%", pad=0.05)
cb = fig.colorbar(cs, cax=cax, orientation='vertical')
cb.ax.tick_params(labelsize=12)
cb.ax.yaxis.offsetText.set_fontsize(12)
cb.set_label('Temperature', fontsize=14)
if PlotGrid == True:
from bathymetry import W2_Bathymetry
filename = '%s\\%s' % (self.workdir, 'Bth_WB1.npt')
WB = W2_Bathymetry(filename)
pat = WB.VisBranch2(branchID)
for sq in pat:
ax.add_patch(sq)
ax.autoscale_view()
timestr = datetime.strftime(self.runtimes[timestep], '%Y-%m-%d')
ax.title.set_text('Time: %s' % timestr)
#ax.set_xlim([4000, 12000])
#ax.set_ylim([135, 150])
ax.set_ylim([135, 160])
ax.set_xlabel('Distance from upstream (m)')
ax.set_ylabel('Water Depth (m)')
#ax.yaxis.grid(True)
#ax.xaxis.grid(True)
#plt.show()
#plt.savefig('particle_tracks_%s.png'%str(timestep))
#plt.savefig('%s\\particle_tracks_%s.png'%(self.workdir, str(timestep)))
#plt.close()
plt.show()