def test_adjoint():
cell = triangle
V1 = FiniteElement("CG", cell, 1)
V2 = FiniteElement("CG", cell, 2)
u = TrialFunction(V1)
v = TestFunction(V2)
assert u.number() > v.number()
u2 = Argument(V1, 2)
v2 = Argument(V2, 3)
assert u2.number() < v2.number()
a = u * v * dx
a_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(a)]
assert a_arg_degrees == [2, 1]
b = adjoint(a)
b_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(b)]
assert b_arg_degrees == [1, 2]
c = adjoint(a, (u2, v2))
c_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(c)]
assert c_arg_degrees == [1, 2]
d = adjoint(b)
d_arg_degrees = [arg.ufl_element().degree() for arg in extract_arguments(d)]
assert d_arg_degrees == [2, 1]
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = form.arguments()
if any(arg.part() is not None for arg in arguments):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"parts not supported")
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces (NB: Order does not
# matter anymore here because number is absolute)
v_1 = Argument(arguments[1].function_space(), arguments[0].number(),
arguments[0].part())
v_0 = Argument(arguments[0].function_space(), arguments[1].number(),
arguments[1].part())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def adjoint(form: ufl.Form, reordered_arguments=None) -> ufl.Form:
"""Compute adjoint of a bilinear form by changing the ordering (count)
of the test and trial functions.
The functions wraps ``ufl.adjoint``, and by default UFL will create new
``Argument`` s. To specify the ``Argument`` s rather than creating new ones,
pass a tuple of ``Argument`` s as ``reordered_arguments``.
See the documentation for ``ufl.adjoint`` for more details.
"""
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = form.arguments()
if len(arguments) != 2:
raise RuntimeError(
"Cannot compute adjoint of form, form is not bilinear")
if any(arg.part() is not None for arg in arguments):
raise RuntimeError(
"Cannot compute adjoint of form, parts not supported")
# Create new Arguments in the same spaces (NB: Order does not matter
# anymore here because number is absolute)
v1 = function.Argument(arguments[1].function_space,
arguments[0].number(), arguments[0].part())
v0 = function.Argument(arguments[0].function_space,
arguments[1].number(), arguments[1].part())
# Return form with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v1, v0))
def vjp_solve_eval_impl(
g: np.array,
fenics_solution: fenics.Function,
fenics_residual: ufl.Form,
fenics_inputs: List[FenicsVariable],
bcs: List[fenics.DirichletBC],
) -> Tuple[np.array]:
"""Computes the gradients of the output with respect to the inputs."""
# Convert tangent covector (adjoint) to a FEniCS variable
adj_value = numpy_to_fenics(g, fenics_solution)
adj_value = adj_value.vector()
F = fenics_residual
u = fenics_solution
V = u.function_space()
dFdu = fenics.derivative(F, u)
adFdu = ufl.adjoint(
dFdu, reordered_arguments=ufl.algorithms.extract_arguments(dFdu)
)
u_adj = fenics.Function(V)
adj_F = ufl.action(adFdu, u_adj)
adj_F = ufl.replace(adj_F, {u_adj: fenics.TrialFunction(V)})
adj_F_assembled = fenics.assemble(adj_F)
if len(bcs) != 0:
for bc in bcs:
bc.homogenize()
hbcs = bcs
for bc in hbcs:
bc.apply(adj_F_assembled)
bc.apply(adj_value)
fenics.solve(adj_F_assembled, u_adj.vector(), adj_value)
fenics_grads = []
for fenics_input in fenics_inputs:
if isinstance(fenics_input, fenics.Function):
V = fenics_input.function_space()
dFdm = fenics.derivative(F, fenics_input, fenics.TrialFunction(V))
adFdm = fenics.adjoint(dFdm)
result = fenics.assemble(-adFdm * u_adj)
if isinstance(fenics_input, fenics.Constant):
fenics_grad = fenics.Constant(result.sum())
else: # fenics.Function
fenics_grad = fenics.Function(V, result)
fenics_grads.append(fenics_grad)
# Convert FEniCS gradients to jax array representation
jax_grads = (
None if fg is None else np.asarray(fenics_to_numpy(fg)) for fg in fenics_grads
)
jax_grad_tuple = tuple(jax_grads)
return jax_grad_tuple
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = extract_arguments(form)
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py", "compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces (NB: Order matters
# here!)
v_1 = Argument(arguments[1].function_space())
v_0 = Argument(arguments[0].function_space())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = extract_arguments(form)
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces (NB: Order matters
# here!)
v_1 = Argument(arguments[1].function_space())
v_0 = Argument(arguments[0].function_space())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = form.arguments()
if any(arg.part() is not None for arg in arguments):
raise RuntimeError("Compute adjoint of form, parts not supported")
if not (len(arguments) == 2):
raise RuntimeError("Compute adjoint of form, form is not bilinear")
# Define new Argument(s) in the same spaces (NB: Order does not
# matter anymore here because number is absolute)
v_1 = Argument(arguments[1].function_space(), arguments[0].number(),
arguments[0].part())
v_0 = Argument(arguments[0].function_space(), arguments[1].number(),
arguments[1].part())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def test_adjoint():
cell = triangle
V1 = FiniteElement("CG", cell, 1)
V2 = FiniteElement("CG", cell, 2)
u = TrialFunction(V1)
v = TestFunction(V2)
assert u.number() > v.number()
u2 = Argument(V1, 2)
v2 = Argument(V2, 3)
assert u2.number() < v2.number()
a = u * v * dx
a_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(a)
]
assert a_arg_degrees == [2, 1]
b = adjoint(a)
b_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(b)
]
assert b_arg_degrees == [1, 2]
c = adjoint(a, (u2, v2))
c_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(c)
]
assert c_arg_degrees == [1, 2]
d = adjoint(b)
d_arg_degrees = [
arg.ufl_element().degree() for arg in extract_arguments(d)
]
assert d_arg_degrees == [2, 1]
def adjoint(form, reordered_arguments=None):
# Call UFL directly if new arguments are provided directly
if reordered_arguments is not None:
return ufl.adjoint(form, reordered_arguments=reordered_arguments)
# Extract form arguments
arguments = extract_arguments(form)
if any(arg.part() != None for arg in arguments):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"parts not supported")
if not (len(arguments) == 2):
cpp.dolfin_error("formmanipulation.py",
"compute adjoint of form",
"Form is not bilinear")
# Define new Argument(s) in the same spaces
# (NB: Order does not matter anymore here because number is absolute)
v_1 = Argument(arguments[1].function_space(), arguments[0].number(), arguments[0].part())
v_0 = Argument(arguments[0].function_space(), arguments[1].number(), arguments[1].part())
# Call ufl.adjoint with swapped arguments as new arguments
return ufl.adjoint(form, reordered_arguments=(v_1, v_0))
def adjoint(form, reordered_arguments=None):
"""UFL form operator:
Given a combined bilinear form, compute the adjoint form by
changing the ordering (count) of the test and trial functions.
By default, new Argument objects will be created with
opposite ordering. However, if the adjoint form is to
be added to other forms later, their arguments must match.
In that case, the user must provide a tuple reordered_arguments=(u2,v2).
"""
# ufl.adjoint creates new Arguments if no reordered_arguments is
# given. To avoid that, always pass reordered_arguments with
# firedrake.Argument objects.
if reordered_arguments is None:
v, u = extract_arguments(form)
reordered_arguments = (Argument(u.element(), u.function_space(),
count=v.count()),
Argument(v.element(), v.function_space(),
count=u.count()))
return ufl.adjoint(form, reordered_arguments)
def adjoint(form, reordered_arguments=None):
"""UFL form operator:
Given a combined bilinear form, compute the adjoint form by
changing the ordering (number) of the test and trial functions.
By default, new Argument objects will be created with
opposite ordering. However, if the adjoint form is to
be added to other forms later, their arguments must match.
In that case, the user must provide a tuple reordered_arguments=(u2,v2).
"""
# ufl.adjoint creates new Arguments if no reordered_arguments is
# given. To avoid that, always pass reordered_arguments with
# firedrake.Argument objects.
if reordered_arguments is None:
v, u = extract_arguments(form)
reordered_arguments = (Argument(u.function_space(),
number=v.number(),
part=v.part()),
Argument(v.function_space(),
number=u.number(),
part=u.part()))
return ufl.adjoint(form, reordered_arguments)
def adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a form.
:arg form: A UFL form, or a Slate tensor.
:arg reordered_arguments: arguments to use when creating the
adjoint. Ignored if form is a Slate tensor.
If the form is a slate tensor, this just returns its transpose.
Otherwise, given a bilinear form, compute the adjoint form by
changing the ordering (number) of the test and trial functions.
By default, new Argument objects will be created with opposite
ordering. However, if the adjoint form is to be added to other
forms later, their arguments must match. In that case, the user
must provide a tuple reordered_arguments=(u2,v2).
"""
if isinstance(form, firedrake.slate.TensorBase):
if reordered_arguments is not None:
firedrake.warning(
"Ignoring arguments for adjoint of Slate tensor.")
if form.rank != 2:
raise ValueError("Expecting rank-2 tensor")
return form.T
else:
if len(form.arguments()) != 2:
raise ValueError("Expecting bilinear form")
# ufl.adjoint creates new Arguments if no reordered_arguments is
# given. To avoid that, always pass reordered_arguments with
# firedrake.Argument objects.
if reordered_arguments is None:
v, u = extract_arguments(form)
reordered_arguments = (Argument(u.function_space(),
number=v.number(),
part=v.part()),
Argument(v.function_space(),
number=u.number(),
part=u.part()))
return ufl.adjoint(form, reordered_arguments)
def adjoint(form, reordered_arguments=None):
"""Compute the adjoint of a form.
:arg form: A UFL form, or a Slate tensor.
:arg reordered_arguments: arguments to use when creating the
adjoint. Ignored if form is a Slate tensor.
If the form is a slate tensor, this just returns its transpose.
Otherwise, given a bilinear form, compute the adjoint form by
changing the ordering (number) of the test and trial functions.
By default, new Argument objects will be created with opposite
ordering. However, if the adjoint form is to be added to other
forms later, their arguments must match. In that case, the user
must provide a tuple reordered_arguments=(u2,v2).
"""
if isinstance(form, firedrake.slate.TensorBase):
if reordered_arguments is not None:
firedrake.warning("Ignoring arguments for adjoint of Slate tensor.")
if form.rank != 2:
raise ValueError("Expecting rank-2 tensor")
return form.T
else:
if len(form.arguments()) != 2:
raise ValueError("Expecting bilinear form")
# ufl.adjoint creates new Arguments if no reordered_arguments is
# given. To avoid that, always pass reordered_arguments with
# firedrake.Argument objects.
if reordered_arguments is None:
v, u = extract_arguments(form)
reordered_arguments = (Argument(u.function_space(),
number=v.number(),
part=v.part()),
Argument(v.function_space(),
number=u.number(),
part=u.part()))
return ufl.adjoint(form, reordered_arguments)
#
# Author: Martin Sandve Alnes
# Date: 2008-10-30
#
from ufl import (Coefficient, FiniteElement, action, adjoint, derivative, dx,
grad, inner, triangle)
element = FiniteElement("Lagrange", triangle, 1)
w = Coefficient(element)
# H1 semi-norm
f = inner(grad(w), grad(w)) / 2 * dx
# grad(w) : grad(v)
b = derivative(f, w)
# stiffness matrix, grad(u) : grad(v)
a = derivative(b, w)
# adjoint, grad(v) : grad(u)
astar = adjoint(a)
# action of adjoint, grad(v) : grad(w)
astaraction = action(astar)
forms = [f, b, a, astar, astaraction]