# -*- coding: utf-8 -*-
"""This module contains the apply_derivatives algorithm which computes the derivatives of a form of expression."""
# Copyright (C) 2008-2016 Martin Sandve Alnæs
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier: LGPL-3.0-or-later
from collections import defaultdict
from ufl_legacy.log import error, warning
from ufl_legacy.core.expr import ufl_err_str
from ufl_legacy.core.terminal import Terminal
from ufl_legacy.core.multiindex import MultiIndex, FixedIndex, indices
from ufl_legacy.tensors import as_tensor, as_scalar, as_scalars, unit_indexed_tensor, unwrap_list_tensor
from ufl_legacy.classes import ConstantValue, Identity, Zero, FloatValue
from ufl_legacy.classes import Coefficient, FormArgument, ReferenceValue
from ufl_legacy.classes import Grad, ReferenceGrad, Variable
from ufl_legacy.classes import Indexed, ListTensor, ComponentTensor
from ufl_legacy.classes import ExprList, ExprMapping
from ufl_legacy.classes import Product, Sum, IndexSum
from ufl_legacy.classes import Conj, Real, Imag
from ufl_legacy.classes import JacobianInverse
from ufl_legacy.classes import SpatialCoordinate
from ufl_legacy.constantvalue import is_true_ufl_scalar, is_ufl_scalar
from ufl_legacy.operators import (conditional, sign,
sqrt, exp, ln, cos, sin, cosh, sinh,
bessel_J, bessel_Y, bessel_I, bessel_K,
cell_avg, facet_avg)
from math import pi
from ufl_legacy.corealg.multifunction import MultiFunction
from ufl_legacy.corealg.map_dag import map_expr_dag
from ufl_legacy.algorithms.map_integrands import map_integrand_dags
from ufl_legacy.checks import is_cellwise_constant
from ufl_legacy.differentiation import CoordinateDerivative
# TODO: Add more rulesets?
# - DivRuleset
# - CurlRuleset
# - ReferenceGradRuleset
# - ReferenceDivRuleset
# Set this to True to enable previously default workaround
# for bug in FFC handling of conditionals, uflacs does not
# have this bug.
CONDITIONAL_WORKAROUND = False
[docs]
class GenericDerivativeRuleset(MultiFunction):
def __init__(self, var_shape):
MultiFunction.__init__(self)
self._var_shape = var_shape
# --- Error checking for missing handlers and unexpected types
[docs]
def expr(self, o):
error("Missing differentiation handler for type {0}. Have you added a new type?".format(o._ufl_class_.__name__))
[docs]
def unexpected(self, o):
error("Unexpected type {0} in AD rules.".format(o._ufl_class_.__name__))
[docs]
def override(self, o):
error("Type {0} must be overridden in specialized AD rule set.".format(o._ufl_class_.__name__))
[docs]
def derivative(self, o):
error("Unhandled derivative type {0}, nested differentiation has failed.".format(o._ufl_class_.__name__))
[docs]
def fixme(self, o):
error("FIXME: Unimplemented differentiation handler for type {0}.".format(o._ufl_class_.__name__))
# --- Some types just don't have any derivative, this is just to
# --- make algorithm structure generic
[docs]
def non_differentiable_terminal(self, o):
"Labels and indices are not differentiable. It's convenient to return the non-differentiated object."
return o
label = non_differentiable_terminal
multi_index = non_differentiable_terminal
# --- Helper functions for creating zeros with the right shapes
[docs]
def independent_terminal(self, o):
"Return a zero with the right shape for terminals independent of differentiation variable."
return Zero(o.ufl_shape + self._var_shape)
[docs]
def independent_operator(self, o):
"Return a zero with the right shape and indices for operators independent of differentiation variable."
return Zero(o.ufl_shape + self._var_shape, o.ufl_free_indices, o.ufl_index_dimensions)
# --- All derivatives need to define grad and averaging
grad = override
cell_avg = override
facet_avg = override
# --- Default rules for terminals
# Literals are by definition independent of any differentiation variable
constant_value = independent_terminal
# Constants are independent of any differentiation
constant = independent_terminal
# Rules for form arguments must be specified in specialized rule set
form_argument = override
# Rules for geometric quantities must be specified in specialized rule set
geometric_quantity = override
# These types are currently assumed independent, but for non-affine domains
# this no longer holds and we want to implement rules for them.
# facet_normal = independent_terminal
# spatial_coordinate = independent_terminal
# cell_coordinate = independent_terminal
# Measures of cell entities, assuming independent although
# this will not be true for all of these for non-affine domains
# cell_volume = independent_terminal
# circumradius = independent_terminal
# facet_area = independent_terminal
# cell_surface_area = independent_terminal
# min_cell_edge_length = independent_terminal
# max_cell_edge_length = independent_terminal
# min_facet_edge_length = independent_terminal
# max_facet_edge_length = independent_terminal
# Other stuff
# cell_orientation = independent_terminal
# quadrature_weigth = independent_terminal
# These types are currently not expected to show up in AD pass.
# To make some of these available to the end-user, they need to be
# implemented here.
# facet_coordinate = unexpected
# cell_origin = unexpected
# facet_origin = unexpected
# cell_facet_origin = unexpected
# jacobian = unexpected
# jacobian_determinant = unexpected
# jacobian_inverse = unexpected
# facet_jacobian = unexpected
# facet_jacobian_determinant = unexpected
# facet_jacobian_inverse = unexpected
# cell_facet_jacobian = unexpected
# cell_facet_jacobian_determinant = unexpected
# cell_facet_jacobian_inverse = unexpected
# cell_vertices = unexpected
# cell_edge_vectors = unexpected
# facet_edge_vectors = unexpected
# reference_cell_edge_vectors = unexpected
# reference_facet_edge_vectors = unexpected
# cell_normal = unexpected # TODO: Expecting rename
# cell_normals = unexpected
# facet_tangents = unexpected
# cell_tangents = unexpected
# cell_midpoint = unexpected
# facet_midpoint = unexpected
# --- Default rules for operators
[docs]
def variable(self, o, df, unused_l):
return df
# --- Indexing and component handling
[docs]
def indexed(self, o, Ap, ii): # TODO: (Partially) duplicated in nesting rules
# Propagate zeros
if isinstance(Ap, Zero):
return self.independent_operator(o)
# Untangle as_tensor(C[kk], jj)[ii] -> C[ll] to simplify
# resulting expression
if isinstance(Ap, ComponentTensor):
B, jj = Ap.ufl_operands
if isinstance(B, Indexed):
C, kk = B.ufl_operands
kk = list(kk)
if all(j in kk for j in jj):
rep = dict(zip(jj, ii))
Cind = [rep.get(k, k) for k in kk]
expr = Indexed(C, MultiIndex(tuple(Cind)))
assert expr.ufl_free_indices == o.ufl_free_indices
assert expr.ufl_shape == o.ufl_shape
return expr
# Otherwise a more generic approach
r = len(Ap.ufl_shape) - len(ii)
if r:
kk = indices(r)
op = Indexed(Ap, MultiIndex(ii.indices() + kk))
op = as_tensor(op, kk)
else:
op = Indexed(Ap, ii)
return op
[docs]
def list_tensor(self, o, *dops):
return ListTensor(*dops)
[docs]
def component_tensor(self, o, Ap, ii):
if isinstance(Ap, Zero):
op = self.independent_operator(o)
else:
Ap, jj = as_scalar(Ap)
op = as_tensor(Ap, ii.indices() + jj)
return op
# --- Algebra operators
[docs]
def index_sum(self, o, Ap, i):
return IndexSum(Ap, i)
[docs]
def sum(self, o, da, db):
return da + db
[docs]
def product(self, o, da, db):
# Even though arguments to o are scalar, da and db may be
# tensor valued
a, b = o.ufl_operands
(da, db), ii = as_scalars(da, db)
pa = Product(da, b)
pb = Product(a, db)
s = Sum(pa, pb)
if ii:
s = as_tensor(s, ii)
return s
[docs]
def division(self, o, fp, gp):
f, g = o.ufl_operands
if not is_ufl_scalar(f):
error("Not expecting nonscalar nominator")
if not is_true_ufl_scalar(g):
error("Not expecting nonscalar denominator")
# do_df = 1/g
# do_dg = -h/g
# op = do_df*fp + do_df*gp
# op = (fp - o*gp) / g
# Get o and gp as scalars, multiply, then wrap as a tensor
# again
so, oi = as_scalar(o)
sgp, gi = as_scalar(gp)
o_gp = so * sgp
if oi or gi:
o_gp = as_tensor(o_gp, oi + gi)
op = (fp - o_gp) / g
return op
[docs]
def power(self, o, fp, gp):
f, g = o.ufl_operands
if not is_true_ufl_scalar(f):
error("Expecting scalar expression f in f**g.")
if not is_true_ufl_scalar(g):
error("Expecting scalar expression g in f**g.")
# Derivation of the general case: o = f(x)**g(x)
# do/df = g * f**(g-1) = g / f * o
# do/dg = ln(f) * f**g = ln(f) * o
# do/df * df + do/dg * dg = o * (g / f * df + ln(f) * dg)
if isinstance(gp, Zero):
# This probably produces better results for the common
# case of f**constant
op = fp * g * f**(g - 1)
else:
# Note: This produces expressions like (1/w)*w**5 instead of w**4
# op = o * (fp * g / f + gp * ln(f)) # This reuses o
op = f**(g - 1) * (g * fp + f * ln(f) * gp) # This gives better accuracy in dolfin integration test
# Example: d/dx[x**(x**3)]:
# f = x
# g = x**3
# df = 1
# dg = 3*x**2
# op1 = o * (fp * g / f + gp * ln(f))
# = x**(x**3) * (x**3/x + 3*x**2*ln(x))
# op2 = f**(g-1) * (g*fp + f*ln(f)*gp)
# = x**(x**3-1) * (x**3 + x*3*x**2*ln(x))
return op
[docs]
def abs(self, o, df):
f, = o.ufl_operands
# return conditional(eq(f, 0), 0, Product(sign(f), df)) abs is
# not complex differentiable, so we workaround the case of a
# real F in complex mode by defensively casting to real inside
# the sign.
return sign(Real(f)) * df
# --- Complex algebra
[docs]
def conj(self, o, df):
return Conj(df)
[docs]
def real(self, o, df):
return Real(df)
[docs]
def imag(self, o, df):
return Imag(df)
# --- Mathfunctions
[docs]
def math_function(self, o, df):
# FIXME: Introduce a UserOperator type instead of this hack
# and define user derivative() function properly
if hasattr(o, 'derivative'):
f, = o.ufl_operands
return df * o.derivative()
error("Unknown math function.")
[docs]
def sqrt(self, o, fp):
return fp / (2 * o)
[docs]
def exp(self, o, fp):
return fp * o
[docs]
def ln(self, o, fp):
f, = o.ufl_operands
if isinstance(f, Zero):
error("Division by zero.")
return fp / f
[docs]
def cos(self, o, fp):
f, = o.ufl_operands
return fp * -sin(f)
[docs]
def sin(self, o, fp):
f, = o.ufl_operands
return fp * cos(f)
[docs]
def tan(self, o, fp):
f, = o.ufl_operands
return 2.0 * fp / (cos(2.0 * f) + 1.0)
[docs]
def cosh(self, o, fp):
f, = o.ufl_operands
return fp * sinh(f)
[docs]
def sinh(self, o, fp):
f, = o.ufl_operands
return fp * cosh(f)
[docs]
def tanh(self, o, fp):
f, = o.ufl_operands
def sech(y):
return (2.0 * cosh(y)) / (cosh(2.0 * y) + 1.0)
return fp * sech(f)**2
[docs]
def acos(self, o, fp):
f, = o.ufl_operands
return -fp / sqrt(1.0 - f**2)
[docs]
def asin(self, o, fp):
f, = o.ufl_operands
return fp / sqrt(1.0 - f**2)
[docs]
def atan(self, o, fp):
f, = o.ufl_operands
return fp / (1.0 + f**2)
[docs]
def atan_2(self, o, fp, gp):
f, g = o.ufl_operands
return (g * fp - f * gp) / (f**2 + g**2)
[docs]
def erf(self, o, fp):
f, = o.ufl_operands
return fp * (2.0 / sqrt(pi) * exp(-f**2))
# --- Bessel functions
[docs]
def bessel_j(self, o, nup, fp):
nu, f = o.ufl_operands
if not (nup is None or isinstance(nup, Zero)):
error("Differentiation of bessel function w.r.t. nu is not supported.")
if isinstance(nu, Zero):
op = -bessel_J(1, f)
else:
op = 0.5 * (bessel_J(nu - 1, f) - bessel_J(nu + 1, f))
return op * fp
[docs]
def bessel_y(self, o, nup, fp):
nu, f = o.ufl_operands
if not (nup is None or isinstance(nup, Zero)):
error("Differentiation of bessel function w.r.t. nu is not supported.")
if isinstance(nu, Zero):
op = -bessel_Y(1, f)
else:
op = 0.5 * (bessel_Y(nu - 1, f) - bessel_Y(nu + 1, f))
return op * fp
[docs]
def bessel_i(self, o, nup, fp):
nu, f = o.ufl_operands
if not (nup is None or isinstance(nup, Zero)):
error("Differentiation of bessel function w.r.t. nu is not supported.")
if isinstance(nu, Zero):
op = bessel_I(1, f)
else:
op = 0.5 * (bessel_I(nu - 1, f) + bessel_I(nu + 1, f))
return op * fp
[docs]
def bessel_k(self, o, nup, fp):
nu, f = o.ufl_operands
if not (nup is None or isinstance(nup, Zero)):
error("Differentiation of bessel function w.r.t. nu is not supported.")
if isinstance(nu, Zero):
op = -bessel_K(1, f)
else:
op = -0.5 * (bessel_K(nu - 1, f) + bessel_K(nu + 1, f))
return op * fp
# --- Restrictions
[docs]
def restricted(self, o, fp):
# Restriction and differentiation commutes
if isinstance(fp, ConstantValue):
return fp # TODO: Add simplification to Restricted instead?
else:
return fp(o._side) # (f+-)' == (f')+-
# --- Conditionals
[docs]
def binary_condition(self, o, dl, dr):
# Should not be used anywhere...
return None
[docs]
def not_condition(self, o, c):
# Should not be used anywhere...
return None
[docs]
def conditional(self, o, unused_dc, dt, df):
global CONDITIONAL_WORKAROUND
if isinstance(dt, Zero) and isinstance(df, Zero):
# Assuming dt and df have the same indices here, which
# should be the case
return dt
elif CONDITIONAL_WORKAROUND:
# Placing t[1],f[1] outside here to avoid getting
# arguments inside conditionals. This will fail when dt
# or df become NaN or Inf in floating point computations!
c = o.ufl_operands[0]
dc = conditional(c, 1, 0)
return dc * dt + (1.0 - dc) * df
else:
# Not placing t[1],f[1] outside, allowing arguments inside
# conditionals. This will make legacy ffc fail, but
# should work with uflacs.
c = o.ufl_operands[0]
return conditional(c, dt, df)
[docs]
def max_value(self, o, df, dg):
# d/dx max(f, g) =
# f > g: df/dx
# f < g: dg/dx
# Placing df,dg outside here to avoid getting arguments inside
# conditionals
f, g = o.ufl_operands
dc = conditional(f > g, 1, 0)
return dc * df + (1.0 - dc) * dg
[docs]
def min_value(self, o, df, dg):
# d/dx min(f, g) =
# f < g: df/dx
# else: dg/dx
# Placing df,dg outside here to avoid getting arguments
# inside conditionals
f, g = o.ufl_operands
dc = conditional(f < g, 1, 0)
return dc * df + (1.0 - dc) * dg
[docs]
class GradRuleset(GenericDerivativeRuleset):
def __init__(self, geometric_dimension):
GenericDerivativeRuleset.__init__(self, var_shape=(geometric_dimension,))
self._Id = Identity(geometric_dimension)
# --- Specialized rules for geometric quantities
[docs]
def geometric_quantity(self, o):
"""Default for geometric quantities is do/dx = 0 if piecewise constant,
otherwise transform derivatives to reference derivatives.
Override for specific types if other behaviour is needed."""
if is_cellwise_constant(o):
return self.independent_terminal(o)
else:
domain = o.ufl_domain()
K = JacobianInverse(domain)
Do = grad_to_reference_grad(o, K)
return Do
[docs]
def jacobian_inverse(self, o):
# grad(K) == K_ji rgrad(K)_rj
if is_cellwise_constant(o):
return self.independent_terminal(o)
if not o._ufl_is_terminal_:
error("ReferenceValue can only wrap a terminal")
Do = grad_to_reference_grad(o, o)
return Do
# TODO: Add more explicit geometry type handlers here, with
# non-affine domains several should be non-zero.
[docs]
def spatial_coordinate(self, o):
"dx/dx = I"
return self._Id
[docs]
def cell_coordinate(self, o):
"dX/dx = inv(dx/dX) = inv(J) = K"
# FIXME: Is this true for manifolds? What about orientation?
return JacobianInverse(o.ufl_domain())
# --- Specialized rules for form arguments
[docs]
def coefficient(self, o):
if is_cellwise_constant(o):
return self.independent_terminal(o)
return Grad(o)
[docs]
def argument(self, o):
# TODO: Enable this after fixing issue#13, unless we move
# simplificat ion to a separate stage?
# if is_cellwise_constant(o):
# # Collapse gradient of cellwise constant function to zero
# # TODO: Missing this type
# return AnnotatedZero(o.ufl_shape + self._var_shape, arguments=(o,))
return Grad(o)
# --- Rules for values or derivatives in reference frame
[docs]
def reference_value(self, o):
# grad(o) == grad(rv(f)) -> K_ji*rgrad(rv(f))_rj
f = o.ufl_operands[0]
if f.ufl_element().mapping() == "physical":
# TODO: Do we need to be more careful for immersed things?
return ReferenceGrad(o)
if not f._ufl_is_terminal_:
error("ReferenceValue can only wrap a terminal")
domain = f.ufl_domain()
K = JacobianInverse(domain)
Do = grad_to_reference_grad(o, K)
return Do
[docs]
def reference_grad(self, o):
# grad(o) == grad(rgrad(rv(f))) -> K_ji*rgrad(rgrad(rv(f)))_rj
f = o.ufl_operands[0]
valid_operand = f._ufl_is_in_reference_frame_ or isinstance(f, (JacobianInverse, SpatialCoordinate))
if not valid_operand:
error("ReferenceGrad can only wrap a reference frame type!")
domain = f.ufl_domain()
K = JacobianInverse(domain)
Do = grad_to_reference_grad(o, K)
return Do
# --- Nesting of gradients
[docs]
def grad(self, o):
"Represent grad(grad(f)) as Grad(Grad(f))."
# Check that o is a "differential terminal"
if not isinstance(o.ufl_operands[0], (Grad, Terminal)):
error("Expecting only grads applied to a terminal.")
return Grad(o)
def _grad(self, o):
pass
# TODO: Not sure how to detect that gradient of f is cellwise constant.
# Can we trust element degrees?
# if is_cellwise_constant(o):
# return self.terminal(o)
# TODO: Maybe we can ask "f.has_derivatives_of_order(n)" to check
# if we should make a zero here?
# 1) n = count number of Grads, get f
# 2) if not f.has_derivatives(n): return zero(...)
cell_avg = GenericDerivativeRuleset.independent_operator
facet_avg = GenericDerivativeRuleset.independent_operator
[docs]
def grad_to_reference_grad(o, K):
"""Relates grad(o) to reference_grad(o) using the Jacobian inverse.
Args
----
o: Operand
K: Jacobian inverse
Returns
-------
Do: grad(o) written in terms of reference_grad(o) and K
"""
r = indices(len(o.ufl_shape))
i, j = indices(2)
# grad(o) == K_ji rgrad(o)_rj
Do = as_tensor(K[j, i] * ReferenceGrad(o)[r + (j,)], r + (i,))
return Do
[docs]
class ReferenceGradRuleset(GenericDerivativeRuleset):
def __init__(self, topological_dimension):
GenericDerivativeRuleset.__init__(self,
var_shape=(topological_dimension,))
self._Id = Identity(topological_dimension)
# --- Specialized rules for geometric quantities
[docs]
def geometric_quantity(self, o):
"dg/dX = 0 if piecewise constant, otherwise ReferenceGrad(g)"
if is_cellwise_constant(o):
return self.independent_terminal(o)
else:
# TODO: Which types does this involve? I don't think the
# form compilers will handle this.
return ReferenceGrad(o)
[docs]
def spatial_coordinate(self, o):
"dx/dX = J"
# Don't convert back to J, otherwise we get in a loop
return ReferenceGrad(o)
[docs]
def cell_coordinate(self, o):
"dX/dX = I"
return self._Id
# TODO: Add more geometry types here, with non-affine domains
# several should be non-zero.
# --- Specialized rules for form arguments
[docs]
def reference_value(self, o):
if not o.ufl_operands[0]._ufl_is_terminal_:
error("ReferenceValue can only wrap a terminal")
return ReferenceGrad(o)
[docs]
def coefficient(self, o):
error("Coefficient should be wrapped in ReferenceValue by now")
[docs]
def argument(self, o):
error("Argument should be wrapped in ReferenceValue by now")
# --- Nesting of gradients
[docs]
def grad(self, o):
error("Grad should have been transformed by this point, but got {0}.".format(type(o).__name__))
[docs]
def reference_grad(self, o):
"Represent ref_grad(ref_grad(f)) as RefGrad(RefGrad(f))."
# Check that o is a "differential terminal"
if not isinstance(o.ufl_operands[0],
(ReferenceGrad, ReferenceValue, Terminal)):
error("Expecting only grads applied to a terminal.")
return ReferenceGrad(o)
cell_avg = GenericDerivativeRuleset.independent_operator
facet_avg = GenericDerivativeRuleset.independent_operator
[docs]
class VariableRuleset(GenericDerivativeRuleset):
def __init__(self, var):
GenericDerivativeRuleset.__init__(self, var_shape=var.ufl_shape)
if var.ufl_free_indices:
error("Differentiation variable cannot have free indices.")
self._variable = var
self._Id = self._make_identity(self._var_shape)
def _make_identity(self, sh):
"Create a higher order identity tensor to represent dv/dv."
res = None
if sh == ():
# Scalar dv/dv is scalar
return FloatValue(1.0)
elif len(sh) == 1:
# Vector v makes dv/dv the identity matrix
return Identity(sh[0])
else:
# TODO: Add a type for this higher order identity?
# II[i0,i1,i2,j0,j1,j2] = 1 if all((i0==j0, i1==j1, i2==j2)) else 0
# Tensor v makes dv/dv some kind of higher rank identity tensor
ind1 = ()
ind2 = ()
for d in sh:
i, j = indices(2)
dij = Identity(d)[i, j]
if res is None:
res = dij
else:
res *= dij
ind1 += (i,)
ind2 += (j,)
fp = as_tensor(res, ind1 + ind2)
return fp
# Explicitly defining dg/dw == 0
geometric_quantity = GenericDerivativeRuleset.independent_terminal
# Explicitly defining da/dw == 0
argument = GenericDerivativeRuleset.independent_terminal
# def _argument(self, o):
# return AnnotatedZero(o.ufl_shape + self._var_shape, arguments=(o,)) # TODO: Missing this type
[docs]
def coefficient(self, o):
"""df/dv = Id if v is f else 0.
Note that if v = variable(f), df/dv is still 0,
but if v == f, i.e. isinstance(v, Coefficient) == True,
then df/dv == df/df = Id.
"""
v = self._variable
if isinstance(v, Coefficient) and o == v:
# dv/dv = identity of rank 2*rank(v)
return self._Id
else:
# df/v = 0
return self.independent_terminal(o)
[docs]
def variable(self, o, df, l):
v = self._variable
if isinstance(v, Variable) and v.label() == l:
# dv/dv = identity of rank 2*rank(v)
return self._Id
else:
# df/v = df
return df
[docs]
def grad(self, o):
"Variable derivative of a gradient of a terminal must be 0."
# Check that o is a "differential terminal"
if not isinstance(o.ufl_operands[0], (Grad, Terminal)):
error("Expecting only grads applied to a terminal.")
return self.independent_terminal(o)
# --- Rules for values or derivatives in reference frame
[docs]
def reference_value(self, o):
# d/dv(o) == d/dv(rv(f)) = 0 if v is not f, or rv(dv/df)
v = self._variable
if isinstance(v, Coefficient) and o.ufl_operands[0] == v:
if v.ufl_element().mapping() != "identity":
# FIXME: This is a bit tricky, instead of Identity it is
# actually inverse(transform), or we should rather not
# convert to reference frame in the first place
error("Missing implementation: To handle derivatives of rv(f) w.r.t. f for" +
" mapped elements, rewriting to reference frame should not happen first...")
# dv/dv = identity of rank 2*rank(v)
return self._Id
else:
# df/v = 0
return self.independent_terminal(o)
[docs]
def reference_grad(self, o):
"Variable derivative of a gradient of a terminal must be 0."
if not isinstance(o.ufl_operands[0],
(ReferenceGrad, ReferenceValue)):
error("Unexpected argument to reference_grad.")
return self.independent_terminal(o)
cell_avg = GenericDerivativeRuleset.independent_operator
facet_avg = GenericDerivativeRuleset.independent_operator
[docs]
class GateauxDerivativeRuleset(GenericDerivativeRuleset):
"""Apply AFD (Automatic Functional Differentiation) to expression.
Implements rules for the Gateaux derivative D_w[v](...) defined as
D_w[v](e) = d/dtau e(w+tau v)|tau=0
"""
def __init__(self, coefficients, arguments, coefficient_derivatives):
GenericDerivativeRuleset.__init__(self, var_shape=())
# Type checking
if not isinstance(coefficients, ExprList):
error("Expecting a ExprList of coefficients.")
if not isinstance(arguments, ExprList):
error("Expecting a ExprList of arguments.")
if not isinstance(coefficient_derivatives, ExprMapping):
error("Expecting a coefficient-coefficient ExprMapping.")
# The coefficient(s) to differentiate w.r.t. and the
# argument(s) s.t. D_w[v](e) = d/dtau e(w+tau v)|tau=0
self._w = coefficients.ufl_operands
self._v = arguments.ufl_operands
self._w2v = {w: v for w, v in zip(self._w, self._v)}
# Build more convenient dict {f: df/dw} for each coefficient f
# where df/dw is nonzero
cd = coefficient_derivatives.ufl_operands
self._cd = {cd[2 * i]: cd[2 * i + 1] for i in range(len(cd) // 2)}
# Explicitly defining dg/dw == 0
geometric_quantity = GenericDerivativeRuleset.independent_terminal
[docs]
def cell_avg(self, o, fp):
# Cell average of a single function and differentiation
# commutes, D_f[v](cell_avg(f)) = cell_avg(v)
return cell_avg(fp)
[docs]
def facet_avg(self, o, fp):
# Facet average of a single function and differentiation
# commutes, D_f[v](facet_avg(f)) = facet_avg(v)
return facet_avg(fp)
# Explicitly defining da/dw == 0
argument = GenericDerivativeRuleset.independent_terminal
[docs]
def coefficient(self, o):
# Define dw/dw := d/ds [w + s v] = v
# Return corresponding argument if we can find o among w
do = self._w2v.get(o)
if do is not None:
return do
# Look for o among coefficient derivatives
dos = self._cd.get(o)
if dos is None:
# If o is not among coefficient derivatives, return
# do/dw=0
do = Zero(o.ufl_shape)
return do
else:
# Compute do/dw_j = do/dw_h : v.
# Since we may actually have a tuple of oprimes and vs in a
# 'mixed' space, sum over them all to get the complete inner
# product. Using indices to define a non-compound inner product.
# Example:
# (f:g) -> (dfdu:v):g + f:(dgdu:v)
# shape(dfdu) == shape(f) + shape(v)
# shape(f) == shape(g) == shape(dfdu : v)
# Make sure we have a tuple to match the self._v tuple
if not isinstance(dos, tuple):
dos = (dos,)
if len(dos) != len(self._v):
error("Got a tuple of arguments, expecting a matching tuple of coefficient derivatives.")
dosum = Zero(o.ufl_shape)
for do, v in zip(dos, self._v):
so, oi = as_scalar(do)
rv = len(v.ufl_shape)
oi1 = oi[:-rv]
oi2 = oi[-rv:]
prod = so * v[oi2]
if oi1:
dosum += as_tensor(prod, oi1)
else:
dosum += prod
return dosum
[docs]
def reference_value(self, o):
error("Currently no support for ReferenceValue in CoefficientDerivative.")
# TODO: This is implementable for regular derivative(M(f),f,v)
# but too messy if customized coefficient derivative
# relations are given by the user. We would only need
# this to allow the user to write
# derivative(...ReferenceValue...,...).
# f, = o.ufl_operands
# if not f._ufl_is_terminal_:
# error("ReferenceValue can only wrap terminals directly.")
# FIXME: check all cases like in coefficient
# if f is w:
# # FIXME: requires that v is an Argument with the same element mapping!
# return ReferenceValue(v)
# else:
# return self.independent_terminal(o)
[docs]
def reference_grad(self, o):
error("Currently no support for ReferenceGrad in CoefficientDerivative.")
# TODO: This is implementable for regular derivative(M(f),f,v)
# but too messy if customized coefficient derivative
# relations are given by the user. We would only need
# this to allow the user to write
# derivative(...ReferenceValue...,...).
[docs]
def grad(self, g):
# If we hit this type, it has already been propagated to a
# coefficient (or grad of a coefficient), # FIXME: Assert
# this! so we need to take the gradient of the variation or
# return zero. Complications occur when dealing with
# derivatives w.r.t. single components...
# Figure out how many gradients are around the inner terminal
ngrads = 0
o = g
while isinstance(o, Grad):
o, = o.ufl_operands
ngrads += 1
if not isinstance(o, FormArgument):
error("Expecting gradient of a FormArgument, not %s" % ufl_err_str(o))
def apply_grads(f):
for i in range(ngrads):
f = Grad(f)
return f
# Find o among all w without any indexing, which makes this
# easy
for (w, v) in zip(self._w, self._v):
if o == w and isinstance(v, FormArgument):
# Case: d/dt [w + t v]
return apply_grads(v)
# If o is not among coefficient derivatives, return do/dw=0
gprimesum = Zero(g.ufl_shape)
def analyse_variation_argument(v):
# Analyse variation argument
if isinstance(v, FormArgument):
# Case: d/dt [w[...] + t v]
vval, vcomp = v, ()
elif isinstance(v, Indexed):
# Case: d/dt [w + t v[...]]
# Case: d/dt [w[...] + t v[...]]
vval, vcomp = v.ufl_operands
vcomp = tuple(vcomp)
else:
error("Expecting argument or component of argument.")
if not all(isinstance(k, FixedIndex) for k in vcomp):
error("Expecting only fixed indices in variation.")
return vval, vcomp
def compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp):
# Apply gradients directly to argument vval, and get the
# right indexed scalar component(s)
kk = indices(ngrads)
Dvkk = apply_grads(vval)[vcomp + kk]
# Place scalar component(s) Dvkk into the right tensor
# positions
if wshape:
Ejj, jj = unit_indexed_tensor(wshape, wcomp)
else:
Ejj, jj = 1, ()
gprimeterm = as_tensor(Ejj * Dvkk, jj + kk)
return gprimeterm
# Accumulate contributions from variations in different
# components
for (w, v) in zip(self._w, self._v):
# Analyse differentiation variable coefficient
if isinstance(w, FormArgument):
if not w == o:
continue
wshape = w.ufl_shape
if isinstance(v, FormArgument):
# Case: d/dt [w + t v]
return apply_grads(v)
elif isinstance(v, ListTensor):
# Case: d/dt [w + t <...,v,...>]
for wcomp, vsub in unwrap_list_tensor(v):
if not isinstance(vsub, Zero):
vval, vcomp = analyse_variation_argument(vsub)
gprimesum = gprimesum + compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp)
else:
if wshape != ():
error("Expecting scalar coefficient in this branch.")
# Case: d/dt [w + t v[...]]
wval, wcomp = w, ()
vval, vcomp = analyse_variation_argument(v)
gprimesum = gprimesum + compute_gprimeterm(ngrads, vval,
vcomp, wshape,
wcomp)
elif isinstance(w, Indexed): # This path is tested in unit tests, but not actually used?
# Case: d/dt [w[...] + t v[...]]
# Case: d/dt [w[...] + t v]
wval, wcomp = w.ufl_operands
if not wval == o:
continue
assert isinstance(wval, FormArgument)
if not all(isinstance(k, FixedIndex) for k in wcomp):
error("Expecting only fixed indices in differentiation variable.")
wshape = wval.ufl_shape
vval, vcomp = analyse_variation_argument(v)
gprimesum = gprimesum + compute_gprimeterm(ngrads, vval, vcomp, wshape, wcomp)
else:
error("Expecting coefficient or component of coefficient.")
# FIXME: Handle other coefficient derivatives: oprimes =
# self._cd.get(o)
if 0:
oprimes = self._cd.get(o)
if oprimes is None:
if self._cd:
# TODO: Make it possible to silence this message
# in particular? It may be good to have for
# debugging...
warning("Assuming d{%s}/d{%s} = 0." % (o, self._w))
else:
# Make sure we have a tuple to match the self._v tuple
if not isinstance(oprimes, tuple):
oprimes = (oprimes,)
if len(oprimes) != len(self._v):
error("Got a tuple of arguments, expecting a"
" matching tuple of coefficient derivatives.")
# Compute dg/dw_j = dg/dw_h : v.
# Since we may actually have a tuple of oprimes and vs
# in a 'mixed' space, sum over them all to get the
# complete inner product. Using indices to define a
# non-compound inner product.
for (oprime, v) in zip(oprimes, self._v):
error("FIXME: Figure out how to do this with ngrads")
so, oi = as_scalar(oprime)
rv = len(v.ufl_shape)
oi1 = oi[:-rv]
oi2 = oi[-rv:]
prod = so * v[oi2]
if oi1:
gprimesum += as_tensor(prod, oi1)
else:
gprimesum += prod
return gprimesum
[docs]
def coordinate_derivative(self, o):
o = o.ufl_operands
return CoordinateDerivative(map_expr_dag(self, o[0]), o[1], o[2], o[3])
[docs]
class DerivativeRuleDispatcher(MultiFunction):
def __init__(self):
MultiFunction.__init__(self)
# caches for reuse in the dispatched transformers
self.vcaches = defaultdict(dict)
self.rcaches = defaultdict(dict)
[docs]
def terminal(self, o):
return o
[docs]
def derivative(self, o):
error("Missing derivative handler for {0}.".format(type(o).__name__))
expr = MultiFunction.reuse_if_untouched
[docs]
def grad(self, o, f):
rules = GradRuleset(o.ufl_shape[-1])
key = (GradRuleset, o.ufl_shape[-1])
return map_expr_dag(rules, f,
vcache=self.vcaches[key],
rcache=self.rcaches[key])
[docs]
def reference_grad(self, o, f):
rules = ReferenceGradRuleset(o.ufl_shape[-1]) # FIXME: Look over this and test better.
key = (ReferenceGradRuleset, o.ufl_shape[-1])
return map_expr_dag(rules, f,
vcache=self.vcaches[key],
rcache=self.rcaches[key])
[docs]
def variable_derivative(self, o, f, dummy_v):
op = o.ufl_operands[1]
rules = VariableRuleset(op)
key = (VariableRuleset, op)
return map_expr_dag(rules, f,
vcache=self.vcaches[key],
rcache=self.rcaches[key])
[docs]
def coefficient_derivative(self, o, f, dummy_w, dummy_v, dummy_cd):
dummy, w, v, cd = o.ufl_operands
rules = GateauxDerivativeRuleset(w, v, cd)
key = (GateauxDerivativeRuleset, w, v, cd)
return map_expr_dag(rules, f,
vcache=self.vcaches[key],
rcache=self.rcaches[key])
[docs]
def coordinate_derivative(self, o, f, dummy_w, dummy_v, dummy_cd):
o_ = o.ufl_operands
key = (CoordinateDerivative, o_[0])
return CoordinateDerivative(map_expr_dag(self, o_[0],
vcache=self.vcaches[key],
rcache=self.rcaches[key]),
o_[1], o_[2], o_[3])
[docs]
def indexed(self, o, Ap, ii): # TODO: (Partially) duplicated in generic rules
# Reuse if untouched
if Ap is o.ufl_operands[0]:
return o
# Untangle as_tensor(C[kk], jj)[ii] -> C[ll] to simplify
# resulting expression
if isinstance(Ap, ComponentTensor):
B, jj = Ap.ufl_operands
if isinstance(B, Indexed):
C, kk = B.ufl_operands
kk = list(kk)
if all(j in kk for j in jj):
rep = dict(zip(jj, ii))
Cind = [rep.get(k, k) for k in kk]
expr = Indexed(C, MultiIndex(tuple(Cind)))
assert expr.ufl_free_indices == o.ufl_free_indices
assert expr.ufl_shape == o.ufl_shape
return expr
# Otherwise a more generic approach
r = len(Ap.ufl_shape) - len(ii)
if r:
kk = indices(r)
op = Indexed(Ap, MultiIndex(ii.indices() + kk))
op = as_tensor(op, kk)
else:
op = Indexed(Ap, ii)
return op
[docs]
def apply_derivatives(expression):
rules = DerivativeRuleDispatcher()
return map_integrand_dags(rules, expression)
[docs]
class CoordinateDerivativeRuleset(GenericDerivativeRuleset):
"""Apply AFD (Automatic Functional Differentiation) to expression.
Implements rules for the Gateaux derivative D_w[v](...) defined as
D_w[v](e) = d/dtau e(w+tau v)|tau=0
where 'e' is a ufl form after pullback and w is a SpatialCoordinate.
"""
def __init__(self, coefficients, arguments, coefficient_derivatives):
GenericDerivativeRuleset.__init__(self, var_shape=())
# Type checking
if not isinstance(coefficients, ExprList):
error("Expecting a ExprList of coefficients.")
if not isinstance(arguments, ExprList):
error("Expecting a ExprList of arguments.")
if not isinstance(coefficient_derivatives, ExprMapping):
error("Expecting a coefficient-coefficient ExprMapping.")
# The coefficient(s) to differentiate w.r.t. and the
# argument(s) s.t. D_w[v](e) = d/dtau e(w+tau v)|tau=0
self._w = coefficients.ufl_operands
self._v = arguments.ufl_operands
self._w2v = {w: v for w, v in zip(self._w, self._v)}
# Build more convenient dict {f: df/dw} for each coefficient f
# where df/dw is nonzero
cd = coefficient_derivatives.ufl_operands
self._cd = {cd[2 * i]: cd[2 * i + 1] for i in range(len(cd) // 2)}
# Explicitly defining dg/dw == 0
geometric_quantity = GenericDerivativeRuleset.independent_terminal
# Explicitly defining da/dw == 0
argument = GenericDerivativeRuleset.independent_terminal
[docs]
def coefficient(self, o):
error("CoordinateDerivative of coefficient in physical space is not implemented.")
[docs]
def grad(self, o):
error("CoordinateDerivative grad in physical space is not implemented.")
[docs]
def spatial_coordinate(self, o):
do = self._w2v.get(o)
# d x /d x => Argument(x.function_space())
if do is not None:
return do
else:
error("Not implemented: CoordinateDerivative found a SpatialCoordinate that is different from the one being differentiated.")
[docs]
def reference_value(self, o):
do = self._cd.get(o)
if do is not None:
return do
else:
return self.independent_terminal(o)
[docs]
def reference_grad(self, g):
# d (grad_X(...(x)) / dx => grad_X(...(Argument(x.function_space()))
o = g
ngrads = 0
while isinstance(o, ReferenceGrad):
o, = o.ufl_operands
ngrads += 1
if not (isinstance(o, SpatialCoordinate) or isinstance(o.ufl_operands[0], FormArgument)):
error("Expecting gradient of a FormArgument, not %s" % ufl_err_str(o))
def apply_grads(f):
for i in range(ngrads):
f = ReferenceGrad(f)
return f
# Find o among all w without any indexing, which makes this
# easy
for (w, v) in zip(self._w, self._v):
if o == w and isinstance(v, ReferenceValue) and isinstance(v.ufl_operands[0], FormArgument):
# Case: d/dt [w + t v]
return apply_grads(v)
return self.independent_terminal(o)
[docs]
def jacobian(self, o):
# d (grad_X(x))/d x => grad_X(Argument(x.function_space())
for (w, v) in zip(self._w, self._v):
if o.ufl_domain() == w.ufl_domain() and isinstance(v.ufl_operands[0], FormArgument):
return ReferenceGrad(v)
return self.independent_terminal(o)
[docs]
class CoordinateDerivativeRuleDispatcher(MultiFunction):
def __init__(self):
MultiFunction.__init__(self)
self.vcache = defaultdict(dict)
self.rcache = defaultdict(dict)
[docs]
def terminal(self, o):
return o
[docs]
def derivative(self, o):
error("Missing derivative handler for {0}.".format(type(o).__name__))
expr = MultiFunction.reuse_if_untouched
[docs]
def grad(self, o):
return o
[docs]
def reference_grad(self, o):
return o
[docs]
def coefficient_derivative(self, o):
return o
[docs]
def coordinate_derivative(self, o, f, w, v, cd):
from ufl_legacy.algorithms import extract_unique_elements
spaces = set(c.family() for c in extract_unique_elements(o))
unsupported_spaces = {"Argyris", "Bell", "Hermite", "Morley"}
if spaces & unsupported_spaces:
error("CoordinateDerivative is not supported for elements of type %s. "
"This is because their pullback is not implemented in UFL." % unsupported_spaces)
_, w, v, cd = o.ufl_operands
rules = CoordinateDerivativeRuleset(w, v, cd)
key = (CoordinateDerivativeRuleset, w, v, cd)
return map_expr_dag(rules, f, vcache=self.vcache[key], rcache=self.rcache[key])
[docs]
def apply_coordinate_derivatives(expression):
rules = CoordinateDerivativeRuleDispatcher()
return map_integrand_dags(rules, expression)