def pytest_runtest_teardown(item, nextitem):
"""Clear Thetis caches after running a test"""
from firedrake.tsfc_interface import TSFCKernel
from pyop2.op2 import Kernel
from pyop2.base import JITModule
from firedrake_adjoint import get_working_tape
# disgusting hack, clear the Class-Cached objects in PyOP2 and
# Firedrake, otherwise these will never be collected. The Kernels
# get very big with bendy on.
Kernel._cache.clear()
TSFCKernel._cache.clear()
JITModule._cache.clear()
# clear the adjoint tape, so subsequent tests don't interfere
get_working_tape().clear_tape()
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 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 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()
# In[3]:
from firedrake import exp, inner, grad, dx
u_true = firedrake.Function(V)
v = firedrake.TestFunction(V)
f = Constant(1.0)
k0 = Constant(0.5)
bc = firedrake.DirichletBC(V, 0, 'on_boundary')
F = (k0 * exp(q_true) * inner(grad(u_true), grad(v)) - f * v) * dx
firedrake.solve(F == 0, u_true, bc)
print('Made fake u_true')
# Clear tape since don't need to have taped above
tape = firedrake_adjoint.get_working_tape()
tape.clear_tape()
# ## Generating Observational Data $u_{obs}$
# We run up in powers of 2 until we have plenty of observations per cell (on average)
# In[4]:
signal_to_noise = 20
U = u_true.dat.data_ro[:]
u_range = U.max() - U.min()
σ = firedrake.Constant(u_range / signal_to_noise)
methods = ['point-cloud', 'nearest', 'linear', 'clough-tocher', 'gaussian']
min_power_of_2 = 2
def heat_exchanger_3D():
parser = argparse.ArgumentParser(description="Level set method parameters")
parser.add_argument(
"--nu",
action="store",
dest="nu",
type=float,
help="Kinematic Viscosity",
default=1.0,
)
parser.add_argument(
"--brinkmann_penalty",
action="store",
dest="brinkmann_penalty",
type=float,
help="Brinkmann term",
default=1e5,
)
parser.add_argument(
"--refinement",
action="store",
dest="refinement",
type=int,
help="Level of refinement",
default=0,
)
parser.add_argument(
"--pressure_drop_constraint",
action="store",
dest="pressure_drop_constraint",
type=float,
help="Pressure drop constraint",
default=10.0,
)
parser.add_argument(
"--n_iters",
dest="n_iters",
type=int,
action="store",
default=1000,
help="Number of optimization iterations",
)
parser.add_argument(
"--output_dir",
dest="output_dir",
type=str,
action="store",
default="./",
help="Output folder",
)
parser.add_argument(
"--type",
dest="type_he",
type=str,
action="store",
default="parallel",
help="Type of heat exchanger: parallel or counter",
)
opts = parser.parse_args()
print(f"Parameters used: {opts}")
beta_param = 0.4
mesh = fd.Mesh("./box_heat_exch.msh")
from parameters_box import (
INLET2,
INLET1,
OUTLET1,
OUTLET2,
INMOUTH1,
INMOUTH2,
OUTMOUTH1,
OUTMOUTH2,
)
if opts.type_he == "counter":
INLET1, OUTLET1 = OUTLET1, INLET1
elif opts.type_he == "u_flow":
INLET2, OUTLET1 = OUTLET1, INLET2
OUTMOUTH1, INMOUTH2 = INMOUTH2, OUTMOUTH1
markers = {
"WALLS": WALLS,
"INLET1": INLET1,
"INLET2": INLET2,
"OUTLET1": OUTLET1,
"OUTLET2": OUTLET2,
}
no_flow_domain_1 = [INMOUTH2, OUTMOUTH2]
no_flow_domain_2 = [INMOUTH1, OUTMOUTH1]
no_flow = no_flow_domain_1.copy()
no_flow.extend(no_flow_domain_2)
mesh = mark_no_flow_regions(mesh, no_flow, no_flow)
mh = fd.MeshHierarchy(mesh, opts.refinement)
mesh = mh[-1]
pressure_drop_constraint = opts.pressure_drop_constraint
pressure_drop_1 = pressure_drop_constraint
pressure_drop_2 = pressure_drop_constraint
S = fd.VectorFunctionSpace(mesh, "CG", 1)
PHI = fd.FunctionSpace(mesh, "CG", 1)
phi = fd.Function(PHI, name="LevelSet")
x, y, z = fd.SpatialCoordinate(mesh)
ω = 0.25
phi_expr = sin(y * pi / ω) * cos(x * pi / ω) * sin(
z * pi / ω) - fd.Constant(0.2)
checkpoints = is_checkpoint(opts.output_dir)
if checkpoints:
current_iter = read_checkpoint(checkpoints, phi)
with open(f"{opts.output_dir}/brinkmann_penalty.txt",
"r") as txt_brinkmann:
brinkmann_penalty_initial = fd.Constant(txt_brinkmann.read())
print(
f"Current brinkmann term: {brinkmann_penalty_initial.values()[0]}")
else:
with stop_annotating():
phi.interpolate(phi_expr)
current_iter = 0
brinkmann_penalty_initial = fd.Constant(opts.brinkmann_penalty)
P2 = fd.VectorElement("CG", mesh.ufl_cell(), 1)
P1 = fd.FiniteElement("CG", mesh.ufl_cell(), 1)
TH = P2 * P1
W = fd.FunctionSpace(mesh, TH)
print(f"DOFS: {W.dim()}")
T = fd.FunctionSpace(mesh, "CG", 1)
global_counter = count(current_iter)
def receive_signal(signum, stack):
iter_current = next(copy(global_counter))
print(f"Received: {signum}, iter: {iter_current}")
with fd.HDF5File(
f"{opts.output_dir}/checkpoint_iter_{iter_current}.h5",
"w") as checkpoint:
checkpoint.write(phi, "/checkpoint")
with open(f"{opts.output_dir}/brinkmann_penalty.txt",
"w") as txt_brinkmann:
txt_brinkmann.write(str(brinkmann_penalty.values()[0]))
signal.signal(signal.SIGHUP, receive_signal)
phi_pvd = fd.File(
f"{opts.output_dir}/phi_evolution.pvd",
target_degree=1,
target_continuity=fd.H1,
mode="a",
)
ns1 = fd.File(f"{opts.output_dir}/navier_stokes_1.pvd", mode="a")
ns2 = fd.File(f"{opts.output_dir}/navier_stokes_2.pvd", mode="a")
temperature = fd.File(f"{opts.output_dir}/temperature.pvd", mode="a")
temperature_pvd = fd.Function(T)
def termination_event_1():
p1_constraint = P1control.tape_value() - 1
p2_constraint = P2control.tape_value() - 1
event_value = max(p1_constraint, p2_constraint)
print(f"Value event: {event_value}")
return event_value
def termination_event_2():
iter_current = next(copy(global_counter))
print(f"Value event iter count: {iter_current}")
return float(iter_current % 500)
brinkmann_penalty_initial_value = brinkmann_penalty_initial.values()[0]
if brinkmann_penalty_initial_value > opts.brinkmann_penalty:
termination_events = [termination_event_2, termination_event_2]
brinkmann_pen_terms = [
brinkmann_penalty_initial,
fd.Constant(brinkmann_penalty_initial_value * 10),
]
else:
termination_events = [termination_event_1, None]
brinkmann_pen_terms = [
brinkmann_penalty_initial,
fd.Constant(brinkmann_penalty_initial_value * 5),
]
for termination_event, brinkmann_penalty in zip(termination_events,
brinkmann_pen_terms):
s = fd.Function(S, name="deform")
mesh.coordinates.assign(mesh.coordinates + s)
# w_sol1, w_sol2, t = forward(brinkmann_penalty)
w_sol1, w_sol2, t = forward(
W,
T,
phi,
opts,
brinkmann_penalty,
no_flow_domain_1=no_flow_domain_1,
no_flow_domain_2=no_flow_domain_2,
markers=markers,
)
w_sol1_control = fda.Control(w_sol1)
w_sol2_control = fda.Control(w_sol2)
t_control = fda.Control(t)
J = heat_flux(w_sol2, t, OUTLET2)
J_hot = heat_flux(w_sol1, t, OUTLET1)
Pdrop1 = pressure_drop(w_sol1, INLET1, OUTLET1, pressure_drop_1)
Pdrop2 = pressure_drop(w_sol2, INLET2, OUTLET2, pressure_drop_2)
print(f"Cold flux: {J}, hot flux: {J_hot}")
c = fda.Control(s)
J_hot_control = fda.Control(J_hot)
def deriv_cb(phi):
with stop_annotating():
iter = next(global_counter)
print(f"Hot flux: {J_hot_control.tape_value()}")
if iter % 15 == 0:
u_sol1, p_sol1 = w_sol1_control.tape_value().split()
u_sol2, p_sol2 = w_sol2_control.tape_value().split()
u_sol1.rename("Velocity1")
p_sol1.rename("Pressure1")
u_sol2.rename("Velocity2")
p_sol2.rename("Pressure2")
ns1.write(u_sol1, p_sol1, time=iter)
ns2.write(u_sol2, p_sol2, time=iter)
phi_pvd.write(phi[0], time=iter)
temperature_pvd.assign(t_control.tape_value())
temperature_pvd.rename("temperature")
temperature.write(temperature_pvd, time=iter)
# Reduced Functionals
Jhat = LevelSetFunctional(J, c, phi, derivative_cb_pre=deriv_cb)
P1hat = LevelSetFunctional(Pdrop1, c, phi)
P1control = fda.Control(Pdrop1)
P2hat = LevelSetFunctional(Pdrop2, c, phi)
P2control = fda.Control(Pdrop2)
print("Pressure drop 1 {:.5f}".format(Pdrop1))
print("Pressure drop 2 {:.5f}".format(Pdrop2))
reg_solver = RegularizationSolver(
S,
mesh,
beta=beta_param,
gamma=1e6,
dx=dx,
design_domain=DESIGN_DOMAIN,
solver_parameters=regularization_solver_parameters,
)
tol = 1e-7
dt = 0.02
params = {
"alphaC": 0.1,
"debug": 5,
"alphaJ": 0.1,
"dt": dt,
"K": 1e-3,
"maxit": opts.n_iters,
"maxtrials": 10,
"itnormalisation": 10,
"tol_merit":
1e-2, # new merit can be within 1% of the previous merit
# "normalize_tol" : -1,
"tol": tol,
}
hj_solver_parameters["ts_dt"] = dt / 50.0
solver_parameters = {
"hj_solver": hj_solver_parameters,
"reinit_solver": reinit_solver_parameters,
}
problem = InfDimProblem(
Jhat,
reg_solver,
ineqconstraints=[
Constraint(P1hat, 1.0, P1control),
Constraint(P2hat, 1.0, P2control),
],
reinit_distance=0.08,
solver_parameters=solver_parameters,
)
problem.set_termination_event(termination_event,
termination_tolerance=1e-1)
_ = nlspace_solve(problem, params)
fda.get_working_tape().clear_tape()