def __init__(
self,
functional,
controls,
level_set,
scale=1.0,
tape=None,
eval_cb_pre=lambda *args: None,
eval_cb_post=lambda *args: None,
derivative_cb_pre=lambda *args: None,
derivative_cb_post=lambda *args: None,
hessian_cb_pre=lambda *args: None,
hessian_cb_post=lambda *args: None,
):
self.functional = functional
self.cost_function = self.functional
self.tape = get_working_tape() if tape is None else tape
self.controls = Enlist(controls)
self.level_set = Enlist(level_set)
self.scale = scale
self.eval_cb_pre = eval_cb_pre
self.eval_cb_post = eval_cb_post
self.derivative_cb_pre = derivative_cb_pre
self.derivative_cb_post = derivative_cb_post
self.hessian_cb_pre = hessian_cb_pre
self.hessian_cb_post = hessian_cb_post
def __init__(self, *args, **kwargs):
super(FunctionAssigner, self).__init__(*args, **kwargs)
self.input_spaces = Enlist(args[1])
self.output_spaces = Enlist(args[0])
self.adj_assigner = backend.FunctionAssigner(args[1],
args[0],
**kwargs)
def sub(self, i, deepcopy=False, **kwargs):
from .function_assigner import FunctionAssigner, FunctionAssignerBlock
annotate = annotate_tape(kwargs)
if deepcopy:
ret = create_overloaded_object(
backend.Function.sub(self, i, deepcopy, **kwargs))
if annotate:
fa = FunctionAssigner(ret.function_space(),
self.function_space())
block = FunctionAssignerBlock(fa, Enlist(self))
tape = get_working_tape()
tape.add_block(block)
block.add_output(ret.block_variable)
else:
extra_kwargs = {}
if annotate:
extra_kwargs = {
"block_class": FunctionSplitBlock,
"_ad_floating_active": True,
"_ad_args": [self, i],
"_ad_output_args": [i],
"output_block_class": FunctionMergeBlock,
"_ad_outputs": [self],
}
ret = compat.create_function(self, i, **extra_kwargs)
return ret
def assign(self, *args, **kwargs):
annotate = annotate_tape(kwargs)
outputs = Enlist(args[0])
inputs = Enlist(args[1])
if annotate:
for i, o in enumerate(outputs):
if not isinstance(o, OverloadedType):
outputs[i] = create_overloaded_object(o)
for j, i in enumerate(outputs):
if not isinstance(i, OverloadedType):
inputs[j] = create_overloaded_object(i)
block = FunctionAssignerBlock(self, inputs)
tape = get_working_tape()
tape.add_block(block)
with stop_annotating():
ret = backend.FunctionAssigner.assign(self, outputs.delist(), inputs.delist(), **kwargs)
if annotate:
for output in outputs:
block.add_output(output.block_variable)
return ret
def __call__(self, values):
"""Computes the reduced functional with supplied control value.
Args:
values ([OverloadedType]): If you have multiple controls this should be a list of
new values for each control in the order you listed the controls to the constructor.
If you have a single control it can either be a list or a single object.
Each new value should have the same type as the corresponding control.
If values has a len(ufl_shape) > 0, we are in a Taylor test and we are updating
self.controls
If values has ufl_shape = (), it is a level set.
Returns:
:obj:`OverloadedType`: The computed value. Typically of instance
of :class:`AdjFloat`.
"""
values = Enlist(values)
if len(values) != len(self.level_set):
raise ValueError(
"values should be a list of same length as level sets.")
# Call callback.
self.eval_cb_pre(self.level_set.delist(values))
# TODO Is there a better way to do this?
if len(values[0].ufl_shape) > 0:
for i, value in enumerate(values):
self.controls[i].update(value)
else:
for i, value in enumerate(values):
self.level_set[i].block_variable.checkpoint = value
self.tape.reset_blocks()
blocks = self.tape.get_blocks()
with self.marked_controls():
with stop_annotating():
for i in range(len(blocks)):
blocks[i].recompute()
func_value = self.scale * self.functional.block_variable.checkpoint
# Call callback
self.eval_cb_post(func_value, self.level_set.delist(values))
return func_value
def __init__(self, lhs, rhs, func, bcs, *args, **kwargs):
super().__init__()
self.adj_cb = kwargs.pop("adj_cb", None)
self.adj_bdy_cb = kwargs.pop("adj_bdy_cb", None)
self.adj2_cb = kwargs.pop("adj2_cb", None)
self.adj2_bdy_cb = kwargs.pop("adj2_bdy_cb", None)
self.adj_sol = None
self.forward_args = []
self.forward_kwargs = {}
self.adj_args = []
self.adj_kwargs = {}
self.assemble_kwargs = {}
# Equation LHS
self.lhs = lhs
# Equation RHS
self.rhs = rhs
# Solution function
self.func = func
self.function_space = self.func.function_space()
# Boundary conditions
self.bcs = []
if bcs is not None:
self.bcs = Enlist(bcs)
if isinstance(self.lhs, ufl.Form) and isinstance(self.rhs, ufl.Form):
self.linear = True
for c in self.rhs.coefficients():
self.add_dependency(c, no_duplicates=True)
else:
self.linear = False
for c in self.lhs.coefficients():
self.add_dependency(c, no_duplicates=True)
for bc in self.bcs:
self.add_dependency(bc, no_duplicates=True)
if self.backend.__name__ != "firedrake":
mesh = self.lhs.ufl_domain().ufl_cargo()
else:
mesh = self.lhs.ufl_domain()
self.add_dependency(mesh)
self._init_solver_parameters(args, kwargs)
def split(self, deepcopy=False, **kwargs):
from .function_assigner import FunctionAssigner, FunctionAssignerBlock
ad_block_tag = kwargs.pop("ad_block_tag", None)
annotate = annotate_tape(kwargs)
num_sub_spaces = backend.Function.function_space(self).num_sub_spaces()
if not annotate:
if deepcopy:
ret = tuple(
create_overloaded_object(
backend.Function.sub(self, i, deepcopy, **kwargs))
for i in range(num_sub_spaces))
else:
ret = tuple(
compat.create_function(self, i)
for i in range(num_sub_spaces))
elif deepcopy:
ret = []
fs = []
for i in range(num_sub_spaces):
f = create_overloaded_object(
backend.Function.sub(self, i, deepcopy, **kwargs))
fs.append(f.function_space())
ret.append(f)
fa = FunctionAssigner(fs, self.function_space())
block = FunctionAssignerBlock(fa,
Enlist(self),
ad_block_tag=ad_block_tag)
tape = get_working_tape()
tape.add_block(block)
for output in ret:
block.add_output(output.block_variable)
ret = tuple(ret)
else:
ret = tuple(
compat.create_function(self,
i,
block_class=FunctionSplitBlock,
_ad_floating_active=True,
_ad_args=[self, i],
_ad_output_args=[i],
output_block_class=FunctionMergeBlock,
_ad_outputs=[self])
for i in range(num_sub_spaces))
return ret
def __call__(self, values):
values = Enlist(values)
if len(values) != len(self.controls):
raise ValueError(
"values should be a list of same length as controls.")
# Call callback.
self.eval_cb_pre(self.controls.delist(values))
for i, value in enumerate(values):
if isinstance(value, Function):
self.controls[
i].block_variable.checkpoint = firedrake.Function(value)
else:
self.controls[i].update(value)
self.tape.reset_blocks()
self._mark_block_variable_lifespan()
blocks = self.tape.get_blocks()
n_timesteps = self.timestep_register.n_timesteps()
self._revolve = pyrevolve.Revolve(self.n_checkpoints, n_timesteps)
with self.marked_controls(), stop_annotating():
while True:
action = self._revolve.next()
if action.type == Action.ADVANCE:
self._revolve_forward_action()
elif action.type == Action.TAKESHOT:
self._revolve_takeshot_action()
elif action.type == Action.RESTORE:
self._revolve_restore_action()
elif action.type == Action.LASTFW:
self._revolve_forward_one_timestep()
break
else:
print(action)
raise SyntaxError(
"Unknown revolve action type in forward mode")
func_value = self.scale * self.functional.block_variable.checkpoint
# Call callback
self.eval_cb_post(func_value, self.controls.delist(values))
return func_value
class LevelSetFunctional(object):
"""Class representing a Lagrangian that depends on a level set
This class is based of pyadjoint.ReducedFunctional and shares many functionalities.
The motivation is to calculate shape derivatives when we evolve a level set.
Args:
functional (:obj:`OverloadedType`): An instance of an OverloadedType,
usually :class:`AdjFloat`. This should be the return value of the
functional you want to reduce.
controls (list[Control]): A list of Control instances, which you want
to map to the functional. It is also possible to supply a single Control
instance instead of a list.
"""
def __init__(
self,
functional,
controls,
level_set,
scale=1.0,
tape=None,
eval_cb_pre=lambda *args: None,
eval_cb_post=lambda *args: None,
derivative_cb_pre=lambda *args: None,
derivative_cb_post=lambda *args: None,
hessian_cb_pre=lambda *args: None,
hessian_cb_post=lambda *args: None,
):
self.functional = functional
self.cost_function = self.functional
self.tape = get_working_tape() if tape is None else tape
self.controls = Enlist(controls)
self.level_set = Enlist(level_set)
self.scale = scale
self.eval_cb_pre = eval_cb_pre
self.eval_cb_post = eval_cb_post
self.derivative_cb_pre = derivative_cb_pre
self.derivative_cb_post = derivative_cb_post
self.hessian_cb_pre = hessian_cb_pre
self.hessian_cb_post = hessian_cb_post
# TODO Check that the level set is in the tape.
# Actually, not even pyadjoint checks if the given Control is in the
# tape.
def derivative(self, options={}):
"""Returns the derivative of the functional w.r.t. the control.
Using the adjoint method, the derivative of the functional with
respect to the control, around the last supplied value of the control,
is computed and returned.
Args:
options (dict): A dictionary of options. To find a list of available options
have a look at the specific control type.
Returns:
OverloadedType: The derivative with respect to the control.
Should be an instance of the same type as the control.
"""
# Call callback
self.derivative_cb_pre(self.level_set)
derivatives = compute_gradient(
self.functional,
self.controls,
options=options,
tape=self.tape,
adj_value=self.scale,
)
# Call callback
self.derivative_cb_post(
self.functional.block_variable.checkpoint,
self.level_set.delist(derivatives),
self.level_set,
)
return self.level_set.delist(derivatives)
@no_annotations
def hessian(self, m_dot, options={}):
"""Returns the action of the Hessian of the functional w.r.t. the control on a vector m_dot.
Using the second-order adjoint method, the action of the Hessian of the
functional with respect to the control, around the last supplied value
of the control, is computed and returned.
Args:
m_dot ([OverloadedType]): The direction in which to compute the
action of the Hessian.
options (dict): A dictionary of options. To find a list of
available options have a look at the specific control type.
Returns:
OverloadedType: The action of the Hessian in the direction m_dot.
Should be an instance of the same type as the control.
"""
# Call callback
self.hessian_cb_pre(self.level_set)
r = compute_hessian(
self.functional,
self.controls,
m_dot,
options=options,
tape=self.tape,
)
# Call callback
self.hessian_cb_post(
self.functional.block_variable.checkpoint,
self.level_set.delist(r),
self.level_set,
)
return self.level_set.delist(r)
@no_annotations
def __call__(self, values):
"""Computes the reduced functional with supplied control value.
Args:
values ([OverloadedType]): If you have multiple controls this should be a list of
new values for each control in the order you listed the controls to the constructor.
If you have a single control it can either be a list or a single object.
Each new value should have the same type as the corresponding control.
If values has a len(ufl_shape) > 0, we are in a Taylor test and we are updating
self.controls
If values has ufl_shape = (), it is a level set.
Returns:
:obj:`OverloadedType`: The computed value. Typically of instance
of :class:`AdjFloat`.
"""
values = Enlist(values)
if len(values) != len(self.level_set):
raise ValueError(
"values should be a list of same length as level sets.")
# Call callback.
self.eval_cb_pre(self.level_set.delist(values))
# TODO Is there a better way to do this?
if len(values[0].ufl_shape) > 0:
for i, value in enumerate(values):
self.controls[i].update(value)
else:
for i, value in enumerate(values):
self.level_set[i].block_variable.checkpoint = value
self.tape.reset_blocks()
blocks = self.tape.get_blocks()
with self.marked_controls():
with stop_annotating():
for i in range(len(blocks)):
blocks[i].recompute()
func_value = self.scale * self.functional.block_variable.checkpoint
# Call callback
self.eval_cb_post(func_value, self.level_set.delist(values))
return func_value
# TODO fix this to avoid deleting the level set
def optimize_tape(self):
self.tape.optimize(
controls=self.controls + self.level_set,
functionals=[self.functional],
)
def marked_controls(self):
return marked_controls(self)
def __init__(
self,
cost_function,
reg_solver,
eqconstraints=None,
ineqconstraints=None,
reinit_distance=0.05,
solver_parameters=None,
output_dir=None,
):
"""Problem interface for the null-space solver
Args:
cost_function ([type]): [description]
reg_solver ([type]): [description]
eqconstraints ([type], optional): [description]. Defaults to None.
ineqconstraints ([type], optional): [description]. Defaults to None.
reinit_distance (int, optional): The reinitialization solver is activated
after the level set is shifted reinit_distance * D,
where D is the max dimensions of a mesh
Defaults to 0.1
solver_parameters ([type], optional): [description]. Defaults to None.
Raises:
TypeError: [description]
TypeError: [description]
TypeError: [description]
Returns:
[type]: [description]
"""
if not isinstance(reg_solver, RegularizationSolver):
raise TypeError(
f"Provided regularization solver '{type(reg_solver).__name__}',\
is not a RegularizationSolver")
self.reg_solver = reg_solver
assert len(cost_function.controls) < 2, "Only one control for now"
self.phi = cost_function.level_set[0]
self.V = self.phi.function_space()
self.Vvec = cost_function.controls[0].control.function_space()
self.delta_x = fd.Function(self.Vvec)
self.max_distance = reinit_distance * max_mesh_dimension(
self.V.ufl_domain())
self.current_max_distance = self.max_distance
self.current_max_distance_at_t0 = self.current_max_distance
self.accum_distance = 0.0
self.last_distance = 0.0
self.output_dir = output_dir
self.accept_iteration = False
self.termination_event = None
V_elem = self.V.ufl_element()
if V_elem.family() in ["TensorProductElement", "Lagrange"]:
if V_elem.family() == "TensorProductElement":
assert (V_elem.sub_elements()[0].family() == "Q"
and V_elem.sub_elements()[1].family()
== "Lagrange"), "Only Lagrange basis"
self.build_cg_solvers(solver_parameters, )
else:
raise RuntimeError(
f"Level set function element {self.V.ufl_element()} not supported."
)
def event(ts, t, X, fvalue):
max_vel = calculate_max_vel(self.delta_x)
fvalue[0] = (self.accum_distance +
max_vel * t) - self.current_max_distance
def postevent(ts, events, t, X, forward):
with self.phi.dat.vec_wo as v:
X.copy(v)
self.phi.assign(self.reinit_solver.solve(self.phi))
with self.phi.dat.vec_wo as v:
v.copy(X)
self.current_max_distance += self.max_distance
direction = [1]
terminate = [False]
self.hj_solver.ts.setEventHandler(direction, terminate, event,
postevent)
self.hj_solver.ts.setEventTolerances(1e-4, vtol=[1e-4])
if eqconstraints:
self.eqconstraints = Enlist(eqconstraints)
for constr in self.eqconstraints:
if not isinstance(constr, Constraint):
raise TypeError(
f"Provided equality constraint '{type(constr).__name__}', not a Constraint"
)
else:
self.eqconstraints = []
if ineqconstraints:
self.ineqconstraints = Enlist(ineqconstraints)
for ineqconstr in self.ineqconstraints:
if not isinstance(ineqconstr, Constraint):
raise TypeError(
f"Provided inequality constraint '{type(ineqconstr).__name__}',\
not a Constraint")
else:
self.ineqconstraints = []
self.n_eqconstraints = len(self.eqconstraints)
self.n_ineqconstraints = len(self.ineqconstraints)
self.gradJ = fd.Function(self.Vvec)
self.gradH = [fd.Function(self.Vvec) for _ in self.ineqconstraints]
self.gradG = [fd.Function(self.Vvec) for _ in self.eqconstraints]
self.cost_function = cost_function
self.i = 0 # iteration count
self.beta_param = reg_solver.beta_param.values()[0]
def __init__(
self,
S,
mesh,
beta=1,
gamma=1.0e4,
bcs=None,
dx=dx,
design_domain=None,
solver_parameters=direct_parameters,
output_dir="./",
):
"""
Solver class to regularize the shape derivatives as explained in
Frédéric de Gournay
Velocity Extension for the Level-set Method and Multiple Eigenvalues in Shape Optimization
SIAM J. Control Optim., 45(1), 343–367. (25 pages)
Args:
S ([type]): Function space of the mesh coordinates
mesh ([type]): Mesh
beta ([type], optional): Regularization parameter.
It should be finite multiple of the mesh size.
Defaults to 1.
gamma ([type], optional): Penalty parameter for the penalization of the normal components
of the regularized shape derivatives on the boundary. Defaults to 1.0e4.
bcs ([type], optional): Dirichlet Boundary conditions.
They should be setting the regularized shape derivatives to zero
wherever there are boundary conditions on the original PDE.
Defaults to None.
dx ([type], optional): [description]. Defaults to dx.
design_domain ([type], optional): If we're interested in setting the shape derivatives to
zero outside of the design_domain,
we pass design_domain marker.
This is convenient when we want to fix certain regions in the domain.
Defaults to None.
solver_parameters ([type], optional): Solver options. Defaults to direct_parameters.
output_dir (str, optional): Plot the output somewhere. Defaults to "./".
"""
n = FacetNormal(mesh)
theta, xi = [TrialFunction(S), TestFunction(S)]
self.xi = xi
hmin = min_mesh_size(mesh)
if beta > 20.0 * hmin:
warning(
f"Length scale parameter beta ({beta}) is much larger than the mesh size {hmin}"
)
self.beta_param = Constant(beta)
self.a = (self.beta_param * inner(grad(theta), grad(xi)) +
inner(theta, xi)) * (dx)
if isinstance(mesh.topology, ExtrudedMeshTopology):
ds_reg = ds_b + ds_v + ds_tb + ds_t
else:
ds_reg = ds
self.a += Constant(gamma) * (inner(dot(theta, n), dot(xi, n))) * ds_reg
# Dirichlet boundary conditions equal to zero for regions where we want
# the domain to be static, i.e. zero velocities
if bcs is None:
self.bcs = []
else:
self.bcs = Enlist(bcs)
if design_domain is not None:
# Heaviside step function in domain of interest
V_DG0_B = FunctionSpace(mesh, "DG", 0)
I_B = Function(V_DG0_B)
I_B.assign(1.0)
# Set to zero all the cells within sim_domain
par_loop(
("{[i] : 0 <= i < f.dofs}", "f[i, 0] = 0.0"),
dx(design_domain),
{"f": (I_B, WRITE)},
is_loopy_kernel=True,
)
I_cg_B = Function(S)
dim = S.mesh().geometric_dimension()
# Assume that `A` is a :class:`.Function` in CG1 and `B` is a
# `.Function` in DG0. Then the following code sets each DoF in
# `A` to the maximum value that `B` attains in the cells adjacent to
# that DoF::
par_loop(
(
"{{[i, j] : 0 <= i < A.dofs and 0 <= j < {0} }}".format(
dim),
"A[i, j] = fmax(A[i, j], B[0, 0])",
),
dx,
{
"A": (I_cg_B, RW),
"B": (I_B, READ)
},
is_loopy_kernel=True,
)
import numpy as np
class MyBC(DirichletBC):
def __init__(self, V, value, markers):
# Call superclass init
# We provide a dummy subdomain id.
super(MyBC, self).__init__(V, value, 0)
# Override the "nodes" property which says where the boundary
# condition is to be applied.
self.nodes = np.unique(
np.where(markers.dat.data_ro_with_halos > 0)[0])
self.bcs.append(MyBC(S, 0, I_cg_B))
self.Av = assemble(self.a, bcs=self.bcs)
self.solver_parameters = solver_parameters
def NavierStokesBrinkmannForm(
W: fd.FunctionSpace,
w: fd.Function,
nu,
phi: Union[fd.Function, Product] = None,
brinkmann_penalty: fd.Constant = None,
brinkmann_min=0.0,
design_domain=None,
hs: Callable = hs,
beta_gls=0.9,
) -> ufl.form:
"""Returns the Galerkin Least Squares formulation for the Navier-Stokes problem with a Brinkmann term
Args:
W (fd.FunctionSpace): [description]
w (fd.Function): [description]
phi (fd.Function): [description]
nu ([type]): [description]
brinkmann_penalty ([type], optional): [description]. Defaults to None.
design_domain ([type], optional): Region where the level set is defined. Defaults to None.
Returns:
ufl.form: Nonlinear form
"""
mesh = w.ufl_domain()
W_elem = W.ufl_element()
assert isinstance(W_elem, fd.MixedElement)
if brinkmann_penalty:
assert isinstance(brinkmann_penalty, fd.Constant)
assert W_elem.num_sub_elements() == 2
for W_sub_elem in W_elem.sub_elements():
assert W_sub_elem.family() == "Lagrange"
assert W_sub_elem.degree() == 1
assert isinstance(W_elem.sub_elements()[0], fd.VectorElement)
v, q = fd.TestFunctions(W)
u, p = fd.split(w)
# Main NS form
F = (nu * inner(grad(u), grad(v)) * dx + inner(dot(grad(u), u), v) * dx -
p * div(v) * dx + div(u) * q * dx)
# Brinkmann terms for design
def add_measures(list_dd, **kwargs):
return sum((dx(dd, kwargs) for dd in list_dd[1::]), dx(list_dd[0]))
def alpha(phi):
return brinkmann_penalty * hs(phi) + fd.Constant(brinkmann_min)
if brinkmann_penalty and phi is not None:
if design_domain is not None:
dx_brinkmann = partial(add_measures, Enlist(design_domain))
else:
dx_brinkmann = dx
F = F + alpha(phi) * inner(u, v) * dx_brinkmann()
# GLS stabilization
R_U = dot(u, grad(u)) - nu * div(grad(u)) + grad(p)
if isinstance(beta_gls, (float, int)):
beta_gls = fd.Constant(beta_gls)
h = fd.CellSize(mesh)
tau_gls = beta_gls * ((4.0 * dot(u, u) / h**2) + 9.0 *
(4.0 * nu / h**2)**2)**(-0.5)
theta_U = dot(u, grad(v)) - nu * div(grad(v)) + grad(q)
F = F + tau_gls * inner(R_U, theta_U) * dx()
if brinkmann_penalty and phi is not None:
tau_gls_alpha = beta_gls * ((4.0 * dot(u, u) / h**2) + 9.0 *
(4.0 * nu / h**2)**2 +
(alpha(phi) / 1.0)**2)**(-0.5)
R_U_alpha = R_U + alpha(phi) * u
theta_alpha = theta_U + alpha(phi) * v
F = F + tau_gls_alpha * inner(R_U_alpha, theta_alpha) * dx_brinkmann()
if (design_domain is not None
): # Substract this domain from the original integral
F = F - tau_gls * inner(R_U, theta_U) * dx_brinkmann()
return F
