Source code for ufl.algorithms.apply_geometry_lowering

"""Algorithm for lowering abstractions of geometric types.

This means replacing high-level types with expressions
of mostly the Jacobian and reference cell data.
"""

# Copyright (C) 2013-2016 Martin Sandve Alnæs
#
# This file is part of UFL (https://www.fenicsproject.org)
#
# SPDX-License-Identifier:    LGPL-3.0-or-later

import warnings
from functools import reduce
from itertools import combinations

from ufl.classes import (
    CellCoordinate,
    CellEdgeVectors,
    CellFacetJacobian,
    CellOrientation,
    CellOrigin,
    CellVertices,
    CellVolume,
    Expr,
    FacetEdgeVectors,
    FacetJacobian,
    FacetJacobianDeterminant,
    FloatValue,
    Form,
    Integral,
    Jacobian,
    JacobianDeterminant,
    JacobianInverse,
    MaxCellEdgeLength,
    ReferenceCellVolume,
    ReferenceFacetVolume,
    ReferenceGrad,
    ReferenceNormal,
    SpatialCoordinate,
)
from ufl.compound_expressions import cross_expr, determinant_expr, inverse_expr
from ufl.core.multiindex import Index, indices
from ufl.corealg.map_dag import map_expr_dag
from ufl.corealg.multifunction import MultiFunction, memoized_handler
from ufl.domain import extract_unique_domain
from ufl.measure import custom_integral_types, point_integral_types
from ufl.operators import conj, max_value, min_value, real, sqrt
from ufl.tensors import as_tensor, as_vector


[docs]class GeometryLoweringApplier(MultiFunction): """Geometry lowering.""" def __init__(self, preserve_types=()): """Initialise.""" MultiFunction.__init__(self) # Store preserve_types as boolean lookup table self._preserve_types = [False] * Expr._ufl_num_typecodes_ for cls in preserve_types: self._preserve_types[cls._ufl_typecode_] = True expr = MultiFunction.reuse_if_untouched
[docs] def terminal(self, t): """Apply to terminal.""" return t
@memoized_handler def jacobian(self, o): """Apply to jacobian.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) if not domain.ufl_coordinate_element().pullback.is_identity: raise ValueError("Piola mapped coordinates are not implemented.") # Note: No longer supporting domain.coordinates(), always # preserving SpatialCoordinate object. However if Jacobians # are not preserved, using # ReferenceGrad(SpatialCoordinate(domain)) to represent them. x = self.spatial_coordinate(SpatialCoordinate(domain)) return ReferenceGrad(x) @memoized_handler def _future_jacobian(self, o): """Apply to _future_jacobian.""" # If we're not using Coefficient to represent the spatial # coordinate, we can just as well just return o here too # unless we add representation of basis functions and dofs to # the ufl layer (which is nice to avoid). return o @memoized_handler def jacobian_inverse(self, o): """Apply to jacobian_inverse.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) J = self.jacobian(Jacobian(domain)) # TODO: This could in principle use # preserve_types[JacobianDeterminant] with minor refactoring: K = inverse_expr(J) return K @memoized_handler def jacobian_determinant(self, o): """Apply to jacobian_determinant.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) J = self.jacobian(Jacobian(domain)) detJ = determinant_expr(J) # TODO: Is "signing" the determinant for manifolds the # cleanest approach? The alternative is to have a # specific type for the unsigned pseudo-determinant. if domain.topological_dimension() < domain.geometric_dimension(): co = CellOrientation(domain) detJ = co * detJ return detJ @memoized_handler def facet_jacobian(self, o): """Apply to facet_jacobian.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) J = self.jacobian(Jacobian(domain)) RFJ = CellFacetJacobian(domain) i, j, k = indices(3) return as_tensor(J[i, k] * RFJ[k, j], (i, j)) @memoized_handler def facet_jacobian_inverse(self, o): """Apply to facet_jacobian_inverse.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) FJ = self.facet_jacobian(FacetJacobian(domain)) # This could in principle use # preserve_types[JacobianDeterminant] with minor refactoring: return inverse_expr(FJ) @memoized_handler def facet_jacobian_determinant(self, o): """Apply to facet_jacobian_determinant.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) FJ = self.facet_jacobian(FacetJacobian(domain)) detFJ = determinant_expr(FJ) # TODO: Should we "sign" the facet jacobian determinant for # manifolds? It's currently used unsigned in # apply_integral_scaling. # if domain.topological_dimension() < domain.geometric_dimension(): # co = CellOrientation(domain) # detFJ = co*detFJ return detFJ @memoized_handler def spatial_coordinate(self, o): """Apply to spatial_coordinate. Fall through to coordinate field of domain if it exists. """ if self._preserve_types[o._ufl_typecode_]: return o if not extract_unique_domain(o).ufl_coordinate_element().pullback.is_identity: raise ValueError("Piola mapped coordinates are not implemented.") # No longer supporting domain.coordinates(), always preserving # SpatialCoordinate object. return o @memoized_handler def cell_coordinate(self, o): """Apply to cell_coordinate. Compute from physical coordinates if they are known, using the appropriate mappings. """ if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) K = self.jacobian_inverse(JacobianInverse(domain)) x = self.spatial_coordinate(SpatialCoordinate(domain)) x0 = CellOrigin(domain) i, j = indices(2) X = as_tensor(K[i, j] * (x[j] - x0[j]), (i,)) return X @memoized_handler def facet_cell_coordinate(self, o): """Apply to facet_cell_coordinate.""" if self._preserve_types[o._ufl_typecode_]: return o raise ValueError( "Missing computation of facet reference coordinates " "from physical coordinates via mappings." ) @memoized_handler def cell_volume(self, o): """Apply to cell_volume.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) if not domain.is_piecewise_linear_simplex_domain(): # Don't lower for non-affine cells, instead leave it to # form compiler warnings.warn("Only know how to compute the cell volume of an affine cell.") return o r = self.jacobian_determinant(JacobianDeterminant(domain)) r0 = ReferenceCellVolume(domain) return abs(r * r0) @memoized_handler def facet_area(self, o): """Apply to facet_area.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) tdim = domain.topological_dimension() if not domain.is_piecewise_linear_simplex_domain(): # Don't lower for non-affine cells, instead leave it to # form compiler warnings.warn("Only know how to compute the facet area of an affine cell.") return o # Area of "facet" of interval (i.e. "area" of a vertex) is defined as 1.0 if tdim == 1: return FloatValue(1.0) r = self.facet_jacobian_determinant(FacetJacobianDeterminant(domain)) r0 = ReferenceFacetVolume(domain) return abs(r * r0) @memoized_handler def circumradius(self, o): """Apply to circumradius.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) if not domain.is_piecewise_linear_simplex_domain(): raise ValueError("Circumradius only makes sense for affine simplex cells") cellname = domain.ufl_cell().cellname() cellvolume = self.cell_volume(CellVolume(domain)) if cellname == "interval": # Optimization for square interval; no square root needed return 0.5 * cellvolume # Compute lengths of cell edges edges = CellEdgeVectors(domain) num_edges = edges.ufl_shape[0] j = Index() elen = [real(sqrt(real(edges[e, j] * conj(edges[e, j])))) for e in range(num_edges)] if cellname == "triangle": return (elen[0] * elen[1] * elen[2]) / (4.0 * cellvolume) elif cellname == "tetrahedron": # la, lb, lc = lengths of the sides of an intermediate triangle # NOTE: Is here some hidden numbering assumption? la = elen[3] * elen[2] lb = elen[4] * elen[1] lc = elen[5] * elen[0] # p = perimeter p = la + lb + lc # s = semiperimeter s = p / 2 # area of intermediate triangle with Herons formula triangle_area = sqrt(s * (s - la) * (s - lb) * (s - lc)) return triangle_area / (6.0 * cellvolume) @memoized_handler def max_cell_edge_length(self, o): """Apply to max_cell_edge_length.""" return self._reduce_cell_edge_length(o, max_value) @memoized_handler def min_cell_edge_length(self, o): """Apply to min_cell_edge_length.""" return self._reduce_cell_edge_length(o, min_value) def _reduce_cell_edge_length(self, o, reduction_op): """Apply to _reduce_cell_edge_length.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) if domain.ufl_coordinate_element().embedded_superdegree > 1: # Don't lower bendy cells, instead leave it to form compiler warnings.warn("Only know how to compute cell edge lengths of P1 or Q1 cell.") return o elif domain.ufl_cell().cellname() == "interval": # Interval optimization, square root not needed return self.cell_volume(CellVolume(domain)) else: # Other P1 or Q1 cells edges = CellEdgeVectors(domain) num_edges = edges.ufl_shape[0] j = Index() elen2 = [real(edges[e, j] * conj(edges[e, j])) for e in range(num_edges)] return real(sqrt(reduce(reduction_op, elen2))) @memoized_handler def cell_diameter(self, o): """Apply to cell_diameter.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) if domain.ufl_coordinate_element().embedded_superdegree > 1: # Don't lower bendy cells, instead leave it to form compiler warnings.warn("Only know how to compute cell diameter of P1 or Q1 cell.") return o elif domain.is_piecewise_linear_simplex_domain(): # Simplices return self.max_cell_edge_length(MaxCellEdgeLength(domain)) else: # Q1 cells, maximal distance between any two vertices verts = CellVertices(domain) verts = [verts[v, ...] for v in range(verts.ufl_shape[0])] j = Index() elen2 = (real((v0 - v1)[j] * conj((v0 - v1)[j])) for v0, v1 in combinations(verts, 2)) return real(sqrt(reduce(max_value, elen2))) @memoized_handler def max_facet_edge_length(self, o): """Apply to max_facet_edge_length.""" return self._reduce_facet_edge_length(o, max_value) @memoized_handler def min_facet_edge_length(self, o): """Apply to min_facet_edge_length.""" return self._reduce_facet_edge_length(o, min_value) def _reduce_facet_edge_length(self, o, reduction_op): """Apply to _reduce_facet_edge_length.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) if domain.ufl_cell().topological_dimension() < 3: raise ValueError("Facet edge lengths only make sense for topological dimension >= 3.") elif domain.ufl_coordinate_element().embedded_superdegree > 1: # Don't lower bendy cells, instead leave it to form compiler warnings.warn("Only know how to compute facet edge lengths of P1 or Q1 cell.") return o else: # P1 tetrahedron or Q1 hexahedron edges = FacetEdgeVectors(domain) num_edges = edges.ufl_shape[0] j = Index() elen2 = [real(edges[e, j] * conj(edges[e, j])) for e in range(num_edges)] return real(sqrt(reduce(reduction_op, elen2))) @memoized_handler def cell_normal(self, o): """Apply to cell_normal.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) gdim = domain.geometric_dimension() tdim = domain.topological_dimension() if tdim == gdim - 1: # n-manifold embedded in n-1 space i = Index() J = self.jacobian(Jacobian(domain)) if tdim == 2: # Surface in 3D t0 = as_vector(J[i, 0], i) t1 = as_vector(J[i, 1], i) cell_normal = cross_expr(t0, t1) elif tdim == 1: # Line in 2D (cell normal is 'up' for a line pointing # to the 'right') cell_normal = as_vector((-J[1, 0], J[0, 0])) else: raise ValueError(f"Cell normal not implemented for tdim {tdim}, gdim {gdim}") # Return normalized vector, sign corrected by cell # orientation co = CellOrientation(domain) return co * cell_normal / sqrt(cell_normal[i] * cell_normal[i]) else: raise ValueError(f"Cell normal undefined for tdim {tdim}, gdim {gdim}") @memoized_handler def facet_normal(self, o): """Apply to facet_normal.""" if self._preserve_types[o._ufl_typecode_]: return o domain = extract_unique_domain(o) tdim = domain.topological_dimension() if tdim == 1: # Special-case 1D (possibly immersed), for which we say # that n is just in the direction of J. J = self.jacobian(Jacobian(domain)) # dx/dX ndir = J[:, 0] gdim = domain.geometric_dimension() if gdim == 1: nlen = abs(ndir[0]) else: i = Index() nlen = sqrt(ndir[i] * ndir[i]) rn = ReferenceNormal(domain) # +/- 1.0 here n = rn[0] * ndir / nlen r = n else: # Recall that the covariant Piola transform u -> J^(-T)*u # preserves tangential components. The normal vector is # characterised by having zero tangential component in # reference and physical space. Jinv = self.jacobian_inverse(JacobianInverse(domain)) i, j = indices(2) rn = ReferenceNormal(domain) # compute signed, unnormalised normal; note transpose ndir = as_vector(Jinv[j, i] * rn[j], i) # normalise i = Index() n = ndir / sqrt(ndir[i] * ndir[i]) r = n if r.ufl_shape != o.ufl_shape: raise ValueError( f"Inconsistent dimensions (in={o.ufl_shape[0]}, out={r.ufl_shape[0]})." ) return r
[docs]def apply_geometry_lowering(form, preserve_types=()): """Change GeometricQuantity objects in expression to the lowest level GeometricQuantity objects. Assumes the expression is preprocessed or at least that derivatives have been expanded. Args: form: An Expr or Form. preserve_types: Preserved types """ if isinstance(form, Form): newintegrals = [ apply_geometry_lowering(integral, preserve_types) for integral in form.integrals() ] return Form(newintegrals) elif isinstance(form, Integral): integral = form if integral.integral_type() in (custom_integral_types + point_integral_types): automatic_preserve_types = [SpatialCoordinate, Jacobian] else: automatic_preserve_types = [CellCoordinate] preserve_types = set(preserve_types) | set(automatic_preserve_types) mf = GeometryLoweringApplier(preserve_types) newintegrand = map_expr_dag(mf, integral.integrand()) return integral.reconstruct(integrand=newintegrand) elif isinstance(form, Expr): expr = form mf = GeometryLoweringApplier(preserve_types) return map_expr_dag(mf, expr) else: raise ValueError(f"Invalid type {form.__class__.__name__}")