Skip to content

Commit

Permalink
Reduced Stokes variant
Browse files Browse the repository at this point in the history
  • Loading branch information
pbrubeck committed Nov 26, 2024
1 parent f95a42f commit 0962837
Show file tree
Hide file tree
Showing 3 changed files with 152 additions and 118 deletions.
3 changes: 3 additions & 0 deletions FIAT/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from FIAT.alfeld_sorokina import AlfeldSorokina
from FIAT.arnold_qin import ArnoldQin
from FIAT.guzman_neilan import GuzmanNeilanFirstKindH1, GuzmanNeilanSecondKindH1, GuzmanNeilanH1div
from FIAT.stokes import Stokes, MacroStokes
from FIAT.christiansen_hu import ChristiansenHu
from FIAT.johnson_mercier import JohnsonMercier
from FIAT.brezzi_douglas_marini import BrezziDouglasMarini
Expand Down Expand Up @@ -104,6 +105,8 @@
"Nedelec 2nd kind H(curl)": NedelecSecondKind,
"Raviart-Thomas": RaviartThomas,
"Regge": Regge,
"Stokes": Stokes,
"Macro Stokes": MacroStokes,
"EnrichedElement": EnrichedElement,
"NodalEnrichedElement": NodalEnrichedElement,
"QuadraticPowellSabin6": QuadraticPowellSabin6,
Expand Down
2 changes: 1 addition & 1 deletion FIAT/functional.py
Original file line number Diff line number Diff line change
Expand Up @@ -264,7 +264,7 @@ def __init__(self, ref_el, s1, s2, pt, comp=(), shp=(), nm=None):
alpha[i] += 1
alpha[j] += 1
alphas.append(tuple(alpha))
tau[cur] = s1[i] * s2[j] / (1 + (i == j))
tau[cur] = s1[i] * s2[j] + (i != j) * s2[i] * s1[j]
cur += 1

dpt_dict = {tuple(pt): [(tau[i], alphas[i], comp) for i in range(len(alphas))]}
Expand Down
265 changes: 148 additions & 117 deletions FIAT/stokes.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,13 @@ def eps(table_u):
d = len(grad_u)
indices = ((i, j) for i in range(d) for j in range(d))
eps_u = [0.5*(grad_u[j][:, i, :] + grad_u[i][:, j, :]) for k, (i, j) in enumerate(indices)]

num_members = table_u[(0,)*d].shape[0]
return numpy.transpose(eps_u, (1, 0, 2)).reshape(num_members, d, d, -1)


def inner(v, u, Qwts):
return numpy.tensordot(numpy.multiply(v, Qwts), u, axes=(range(1, v.ndim), range(1, u.ndim)))
return numpy.tensordot(numpy.multiply(v, Qwts), u,
axes=(range(1, v.ndim), range(1, u.ndim)))


def map_duals(ref_el, dim, entity, mapping, Q_ref, Phis):
Expand Down Expand Up @@ -60,80 +60,84 @@ def __init__(self, ref_el, degree):
top = ref_el.get_topology()
sd = ref_el.get_spatial_dimension()
shp = (sd,)
comps = list(numpy.ndindex(shp))

moment_degree = degree - 2*sd
edge_weight = sd-1
Q_edge, phis_edge = jacobi_duals(ref_el, edge_weight, degree-4, degree)

mapping = None
self._reduced_dofs = {}
for dim in sorted(top):
entity_ids[dim] = {}
if dim == 0:
for entity in sorted(top[dim]):
cur = len(nodes)
pts = ref_el.make_points(dim, entity, degree)
nodes.extend(ComponentPointEvaluation(ref_el, (k,), shp, pt)
for pt in pts for k in range(sd))
self._reduced_dofs[dim] = None

if dim == 1:
Q_ref, Phis = Q_edge, phis_edge[:moment_degree+1]
elif dim < sd:
Q_ref, Phis = self._reference_duals(ref_el, dim, degree, moment_degree)

for entity in sorted(top[dim]):
cur = len(nodes)
if dim == sd:
# Interior dofs
Q, phis, nullspace_dim = self._interior_duals(ref_el, degree)
self._reduced_dofs[dim] = nullspace_dim
nodes.extend(FrobeniusIntegralMoment(ref_el, Q, phi) for phi in phis)
entity_ids[dim][entity] = list(range(cur, len(nodes)))
else:
if dim == 1:
Q_ref, Phis = Q_edge, phis_edge[:moment_degree+1]
elif dim < sd:
Q_ref, Phis = self._reference_duals(ref_el, dim, degree, moment_degree)

for entity in sorted(top[dim]):
cur = len(nodes)
if dim == sd:
# Interior dofs
Q, phis = self._interior_duals(ref_el, degree)
nodes.extend(FrobeniusIntegralMoment(ref_el, Q, phi) for phi in phis)
entity_ids[dim][entity] = list(range(cur, len(nodes)))
continue

if dim == 1:
# Vertex-edge dofs
verts = ref_el.get_vertices_of_subcomplex(top[dim][entity])
ells = [PointTangentialDerivative]
if sd == 3:
ells.append(PointTangentialSecondDerivative)
nodes.extend(ell(ref_el, entity, pt, comp=(i,), shp=shp)
for pt in verts for ell in ells for i in range(sd))

elif dim == 2:
# Face-vertex dofs
verts = numpy.array(ref_el.get_vertices_of_subcomplex(top[dim][entity]))
for i in range(len(verts)):
tangents = [verts[j] - verts[i] for j in range(len(verts)) if j != i]
nodes.extend(PointSecondDerivative(ref_el, *tangents, verts[i], comp=(k,), shp=shp)
for k in range(sd))

# Face-edge dofs
face = set(top[dim][entity])
edges = [e for e in top[1] if set(top[1][e]) < face]
n = ref_el.compute_scaled_normal(entity)
for e in edges:
t = ref_el.compute_edge_tangent(e)
s = numpy.cross(n, t)
Q, phis = map_duals(ref_el, 1, e, mapping, Q_edge, phis_edge[:degree-5+1])
nodes.extend(IntegralMomentOfDerivative(ref_el, s, Q, phi, comp=(k,), shp=shp)
for phi in phis for k in range(sd))

# Rest of the facet moments
Q, phis = map_duals(ref_el, dim, entity, mapping, Q_ref, Phis)
nodes.extend(IntegralMoment(ref_el, Q, phi, comp=(k,), shp=shp)
for phi in phis for k in range(sd))
continue

elif dim == 0:
# Vertex dofs
pts = ref_el.make_points(dim, entity, degree)
nodes.extend(ComponentPointEvaluation(ref_el, comp, shp, pt)
for pt in pts for comp in comps)
entity_ids[dim][entity] = list(range(cur, len(nodes)))
continue

elif dim == 1:
# Vertex-edge dofs
verts = ref_el.get_vertices_of_subcomplex(top[dim][entity])
ells = [PointTangentialDerivative]
if sd == 3:
ells.append(PointTangentialSecondDerivative)
nodes.extend(ell(ref_el, entity, pt, comp=comp, shp=shp)
for pt in verts for ell in ells for comp in comps)

elif dim == 2:
# Face-vertex dofs
verts = numpy.array(ref_el.get_vertices_of_subcomplex(top[dim][entity]))
for i in range(len(verts)):
tangents = [verts[j] - verts[i] for j in range(len(verts)) if j != i]
nodes.extend(PointSecondDerivative(ref_el, *tangents, verts[i], comp=comp, shp=shp)
for comp in comps)

# Face-edge dofs
mid_face, = numpy.asarray(ref_el.make_points(dim, entity, dim+1))
edges = ref_el.connectivity[(dim, dim-1)][entity]
for e in edges:
mid_edge, = numpy.asarray(ref_el.make_points(dim-1, e, dim))
s = mid_face - mid_edge
Q, phis = map_duals(ref_el, dim-1, e, mapping, Q_edge, phis_edge[:degree-5+1])
nodes.extend(IntegralMomentOfDerivative(ref_el, s, Q, phi, comp=comp, shp=shp)
for phi in phis for comp in comps)

# Rest of the facet moments
Q, phis = map_duals(ref_el, dim, entity, mapping, Q_ref, Phis)
nodes.extend(IntegralMoment(ref_el, Q, phi, comp=comp, shp=shp)
for phi in phis for comp in comps)

entity_ids[dim][entity] = list(range(cur, len(nodes)))

super().__init__(nodes, ref_el, entity_ids)

def _reference_duals(self, ref_el, dim, degree, moment_degree):
"""Compute moments of trial space V against a subset of itself.
"""
sd = ref_el.get_spatial_dimension()
assert dim != sd
facet = ref_el.construct_subelement(dim)
if dim == 0:
return None, []

facet = ref_el.construct_subelement(dim)
V = ONPolynomialSet(facet, moment_degree, scale="orthonormal")
Q = create_quadrature(facet, degree + V.degree)
phis = V.tabulate(Q.get_points())[(0,)*dim]
Expand All @@ -146,7 +150,7 @@ def _interior_duals(self, ref_el, degree):
sd = ref_el.get_spatial_dimension()
shp = (sd,)

Q = create_quadrature(ref_el, degree * 2)
Q = create_quadrature(ref_el, 2*degree)
Qpts, Qwts = Q.get_points(), Q.get_weights()

# Test space: bubbles
Expand Down Expand Up @@ -177,13 +181,13 @@ def _interior_duals(self, ref_el, degree):
tol = sig[-1] * 1E-12
nullspace_dim = len([s for s in sig if abs(s) <= tol])

S2 = S[:, :nullspace_dim]
S1 = S[:, nullspace_dim:]
S1 *= numpy.sqrt(1 / sig[None, nullspace_dim:])
S1 = S[:, :nullspace_dim]
S2 = S[:, nullspace_dim:]
S2 *= numpy.sqrt(1 / sig[None, nullspace_dim:])

# Apply change of basis
div_test = numpy.tensordot(S1, div_test, axes=(0, 0))
eps_test = numpy.tensordot(S2, eps_test, axes=(0, 0))
eps_test = numpy.tensordot(S1, eps_test, axes=(0, 0))
div_test = numpy.tensordot(S2, div_test, axes=(0, 0))

# Trial space
V = ONPolynomialSet(ref_el, degree, shp, scale="orthonormal")
Expand All @@ -192,10 +196,31 @@ def _interior_duals(self, ref_el, degree):
eps_trial = eps(V_at_qpts)
div_trial = numpy.trace(eps_trial, axis1=1, axis2=2)

K = numpy.concatenate((inner(div_test, div_trial, Qwts),
inner(eps_test, eps_trial, Qwts)), axis=0)
K = numpy.concatenate((inner(eps_test, eps_trial, Qwts),
inner(div_test, div_trial, Qwts),
), axis=0)
phis = numpy.tensordot(K, trial, axes=(1, 0))
return Q, phis
return Q, phis, nullspace_dim

def get_indices(self, restriction_domain, take_closure=True):
"""Return the list of dofs with support on the given restriction domain.
Allows for reduced Demkowicz elements, excluding the exterior
derivative of the previous space in the de Rham complex.
:arg restriction_domain: can be 'reduced', 'interior', 'vertex',
'edge', 'face' or 'facet'
:kwarg take_closure: Are we taking the closure of the restriction domain?
"""
if restriction_domain == "reduced":
indices = []
entity_ids = self.get_entity_ids()
for dim in entity_ids:
reduced_dofs = self._reduced_dofs[dim]
for entity, ids in entity_ids[dim].items():
indices.extend(ids[:reduced_dofs])
return indices
else:
return dual_set.DualSet.get_indices(self, restriction_domain, take_closure=take_closure)


class Stokes(finite_element.CiarletElement):
Expand All @@ -210,7 +235,7 @@ def __init__(self, ref_el, degree):
poly_set = ONPolynomialSet(ref_el, degree, shape=(sd,), variant="bubble")
dual = StokesDual(ref_el, degree)
formdegree = sd-1 # (n-1)-form
super().__init__(poly_set, dual, degree, formdegree)
super().__init__(poly_set, dual, degree, formdegree, mapping="contravariant piola")


class MacroStokesDual(StokesDual):
Expand All @@ -222,34 +247,35 @@ def __init__(self, ref_complex, degree):
top = ref_el.get_topology()
sd = ref_el.get_spatial_dimension()
shp = (sd,)
comps = list(numpy.ndindex(shp))

mapping = None
self._reduced_dofs = {}
for dim in sorted(top):
entity_ids[dim] = {}
if dim == 0:
for entity in sorted(top[dim]):
cur = len(nodes)
if dim > 0 and dim < sd:
moment_degree = degree - (dim+1)
Q_ref, Phis = self._reference_duals(ref_complex, dim, degree, moment_degree)

self._reduced_dofs[dim] = None
for entity in sorted(top[dim]):
cur = len(nodes)
if dim == 0:
# Vertex dofs
pts = ref_el.make_points(dim, entity, degree)
nodes.extend(ComponentPointEvaluation(ref_el, (k,), shp, pt)
for pt in pts for k in range(sd))
entity_ids[dim][entity] = list(range(cur, len(nodes)))
else:
if dim < sd:
moment_degree = degree - (1+dim)
Q_ref, Phis = self._reference_duals(ref_complex, dim, degree, moment_degree)

for entity in sorted(top[dim]):
cur = len(nodes)
if dim == sd:
# Interior dofs
Q, phis = self._interior_duals(ref_complex, degree)
nodes.extend(FrobeniusIntegralMoment(ref_el, Q, phi) for phi in phis)
else:
# Facet dofs
Q, phis = map_duals(ref_el, dim, entity, mapping, Q_ref, Phis)
nodes.extend(IntegralMoment(ref_el, Q, phi, comp=(k,), shp=shp)
for phi in phis for k in range(sd))
entity_ids[dim][entity] = list(range(cur, len(nodes)))
nodes.extend(ComponentPointEvaluation(ref_el, comp, shp, pt)
for pt in pts for comp in comps)
elif dim == sd:
# Interior dofs
Q, phis, nullspace_dim = self._interior_duals(ref_complex, degree)
nodes.extend(FrobeniusIntegralMoment(ref_el, Q, phi) for phi in phis)
self._reduced_dofs[dim] = nullspace_dim
else:
# Facet dofs
Q, phis = map_duals(ref_el, dim, entity, mapping, Q_ref, Phis)
nodes.extend(IntegralMoment(ref_el, Q, phi, comp=comp, shp=shp)
for phi in phis for comp in comps)
entity_ids[dim][entity] = list(range(cur, len(nodes)))
super(StokesDual, self).__init__(nodes, ref_el, entity_ids)


Expand All @@ -266,42 +292,47 @@ def __init__(self, ref_el, degree):
poly_set = ONPolynomialSet(ref_complex, degree, shape=(sd,), variant="bubble")
dual = MacroStokesDual(ref_complex, degree)
formdegree = sd-1 # (n-1)-form
super().__init__(poly_set, dual, degree, formdegree)
super().__init__(poly_set, dual, degree, formdegree, mapping="contravariant piola")


if __name__ == "__main__":
import matplotlib.pyplot as plt
import FIAT

dim = 3
degree = 3
dim = 2
degree = 8
ref_el = FIAT.reference_element.symmetric_simplex(dim)
# fe = Stokes(ref_el, degree)
fe = MacroStokes(ref_el, degree)

Q = create_quadrature(fe.ref_complex, 2 * degree)
Qpts, Qwts = Q.get_points(), Q.get_weights()
phi_at_qpts = fe.tabulate(1, Qpts)

Veps = eps(phi_at_qpts)
Vdiv = numpy.trace(Veps, axis1=1, axis2=2)
Aeps = inner(Veps, Veps, Qwts)
Adiv = inner(Vdiv, Vdiv, Qwts)

fig, axes = plt.subplots(ncols=2, nrows=1, figsize=(12, 6))
axes = axes.flat
title = f"{type(fe).__name__}({degree})"
names = ("eps", "div")
mats = (Aeps, Adiv)
for name, A in zip(names, mats):
A[abs(A) < 1E-10] = 0.0
scipy_mat = scipy.sparse.csr_matrix(A)
nnz = scipy_mat.count_nonzero()
ms = 0
ax = next(axes)
ax.spy(A, markersize=ms)
if ms == 0:
ax.pcolor(numpy.log(abs(A)))
ax.set_title(f"{title} {name} nnz {nnz}")

family = type(fe).__name__
domains = (None, "reduced", "facet")
fig, axes = plt.subplots(nrows=2, ncols=len(domains), figsize=(6*len(domains), 6*2))
axes = axes.T.flat
for domain in domains:
if domain:
fe = FIAT.RestrictedElement(fe, restriction_domain=domain)

phi_at_qpts = fe.tabulate(1, Qpts)
Veps = eps(phi_at_qpts)
Vdiv = numpy.trace(Veps, axis1=1, axis2=2)
Aeps = inner(Veps, Veps, Qwts)
Adiv = inner(Vdiv, Vdiv, Qwts)

title = f"{family}({degree}, {domain})"
names = ("eps", "div")
mats = (Aeps, Adiv)
for name, A in zip(names, mats):
A[abs(A) < 1E-10] = 0.0
scipy_mat = scipy.sparse.csr_matrix(A)
nnz = scipy_mat.count_nonzero()
ms = 0
ax = next(axes)
ax.spy(A, markersize=ms)
if ms == 0:
ax.pcolor(numpy.log(abs(A)))
ax.set_title(f"{title} {name} nnz {nnz}")

plt.show()

0 comments on commit 0962837

Please sign in to comment.