def get_aspect_ratios2d(mesh, python=False):
"""
Computes the aspect ratio of each cell in a 2D triangular mesh
:arg mesh: the input mesh to do computations on
:kwarg python: compute the measure using Python?
:rtype: firedrake.function.Function aspect_ratios with
aspect ratio data
"""
P0 = firedrake.FunctionSpace(mesh, "DG", 0)
if python:
P0_ten = firedrake.TensorFunctionSpace(mesh, "DG", 0)
J = firedrake.interpolate(ufl.Jacobian(mesh), P0_ten)
edge1 = ufl.as_vector([J[0, 0], J[1, 0]])
edge2 = ufl.as_vector([J[0, 1], J[1, 1]])
edge3 = edge1 - edge2
a = ufl.sqrt(ufl.dot(edge1, edge1))
b = ufl.sqrt(ufl.dot(edge2, edge2))
c = ufl.sqrt(ufl.dot(edge3, edge3))
aspect_ratios = firedrake.interpolate(
a * b * c / ((a + b - c) * (b + c - a) * (c + a - b)), P0
)
else:
coords = mesh.coordinates
aspect_ratios = firedrake.Function(P0)
op2.par_loop(
get_pyop2_kernel("get_aspect_ratio", 2),
mesh.cell_set,
aspect_ratios.dat(op2.WRITE, aspect_ratios.cell_node_map()),
coords.dat(op2.READ, coords.cell_node_map()),
)
return aspect_ratios
def test_de_rahm_2D(order):
mesh = create_unit_square(MPI.COMM_WORLD, 3, 4)
W = FunctionSpace(mesh, ("Lagrange", order))
w = Function(W)
w.interpolate(lambda x: x[0] + x[0] * x[1] + 2 * x[1]**2)
g = ufl.grad(w)
Q = FunctionSpace(mesh, ("N2curl", order - 1))
q = Function(Q)
q.interpolate(Expression(g, Q.element.interpolation_points))
x = ufl.SpatialCoordinate(mesh)
g_ex = ufl.as_vector((1 + x[1], 4 * x[1] + x[0]))
assert np.isclose(
np.abs(assemble_scalar(form(ufl.inner(q - g_ex, q - g_ex) * ufl.dx))),
0)
V = FunctionSpace(mesh, ("BDM", order - 1))
v = Function(V)
def curl2D(u):
return ufl.as_vector((ufl.Dx(u[1], 0), -ufl.Dx(u[0], 1)))
v.interpolate(
Expression(curl2D(ufl.grad(w)), V.element.interpolation_points))
h_ex = ufl.as_vector((1, -1))
assert np.isclose(
np.abs(assemble_scalar(form(ufl.inner(v - h_ex, v - h_ex) * ufl.dx))),
0)
def grad(self, o, expr):
shape = o.ufl_shape
if len(shape) == 1:
return as_vector([expr] * shape[0])
elif len(shape) == 2:
return as_vector([[expr[i]] * shape[1] for i in range(shape[0])])
else:
raise NotImplementedError()
def UFLFunction(grid, name, order, expr, renumbering=None, virtualize=True, tempVars=True,
predefined=None, **kwargs):
scalar = False
if type(expr) == list or type(expr) == tuple:
expr = ufl.as_vector(expr)
elif type(expr) == int or type(expr) == float:
expr = ufl.as_vector( [expr] )
scalar = True
try:
if expr.ufl_shape == ():
expr = ufl.as_vector([expr])
scalar = True
except:
return None
_, coeff_ = ufl.algorithms.analysis.extract_arguments_and_coefficients(expr)
coeff = {c : c.toVectorCoefficient()[0] for c in coeff_ if len(c.ufl_shape) == 0 and not c.is_cellwise_constant()}
expr = replace(expr,coeff)
if len(expr.ufl_shape) > 1:
raise AttributeError("can only generate grid functions from vector values UFL expressions not from expressions with shape=",expr.ufl_shape)
# set up the source class
source = UFLFunctionSource(grid, expr,
name,order,
tempVars=tempVars,virtualize=virtualize,
predefined=predefined)
coefficients = source.coefficientList
numCoefficients = len(coefficients)
if renumbering is None:
renumbering = dict()
renumbering.update((c, i) for i, c in enumerate(sorted((c for c in coefficients if not c.is_cellwise_constant()), key=lambda c: c.count())))
renumbering.update((c, i) for i, c in enumerate(c for c in coefficients if c.is_cellwise_constant()))
coefficientNames = ['coefficient' + str(i) if n is None else n for i, n in enumerate(getattr(c, 'name', None) for c in coefficients if not c.is_cellwise_constant())]
# call code generator
from dune.generator import builder
module = builder.load(source.name(), source, "UFLLocalFunction")
assert hasattr(module,"UFLLocalFunction"),\
"GridViews of coefficients need to be compatible with the grid view of the ufl local functions"
class LocalFunction(module.UFLLocalFunction):
def __init__(self, gridView, name, order, *args, **kwargs):
self.base = module.UFLLocalFunction
self._coefficientNames = {n: i for i, n in enumerate(source.coefficientNames)}
if renumbering is not None:
self._renumbering = renumbering
self._setConstant = self.setConstant # module.UFLLocalFunction.__dict__['setConstant']
self.setConstant = lambda *args: setConstant(self,*args)
self.constantShape = source._constantShapes
self._constants = [c for c in source.constantList if isinstance(c,Constant)]
self.scalar = scalar
init(self, gridView, name, order, *args, **kwargs)
return LocalFunction
def argument(self, o):
shape = o.ufl_shape
if len(shape) == 0:
return 1
elif len(shape) == 1:
return as_vector([1] * shape[0])
elif len(shape) == 2:
return as_vector([as_vector([1] * shape[1]) for _ in range(shape[0])])
else:
raise NotImplementedError()
def list_tensor(self, o):
shape = o.ufl_shape
if len(shape) == 1:
return as_vector([self.visit(op) for op in o.ufl_operands])
elif len(shape) == 2:
return as_vector(
[as_vector([self.visit(op) for op in row.ufl_operands]) for row in o.ufl_operands]
)
else:
raise NotImplementedError()
def codeDG(self):
code = self._code()
u = self.trialFunction
ubar = Coefficient(u.ufl_function_space())
penalty = self.penalty
if penalty is None:
penalty = 1
if isinstance(penalty, Expr):
if penalty.ufl_shape == ():
penalty = as_vector([penalty])
try:
penalty = expand_indices(
expand_derivatives(expand_compounds(penalty)))
except:
pass
assert penalty.ufl_shape == u.ufl_shape
dmPenalty = as_vector([
replace(
expand_derivatives(diff(replace(penalty, {u: ubar}),
ubar))[i, i], {ubar: u})
for i in range(u.ufl_shape[0])
])
else:
dmPenalty = None
code.append(AccessModifier("public"))
x = SpatialCoordinate(self.space.cell())
predefined = {}
self.predefineCoefficients(predefined, x)
spatial = Variable('const auto', 'y')
predefined.update({
x:
UnformattedExpression(
'auto', 'entity().geometry().global( Dune::Fem::coordinate( x ) )')
})
generateMethod(code,
penalty,
'RRangeType',
'penalty',
args=['const Point &x', 'const DRangeType &u'],
targs=['class Point', 'class DRangeType'],
static=False,
const=True,
predefined=predefined)
generateMethod(code,
dmPenalty,
'RRangeType',
'linPenalty',
args=['const Point &x', 'const DRangeType &u'],
targs=['class Point', 'class DRangeType'],
static=False,
const=True,
predefined=predefined)
return code
def __init__(self, functionSpace, value, subDomain=None):
if functionSpace.scalar:
self.functionSpace = functionSpace.toVectorSpace()
value = ufl.as_vector([value])
else:
self.functionSpace = functionSpace
self.value = value
self.subDomain = subDomain
if type(value) is list:
self.ufl_value = value # DirichletBC.flatten(value)
self.ufl_value = [0 if v is None else v for v in self.ufl_value]
self.ufl_value = ufl.as_vector(self.ufl_value)
else:
self.ufl_value = value
assert self.ufl_value.ufl_shape[0] == functionSpace.dimRange
def mark(indicator,
refineTolerance,
coarsenTolerance=0,
minLevel=0,
maxLevel=None,
gridView=None):
if ufl and (isinstance(indicator, list) or isinstance(indicator, tuple)):
indicator = ufl.as_vector(indicator)
if ufl and isinstance(indicator, ufl.core.expr.Expr): #
if gridView is None:
_, coeff_ = extract_arguments_and_coefficients(indicator)
gridView = [c.grid for c in coeff_ if hasattr(c, "grid")]
gridView += [c.gridView for c in coeff_ if hasattr(c, "gridView")]
if len(gridView) == 0:
raise ValueError(
"if a ufl expression is passed as indicator then the 'gridView' must also be provided"
)
gridView = gridView[0]
indicator = expression2GF(gridView, indicator, 0)
if gridView is None:
gridView = indicator.gridView
try:
if not gridView.canAdapt:
raise AttributeError(
"indicator function must be over grid view that supports adaptation"
)
except AttributeError:
raise AttributeError(
"indicator function must be over grid view that supports adaptation"
)
if maxLevel == None:
maxLevel = -1
return gridView.mark(indicator, refineTolerance, coarsenTolerance,
minLevel, maxLevel)
def argument(self, o):
from ufl import as_vector
from ufl.constantvalue import Zero
from dolfin import Argument
from numpy import ndindex
V = o.function_space()
if V.num_sub_spaces() == 0:
# Not on a mixed space, just return ourselves.
return o
# Already seen this argument, return the cached version.
if o in self._args:
return self._args[o]
args = []
for i, V_i_sub in enumerate(V.split()):
# Walk over the subspaces and build a vector that is zero
# where the argument does not match the one we're looking
# for and is just the non-mixed argument when we do want
# it.
V_i = V_i_sub.collapse()
a = Argument(V_i, o.number(), part=o.part())
indices = ndindex(a.ufl_shape)
if self.idx[o.number()] == i:
args += [a[j] for j in indices]
else:
args += [Zero() for j in indices]
self._args[o] = as_vector(args)
return self._args[o]
def argument(self, arg):
"""Split an argument into its constituent spaces."""
from itertools import product
if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
# Look up the split argument in cache since we want it unique
if arg in self._args:
return self._args[arg]
args = []
for i, fs in enumerate(arg.function_space().split()):
# Build the sub-space Argument (not part of the mixed
# space).
a = Argument(fs, arg.number(), part=arg.part())
# Produce indexing iterator.
# For scalar-valued spaces this results in the empty
# tuple (), which returns the Argument when indexing.
# For vector-valued spaces, this is just
# range(a.ufl_shape[0])
# For tensor-valued spaces, it's the nested loop of
# range(x for x in a.ufl_shape).
#
# Each of the indexed things is then scalar-valued
# which we turn into a ufl Vector.
indices = product(*map(range, a.ufl_shape))
if self._idx[arg.number()] == i:
args += [a[idx] for idx in indices]
else:
args += [Zero() for _ in indices]
self._args[arg] = as_vector(args)
return self._args[arg]
return arg
def __init__(self, grid, expr,
name,order,
tempVars=True, virtualize=True,
predefined=None):
gridType = grid.cppTypeName
gridIncludes = grid.cppIncludes
if len(expr.ufl_shape) == 0:
expr = as_vector( [ expr ] )
dimRange = expr.ufl_shape[0]
predefined = {} if predefined is None else predefined
codegen.ModelClass.__init__(self,"UFLLocalFunction",[expr],
virtualize,dimRange=dimRange, predefined=predefined)
self.evalCode = []
self.jacCode = []
self.hessCode = []
self.gridType = gridType
self.gridIncludes = gridIncludes
self.functionName = name
self.functionOrder = order
self.expr = expr
self.tempVars = tempVars
self.virtualize = virtualize
self.codeString = self.toString()
def argument(self, obj):
Q = obj.ufl_function_space()
dom = Q.ufl_domain()
sub_elements = obj.ufl_element().sub_elements()
# If not a mixed element, do nothing
if not isinstance(obj.ufl_element(), MixedElement):
return obj
# Split into sub-elements, creating appropriate space for each
args = []
for i, sub_elem in enumerate(sub_elements):
Q_i = FunctionSpace(dom, sub_elem)
a = Argument(Q_i, obj.number(), part=obj.part())
indices = [()]
for m in a.ufl_shape:
indices = [(k + (j, )) for k in indices for j in range(m)]
if (i == self.idx[obj.number()]):
args += [a[j] for j in indices]
else:
args += [Zero() for j in indices]
return as_vector(args)
def amg_solve(N, method):
# Elasticity parameters
E = 1.0e9
nu = 0.3
mu = E / (2.0 * (1.0 + nu))
lmbda = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))
# Stress computation
def sigma(v):
return 2.0 * mu * sym(grad(v)) + lmbda * tr(sym(
grad(v))) * Identity(2)
# Define problem
mesh = UnitSquareMesh(MPI.comm_world, N, N)
V = VectorFunctionSpace(mesh, 'Lagrange', 1)
bc0 = Function(V)
with bc0.vector().localForm() as bc_local:
bc_local.set(0.0)
def boundary(x, only_boundary):
return [only_boundary] * x.shape(0)
bc = DirichletBC(V.sub(0), bc0, boundary)
u = TrialFunction(V)
v = TestFunction(V)
# Forms
a, L = inner(sigma(u), grad(v)) * dx, dot(ufl.as_vector(
(1.0, 1.0)), v) * dx
# Assemble linear algebra objects
A = assemble_matrix(a, [bc])
A.assemble()
b = assemble_vector(L)
apply_lifting(b, [a], [[bc]])
b.ghostUpdate(addv=PETSc.InsertMode.ADD,
mode=PETSc.ScatterMode.REVERSE)
set_bc(b, [bc])
# Create solution function
u = Function(V)
# Create near null space basis and orthonormalize
null_space = build_nullspace(V, u.vector())
# Attached near-null space to matrix
A.set_near_nullspace(null_space)
# Test that basis is orthonormal
assert null_space.is_orthonormal()
# Create PETSC smoothed aggregation AMG preconditioner, and
# create CG solver
solver = PETScKrylovSolver("cg", method)
# Set matrix operator
solver.set_operator(A)
# Compute solution and return number of iterations
return solver.solve(u.vector(), b)
def argument(self, arg):
"""Split an argument into its constituent spaces."""
from itertools import product
if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
# Look up the split argument in cache since we want it unique
if arg in self._args:
return self._args[arg]
args = []
for i, fs in enumerate(arg.function_space().split()):
# Build the sub-space Argument (not part of the mixed
# space).
a = Argument(fs, arg.number(), part=arg.part())
# Produce indexing iterator.
# For scalar-valued spaces this results in the empty
# tuple (), which returns the Argument when indexing.
# For vector-valued spaces, this is just
# range(a.ufl_shape[0])
# For tensor-valued spaces, it's the nested loop of
# range(x for x in a.ufl_shape).
#
# Each of the indexed things is then scalar-valued
# which we turn into a ufl Vector.
indices = product(*map(range, a.ufl_shape))
if self._idx[arg.number()] == i:
args += [a[idx] for idx in indices]
else:
args += [Zero() for _ in indices]
self._args[arg] = as_vector(args)
return self._args[arg]
return arg
def argument(self, arg):
"""
Argument is UFL lingo for test (num=0) and trial (num=1) functions
"""
# Do not make several new args for the same original arg
if arg in self._cache:
return self._cache[arg]
# Get some data about the argument and get our target index
V = arg.function_space()
N = V.num_sub_spaces()
num = arg.number()
idx_wanted = [self._index_v, self._index_u][num]
new_arg = []
for idx in range(N):
# Construct non-coupled argument
Vi = V.sub(idx).collapse()
a = dolfin.function.argument.Argument(Vi, num, part=arg.part())
indices = numpy.ndindex(a.ufl_shape)
if idx == idx_wanted:
# This is a wanted index
new_arg.extend(a[I] for I in indices)
else:
# This index should be pruned
new_arg.extend(Zero() for _ in indices)
new_arg = as_vector(new_arg)
self._cache[arg] = new_arg
return new_arg
def amg_solve(N, method):
# Elasticity parameters
E = 1.0e9
nu = 0.3
mu = E / (2.0 * (1.0 + nu))
lmbda = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))
# Stress computation
def sigma(v):
return 2.0 * mu * sym(grad(v)) + lmbda * tr(sym(
grad(v))) * Identity(2)
# Define problem
mesh = create_unit_square(MPI.COMM_WORLD, N, N)
V = VectorFunctionSpace(mesh, 'Lagrange', 1)
u = TrialFunction(V)
v = TestFunction(V)
facetdim = mesh.topology.dim - 1
bndry_facets = locate_entities_boundary(
mesh, facetdim, lambda x: np.full(x.shape[1], True))
bdofs = locate_dofs_topological(V.sub(0), V, facetdim, bndry_facets)
bc = dirichletbc(PETSc.ScalarType(0), bdofs, V.sub(0))
# Forms
a, L = inner(sigma(u), grad(v)) * dx, dot(ufl.as_vector(
(1.0, 1.0)), v) * dx
# Assemble linear algebra objects
A = assemble_matrix(a, [bc])
A.assemble()
b = assemble_vector(L)
apply_lifting(b, [a], [[bc]])
b.ghostUpdate(addv=PETSc.InsertMode.ADD,
mode=PETSc.ScatterMode.REVERSE)
set_bc(b, [bc])
# Create solution function
u = Function(V)
# Create near null space basis and orthonormalize
null_space = build_nullspace(V, u.vector)
# Attached near-null space to matrix
A.set_near_nullspace(null_space)
# Test that basis is orthonormal
assert null_space.is_orthonormal()
# Create PETSC smoothed aggregation AMG preconditioner, and
# create CG solver
solver = PETSc.KSP().create(mesh.comm)
solver.setType("cg")
# Set matrix operator
solver.setOperators(A)
# Compute solution and return number of iterations
return solver.solve(b, u.vector)
def rebuild_expression_from_graph(G):
"This is currently only used by tests."
w = rebuild_with_scalar_subexpressions(G)
# Find expressions of final v
if len(w) == 1:
return w[0]
else:
return as_vector(w) # TODO: Consider shape of initial v
def split(self, fields):
from ufl import as_vector, replace
from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
try:
if len(field) > 1:
raise NotImplementedError("Can't split into subblock")
except TypeError:
# Just a single field, we can handle that
pass
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
subu = us[field]
vec = []
for i, u in enumerate(us):
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp,
argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
if bc.function_space().index == field:
V = FunctionSpace(subu.ufl_domain(), subu.ufl_element())
bcs.append(
type(bc)(V,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(
F,
subu,
bcs=bcs,
J=J,
Jp=None,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(
type(self)(new_problem,
mat_type=self.mat_type,
pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def test_2D_lagrange_to_curl(order):
mesh = create_unit_square(MPI.COMM_WORLD, 3, 4)
V = FunctionSpace(mesh, ("N1curl", order))
u = Function(V)
W = FunctionSpace(mesh, ("Lagrange", order))
u0 = Function(W)
u0.interpolate(lambda x: -x[1])
u1 = Function(W)
u1.interpolate(lambda x: x[0])
f = ufl.as_vector((u0, u1))
f_expr = Expression(f, V.element.interpolation_points)
u.interpolate(f_expr)
x = ufl.SpatialCoordinate(mesh)
f_ex = ufl.as_vector((-x[1], x[0]))
assert np.isclose(
np.abs(assemble_scalar(form(ufl.inner(u - f_ex, u - f_ex) * ufl.dx))),
0)
def test_ufl_only_to_contravariant_piola():
mesh = ufl.Mesh(ufl.VectorElement("P", ufl.triangle, 1))
V = ufl.FunctionSpace(mesh, ufl.FiniteElement("P", ufl.triangle, 2))
v = ufl.Coefficient(V)
expr = ufl.as_vector([v, v])
W = ufl.FunctionSpace(mesh, ufl.FiniteElement("RT", ufl.triangle, 1))
to_element = create_element(W.ufl_element())
ast, oriented, needs_cell_sizes, coefficients, first_coeff_fake_coords, *_ = compile_expression_dual_evaluation(
expr, to_element, coffee=False)
assert first_coeff_fake_coords is True
def test_rank1(compile_args):
"""Tests evaluation of rank-1 form.
Builds a linear operator which takes vector-valued functions in P1 space
and evaluates expression [u_y, u_x] + grad(u_x) at specified points.
"""
e = ufl.VectorElement("P", "triangle", 1)
mesh = ufl.Mesh(e)
V = ufl.FunctionSpace(mesh, e)
u = ufl.TrialFunction(V)
expr = ufl.as_vector([u[1], u[0]]) + ufl.grad(u[0])
points = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 1.0]])
obj, module = ffcx.codegeneration.jit.compile_expressions(
[(expr, points)], cffi_extra_compile_args=compile_args)
ffi = cffi.FFI()
kernel = obj[0][0]
c_type, np_type = float_to_type("double")
# 2 components for vector components of expression
# 3 points of evaluation
# 6 degrees of freedom for rank1 form
A = np.zeros((3, 2, 6), dtype=np_type)
# Coefficient storage XYXYXY
w = np.array([0.0], dtype=np_type)
c = np.array([0.0], dtype=np_type)
# Coords storage XYXYXY
coords = np.array(points.flatten(), dtype=np.float64)
kernel.tabulate_expression(
ffi.cast('{type} *'.format(type=c_type), A.ctypes.data),
ffi.cast('{type} *'.format(type=c_type), w.ctypes.data),
ffi.cast('{type} *'.format(type=c_type), c.ctypes.data),
ffi.cast('double *', coords.ctypes.data))
f = np.array([[1.0, 2.0, 3.0], [-4.0, -5.0, 6.0]])
# Apply the operator on some test input data
u_ffcx = np.einsum("ijk,k", A, f.T.flatten())
# Compute the correct values using NumPy
# Gradf0 is gradient of f[0], each component of the gradient is constant
gradf0 = np.array(
[[f[0, 1] - f[0, 0], f[0, 1] - f[0, 0], f[0, 1] - f[0, 0]],
[f[0, 2] - f[0, 0], f[0, 2] - f[0, 0], f[0, 2] - f[0, 0]]])
u_correct = np.array([f[1], f[0]]) + gradf0
assert np.allclose(u_ffcx, u_correct.T)
def _uflToExpr(grid, order, f):
if not ufl: return f
if isinstance(f, list) or isinstance(f, tuple):
if isinstance(f[0], ufl.core.expr.Expr):
f = ufl.as_vector(f)
if isinstance(f, GridFunction):
return f
if isinstance(f, ufl.core.expr.Expr):
return expression2GF(grid, f, order).as_ufl()
else:
return f
def a(phi):
# phi = 1e5*phi
n = ufl.grad(phi) / ufl.sqrt(ufl.dot(ufl.grad(phi), ufl.grad(phi)) + 1e-5)
theta = ufl.atan_2(n[1], n[0])
n = 1e5 * n + ufl.as_vector([1e-5, 1e-5])
xx = n[1] / n[0]
# theta = ufl.asin(xx/ ufl.sqrt(1 + xx**2))
theta = ufl.atan(xx)
return 1.0 + epsilon_4 * ufl.cos(m * (theta - theta_0))
def test_curl(space_type, order):
"""Test that curl is consistent for different cell permutations of a tetrahedron."""
tdim = cpp.mesh.cell_dim(CellType.tetrahedron)
points = unit_cell_points(CellType.tetrahedron)
spaces = []
results = []
cell = list(range(len(points)))
# Assemble vector on 5 randomly numbered cells
for i in range(5):
shuffle(cell)
domain = ufl.Mesh(
ufl.VectorElement("Lagrange",
cpp.mesh.to_string(CellType.tetrahedron), 1))
mesh = create_mesh(MPI.COMM_WORLD, [cell], points, domain)
mesh.topology.create_connectivity_all()
V = FunctionSpace(mesh, (space_type, order))
v = ufl.TestFunction(V)
f = ufl.as_vector(tuple(1 if i == 0 else 0 for i in range(tdim)))
form = ufl.inner(f, ufl.curl(v)) * ufl.dx
result = fem.assemble_vector(form)
spaces.append(V)
results.append(result.array)
# Check that all DOFs on edges agree
V = spaces[0]
result = results[0]
connectivity = V.mesh.topology.connectivity(1, 0)
for i, edge in enumerate(V.mesh.topology.connectivity(tdim, 1).links(0)):
vertices = connectivity.links(edge)
values = sorted([
result[V.dofmap.cell_dofs(0)[a]]
for a in V.dofmap.dof_layout.entity_dofs(1, i)
])
for s, r in zip(spaces[1:], results[1:]):
c = s.mesh.topology.connectivity(1, 0)
for j, e in enumerate(
s.mesh.topology.connectivity(tdim, 1).links(0)):
if sorted(c.links(e)) == sorted(vertices):
v = sorted([
r[s.dofmap.cell_dofs(0)[a]]
for a in s.dofmap.dof_layout.entity_dofs(1, j)
])
assert np.allclose(values, v)
break
else:
continue
break
def generateMethodBody(cppType, expr, returnResult, default, predefined):
if expr is not None and not expr == [None]:
try:
dimR = expr.ufl_shape[0]
except:
if isinstance(expr,list) or isinstance(expr,tuple):
expr = as_vector(expr)
else:
expr = as_vector([expr])
dimR = expr.ufl_shape[0]
_, coeff = extract_arguments_and_coefficients(expr)
coeff = {c : c.toVectorCoefficient()[0] for c in coeff if len(c.ufl_shape) == 0 and not c.is_cellwise_constant()}
expr = replace(expr, coeff)
t = ExprTensor(expr.ufl_shape) # , exprTensorApply(lambda u: u, expr.ufl_shape, expr))
expression = [expr[i] for i in t.keys()]
u = extract_arguments_and_coefficients(expr)[0]
if u != []:
u = u[0]
du = Grad(u)
d2u = Grad(du)
arg_u = Variable("const RangeType &", "u")
arg_du = Variable("const JacobianRangeType &", "du")
arg_d2u = Variable("const HessianRangeType &", "d2u")
predefined.update( {u: arg_u, du: arg_du, d2u: arg_d2u} )
code, results = generateCode(predefined, expression, tempVars=False)
result = Variable(cppType, 'result')
if cppType == 'double':
code = code + [assign(result, results[0])]
else:
code = code + [assign(result[i], r) for i, r in zip(t.keys(), results)]
if returnResult:
code = [Declaration(result)] + code + [return_(result)]
else:
result = Variable(cppType, 'result')
code = [assign(result, construct(cppType,default) )]
if returnResult:
code = [Declaration(result)] + code + [return_(result)]
return code
def dfInterpolate(self, f):
if isinstance(f, list) or isinstance(f, tuple):
if isinstance(f[0], ufl.core.expr.Expr):
f = ufl.as_vector(f)
dimExpr = 0
if isinstance(f, GridFunction):
func = f.gf
dimExpr = func.dimRange
elif isinstance(f, ufl.core.expr.Expr):
func = expression2GF(self.space.gridView, f, self.space.order).as_ufl()
if func.ufl_shape == ():
dimExpr = 1
else:
dimExpr = func.ufl_shape[0]
else:
try:
gl = len(inspect.getargspec(f)[0])
func = None
except TypeError:
func = f
if isinstance(func, int) or isinstance(func, float):
dimExpr = 1
else:
dimExpr = len(func)
if func is None:
if gl == 1: # global function
func = function.globalFunction(self.space.gridView, "tmp",
self.space.order, f).gf
elif gl == 2: # local function
func = function.localFunction(self.space.gridView, "tmp",
self.space.order, f).gf
elif gl == 3: # local function with self argument (i.e. from @gridFunction)
func = function.localFunction(self.space.gridView, "tmp",
self.space.order,
lambda en, x: f(en, x)).gf
dimExpr = func.dimRange
if dimExpr == 0:
raise AttributeError("can not determine if expression shape"\
" fits the space's range dimension")
elif dimExpr != self.space.dimRange:
raise AttributeError("trying to interpolate an expression"\
" of size "+str(dimExpr)+" into a space with range dimension = "\
+ str(self.space.dimRange))
try:
# API GridView: assert func.gridView == self.gridView, "can only interpolate with same grid views"
assert func.gridView == self.gridView, "can only interpolate with same grid views"
assert func.dimRange == self.dimRange, "range dimension mismatch"
except AttributeError:
pass
return self._interpolate(func)
def physical_edge_lengths(self):
expr = ufl.classes.CellEdgeVectors(self.mt.terminal.ufl_domain())
if self.mt.restriction == '+':
expr = PositiveRestricted(expr)
elif self.mt.restriction == '-':
expr = NegativeRestricted(expr)
expr = ufl.as_vector([ufl.sqrt(ufl.dot(expr[i, :], expr[i, :])) for i in range(3)])
expr = preprocess_expression(expr)
config = {"point_set": PointSingleton([1/3, 1/3])}
config.update(self.config)
context = PointSetContext(**config)
return map_expr_dag(context.translator, expr)
def rotate(v, angle, origin=None):
"""
Rotate a UFL :class:`as_vector` ``v``
by ``angle`` about an ``origin``.
"""
dim = len(v)
origin = origin or ufl.as_vector(np.zeros(dim))
assert len(origin) == dim, "Origin does not match dimension"
if dim == 2:
R = rotation_matrix_2d(angle)
else:
raise NotImplementedError
return ufl.dot(R, v - origin) + origin
def argument(self, arg):
"""Split an argument into its constituent spaces."""
if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
if arg in self._args:
return self._args[arg]
args = []
for i, fs in enumerate(arg.function_space().split()):
# Look up the split argument in cache since we want it unique
a = Argument(fs, arg.number(), part=arg.part())
if a.shape():
if self._idx[arg.number()] == i:
args += [a[j] for j in range(a.shape()[0])]
else:
args += [Zero() for j in range(a.shape()[0])]
else:
if self._idx[arg.number()] == i:
args.append(a)
else:
args.append(Zero())
self._args[arg] = as_vector(args)
return as_vector(args)
return arg
def argument(self, arg):
"""Split an argument into its constituent spaces."""
if isinstance(arg.function_space(), functionspace.MixedFunctionSpace):
if arg in self._args:
return self._args[arg]
args = []
for i, fs in enumerate(arg.function_space().split()):
# Look up the split argument in cache since we want it unique
a = Argument(fs.ufl_element(), fs, arg.number())
if a.shape():
if self._idx[arg.number()] == i:
args += [a[j] for j in range(a.shape()[0])]
else:
args += [Zero() for j in range(a.shape()[0])]
else:
if self._idx[arg.number()] == i:
args.append(a)
else:
args.append(Zero())
self._args[arg] = as_vector(args)
return as_vector(args)
return arg
def get_scaled_jacobians2d(mesh, python=False):
"""
Computes the scaled Jacobian of each cell in a 2D triangular mesh
:arg mesh: the input mesh to do computations on
:kwarg python: compute the measure using Python?
:rtype: firedrake.function.Function scaled_jacobians with scaled
jacobian data.
"""
P0 = firedrake.FunctionSpace(mesh, "DG", 0)
if python:
P0_ten = firedrake.TensorFunctionSpace(mesh, "DG", 0)
J = firedrake.interpolate(ufl.Jacobian(mesh), P0_ten)
edge1 = ufl.as_vector([J[0, 0], J[1, 0]])
edge2 = ufl.as_vector([J[0, 1], J[1, 1]])
edge3 = edge1 - edge2
a = ufl.sqrt(ufl.dot(edge1, edge1))
b = ufl.sqrt(ufl.dot(edge2, edge2))
c = ufl.sqrt(ufl.dot(edge3, edge3))
detJ = ufl.JacobianDeterminant(mesh)
jacobian_sign = ufl.sign(detJ)
max_product = ufl.Max(
ufl.Max(ufl.Max(a * b, a * c), ufl.Max(b * c, b * a)), ufl.Max(c * a, c * b)
)
scaled_jacobians = firedrake.interpolate(detJ / max_product * jacobian_sign, P0)
else:
coords = mesh.coordinates
scaled_jacobians = firedrake.Function(P0)
op2.par_loop(
get_pyop2_kernel("get_scaled_jacobian", 2),
mesh.cell_set,
scaled_jacobians.dat(op2.WRITE, scaled_jacobians.cell_node_map()),
coords.dat(op2.READ, coords.cell_node_map()),
)
return scaled_jacobians
def __init__(self, angular_quad, L):
from transport_data import angular_tensors_ext_module
i,j,k1,k2,p,q = ufl.indices(6)
tensors = angular_tensors_ext_module.AngularTensors(angular_quad, L)
self.Y = ufl.as_tensor( numpy.reshape(tensors.Y(), tensors.shape_Y()) )
self.Q = ufl.as_tensor( numpy.reshape(tensors.Q(), tensors.shape_Q()) )
self.QT = ufl.transpose(self.Q)
self.Qt = ufl.as_tensor( numpy.reshape(tensors.Qt(), tensors.shape_Qt()) )
self.QtT = ufl.as_tensor( self.Qt[k1,p,i], (p,i,k1) )
self.G = ufl.as_tensor( numpy.reshape(tensors.G(), tensors.shape_G()) )
self.T = ufl.as_tensor( numpy.reshape(tensors.T(), tensors.shape_T()) )
self.Wp = ufl.as_vector( angular_quad.get_pw() )
self.W = ufl.diag(self.Wp)
def split(self, fields):
from ufl import as_vector, replace
from firedrake import NonlinearVariationalProblem as NLVP, FunctionSpace
splits = self._splits.get(tuple(fields))
if splits is not None:
return splits
splits = []
problem = self._problem
splitter = ExtractSubBlock()
for field in fields:
try:
if len(field) > 1:
raise NotImplementedError("Can't split into subblock")
except TypeError:
# Just a single field, we can handle that
pass
F = splitter.split(problem.F, argument_indices=(field, ))
J = splitter.split(problem.J, argument_indices=(field, field))
us = problem.u.split()
subu = us[field]
vec = []
for i, u in enumerate(us):
for idx in numpy.ndindex(u.ufl_shape):
vec.append(u[idx])
u = as_vector(vec)
F = replace(F, {problem.u: u})
J = replace(J, {problem.u: u})
if problem.Jp is not None:
Jp = splitter.split(problem.Jp, argument_indices=(field, field))
Jp = replace(Jp, {problem.u: u})
else:
Jp = None
bcs = []
for bc in problem.bcs:
if bc.function_space().index == field:
V = FunctionSpace(subu.ufl_domain(), subu.ufl_element())
bcs.append(type(bc)(V,
bc.function_arg,
bc.sub_domain,
method=bc.method))
new_problem = NLVP(F, subu, bcs=bcs, J=J, Jp=None,
form_compiler_parameters=problem.form_compiler_parameters)
new_problem._constant_jacobian = problem._constant_jacobian
splits.append(type(self)(new_problem, mat_type=self.mat_type, pmat_type=self.pmat_type,
appctx=self.appctx))
return self._splits.setdefault(tuple(fields), splits)
def solve(self, it=0):
if self.group_GS:
self.solve_group_GS(it)
else:
super(EllipticSNFluxModule, self).solve(it)
self.slns_mg = split(self.sln)
i,p,q,k1,k2 = ufl.indices(5)
sol_timer = Timer("-- Complete solution")
aux_timer = Timer("---- SOL: Computing angular flux + adjoint")
# TODO: Move to Discretization
V11 = FunctionSpace(self.DD.mesh, "CG", self.DD.parameters["p"])
for gto in range(self.DD.G):
self.PD.get_xs('D', self.D, gto)
form = self.D * ufl.diag_vector(ufl.as_matrix(self.DD.ordinates_matrix[i,p]*self.slns_mg[gto][q].dx(i), (p,q)))
for gfrom in range(self.DD.G):
pres_Ss = False
# TODO: Enlarge self.S and self.C to (L+1)^2 (or 1./2.*(L+1)*(L+2) in 2D) to accomodate for anisotropic
# scattering (lines below using CC, SS are correct only for L = 0, when the following inner loop runs only
# once.
for l in range(self.L+1):
for m in range(-l, l+1):
if self.DD.angular_quad.get_D() == 2 and divmod(l+m, 2)[1] == 0:
continue
pres_Ss |= self.PD.get_xs('Ss', self.S[l], gto, gfrom, l)
self.PD.get_xs('C', self.C[l], gto, gfrom, l)
if pres_Ss:
Cd = ufl.diag(self.C)
CC = self.tensors.Y[p,k1] * Cd[k1,k2] * self.tensors.Qt[k2,q,i]
form += ufl.as_vector(CC[p,q,i] * self.slns_mg[gfrom][q].dx(i), p)
# project(form, self.DD.Vpsi1, function=self.aux_slng, preconditioner_type="petsc_amg")
# FASTER, but requires form compilation for each dir.:
for pp in range(self.DD.M):
assign(self.aux_slng.sub(pp), project(form[pp], V11, preconditioner_type="petsc_amg"))
self.psi_mg[gto].assign(self.slns_mg[gto] + self.aux_slng)
self.adj_psi_mg[gto].assign(self.slns_mg[gto] - self.aux_slng)
def test_dolfin_expression_compilation_of_vector_literal():
# Define some literal ufl expression
uexpr = ufl.as_vector((1.23, 7.89))
# Define expected output from compilation
expected_lines = ['double s[2];',
's[0] = 1.23;',
's[1] = 7.89;',
'values[0] = s[0];',
'values[1] = s[1];']
# Define expected evaluation values: [(x,value), (x,value), ...]
expected_values = [((0.0, 0.0), (1.23, 7.89)),
((0.6, 0.7), (1.23, 7.89)),
]
# Execute all tests
check_dolfin_expression_compilation(uexpr, expected_lines, expected_values)
def argument(self, o):
from ufl import split
from firedrake import MixedFunctionSpace, FunctionSpace
V = o.function_space()
if len(V) == 1:
# Not on a mixed space, just return ourselves.
return o
if o in self._arg_cache:
return self._arg_cache[o]
V_is = V.split()
indices = self.blocks[o.number()]
try:
indices = tuple(indices)
nidx = len(indices)
except TypeError:
# Only one index provided.
indices = (indices, )
nidx = 1
if nidx == 1:
W = V_is[indices[0]]
W = FunctionSpace(W.mesh(), W.ufl_element())
a = (Argument(W, o.number(), part=o.part()), )
else:
W = MixedFunctionSpace([V_is[i] for i in indices])
a = split(Argument(W, o.number(), part=o.part()))
args = []
for i in range(len(V_is)):
if i in indices:
c = indices.index(i)
a_ = a[c]
if len(a_.ufl_shape) == 0:
args += [a_]
else:
args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)]
else:
args += [Zero()
for j in numpy.ndindex(
V_is[i].ufl_element().value_shape())]
return self._arg_cache.setdefault(o, as_vector(args))
def _find_linear_terms(rhs_exprs, u):
"""Help function that takes a list of rhs expressions and return a
list of bools determining what component, rhs_exprs[i], is linear
wrt u[i].
"""
uu = [Constant(1.0) for _ in rhs_exprs]
if len(rhs_exprs) > 1:
repl = {u: ufl.as_vector(uu)}
else:
repl = {u: uu[0]}
linear_terms = []
for i, ui in enumerate(uu):
comp_i_s = expand_indices(ufl.replace(rhs_exprs[i], repl))
linear_terms.append(ui in extract_coefficients(comp_i_s) and
ui not in extract_coefficients(
expand_derivatives(ufl.diff(comp_i_s, ui))))
return linear_terms
def argument(self, o):
V = o.function_space()
if len(V) == 1:
# Not on a mixed space, just return ourselves.
return o
# Already seen this argument, return the cached version.
if o in self._args:
return self._args[o]
args = []
for i, V_i in enumerate(V.split()):
# Walk over the subspaces and build a vector that is zero
# where the argument does not match the one we're looking
# for and is just the non-mixed argument when we do want
# it.
a = Argument(V_i, o.number(), part=o.part())
indices = numpy.ndindex(a.ufl_shape)
if self.idx[o.number()] == i:
args += [a[j] for j in indices]
else:
args += [Zero() for j in indices]
self._args[o] = as_vector(args)
return self._args[o]
def _rush_larsen_scheme_generator(rhs_form, solution, time, order, generalized):
"""Generates a list of forms and solutions for a given Butcher
tableau
*Arguments*
rhs_form (ufl.Form)
A UFL form representing the rhs for a time differentiated equation
solution (_Function_)
The prognostic variable
time (_Constant_)
A Constant holding the time at the start of the time step
order (int)
The order of the scheme
generalized (bool)
If True generate a generalized Rush Larsen scheme, linearizing all
components.
"""
DX = _check_form(rhs_form)
if DX != ufl.dP:
raise TypeError("Expected a form with a Pointintegral.")
# Create time step
dt = Constant(0.1)
# Get test function
# arguments = rhs_form.arguments()
# coefficients = rhs_form.coefficients()
# Get time dependent expressions
time_dep_expressions = _time_dependent_expressions(rhs_form, time)
# Extract rhs expressions from form
rhs_integrand = rhs_form.integrals()[0].integrand()
rhs_exprs, v = extract_tested_expressions(rhs_integrand)
vector_rhs = len(v.ufl_shape) > 0 and v.ufl_shape[0] > 1
system_size = v.ufl_shape[0] if vector_rhs else 1
# Fix for indexing of v for scalar expressions
v = v if vector_rhs else [v]
# Extract linear terms if not using generalized Rush Larsen
if not generalized:
linear_terms = _find_linear_terms(rhs_exprs, solution)
else:
linear_terms = [True for _ in range(system_size)]
# Wrap the rhs expressions into a ufl vector type
rhs_exprs = ufl.as_vector([rhs_exprs[i] for i in range(system_size)])
rhs_jac = ufl.diff(rhs_exprs, solution)
# Takes time!
if vector_rhs:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[ind, ind]))
for ind in range(system_size)]
else:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[0]))]
solution = [solution]
ufl_stage_forms = []
dolfin_stage_forms = []
dt_stage_offsets = []
# Stage solutions (3 per order rhs, linearized, and final step)
# If 2nd order the final step for 1 step is a stage
if order == 1:
stage_solutions = []
rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs, linear_terms,
system_size, solution, None, dt, time, 1.0,
0.0, v, DX, time_dep_expressions)
elif order == 2:
# Stage solution for order 2
if vector_rhs:
stage_solutions = [Function(solution.function_space(), name="y_1/2")]
else:
stage_solutions = [Function(solution[0].function_space(), name="y_1/2")]
stage_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 0.5,
0.0, v, DX,
time_dep_expressions)
rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, stage_solutions[0],
dt, time, 1.0, 0.5, v, DX,
time_dep_expressions)
ufl_stage_forms.append([stage_form])
dolfin_stage_forms.append([Form(stage_form)])
# Get last stage form
last_stage = Form(rl_ufl_form)
human_form = "%srush larsen %s" % ("generalized " if generalized else "",
str(order))
return rhs_form, linear_terms, ufl_stage_forms, dolfin_stage_forms, last_stage, \
stage_solutions, dt, dt_stage_offsets, human_form, None
def _rush_larsen_scheme_generator_adm(rhs_form, solution, time, order,
generalized, perturbation):
"""Generates a list of forms and solutions for the adjoint
linearisation of the Rush-Larsen scheme
*Arguments*
rhs_form (ufl.Form)
A UFL form representing the rhs for a time differentiated equation
solution (_Function_)
The prognostic variable
time (_Constant_)
A Constant holding the time at the start of the time step
order (int)
The order of the scheme
generalized (bool)
If True generate a generalized Rush Larsen scheme, linearizing all
components.
perturbation (Function)
The vector on which we compute the adjoint action.
"""
DX = _check_form(rhs_form)
if DX != ufl.dP:
raise TypeError("Expected a form with a Pointintegral.")
# Create time step
dt = Constant(0.1)
# Get test function
# arguments = rhs_form.arguments()
# coefficients = rhs_form.coefficients()
# Get time dependent expressions
time_dep_expressions = _time_dependent_expressions(rhs_form, time)
# Extract rhs expressions from form
rhs_integrand = rhs_form.integrals()[0].integrand()
rhs_exprs, v = extract_tested_expressions(rhs_integrand)
vector_rhs = len(v.ufl_shape) > 0 and v.ufl_shape[0] > 1
system_size = v.ufl_shape[0] if vector_rhs else 1
# Fix for indexing of v for scalar expressions
v = v if vector_rhs else [v]
# Extract linear terms if not using generalized Rush Larsen
if not generalized:
linear_terms = _find_linear_terms(rhs_exprs, solution)
else:
linear_terms = [True for _ in range(system_size)]
# Wrap the rhs expressions into a ufl vector type
rhs_exprs = ufl.as_vector([rhs_exprs[i] for i in range(system_size)])
rhs_jac = ufl.diff(rhs_exprs, solution)
# Takes time!
if vector_rhs:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[ind, ind]))
for ind in range(system_size)]
soln = solution
else:
diff_rhs_exprs = [expand_indices(expand_derivatives(rhs_jac[0]))]
solution = [solution]
soln = solution[0]
ufl_stage_forms = []
dolfin_stage_forms = []
dt_stage_offsets = []
trial = TrialFunction(soln.function_space())
# Stage solutions (3 per order rhs, linearized, and final step)
# If 2nd order the final step for 1 step is a stage
if order == 1:
stage_solutions = []
# Fetch the original step
rl_ufl_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 1.0,
0.0, v, DX,
time_dep_expressions)
# If this is commented out, we don't get NaNs. Yhy is
# solution a list of length zero anyway?
rl_ufl_form = safe_action(safe_adjoint(derivative(rl_ufl_form, soln, trial)),
perturbation)
elif order == 2:
# Stage solution for order 2
fn_space = soln.function_space()
stage_solutions = [Function(fn_space, name="y_1/2"),
Function(fn_space, name="y_1"),
Function(fn_space, name="y_bar_1/2")]
y_half_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, None, dt, time, 0.5,
0.0, v, DX,
time_dep_expressions)
y_one_form = _rush_larsen_step(rhs_exprs, diff_rhs_exprs,
linear_terms, system_size,
solution, stage_solutions[0],
dt, time, 1.0, 0.5, v, DX,
time_dep_expressions)
y_bar_half_form = safe_action(safe_adjoint(derivative(y_one_form,
stage_solutions[0], trial)), perturbation)
rl_ufl_form = safe_action(safe_adjoint(derivative(y_one_form, soln, trial)), perturbation) + \
safe_action(safe_adjoint(derivative(y_half_form, soln, trial)), stage_solutions[2])
ufl_stage_forms.append([y_half_form])
ufl_stage_forms.append([y_one_form])
ufl_stage_forms.append([y_bar_half_form])
dolfin_stage_forms.append([Form(y_half_form)])
dolfin_stage_forms.append([Form(y_one_form)])
dolfin_stage_forms.append([Form(y_bar_half_form)])
# Get last stage form
last_stage = Form(rl_ufl_form)
human_form = "%srush larsen %s" % ("generalized " if generalized else "",
str(order))
return rhs_form, linear_terms, ufl_stage_forms, dolfin_stage_forms, last_stage, \
stage_solutions, dt, dt_stage_offsets, human_form, perturbation
