def __init__(self, V, fixed_dims=[], direct_solve=False):
if isinstance(fixed_dims, int):
fixed_dims = [fixed_dims]
self.V = V
self.fixed_dims = fixed_dims
self.direct_solve = direct_solve
self.zero = fd.Constant(V.mesh().topological_dimension() * (0, ))
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
self.zero_fun = fd.Function(V)
self.a = 1e-2 * \
fd.inner(u, v) * fd.dx + fd.inner(fd.sym(fd.grad(u)),
fd.sym(fd.grad(v))) * fd.dx
self.bc_fun = fd.Function(V)
if len(self.fixed_dims) == 0:
bcs = [fd.DirichletBC(self.V, self.bc_fun, "on_boundary")]
else:
bcs = []
for i in range(self.V.mesh().topological_dimension()):
if i in self.fixed_dims:
bcs.append(fd.DirichletBC(self.V.sub(i), 0, "on_boundary"))
else:
bcs.append(
fd.DirichletBC(self.V.sub(i), self.bc_fun.sub(i),
"on_boundary"))
self.A_ext = fd.assemble(self.a, bcs=bcs, mat_type="aij")
self.ls_ext = fd.LinearSolver(self.A_ext,
solver_parameters=self.get_params())
self.A_adj = fd.assemble(self.a,
bcs=fd.DirichletBC(self.V, self.zero,
"on_boundary"),
mat_type="aij")
self.ls_adj = fd.LinearSolver(self.A_adj,
solver_parameters=self.get_params())
def solve(self, phi: fd.Function, iters: int = 5) -> fd.Function:
marking = fd.Function(self.DG0)
marking_bc_nodes = fd.Function(self.V)
# Mark cells cut by phi(x) = 0
domain = "{[i, j]: 0 <= i < b.dofs}"
instructions = """
<float64> min_value = 1e20
<float64> max_value = -1e20
for i
min_value = fmin(min_value, b[i, 0])
max_value = fmax(max_value, b[i, 0])
end
a[0, 0] = 1.0 if (min_value < 0 and max_value > 0) else 0.0
"""
fd.par_loop(
(domain, instructions),
dx,
{
"a": (marking, RW),
"b": (phi, READ)
},
is_loopy_kernel=True,
)
# Mark the nodes in the marked cells
fd.par_loop(
("{[i] : 0 <= i < A.dofs}", "A[i, 0] = fmax(A[i, 0], B[0, 0])"),
dx,
{
"A": (marking_bc_nodes, RW),
"B": (marking, READ)
},
is_loopy_kernel=True,
)
# Mark the nodes in the marked cells
# Project the gradient of phi on the cut cells
self.phi.assign(phi)
V = self.V
rho, sigma = fd.TrialFunction(V), fd.TestFunction(V)
a = rho * sigma * marking * dx
L_proj = (self.phi / sqrt(inner(grad(self.phi), grad(self.phi))) *
marking * sigma * dx)
bc_proj = BCOut(V, fd.Constant(0.0), marking_bc_nodes)
self.A_proj = fd.assemble(a, tensor=self.A_proj, bcs=bc_proj)
b_proj = fd.assemble(L_proj, bcs=bc_proj)
solver_proj = fd.LinearSolver(self.A_proj,
solver_parameters=self.solver_parameters)
solver_proj.solve(self.phi_int, b_proj)
def nabla_phi_bar(phi):
return sqrt(inner(grad(phi), grad(phi)))
def d1(s):
return fd.Constant(1.0) - fd.Constant(1.0) / s
def d2(s):
return conditional(le(s, fd.Constant(1.0)), s - fd.Constant(1.0),
d1(s))
def residual_phi(phi):
return fd.norm(
fd.assemble(
d2(nabla_phi_bar(phi)) * inner(grad(phi), grad(sigma)) *
dx))
a = inner(grad(rho), grad(sigma)) * dx
L = (inner(
(-d2(nabla_phi_bar(self.phi)) + fd.Constant(1.0)) * grad(self.phi),
grad(sigma),
) * dx)
bc = BCInt(V, self.phi_int, marking_bc_nodes)
phi_sol = fd.Function(V)
A = fd.assemble(a, bcs=bc)
b = fd.assemble(L, bcs=bc)
solver = fd.LinearSolver(A, solver_parameters=self.solver_parameters)
# Solve the Signed distance equation with Picard iteration
bc.apply(phi_sol)
init_res = residual_phi(phi_sol)
res = 1e10
it = 0
while res > init_res or it < iters:
solver.solve(phi_sol, b)
self.phi.assign(phi_sol)
b = fd.assemble(L, bcs=bc, tensor=b)
it += 1
res = residual_phi(phi_sol)
if res > init_res:
fd.warning(
f"Residual in signed distance function increased: {res}, before: {init_res}"
)
return self.phi
def __init__(self, Q, fixed_bids=[], extra_bcs=[], direct_solve=False):
if isinstance(extra_bcs, fd.DirichletBC):
extra_bcs = [extra_bcs]
self.direct_solve = direct_solve
self.fixed_bids = fixed_bids # fixed parts of bdry
self.params = self.get_params() # solver parameters
self.Q = Q
"""
V: type fd.FunctionSpace
I: type PETSc.Mat, interpolation matrix between V and ControlSpace
"""
(V, I_interp) = Q.get_space_for_inner()
free_bids = list(V.mesh().topology.exterior_facets.unique_markers)
self.free_bids = [int(i) for i in free_bids] # np.int->int
for bid in self.fixed_bids:
self.free_bids.remove(bid)
# Some weak forms have a nullspace. We import the nullspace if no
# parts of the bdry are fixed (we assume that a DirichletBC is
# sufficient to empty the nullspace).
nsp = None
if len(self.fixed_bids) == 0:
nsp_functions = self.get_nullspace(V)
if nsp_functions is not None:
nsp = fd.VectorSpaceBasis(nsp_functions)
nsp.orthonormalize()
bcs = []
# impose homogeneous Dirichlet bcs on bdry parts that are fixed.
if len(self.fixed_bids) > 0:
dim = V.value_size
if dim == 2:
zerovector = fd.Constant((0, 0))
elif dim == 3:
zerovector = fd.Constant((0, 0, 0))
else:
raise NotImplementedError
bcs.append(fd.DirichletBC(V, zerovector, self.fixed_bids))
if len(extra_bcs) > 0:
bcs += extra_bcs
if len(bcs) == 0:
bcs = None
a = self.get_weak_form(V)
A = fd.assemble(a, mat_type='aij', bcs=bcs)
ls = fd.LinearSolver(A,
solver_parameters=self.params,
nullspace=nsp,
transpose_nullspace=nsp)
self.ls = ls
self.A = A.petscmat
self.interpolated = False
# If the matrix I is passed, replace A with transpose(I)*A*I
# and set up a ksp solver for self.riesz_map
if I_interp is not None:
self.interpolated = True
ITAI = self.A.PtAP(I_interp)
from firedrake.petsc import PETSc
import numpy as np
zero_rows = []
# if there are zero-rows, replace them with rows that
# have 1 on the diagonal entry
for row in range(*ITAI.getOwnershipRange()):
(cols, vals) = ITAI.getRow(row)
valnorm = np.linalg.norm(vals)
if valnorm < 1e-13:
zero_rows.append(row)
for row in zero_rows:
ITAI.setValue(row, row, 1.0)
ITAI.assemble()
# overwrite the self.A created by get_impl
self.A = ITAI
# create ksp solver for self.riesz_map
Aksp = PETSc.KSP().create(comm=V.comm)
Aksp.setOperators(ITAI)
Aksp.setType("preonly")
Aksp.pc.setType("cholesky")
Aksp.pc.setFactorSolverType("mumps")
Aksp.setFromOptions()
Aksp.setUp()
self.Aksp = Aksp
def solver(self):
return firedrake.LinearSolver(self.A, solver_parameters=self.solver_parameters)
def solver_CG(mesh, el, space, deg, T, dt=0.001, warm_up=False):
"""Solve the scalar wave equation on a unit square/cube using a
CG FEM formulation with several different element types.
Parameters
----------
mesh: Firedrake.mesh
A utility mesh from the Firedrake package
el: string
The type of element either "tria" or "quad".
`tria` in 3d implies tetrahedra and
`quad` in 3d implies hexahedral elements.
space: string
The space of the FEM. Available options are:
"CG": Continuous Galerkin Finite Elements,
"KMV": Kong-Mulder-Veldhuzien higher-order mass lumped elements
"S" (for Serendipity) (NB: quad/hexs only)
"spectral": spectral elements using GLL quad
points (NB: quads/hexs only)
deg: int
The spatial polynomial degree.
T: float
The simulation duration in simulation seconds.
dt: float, optional
Simulation timestep
warm_up: boolean, optional
Warm up symbolics by running one timestep and shutting down.
Returns
-------
u_n: Firedrake.Function
The solution at time `T`
"""
sd = mesh.geometric_dimension()
V = _build_space(mesh, el, space, deg)
quad_rule1, quad_rule2 = _build_quad_rule(el, V, space)
params = _select_params(space)
# DEBUG
# outfile = fd.File(os.getcwd() + "/results/simple_shots.pvd")
# END DEBUG
tot_dof = COMM_WORLD.allreduce(V.dof_dset.total_size, op=MPI.SUM)
# if COMM_WORLD.rank == 0:
# print("------------------------------------------")
# print("The problem has " + str(tot_dof) + " degrees of freedom.")
# print("------------------------------------------")
nt = int(T / dt) # number of timesteps
u = fd.TrialFunction(V)
v = fd.TestFunction(V)
u_np1 = fd.Function(V) # n+1
u_n = fd.Function(V) # n
u_nm1 = fd.Function(V) # n-1
# constant speed
c = Constant(1.5)
m = (
(1.0 / (c * c))
* (u - 2.0 * u_n + u_nm1)
/ Constant(dt * dt)
* v
* dx(rule=quad_rule1)
) # mass-like matrix
a = dot(grad(u_n), grad(v)) * dx(rule=quad_rule2) # stiffness matrix
# injection of source into mesh
ricker = Constant(0.0)
source = Constant([0.5] * sd)
coords = fd.SpatialCoordinate(mesh)
F = m + a - delta_expr(source, *coords) * ricker * v * dx(rule=quad_rule2)
a, r = fd.lhs(F), fd.rhs(F)
A, R = fd.assemble(a), fd.assemble(r)
solver = fd.LinearSolver(A, solver_parameters=params, options_prefix="")
# timestepping loop
results = []
t = 0.0
for step in range(nt):
with PETSc.Log.Stage("{el}{deg}".format(el=el, deg=deg)):
ricker.assign(RickerWavelet(t, freq=6))
R = fd.assemble(r, tensor=R)
solver.solve(u_np1, R)
snes = _get_time("SNESSolve")
ksp = _get_time("KSPSolve")
pcsetup = _get_time("PCSetUp")
pcapply = _get_time("PCApply")
jac = _get_time("SNESJacobianEval")
residual = _get_time("SNESFunctionEval")
sparsity = _get_time("CreateSparsity")
results.append(
[tot_dof, snes, ksp, pcsetup, pcapply, jac, residual, sparsity]
)
if warm_up:
# Warm up symbolics/disk cache
solver.solve(u_np1, R)
sys.exit("Warming up...")
u_nm1.assign(u_n)
u_n.assign(u_np1)
t = step * float(dt)
# if step % 10 == 0:
# outfile.write(u_n)
# print("Time is " + str(t), flush=True)
results = np.asarray(results)
if mesh.comm.rank == 0:
with open(
"data/scalar_wave.{el}.{deg}.{space}.csv".format(
el=el, deg=deg, space=space
),
"w",
) as f:
np.savetxt(
f,
results,
fmt=["%d"] + ["%e"] * 7,
delimiter=",",
header="tot_dof,SNESSolve,KSPSolve,PCSetUp,PCApply,SNESJacobianEval,SNESFunctionEval,CreateSparsity",
comments="",
)
return u_n
def initialize(self, pc):
if pc.getType() != "python":
raise ValueError("Expecting PC type python")
prefix = pc.getOptionsPrefix() + "diagfft_"
# we assume P has things stuffed inside of it
_, P = pc.getOperators()
context = P.getPythonContext()
appctx = context.appctx
self.appctx = appctx
# all at once solution passed through the appctx
self.w_all = appctx.get("w_all", None)
# FunctionSpace checks
test, trial = context.a.arguments()
if test.function_space() != trial.function_space():
raise ValueError("Pressure space test and trial space differ")
W = test.function_space()
# basic model function space
get_blockV = appctx.get("get_blockV", None)
self.blockV = get_blockV()
M = int(W.dim() / self.blockV.dim())
assert (self.blockV.dim() * M == W.dim())
self.M = M
self.NM = W.dim()
# Input/Output wrapper Functions
self.xf = fd.Function(W) # input
self.yf = fd.Function(W) # output
# Gamma coefficients
Nt = M
exponents = np.arange(Nt) / Nt
alphav = appctx.get("alpha", None)
self.Gam = alphav**exponents
# Di coefficients
thetav = appctx.get("theta", None)
Dt = appctx.get("dt", None)
C1col = np.zeros(Nt)
C2col = np.zeros(Nt)
C1col[:2] = np.array([1, -1]) / Dt
C2col[:2] = np.array([thetav, 1 - thetav])
self.D1 = np.sqrt(Nt) * fft(self.Gam * C1col)
self.D2 = np.sqrt(Nt) * fft(self.Gam * C2col)
# Block system setup
# First need to build the vector function space version of
# blockV
mesh = self.blockV.mesh()
Ve = self.blockV.ufl_element()
if isinstance(Ve, fd.MixedElement):
MixedCpts = []
self.ncpts = Ve.num_sub_elements()
for cpt in range(Ve.num_sub_elements()):
SubV = Ve.sub_elements()[cpt]
if isinstance(SubV, fd.FiniteElement):
MixedCpts.append(fd.VectorElement(SubV, dim=2))
elif isinstance(SubV, fd.VectorElement):
shape = (2, SubV.num_sub_elements())
MixedCpts.append(fd.TensorElement(SubV, shape))
elif isinstance(SubV, fd.TensorElement):
shape = (2, ) + SubV._shape
MixedCpts.append(fd.TensorElement(SubV, shape))
else:
raise NotImplementedError
dim = len(MixedCpts)
self.CblockV = np.prod(
[fd.FunctionSpace(mesh, MixedCpts[i]) for i in range(dim)])
else:
self.ncpts = 1
if isinstance(Ve, fd.FiniteElement):
self.CblockV = fd.FunctionSpace(mesh,
fd.VectorElement(Ve, dim=2))
elif isinstance(Ve, fd.VectorElement):
shape = (2, Ve.num_sub_elements())
self.CblockV = fd.FunctionSpace(mesh,
fd.TensorElement(Ve, shape))
elif isinstance(Ve, fd.TensorElement):
shape = (2, ) + Ve._shape
self.CblockV = fd.FunctionSpace(mesh,
fd.TensorElement(Ve, shape))
else:
raise NotImplementedError
# Now need to build the block solver
vs = fd.TestFunctions(self.CblockV)
self.u0 = fd.Function(self.CblockV) # we will create a linearisation
us = fd.split(self.u0)
# extract the real and imaginary parts
vsr = []
vsi = []
usr = []
usi = []
if isinstance(Ve, fd.MixedElement):
N = Ve.num_sub_elements()
for i in range(N):
SubV = Ve.sub_elements()[i]
if len(SubV.value_shape()) == 0:
vsr.append(vs[i][0])
vsi.append(vs[i][1])
usr.append(us[i][0])
usi.append(us[i][1])
elif len(SubV.value_shape()) == 1:
vsr.append(vs[i][0, :])
vsi.append(vs[i][1, :])
usr.append(us[i][0, :])
usi.append(us[i][1, :])
elif len(SubV.value_shape()) == 2:
vsr.append(vs[i][0, :, :])
vsi.append(vs[i][1, :, :])
usr.append(us[i][0, :, :])
usi.append(us[i][1, :, :])
else:
raise NotImplementedError
else:
if isinstance(Ve, fd.FiniteElement):
vsr.append(vs[0])
vsi.append(vs[1])
usr.append(us[0])
usi.append(us[1])
elif isinstance(Ve, fd.VectorElement):
vsr.append(vs[0, :])
vsi.append(vs[1, :])
usr.append(self.u0[0, :])
usi.append(self.u0[1, :])
elif isinstance(Ve, fd.TensorElement):
vsr.append(vs[0, :])
vsi.append(vs[1, :])
usr.append(self.u0[0, :])
usi.append(self.u0[1, :])
else:
raise NotImplementedError
# input and output functions
self.Jprob_in = fd.Function(self.CblockV)
self.Jprob_out = fd.Function(self.CblockV)
# A place to store all the inputs to the block problems
self.xfi = fd.Function(W)
self.xfr = fd.Function(W)
# Building the nonlinear operator
self.Jsolvers = []
self.Js = []
form_mass = appctx.get("form_mass", None)
form_function = appctx.get("form_function", None)
# setting up the Riesz map
# input for the Riesz map
self.xtemp = fd.Function(self.CblockV)
v = fd.TestFunction(self.CblockV)
u = fd.TrialFunction(self.CblockV)
a = fd.assemble(fd.inner(u, v) * fd.dx)
self.Proj = fd.LinearSolver(a, options_prefix=prefix + "mass_")
# building the block problem solvers
for i in range(M):
D1i = fd.Constant(np.imag(self.D1[i]))
D1r = fd.Constant(np.real(self.D1[i]))
D2i = fd.Constant(np.imag(self.D2[i]))
D2r = fd.Constant(np.real(self.D2[i]))
# pass sigma into PC:
sigma = self.D1[i]**2 / self.D2[i]
sigma_inv = self.D2[i]**2 / self.D1[i]
appctx_h = appctx.copy()
appctx_h["sr"] = fd.Constant(np.real(sigma))
appctx_h["si"] = fd.Constant(np.imag(sigma))
appctx_h["sinvr"] = fd.Constant(np.real(sigma_inv))
appctx_h["sinvi"] = fd.Constant(np.imag(sigma_inv))
appctx_h["D2r"] = D2r
appctx_h["D2i"] = D2i
appctx_h["D1r"] = D1r
appctx_h["D1i"] = D1i
A = (D1r * form_mass(*usr, *vsr) - D1i * form_mass(*usi, *vsr) +
D2r * form_function(*usr, *vsr) -
D2i * form_function(*usi, *vsr) +
D1r * form_mass(*usi, *vsi) + D1i * form_mass(*usr, *vsi) +
D2r * form_function(*usi, *vsi) +
D2i * form_function(*usr, *vsi))
# The linear operator
J = fd.derivative(A, self.u0)
# The rhs
v = fd.TestFunction(self.CblockV)
L = fd.inner(v, self.Jprob_in) * fd.dx
block_prefix = prefix + str(i) + '_'
jprob = fd.LinearVariationalProblem(J, L, self.Jprob_out)
Jsolver = fd.LinearVariationalSolver(jprob,
appctx=appctx_h,
options_prefix=block_prefix)
self.Jsolvers.append(Jsolver)