def comp_mom_eq(grids, point):
# Roll coordinate system if necessary
s = grids[0].shape
point = ndimed.roller(point, s)
ndim = len(grids)
# Stupid Check!
for grid in grids:
if grid.shape != s or len(grid.shape) != len(point):
print "Shape mismatch!"
raise TypeError
# Initialize the equations
mom_eqs = []
for x in range(ndim):
mom_eqs.append(Equation())
# Iterate through the dimensions
for dim in range(ndim):
# Make a point
# Shifted by 1
# In the dimension we are working in
shifted_point = array(point)
shifted_point[dim] += 1
shifted_point = tuple(shifted_point)
shifted_point = ndimed.roller(shifted_point, s)
# Think M(right) - M(left) (pg.84 in Dr. E. Thesis)
# TODO: implement __iadd__ in Equation
mom_eqs[dim] = mom_eqs[dim] + momentum_eq( grids[dim], shifted_point )
mom_eqs[dim] = mom_eqs[dim] - momentum_eq( grids[dim], point )
# Not obvious but these are the U, V, (W) components . . .
return mom_eqs
def lap_div(pdof, grids, point):
s = grids[0].shape
ndim = len(grids)
# Center div equation
center_eq = div_eq( grids, point )
# Initialize the equations
div_lap_eqs = []
for x in range( ndim ):
div_lap_eqs.append( Equation() )
for pp in ndimed.perturb(point):
# PBC correction
pp = ndimed.roller(pp, s)
# If there is a single -3 dof, it must be a solid square
if pdof[pp] < 0:
continue
pert_eq = div_eq(grids, pp)
for dim in range(ndim):
div_lap_eqs[dim] = div_lap_eqs[dim] + ( pert_eq[dim] - center_eq[dim] )
return div_lap_eqs
def momentum_eq(dof_grid, point):
# Correct for pbc's
shape = dof_grid.shape
ndim = len(shape)
point = ndimed.roller(point, shape)
# If solid . . . there is no momentum equation
if dof_grid[ point ] < 0:
return Equation()
# DOF at the center point
center_dof_num = dof_grid[ point ]
m_eq = Equation()
m_eq[ center_dof_num ] = -2 * ndim
# Move one step in each direction
for p in ndimed.perturb( point ):
# Wrap where necessary
p = ndimed.roller(p, shape)
# Get the current dof number
current_dof = dof_grid[p]
# If moving in one direction is a dof
if current_dof >= 0:
# This could glitch if the domain is small enough
# For the laplacian est. to touch itself (3x3 or smaller . . .)
m_eq[ current_dof ] = 1
# If it is solid
elif current_dof == -2:
continue
# If it is completely enclosed solid
elif current_dof == -3:
m_eq[ center_dof_num ] -= 1
return m_eq
def div_eq(grids, point):
# Because any time this is called on a solid centered cell
# it will return all zeros, we don't worry about that
s = grids[0].shape
ndim = len(grids)
# Stupid Check!
for grid in grids:
if grid.shape != s or len(grid.shape) != len(point):
print "Shape mismatch!"
raise TypeError
# PCB correction
point = ndimed.roller(point, s)
# Initialize the equations
div_eqs = []
for x in range( ndim ):
div_eqs.append( Equation() )
#iterate through the dimensions, grids, and corresponding momentum equations
for dim in range(ndim):
# Shifted by 1 in the n-th axis
shifted_point = array(point)
shifted_point[dim] += 1
shifted_point = tuple(shifted_point)
shifted_point = ndimed.roller(shifted_point, s)
# Inflow
dof = grids[dim][point]
if dof >= 0:
div_eqs[dim][dof] = -1
# outflow
dof = grids[dim][shifted_point]
if dof >= 0:
div_eqs[dim][dof] = 1
return div_eqs
