diff --git a/generate_all_operators.py b/generate_all_operators.py
index 6ce4a12dcab8ddddd60fb2e05107591d8415a84a..af0df0d527b10e42fa51c1b0e0ed256b134066b5 100644
--- a/generate_all_operators.py
+++ b/generate_all_operators.py
@@ -1,910 +1,915 @@
-# HyTeG Operator Generator
-# Copyright (C) 2024 HyTeG Team
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program. If not, see <https://www.gnu.org/licenses/>.
-
-import argparse
-import sys
-from argparse import ArgumentParser
-from dataclasses import dataclass, field
-from functools import partial
-import logging
-import os
-import re
-from typing import Callable, List, Optional, Sequence, Set, Tuple, Union
-
-import numpy as np
-import sympy as sp
-from sympy.core.cache import clear_cache
-from tabulate import tabulate
-
-from hog.blending import GeometryMap, IdentityMap, AnnulusMap, IcosahedralShellMap, ParametricMap
-from hog.cse import CseImplementation
-from hog.element_geometry import (
- ElementGeometry,
- TriangleElement,
- TetrahedronElement,
- LineElement,
-)
-from hog.exception import HOGException
-from hog.forms import (
- diffusion,
- divergence,
- gradient,
- div_k_grad,
- shear_heating,
- epsilon,
- full_stokes,
- nonlinear_diffusion,
- nonlinear_diffusion_newton_galerkin,
- supg_diffusion,
- advection,
- supg_advection,
- grad_rho_by_rho_dot_u,
-)
-from hog.forms_boundary import (
- mass_boundary,
- freeslip_gradient_weak_boundary,
- freeslip_divergence_weak_boundary,
- freeslip_momentum_weak_boundary,
-)
-from hog.forms_vectorial import curl_curl, curl_curl_plus_mass, mass_n1e1
-from hog.function_space import (
- LagrangianFunctionSpace,
- N1E1Space,
- TensorialVectorFunctionSpace,
- TrialSpace,
- TestSpace,
-)
-from hog.logger import get_logger, TimedLogger
-from hog.operator_generation.kernel_types import (
- ApplyWrapper,
- AssembleWrapper,
- AssembleDiagonalWrapper,
- KernelWrapperType,
-)
-from hog.operator_generation.loop_strategies import (
- LoopStrategy,
- SAWTOOTH,
- CUBES,
- FUSEDROWS,
-)
-from hog.operator_generation.operators import (
- HyTeGElementwiseOperator,
- MacroIntegrationDomain,
-)
-
-from operator import itemgetter
-from hog.operator_generation.optimizer import (
- Opts,
- opts_arg_mapping,
- ordered_opts_suffix,
-)
-from hog.quadrature import Quadrature
-from hog.symbolizer import Symbolizer
-from generate_all_hyteg_forms import valid_base_dir
-from hog.operator_generation.types import parse_argument_type, HOGType, HOGPrecision
-import quadpy
-from hog.quadrature.quadrature import select_quadrule
-from hog.integrand import Form
-
-ALL_GEOMETRY_MAPS = [
- IdentityMap(),
- AnnulusMap(),
- IcosahedralShellMap(),
- ParametricMap(1),
- ParametricMap(2),
-]
-
-
-def parse_arguments():
- parser = ArgumentParser(
- description="Generates a collection of elementwise operators."
- )
-
- list_or_gen = parser.add_mutually_exclusive_group(required=True)
- list_or_gen.add_argument(
- "-l",
- "--list-operators",
- action="store_true",
- help="List all available operators.",
- )
- list_or_gen.add_argument(
- "--list-quadratures",
- nargs=1,
- help="List all available quadrature rules for a given geometry (t1: Line, t2: Triangle, t3: Tetrahedron)",
- )
-
- list_or_gen.add_argument(
- "--hyteg",
- help=(
- "Path to the HyTeG repository. The generated code is written to the correct directory in the source tree."
- ),
- )
- list_or_gen.add_argument(
- "--output",
- help=("Path to output directory."),
- )
-
- parser.add_argument(
- "-s",
- "--space-mapping",
- default=".*?",
- help="Filter by function space, e.g. 'P1' or 'P1toP2' (regex).",
- # TODO add choices list and type=string.lower
- )
-
- parser.add_argument(
- "-f",
- "--form",
- default=".*?",
- help="Filter by form, e.g. 'CurlCurl' or 'Diffusion' (regex).",
- # TODO add choices list and type=string.lower
- )
-
- parser.add_argument(
- "-o",
- "--opts",
- default="",
- nargs="+",
- help=f"""
- Enter a space separated list of optimizations.
- Not all combinations are currently supported. Available optimizations are: {Opts.__members__.keys()}""",
- # TODO add choices list and type=string.lower
- )
- # TODO explanations for optimizations e.g. CUTLOOPS is not straight-forward
-
- parser.add_argument(
- "--loop-strategy",
- default="SAWTOOTH",
- help="Enter a loop strategy for the operator. Available strategies are SAWTOOTH, CUBES.",
- # TODO add choices list and type=string.lower
- )
-
- quad_rule_or_degree = parser.add_mutually_exclusive_group(required=False)
- quad_rule_or_degree.add_argument(
- "--quad-rule",
- type=str,
- nargs=2,
- default=None,
- help="""Supply two rules by their identifier (leftmost column when calling with --list-quadratures),
- the first will be used for triangles and the second for tetrahedra.""",
- )
- quad_rule_or_degree.add_argument(
- "--quad-degree",
- type=int,
- nargs=2,
- default=[-1, -1],
- help="""Supply two polynomial degrees (as integers) up to which the quadrature should be exact.
- The generator will chose the rule with minimal points that integrates the specified
- degree exactly. The first integer will be used for triangles and the second for tetrahedra.""",
- )
-
- prec_or_mpfr = parser.add_mutually_exclusive_group(required=False)
- prec_or_mpfr.add_argument(
- "-p",
- "--precision",
- type=str.lower,
- help=(
- "Defines the precision in which the kernels are calculated; "
- "If not specified we choose fp64(double) precision"
- ),
- choices=list(choice.value for choice in HOGPrecision),
- default="fp64",
- )
-
- parser.add_argument(
- "-b",
- "--blending",
- default=str(IdentityMap()),
- help=f"Enter a blending type for the operator. Note some geometry maps only work in 2D or 3D respectively. "
- f"Available geometry maps: {[str(m) for m in ALL_GEOMETRY_MAPS]}",
- # TODO add choices list and type=string.lower
- )
-
- parser.add_argument(
- "--name",
- help=f"Specify the name of the generated C++ class.",
- )
-
- parser.add_argument(
- "--dimensions",
- type=int,
- nargs="+",
- default=[2, 3],
- help=f"Domain dimensions for which operators shall be generated.",
- )
-
- logging_group = parser.add_mutually_exclusive_group(required=False)
- logging_group.add_argument(
- "-v",
- "--verbose",
- type=int,
- help="specify verbosity level (0 = silent, 1 = info, 2 = debug)",
- choices=[0, 1, 2],
- )
- logging_group.add_argument(
- "--loglevel",
- type=str.lower,
- help=(
- "Defines which messages are shown and which not; "
- "If not specified 'warning' is chosen."
- ),
- choices=["debug", "info", "warning", "error", "critical"],
- )
-
- clang_format_group = parser.add_mutually_exclusive_group(required=False)
- clang_format_group.add_argument(
- "--no-clang-format",
- dest="clang_format",
- action="store_false",
- help="Do not apply clang-format on generated files.",
- )
-
- clang_format_group.add_argument(
- "--clang-format-binary",
- default="clang-format",
- help=f"Allows to specify the name of the clang-format binary and/or optionally its full path."
- f" By default 'clang-format' will be used.",
- )
-
- args = parser.parse_args()
- return parser, args
-
-
-strategy_args_mapping = {
- "CUBES": CUBES(),
- "SAWTOOTH": SAWTOOTH(),
- "FUSEDROWS": FUSEDROWS(),
-}
-
-
-def main():
- clear_cache()
-
- parser, args = parse_arguments()
-
- # Adapt clang_format_binary according to the value of clang_format
- if "clang_format" in vars(args):
- if not args.clang_format:
- args.clang_format_binary = None
-
- if args.verbose:
- if args.verbose == 0:
- log_level = logging.CRITICAL
- elif args.verbose == 1:
- log_level = logging.INFO
- elif args.verbose == 2:
- log_level = logging.DEBUG
- else:
- raise HOGException(f"Invalid verbosity level {args.verbose}.")
- elif args.loglevel:
- if args.loglevel == "debug":
- log_level = logging.DEBUG
- elif args.loglevel == "info":
- log_level = logging.INFO
- elif args.loglevel == "warning":
- log_level = logging.WARNING
- elif args.loglevel == "error":
- log_level = logging.ERROR
- elif args.loglevel == "critical":
- log_level = logging.CRITICAL
- else:
- raise HOGException(f"Invalid log level {args.loglevel}.")
- else:
- # The default log level is chosen here
- log_level = logging.DEBUG
-
- logger = get_logger(log_level)
- TimedLogger.set_log_level(log_level)
-
- logger.info(f"Running {__file__} with arguments: {' '.join(sys.argv[1:])}")
-
- symbolizer = Symbolizer()
-
- if args.list_quadratures:
- geometry = args.list_quadratures[0]
- logger.debug(f"Listing quadrature rules for geometry: {geometry}.")
- if geometry == "t3":
- schemes = quadpy.t3.schemes
- elif geometry == "t2":
- schemes = quadpy.t2.schemes
- else:
- raise HOGException(f"Unexpected geometry: {geometry}")
-
- all_schemes = []
- for key, s in schemes.items():
- try:
- scheme = s()
- entry = {
- "key": key,
- "name": scheme.name,
- "points": scheme.points.shape[1],
- "degree": scheme.degree,
- "test_tolerance": scheme.test_tolerance,
- }
- all_schemes.append(entry)
- except TypeError as e:
- # Some schemes seem to require parameters, we simply exclude these with this litte try-except hack
- pass
-
- all_schemes = sorted(all_schemes, key=itemgetter("degree", "points"))
- table = tabulate(all_schemes, headers="keys", tablefmt="grid")
- print(table)
- exit()
-
- blending = None
- for m in ALL_GEOMETRY_MAPS:
- if str(m) == args.blending:
- blending = m
- if blending is None:
- raise HOGException("Invalid geometry map.")
- logger.debug(f"Blending: {blending}")
-
- type_descriptor = parse_argument_type(args)
-
- logger.debug(f"Data type name is {type_descriptor.pystencils_type}.")
-
- opts = {opts_arg_mapping[opt] for opt in args.opts}
- loop_strategy = strategy_args_mapping[args.loop_strategy]
- re_form = re.compile(args.form)
-
- if (
- not args.list_operators
- and args.quad_degree == [-1, -1]
- and args.quad_rule is None
- ):
- parser.error("Either --quad-degree or --quad-rule are required.")
- quad = {
- geometry: info
- for geometry, info in zip(
- [TriangleElement(), TetrahedronElement()],
- args.quad_degree if args.quad_rule is None else args.quad_rule,
- )
- }
-
- enabled_geometries: Set[TriangleElement | TetrahedronElement] = set()
- if 2 in args.dimensions:
- enabled_geometries.add(TriangleElement())
- if 3 in args.dimensions:
- enabled_geometries.add(TetrahedronElement())
-
- operators = all_operators(
- symbolizer,
- [(opts, loop_strategy, ordered_opts_suffix(loop_strategy, opts))],
- type_descriptor,
- blending,
- geometries=enabled_geometries,
- )
-
- filtered_operators = list(
- filter(
- lambda o: re.fullmatch(args.space_mapping, o.mapping, re.IGNORECASE)
- and re.fullmatch(args.form, o.name, re.IGNORECASE),
- operators,
- )
- )
-
- filtered_operators = [f for f in filtered_operators if f.geometries]
-
- if len(filtered_operators) == 0:
- raise HOGException(
- f"Form {args.form} does not match any available forms. Run --list for a list."
- )
-
- if args.list_operators:
- print(
- tabulate(
- sorted(
- [
- (
- o.mapping,
- o.name,
- ", ".join(
- str(geo.dimensions) + "D" for geo in o.geometries
- ),
- len(o.opts),
- )
- for o in filtered_operators
- ]
- ),
- headers=(
- "Space mapping",
- "Form",
- "Dimensions",
- "Optimizations",
- ),
- )
- )
- exit()
-
- if args.hyteg:
- args.output = os.path.join(args.hyteg, "src/hyteg/operatorgeneration/generated")
- if not valid_base_dir(args.hyteg):
- raise HOGException(
- "The specified directory does not seem to be the HyTeG directory."
- )
-
- for operator in filtered_operators:
- for opt, loop, opt_name in operator.opts:
- blending_str = (
- "_" + str(blending) if not isinstance(blending, IdentityMap) else ""
- )
- if args.name:
- name = args.name
- else:
- name = (
- f"{operator.mapping}Elementwise{operator.name}{blending_str}{opt_name}"
- f"{'_' + str(type_descriptor) if type_descriptor.add_file_suffix else ''}"
- )
- with TimedLogger(f"Generating {name}", level=logging.INFO):
- generate_elementwise_op(
- args,
- symbolizer,
- operator,
- name,
- opt,
- loop,
- blending,
- quad,
- quad,
- type_descriptor=type_descriptor,
- )
-
-
-def all_opts_sympy_cse() -> List[Tuple[Set[Opts], LoopStrategy, str]]:
- return [
- (
- set(
- {
- Opts.MOVECONSTANTS,
- Opts.REORDER,
- Opts.ELIMTMPS,
- Opts.QUADLOOPS,
- Opts.VECTORIZE,
- }
- ),
- SAWTOOTH(),
- "_const_vect_fused_quadloops_reordered_elimtmps",
- ),
- ]
-
-
-def all_opts_poly_cse() -> List[Tuple[Set[Opts], LoopStrategy, str]]:
- return [
- (o | {Opts.POLYCSE}, l, n + "_polycse") for (o, l, n) in all_opts_sympy_cse()
- ]
-
-
-def all_opts_both_cses() -> List[Tuple[Set[Opts], LoopStrategy, str]]:
- return all_opts_sympy_cse() + all_opts_poly_cse()
-
-
-@dataclass
-class OperatorInfo:
- mapping: str
- name: str
- trial_space: TrialSpace
- test_space: TestSpace
- form: (
- Callable[
- [
- TrialSpace,
- TestSpace,
- ElementGeometry,
- Symbolizer,
- GeometryMap,
- ],
- Form,
- ]
- | None
- )
- type_descriptor: HOGType
- form_boundary: (
- Callable[
- [
- TrialSpace,
- TestSpace,
- ElementGeometry,
- ElementGeometry,
- Symbolizer,
- GeometryMap,
- ],
- Form,
- ]
- | None
- ) = None
- kernel_types: List[KernelWrapperType] = None # type: ignore[assignment] # will definitely be initialized in __post_init__
- geometries: Sequence[ElementGeometry] = field(
- default_factory=lambda: [TriangleElement(), TetrahedronElement()]
- )
- opts: List[Tuple[Set[Opts], LoopStrategy, str]] = field(
- default_factory=all_opts_sympy_cse
- )
- blending: GeometryMap = field(default_factory=lambda: IdentityMap())
-
- def __post_init__(self):
- # Removing geometries that are not supported by blending.
- # I don't like this, but it looks like the collection of operators has to be refactored anyway.
- self.geometries = list(
- set(self.geometries) & set(self.blending.supported_geometries())
- )
-
- if self.kernel_types is None:
- dims = [g.dimensions for g in self.geometries]
- self.kernel_types = [
- ApplyWrapper(
- self.trial_space,
- self.test_space,
- type_descriptor=self.type_descriptor,
- dims=dims,
- )
- ]
-
- all_opts = set().union(*[o for (o, _, _) in self.opts])
- if not ({Opts.VECTORIZE, Opts.VECTORIZE512}.intersection(all_opts)):
- self.kernel_types.append(
- AssembleWrapper(
- self.trial_space,
- self.test_space,
- type_descriptor=self.type_descriptor,
- dims=dims,
- )
- )
-
- if self.test_space == self.trial_space:
- self.kernel_types.append(
- AssembleDiagonalWrapper(
- self.trial_space,
- type_descriptor=self.type_descriptor,
- dims=dims,
- )
- )
-
-
-def all_operators(
- symbolizer: Symbolizer,
- opts: List[Tuple[Set[Opts], LoopStrategy, str]],
- type_descriptor: HOGType,
- blending: GeometryMap,
- geometries: Set[Union[TriangleElement, TetrahedronElement]],
-) -> List[OperatorInfo]:
- P1 = LagrangianFunctionSpace(1, symbolizer)
- P1Vector = TensorialVectorFunctionSpace(P1)
- P2 = LagrangianFunctionSpace(2, symbolizer)
- P2Vector = TensorialVectorFunctionSpace(P2)
- N1E1 = N1E1Space(symbolizer)
-
- two_d = list(geometries & {TriangleElement()})
- three_d = list(geometries & {TetrahedronElement()})
-
- ops: List[OperatorInfo] = []
-
- # fmt: off
- # TODO switch to manual specification of opts for now/developement, later use default factory
- ops.append(OperatorInfo("N1E1", "CurlCurl", TrialSpace(N1E1), TestSpace(N1E1), form=curl_curl,
- type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
- ops.append(OperatorInfo("N1E1", "Mass", TrialSpace(N1E1), TestSpace(N1E1), form=mass_n1e1,
- type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
- ops.append(OperatorInfo("N1E1", "CurlCurlPlusMass", TrialSpace(N1E1), TestSpace(N1E1),
- form=partial(curl_curl_plus_mass, alpha_fem_space=P1, beta_fem_space=P1),
- type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
- ops.append(OperatorInfo("P1", "Diffusion", TrialSpace(P1), TestSpace(P1), form=diffusion,
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
- ops.append(OperatorInfo("P1", "DivKGrad", TrialSpace(P1), TestSpace(P1),
- form=partial(div_k_grad, coefficient_function_space=P1),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P2", "Diffusion", TrialSpace(P2), TestSpace(P2), form=diffusion,
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
- ops.append(OperatorInfo("P2", "DivKGrad", TrialSpace(P2), TestSpace(P2),
- form=partial(div_k_grad, coefficient_function_space=P2),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P2", "ShearHeating", TrialSpace(P2), TestSpace(P2),
- form=partial(shear_heating, viscosity_function_space=P2, velocity_function_space=P2),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P1", "NonlinearDiffusion", TrialSpace(P1), TestSpace(P1),
- form=partial(nonlinear_diffusion, coefficient_function_space=P1),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
- ops.append(OperatorInfo("P1", "NonlinearDiffusionNewtonGalerkin", TrialSpace(P1),
- TestSpace(P1), form=partial(nonlinear_diffusion_newton_galerkin,
- coefficient_function_space=P1, only_newton_galerkin_part_of_form=False),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P1Vector", "Diffusion", TrialSpace(P1Vector), TestSpace(P1Vector),
- form=diffusion, type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P2", "SUPGDiffusion", TrialSpace(P2), TestSpace(P2),
- form=partial(supg_diffusion, velocity_function_space=P2, diffusivityXdelta_function_space=P2),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P2", "AdvectionConstCp", TrialSpace(P2), TestSpace(P2),
- form=partial(advection, velocity_function_space=P2, coefficient_function_space=P2, constant_cp = True),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P2", "AdvectionVarCp", TrialSpace(P2), TestSpace(P2),
- form=partial(advection, velocity_function_space=P2, coefficient_function_space=P2),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P2", "SUPGAdvection", TrialSpace(P2), TestSpace(P2),
- form=partial(supg_advection, velocity_function_space=P2, coefficient_function_space=P2),
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
-
- ops.append(OperatorInfo("P2", "BoundaryMass", TrialSpace(P2), TestSpace(P2), form=None,
- form_boundary=mass_boundary, type_descriptor=type_descriptor, geometries=list(geometries),
- opts=opts, blending=blending))
-
- # fmt: on
-
- p2vec_epsilon = partial(
- epsilon,
- variable_viscosity=True,
- coefficient_function_space=P2,
- )
- ops.append(
- OperatorInfo(
- f"P2Vector",
- f"Epsilon",
- TrialSpace(P2Vector),
- TestSpace(P2Vector),
- form=p2vec_epsilon,
- type_descriptor=type_descriptor,
- geometries=list(geometries),
- opts=opts,
- blending=blending,
- )
- )
-
- p2vec_freeslip_momentum_weak_boundary = partial(
- freeslip_momentum_weak_boundary, function_space_mu=P2
- )
-
- ops.append(
- OperatorInfo(
- f"P2Vector",
- f"EpsilonFreeslip",
- TrialSpace(P2Vector),
- TestSpace(P2Vector),
- form=p2vec_epsilon,
- form_boundary=p2vec_freeslip_momentum_weak_boundary,
- type_descriptor=type_descriptor,
- geometries=list(geometries),
- opts=opts,
- blending=blending,
- )
- )
-
- ops.append(
- OperatorInfo(
- f"P2VectorToP1",
- f"DivergenceFreeslip",
- TrialSpace(P2Vector),
- TestSpace(P1),
- form=divergence,
- form_boundary=freeslip_divergence_weak_boundary,
- type_descriptor=type_descriptor,
- geometries=list(geometries),
- opts=opts,
- blending=blending,
- )
- )
-
- ops.append(
- OperatorInfo(
- f"P1ToP2Vector",
- f"GradientFreeslip",
- TrialSpace(P1),
- TestSpace(P2Vector),
- form=gradient,
- form_boundary=freeslip_gradient_weak_boundary,
- type_descriptor=type_descriptor,
- geometries=list(geometries),
- opts=opts,
- blending=blending,
- )
- )
-
- p2vec_grad_rho = partial(
- grad_rho_by_rho_dot_u,
- density_function_space=P2,
- )
- ops.append(
- OperatorInfo(
- f"P2VectorToP1",
- f"GradRhoByRhoDotU",
- TrialSpace(P1),
- TestSpace(P2Vector),
- form=p2vec_grad_rho,
- type_descriptor=type_descriptor,
- geometries=list(geometries),
- opts=opts,
- blending=blending,
- )
- )
-
- for c in [0, 1, 2]:
- # fmt: off
- if c == 2:
- div_geometries = three_d
- else:
- div_geometries = list(geometries)
- ops.append(OperatorInfo(f"P2ToP1", f"Div_{c}",
- TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)), TestSpace(P1),
- form=partial(divergence, component_index=c),
- type_descriptor=type_descriptor, opts=opts, geometries=div_geometries,
- blending=blending))
- ops.append(OperatorInfo(f"P1ToP2", f"DivT_{c}", TrialSpace(P1),
- TestSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
- form=partial(gradient, component_index=c),
- type_descriptor=type_descriptor, opts=opts, geometries=div_geometries,
- blending=blending))
- # fmt: on
-
- for c in [0, 1]:
- for r in [0, 1]:
- p2_epsilon = partial(
- epsilon,
- variable_viscosity=True,
- coefficient_function_space=P2,
- component_trial=c,
- component_test=r,
- )
-
- p2_full_stokes = partial(
- full_stokes,
- variable_viscosity=True,
- coefficient_function_space=P2,
- component_trial=c,
- component_test=r,
- )
- # fmt: off
- ops.append(
- OperatorInfo(f"P2", f"Epsilon_{r}_{c}",
- TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
- TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)), form=p2_epsilon,
- type_descriptor=type_descriptor, geometries=list(geometries), opts=opts,
- blending=blending))
- ops.append(OperatorInfo(f"P2", f"FullStokes_{r}_{c}",
- TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
- TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)),
- form=p2_full_stokes, type_descriptor=type_descriptor, geometries=list(geometries),
- opts=opts, blending=blending))
- # fmt: on
- for c, r in [(0, 2), (1, 2), (2, 2), (2, 1), (2, 0)]:
- p2_epsilon = partial(
- epsilon,
- variable_viscosity=True,
- coefficient_function_space=P2,
- component_trial=c,
- component_test=r,
- )
-
- p2_full_stokes = partial(
- full_stokes,
- variable_viscosity=True,
- coefficient_function_space=P2,
- component_trial=c,
- component_test=r,
- )
- # fmt: off
- ops.append(
- OperatorInfo(f"P2", f"Epsilon_{r}_{c}", TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
- TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)), form=p2_epsilon,
- type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
- ops.append(
- OperatorInfo(f"P2", f"FullStokes_{r}_{c}", TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
- TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)), form=p2_full_stokes,
- type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
- # fmt: on
-
- # Removing all operators without viable element types (e.g.: some ops only support 2D, but a blending map maybe only
- # supports 3D).
- ops = [o for o in ops if o.geometries]
-
- return ops
-
-
-def generate_elementwise_op(
- args: argparse.Namespace,
- symbolizer: Symbolizer,
- op_info: OperatorInfo,
- name: str,
- optimizations: Set[Opts],
- loop_strategy: LoopStrategy,
- blending: GeometryMap,
- quad_info: dict[ElementGeometry, int | str],
- quad_info_boundary: dict[ElementGeometry, int | str],
- type_descriptor: HOGType,
-) -> None:
- """Generates a single operator and writes it to the HyTeG directory."""
-
- operator = HyTeGElementwiseOperator(
- name,
- symbolizer,
- kernel_wrapper_types=op_info.kernel_types,
- type_descriptor=type_descriptor,
- )
-
- for geometry in op_info.geometries:
- # Is this necessary? Currently, the decision of the geometry is all over the place.
- if geometry not in blending.supported_geometries():
- continue
-
- if op_info.form is not None:
- quad = Quadrature(
- select_quadrule(quad_info[geometry], geometry),
- geometry,
- )
-
- form = op_info.form(
- op_info.trial_space,
- op_info.test_space,
- geometry,
- symbolizer,
- blending=blending, # type: ignore[call-arg] # kw-args are not supported by Callable
- )
-
- operator.add_volume_integral(
- name="".join(name.split()),
- volume_geometry=geometry,
- quad=quad,
- blending=blending,
- form=form,
- loop_strategy=loop_strategy,
- optimizations=optimizations,
- )
-
- if op_info.form_boundary is not None:
- boundary_geometry: ElementGeometry
- if geometry == TriangleElement():
- boundary_geometry = LineElement(space_dimension=2)
- elif geometry == TetrahedronElement():
- boundary_geometry = TriangleElement(space_dimension=3)
- else:
- raise HOGException("Invalid volume geometry for boundary integral.")
-
- quad = Quadrature(
- select_quadrule(quad_info_boundary[geometry], boundary_geometry),
- boundary_geometry,
- )
-
- form_boundary = op_info.form_boundary(
- op_info.trial_space,
- op_info.test_space,
- geometry,
- boundary_geometry,
- symbolizer,
- blending=blending, # type: ignore[call-arg] # kw-args are not supported by Callable
- )
-
- operator.add_boundary_integral(
- name="".join(name.split()),
- volume_geometry=geometry,
- quad=quad,
- blending=blending,
- form=form_boundary,
- optimizations=set(),
- )
-
- dir_path = os.path.join(args.output, op_info.name.split("_")[0])
- operator.generate_class_code(
- dir_path,
- clang_format_binary=args.clang_format_binary,
- )
-
-
-if __name__ == "__main__":
- main()
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+import argparse
+import sys
+from argparse import ArgumentParser
+from dataclasses import dataclass, field
+from functools import partial
+import logging
+import os
+import re
+from typing import Callable, List, Optional, Sequence, Set, Tuple, Union
+
+import numpy as np
+import sympy as sp
+from sympy.core.cache import clear_cache
+from tabulate import tabulate
+
+from hog.blending import GeometryMap, IdentityMap, AnnulusMap, IcosahedralShellMap, ParametricMap
+from hog.cse import CseImplementation
+from hog.element_geometry import (
+ ElementGeometry,
+ TriangleElement,
+ TetrahedronElement,
+ LineElement,
+)
+from hog.exception import HOGException
+from hog.forms import (
+ diffusion,
+ divergence,
+ gradient,
+ div_k_grad,
+ shear_heating,
+ epsilon,
+ full_stokes,
+ nonlinear_diffusion,
+ nonlinear_diffusion_newton_galerkin,
+ supg_diffusion,
+ advection,
+ supg_advection,
+ grad_rho_by_rho_dot_u,
+ linear_elasticity,
+)
+from hog.forms_boundary import (
+ mass_boundary,
+ freeslip_gradient_weak_boundary,
+ freeslip_divergence_weak_boundary,
+ freeslip_momentum_weak_boundary,
+)
+from hog.forms_vectorial import curl_curl, curl_curl_plus_mass, mass_n1e1
+from hog.function_space import (
+ LagrangianFunctionSpace,
+ N1E1Space,
+ TensorialVectorFunctionSpace,
+ TrialSpace,
+ TestSpace,
+)
+from hog.logger import get_logger, TimedLogger
+from hog.operator_generation.kernel_types import (
+ ApplyWrapper,
+ AssembleWrapper,
+ AssembleDiagonalWrapper,
+ KernelWrapperType,
+)
+from hog.operator_generation.loop_strategies import (
+ LoopStrategy,
+ SAWTOOTH,
+ CUBES,
+ FUSEDROWS,
+)
+from hog.operator_generation.operators import (
+ HyTeGElementwiseOperator,
+ MacroIntegrationDomain,
+)
+
+from operator import itemgetter
+from hog.operator_generation.optimizer import (
+ Opts,
+ opts_arg_mapping,
+ ordered_opts_suffix,
+)
+from hog.quadrature import Quadrature
+from hog.symbolizer import Symbolizer
+from generate_all_hyteg_forms import valid_base_dir
+from hog.operator_generation.types import parse_argument_type, HOGType, HOGPrecision
+import quadpy
+from hog.quadrature.quadrature import select_quadrule
+from hog.integrand import Form
+
+ALL_GEOMETRY_MAPS = [
+ IdentityMap(),
+ AnnulusMap(),
+ IcosahedralShellMap(),
+ ParametricMap(1),
+ ParametricMap(2),
+]
+
+
+def parse_arguments():
+ parser = ArgumentParser(
+ description="Generates a collection of elementwise operators."
+ )
+
+ list_or_gen = parser.add_mutually_exclusive_group(required=True)
+ list_or_gen.add_argument(
+ "-l",
+ "--list-operators",
+ action="store_true",
+ help="List all available operators.",
+ )
+ list_or_gen.add_argument(
+ "--list-quadratures",
+ nargs=1,
+ help="List all available quadrature rules for a given geometry (t1: Line, t2: Triangle, t3: Tetrahedron)",
+ )
+
+ list_or_gen.add_argument(
+ "--hyteg",
+ help=(
+ "Path to the HyTeG repository. The generated code is written to the correct directory in the source tree."
+ ),
+ )
+ list_or_gen.add_argument(
+ "--output",
+ help=("Path to output directory."),
+ )
+
+ parser.add_argument(
+ "-s",
+ "--space-mapping",
+ default=".*?",
+ help="Filter by function space, e.g. 'P1' or 'P1toP2' (regex).",
+ # TODO add choices list and type=string.lower
+ )
+
+ parser.add_argument(
+ "-f",
+ "--form",
+ default=".*?",
+ help="Filter by form, e.g. 'CurlCurl' or 'Diffusion' (regex).",
+ # TODO add choices list and type=string.lower
+ )
+
+ parser.add_argument(
+ "-o",
+ "--opts",
+ default="",
+ nargs="+",
+ help=f"""
+ Enter a space separated list of optimizations.
+ Not all combinations are currently supported. Available optimizations are: {Opts.__members__.keys()}""",
+ # TODO add choices list and type=string.lower
+ )
+ # TODO explanations for optimizations e.g. CUTLOOPS is not straight-forward
+
+ parser.add_argument(
+ "--loop-strategy",
+ default="SAWTOOTH",
+ help="Enter a loop strategy for the operator. Available strategies are SAWTOOTH, CUBES.",
+ # TODO add choices list and type=string.lower
+ )
+
+ quad_rule_or_degree = parser.add_mutually_exclusive_group(required=False)
+ quad_rule_or_degree.add_argument(
+ "--quad-rule",
+ type=str,
+ nargs=2,
+ default=None,
+ help="""Supply two rules by their identifier (leftmost column when calling with --list-quadratures),
+ the first will be used for triangles and the second for tetrahedra.""",
+ )
+ quad_rule_or_degree.add_argument(
+ "--quad-degree",
+ type=int,
+ nargs=2,
+ default=[-1, -1],
+ help="""Supply two polynomial degrees (as integers) up to which the quadrature should be exact.
+ The generator will chose the rule with minimal points that integrates the specified
+ degree exactly. The first integer will be used for triangles and the second for tetrahedra.""",
+ )
+
+ prec_or_mpfr = parser.add_mutually_exclusive_group(required=False)
+ prec_or_mpfr.add_argument(
+ "-p",
+ "--precision",
+ type=str.lower,
+ help=(
+ "Defines the precision in which the kernels are calculated; "
+ "If not specified we choose fp64(double) precision"
+ ),
+ choices=list(choice.value for choice in HOGPrecision),
+ default="fp64",
+ )
+
+ parser.add_argument(
+ "-b",
+ "--blending",
+ default=str(IdentityMap()),
+ help=f"Enter a blending type for the operator. Note some geometry maps only work in 2D or 3D respectively. "
+ f"Available geometry maps: {[str(m) for m in ALL_GEOMETRY_MAPS]}",
+ # TODO add choices list and type=string.lower
+ )
+
+ parser.add_argument(
+ "--name",
+ help=f"Specify the name of the generated C++ class.",
+ )
+
+ parser.add_argument(
+ "--dimensions",
+ type=int,
+ nargs="+",
+ default=[2, 3],
+ help=f"Domain dimensions for which operators shall be generated.",
+ )
+
+ logging_group = parser.add_mutually_exclusive_group(required=False)
+ logging_group.add_argument(
+ "-v",
+ "--verbose",
+ type=int,
+ help="specify verbosity level (0 = silent, 1 = info, 2 = debug)",
+ choices=[0, 1, 2],
+ )
+ logging_group.add_argument(
+ "--loglevel",
+ type=str.lower,
+ help=(
+ "Defines which messages are shown and which not; "
+ "If not specified 'warning' is chosen."
+ ),
+ choices=["debug", "info", "warning", "error", "critical"],
+ )
+
+ clang_format_group = parser.add_mutually_exclusive_group(required=False)
+ clang_format_group.add_argument(
+ "--no-clang-format",
+ dest="clang_format",
+ action="store_false",
+ help="Do not apply clang-format on generated files.",
+ )
+
+ clang_format_group.add_argument(
+ "--clang-format-binary",
+ default="clang-format",
+ help=f"Allows to specify the name of the clang-format binary and/or optionally its full path."
+ f" By default 'clang-format' will be used.",
+ )
+
+ args = parser.parse_args()
+ return parser, args
+
+
+strategy_args_mapping = {
+ "CUBES": CUBES(),
+ "SAWTOOTH": SAWTOOTH(),
+ "FUSEDROWS": FUSEDROWS(),
+}
+
+
+def main():
+ clear_cache()
+
+ parser, args = parse_arguments()
+
+ # Adapt clang_format_binary according to the value of clang_format
+ if "clang_format" in vars(args):
+ if not args.clang_format:
+ args.clang_format_binary = None
+
+ if args.verbose:
+ if args.verbose == 0:
+ log_level = logging.CRITICAL
+ elif args.verbose == 1:
+ log_level = logging.INFO
+ elif args.verbose == 2:
+ log_level = logging.DEBUG
+ else:
+ raise HOGException(f"Invalid verbosity level {args.verbose}.")
+ elif args.loglevel:
+ if args.loglevel == "debug":
+ log_level = logging.DEBUG
+ elif args.loglevel == "info":
+ log_level = logging.INFO
+ elif args.loglevel == "warning":
+ log_level = logging.WARNING
+ elif args.loglevel == "error":
+ log_level = logging.ERROR
+ elif args.loglevel == "critical":
+ log_level = logging.CRITICAL
+ else:
+ raise HOGException(f"Invalid log level {args.loglevel}.")
+ else:
+ # The default log level is chosen here
+ log_level = logging.DEBUG
+
+ logger = get_logger(log_level)
+ TimedLogger.set_log_level(log_level)
+
+ logger.info(f"Running {__file__} with arguments: {' '.join(sys.argv[1:])}")
+
+ symbolizer = Symbolizer()
+
+ if args.list_quadratures:
+ geometry = args.list_quadratures[0]
+ logger.debug(f"Listing quadrature rules for geometry: {geometry}.")
+ if geometry == "t3":
+ schemes = quadpy.t3.schemes
+ elif geometry == "t2":
+ schemes = quadpy.t2.schemes
+ else:
+ raise HOGException(f"Unexpected geometry: {geometry}")
+
+ all_schemes = []
+ for key, s in schemes.items():
+ try:
+ scheme = s()
+ entry = {
+ "key": key,
+ "name": scheme.name,
+ "points": scheme.points.shape[1],
+ "degree": scheme.degree,
+ "test_tolerance": scheme.test_tolerance,
+ }
+ all_schemes.append(entry)
+ except TypeError as e:
+ # Some schemes seem to require parameters, we simply exclude these with this litte try-except hack
+ pass
+
+ all_schemes = sorted(all_schemes, key=itemgetter("degree", "points"))
+ table = tabulate(all_schemes, headers="keys", tablefmt="grid")
+ print(table)
+ exit()
+
+ blending = None
+ for m in ALL_GEOMETRY_MAPS:
+ if str(m) == args.blending:
+ blending = m
+ if blending is None:
+ raise HOGException("Invalid geometry map.")
+ logger.debug(f"Blending: {blending}")
+
+ type_descriptor = parse_argument_type(args)
+
+ logger.debug(f"Data type name is {type_descriptor.pystencils_type}.")
+
+ opts = {opts_arg_mapping[opt] for opt in args.opts}
+ loop_strategy = strategy_args_mapping[args.loop_strategy]
+ re_form = re.compile(args.form)
+
+ if (
+ not args.list_operators
+ and args.quad_degree == [-1, -1]
+ and args.quad_rule is None
+ ):
+ parser.error("Either --quad-degree or --quad-rule are required.")
+ quad = {
+ geometry: info
+ for geometry, info in zip(
+ [TriangleElement(), TetrahedronElement()],
+ args.quad_degree if args.quad_rule is None else args.quad_rule,
+ )
+ }
+
+ enabled_geometries: Set[TriangleElement | TetrahedronElement] = set()
+ if 2 in args.dimensions:
+ enabled_geometries.add(TriangleElement())
+ if 3 in args.dimensions:
+ enabled_geometries.add(TetrahedronElement())
+
+ operators = all_operators(
+ symbolizer,
+ [(opts, loop_strategy, ordered_opts_suffix(loop_strategy, opts))],
+ type_descriptor,
+ blending,
+ geometries=enabled_geometries,
+ )
+
+ filtered_operators = list(
+ filter(
+ lambda o: re.fullmatch(args.space_mapping, o.mapping, re.IGNORECASE)
+ and re.fullmatch(args.form, o.name, re.IGNORECASE),
+ operators,
+ )
+ )
+
+ filtered_operators = [f for f in filtered_operators if f.geometries]
+
+ if len(filtered_operators) == 0:
+ raise HOGException(
+ f"Form {args.form} does not match any available forms. Run --list for a list."
+ )
+
+ if args.list_operators:
+ print(
+ tabulate(
+ sorted(
+ [
+ (
+ o.mapping,
+ o.name,
+ ", ".join(
+ str(geo.dimensions) + "D" for geo in o.geometries
+ ),
+ len(o.opts),
+ )
+ for o in filtered_operators
+ ]
+ ),
+ headers=(
+ "Space mapping",
+ "Form",
+ "Dimensions",
+ "Optimizations",
+ ),
+ )
+ )
+ exit()
+
+ if args.hyteg:
+ args.output = os.path.join(args.hyteg, "src/hyteg/operatorgeneration/generated")
+ if not valid_base_dir(args.hyteg):
+ raise HOGException(
+ "The specified directory does not seem to be the HyTeG directory."
+ )
+
+ for operator in filtered_operators:
+ for opt, loop, opt_name in operator.opts:
+ blending_str = (
+ "_" + str(blending) if not isinstance(blending, IdentityMap) else ""
+ )
+ if args.name:
+ name = args.name
+ else:
+ name = (
+ f"{operator.mapping}Elementwise{operator.name}{blending_str}{opt_name}"
+ f"{'_' + str(type_descriptor) if type_descriptor.add_file_suffix else ''}"
+ )
+ with TimedLogger(f"Generating {name}", level=logging.INFO):
+ generate_elementwise_op(
+ args,
+ symbolizer,
+ operator,
+ name,
+ opt,
+ loop,
+ blending,
+ quad,
+ quad,
+ type_descriptor=type_descriptor,
+ )
+
+
+def all_opts_sympy_cse() -> List[Tuple[Set[Opts], LoopStrategy, str]]:
+ return [
+ (
+ set(
+ {
+ Opts.MOVECONSTANTS,
+ Opts.REORDER,
+ Opts.ELIMTMPS,
+ Opts.QUADLOOPS,
+ Opts.VECTORIZE,
+ }
+ ),
+ SAWTOOTH(),
+ "_const_vect_fused_quadloops_reordered_elimtmps",
+ ),
+ ]
+
+
+def all_opts_poly_cse() -> List[Tuple[Set[Opts], LoopStrategy, str]]:
+ return [
+ (o | {Opts.POLYCSE}, l, n + "_polycse") for (o, l, n) in all_opts_sympy_cse()
+ ]
+
+
+def all_opts_both_cses() -> List[Tuple[Set[Opts], LoopStrategy, str]]:
+ return all_opts_sympy_cse() + all_opts_poly_cse()
+
+
+@dataclass
+class OperatorInfo:
+ mapping: str
+ name: str
+ trial_space: TrialSpace
+ test_space: TestSpace
+ form: (
+ Callable[
+ [
+ TrialSpace,
+ TestSpace,
+ ElementGeometry,
+ Symbolizer,
+ GeometryMap,
+ ],
+ Form,
+ ]
+ | None
+ )
+ type_descriptor: HOGType
+ form_boundary: (
+ Callable[
+ [
+ TrialSpace,
+ TestSpace,
+ ElementGeometry,
+ ElementGeometry,
+ Symbolizer,
+ GeometryMap,
+ ],
+ Form,
+ ]
+ | None
+ ) = None
+ kernel_types: List[KernelWrapperType] = None # type: ignore[assignment] # will definitely be initialized in __post_init__
+ geometries: Sequence[ElementGeometry] = field(
+ default_factory=lambda: [TriangleElement(), TetrahedronElement()]
+ )
+ opts: List[Tuple[Set[Opts], LoopStrategy, str]] = field(
+ default_factory=all_opts_sympy_cse
+ )
+ blending: GeometryMap = field(default_factory=lambda: IdentityMap())
+
+ def __post_init__(self):
+ # Removing geometries that are not supported by blending.
+ # I don't like this, but it looks like the collection of operators has to be refactored anyway.
+ self.geometries = list(
+ set(self.geometries) & set(self.blending.supported_geometries())
+ )
+
+ if self.kernel_types is None:
+ dims = [g.dimensions for g in self.geometries]
+ self.kernel_types = [
+ ApplyWrapper(
+ self.trial_space,
+ self.test_space,
+ type_descriptor=self.type_descriptor,
+ dims=dims,
+ )
+ ]
+
+ all_opts = set().union(*[o for (o, _, _) in self.opts])
+ if not ({Opts.VECTORIZE, Opts.VECTORIZE512}.intersection(all_opts)):
+ self.kernel_types.append(
+ AssembleWrapper(
+ self.trial_space,
+ self.test_space,
+ type_descriptor=self.type_descriptor,
+ dims=dims,
+ )
+ )
+
+ if self.test_space == self.trial_space:
+ self.kernel_types.append(
+ AssembleDiagonalWrapper(
+ self.trial_space,
+ type_descriptor=self.type_descriptor,
+ dims=dims,
+ )
+ )
+
+
+def all_operators(
+ symbolizer: Symbolizer,
+ opts: List[Tuple[Set[Opts], LoopStrategy, str]],
+ type_descriptor: HOGType,
+ blending: GeometryMap,
+ geometries: Set[Union[TriangleElement, TetrahedronElement]],
+) -> List[OperatorInfo]:
+ P1 = LagrangianFunctionSpace(1, symbolizer)
+ P1Vector = TensorialVectorFunctionSpace(P1)
+ P2 = LagrangianFunctionSpace(2, symbolizer)
+ P2Vector = TensorialVectorFunctionSpace(P2)
+ N1E1 = N1E1Space(symbolizer)
+
+ two_d = list(geometries & {TriangleElement()})
+ three_d = list(geometries & {TetrahedronElement()})
+
+ ops: List[OperatorInfo] = []
+
+ # fmt: off
+ # TODO switch to manual specification of opts for now/developement, later use default factory
+ ops.append(OperatorInfo("N1E1", "CurlCurl", TrialSpace(N1E1), TestSpace(N1E1), form=curl_curl,
+ type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
+ ops.append(OperatorInfo("N1E1", "Mass", TrialSpace(N1E1), TestSpace(N1E1), form=mass_n1e1,
+ type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
+ ops.append(OperatorInfo("N1E1", "CurlCurlPlusMass", TrialSpace(N1E1), TestSpace(N1E1),
+ form=partial(curl_curl_plus_mass, alpha_fem_space=P1, beta_fem_space=P1),
+ type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
+ ops.append(OperatorInfo("P1", "Diffusion", TrialSpace(P1), TestSpace(P1), form=diffusion,
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+ ops.append(OperatorInfo("P1", "DivKGrad", TrialSpace(P1), TestSpace(P1),
+ form=partial(div_k_grad, coefficient_function_space=P1),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P2", "Diffusion", TrialSpace(P2), TestSpace(P2), form=diffusion,
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+ ops.append(OperatorInfo("P2", "DivKGrad", TrialSpace(P2), TestSpace(P2),
+ form=partial(div_k_grad, coefficient_function_space=P2),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P2", "ShearHeating", TrialSpace(P2), TestSpace(P2),
+ form=partial(shear_heating, viscosity_function_space=P2, velocity_function_space=P2),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P1", "NonlinearDiffusion", TrialSpace(P1), TestSpace(P1),
+ form=partial(nonlinear_diffusion, coefficient_function_space=P1),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+ ops.append(OperatorInfo("P1", "NonlinearDiffusionNewtonGalerkin", TrialSpace(P1),
+ TestSpace(P1), form=partial(nonlinear_diffusion_newton_galerkin,
+ coefficient_function_space=P1, only_newton_galerkin_part_of_form=False),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P1Vector", "Diffusion", TrialSpace(P1Vector), TestSpace(P1Vector),
+ form=diffusion, type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P2", "SUPGDiffusion", TrialSpace(P2), TestSpace(P2),
+ form=partial(supg_diffusion, velocity_function_space=P2, diffusivityXdelta_function_space=P2),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P2", "AdvectionConstCp", TrialSpace(P2), TestSpace(P2),
+ form=partial(advection, velocity_function_space=P2, coefficient_function_space=P2, constant_cp = True),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P2", "AdvectionVarCp", TrialSpace(P2), TestSpace(P2),
+ form=partial(advection, velocity_function_space=P2, coefficient_function_space=P2),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P2", "SUPGAdvection", TrialSpace(P2), TestSpace(P2),
+ form=partial(supg_advection, velocity_function_space=P2, coefficient_function_space=P2),
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P2", "BoundaryMass", TrialSpace(P2), TestSpace(P2), form=None,
+ form_boundary=mass_boundary, type_descriptor=type_descriptor, geometries=list(geometries),
+ opts=opts, blending=blending))
+
+ ops.append(OperatorInfo("P1Vector", "LinearElasticity", TrialSpace(P1Vector), TestSpace(P1Vector), form=linear_elasticity
+ , form_boundary=mass_boundary, type_descriptor=type_descriptor,
+ geometries=list(geometries), opts=opts, blending=blending))
+
+ # fmt: on
+
+ p2vec_epsilon = partial(
+ epsilon,
+ variable_viscosity=True,
+ coefficient_function_space=P2,
+ )
+ ops.append(
+ OperatorInfo(
+ f"P2Vector",
+ f"Epsilon",
+ TrialSpace(P2Vector),
+ TestSpace(P2Vector),
+ form=p2vec_epsilon,
+ type_descriptor=type_descriptor,
+ geometries=list(geometries),
+ opts=opts,
+ blending=blending,
+ )
+ )
+
+ p2vec_freeslip_momentum_weak_boundary = partial(
+ freeslip_momentum_weak_boundary, function_space_mu=P2
+ )
+
+ ops.append(
+ OperatorInfo(
+ f"P2Vector",
+ f"EpsilonFreeslip",
+ TrialSpace(P2Vector),
+ TestSpace(P2Vector),
+ form=p2vec_epsilon,
+ form_boundary=p2vec_freeslip_momentum_weak_boundary,
+ type_descriptor=type_descriptor,
+ geometries=list(geometries),
+ opts=opts,
+ blending=blending,
+ )
+ )
+
+ ops.append(
+ OperatorInfo(
+ f"P2VectorToP1",
+ f"DivergenceFreeslip",
+ TrialSpace(P2Vector),
+ TestSpace(P1),
+ form=divergence,
+ form_boundary=freeslip_divergence_weak_boundary,
+ type_descriptor=type_descriptor,
+ geometries=list(geometries),
+ opts=opts,
+ blending=blending,
+ )
+ )
+
+ ops.append(
+ OperatorInfo(
+ f"P1ToP2Vector",
+ f"GradientFreeslip",
+ TrialSpace(P1),
+ TestSpace(P2Vector),
+ form=gradient,
+ form_boundary=freeslip_gradient_weak_boundary,
+ type_descriptor=type_descriptor,
+ geometries=list(geometries),
+ opts=opts,
+ blending=blending,
+ )
+ )
+
+ p2vec_grad_rho = partial(
+ grad_rho_by_rho_dot_u,
+ density_function_space=P2,
+ )
+ ops.append(
+ OperatorInfo(
+ f"P2VectorToP1",
+ f"GradRhoByRhoDotU",
+ TrialSpace(P1),
+ TestSpace(P2Vector),
+ form=p2vec_grad_rho,
+ type_descriptor=type_descriptor,
+ geometries=list(geometries),
+ opts=opts,
+ blending=blending,
+ )
+ )
+
+ for c in [0, 1, 2]:
+ # fmt: off
+ if c == 2:
+ div_geometries = three_d
+ else:
+ div_geometries = list(geometries)
+ ops.append(OperatorInfo(f"P2ToP1", f"Div_{c}",
+ TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)), TestSpace(P1),
+ form=partial(divergence, component_index=c),
+ type_descriptor=type_descriptor, opts=opts, geometries=div_geometries,
+ blending=blending))
+ ops.append(OperatorInfo(f"P1ToP2", f"DivT_{c}", TrialSpace(P1),
+ TestSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
+ form=partial(gradient, component_index=c),
+ type_descriptor=type_descriptor, opts=opts, geometries=div_geometries,
+ blending=blending))
+ # fmt: on
+
+ for c in [0, 1]:
+ for r in [0, 1]:
+ p2_epsilon = partial(
+ epsilon,
+ variable_viscosity=True,
+ coefficient_function_space=P2,
+ component_trial=c,
+ component_test=r,
+ )
+
+ p2_full_stokes = partial(
+ full_stokes,
+ variable_viscosity=True,
+ coefficient_function_space=P2,
+ component_trial=c,
+ component_test=r,
+ )
+ # fmt: off
+ ops.append(
+ OperatorInfo(f"P2", f"Epsilon_{r}_{c}",
+ TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
+ TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)), form=p2_epsilon,
+ type_descriptor=type_descriptor, geometries=list(geometries), opts=opts,
+ blending=blending))
+ ops.append(OperatorInfo(f"P2", f"FullStokes_{r}_{c}",
+ TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
+ TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)),
+ form=p2_full_stokes, type_descriptor=type_descriptor, geometries=list(geometries),
+ opts=opts, blending=blending))
+ # fmt: on
+ for c, r in [(0, 2), (1, 2), (2, 2), (2, 1), (2, 0)]:
+ p2_epsilon = partial(
+ epsilon,
+ variable_viscosity=True,
+ coefficient_function_space=P2,
+ component_trial=c,
+ component_test=r,
+ )
+
+ p2_full_stokes = partial(
+ full_stokes,
+ variable_viscosity=True,
+ coefficient_function_space=P2,
+ component_trial=c,
+ component_test=r,
+ )
+ # fmt: off
+ ops.append(
+ OperatorInfo(f"P2", f"Epsilon_{r}_{c}", TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
+ TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)), form=p2_epsilon,
+ type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
+ ops.append(
+ OperatorInfo(f"P2", f"FullStokes_{r}_{c}", TrialSpace(TensorialVectorFunctionSpace(P2, single_component=c)),
+ TestSpace(TensorialVectorFunctionSpace(P2, single_component=r)), form=p2_full_stokes,
+ type_descriptor=type_descriptor, geometries=three_d, opts=opts, blending=blending))
+ # fmt: on
+
+ # Removing all operators without viable element types (e.g.: some ops only support 2D, but a blending map maybe only
+ # supports 3D).
+ ops = [o for o in ops if o.geometries]
+
+ return ops
+
+
+def generate_elementwise_op(
+ args: argparse.Namespace,
+ symbolizer: Symbolizer,
+ op_info: OperatorInfo,
+ name: str,
+ optimizations: Set[Opts],
+ loop_strategy: LoopStrategy,
+ blending: GeometryMap,
+ quad_info: dict[ElementGeometry, int | str],
+ quad_info_boundary: dict[ElementGeometry, int | str],
+ type_descriptor: HOGType,
+) -> None:
+ """Generates a single operator and writes it to the HyTeG directory."""
+
+ operator = HyTeGElementwiseOperator(
+ name,
+ symbolizer,
+ kernel_wrapper_types=op_info.kernel_types,
+ type_descriptor=type_descriptor,
+ )
+
+ for geometry in op_info.geometries:
+ # Is this necessary? Currently, the decision of the geometry is all over the place.
+ if geometry not in blending.supported_geometries():
+ continue
+
+ if op_info.form is not None:
+ quad = Quadrature(
+ select_quadrule(quad_info[geometry], geometry),
+ geometry,
+ )
+
+ form = op_info.form(
+ op_info.trial_space,
+ op_info.test_space,
+ geometry,
+ symbolizer,
+ blending=blending, # type: ignore[call-arg] # kw-args are not supported by Callable
+ )
+
+ operator.add_volume_integral(
+ name="".join(name.split()),
+ volume_geometry=geometry,
+ quad=quad,
+ blending=blending,
+ form=form,
+ loop_strategy=loop_strategy,
+ optimizations=optimizations,
+ )
+
+ if op_info.form_boundary is not None:
+ boundary_geometry: ElementGeometry
+ if geometry == TriangleElement():
+ boundary_geometry = LineElement(space_dimension=2)
+ elif geometry == TetrahedronElement():
+ boundary_geometry = TriangleElement(space_dimension=3)
+ else:
+ raise HOGException("Invalid volume geometry for boundary integral.")
+
+ quad = Quadrature(
+ select_quadrule(quad_info_boundary[geometry], boundary_geometry),
+ boundary_geometry,
+ )
+
+ form_boundary = op_info.form_boundary(
+ op_info.trial_space,
+ op_info.test_space,
+ geometry,
+ boundary_geometry,
+ symbolizer,
+ blending=blending, # type: ignore[call-arg] # kw-args are not supported by Callable
+ )
+
+ operator.add_boundary_integral(
+ name="".join(name.split()),
+ volume_geometry=geometry,
+ quad=quad,
+ blending=blending,
+ form=form_boundary,
+ optimizations=set(),
+ )
+
+ dir_path = os.path.join(args.output, op_info.name.split("_")[0])
+ operator.generate_class_code(
+ dir_path,
+ clang_format_binary=args.clang_format_binary,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/hog/forms.py b/hog/forms.py
index 73a03f83620437b35ed6bcf9ca2ef285ed41f91e..58e7ec47f8cb8e84464bc32afb670c641493122b 100644
--- a/hog/forms.py
+++ b/hog/forms.py
@@ -1,1180 +1,1180 @@
-# HyTeG Operator Generator
-# Copyright (C) 2024 HyTeG Team
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program. If not, see <https://www.gnu.org/licenses/>.
-
-import logging
-import sympy as sp
-from typing import Optional, Callable, Any
-
-from hog.element_geometry import ElementGeometry
-from hog.exception import HOGException
-from hog.fem_helpers import (
- trafo_ref_to_affine,
- jac_ref_to_affine,
- jac_affine_to_physical,
- hessian_shape_affine_ref_pullback,
- hessian_shape_blending_ref_pullback,
- create_empty_element_matrix,
- element_matrix_iterator,
- scalar_space_dependent_coefficient,
- vector_space_dependent_coefficient,
- fem_function_on_element,
- fem_function_gradient_on_element,
-)
-from hog.function_space import FunctionSpace, N1E1Space, TrialSpace, TestSpace
-from hog.math_helpers import dot, inv, abs, det, double_contraction
-from hog.quadrature import Quadrature, Tabulation
-from hog.symbolizer import Symbolizer
-from hog.logger import TimedLogger
-from hog.blending import GeometryMap, ExternalMap, IdentityMap
-from hog.integrand import process_integrand, Form
-
-
-def diffusion(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-Diffusion operator without coefficients.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (space: {trial})
- v: test function (space: {test})
-
- ∫ ∇u : ∇v
-
- Note that the double contraction (:) reduces to the dot product for scalar function spaces, i.e. the form becomes
-
- ∫ ∇u · ∇v
-"""
-
- from hog.recipes.integrands.volume.diffusion import integrand
- from hog.recipes.integrands.volume.diffusion_affine import (
- integrand as integrand_affine,
- )
-
- # mypy type checking supposedly cannot figure out ternaries.
- # https://stackoverflow.com/a/70832391
- integr: Callable[..., Any] = integrand
- if blending == IdentityMap():
- integr = integrand_affine
-
- return process_integrand(
- integr,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- )
-
-
-def mass(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-Mass operator.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (space: {trial})
- v: test function (space: {test})
-
- ∫ uv
-"""
-
- from hog.recipes.integrands.volume.mass import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- )
-
-
-def div_k_grad(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- coefficient_function_space: Optional[FunctionSpace] = None,
-) -> Form:
- docstring = f"""
-Diffusion operator with a scalar coefficient.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (space: {trial})
- v: test function (space: {test})
- k: coefficient (space: {coefficient_function_space})
-
- ∫ k ∇u · ∇v
-"""
-
- from hog.recipes.integrands.volume.div_k_grad import integrand
- from hog.recipes.integrands.volume.div_k_grad_affine import (
- integrand as integrand_affine,
- )
-
- # mypy type checking supposedly cannot figure out ternaries.
- # https://stackoverflow.com/a/70832391
- integr: Callable[..., Any] = integrand
- if blending == IdentityMap():
- integr = integrand_affine
-
- return process_integrand(
- integr,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- fe_coefficients={"k": coefficient_function_space},
- is_symmetric=trial == test,
- docstring=docstring,
- )
-
-
-def nonlinear_diffusion(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- coefficient_function_space: FunctionSpace,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-Diffusion operator with coefficient function.
-
-Weak formulation
-
- u: trial function (space: {trial})
- v: test function (space: {test})
- c: FE coefficient function (space: {coefficient_function_space})
-
- ∫ a(c) ∇u · ∇v
-
-Note: :math:`a(c) = 1/8 + u^2` is currently hard-coded and the form is intended for :math:`c = u`,
- hence, the naming, but will work for other coefficients, too.
-"""
-
- if blending != IdentityMap():
- raise HOGException(
- "The nonlinear-diffusion form does currently not support blending."
- )
-
- if trial != test:
- raise HOGException(
- "Trial space must be equal to test space to assemble non-linear diffusion matrix."
- )
-
- def integrand(
- *,
- jac_a_inv,
- jac_a_abs_det,
- grad_u,
- grad_v,
- k,
- tabulate,
- **_,
- ):
- a = sp.Rational(1, 8) + k["u"] * k["u"]
- return a * tabulate(
- double_contraction(
- jac_a_inv.T * grad_u,
- jac_a_inv.T * grad_v,
- )
- * jac_a_abs_det
- )
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- fe_coefficients={"u": coefficient_function_space},
- )
-
-
-def nonlinear_diffusion_newton_galerkin(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- coefficient_function_space: FunctionSpace,
- blending: GeometryMap = IdentityMap(),
- only_newton_galerkin_part_of_form: Optional[bool] = True,
-) -> Form:
- docstring = f"""
-
-Bi-linear form for the solution of the non-linear diffusion equation by a Newton-Galerkin approach
-
-Weak formulation
-
- u: trial function (space: {trial})
- v: test function (space: {test})
- k: FE coefficient function (space: {coefficient_function_space})
-
- ∫ a(k) ∇u · ∇v + ∫ a'(k) u ∇k · ∇v
-
-Note: :math:`a(k) = 1/8 + k^2` is currently hard-coded and the form is intended for :math:`k = u`.
-"""
- if trial != test:
- raise HOGException(
- "Trial space must be equal to test space to assemble diffusion matrix."
- )
-
- if blending != IdentityMap():
- raise HOGException(
- "The nonlinear_diffusion_newton_galerkin form does currently not support blending."
- )
-
- def integrand(
- *,
- jac_a_inv,
- jac_a_abs_det,
- u,
- grad_u,
- grad_v,
- k,
- grad_k,
- tabulate,
- **_,
- ):
- a = sp.Rational(1, 8) + k["k"] * k["k"]
- a_prime = 2 * k["k"]
-
- diffusion_term = a * tabulate(
- dot(jac_a_inv.T * grad_u, jac_a_inv.T * grad_v) * jac_a_abs_det
- )
-
- newton_galerkin_term = (
- a_prime
- * u
- * dot(jac_a_inv.T * grad_k["k"], tabulate(jac_a_inv.T * grad_v))
- * tabulate(jac_a_abs_det)
- )
-
- if only_newton_galerkin_part_of_form:
- return newton_galerkin_term
- else:
- return diffusion_term + newton_galerkin_term
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=False,
- docstring=docstring,
- fe_coefficients={"k": coefficient_function_space},
- )
-
-
-def epsilon(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- component_trial: int = 0,
- component_test: int = 0,
- variable_viscosity: bool = True,
- coefficient_function_space: Optional[FunctionSpace] = None,
-) -> Form:
- docstring = f"""
-"Epsilon" operator.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (vectorial space: {trial})
- v: test function (vectorial space: {test})
- μ: coefficient (scalar space: {coefficient_function_space})
-
- ∫ 2 μ ε(u) : ε(v)
-
-where
-
- ε(w) := (1/2) (∇w + (∇w)ᵀ)
-"""
- if not variable_viscosity:
- raise HOGException("Constant viscosity currently not supported.")
-
- from hog.recipes.integrands.volume.epsilon import integrand
- from hog.recipes.integrands.volume.epsilon_affine import (
- integrand as integrand_affine,
- )
-
- # mypy type checking supposedly cannot figure out ternaries.
- # https://stackoverflow.com/a/70832391
- integr: Callable[..., Any] = integrand
- if blending == IdentityMap():
- integr = integrand_affine
-
- return process_integrand(
- integr,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- fe_coefficients={"mu": coefficient_function_space},
- )
-
-
-def k_mass(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- coefficient_function_space: Optional[FunctionSpace] = None,
-) -> Form:
- docstring = f"""
-Diffusion operator with a scalar coefficient.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (space: {trial})
- v: test function (space: {test})
- k: coefficient (space: {coefficient_function_space})
-
- ∫ k uv
-"""
-
- from hog.recipes.integrands.volume.k_mass import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- )
-
-
-def pspg(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- quad: Quadrature,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-PSPG stabilisation.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (space: {trial})
- v: test function (space: {test})
-
- ∫ τ ∇u · ∇v
-
- where tau is an element-wise factor given by
-
- 2D:
- τ = -CellVolume(triangle) / 5
- 3D:
- τ = -(CellVolume(tetrahedron))**(2/3) / 12
-
-See e.g.
-
- Brezzi, F., & Douglas, J. (1988).
- Stabilized mixed methods for the Stokes problem.
- Numerische Mathematik, 53, 225-235.
-
-or
-
- Hughes, T. J., Franca, L. P., & Balestra, M. (1986).
- A new finite element formulation for computational fluid dynamics: V.
- Circumventing the Babuška-Brezzi condition: A stable Petrov-Galerkin formulation of the Stokes problem accommodating
- equal-order interpolations.
- Computer Methods in Applied Mechanics and Engineering, 59(1), 85-99.
-
-for details.
-"""
-
- from hog.recipes.integrands.volume.pspg import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- )
-
-
-def linear_form(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- quad: Quadrature,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
-) -> sp.Matrix:
- """
- Implements a linear form of type:
-
- (k(x), psi)
-
- where psi a test function and k = k(x) a scalar, external function.
- """
-
- if trial != test:
- raise HOGException(
- "Trial space must be equal to test space to assemble linear form (jep this is weird, but linear forms are implemented as diagonal matrices)."
- )
-
- if quad.is_exact() and isinstance(blending, ExternalMap):
- raise HOGException(
- "Exact integration is not supported for externally defined blending functions."
- )
-
- with TimedLogger("assembling linear form", level=logging.DEBUG):
- jac_affine = jac_ref_to_affine(geometry, symbolizer)
- jac_affine_det = abs(det(jac_affine))
-
- if isinstance(blending, ExternalMap):
- jac_blending = jac_affine_to_physical(geometry, symbolizer)
- else:
- affine_coords = trafo_ref_to_affine(geometry, symbolizer)
- jac_blending = blending.jacobian(affine_coords)
-
- jac_blending_det = abs(det(jac_blending))
-
- mat = create_empty_element_matrix(trial, test, geometry)
-
- it = element_matrix_iterator(trial, test, geometry)
-
- if isinstance(trial, N1E1Space):
- jac_affine_inv_T = inv(jac_affine).T
- jac_blending_inv_T = inv(jac_blending).T
- space_dependent_coefficient = vector_space_dependent_coefficient
- else:
- space_dependent_coefficient = scalar_space_dependent_coefficient
- k = space_dependent_coefficient("k", geometry, symbolizer, blending=blending)
-
- with TimedLogger(
- f"integrating {mat.shape[0]} expressions",
- level=logging.DEBUG,
- ):
- for data in it:
- if data.row == data.col:
- psi = data.test_shape
-
- if isinstance(trial, N1E1Space):
- form = (
- dot(k, jac_blending_inv_T * jac_affine_inv_T * psi)
- * jac_affine_det
- * jac_blending_det
- )
- else:
- form = k * psi * jac_affine_det * jac_blending_det
-
- mat[data.row, data.col] = quad.integrate(form, symbolizer)
-
- return mat
-
-
-def divergence(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- component_index: int = 0,
-) -> Form:
- docstring = f"""
-Divergence.
-
-Component: {component_index}
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (vectorial space: {trial})
- v: test function (scalar space: {test})
-
- ∫ - ( ∇ · u ) v
-"""
-
- from hog.recipes.integrands.volume.divergence import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=False,
- docstring=docstring,
- )
-
-
-def gradient(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- component_index: int = 0,
-) -> Form:
- docstring = f"""
- Gradient.
-
- Component: {component_index}
- Geometry map: {blending}
-
- Weak formulation
-
- u: trial function (scalar space: {trial})
- v: test function (vectorial space: {test})
-
- ∫ - ( ∇ · v ) u
- """
-
- from hog.recipes.integrands.volume.gradient import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=False,
- docstring=docstring,
- )
-
-
-def full_stokes(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- component_trial: int = 0,
- component_test: int = 0,
- variable_viscosity: bool = True,
- coefficient_function_space: Optional[FunctionSpace] = None,
-) -> Form:
- docstring = f"""
-Implements the fully coupled viscous operator of the Stokes problem.
-The latter is the extension of the Epsilon operator to the case where
-the velocity field need not be divergence-free. This is e.g. the case
-in the (truncated) anelastic liquid approximation of mantle convection.
-
-The strong representation of the operator is given by:
-
- - div[ μ (grad(u)+grad(u)ᵀ) ] + 2/3 grad[ μ div(u) ]
-
-Note that the factor 2/3 means that for 2D this is the pseudo-3D form
-of the operator.
-
-Component trial: {component_trial}
-Component test: {component_test}
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (vectorial space: {trial})
- v: test function (vectorial space: {test})
- μ: coefficient (scalar space: {coefficient_function_space})
-
- ∫ μ {{ ( 2 ε(u) : ε(v) ) - (2/3) [ ( ∇ · u ) · ( ∇ · v ) ] }}
-
-where
-
- ε(w) := (1/2) (∇w + (∇w)ᵀ)
-"""
-
- from hog.recipes.integrands.volume.full_stokes import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- fe_coefficients={"mu": coefficient_function_space},
- )
-
-
-def shear_heating(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- component_trial: int = 0,
- component_test: int = 0,
- variable_viscosity: bool = True,
- viscosity_function_space: Optional[FunctionSpace] = None,
- velocity_function_space: Optional[FunctionSpace] = None,
-) -> Form:
- docstring = f"""
-Implements the fully coupled viscous operator for the shear heating term.
-The latter is the extension of the Epsilon operator to the case where
-the velocity field need not be divergence-free. This is e.g. the case
-in the (truncated) anelastic liquid approximation of mantle convection.
-
-https://doi.org/10.1111/j.1365-246X.2009.04413.x
-(3) and (5)
-
-https://doi.org/10.5194/gmd-15-5127-2022
-Listing 2
-
-The strong representation of the operator is given by:
-
- 𝜏(w) : grad(w)
- 2 {{[ μ (grad(w)+grad(w)ᵀ) / 2 ] - 1/dim [ μ div(w) ]I}} : grad(w)
-
-Note that the factor 1/dim means that for 2D this is the pseudo-3D form
-of the operator.
-
-Component trial: {component_trial}
-Component test: {component_test}
-Geometry map: {blending}
-
-Weak formulation
-
- T: trial function (scalar space: {trial})
- s: test function (scalar space: {test})
- μ: coefficient (scalar space: {viscosity_function_space})
- w: velocity (vectorial space: {velocity_function_space})
-
- ∫ {{ 2 {{[ μ (grad(w)+grad(w)ᵀ) / 2 ] - 1/dim [ μ div(w) ]I}} : grad(w) }} T_h s_h
-
-The resulting matrix must be multiplied with a vector of ones to be used as the shear heating term in the RHS
-"""
-
- def integrand(
- *,
- jac_a_inv,
- jac_b_inv,
- jac_a_abs_det,
- jac_b_abs_det,
- u,
- v,
- k,
- grad_k,
- volume_geometry,
- tabulate,
- **_,
- ):
- """First function: mu, other functions: ux, uy, uz."""
-
- mu = k["mu"]
-
- # grad_k[0] is grad_mu_ref
- grad_wx = jac_b_inv.T * jac_a_inv.T * grad_k["wx"]
- grad_wy = jac_b_inv.T * jac_a_inv.T * grad_k["wy"]
- grad_wz = jac_b_inv.T * jac_a_inv.T * grad_k["wz"]
-
- grad_w = grad_wx.row_join(grad_wy)
- dim = volume_geometry.dimensions
- if dim == 3:
- grad_w = grad_w.row_join(grad_wz)
-
- sym_grad_w = 0.5 * (grad_w + grad_w.T)
-
- divdiv = grad_w.trace() * sp.eye(dim)
-
- tau = 2 * (sym_grad_w - sp.Rational(1, dim) * divdiv)
-
- return (
- mu
- * (double_contraction(tau, grad_w)[0])
- * jac_b_abs_det
- * tabulate(jac_a_abs_det * u * v)
- )
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- fe_coefficients={
- "mu": viscosity_function_space,
- "wx": velocity_function_space,
- "wy": velocity_function_space,
- "wz": velocity_function_space,
- },
- is_symmetric=trial == test,
- docstring=docstring,
- )
-
-
-def divdiv(
- trial: TrialSpace,
- test: TestSpace,
- component_trial: int,
- component_test: int,
- geometry: ElementGeometry,
- quad: Quadrature,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-divdiv operator which is used as a stabilization term that is taken from
-
- Blank, L. (2014).
- On Divergence-Free Finite Element Methods for the Stokes Equations (Freie Universität Berlin).
- p. 84, eq. (6.2)
-
-for the P1-P0 stabilized Stokes discretization.
-
-Component trial: {component_trial}
-Component test: {component_test}
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (vectorial space: {trial})
- v: test function (vectorial space: {test})
-
- ∫ ( ∇ · u ) · ( ∇ · v )
-"""
-
- from hog.recipes.integrands.volume.divdiv import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- )
-
-
-def advection(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- velocity_function_space: FunctionSpace,
- coefficient_function_space: FunctionSpace,
- constant_cp: bool = False,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-advection operator which needs to be used in combination with SUPG
-
-Geometry map: {blending}
-
-Weak formulation
-
- T: trial function (scalar space: {trial})
- s: test function (scalar space: {test})
- u: velocity function (vectorial space: {velocity_function_space})
-
- ∫ cp ( u · ∇T ) s
-"""
-
- from hog.recipes.integrands.volume.advection import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=False,
- docstring=docstring,
- fe_coefficients={
- "ux": velocity_function_space,
- "uy": velocity_function_space,
- "uz": velocity_function_space,
- } if constant_cp else {
- "ux": velocity_function_space,
- "uy": velocity_function_space,
- "uz": velocity_function_space,
- "cp": coefficient_function_space,
- },
- )
-
-
-def supg_advection(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- velocity_function_space: FunctionSpace,
- coefficient_function_space: FunctionSpace,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-advection operator which needs to be used in combination with SUPG
-
-Geometry map: {blending}
-
-Weak formulation
-
- T: trial function (scalar space: {trial})
- s: test function (scalar space: {test})
- u: velocity function (vectorial space: {velocity_function_space})
-
- ∫ cp ( u · ∇T ) 𝛿(u · ∇s)
-"""
-
- from hog.recipes.integrands.volume.supg_advection import integrand
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=False,
- docstring=docstring,
- fe_coefficients={
- "ux": velocity_function_space,
- "uy": velocity_function_space,
- "uz": velocity_function_space,
- "cp_times_delta": coefficient_function_space,
- },
- )
-
-
-def supg_diffusion(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- velocity_function_space: FunctionSpace,
- diffusivityXdelta_function_space: FunctionSpace,
- blending: GeometryMap = IdentityMap(),
-) -> Form:
- docstring = f"""
-Form for SUPGDiffusion operator used for SUPG stabilisation
-
-Geometry map: {blending}
-
-Weak formulation
-
- T: trial function (space: {trial})
- s: test function (space: {test})
- w: velocity function (space: {velocity_function_space})
- k𝛿: FE function representing k·𝛿 (space: {diffusivityXdelta_function_space})
-
- For OpGen,
-
- ∫ k(ΔT) · 𝛿(w · ∇s)
-
- -------------------
-
- For ExternalMap (only for testing, currently not supported),
-
- ∫ (ΔT) s
-"""
-
- def integrand(
- *,
- jac_a_inv,
- jac_b_inv,
- hessian_b,
- jac_a_abs_det,
- jac_b_abs_det,
- grad_u,
- grad_v,
- hessian_u,
- k,
- volume_geometry,
- tabulate,
- **_,
- ):
- """First function: k𝛿, other functions: ux, uy, uz."""
-
- k_times_delta = k["diffusivity_times_delta"]
- wx = k["wx"]
- wy = k["wy"]
- wz = k["wz"]
-
- dim = volume_geometry.dimensions
- if dim == 2:
- w = sp.Matrix([[wx], [wy]])
- elif dim == 3:
- w = sp.Matrix([[wx], [wy], [wz]])
-
- if isinstance(blending, IdentityMap):
- hessian_affine = hessian_shape_affine_ref_pullback(hessian_u, jac_a_inv)
-
- laplacian = sum([hessian_affine[i, i] for i in range(geometry.dimensions)])
-
- form = (
- k_times_delta
- * tabulate(laplacian)
- * dot(w, tabulate(jac_a_inv.T * grad_v))
- * jac_a_abs_det
- )
-
- else:
- hessian_blending = hessian_shape_blending_ref_pullback(
- geometry,
- grad_u,
- hessian_u,
- jac_a_inv,
- hessian_b,
- jac_b_inv,
- )
-
- laplacian = sum(
- [hessian_blending[i, i] for i in range(geometry.dimensions)]
- )
-
- form = (
- laplacian
- * dot(w, jac_b_inv.T * tabulate(jac_a_inv.T * grad_v))
- * k_times_delta
- * jac_a_abs_det
- * jac_b_abs_det
- )
-
- return form
-
- return process_integrand(
- integrand,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- fe_coefficients={
- "diffusivity_times_delta": diffusivityXdelta_function_space,
- "wx": velocity_function_space,
- "wy": velocity_function_space,
- "wz": velocity_function_space,
- },
- )
-
-
-def grad_rho_by_rho_dot_u(
- trial: TrialSpace,
- test: TestSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- density_function_space: Optional[FunctionSpace] = None,
-) -> Form:
- docstring = f"""
-RHS operator for the frozen velocity approach.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (vectorial space: {trial})
- v: test function (space: {test})
- rho: coefficient (space: {density_function_space})
-
- ∫ ((∇ρ / ρ) · u) v
-"""
-
- with TimedLogger(
- "assembling grad_rho_by_rho_dot_u-mass matrix", level=logging.DEBUG
- ):
- tabulation = Tabulation(symbolizer)
-
- jac_affine_det = symbolizer.abs_det_jac_ref_to_affine()
- jac_affine_inv = symbolizer.jac_ref_to_affine_inv(geometry)
-
- if isinstance(blending, ExternalMap):
- jac_blending = jac_affine_to_physical(geometry, symbolizer)
- else:
- affine_coords = trafo_ref_to_affine(geometry, symbolizer)
- # jac_blending = blending.jacobian(affine_coords)
- jac_blending_inv = symbolizer.jac_affine_to_blending_inv(
- geometry.dimensions
- )
- jac_blending_det = symbolizer.abs_det_jac_affine_to_blending()
-
- # jac_blending_det = abs(det(jac_blending))
-
- mat = create_empty_element_matrix(trial, test, geometry)
-
- it = element_matrix_iterator(trial, test, geometry)
-
- if density_function_space:
- rho, dof_symbols = fem_function_on_element(
- density_function_space,
- geometry,
- symbolizer,
- domain="reference",
- function_id="rho",
- )
-
- grad_rho, _ = fem_function_gradient_on_element(
- density_function_space,
- geometry,
- symbolizer,
- domain="reference",
- function_id="grad_rho",
- dof_symbols=dof_symbols,
- )
-
- with TimedLogger(
- f"integrating {mat.shape[0] * mat.shape[1]} expressions",
- level=logging.DEBUG,
- ):
- for data in it:
- phi = data.trial_shape
- psi = data.test_shape
- # print("phi = ", psi.shape)
- # phi_psi_jac_affine_det = tabulation.register_factor(
- # "phi_psi_jac_affine_det",
- # sp.Matrix([phi * psi * jac_affine_det]),
- # )[0]
- if blending == IdentityMap():
- form = (
- dot(((jac_affine_inv.T * grad_rho) / rho[0]), phi)
- * psi
- * jac_affine_det
- )
- else:
- form = (
- dot(
- (
- (jac_blending_inv.T * jac_affine_inv.T * grad_rho)
- / rho[0]
- ),
- phi,
- )
- * psi
- * jac_affine_det
- * jac_blending_det
- )
- mat[data.row, data.col] = form
-
- return Form(mat, tabulation, symmetric=trial == test, docstring=docstring)
-
-
-def zero_form(
- trial: FunctionSpace, test: FunctionSpace, geometry: ElementGeometry
-) -> sp.Matrix:
- rows = test.num_dofs(geometry)
- cols = trial.num_dofs(geometry) if trial is not None else 1
- return sp.zeros(rows, cols)
-
-
-def linear_elasticity(
- trial: FunctionSpace,
- test: FunctionSpace,
- geometry: ElementGeometry,
- symbolizer: Symbolizer,
- blending: GeometryMap = IdentityMap(),
- component_trial: int = 0,
- component_test: int = 0,
- variable_viscosity: bool = True,
- coefficient_function_space: Optional[FunctionSpace] = None,
-) -> Form:
- docstring = f"""
-"Linear elasticity" operator.
-
-Geometry map: {blending}
-
-Weak formulation
-
- u: trial function (vectorial space: {trial})
- v: test function (vectorial space: {test})
- λ: coefficient (scalar space: {coefficient_function_space})
- μ: coefficient (scalar space: {coefficient_function_space})
-
- ∫ σ(u) : ε(v)
-
-where
-
- σ(w) := λ(∇.w)I + μ(∇w + (∇w)ᵀ)
- ε(w) := (1/2) (∇w + (∇w)ᵀ)
-"""
- if not variable_viscosity:
- raise HOGException("Constant viscosity currently not supported.")
-
- from hog.recipes.integrands.volume.elasticity import integrand
- from hog.recipes.integrands.volume.elasticity_affine import (
- integrand as integrand_affine,
- )
-
- # mypy type checking supposedly cannot figure out ternaries.
- # https://stackoverflow.com/a/70832391
- integr: Callable[..., Any] = integrand
- if blending == IdentityMap():
- integr = integrand_affine
-
- return process_integrand(
- integr,
- trial,
- test,
- geometry,
- symbolizer,
- blending=blending,
- is_symmetric=trial == test,
- docstring=docstring,
- fe_coefficients={"mu": coefficient_function_space, "lambda": coefficient_function_space,},
- )
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+import logging
+import sympy as sp
+from typing import Optional, Callable, Any
+
+from hog.element_geometry import ElementGeometry
+from hog.exception import HOGException
+from hog.fem_helpers import (
+ trafo_ref_to_affine,
+ jac_ref_to_affine,
+ jac_affine_to_physical,
+ hessian_shape_affine_ref_pullback,
+ hessian_shape_blending_ref_pullback,
+ create_empty_element_matrix,
+ element_matrix_iterator,
+ scalar_space_dependent_coefficient,
+ vector_space_dependent_coefficient,
+ fem_function_on_element,
+ fem_function_gradient_on_element,
+)
+from hog.function_space import FunctionSpace, N1E1Space, TrialSpace, TestSpace
+from hog.math_helpers import dot, inv, abs, det, double_contraction
+from hog.quadrature import Quadrature, Tabulation
+from hog.symbolizer import Symbolizer
+from hog.logger import TimedLogger
+from hog.blending import GeometryMap, ExternalMap, IdentityMap
+from hog.integrand import process_integrand, Form
+
+
+def diffusion(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+Diffusion operator without coefficients.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (space: {trial})
+ v: test function (space: {test})
+
+ ∫ ∇u : ∇v
+
+ Note that the double contraction (:) reduces to the dot product for scalar function spaces, i.e. the form becomes
+
+ ∫ ∇u · ∇v
+"""
+
+ from hog.recipes.integrands.volume.diffusion import integrand
+ from hog.recipes.integrands.volume.diffusion_affine import (
+ integrand as integrand_affine,
+ )
+
+ # mypy type checking supposedly cannot figure out ternaries.
+ # https://stackoverflow.com/a/70832391
+ integr: Callable[..., Any] = integrand
+ if blending == IdentityMap():
+ integr = integrand_affine
+
+ return process_integrand(
+ integr,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ )
+
+
+def mass(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+Mass operator.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (space: {trial})
+ v: test function (space: {test})
+
+ ∫ uv
+"""
+
+ from hog.recipes.integrands.volume.mass import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ )
+
+
+def div_k_grad(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ coefficient_function_space: Optional[FunctionSpace] = None,
+) -> Form:
+ docstring = f"""
+Diffusion operator with a scalar coefficient.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (space: {trial})
+ v: test function (space: {test})
+ k: coefficient (space: {coefficient_function_space})
+
+ ∫ k ∇u · ∇v
+"""
+
+ from hog.recipes.integrands.volume.div_k_grad import integrand
+ from hog.recipes.integrands.volume.div_k_grad_affine import (
+ integrand as integrand_affine,
+ )
+
+ # mypy type checking supposedly cannot figure out ternaries.
+ # https://stackoverflow.com/a/70832391
+ integr: Callable[..., Any] = integrand
+ if blending == IdentityMap():
+ integr = integrand_affine
+
+ return process_integrand(
+ integr,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ fe_coefficients={"k": coefficient_function_space},
+ is_symmetric=trial == test,
+ docstring=docstring,
+ )
+
+
+def nonlinear_diffusion(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ coefficient_function_space: FunctionSpace,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+Diffusion operator with coefficient function.
+
+Weak formulation
+
+ u: trial function (space: {trial})
+ v: test function (space: {test})
+ c: FE coefficient function (space: {coefficient_function_space})
+
+ ∫ a(c) ∇u · ∇v
+
+Note: :math:`a(c) = 1/8 + u^2` is currently hard-coded and the form is intended for :math:`c = u`,
+ hence, the naming, but will work for other coefficients, too.
+"""
+
+ if blending != IdentityMap():
+ raise HOGException(
+ "The nonlinear-diffusion form does currently not support blending."
+ )
+
+ if trial != test:
+ raise HOGException(
+ "Trial space must be equal to test space to assemble non-linear diffusion matrix."
+ )
+
+ def integrand(
+ *,
+ jac_a_inv,
+ jac_a_abs_det,
+ grad_u,
+ grad_v,
+ k,
+ tabulate,
+ **_,
+ ):
+ a = sp.Rational(1, 8) + k["u"] * k["u"]
+ return a * tabulate(
+ double_contraction(
+ jac_a_inv.T * grad_u,
+ jac_a_inv.T * grad_v,
+ )
+ * jac_a_abs_det
+ )
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ fe_coefficients={"u": coefficient_function_space},
+ )
+
+
+def nonlinear_diffusion_newton_galerkin(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ coefficient_function_space: FunctionSpace,
+ blending: GeometryMap = IdentityMap(),
+ only_newton_galerkin_part_of_form: Optional[bool] = True,
+) -> Form:
+ docstring = f"""
+
+Bi-linear form for the solution of the non-linear diffusion equation by a Newton-Galerkin approach
+
+Weak formulation
+
+ u: trial function (space: {trial})
+ v: test function (space: {test})
+ k: FE coefficient function (space: {coefficient_function_space})
+
+ ∫ a(k) ∇u · ∇v + ∫ a'(k) u ∇k · ∇v
+
+Note: :math:`a(k) = 1/8 + k^2` is currently hard-coded and the form is intended for :math:`k = u`.
+"""
+ if trial != test:
+ raise HOGException(
+ "Trial space must be equal to test space to assemble diffusion matrix."
+ )
+
+ if blending != IdentityMap():
+ raise HOGException(
+ "The nonlinear_diffusion_newton_galerkin form does currently not support blending."
+ )
+
+ def integrand(
+ *,
+ jac_a_inv,
+ jac_a_abs_det,
+ u,
+ grad_u,
+ grad_v,
+ k,
+ grad_k,
+ tabulate,
+ **_,
+ ):
+ a = sp.Rational(1, 8) + k["k"] * k["k"]
+ a_prime = 2 * k["k"]
+
+ diffusion_term = a * tabulate(
+ dot(jac_a_inv.T * grad_u, jac_a_inv.T * grad_v) * jac_a_abs_det
+ )
+
+ newton_galerkin_term = (
+ a_prime
+ * u
+ * dot(jac_a_inv.T * grad_k["k"], tabulate(jac_a_inv.T * grad_v))
+ * tabulate(jac_a_abs_det)
+ )
+
+ if only_newton_galerkin_part_of_form:
+ return newton_galerkin_term
+ else:
+ return diffusion_term + newton_galerkin_term
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=False,
+ docstring=docstring,
+ fe_coefficients={"k": coefficient_function_space},
+ )
+
+
+def epsilon(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ component_trial: int = 0,
+ component_test: int = 0,
+ variable_viscosity: bool = True,
+ coefficient_function_space: Optional[FunctionSpace] = None,
+) -> Form:
+ docstring = f"""
+"Epsilon" operator.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (vectorial space: {trial})
+ v: test function (vectorial space: {test})
+ μ: coefficient (scalar space: {coefficient_function_space})
+
+ ∫ 2 μ ε(u) : ε(v)
+
+where
+
+ ε(w) := (1/2) (∇w + (∇w)ᵀ)
+"""
+ if not variable_viscosity:
+ raise HOGException("Constant viscosity currently not supported.")
+
+ from hog.recipes.integrands.volume.epsilon import integrand
+ from hog.recipes.integrands.volume.epsilon_affine import (
+ integrand as integrand_affine,
+ )
+
+ # mypy type checking supposedly cannot figure out ternaries.
+ # https://stackoverflow.com/a/70832391
+ integr: Callable[..., Any] = integrand
+ if blending == IdentityMap():
+ integr = integrand_affine
+
+ return process_integrand(
+ integr,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ fe_coefficients={"mu": coefficient_function_space},
+ )
+
+
+def k_mass(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ coefficient_function_space: Optional[FunctionSpace] = None,
+) -> Form:
+ docstring = f"""
+Diffusion operator with a scalar coefficient.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (space: {trial})
+ v: test function (space: {test})
+ k: coefficient (space: {coefficient_function_space})
+
+ ∫ k uv
+"""
+
+ from hog.recipes.integrands.volume.k_mass import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ )
+
+
+def pspg(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ quad: Quadrature,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+PSPG stabilisation.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (space: {trial})
+ v: test function (space: {test})
+
+ ∫ τ ∇u · ∇v
+
+ where tau is an element-wise factor given by
+
+ 2D:
+ τ = -CellVolume(triangle) / 5
+ 3D:
+ τ = -(CellVolume(tetrahedron))**(2/3) / 12
+
+See e.g.
+
+ Brezzi, F., & Douglas, J. (1988).
+ Stabilized mixed methods for the Stokes problem.
+ Numerische Mathematik, 53, 225-235.
+
+or
+
+ Hughes, T. J., Franca, L. P., & Balestra, M. (1986).
+ A new finite element formulation for computational fluid dynamics: V.
+ Circumventing the Babuška-Brezzi condition: A stable Petrov-Galerkin formulation of the Stokes problem accommodating
+ equal-order interpolations.
+ Computer Methods in Applied Mechanics and Engineering, 59(1), 85-99.
+
+for details.
+"""
+
+ from hog.recipes.integrands.volume.pspg import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ )
+
+
+def linear_form(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ quad: Quadrature,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+) -> sp.Matrix:
+ """
+ Implements a linear form of type:
+
+ (k(x), psi)
+
+ where psi a test function and k = k(x) a scalar, external function.
+ """
+
+ if trial != test:
+ raise HOGException(
+ "Trial space must be equal to test space to assemble linear form (jep this is weird, but linear forms are implemented as diagonal matrices)."
+ )
+
+ if quad.is_exact() and isinstance(blending, ExternalMap):
+ raise HOGException(
+ "Exact integration is not supported for externally defined blending functions."
+ )
+
+ with TimedLogger("assembling linear form", level=logging.DEBUG):
+ jac_affine = jac_ref_to_affine(geometry, symbolizer)
+ jac_affine_det = abs(det(jac_affine))
+
+ if isinstance(blending, ExternalMap):
+ jac_blending = jac_affine_to_physical(geometry, symbolizer)
+ else:
+ affine_coords = trafo_ref_to_affine(geometry, symbolizer)
+ jac_blending = blending.jacobian(affine_coords)
+
+ jac_blending_det = abs(det(jac_blending))
+
+ mat = create_empty_element_matrix(trial, test, geometry)
+
+ it = element_matrix_iterator(trial, test, geometry)
+
+ if isinstance(trial, N1E1Space):
+ jac_affine_inv_T = inv(jac_affine).T
+ jac_blending_inv_T = inv(jac_blending).T
+ space_dependent_coefficient = vector_space_dependent_coefficient
+ else:
+ space_dependent_coefficient = scalar_space_dependent_coefficient
+ k = space_dependent_coefficient("k", geometry, symbolizer, blending=blending)
+
+ with TimedLogger(
+ f"integrating {mat.shape[0]} expressions",
+ level=logging.DEBUG,
+ ):
+ for data in it:
+ if data.row == data.col:
+ psi = data.test_shape
+
+ if isinstance(trial, N1E1Space):
+ form = (
+ dot(k, jac_blending_inv_T * jac_affine_inv_T * psi)
+ * jac_affine_det
+ * jac_blending_det
+ )
+ else:
+ form = k * psi * jac_affine_det * jac_blending_det
+
+ mat[data.row, data.col] = quad.integrate(form, symbolizer)
+
+ return mat
+
+
+def divergence(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ component_index: int = 0,
+) -> Form:
+ docstring = f"""
+Divergence.
+
+Component: {component_index}
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (vectorial space: {trial})
+ v: test function (scalar space: {test})
+
+ ∫ - ( ∇ · u ) v
+"""
+
+ from hog.recipes.integrands.volume.divergence import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=False,
+ docstring=docstring,
+ )
+
+
+def gradient(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ component_index: int = 0,
+) -> Form:
+ docstring = f"""
+ Gradient.
+
+ Component: {component_index}
+ Geometry map: {blending}
+
+ Weak formulation
+
+ u: trial function (scalar space: {trial})
+ v: test function (vectorial space: {test})
+
+ ∫ - ( ∇ · v ) u
+ """
+
+ from hog.recipes.integrands.volume.gradient import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=False,
+ docstring=docstring,
+ )
+
+
+def full_stokes(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ component_trial: int = 0,
+ component_test: int = 0,
+ variable_viscosity: bool = True,
+ coefficient_function_space: Optional[FunctionSpace] = None,
+) -> Form:
+ docstring = f"""
+Implements the fully coupled viscous operator of the Stokes problem.
+The latter is the extension of the Epsilon operator to the case where
+the velocity field need not be divergence-free. This is e.g. the case
+in the (truncated) anelastic liquid approximation of mantle convection.
+
+The strong representation of the operator is given by:
+
+ - div[ μ (grad(u)+grad(u)ᵀ) ] + 2/3 grad[ μ div(u) ]
+
+Note that the factor 2/3 means that for 2D this is the pseudo-3D form
+of the operator.
+
+Component trial: {component_trial}
+Component test: {component_test}
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (vectorial space: {trial})
+ v: test function (vectorial space: {test})
+ μ: coefficient (scalar space: {coefficient_function_space})
+
+ ∫ μ {{ ( 2 ε(u) : ε(v) ) - (2/3) [ ( ∇ · u ) · ( ∇ · v ) ] }}
+
+where
+
+ ε(w) := (1/2) (∇w + (∇w)ᵀ)
+"""
+
+ from hog.recipes.integrands.volume.full_stokes import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ fe_coefficients={"mu": coefficient_function_space},
+ )
+
+
+def shear_heating(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ component_trial: int = 0,
+ component_test: int = 0,
+ variable_viscosity: bool = True,
+ viscosity_function_space: Optional[FunctionSpace] = None,
+ velocity_function_space: Optional[FunctionSpace] = None,
+) -> Form:
+ docstring = f"""
+Implements the fully coupled viscous operator for the shear heating term.
+The latter is the extension of the Epsilon operator to the case where
+the velocity field need not be divergence-free. This is e.g. the case
+in the (truncated) anelastic liquid approximation of mantle convection.
+
+https://doi.org/10.1111/j.1365-246X.2009.04413.x
+(3) and (5)
+
+https://doi.org/10.5194/gmd-15-5127-2022
+Listing 2
+
+The strong representation of the operator is given by:
+
+ 𝜏(w) : grad(w)
+ 2 {{[ μ (grad(w)+grad(w)ᵀ) / 2 ] - 1/dim [ μ div(w) ]I}} : grad(w)
+
+Note that the factor 1/dim means that for 2D this is the pseudo-3D form
+of the operator.
+
+Component trial: {component_trial}
+Component test: {component_test}
+Geometry map: {blending}
+
+Weak formulation
+
+ T: trial function (scalar space: {trial})
+ s: test function (scalar space: {test})
+ μ: coefficient (scalar space: {viscosity_function_space})
+ w: velocity (vectorial space: {velocity_function_space})
+
+ ∫ {{ 2 {{[ μ (grad(w)+grad(w)ᵀ) / 2 ] - 1/dim [ μ div(w) ]I}} : grad(w) }} T_h s_h
+
+The resulting matrix must be multiplied with a vector of ones to be used as the shear heating term in the RHS
+"""
+
+ def integrand(
+ *,
+ jac_a_inv,
+ jac_b_inv,
+ jac_a_abs_det,
+ jac_b_abs_det,
+ u,
+ v,
+ k,
+ grad_k,
+ volume_geometry,
+ tabulate,
+ **_,
+ ):
+ """First function: mu, other functions: ux, uy, uz."""
+
+ mu = k["mu"]
+
+ # grad_k[0] is grad_mu_ref
+ grad_wx = jac_b_inv.T * jac_a_inv.T * grad_k["wx"]
+ grad_wy = jac_b_inv.T * jac_a_inv.T * grad_k["wy"]
+ grad_wz = jac_b_inv.T * jac_a_inv.T * grad_k["wz"]
+
+ grad_w = grad_wx.row_join(grad_wy)
+ dim = volume_geometry.dimensions
+ if dim == 3:
+ grad_w = grad_w.row_join(grad_wz)
+
+ sym_grad_w = 0.5 * (grad_w + grad_w.T)
+
+ divdiv = grad_w.trace() * sp.eye(dim)
+
+ tau = 2 * (sym_grad_w - sp.Rational(1, dim) * divdiv)
+
+ return (
+ mu
+ * (double_contraction(tau, grad_w)[0])
+ * jac_b_abs_det
+ * tabulate(jac_a_abs_det * u * v)
+ )
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ fe_coefficients={
+ "mu": viscosity_function_space,
+ "wx": velocity_function_space,
+ "wy": velocity_function_space,
+ "wz": velocity_function_space,
+ },
+ is_symmetric=trial == test,
+ docstring=docstring,
+ )
+
+
+def divdiv(
+ trial: TrialSpace,
+ test: TestSpace,
+ component_trial: int,
+ component_test: int,
+ geometry: ElementGeometry,
+ quad: Quadrature,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+divdiv operator which is used as a stabilization term that is taken from
+
+ Blank, L. (2014).
+ On Divergence-Free Finite Element Methods for the Stokes Equations (Freie Universität Berlin).
+ p. 84, eq. (6.2)
+
+for the P1-P0 stabilized Stokes discretization.
+
+Component trial: {component_trial}
+Component test: {component_test}
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (vectorial space: {trial})
+ v: test function (vectorial space: {test})
+
+ ∫ ( ∇ · u ) · ( ∇ · v )
+"""
+
+ from hog.recipes.integrands.volume.divdiv import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ )
+
+
+def advection(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ velocity_function_space: FunctionSpace,
+ coefficient_function_space: FunctionSpace,
+ constant_cp: bool = False,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+advection operator which needs to be used in combination with SUPG
+
+Geometry map: {blending}
+
+Weak formulation
+
+ T: trial function (scalar space: {trial})
+ s: test function (scalar space: {test})
+ u: velocity function (vectorial space: {velocity_function_space})
+
+ ∫ cp ( u · ∇T ) s
+"""
+
+ from hog.recipes.integrands.volume.advection import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=False,
+ docstring=docstring,
+ fe_coefficients={
+ "ux": velocity_function_space,
+ "uy": velocity_function_space,
+ "uz": velocity_function_space,
+ } if constant_cp else {
+ "ux": velocity_function_space,
+ "uy": velocity_function_space,
+ "uz": velocity_function_space,
+ "cp": coefficient_function_space,
+ },
+ )
+
+
+def supg_advection(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ velocity_function_space: FunctionSpace,
+ coefficient_function_space: FunctionSpace,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+advection operator which needs to be used in combination with SUPG
+
+Geometry map: {blending}
+
+Weak formulation
+
+ T: trial function (scalar space: {trial})
+ s: test function (scalar space: {test})
+ u: velocity function (vectorial space: {velocity_function_space})
+
+ ∫ cp ( u · ∇T ) 𝛿(u · ∇s)
+"""
+
+ from hog.recipes.integrands.volume.supg_advection import integrand
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=False,
+ docstring=docstring,
+ fe_coefficients={
+ "ux": velocity_function_space,
+ "uy": velocity_function_space,
+ "uz": velocity_function_space,
+ "cp_times_delta": coefficient_function_space,
+ },
+ )
+
+
+def supg_diffusion(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ velocity_function_space: FunctionSpace,
+ diffusivityXdelta_function_space: FunctionSpace,
+ blending: GeometryMap = IdentityMap(),
+) -> Form:
+ docstring = f"""
+Form for SUPGDiffusion operator used for SUPG stabilisation
+
+Geometry map: {blending}
+
+Weak formulation
+
+ T: trial function (space: {trial})
+ s: test function (space: {test})
+ w: velocity function (space: {velocity_function_space})
+ k𝛿: FE function representing k·𝛿 (space: {diffusivityXdelta_function_space})
+
+ For OpGen,
+
+ ∫ k(ΔT) · 𝛿(w · ∇s)
+
+ -------------------
+
+ For ExternalMap (only for testing, currently not supported),
+
+ ∫ (ΔT) s
+"""
+
+ def integrand(
+ *,
+ jac_a_inv,
+ jac_b_inv,
+ hessian_b,
+ jac_a_abs_det,
+ jac_b_abs_det,
+ grad_u,
+ grad_v,
+ hessian_u,
+ k,
+ volume_geometry,
+ tabulate,
+ **_,
+ ):
+ """First function: k𝛿, other functions: ux, uy, uz."""
+
+ k_times_delta = k["diffusivity_times_delta"]
+ wx = k["wx"]
+ wy = k["wy"]
+ wz = k["wz"]
+
+ dim = volume_geometry.dimensions
+ if dim == 2:
+ w = sp.Matrix([[wx], [wy]])
+ elif dim == 3:
+ w = sp.Matrix([[wx], [wy], [wz]])
+
+ if isinstance(blending, IdentityMap):
+ hessian_affine = hessian_shape_affine_ref_pullback(hessian_u, jac_a_inv)
+
+ laplacian = sum([hessian_affine[i, i] for i in range(geometry.dimensions)])
+
+ form = (
+ k_times_delta
+ * tabulate(laplacian)
+ * dot(w, tabulate(jac_a_inv.T * grad_v))
+ * jac_a_abs_det
+ )
+
+ else:
+ hessian_blending = hessian_shape_blending_ref_pullback(
+ geometry,
+ grad_u,
+ hessian_u,
+ jac_a_inv,
+ hessian_b,
+ jac_b_inv,
+ )
+
+ laplacian = sum(
+ [hessian_blending[i, i] for i in range(geometry.dimensions)]
+ )
+
+ form = (
+ laplacian
+ * dot(w, jac_b_inv.T * tabulate(jac_a_inv.T * grad_v))
+ * k_times_delta
+ * jac_a_abs_det
+ * jac_b_abs_det
+ )
+
+ return form
+
+ return process_integrand(
+ integrand,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ fe_coefficients={
+ "diffusivity_times_delta": diffusivityXdelta_function_space,
+ "wx": velocity_function_space,
+ "wy": velocity_function_space,
+ "wz": velocity_function_space,
+ },
+ )
+
+
+def grad_rho_by_rho_dot_u(
+ trial: TrialSpace,
+ test: TestSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ density_function_space: Optional[FunctionSpace] = None,
+) -> Form:
+ docstring = f"""
+RHS operator for the frozen velocity approach.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (vectorial space: {trial})
+ v: test function (space: {test})
+ rho: coefficient (space: {density_function_space})
+
+ ∫ ((∇ρ / ρ) · u) v
+"""
+
+ with TimedLogger(
+ "assembling grad_rho_by_rho_dot_u-mass matrix", level=logging.DEBUG
+ ):
+ tabulation = Tabulation(symbolizer)
+
+ jac_affine_det = symbolizer.abs_det_jac_ref_to_affine()
+ jac_affine_inv = symbolizer.jac_ref_to_affine_inv(geometry)
+
+ if isinstance(blending, ExternalMap):
+ jac_blending = jac_affine_to_physical(geometry, symbolizer)
+ else:
+ affine_coords = trafo_ref_to_affine(geometry, symbolizer)
+ # jac_blending = blending.jacobian(affine_coords)
+ jac_blending_inv = symbolizer.jac_affine_to_blending_inv(
+ geometry.dimensions
+ )
+ jac_blending_det = symbolizer.abs_det_jac_affine_to_blending()
+
+ # jac_blending_det = abs(det(jac_blending))
+
+ mat = create_empty_element_matrix(trial, test, geometry)
+
+ it = element_matrix_iterator(trial, test, geometry)
+
+ if density_function_space:
+ rho, dof_symbols = fem_function_on_element(
+ density_function_space,
+ geometry,
+ symbolizer,
+ domain="reference",
+ function_id="rho",
+ )
+
+ grad_rho, _ = fem_function_gradient_on_element(
+ density_function_space,
+ geometry,
+ symbolizer,
+ domain="reference",
+ function_id="grad_rho",
+ dof_symbols=dof_symbols,
+ )
+
+ with TimedLogger(
+ f"integrating {mat.shape[0] * mat.shape[1]} expressions",
+ level=logging.DEBUG,
+ ):
+ for data in it:
+ phi = data.trial_shape
+ psi = data.test_shape
+ # print("phi = ", psi.shape)
+ # phi_psi_jac_affine_det = tabulation.register_factor(
+ # "phi_psi_jac_affine_det",
+ # sp.Matrix([phi * psi * jac_affine_det]),
+ # )[0]
+ if blending == IdentityMap():
+ form = (
+ dot(((jac_affine_inv.T * grad_rho) / rho[0]), phi)
+ * psi
+ * jac_affine_det
+ )
+ else:
+ form = (
+ dot(
+ (
+ (jac_blending_inv.T * jac_affine_inv.T * grad_rho)
+ / rho[0]
+ ),
+ phi,
+ )
+ * psi
+ * jac_affine_det
+ * jac_blending_det
+ )
+ mat[data.row, data.col] = form
+
+ return Form(mat, tabulation, symmetric=trial == test, docstring=docstring)
+
+
+def zero_form(
+ trial: FunctionSpace, test: FunctionSpace, geometry: ElementGeometry
+) -> sp.Matrix:
+ rows = test.num_dofs(geometry)
+ cols = trial.num_dofs(geometry) if trial is not None else 1
+ return sp.zeros(rows, cols)
+
+
+def linear_elasticity(
+ trial: FunctionSpace,
+ test: FunctionSpace,
+ geometry: ElementGeometry,
+ symbolizer: Symbolizer,
+ blending: GeometryMap = IdentityMap(),
+ component_trial: int = 0,
+ component_test: int = 0,
+ variable_viscosity: bool = True,
+ coefficient_function_space: Optional[FunctionSpace] = None,
+) -> Form:
+ docstring = f"""
+"Linear elasticity" operator.
+
+Geometry map: {blending}
+
+Weak formulation
+
+ u: trial function (vectorial space: {trial})
+ v: test function (vectorial space: {test})
+ λ: coefficient (scalar space: {coefficient_function_space})
+ μ: coefficient (scalar space: {coefficient_function_space})
+
+ ∫ σ(u) : ε(v)
+
+where
+
+ σ(w) := λ(∇.w)I + μ(∇w + (∇w)ᵀ)
+ ε(w) := (1/2) (∇w + (∇w)ᵀ)
+"""
+ if not variable_viscosity:
+ raise HOGException("Constant viscosity currently not supported.")
+
+ from hog.recipes.integrands.volume.elasticity import integrand
+ from hog.recipes.integrands.volume.elasticity_affine import (
+ integrand as integrand_affine,
+ )
+
+ # mypy type checking supposedly cannot figure out ternaries.
+ # https://stackoverflow.com/a/70832391
+ integr: Callable[..., Any] = integrand
+ if blending == IdentityMap():
+ integr = integrand_affine
+
+ return process_integrand(
+ integr,
+ trial,
+ test,
+ geometry,
+ symbolizer,
+ blending=blending,
+ is_symmetric=trial == test,
+ docstring=docstring,
+ fe_coefficients={"mu": coefficient_function_space, "lambda": coefficient_function_space,},
+ )
diff --git a/hog/recipes/integrands/volume/elasticity.py b/hog/recipes/integrands/volume/elasticity.py
index 35ea06b5c93174e5b023c62e8334c3858dce6af7..5125580da62fbc5ef1a6a95a1b2a9dbcfff20d8f 100644
--- a/hog/recipes/integrands/volume/elasticity.py
+++ b/hog/recipes/integrands/volume/elasticity.py
@@ -1,47 +1,47 @@
-# HyTeG Operator Generator
-# Copyright (C) 2024 HyTeG Team
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program. If not, see <https://www.gnu.org/licenses/>.
-
-from hog.recipes.common import *
-import sympy as sp
-
-def integrand(
- *,
- jac_a_inv,
- jac_a_abs_det,
- jac_b_inv,
- jac_b_abs_det,
- grad_u,
- grad_v,
- k,
- tabulate,
- **_,
-):
-
- grad_u_chain = jac_b_inv.T * tabulate(jac_a_inv.T * grad_u)
- grad_v_chain = jac_b_inv.T * tabulate(jac_a_inv.T * grad_v)
-
- def symm_grad(w):
- return 0.5 * (w + w.T)
-
- symm_grad_u = symm_grad(grad_u_chain)
- sigma_u = k["lambda"] * grad_u_chain.trace() * sp.eye(grad_u_chain.shape[0]) + k["mu"] * 2 * symm_grad_u
- symm_grad_v = symm_grad(grad_v_chain)
-
- return (
- double_contraction(sigma_u, symm_grad_v)
- * tabulate(jac_a_abs_det)
- * jac_b_abs_det
- )
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+from hog.recipes.common import *
+import sympy as sp
+
+def integrand(
+ *,
+ jac_a_inv,
+ jac_a_abs_det,
+ jac_b_inv,
+ jac_b_abs_det,
+ grad_u,
+ grad_v,
+ k,
+ tabulate,
+ **_,
+):
+
+ grad_u_chain = jac_b_inv.T * tabulate(jac_a_inv.T * grad_u)
+ grad_v_chain = jac_b_inv.T * tabulate(jac_a_inv.T * grad_v)
+
+ def symm_grad(w):
+ return 0.5 * (w + w.T)
+
+ symm_grad_u = symm_grad(grad_u_chain)
+ sigma_u = k["lambda"] * grad_u_chain.trace() * sp.eye(grad_u_chain.shape[0]) + k["mu"] * 2 * symm_grad_u
+ symm_grad_v = symm_grad(grad_v_chain)
+
+ return (
+ double_contraction(sigma_u, symm_grad_v)
+ * tabulate(jac_a_abs_det)
+ * jac_b_abs_det
+ )
diff --git a/hog/recipes/integrands/volume/elasticity_affine.py b/hog/recipes/integrands/volume/elasticity_affine.py
index b51bd8a78603f90cc37ae5d991ed982c58355132..8d9c4edf9e1082b9b9dc448c2281b4f0009e14d1 100644
--- a/hog/recipes/integrands/volume/elasticity_affine.py
+++ b/hog/recipes/integrands/volume/elasticity_affine.py
@@ -1,44 +1,44 @@
-# HyTeG Operator Generator
-# Copyright (C) 2024 HyTeG Team
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program. If not, see <https://www.gnu.org/licenses/>.
-
-from hog.recipes.common import *
-import sympy as sp
-
-
-def integrand(
- *,
- jac_a_inv,
- jac_a_abs_det,
- grad_u,
- grad_v,
- k,
- tabulate,
- **_,
-):
-
- grad_u_chain = jac_a_inv.T * grad_u
- grad_v_chain = jac_a_inv.T * grad_v
-
- def symm_grad(w):
- return 0.5 * (w + w.T)
-
- symm_grad_u = symm_grad(grad_u_chain)
- sigma_u = k["lambda"] * grad_u_chain.trace() * sp.eye(grad_u_chain.shape[0]) + k["mu"] * 2 * symm_grad_u
- symm_grad_v = symm_grad(grad_v_chain)
-
- return tabulate(
- double_contraction(sigma_u, symm_grad_v) * jac_a_abs_det
- )
+# HyTeG Operator Generator
+# Copyright (C) 2024 HyTeG Team
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+from hog.recipes.common import *
+import sympy as sp
+
+
+def integrand(
+ *,
+ jac_a_inv,
+ jac_a_abs_det,
+ grad_u,
+ grad_v,
+ k,
+ tabulate,
+ **_,
+):
+
+ grad_u_chain = jac_a_inv.T * grad_u
+ grad_v_chain = jac_a_inv.T * grad_v
+
+ def symm_grad(w):
+ return 0.5 * (w + w.T)
+
+ symm_grad_u = symm_grad(grad_u_chain)
+ sigma_u = k["lambda"] * grad_u_chain.trace() * sp.eye(grad_u_chain.shape[0]) + k["mu"] * 2 * symm_grad_u
+ symm_grad_v = symm_grad(grad_v_chain)
+
+ return tabulate(
+ double_contraction(sigma_u, symm_grad_v) * jac_a_abs_det
+ )