"""Pull back and push forward maps."""
# Copyright (C) 2023 Matthew Scroggs, David Ham, Garth Wells
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier: LGPL-3.0-or-later
from __future__ import annotations
import typing
from abc import ABC, abstractmethod, abstractproperty
from typing import TYPE_CHECKING
import numpy as np
from ufl.core.expr import Expr
from ufl.core.multiindex import indices
from ufl.domain import extract_unique_domain
from ufl.functionspace import FunctionSpace
from ufl.tensors import as_tensor
if TYPE_CHECKING:
from ufl.finiteelement import AbstractFiniteElement as _AbstractFiniteElement
__all_classes__ = [
"NonStandardPullbackException",
"AbstractPullback",
"IdentityPullback",
"ContravariantPiola",
"CovariantPiola",
"L2Piola",
"DoubleContravariantPiola",
"DoubleCovariantPiola",
"MixedPullback",
"SymmetricPullback",
"PhysicalPullback",
"CustomPullback",
"UndefinedPullback",
]
[docs]class NonStandardPullbackException(BaseException):
"""Exception to raise if a map is non-standard."""
pass
[docs]class AbstractPullback(ABC):
"""An abstract pull back."""
@abstractmethod
def __repr__(self) -> str:
"""Return a representation of the object."""
[docs] @abstractmethod
def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
@abstractproperty
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
[docs] def apply(self, expr: Expr) -> Expr:
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
raise NonStandardPullbackException()
[docs]class IdentityPullback(AbstractPullback):
"""The identity pull back."""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "IdentityPullback()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return True
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
return expr
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
return element.reference_value_shape
[docs]class ContravariantPiola(AbstractPullback):
"""The contravariant Piola pull back."""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "ContravariantPiola()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return False
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
from ufl.classes import Jacobian, JacobianDeterminant
domain = extract_unique_domain(expr)
J = Jacobian(domain)
detJ = JacobianDeterminant(J)
transform = (1.0 / detJ) * J
# Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
*k, i, j = indices(len(expr.ufl_shape) + 1)
kj = (*k, j)
return as_tensor(transform[i, j] * expr[kj], (*k, i))
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
gdim = domain.geometric_dimension()
return (gdim,) + element.reference_value_shape[1:]
[docs]class CovariantPiola(AbstractPullback):
"""The covariant Piola pull back."""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "CovariantPiola()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return False
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
from ufl.classes import JacobianInverse
domain = extract_unique_domain(expr)
K = JacobianInverse(domain)
# Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
*k, i, j = indices(len(expr.ufl_shape) + 1)
kj = (*k, j)
return as_tensor(K[j, i] * expr[kj], (*k, i))
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
gdim = domain.geometric_dimension()
return (gdim,) + element.reference_value_shape[1:]
[docs]class L2Piola(AbstractPullback):
"""The L2 Piola pull back."""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "L2Piola()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return False
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
from ufl.classes import JacobianDeterminant
domain = extract_unique_domain(expr)
detJ = JacobianDeterminant(domain)
return expr / detJ
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
return element.reference_value_shape
[docs]class DoubleContravariantPiola(AbstractPullback):
"""The double contravariant Piola pull back."""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "DoubleContravariantPiola()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return False
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
from ufl.classes import Jacobian, JacobianDeterminant
domain = extract_unique_domain(expr)
J = Jacobian(domain)
detJ = JacobianDeterminant(J)
# Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
*k, i, j, m, n = indices(len(expr.ufl_shape) + 2)
kmn = (*k, m, n)
return as_tensor((1.0 / detJ) ** 2 * J[i, m] * expr[kmn] * J[j, n], (*k, i, j))
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
gdim = domain.geometric_dimension()
return (gdim, gdim)
[docs]class DoubleCovariantPiola(AbstractPullback):
"""The double covariant Piola pull back."""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "DoubleCovariantPiola()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return False
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
from ufl.classes import JacobianInverse
domain = extract_unique_domain(expr)
K = JacobianInverse(domain)
# Apply transform "row-wise" to TensorElement(PiolaMapped, ...)
*k, i, j, m, n = indices(len(expr.ufl_shape) + 2)
kmn = (*k, m, n)
return as_tensor(K[m, i] * expr[kmn] * K[n, j], (*k, i, j))
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
gdim = domain.geometric_dimension()
return (gdim, gdim)
[docs]class MixedPullback(AbstractPullback):
"""Pull back for a mixed element."""
def __init__(self, element: _AbstractFiniteElement):
"""Initalise.
Args:
element: The mixed element
"""
self._element = element
def __repr__(self) -> str:
"""Return a representation of the object."""
return f"MixedPullback({self._element!r})"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return all(e.pullback.is_identity for e in self._element.sub_elements)
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
domain = expr.ufl_domain()
space = FunctionSpace(domain, self._element)
rflat = [expr[idx] for idx in np.ndindex(expr.ufl_shape)]
g_components = []
offset = 0
# For each unique piece in reference space, apply the appropriate pullback
for subelem in self._element.sub_elements:
rsub = as_tensor(
np.asarray(rflat[offset : offset + subelem.reference_value_size]).reshape(
subelem.reference_value_shape
)
)
rmapped = subelem.pullback.apply(rsub)
# Flatten into the pulled back expression for the whole thing
g_components.extend([rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape)])
offset += subelem.reference_value_size
# And reshape appropriately
f = as_tensor(np.asarray(g_components).reshape(space.value_shape))
if f.ufl_shape != space.value_shape:
raise ValueError(
"Expecting pulled back expression with shape "
f"'{space.value_shape}', got '{f.ufl_shape}'"
)
return f
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
assert element == self._element
dim = sum(FunctionSpace(domain, e).value_size for e in self._element.sub_elements)
return (dim,)
[docs]class SymmetricPullback(AbstractPullback):
"""Pull back for an element with symmetry."""
def __init__(
self, element: _AbstractFiniteElement, symmetry: typing.Dict[typing.tuple[int, ...], int]
):
"""Initalise.
Args:
element: The element
symmetry: A dictionary mapping from the component in
physical space to the local component
"""
self._element = element
self._symmetry = symmetry
self._sub_element_value_shape = element.sub_elements[0].reference_value_shape
for e in element.sub_elements:
if e.reference_value_shape != self._sub_element_value_shape:
raise ValueError("Sub-elements must all have the same value shape.")
self._block_shape = tuple(i + 1 for i in max(symmetry.keys()))
def __repr__(self) -> str:
"""Return a representation of the object."""
return f"SymmetricPullback({self._element!r}, {self._symmetry!r})"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return all(e.pullback.is_identity for e in self._element.sub_elements)
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
domain = expr.ufl_domain()
space = FunctionSpace(domain, self._element)
rflat = [expr[idx] for idx in np.ndindex(expr.ufl_shape)]
g_components = []
offsets = [0]
for subelem in self._element.sub_elements:
offsets.append(offsets[-1] + subelem.reference_value_size)
# For each unique piece in reference space, apply the appropriate pullback
for component in np.ndindex(self._block_shape):
i = self._symmetry[component]
subelem = self._element.sub_elements[i]
rsub = as_tensor(
np.asarray(rflat[offsets[i] : offsets[i + 1]]).reshape(
subelem.reference_value_shape
)
)
rmapped = subelem.pullback.apply(rsub)
# Flatten into the pulled back expression for the whole thing
g_components.extend([rmapped[idx] for idx in np.ndindex(rmapped.ufl_shape)])
# And reshape appropriately
f = as_tensor(np.asarray(g_components).reshape(space.value_shape))
if f.ufl_shape != space.value_shape:
raise ValueError(
f"Expecting pulled back expression with shape "
f"'{space.value_shape}', got '{f.ufl_shape}'"
)
return f
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
assert element == self._element
return tuple(i + 1 for i in max(self._symmetry.keys()))
[docs]class PhysicalPullback(AbstractPullback):
"""Physical pull back.
This should probably be removed.
"""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "PhysicalPullback()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return True
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
return expr
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
raise NotImplementedError()
[docs]class CustomPullback(AbstractPullback):
"""Custom pull back.
This should probably be removed.
"""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "CustomPullback()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return True
[docs] def apply(self, expr):
"""Apply the pull back.
Args:
expr: A function on a physical cell
Returns: The function pulled back to the reference cell
"""
return expr
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
if element.reference_value_shape == ():
return ()
raise NotImplementedError()
[docs]class UndefinedPullback(AbstractPullback):
"""Undefined pull back.
This should probably be removed.
"""
def __repr__(self) -> str:
"""Return a representation of the object."""
return "UndefinedPullback()"
@property
def is_identity(self) -> bool:
"""Is this pull back the identity (or the identity applied to mutliple components)."""
return True
[docs] def physical_value_shape(self, element, domain) -> typing.Tuple[int, ...]:
"""Get the physical value shape when this pull back is applied to an element on a domain.
Args:
element: The element that the pull back is applied to
domain: The domain
Returns:
The value shape when the pull back is applied to the given element
"""
raise NotImplementedError()
identity_pullback = IdentityPullback()
covariant_piola = CovariantPiola()
contravariant_piola = ContravariantPiola()
l2_piola = L2Piola()
double_covariant_piola = DoubleCovariantPiola()
double_contravariant_piola = DoubleContravariantPiola()
physical_pullback = PhysicalPullback()
custom_pullback = CustomPullback()
undefined_pullback = UndefinedPullback()