Improve typing in models.em

.github/workflows/ci.yml

@@ -31,7 +31,7 @@ jobs:
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
                 build_py_project_in_conda_env
-                cipip install pytest modepy
+                cipip install pytest modepy optype
                 cipip install basedpyright
                 basedpyright
 

doc/conf.py

@@ -40,8 +40,8 @@
 
 sphinxconfig_missing_reference_aliases = {
     # numpy
-    "np.floating": "class:numpy.floating",
     "DTypeLike": "obj:numpy.typing.DTypeLike",
+    "np.floating": "class:numpy.floating",
 
     # mpi4py
     "Intracomm": "mpi4py.MPI.Intracomm",
@@ -54,26 +54,26 @@
     "cl_array.Allocator": "class:pyopencl.array.Allocator",
 
     # actx
-    "arraycontext.typing.ArithArrayContainerT":
-        "obj:arraycontext.ArithArrayContainerT",
-    "ScalarLike": "obj:arraycontext.ScalarLike",
-    "Array": "obj:arraycontext.Array",
-    "ArrayContainer": "obj:arraycontext.ArrayContainer",
     "ArithArrayContainer": "obj:arraycontext.ArithArrayContainer",
     "ArithArrayContainerT": "obj:arraycontext.ArithArrayContainerT",
-    "ArrayOrContainer": "obj:arraycontext.ArrayOrContainer",
+    "Array": "obj:arraycontext.Array",
+    "ArrayContainer": "obj:arraycontext.ArrayContainer",
     "ArrayOrArithContainer": "obj:arraycontext.ArrayOrArithContainer",
+    "ArrayOrContainer": "obj:arraycontext.ArrayOrContainer",
+    "ArrayOrContainerOrScalar": "obj:arraycontext.ArrayOrContainerOrScalar",
     "ArrayOrContainerOrScalarT": "obj:arraycontext.ArrayOrContainerOrScalarT",
+    "ScalarLike": "obj:arraycontext.ScalarLike",
+    "arraycontext.typing.ArithArrayContainerT": "obj:arraycontext.ArithArrayContainerT",
 
     # meshmode
     "DOFArray": "meshmode.dof_array.DOFArray",
-    "Mesh": "meshmode.mesh.Mesh",
     "Discretization": "class:meshmode.discretization.Discretization",
+    "Mesh": "meshmode.mesh.Mesh",
 
     # grudge
     "DiscretizationTag": "obj:grudge.dof_desc.DiscretizationTag",
-    "VolumeTag": "obj:grudge.dof_desc.VolumeTag",
     "TracePair": "class:grudge.trace_pair.TracePair",
+    "VolumeTag": "obj:grudge.dof_desc.VolumeTag",
 }
 
 

grudge/geometry/metrics.py

@@ -60,8 +60,8 @@
 """
 
 
-from functools import reduce
-from operator import add
+from functools import partial, reduce
+from operator import add, xor
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -75,8 +75,8 @@
     DiscretizationAmbientDimAxisTag,
     DiscretizationTopologicalDimAxisTag,
 )
-from pymbolic.geometric_algebra import MultiVector
-from pytools import memoize_in
+from pymbolic.geometric_algebra import MultiVector, get_euclidean_space
+from pytools import flatten, memoize_in, product
 
 import grudge.dof_desc as dof_desc
 from grudge import DiscretizationCollection
@@ -96,9 +96,9 @@ def _geometry_to_quad_if_requested(
             dd: DOFDesc,
             vec: DOFArray,
             _use_geoderiv_connection: bool
-        ):
+        ) -> DOFArray:
 
-    def to_quad(vec: DOFArray):
+    def to_quad(vec: DOFArray) -> DOFArray:
         if not dd.uses_quadrature():
             return vec
         return dcoll.connection_from_dds(inner_dd, dd)(vec)
@@ -172,15 +172,15 @@ def forward_metric_nth_derivative(
         ref_axes = ((ref_axes, 1),)
 
     if not isinstance(ref_axes, tuple):
-        raise ValueError("ref_axes must be a tuple")
+        raise ValueError(f"'ref_axes' must be a tuple: {type(ref_axes)}")
 
     if tuple(sorted(ref_axes)) != ref_axes:
-        raise ValueError("ref_axes must be sorted")
+        raise ValueError(f"'ref_axes' must be sorted: {ref_axes}")
 
     if len(set(ref_axes)) != len(ref_axes):
-        raise ValueError("ref_axes must not contain an axis more than once")
+        raise ValueError(
+            f"'ref_axes' must not contain an axis more than once: {ref_axes}")
 
-    from pytools import flatten
     flat_ref_axes = flatten([rst_axis] * n for rst_axis, n in ref_axes)
 
     from meshmode.discretization import num_reference_derivative
@@ -453,11 +453,9 @@ def inverse_metric_derivative(
             for rst in range(dim)]
 
     # Yay Cramer's rule!
-    from functools import partial, reduce
-    from operator import xor as outerprod_op
-    outerprod = partial(reduce, outerprod_op)
+    outerprod = partial(reduce, xor)
 
-    def outprod_with_unit(i: int, at: int):
+    def outprod_with_unit(i: int, at: int) -> MultiVector[DOFArray]:
         unit_vec = np.zeros(dim)
         unit_vec[i] = 1
 
@@ -532,7 +530,7 @@ def inverse_surface_metric_derivative_mat(
             dd: ToDOFDescConvertible = None,
             *, times_area_element: bool = False,
             _use_geoderiv_connection: bool = False
-        ):
+        ) -> DOFArray:
     r"""Computes the matrix of inverse surface metric derivatives, indexed by
     ``(xyz_axis, rst_axis)``. It returns all values of
     :func:`inverse_surface_metric_derivative_mat` in cached matrix form.
@@ -557,7 +555,7 @@ def inverse_surface_metric_derivative_mat(
 
     @memoize_in(dcoll, (inverse_surface_metric_derivative_mat, dd,
         times_area_element, _use_geoderiv_connection))
-    def _inv_surf_metric_deriv():
+    def _inv_surf_metric_deriv() -> DOFArray:
         if times_area_element:
             multiplier = area_element(actx, dcoll, dd=dd,
                     _use_geoderiv_connection=_use_geoderiv_connection)
@@ -600,7 +598,7 @@ def _signed_face_ones(
 
     discr = dcoll.discr_from_dd(dd.with_discr_tag(DISCR_TAG_BASE))
 
-    new_group_arrays = []
+    new_group_arrays: list[np.ndarray[tuple[int, int], np.dtype[np.floating]]] = []
 
     for dgrp, grp in zip(discr.groups, all_faces_conn.groups, strict=True):
         sign = np.ones((dgrp.nelements, dgrp.nunit_dofs),
@@ -613,7 +611,6 @@ def _signed_face_ones(
 
         new_group_arrays.append(sign)
 
-    from meshmode.dof_array import DOFArray
     return actx.from_numpy(DOFArray(actx, tuple(new_group_arrays)))
 
 
@@ -638,15 +635,11 @@ def parametrization_derivative(
 
     dim = dcoll.discr_from_dd(dd).dim
     if dim == 0:
-        from pymbolic.geometric_algebra import get_euclidean_space
-
         return MultiVector(
             _signed_face_ones(actx, dcoll, dd),
             space=get_euclidean_space(dcoll.ambient_dim)
         )
 
-    from pytools import product
-
     res = product(
         forward_metric_derivative_mv(
             actx, dcoll, rst_axis, dd,
@@ -655,6 +648,7 @@ def parametrization_derivative(
     )
     if isinstance(res, int):
         return MultiVector(res)
+
     return res
 
 
@@ -703,11 +697,12 @@ def area_element(
         dd = DD_VOLUME_ALL
 
     @memoize_in(dcoll, (area_element, dd, _use_geoderiv_connection))
-    def _area_elements():
+    def _area_elements() -> DOFArray:
         result = actx.np.sqrt(
             pseudoscalar(
                 actx, dcoll, dd=dd,
                 _use_geoderiv_connection=_use_geoderiv_connection).norm_squared())
+        assert isinstance(result, DOFArray)
 
         return actx.freeze(
                 actx.tag(NameHint(f"area_el_{dd.as_identifier()}"), result))
@@ -783,7 +778,7 @@ def mv_normal(
         use_geoderiv_connection = actx.supports_nonscalar_broadcasting
 
     @memoize_in(dcoll, (mv_normal, dd, use_geoderiv_connection))
-    def _normal():
+    def _normal() -> MultiVector[DOFArray]:
         dim = dcoll.discr_from_dd(dd).dim
         ambient_dim = dcoll.ambient_dim
 
@@ -834,8 +829,11 @@ def _normal():
     return actx.thaw(_normal())
 
 
-def normal(actx: ArrayContext, dcoll: DiscretizationCollection, dd: DOFDesc,
-        *, _use_geoderiv_connection: bool | None = None):
+def normal(actx: ArrayContext,
+           dcoll: DiscretizationCollection,
+           dd: DOFDesc,
+           *, _use_geoderiv_connection: bool | None = None,
+           ) -> obj_array.ObjectArray1D[DOFArray]:
     """Get the unit normal to the specified surface discretization, *dd*.
     This supports both volume discretizations
     (where ambient == topological dimension) and surface discretizations
@@ -897,8 +895,6 @@ def second_fundamental_form(
     else:
         raise ValueError(f"{dim}D surfaces not supported")
 
-    from pytools import flatten
-
     form2 = np.empty((dim, dim), dtype=object)
 
     for ref_axes in second_ref_axes:

grudge/models/__init__.py

@@ -28,6 +28,13 @@
 """
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from arraycontext import Array, ArrayContext
+
+    from grudge.discretization import DiscretizationCollection
 
 
 class Operator(ABC):  # noqa: B024
@@ -40,11 +47,11 @@ class Operator(ABC):  # noqa: B024
     """
 
 
-class HyperbolicOperator(Operator):
+class HyperbolicOperator(Operator, ABC):
     """A base class for hyperbolic Discontinuous Galerkin operators."""
 
     @abstractmethod
-    def max_characteristic_velocity(self, actx, **kwargs):
+    def max_characteristic_velocity(self, actx: ArrayContext, **kwargs: Any) -> Array:
         r"""Return a maximum characteristic wavespeed for the operator.
 
         :arg actx: a :class:`arraycontext.ArrayContext`.
@@ -56,7 +63,10 @@ def max_characteristic_velocity(self, actx, **kwargs):
         the operator.
         """
 
-    def estimate_rk4_timestep(self, actx, dcoll, **kwargs):
+    def estimate_rk4_timestep(
+            self,
+            actx: ArrayContext,
+            dcoll: DiscretizationCollection, **kwargs: Any) -> Array:
         r"""Estimate the largest stable timestep for an RK4 method.
 
         :arg actx: a :class:`arraycontext.ArrayContext`.