def corr_uvw(uvw, p, rho, dt, dxyz, obst):
# --------------------------------------------------------------------------
"""
Docstring.
"""
# Unpack received tuples
dx, dy, dz = dxyz
# Compute pressure correction gradients
p_x = dif(X, p.val) / avg(X, dx)
p_y = dif(Y, p.val) / avg(Y, dy)
p_z = dif(Z, p.val) / avg(Z, dz)
# Set to zero in obst
if obst.any() != 0:
p_x = obst_zero_val(X, p_x, obst)
p_y = obst_zero_val(Y, p_y, obst)
p_z = obst_zero_val(Z, p_z, obst)
# Pad with boundary values by expanding from interior
# (This is done only for collocated formulation)
if uvw[X].pos == C:
p_x = avg(X, cat(X, (p_x[:1, :, :], p_x, p_x[-1:, :, :])))
p_y = avg(Y, cat(Y, (p_y[:, :1, :], p_y, p_y[:, -1:, :])))
p_z = avg(Z, cat(Z, (p_z[:, :, :1], p_z, p_z[:, :, -1:])))
# Correct the velocities
uvw[X].val[:] = uvw[X].val[:] - dt / avg(uvw[X].pos, rho) * p_x
uvw[Y].val[:] = uvw[Y].val[:] - dt / avg(uvw[Y].pos, rho) * p_y
uvw[Z].val[:] = uvw[Z].val[:] - dt / avg(uvw[Z].pos, rho) * p_z
return # end of function
def plot_isolines(phi, uvw, xyzn, d):
"""
Docstring.
"""
# Unpack tuples
u, v, w = uvw
xn, yn, zn = xyzn
# Cell coordinates
xc = avg(xn)
yc = avg(yn)
zc = avg(zn)
# Collocated velocity components
if u.pos == C:
uc = u.val
vc = v.val
wc = w.val
else:
uc = avg(X, cat(X, (u.bnd[W].val[:1, :, :], \
u.val, \
u.bnd[E].val[:1, :, :])))
vc = avg(Y, cat(Y, (v.bnd[S].val[:, :1, :], \
v.val, \
v.bnd[N].val[:, :1, :])))
wc = avg(Z, cat(Z, (w.bnd[B].val[:, :, :1], \
w.val, \
w.bnd[T].val[:, :, :1])))
# Pick coordinates for plotting (xp, yp) and values for plotting
if d == Y:
jp = floor(yc.size / 2)
xp, yp = meshgrid(xc, zc)
zp = transpose(phi[:, jp, :], (1, 0))
up = transpose(uc[:, jp, :], (1, 0))
vp = transpose(wc[:, jp, :], (1, 0))
if d == Z:
kp = floor(zc.size / 2)
xp, yp = meshgrid(xc, yc)
zp = transpose(phi[:, :, kp], (1, 0))
up = transpose(uc[:, :, kp], (1, 0))
vp = transpose(vc[:, :, kp], (1, 0))
# Set levels and normalize the colors
levels = linspace(zp.min(), zp.max(), 11)
norm = cm.colors.Normalize(vmax=zp.max(), vmin=zp.min())
plt.figure()
plt.gca(aspect='equal')
plt.contour(xp, yp, zp, levels, cmap=plt.cm.rainbow, norm=norm)
plt.quiver(xp, yp, up, vp)
plt.axis([min(xn), max(xn), min(yn), max(yn)])
plt.show()
return # end of function
def vol_balance(uvwf, dxyz, obst):
"""
Computes the volume balance, which is essentially a right hand side in
the Poisson's equation for pressure.
Note that "obst" is an optional parameter. If it is not sent, the source
won't be zeroed inside the obstacle. That is important for calculation
of pressure, see function "calc_p".
"""
# Unpack tuples
dx, dy, dz = dxyz
uf, vf, wf = uvwf
# Compute it throughout the domain
src = - dif(X, cat(X, (uf.bnd[W].val, uf.val, uf.bnd[E].val)))*dy*dz \
- dif(Y, cat(Y, (vf.bnd[S].val, vf.val, vf.bnd[N].val)))*dx*dz \
- dif(Z, cat(Z, (wf.bnd[B].val, wf.val, wf.bnd[T].val)))*dx*dy
# Zero it inside obstacles, if obstacle is sent as parameter
if obst.any() != 0:
src = obst_zero_val(C, src, obst)
return src # end of function
def calc_uvw(uvw, uvwf, rho, mu, p_tot, e_f, dt, dxyz, obst):
"""
Calculate uvw
TODO: Write useful docstring
:param uvw:
:param uvwf:
:param rho:
:param mu:
:param p_tot:
:param e_f:
:param dt:
:param dxyz:
:param obst:
:return:
"""
# pylint: disable=invalid-name, too-many-locals, too-many-statements
# Unpack tuples
u, v, w = uvw
uf, vf, wf = uvwf
dx, dy, dz = dxyz
e_x, e_y, e_z = e_f
# Fetch resolutions
ru = u.val.shape
rv = v.val.shape
rw = w.val.shape
# Pre-compute geometrical quantities
dv = dx * dy * dz
d = u.pos
# Create linear systems
A_u, b_u = create_matrix(u, rho / dt, mu, dxyz, obst, 'd')
A_v, b_v = create_matrix(v, rho / dt, mu, dxyz, obst, 'd')
A_w, b_w = create_matrix(w, rho / dt, mu, dxyz, obst, 'd')
# Advection terms for momentum
c_u = advection(rho, u, uvwf, dxyz, dt, 'superbee')
c_v = advection(rho, v, uvwf, dxyz, dt, 'superbee')
c_w = advection(rho, w, uvwf, dxyz, dt, 'superbee')
# Innertial term for momentum (this works for collocated and staggered)
i_u = u.old * avg(u.pos, rho) * avg(u.pos, dv) / dt
i_v = v.old * avg(v.pos, rho) * avg(v.pos, dv) / dt
i_w = w.old * avg(w.pos, rho) * avg(w.pos, dv) / dt
# Compute staggered pressure gradients
p_tot_x = dif(X, p_tot) / avg(X, dx)
p_tot_y = dif(Y, p_tot) / avg(Y, dy)
p_tot_z = dif(Z, p_tot) / avg(Z, dz)
# Make pressure gradients cell-centered
if d == C:
p_tot_x = avg(X,
cat(X, (p_tot_x[:1, :, :], p_tot_x, p_tot_x[-1:, :, :])))
p_tot_y = avg(Y,
cat(Y, (p_tot_y[:, :1, :], p_tot_y, p_tot_y[:, -1:, :])))
p_tot_z = avg(Z,
cat(Z, (p_tot_z[:, :, :1], p_tot_z, p_tot_z[:, :, -1:])))
# Total pressure gradients (this works for collocated and staggered)
p_st_u = p_tot_x * avg(u.pos, dv)
p_st_v = p_tot_y * avg(v.pos, dv)
p_st_w = p_tot_z * avg(w.pos, dv)
# Full force terms for momentum equations (collocated and staggered)
f_u = b_u - c_u + i_u - p_st_u + e_x * avg(u.pos, dv)
f_v = b_v - c_v + i_v - p_st_v + e_y * avg(v.pos, dv)
f_w = b_w - c_w + i_w - p_st_w + e_z * avg(w.pos, dv)
# Take care of obsts in the domian
if obst.any() != 0:
f_u = obst_zero_val(u.pos, f_u, obst)
f_v = obst_zero_val(v.pos, f_v, obst)
f_w = obst_zero_val(w.pos, f_w, obst)
# Solve for velocities
res0 = bicgstab(A_u, reshape(f_u, prod(ru)), tol=TOL)
res1 = bicgstab(A_v, reshape(f_v, prod(rv)), tol=TOL)
res2 = bicgstab(A_w, reshape(f_w, prod(rw)), tol=TOL)
u.val[:] = reshape(res0[0], ru)
v.val[:] = reshape(res1[0], rv)
w.val[:] = reshape(res2[0], rw)
# Update velocities in boundary cells
adj_o_bnds((u, v, w), (dx, dy, dz), dt)
# Update face velocities (also substract cell-centered pressure gradients
# and add staggered pressure gradients)
if d == C:
uf.val[:] = (avg(X, u.val + dt / rho * (p_tot_x)) - dt / avg(X, rho) *
(dif(X, p_tot) / avg(X, dx)))
vf.val[:] = (avg(Y, v.val + dt / rho * (p_tot_y)) - dt / avg(Y, rho) *
(dif(Y, p_tot) / avg(Y, dy)))
wf.val[:] = (avg(Z, w.val + dt / rho * (p_tot_z)) - dt / avg(Z, rho) *
(dif(Z, p_tot) / avg(Z, dz)))
for j in (W, E):
uf.bnd[j].val[:] = u.bnd[j].val[:]
vf.bnd[j].val[:] = avg(Y, v.bnd[j].val[:])
wf.bnd[j].val[:] = avg(Z, w.bnd[j].val[:])
for j in (S, N):
uf.bnd[j].val[:] = avg(X, u.bnd[j].val[:])
vf.bnd[j].val[:] = v.bnd[j].val[:]
wf.bnd[j].val[:] = avg(Z, w.bnd[j].val[:])
for j in (B, T):
uf.bnd[j].val[:] = avg(X, u.bnd[j].val[:])
vf.bnd[j].val[:] = avg(Y, v.bnd[j].val[:])
wf.bnd[j].val[:] = w.bnd[j].val[:]
else:
uf.val[:] = u.val[:]
vf.val[:] = v.val[:]
wf.val[:] = w.val[:]
for j in (W, E, S, N, B, T):
uf.bnd[j].val[:] = u.bnd[j].val[:]
vf.bnd[j].val[:] = v.bnd[j].val[:]
wf.bnd[j].val[:] = w.bnd[j].val[:]
if obst.any() != 0:
uf.val[:] = obst_zero_val(X, uf.val, obst)
vf.val[:] = obst_zero_val(Y, vf.val, obst)
wf.val[:] = obst_zero_val(Z, wf.val, obst)
return # end of function
def advection(rho, phi, uvwf, dxyz, dt, lim_name):
"""
Docstring.
"""
res = phi.val.shape
nx, ny, nz = res
# Unpack tuples
uf, vf, wf = uvwf
dx, dy, dz = dxyz
d = phi.pos
# Pre-compute geometrical quantities
sx = dy * dz
sy = dx * dz
sz = dx * dy
# -------------------------------------------------
# Specific for cell-centered transported variable
# -------------------------------------------------
if d == C:
# Facial values of physical properties including boundary cells
rho_x_fac = cat(
X, (rho[:1, :, :], avg(X, rho), rho[-1:, :, :])) # nxp,ny, nz
rho_y_fac = cat(
Y, (rho[:, :1, :], avg(Y, rho), rho[:, -1:, :])) # nx, nyp,nz
rho_z_fac = cat(
Z, (rho[:, :, :1], avg(Z, rho), rho[:, :, -1:])) # nx, ny, nzp
# Facial values of areas including boundary cells
a_x_fac = cat(X, (sx[:1, :, :], avg(X, sx), sx[-1:, :, :]))
a_y_fac = cat(Y, (sy[:, :1, :], avg(Y, sy), sy[:, -1:, :]))
a_z_fac = cat(Z, (sz[:, :, :1], avg(Z, sz), sz[:, :, -1:]))
del_x = avg(X, dx)
del_y = avg(Y, dy)
del_z = avg(Z, dz)
# Facial values of velocities without boundary values
u_fac = uf.val # nxm,ny, nz
v_fac = vf.val # nx, nym,nz
w_fac = wf.val # nx, ny, nzm
# Boundary velocity values
u_bnd_W = uf.bnd[W].val
u_bnd_E = uf.bnd[E].val
v_bnd_S = vf.bnd[S].val
v_bnd_N = vf.bnd[N].val
w_bnd_B = wf.bnd[B].val
w_bnd_T = wf.bnd[T].val
# ------------------------------------------------------------
# Specific for transported variable staggered in x direction
# ------------------------------------------------------------
if d == X:
# Facial values of physical properties including boundary cells
rho_x_fac = rho # nx, ny, nz
rho_nod_y = avg(X, avg(Y, rho)) # nxm,nym,nz
rho_y_fac = cat(Y, (rho_nod_y[:, :1, :], \
rho_nod_y[:, :, :], \
rho_nod_y[:, -1:, :])) # nxm,nyp,nz
rho_nod_z = avg(X, avg(Z, rho)) # nxm,ny,nzm
rho_z_fac = cat(Z, (rho_nod_z[:, :, :1], \
rho_nod_z[:, :, :], \
rho_nod_z[:, :, -1:])) # nxm,ny,nzp
# Facial values of areas including boundary cells
a_x_fac = sx
a_y_fac = cat(Y, ( \
avg(X, sy[:, :1, :]),\
avg(X, avg(Y, sy)), \
avg(X, sy[:, -1:, :])))
a_z_fac = cat(Z, ( \
avg(X, sz[:, :, :1]), \
avg(X, avg(Z, sz)), \
avg(X, sz[:, :, -1:])))
del_x = dx[1:-1, :, :]
del_y = avg(X, avg(Y, dy))
del_z = avg(X, avg(Z, dz))
# Facial values of velocities without boundary values
u_fac = avg(X, uf.val) # nxmm,ny, nz
v_fac = avg(X, vf.val) # nxm, nym,nz
w_fac = avg(X, wf.val) # nxm, ny, nzm
# Boundary velocity values
u_bnd_W = uf.bnd[W].val
u_bnd_E = uf.bnd[E].val
v_bnd_S = avg(X, vf.bnd[S].val)
v_bnd_N = avg(X, vf.bnd[N].val)
w_bnd_B = avg(X, wf.bnd[B].val)
w_bnd_T = avg(X, wf.bnd[T].val)
# ------------------------------------------------------------
# Specific for transported variable staggered in y direction
# ------------------------------------------------------------
if d == Y:
# Facial values of physical properties including boundary cells
rho_nod_x = avg(Y, avg(X, rho)) # nxm,nym,nz
rho_x_fac = cat(X, (rho_nod_x[:1, :, :], \
rho_nod_x[:, :, :], \
rho_nod_x[-1:, :, :])) # nxp,nym,nz
rho_y_fac = rho # nx, ny, nz
rho_nod_z = avg(Y, avg(Z, rho)) # nx, nym,nzm
rho_z_fac = cat(Z, (rho_nod_z[:, :, :1], \
rho_nod_z[:, :, :], \
rho_nod_z[:, :, -1:])) # nx, nym,nzp
# Facial values of areas including boundary cells
a_x_fac = cat(X, ( \
avg(Y, sx[:1, :, :]), \
avg(Y, avg(X, sx)), \
avg(Y, sx[-1:, :, :])))
a_y_fac = sy
a_z_fac = cat(Z, ( \
avg(Y, sz[:, :, :1]), \
avg(Y, avg(Z, sz)), \
avg(Y, sz[:, :, -1:])))
del_x = avg(Y, avg(X, dx))
del_y = dy[:, 1:-1, :]
del_z = avg(Y, avg(Z, dz))
# Facial values of velocities without boundary values
u_fac = avg(Y, uf.val) # nxm,nym, nz
v_fac = avg(Y, vf.val) # nx, nymm,nz
w_fac = avg(Y, wf.val) # nx, nym, nzm
# Facial values of velocities with boundary values
u_bnd_W = avg(Y, uf.bnd[W].val)
u_bnd_E = avg(Y, uf.bnd[E].val)
v_bnd_S = vf.bnd[S].val
v_bnd_N = vf.bnd[N].val
w_bnd_B = avg(Y, wf.bnd[B].val)
w_bnd_T = avg(Y, wf.bnd[T].val)
# ------------------------------------------------------------
# Specific for transported variable staggered in z direction
# ------------------------------------------------------------
if d == Z:
# Facial values of physical properties including boundary cells
rho_nod_x = avg(Z, avg(X, rho)) # nxm,ny, nzm
rho_x_fac = cat(X, (rho_nod_x[ :1, :, :], \
rho_nod_x[ :, :, :], \
rho_nod_x[-1:, :, :])) # nxp,ny, nzm
rho_nod_y = avg(Z, avg(Y, rho)) # nx, nym,nzm
rho_y_fac = cat(Y, (rho_nod_y[:, :1, :], \
rho_nod_y[:, :, :], \
rho_nod_y[:, -1:, :])) # nx, nyp,nzm
rho_z_fac = rho # nx, ny, nz
# Facial values of areas including boundary cells
a_x_fac = cat(X, ( \
avg(Z, sx[:1, :, :]), \
avg(Z, avg(X, sx)), \
avg(Z, sx[-1:, :, :])))
a_y_fac = cat(Y, ( \
avg(Z, sy[:, :1, :]), \
avg(Z, avg(Y, sy)), \
avg(Z, sy[:, -1:, :])))
a_z_fac = sz
del_x = avg(Z, avg(X, dx))
del_y = avg(Z, avg(Y, dy))
del_z = dz[:, :, 1:-1]
# Facial values of velocities without boundary values
u_fac = avg(Z, uf.val) # nxm,ny, nzm
v_fac = avg(Z, vf.val) # nx, nym, nzm
w_fac = avg(Z, wf.val) # nx, ny, nzmm
# Facial values of velocities with boundary values
u_bnd_W = avg(Z, uf.bnd[W].val)
u_bnd_E = avg(Z, uf.bnd[E].val)
v_bnd_S = avg(Z, vf.bnd[S].val)
v_bnd_N = avg(Z, vf.bnd[N].val)
w_bnd_B = wf.bnd[B].val
w_bnd_T = wf.bnd[T].val
# ------------------------------
# Common part of the algorithm
# ------------------------------
# ------------------------------------------------------------
#
# |-o-|-o-|-o-|-o-|-o-|-o-|-o-|-o-|-o-|-o-|
# 1 2 3 4 5 6 7 8 9 10 phi
# x---x---x---x---x---x---x---x---x
# 1 2 3 4 5 6 7 8 9 d_x initial
# 0---x---x---x---x---x---x---x---x---x---0
# 1 2 3 4 5 6 7 8 9 10 11 d_x padded
#
# ------------------------------------------------------------
# Compute consecutive differences (and avoid division by zero)
d_x = dif(X, phi.val) # nxm, ny, nz
d_x[(d_x > -TINY) & (d_x <= 0.0)] = -TINY
d_x[(d_x >= 0.0) & (d_x < +TINY)] = +TINY
d_x = cat(X, (d_x[:1, :, :], d_x, d_x[-1:, :, :])) # nxp, ny, nz
d_y = dif(Y, phi.val) # nx, nym, nz
d_y[(d_y > -TINY) & (d_y <= 0.0)] = -TINY
d_y[(d_y >= 0.0) & (d_y < +TINY)] = +TINY
d_y = cat(Y, (d_y[:, :1, :], d_y, d_y[:, -1:, :])) # nx, nyp, nz
d_z = dif(Z, phi.val) # nx, ny, nzm
d_z[(d_z > -TINY) & (d_z <= 0.0)] = -TINY
d_z[(d_z >= 0.0) & (d_z < +TINY)] = +TINY
d_z = cat(Z, (d_z[:, :, :1], d_z, d_z[:, :, -1:])) # nx, ny, nzp
# Ratio of consecutive gradients for positive and negative flow
r_x_we = d_x[1:-1, :, :] / d_x[0:-2, :, :] # nxm,ny, nz
r_x_ew = d_x[2:, :, :] / d_x[1:-1, :, :] # nxm,ny, nz
r_y_sn = d_y[:, 1:-1, :] / d_y[:, 0:-2, :] # nx, nym,nz
r_y_ns = d_y[:, 2:, :] / d_y[:, 1:-1, :] # nx, nym,nz
r_z_bt = d_z[:, :, 1:-1] / d_z[:, :, 0:-2] # nx, ny, nzm
r_z_tb = d_z[:, :, 2:] / d_z[:, :, 1:-1] # nx, ny, nzm
flow_we = u_fac >= 0
flow_ew = lnot(flow_we)
flow_sn = v_fac >= 0
flow_ns = lnot(flow_sn)
flow_bt = w_fac >= 0
flow_tb = lnot(flow_bt)
r_x = r_x_we * flow_we + r_x_ew * flow_ew
r_y = r_y_sn * flow_sn + r_y_ns * flow_ns
r_z = r_z_bt * flow_bt + r_z_tb * flow_tb
# Apply a limiter
if lim_name == 'upwind':
psi_x = r_x * 0.0
psi_y = r_y * 0.0
psi_z = r_z * 0.0
elif lim_name == 'minmod':
psi_x = mx(zeros(r_x.shape), mn(r_x, ones(r_x.shape)))
psi_y = mx(zeros(r_y.shape), mn(r_y, ones(r_y.shape)))
psi_z = mx(zeros(r_z.shape), mn(r_z, ones(r_z.shape)))
elif lim_name == 'superbee':
psi_x = mx(zeros(r_x.shape), mn(2. * r_x, ones(r_x.shape)),
mn(r_x, 2.))
psi_y = mx(zeros(r_y.shape), mn(2. * r_y, ones(r_y.shape)),
mn(r_y, 2.))
psi_z = mx(zeros(r_z.shape), mn(2. * r_z, ones(r_z.shape)),
mn(r_z, 2.))
elif lim_name == 'koren':
psi_x = mx(zeros(r_x.shape), mn(2.*r_x, (2.+r_x)/3., \
2.*ones(r_x.shape)))
psi_y = mx(zeros(r_y.shape), mn(2.*r_y, (2.+r_y)/3., \
2.*ones(r_y.shape)))
psi_z = mx(zeros(r_z.shape), mn(2.*r_z, (2.+r_z)/3., \
2.*ones(r_z.shape)))
flux_fac_lim_x = phi.val[0:-1, :, :] * u_fac * flow_we \
+ phi.val[1:, :, :] * u_fac * flow_ew \
+ 0.5 * abs(u_fac) * (1 - abs(u_fac) * dt / del_x) \
* (psi_x[:, :, :] * d_x[0:nx-1, :, :] * flow_we \
+ psi_x[:, :, :] * d_x[1:nx, :, :] * flow_ew)
flux_fac_lim_y = phi.val[:, 0:-1, :] * v_fac * flow_sn \
+ phi.val[:, 1:, :] * v_fac * flow_ns \
+ 0.5 * abs(v_fac) * (1 - abs(v_fac) * dt / del_y) \
* (psi_y[:, :, :] * d_y[:, 0:ny-1, :] * flow_sn \
+ psi_y[:, :, :] * d_y[:, 1:ny, :] * flow_ns)
flux_fac_lim_z = phi.val[:, :, 0:-1] * w_fac * flow_bt \
+ phi.val[:, :, 1: ] * w_fac * flow_tb \
+ 0.5 * abs(w_fac) * (1 - abs(w_fac) * dt / del_z) \
* (psi_z[:, :, :] * d_z[:, :, 0:nz-1] * flow_bt \
+ psi_z[:, :, :] * d_z[:, :, 1:nz] * flow_tb)
# Pad with boundary values
flux_fac_lim_x = cat(X, (phi.bnd[W].val * u_bnd_W, \
flux_fac_lim_x, \
phi.bnd[E].val * u_bnd_E))
flux_fac_lim_y = cat(Y, (phi.bnd[S].val * v_bnd_S, \
flux_fac_lim_y, \
phi.bnd[N].val * v_bnd_N))
flux_fac_lim_z = cat(Z, (phi.bnd[B].val * w_bnd_B, \
flux_fac_lim_z, \
phi.bnd[T].val * w_bnd_T))
# Multiply with face areas
flux_fac_lim_x = rho_x_fac * flux_fac_lim_x * a_x_fac
flux_fac_lim_y = rho_y_fac * flux_fac_lim_y * a_y_fac
flux_fac_lim_z = rho_z_fac * flux_fac_lim_z * a_z_fac
# Sum contributions from all directions up
c = dif(X, flux_fac_lim_x) + \
dif(Y, flux_fac_lim_y) + \
dif(Z, flux_fac_lim_z)
return c # end of function
def create_matrix(phi, inn, mu, dxyz, obst, obc):
"""
pos - position of variable (C - central,
X - staggered in x direction,
Y - staggered in y direction,
Z - staggered in z direction)
inn - innertial term
mu - viscous coefficient
dx, dy, dz - cell size in x, y and z directions
obc - obstacles's boundary condition, ('n' - Neumann,
'd' - Dirichlet)
-------------------------------------------------------------------------
"""
# Unpack tuples
dx, dy, dz = dxyz
res = phi.val.shape
# -------------------------------
# Create right hand side vector
# -------------------------------
b = zeros(res)
# ------------------------------------
# Create default matrix coefficients
# ------------------------------------
coefficients = namedtuple('matrix_diagonal', 'W E S N B T P')
c = coefficients(zeros(res), zeros(res), zeros(res), \
zeros(res), zeros(res), zeros(res), \
zeros(res))
d = phi.pos
# Handle central coefficient due to innertia
c.P[:] = avg(d, inn) * avg(d, dx * dy * dz)
# Pre-compute geometrical quantities
sx = dy * dz
sy = dx * dz
sz = dx * dy
if d != X:
c.W[:] = cat(X, ( \
avg(d, mu[:1, :, :]) \
* avg(d, sx[:1, :, :]) / avg(d, (dx[:1, :, :])/2.0), \
avg(d, avg(X, mu)) * avg(d, avg(X, sx)) / avg(d, avg(X, dx))))
c.E[:] = cat(X, ( \
avg(d, avg(X, mu)) * avg(d, avg(X, sx)) / avg(d, avg(X, dx)), \
avg(d, mu[-1:, :, :]) \
* avg(d, sx[-1:, :, :])/ avg(d, (dx[-1:, :, :])/2.0)))
if d != Y:
c.S[:] = cat(Y, ( \
avg(d, mu[:, :1, :]) \
* avg(d, sy[:, :1, :]) / avg(d, (dy[:, :1, :])/2.0), \
avg(d, avg(Y, mu)) * avg(d, avg(Y, sy)) / avg(d, avg(Y, dy))))
c.N[:] = cat(Y, ( \
avg(d, avg(Y, mu)) * avg(d, avg(Y, sy)) / avg(d, avg(Y, dy)), \
avg(d, mu[:, -1:, :]) \
* avg(d, sy[:, -1:, :]) / avg(d, (dy[:, -1:, :])/2.0)))
if d != Z:
c.B[:] = cat(Z, ( \
avg(d, mu[:, :, :1]) \
* avg(d, sz[:, :, :1]) / avg(d, (dz[:, :, :1])/2.0), \
avg(d, avg(Z, mu)) * avg(d, avg(Z, sz)) / avg(d, avg(Z, dz))))
c.T[:] = cat(Z, ( \
avg(d, avg(Z, mu)) * avg(d, avg(Z, sz)) / avg(d, avg(Z, dz)), \
avg(d, mu[:, :, -1:]) \
* avg(d, sz[:, :, -1:]) / avg(d, (dz[:, :, -1:])/2.0)))
# ---------------------------------
# Correct for staggered variables
# ---------------------------------
if d == X:
c.W[:] = mu[0:-1, :, :] * sx[0:-1, :, :] / dx[0:-1, :, :]
c.E[:] = mu[1:, :, :] * sx[1:, :, :] / dx[1:, :, :]
elif d == Y:
c.S[:] = mu[:, 0:-1, :] * sy[:, 0:-1, :] / dy[:, 0:-1, :]
c.N[:] = mu[:, 1:, :] * sy[:, 1:, :] / dy[:, 1:, :]
elif d == Z:
c.B[:] = mu[:, :, 0:-1] * sz[:, :, 0:-1] / dz[:, :, 0:-1]
c.T[:] = mu[:, :, 1:] * sz[:, :, 1:] / dz[:, :, 1:]
# -----------------------------------------------------------------------
# Zero them (correct them) for vanishing derivative boundary condition.
# -----------------------------------------------------------------------
# The values defined here will be false (numerical value 0)
# wherever there is or Neumann boundary condition.
c.W[:1, :, :] *= (phi.bnd[W].typ[:] == DIRICHLET)
c.E[-1:, :, :] *= (phi.bnd[E].typ[:] == DIRICHLET)
c.S[:, :1, :] *= (phi.bnd[S].typ[:] == DIRICHLET)
c.N[:, -1:, :] *= (phi.bnd[N].typ[:] == DIRICHLET)
c.B[:, :, :1] *= (phi.bnd[B].typ[:] == DIRICHLET)
c.T[:, :, -1:] *= (phi.bnd[T].typ[:] == DIRICHLET)
# --------------------------------------------
# Fill the source terms with boundary values
# --------------------------------------------
b[:1, :, :] += c.W[:1, :, :] * phi.bnd[W].val[:1, :, :]
b[-1:, :, :] += c.E[-1:, :, :] * phi.bnd[E].val[:1, :, :]
b[:, :1, :] += c.S[:, :1, :] * phi.bnd[S].val[:, :1, :]
b[:, -1:, :] += c.N[:, -1:, :] * phi.bnd[N].val[:, :1, :]
b[:, :, :1] += c.B[:, :, :1] * phi.bnd[B].val[:, :, :1]
b[:, :, -1:] += c.T[:, :, -1:] * phi.bnd[T].val[:, :, :1]
# ---------------------------------------
# Correct system matrices for obstacles
# ---------------------------------------
if obst.any() != 0:
c = obst_mod_matrix(phi, c, obst, obc)
# -----------------------------------------------
# Add all neighbours to the central matrix,
# and zero the coefficients towards boundaries
# -----------------------------------------------
c.P[:] += c.W[:] + c.E[:] + c.S[:] + c.N[:] + c.B[:] + c.T[:]
c.W[:1, :, :] = 0.0
c.E[-1:, :, :] = 0.0
c.S[:, :1, :] = 0.0
c.N[:, -1:, :] = 0.0
c.B[:, :, :1] = 0.0
c.T[:, :, -1:] = 0.0
# ----------------------
# Create sparse matrix
# ----------------------
nx, ny, nz = res
n = nx * ny * nz
data = array([reshape(c.P, n), \
reshape(-c.W, n), reshape(-c.E, n), \
reshape(-c.S, n), reshape(-c.N, n), \
reshape(-c.B, n), reshape(-c.T, n)])
diag = array([0, +ny * nz, -ny * nz, +nz, -nz, +1, -1])
A = spdiags(data, diag, n, n)
return A, b # end of function