def cartesian_grid(xn, yn, zn):
"""
Spawns a 3D Cartesian grid from three arrays with node coordinates.
"""
# Compute cell resolutions
nx = len(xn)-1
ny = len(yn)-1
nz = len(zn)-1
# Create matrices for cell dimensions ...
dx = empty((nx, ny, nz))
dy = empty((nx, ny, nz))
dz = empty((nx, ny, nz))
# ... and fill them up!
dx[:] = dif(xn).reshape(nx, 1, 1)
dy[:] = dif(yn).reshape(1, ny, 1)
dz[:] = dif(zn).reshape(1, 1, nz)
# Compute resolutions for cell-centered and all collocated variables
rc = nx, ny, nz
ru = nx-1, ny, nz
rv = nx, ny-1, nz
rw = nx, ny, nz-1
return nx, ny, nz, \
dx, dy, dz, \
rc, ru, rv, rw # end of function
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 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 obst_mod_matrix(phi, c, obst, obc):
"""
Adjusts the system matrix for obstacles and cell centered varaibles
(such as pressure)
phi - variable
c - coefficients in system matrix
obst - obstacle array
obc - obstacles's boundary condition, ('n' - Neumann, 'd' - Dirichlet)
"""
pos = phi.pos
#--------------------------
#
# For collocated variables
#
#--------------------------
if pos == C:
#------------------------------------
# Neumann's boundary on the obstacle
#------------------------------------
if obc == 'n':
# Correct west and east
sol_x = dif(X, obst) # will be +1 east of obst, -1 west of obst
corr = 1 - (sol_x < 0)
c.W[1:, :, :] = c.W[1:, :, :] * corr
corr = 1 - (sol_x > 0)
c.E[:-1, :, :] = c.E[:-1, :, :] * corr
# Correct south and north
sol_y = dif(Y, obst) # will be +1 north of obst, -1 south of obst
corr = 1 - (sol_y < 0)
c.S[:, 1:, :] = c.S[:, 1:, :] * corr
corr = 1 - (sol_y > 0)
c.N[:, :-1, :] = c.N[:, :-1, :] * corr
# Correct bottom and top
sol_z = dif(Z, obst) # will be +1 north of obst, -1 south of obst
corr = 1 - (sol_z < 0)
c.B[:, :, 1:] = c.B[:, :, 1:] * corr
corr = 1 - (sol_z > 0)
c.T[:, :, :-1] = c.T[:, :, :-1] * corr
#--------------------------------------
# Dirichlet's boundary on the obstacle
#--------------------------------------
elif obc == 'd':
# Set central coefficient to 1 in obst, unchanged elsewhere
c.P[:] = c.P[:] * lnot(obst) + obst
# Set neighbour coefficients to zero in obst
c.W[:] = c.W[:] * lnot(obst)
c.E[:] = c.E[:] * lnot(obst)
c.S[:] = c.S[:] * lnot(obst)
c.N[:] = c.N[:] * lnot(obst)
c.B[:] = c.B[:] * lnot(obst)
c.T[:] = c.T[:] * lnot(obst)
# Increase coefficients close to obst (makes sense for momentum)
sol_x = dif(X, obst) # will be +1 east of obst, -1 west of obst
corr = 1 + (sol_x > 0)
c.E[:-1, :, :] = c.E[:-1, :, :] * corr
corr = 1 + (sol_x < 0)
c.W[1:, :, :] = c.W[1:, :, :] * corr
sol_y = dif(Y, obst) # will be +1 north of obst, -1 south of obst
corr = 1 + (sol_y > 0)
c.N[:, :-1, :] = c.N[:, :-1, :] * corr
corr = 1 + (sol_y < 0)
c.S[:, 1:, :] = c.S[:, 1:, :] * corr
sol_z = dif(Z, obst) # will be +1 top of obst, -1 bottom of obst
corr = 1 + (sol_z > 0)
c.T[:, :, :-1] = c.T[:, :, :-1] * corr
corr = 1 + (sol_z < 0)
c.B[:, :, 1:] = c.B[:, :, 1:] * corr
#-------------------------
#
# For staggered variables
#
#-------------------------
elif pos == X:
# Set central coefficient to 1 in obst, unchanged elsewhere
obst_x = mx(obst[:-1, :, :], obst[1:, :, :])
c.P[:] = c.P[:] * lnot(obst_x) + obst_x
# Set neighbour coefficients to zero in obst
c.W[:] = c.W[:] * lnot(obst_x)
c.E[:] = c.E[:] * lnot(obst_x)
c.S[:] = c.S[:] * lnot(obst_x)
c.N[:] = c.N[:] * lnot(obst_x)
c.B[:] = c.B[:] * lnot(obst_x)
c.T[:] = c.T[:] * lnot(obst_x)
# Increase coefficients close to obst (makes sense for momentum)
sol_y = dif(Y, obst_x) # will be +1 north of obst,-1 south of obst
corr = 1 + (sol_y > 0)
c.N[:, :-1, :] = c.N[:, :-1, :] * corr
corr = 1 + (sol_y < 0)
c.S[:, 1:, :] = c.S[:, 1:, :] * corr
sol_z = dif(Z, obst_x) # will be +1 top of obst, -1 bottom of obst
corr = 1 + (sol_z > 0)
c.T[:, :, :-1] = c.T[:, :, :-1] * corr
corr = 1 + (sol_z < 0)
c.B[:, :, 1:] = c.B[:, :, 1:] * corr
elif pos == Y:
# Set central coefficient to 1 in obst, unchanged elsewhere
obst_y = mx(obst[:, :-1, :], obst[:, 1:, :])
c.P[:] = c.P[:] * lnot(obst_y) + obst_y
# Set neighbour coefficients to zero in obst
c.W[:] = c.W[:] * lnot(obst_y)
c.E[:] = c.E[:] * lnot(obst_y)
c.S[:] = c.S[:] * lnot(obst_y)
c.N[:] = c.N[:] * lnot(obst_y)
c.B[:] = c.B[:] * lnot(obst_y)
c.T[:] = c.T[:] * lnot(obst_y)
# Increase coefficients close to obst (makes sense for momentum)
sol_x = dif(X, obst_y) # will be +1 north of obst,-1 south of obst
corr = 1 + (sol_x > 0)
c.E[:-1, :, :] = c.E[:-1, :, :] * corr
corr = 1 + (sol_x < 0)
c.W[1:, :, :] = c.W[1:, :, :] * corr
sol_z = dif(Z, obst_y) # will be +1 north of obst,-1 south of obst
corr = 1 + (sol_z > 0)
c.T[:, :, :-1] = c.T[:, :, :-1] * corr
corr = 1 + (sol_z < 0)
c.B[:, :, 1:] = c.B[:, :, 1:] * corr
elif pos == Z:
# Set central coefficient to 1 in obst, unchanged elsewhere
obst_z = mx(obst[:, :, :-1], obst[:, :, 1:])
c.P[:] = c.P[:] * lnot(obst_z) + obst_z
# Set neighbour coefficients to zero in obst
c.W[:] = c.W[:] * lnot(obst_z)
c.E[:] = c.E[:] * lnot(obst_z)
c.S[:] = c.S[:] * lnot(obst_z)
c.N[:] = c.N[:] * lnot(obst_z)
c.B[:] = c.B[:] * lnot(obst_z)
c.T[:] = c.T[:] * lnot(obst_z)
# Increase coefficients close to obst (makes sense for momentum)
sol_x = dif(X, obst_z) # will be +1 north of obst,-1 south of obst
corr = 1 + (sol_x > 0)
c.E[:-1, :, :] = c.E[:-1, :, :] * corr
corr = 1 + (sol_x < 0)
c.W[1:, :, :] = c.W[1:, :, :] * corr
sol_y = dif(Y, obst_z) # will be +1 north of obst,-1 south of obst
corr = 1 + (sol_y > 0)
c.N[:, :-1, :] = c.N[:, :-1, :] * corr
corr = 1 + (sol_y < 0)
c.S[:, 1:, :] = c.S[:, 1:, :] * corr
return c # end of function