def update(self, x, flag, iteration):
"""Update domain and solution to state and adjoint equation."""
if self.Q.update_domain(x):
try:
# We use pyadjoint to calculate adjoint and shape derivatives,
# in order to do this we need to "record a tape of the forward
# solve", pyadjoint will then figure out all necessary
# adjoints.
tape = fda.get_working_tape()
tape.clear_tape()
fda.continue_annotation()
mesh_m = self.J.Q.mesh_m
s = fda.Function(self.J.V_m)
mesh_m.coordinates.assign(mesh_m.coordinates + s)
self.s = s
self.c = fda.Control(s)
self.e.solve()
Jpyadj = fda.assemble(self.J.value_form())
self.Jred = fda.ReducedFunctional(Jpyadj, self.c)
fda.pause_annotation()
except fd.ConvergenceError:
if self.cb is not None:
self.cb()
raise
if iteration >= 0 and self.cb is not None:
self.cb()
def update(self, x, flag, iteration):
"""Update domain and solution to state and adjoint equation."""
if self.Q.update_domain(x):
try:
# We use pyadjoint to calculate adjoint and shape derivatives,
# in order to do this we need to "record a tape of the forward
# solve", pyadjoint will then figure out all necessary
# adjoints.
import firedrake_adjoint as fda
tape = fda.get_working_tape()
tape.clear_tape()
# ensure we are annotating
from pyadjoint.tape import annotate_tape
safety_counter = 0
while not annotate_tape():
safety_counter += 1
fda.continue_annotation()
if safety_counter > 1e2:
import sys
sys.exit('Cannot annotate even after 100 attempts.')
mesh_m = self.J.Q.mesh_m
s = fd.Function(self.J.V_m)
mesh_m.coordinates.assign(mesh_m.coordinates + s)
self.s = s
self.c = fda.Control(s)
self.e.solve()
Jpyadj = fd.assemble(self.J.value_form())
self.Jred = fda.ReducedFunctional(Jpyadj, self.c)
fda.pause_annotation()
except fd.ConvergenceError:
if self.cb is not None:
self.cb()
raise
if iteration >= 0 and self.cb is not None:
self.cb()
def gradient_test_acoustic(model,
mesh,
V,
comm,
vp_exact,
vp_guess,
mask=None): #{{{
import firedrake_adjoint as fire_adj
with fire_adj.stop_annotating():
if comm.comm.rank == 0:
print('######## Starting gradient test ########', flush=True)
sources = spyro.Sources(model, mesh, V, comm)
receivers = spyro.Receivers(model, mesh, V, comm)
wavelet = spyro.full_ricker_wavelet(
model["timeaxis"]["dt"],
model["timeaxis"]["tf"],
model["acquisition"]["frequency"],
)
point_cloud = receivers.set_point_cloud(comm)
# simulate the exact model
if comm.comm.rank == 0:
print('######## Running the exact model ########', flush=True)
p_exact_recv = forward(model, mesh, comm, vp_exact, sources, wavelet,
point_cloud)
# simulate the guess model
if comm.comm.rank == 0:
print('######## Running the guess model ########', flush=True)
p_guess_recv, Jm = forward(model,
mesh,
comm,
vp_guess,
sources,
wavelet,
point_cloud,
fwi=True,
true_rec=p_exact_recv)
if comm.comm.rank == 0:
print("\n Cost functional at fixed point : " + str(Jm) + " \n ",
flush=True)
# compute the gradient of the control (to be verified)
if comm.comm.rank == 0:
print(
'######## Computing the gradient by automatic differentiation ########',
flush=True)
control = fire_adj.Control(vp_guess)
dJ = fire_adj.compute_gradient(Jm, control)
if mask:
dJ *= mask
# File("gradient.pvd").write(dJ)
#steps = [1e-3, 1e-4, 1e-5, 1e-6, 1e-7] # step length
#steps = [1e-4, 1e-5, 1e-6, 1e-7] # step length
steps = [1e-5, 1e-6, 1e-7, 1e-8] # step length
with fire_adj.stop_annotating():
delta_m = Function(V) # model direction (random)
delta_m.assign(dJ)
Jhat = fire_adj.ReducedFunctional(Jm, control)
derivative = enlisting.Enlist(Jhat.derivative())
hs = enlisting.Enlist(delta_m)
projnorm = sum(hi._ad_dot(di) for hi, di in zip(hs, derivative))
# this deepcopy is important otherwise pertubations accumulate
vp_original = vp_guess.copy(deepcopy=True)
if comm.comm.rank == 0:
print(
'######## Computing the gradient by finite diferences ########',
flush=True)
errors = []
for step in steps: # range(3):
# steps.append(step)
# perturb the model and calculate the functional (again)
# J(m + delta_m*h)
vp_guess = vp_original + step * delta_m
p_guess_recv, Jp = forward(model,
mesh,
comm,
vp_guess,
sources,
wavelet,
point_cloud,
fwi=True,
true_rec=p_exact_recv)
fd_grad = (Jp - Jm) / step
if comm.comm.rank == 0:
print("\n Cost functional for step " + str(step) + " : " +
str(Jp) + ", fd approx.: " + str(fd_grad) +
", grad'*dir : " + str(projnorm) + " \n ",
flush=True)
errors.append(100 * ((fd_grad - projnorm) / projnorm))
fire_adj.get_working_tape().clear_tape()
# all errors less than 1 %
errors = np.array(errors)
assert (np.abs(errors) < 5.0).all()
def compliance():
parser = argparse.ArgumentParser(description="Compliance problem with MMA")
parser.add_argument(
"--nref",
action="store",
dest="nref",
type=int,
help="Number of mesh refinements",
default=2,
)
parser.add_argument(
"--uniform",
action="store",
dest="uniform",
type=int,
help="Use uniform mesh",
default=0,
)
parser.add_argument(
"--inner_product",
action="store",
dest="inner_product",
type=str,
help="Inner product, euclidean or L2",
default="L2",
)
parser.add_argument(
"--output_dir",
action="store",
dest="output_dir",
type=str,
help="Directory for all the output",
default="./",
)
args = parser.parse_args()
nref = args.nref
inner_product = args.inner_product
output_dir = args.output_dir
assert inner_product == "L2" or inner_product == "euclidean"
mesh = fd.Mesh("./beam_uniform.msh")
#mh = fd.MeshHierarchy(mesh, 2)
#mesh = mh[-1]
if nref > 0:
mh = fd.MeshHierarchy(mesh, nref)
mesh = mh[-1]
elif nref < 0:
raise RuntimeError("Non valid mesh argument")
V = fd.VectorFunctionSpace(mesh, "CG", 1)
u, v = fd.TrialFunction(V), fd.TestFunction(V)
# Elasticity parameters
E, nu = 1e0, 0.3
mu, lmbda = fd.Constant(E / (2 * (1 + nu))), fd.Constant(
E * nu / ((1 + nu) * (1 - 2 * nu))
)
# Helmholtz solver
RHO = fd.FunctionSpace(mesh, "CG", 1)
rho = fd.interpolate(fd.Constant(0.1), RHO)
af, b = fd.TrialFunction(RHO), fd.TestFunction(RHO)
#filter_radius = fd.Constant(0.2)
#x, y = fd.SpatialCoordinate(mesh)
#x_ = fd.interpolate(x, RHO)
#y_ = fd.interpolate(y, RHO)
#aH = filter_radius * inner(grad(af), grad(b)) * dx + af * b * dx
#LH = rho * b * dx
rhof = fd.Function(RHO)
solver_params = {
"ksp_type": "preonly",
"pc_type": "lu",
"pc_factor_mat_solver_type": "mumps",
"mat_mumps_icntl_14": 200,
"mat_mumps_icntl_24": 1,
}
#fd.solve(aH == LH, rhof, solver_parameters=solver_params)
rhof.assign(rho)
rhofControl = fda.Control(rhof)
eps = fd.Constant(1e-5)
p = fd.Constant(3.0)
def simp(rho):
return eps + (fd.Constant(1.0) - eps) * rho ** p
def epsilon(v):
return sym(nabla_grad(v))
def sigma(v):
return 2.0 * mu * epsilon(v) + lmbda * tr(epsilon(v)) * Identity(2)
DIRICHLET = 3
NEUMANN = 4
a = inner(simp(rhof) * sigma(u), epsilon(v)) * dx
load = fd.Constant((0.0, -1.0))
L = inner(load, v) * ds(NEUMANN)
u_sol = fd.Function(V)
bcs = fd.DirichletBC(V, fd.Constant((0.0, 0.0)), DIRICHLET)
fd.solve(a == L, u_sol, bcs=bcs, solver_parameters=solver_params)
c = fda.Control(rho)
J = fd.assemble(fd.Constant(1e-4) * inner(u_sol, load) * ds(NEUMANN))
Vol = fd.assemble(rhof * dx)
VolControl = fda.Control(Vol)
with fda.stop_annotating():
Vlimit = fd.assemble(fd.Constant(1.0) * dx(domain=mesh)) * 0.5
rho_viz_f = fd.Function(RHO, name="rho")
plot_file = f"{output_dir}/design_{inner_product}.pvd"
controls_f = fd.File(plot_file)
def deriv_cb(j, dj, rho):
with fda.stop_annotating():
rho_viz_f.assign(rhofControl.tape_value())
controls_f.write(rho_viz_f)
Jhat = fda.ReducedFunctional(J, c, derivative_cb_post=deriv_cb)
Volhat = fda.ReducedFunctional(Vol, c)
class VolumeConstraint(fda.InequalityConstraint):
def __init__(self, Vhat, Vlimit, VolControl):
self.Vhat = Vhat
self.Vlimit = float(Vlimit)
self.VolControl = VolControl
def function(self, m):
# Compute the integral of the control over the domain
integral = self.VolControl.tape_value()
with fda.stop_annotating():
value = -integral / self.Vlimit + 1.0
return [value]
def jacobian(self, m):
with fda.stop_annotating():
gradients = self.Vhat.derivative()
with gradients.dat.vec as v:
v.scale(-1.0 / self.Vlimit)
return [gradients]
def output_workspace(self):
return [0.0]
def length(self):
"""Return the number of components in the constraint vector (here, one)."""
return 1
lb = 1e-5
ub = 1.0
problem = fda.MinimizationProblem(
Jhat,
bounds=(lb, ub),
constraints=[VolumeConstraint(Volhat, Vlimit, VolControl)],
)
parameters_mma = {
"move": 0.2,
"maximum_iterations": 200,
"m": 1,
"IP": 0,
"tol": 1e-6,
"accepted_tol": 1e-4,
"norm": inner_product,
#"norm": "euclidean",
"gcmma": False,
}
solver = MMASolver(problem, parameters=parameters_mma)
rho_opt = solver.solve()
with open(f"{output_dir}/finished_{inner_product}.txt", "w") as f:
f.write("Done")
u_interpolated.dat.data[:] = interpolator(X[:, 0], X[:, 1])
# Two terms in the functional - note difference in misfit term!
misfit_expr = 0.5 * ((u_interpolated - u) / σ)**2
α = firedrake.Constant(0.5)
regularisation_expr = 0.5 * α**2 * inner(grad(q), grad(q))
# Should be able to write firedrake.assemble(misfit + regularisation * dx) but can't yet
# because of the meshes being different in the point-cloud case
print('Assembling J')
J = firedrake.assemble(misfit_expr * dx) + firedrake.assemble(regularisation_expr * dx)
# Create reduced functional
print('Creating q̂ and Ĵ')
q̂ = firedrake_adjoint.Control(q)
Ĵ = firedrake_adjoint.ReducedFunctional(J, q̂)
# Minimise reduced functional
print('Minimising Ĵ to get q_min')
q_min = firedrake_adjoint.minimize(
Ĵ, method='Newton-CG', options={'disp': True}
)
q_min.rename(name=f'q_min_{method}_{num_points}')
# Clear tape to avoid memory leak
print('Clearing tape')
tape.clear_tape()
# Calculate error terms
print('Calculating error')
q_err = firedrake.Function(Q).assign(q_min-q_true)