Start adding some typing and documentation

Author: jorgensd

ffcx/codegeneration/access.py

@@ -13,6 +13,7 @@
 import ufl
 
 import ffcx.codegeneration.lnodes as L
+from ffcx.definitions import entity_types
 from ffcx.ir.analysis.modified_terminals import ModifiedTerminal
 from ffcx.ir.elementtables import UniqueTableReferenceT
 from ffcx.ir.representationutils import QuadratureRule
@@ -23,7 +24,9 @@
 class FFCXBackendAccess:
     """FFCx specific formatter class."""
 
-    def __init__(self, entity_type: str, integral_type: str, symbols, options):
+    entity_type: entity_types
+
+    def __init__(self, entity_type: entity_types, integral_type: str, symbols, options):
         """Initialise."""
         # Store ir and options
         self.entity_type = entity_type
@@ -399,7 +402,7 @@ def _pass(self, *args, **kwargs):
     def table_access(
         self,
         tabledata: UniqueTableReferenceT,
-        entity_type: str,
+        entity_type: entity_types,
         restriction: str,
         quadrature_index: L.MultiIndex,
         dof_index: L.MultiIndex,
@@ -408,7 +411,7 @@ def table_access(
 
         Args:
             tabledata: Table data object
-            entity_type: Entity type ("cell", "facet", "vertex")
+            entity_type: Entity type
             restriction: Restriction ("+", "-")
             quadrature_index: Quadrature index
             dof_index: Dof index

ffcx/codegeneration/definitions.py

@@ -11,6 +11,7 @@
 import ufl
 
 import ffcx.codegeneration.lnodes as L
+from ffcx.definitions import entity_types
 from ffcx.ir.analysis.modified_terminals import ModifiedTerminal
 from ffcx.ir.elementtables import UniqueTableReferenceT
 from ffcx.ir.representationutils import QuadratureRule
@@ -50,7 +51,9 @@ def create_dof_index(tabledata, dof_index_symbol):
 class FFCXBackendDefinitions:
     """FFCx specific code definitions."""
 
-    def __init__(self, entity_type: str, integral_type: str, access, options):
+    entity_type: entity_types
+
+    def __init__(self, entity_type: entity_types, integral_type: str, access, options):
         """Initialise."""
         # Store ir and options
         self.integral_type = integral_type

ffcx/codegeneration/symbols.py

@@ -10,6 +10,7 @@
 import ufl
 
 import ffcx.codegeneration.lnodes as L
+from ffcx.definitions import entity_types
 
 logger = logging.getLogger("ffcx")
 
@@ -95,7 +96,7 @@ def __init__(self, coefficient_numbering, coefficient_offsets, original_constant
         # Table for chunk of custom quadrature points (physical coordinates).
         self.custom_points_table = L.Symbol("points_chunk", dtype=L.DataType.REAL)
 
-    def entity(self, entity_type, restriction):
+    def entity(self, entity_type: entity_types, restriction):
         """Entity index for lookup in element tables."""
         if entity_type == "cell":
             # Always 0 for cells (even with restriction)
@@ -175,7 +176,7 @@ def constant_index_access(self, constant, index):
         return c[offset + index]
 
     # TODO: Remove this, use table_access instead
-    def element_table(self, tabledata, entity_type, restriction):
+    def element_table(self, tabledata, entity_type: entity_types, restriction):
         """Get an element table."""
         entity = self.entity(entity_type, restriction)
 

ffcx/definitions.py

@@ -0,0 +1,5 @@
+"""Module for storing type definitions used in the FFCx code base."""
+
+from typing import Literal
+
+entity_types = Literal["cell", "facet", "vertex"]

ffcx/ir/elementtables.py

@@ -13,8 +13,11 @@
 import numpy.typing as npt
 import ufl
 
+from ffcx.definitions import entity_types
 from ffcx.element_interface import basix_index
+from ffcx.ir.analysis.modified_terminals import ModifiedTerminal
 from ffcx.ir.representationutils import (
+    QuadratureRule,
     create_quadrature_points_and_weights,
     integral_type_to_entity_dim,
     map_integral_points,
@@ -82,7 +85,7 @@ def get_ffcx_table_values(
     integral_type,
     element,
     avg,
-    entity_type,
+    entity_type: entity_types,
     derivative_counts,
     flat_component,
     codim,
@@ -175,7 +178,12 @@ def get_ffcx_table_values(
 
 
 def generate_psi_table_name(
-    quadrature_rule, element_counter, averaged: str, entity_type, derivative_counts, flat_component
+    quadrature_rule: QuadratureRule,
+    element_counter,
+    averaged: str,
+    entity_type: entity_types,
+    derivative_counts,
+    flat_component,
 ):
     """Generate a name for the psi table.
 
@@ -293,26 +301,33 @@ def permute_quadrature_quadrilateral(points, reflections=0, rotations=0):
 
 
 def build_optimized_tables(
-    quadrature_rule,
-    cell,
-    integral_type,
-    entity_type,
-    modified_terminals,
-    existing_tables,
-    use_sum_factorization,
-    is_mixed_dim,
-    rtol=default_rtol,
-    atol=default_atol,
-):
+    quadrature_rule: QuadratureRule,
+    cell: ufl.Cell,
+    integral_type: str,
+    entity_type: entity_types,
+    modified_terminals: ModifiedTerminal,
+    existing_tables: dict[str, np.ndarray],
+    use_sum_factorization: bool,
+    is_mixed_dim: bool,
+    rtol: float = default_rtol,
+    atol: float = default_atol,
+) -> dict[ModifiedTerminal, UniqueTableReferenceT]:
     """Build the element tables needed for a list of modified terminals.
 
-    Input:
-      entity_type - str
-      modified_terminals - ordered sequence of unique modified terminals
-      FIXME: Document
-
-    Output:
-      mt_tables - dict(ModifiedTerminal: table data)
+    Args:
+        quadrature_rule: The quadrature rule relating to the tables.
+        cell: The cell type of the domain the tables will be used with.
+        entity_type: On what entity (vertex,edge,facet,cell) the tables are evaluated at.
+        integral_type: The type of integral the tables are used for.
+        modified_terminals: ordered sequence of unique modified terminals
+        existing_tables: Register of tables that already exist and reused.
+        use_sum_factorization: Use sum factorization for tensor product elements.
+        is_mixed_dim: Mixed dimensionality of the domain.
+        rtol: Relative tolerance for comparing tables.
+        atol: Absolute tolerance for comparing tables.
+
+    Returns:
+      mt_tables - Dictionary mapping each modified terminal to the a unique table reference.
     """
     # Add to element tables
     analysis = {}

ffcx/ir/integral.py

@@ -16,6 +16,7 @@
 from ufl.checks import is_cellwise_constant
 from ufl.classes import QuadratureWeight
 
+from ffcx.definitions import entity_types
 from ffcx.ir.analysis.factorization import compute_argument_factorization
 from ffcx.ir.analysis.graph import build_scalar_graph
 from ffcx.ir.analysis.modified_terminals import analyse_modified_terminal, is_modified_terminal
@@ -46,7 +47,9 @@ class BlockDataT(typing.NamedTuple):
     is_permuted: bool  # Do quad points on facets need to be permuted?
 
 
-def compute_integral_ir(cell, integral_type, entity_type, integrands, argument_shape, p, visualise):
+def compute_integral_ir(
+    cell, integral_type: str, entity_type: entity_types, integrands, argument_shape, p, visualise
+):
     """Compute intermediate representation for an integral."""
     # The intermediate representation dict we're building and returning
     # here