Skip to content

Commit

Permalink
TSFC: cache QuadratureRule
Browse files Browse the repository at this point in the history
  • Loading branch information
pbrubeck committed Dec 13, 2024
1 parent 9de4443 commit dcb322d
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 16 deletions.
23 changes: 13 additions & 10 deletions tsfc/fem.py
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,18 @@ def needs_coordinate_mapping(element):
return isinstance(create_element(element), NeedsCoordinateMappingElement)


@serial_cache(hashkey=lambda *args: args)
def get_quadrature_rule(fiat_cell, integration_dim, quadrature_degree, scheme):
integration_cell = fiat_cell.construct_subcomplex(integration_dim)
return make_quadrature(integration_cell, quadrature_degree, scheme=scheme)


def make_basis_evaluation_key(ctx, finat_element, mt, entity_id):
domain = extract_unique_domain(mt.terminal)
coordinate_element = domain.ufl_coordinate_element()
return (finat_element, mt.local_derivatives, ctx.point_set, ctx.integration_dim, entity_id, coordinate_element, mt.restriction)


class PointSetContext(ContextBase):
"""Context for compile-time known evaluation points."""

Expand All @@ -273,13 +285,9 @@ class PointSetContext(ContextBase):
'weight_expr',
)

@cached_property
def integration_cell(self):
return self.fiat_cell.construct_subelement(self.integration_dim)

@cached_property
def quadrature_rule(self):
return make_quadrature(self.integration_cell, self.quadrature_degree)
return get_quadrature_rule(self.fiat_cell, self.integration_dim, self.quadrature_degree, "default")

@cached_property
def point_set(self):
Expand All @@ -297,11 +305,6 @@ def point_expr(self):
def weight_expr(self):
return self.quadrature_rule.weight_expression

def make_basis_evaluation_key(self, finat_element, mt, entity_id):
domain = extract_unique_domain(mt.terminal)
coordinate_element = domain.ufl_coordinate_element()
return (self, finat_element, mt.local_derivatives, coordinate_element, mt.restriction, entity_id)

@serial_cache(hashkey=make_basis_evaluation_key)
def basis_evaluation(self, finat_element, mt, entity_id):
return finat_element.basis_evaluation(mt.local_derivatives,
Expand Down
11 changes: 5 additions & 6 deletions tsfc/kernel_interface/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import numpy
from FIAT.reference_element import TensorProductCell
from finat.cell_tools import max_complex
from finat.quadrature import AbstractQuadratureRule, make_quadrature
from finat.quadrature import AbstractQuadratureRule
from gem.node import traversal
from gem.optimise import constant_fold_zero
from gem.optimise import remove_componenttensors as prune
Expand Down Expand Up @@ -136,7 +136,7 @@ def compile_integrand(self, integrand, params, ctx):
integrand = ufl_utils.split_coefficients(integrand, self.coefficient_split)
# Compile: ufl -> gem
info = self.integral_data_info
functions = list(info.arguments) + [self.coordinate(info.domain)] + list(info.coefficients)
functions = [*info.arguments, self.coordinate(info.domain), *info.coefficients]
set_quad_rule(params, info.domain.ufl_cell(), info.integral_type, functions)
quad_rule = params["quadrature_rule"]
config = self.fem_config()
Expand Down Expand Up @@ -319,8 +319,7 @@ def set_quad_rule(params, cell, integral_type, functions):
fiat_cell = max_complex(fiat_cells)

integration_dim, _ = lower_integral_type(fiat_cell, integral_type)
integration_cell = fiat_cell.construct_subcomplex(integration_dim)
quad_rule = make_quadrature(integration_cell, quadrature_degree, scheme=scheme)
quad_rule = fem.get_quadrature_rule(fiat_cell, integration_dim, quadrature_degree, scheme)
params["quadrature_rule"] = quad_rule

if not isinstance(quad_rule, AbstractQuadratureRule):
Expand All @@ -329,8 +328,8 @@ def set_quad_rule(params, cell, integral_type, functions):


def get_index_ordering(quadrature_indices, return_variables):
split_argument_indices = tuple(chain(*[var.index_ordering()
for var in return_variables]))
split_argument_indices = tuple(chain(*(var.index_ordering()
for var in return_variables)))
return tuple(quadrature_indices) + split_argument_indices


Expand Down

0 comments on commit dcb322d

Please sign in to comment.