In [131]:
import python_subgrid.wrapper
import numpy as np
import python_subgrid.plotting
import matplotlib.pyplot as plt
%matplotlib inline
In [2]:
# read the san francisco model
subgrid = python_subgrid.wrapper.SubgridWrapper('/Users/baart_f/models/sfo/sfo-3di/sanfranciscobay2.mdu')
subgrid.start()
In [3]:
# do 2 timesteps
subgrid.update(-1)
subgrid.update(-1)
# read the grid info
pixels = {}
pixels['dps'] = subgrid.get_nd('dps', sliced=True)
pixels['dsnop'] = subgrid.get_nd('dsnop')
pixels['dps'] = np.ma.masked_array(-pixels['dps'], mask=pixels['dps']==pixels['dsnop'])
pixels['x0p'] = subgrid.get_nd('x0p')
pixels['dxp'] = subgrid.get_nd('dxp')
pixels['y0p'] = subgrid.get_nd('y0p')
pixels['dyp'] = subgrid.get_nd('dyp')
pixels['imax'] = subgrid.get_nd('imax')
pixels['jmax'] = subgrid.get_nd('jmax')
pixels['x'] = np.arange(pixels['x0p'], pixels['x0p'] + pixels['imax']*pixels['dxp'], pixels['dxp'])
pixels['y'] = np.arange(pixels['y0p'], pixels['y0p'] + pixels['jmax']*pixels['dyp'], pixels['dyp'])
pixels['Y'], pixels['X'] = np.meshgrid(pixels['y'], pixels['x'])
# create a quad grid for indexing
quad_grid = python_subgrid.plotting.make_quad_grid(subgrid)
# lookup variables
u1 = subgrid.get_nd('u1', sliced=False)[:]
line = subgrid.get_nd('FlowLink')
xu = subgrid.get_nd('FlowLink_xu')
yu = subgrid.get_nd('FlowLink_yu')
xcc = subgrid.get_nd('FlowElem_xcc')
ycc = subgrid.get_nd('FlowElem_ycc')
uc = subgrid.get_nd('uc')
ucx, ucy = uc
dps = subgrid.get_nd('dps')
cf = subgrid.get_nd('cf')
s1 = subgrid.get_nd('s1')
# cell variables
nod_type = subgrid.get_nd('nod_type')
uc = uc[:,nod_type==1]
ucx = ucx[nod_type==1]
ucy = ucy[nod_type==1]
xcc = xcc[nod_type==1]
ycc = ycc[nod_type==1]
# edge variables
link_type = subgrid.get_nd('link_type')
u1 = u1[link_type==1]
line = line[:,link_type[1:]==1]
xu = xu[link_type==1]
yu = yu[link_type==1]
# convert to pixel grid
if cf is None:
frictcoefuser = subgrid.get_nd('frictcoefuser')
pixels['cf'] = np.zeros_like(pixels['dps']) + frictcoefuser
else:
pixels['cf'] = cf
pixels['s1'] = s1[quad_grid]
# do another timestep
subgrid.update(-1)
Out[3]:
In [4]:
# split up u and v
is_u = np.zeros_like(u1, dtype='bool')
is_v = np.zeros_like(u1, dtype='bool')
# fill the boolean indices
liutot = subgrid.get_nd('liutot')
is_u[:liutot] = True
livtot = subgrid.get_nd('livtot')
is_v[(liutot):(liutot+livtot)] = True
np.all(is_u == ~is_v)
Out[4]:
In [12]:
# wait for velocities to reach the bay
for i in range(100):
subgrid.update(-1)
In [29]:
# 10 minutes
scale = 1/60.0
# split up u1 in u;v
u = is_u*u1
v = is_v*u1
In [44]:
# plot the grid and u and v
plt.figure(figsize=(20,13))
plt.plot(xu, yu, 'k+')
plt.plot(xcc, ycc, 'k.')
plt.quiver(xu, yu, u, v, np.sqrt(u**2 + v**2),
units='xy',
angles='xy',
scale_units='xy',
scale=scale,
alpha=0.5,
cmap='Blues')
plt.quiver(xcc, ycc, ucx, ucy, np.sqrt(ucx**2 + ucy**2),
units='xy',
angles='xy',
scale_units='xy',
scale=scale,
alpha=0.5,
cmap='Blues')
plt.colorbar()
plt.title('u max {}, v max {}'.format(u.max(), v.max()))
plt.axis('equal')
plt.xlim(544000, 548000)
plt.ylim(4183000, 4188000)
Out[44]:
In [33]:
In [41]:
# V1. px = F(x,y,ucc,vcc)
import scipy.interpolate
u1 = subgrid.get_nd('u1', sliced=False)[:]
u1 = u1[link_type==1]
uc = subgrid.get_nd('uc')
ucx, ucy = uc
ucx = ucx[nod_type==1]
ucy = ucy[nod_type==1]
F_u = scipy.interpolate.LinearNDInterpolator(np.c_[xu[is_u], yu[is_u]], u1[is_u])
F_v = scipy.interpolate.LinearNDInterpolator(np.c_[xu[is_v], yu[is_v]], u1[is_v])
pixels['U'] = F_u(np.c_[pixels['X'].ravel(), pixels['Y'].ravel()]).reshape(pixels['X'].shape).T
pixels['V'] = F_v(np.c_[pixels['X'].ravel(), pixels['Y'].ravel()]).reshape(pixels['X'].shape).T
pixels['u1'] = np.sqrt(ucx ** 2 + ucy **2)[quad_grid]
pixels['urad'] = np.arctan2(ucy, ucx)[quad_grid]
In [42]:
fig, axes = plt.subplots(3, 3, figsize=(30,24))
thin = 10
axes[0,0].imshow(pixels['s1'][::thin,::thin])
axes[0,1].imshow(dps[::thin,::thin], vmin=-100, vmax=100, cmap='gist_earth_r')
axes[0,2].imshow(pixels['cf'][::thin,::thin])
# Add the grid
im = axes[1,0].imshow(pixels['u1'][::thin,::thin], cmap='Greys', vmax=1.0)
plt.colorbar(im, ax=axes[1,0])
axes[1,1].imshow(pixels['urad'][::thin,::thin], vmin=-np.pi, vmax=np.pi, cmap='hsv')
im = axes[2,0].imshow(np.sqrt(pixels['U']**2 + pixels['V']**2)[::thin, ::thin], cmap='Greys', vmax=1.0)
plt.colorbar(im, ax=axes[2,0])
axes[2,1].imshow(np.arctan2(pixels['V'],pixels['U'])[::thin, ::thin], vmin=-np.pi, vmax=np.pi, cmap='hsv')
Out[42]:
In [268]:
U = pixels['U']/pixels['dxp']
V = pixels['V']/pixels['dyp']
U = U[::thin,::thin].astype('float32')*thin
V = V[::thin,::thin].astype('float32')*thin
img = np.zeros(U.shape + (3,), dtype='uint8') + 255
import skimage.draw
brush = 8
for i in range(1000):
cy = np.random.random(1)*(U.shape[0] )
cx = np.random.random(1)*(U.shape[1] )
rr, cc = skimage.draw.circle(cy, cx, brush, shape=U.shape)
img[rr, cc,:] = 128
plt.imshow(img)
Out[268]:
In [262]:
#plt.imshow(U, cmap='Greys')
U_rgba = np.rollaxis(U.view('uint8').reshape(U.shape[0], 4, U.shape[1]), 1, start=3)
V_rgba = np.rollaxis(V.view('uint8').reshape(V.shape[0], 4, V.shape[1]), 1, start=3)
#U_rgba = np.rollaxis(U.view('uint8'), 0, start=3)
plt.imshow(U, cmap='Greys')
plt.imshow(U_rgba)
import skimage.io
import os
skimage.io.imsave(os.path.join(subgrid.original_dir, 'U.png'), U_rgba)
skimage.io.imsave(os.path.join(subgrid.original_dir, 'V.png'), V_rgba)
#skimage.io.imsave('V.png', V_rgba)
In [263]:
U.shape, V.shape, UV.shape, img.shape, U.dtype
Out[263]:
In [266]:
import cv2
# convert to pixels/s
#img = np.random.randint(0,255, U.shape + (3,), ).astype('uint8')
#img = np.random.randint(0, 256, size=U.shape).astype('uint8')
fig, axes = plt.subplots(4,5, figsize=(15,20))
h, w = U.shape
UV = np.concatenate([U[:,:,np.newaxis], V[:,:, np.newaxis]], axis=2)
UV = -UV * 5
UV[:,:,0] += np.arange(w)
UV[:,:,1] += np.arange(h)[:,np.newaxis]
for i in range(100):
img = cv2.remap(img, UV[:,:,0], UV[:,:,1], interpolation=cv2.INTER_LINEAR)
for i, ax in zip(range(np.prod(axes.shape)), iter(axes.flat)):
ax.imshow(img)
img = cv2.remap(img, UV[:,:,0], UV[:,:,1], interpolation=cv2.INTER_LINEAR)
In [269]:
import matplotlib.animation
fig, ax = plt.subplots(figsize=(8,6))
def frames(img):
rr, cc = skimage.draw.circle(180, 430, 10)
for i in range(2000):
img[rr,cc,:2] = (160,128)
img = cv2.remap(img, UV, None, interpolation=cv2.INTER_LINEAR)
yield img
def func(img, *args, **kwargs):
ax.imshow(img)
anim = matplotlib.animation.FuncAnimation(fig=fig, func=func, frames=frames(img))
anim.save('a.mp4', fps=12, extra_args=['-vcodec', 'libx264'])
import os
os.getcwd()
#os.path.exists('a.webm')
Out[269]:
In [36]:
for i in range(200):
subgrid.update(-1)
In [108]:
plt.figure(figsize=(13,20))
plt.imshow(img)
plt.colorbar()
Out[108]:
In [208]:
import IPython.display
In [209]:
IPython.display.HTML('<img src="U.png" id="U"><img src="V.png" id="V">')
Out[209]:
In [ ]: