@@ -0,0 +1,314 @@
from dolfin import *
import numpy as np
import mshr
from IPython import embed
from cbcpost.utils.mpi_utils import gather, broadcast
def vec(x): return as_backend_type(x).vec()
def mat(A): return as_backend_type(A).mat()
class NormalRotation(PETScMatrix):
def __init__(self, V):
PETScMatrix.__init__(self)
self.V = V
self.D = V.num_sub_spaces()
self.U = Function(self.V)
dm = V.dofmap()
mesh = self.V.mesh()
dim = mesh.topology().dim()
assert dim == self.D
#index_blocks = []
indices = None
dof_max = dm.local_dimension("owned")
for d in range(V.num_sub_spaces()):
d_indices = DirichletBC(V.sub(d), Constant(0), "on_boundary").get_boundary_values().keys()
d_indices = np.array(d_indices)
d_indices.sort()
cut = np.argmax(d_indices >= dof_max)
d_indices = d_indices[:cut if cut!=0 else None]
if indices == None:
indices = d_indices
else:
indices = np.append(indices, d_indices)
self.indices = indices.reshape((dim, -1)).T
all_coords = dm.tabulate_all_coordinates(V.mesh()).reshape(-1,dim)
assert np.all(all_coords[self.indices[:,0]] == all_coords[self.indices.T])
# Create sparsity pattern
u, v = TrialFunction(V), TestFunction(V)
form = Constant(0)*inner(u,v)*ds()
assemble(form, tensor=self, keep_diagonal=True)
self._T = PETScMatrix(self)
@property
def T(self):
self._T.apply("insert")
return self._T
def build(self, N):
R = np.zeros((self.D,self.D))
RT = np.zeros((self.D,self.D))
def fill_2D(R,n):
R[0,:] = n
R[1,:] = R[0,1], -R[0,0]
def fill_3D(R,n):
x,y,z = n
theta = np.arctan2(-z,y)
alpha = np.arctan2(-(cos(theta)*y-sin(theta)*z), x)
ct = np.cos(theta)
st = np.sin(theta)
ca = np.cos(alpha)
sa = np.sin(alpha)
R[0,0] = ca
R[0,1] = -sa*ct
R[0,2] = sa*st
R[1,0] = sa
R[1,1] = ca*ct
R[1,2] = -ca*st
R[2,1] = st
R[2,2] = ct
if self.D == 2:
fill = fill_2D
else:
fill = fill_3D
self.zero()
identity = np.eye(self.D)
for i in self.indices:
_n = N.vector()[i]
_n_len = np.linalg.norm(_n)
if _n_len < 1e-12:
R[:] = identity
continue
_n /= _n_len
fill(R,_n)
i = i.astype(np.intc)
#assert np.max(abs(np.dot(R,R.T)-I)) < 1e-12
test = np.dot(R,_n)
#assert test[0] > 1-1e-12
#assert test[1] < 1e-12
#assert test[2] < 1e-12
RT[:] = R.T
self.set_local(R, i, i)
self._T.set_local(RT, i, i)
self.apply("insert")
self._T.apply("insert")
self.ident_zeros()
self._T.ident_zeros()
self.apply("insert")
self._T.apply("insert")
class NoSlipBC(DirichletBC):
def __init__(self, V, normal, *args, **kwargs):
DirichletBC.__init__(self, V, Constant((0.0,)*V.num_sub_spaces()), *args, **kwargs)
self.normal = normal
self.R = NormalRotation(V)
self.R.build(self.normal)
self.bcs = [DirichletBC(V.sub(d), Constant(0), *args, **kwargs) for d in range(1, V.num_sub_spaces())]
def apply(self, *args, **kwargs):
""" This will perform a change of basis on the dofs associated with the boundary.
The x-component will be replaced with the normal direction, and the y,z-components
associated with the tangential directions.
NOTE: The meaning of matrix A and vector b will be changed.
"""
assert 1<=len(args)<=2
recompute_normal = kwargs.get("recompute_normal", False)
A,b = None, None
if len(args) == 1 and isinstance(args[0], GenericMatrix):
A = args[0]
elif len(args) == 1 and isinstance(args[0], GenericVector):
b = args[0]
elif len(args) == 2:
assert isinstance(args[0], GenericMatrix)
assert isinstance(args[1], GenericVector)
A = args[0]
b = args[1]
if recompute_normal:
if isinstance(self.normal, Normal):
self.normal.compute()
else:
print "Unable to recompute normal: Not of class Normal"
self.R.build(self.normal)
print A,b
if b != None:
b[:] = self.R*b
for bc in self.bcs:
bc.apply(b)
if A != None:
# R.T*A*R
B = PETScMatrix((mat(self.R).matMult(mat(A)).matMult(mat(self.R).transpose())))
tic()
A.zero()
A.axpy(1.0, B, False)
for bc in self.bcs:
"""
dofs = np.array(bc.get_boundary_values().keys(), dtype=np.intc)
dofs.sort()
dof_max = self.function_space().dofmap().local_dimension("owned")
cut = np.argmax(dofs >= dof_max)
dofs = dofs[:cut if cut!=0 else None]
A.ident(dofs)
"""
bc.apply(A)
class Normal(Function):
def __init__(self, V, priority_order=None, facet_domains=None):
Function.__init__(self, V)
#facet_domains = facet_domains
priority_order = [] if priority_order == None else priority_order
if len(priority_order) > 0:
assert facet_domains
self.ds, self.bcs = self._priority(priority_order, facet_domains)
self.indices = self._get_boundary_indices()
self.compute()
def _get_boundary_indices(self):
# Used for normalization
V = self.function_space()
indices = None
dm = V.dofmap()
dim = V.num_sub_spaces()
dof_max = dm.local_dimension("owned")
for d in range(V.num_sub_spaces()):
d_indices = DirichletBC(V.sub(d), Constant(0), "on_boundary").get_boundary_values().keys()
d_indices = np.array(d_indices)
d_indices.sort()
cut = np.argmax(d_indices >= dof_max)
d_indices = d_indices[:cut if cut!=0 else None]
if indices == None:
indices = d_indices
else:
indices = np.append(indices, d_indices)
return indices.reshape((dim, -1)).T
def _priority(self, priority_order, facet_domains):
if len(priority_order) == 0:
ds = [Measure("ds")]
bcs = [None]
return ds, bcs
ds = Measure("ds")[facet_domains]
unprioritized = list(set(np.unique(facet_domains.array()))-set(priority_order))
unprioritized = gather(unprioritized,0)
unprioritized = np.unique(np.array(unprioritized))
unprioritized = broadcast(unprioritized, 0)
unprioritized = unprioritized.astype(np.int)
_ds = []
if len(unprioritized) > 0:
nds = ds(unprioritized[0])
for i in unprioritized[1:]: nds += ds(i)
_ds.append(nds)
_ds += [ds(i) for i in priority_order]
bcs = [DirichletBC(V, Constant((0,)*V.num_sub_spaces()), facet_domains, i) for i in priority_order]
if len(unprioritized) > 0:
bcs = [None]+bcs
return _ds, bcs
def compute(self, normalize=True):
n = FacetNormal(self.function_space().mesh())
v = TestFunction(self.function_space())
for ds, bc in zip(self.ds, self.bcs):
if bc != None:
bc.apply(self.vector())
fcp = dict()
self.vector()[:] += assemble(inner(v, n) * ds, form_compiler_parameters=fcp)
self.vector()[:] += assemble(inner(v, n) * ds, form_compiler_parameters=dict(quadrature_rule="vertex", quadrature_degree=0))
if normalize:
self.normalize()
def normalize(self):
tmp = self.vector().array()
#embed()
scale = np.linalg.norm(tmp[self.indices], axis=1)[:,np.newaxis]
tmp[self.indices] /= scale
self.vector()[:] = tmp
if __name__ == '__main__':
mesh = UnitCubeMesh(5,5,5)
#mesh = UnitSquareMesh(15,15)
facet_domains = FacetFunction("size_t", mesh)
facet_domains.set_all(0)
CompiledSubDomain("x[0]<1e-12").mark(facet_domains, 1)
CompiledSubDomain("x[1]<1e-12").mark(facet_domains, 2)
V = VectorFunctionSpace(mesh, "CG", 2)
U = Function(V)
n = Normal(V)
noslip = NoSlipBC(V, n, "on_boundary")
#noslip.apply(u.vector())
#File("u.xdmf") << u
u,v = TrialFunction(V), TestFunction(V)
a = inner(u,v)*dx()
L = inner(Constant((1,)*V.num_sub_spaces()), v)*dx()
tic()
A = assemble(a)
print "Assemble a: ", toc()
tic()
b = assemble(L)
print "Assemble b: ", toc()
tic()
noslip.apply(A,b)
#U.vector()[:] = 1.0; noslip.apply(U.vector()); U.vector()[:] = noslip.R.T*U.vector()
print "Noslip apply: ", toc()
print A.norm("frobenius")
print norm(b)
solve(A, U.vector(), b)
U.vector()[:] = noslip.R.T*U.vector()
print norm(U)
print assemble(dot(U,n)*ds)
#File("u.xdmf") << U
plot(U)
interactive()
oyvinev created this gist
oyvinev
created
this gist