Add more types to linalg

doc/conf.py

@@ -13,8 +13,8 @@
 
 autodoc_type_aliases = {
         "GeometryLike": "pytential.collection.GeometryLike",
-        "DiscretizationStages": "pytential.symbolic.dof_desc.DiscretizationStages",
-        "DOFGranularities": "pytential.symbolic.dof_desc.DOFGranularities",
+        "DiscretizationStage": "pytential.symbolic.dof_desc.DiscretizationStage",
+        "DOFGranularity": "pytential.symbolic.dof_desc.DOFGranularity",
         "DOFDescriptorLike": "pytential.symbolic.dof_desc.DOFDescriptorLike",
         }
 
@@ -40,41 +40,54 @@
     ["py:class", r"TypeAliasForwardRef"],
     ["py:class", r"arraycontext.typing._UserDefinedArrayContainer"],
     ["py:class", r"arraycontext.typing._UserDefinedArithArrayContainer"],
-    ["py:class", r"T"],
+    ["py:class", r"TreeKind"],
 ]
 
 
 sphinxconfig_missing_reference_aliases = {
     # numpy
     "NDArray": "obj:numpy.typing.NDArray",
+    "np.integer": "obj:numpy.integer",
+    "np.floating": "obj:numpy.floating",
     # pytools
     "ObjectArrayND": "obj:pytools.obj_array.ObjectArrayND",
+    "T": "obj:pytools.T",
+    "obj_array.ObjectArray1D": "obj:pytools.obj_array.ObjectArray1D",
     # pyopencl
     "WaitList": "obj:pyopencl.WaitList",
     # pymbolic
-    "Variable": "class:pymbolic.primitives.Variable",
-    "Expression": "obj:pymbolic.typing.Expression",
     "ArithmeticExpression": "obj:pymbolic.ArithmeticExpression",
+    "Expression": "obj:pymbolic.typing.Expression",
     "MultiVector": "obj:pymbolic.geometric_algebra.MultiVector",
+    "Variable": "class:pymbolic.primitives.Variable",
     # arraycontext
-    "ScalarLike": "obj:arraycontext.ScalarLike",
     "ArrayOrContainerOrScalar": "obj:arraycontext.ArrayOrContainerOrScalar",
     "PyOpenCLArrayContext": "class:arraycontext.PyOpenCLArrayContext",
+    "ScalarLike": "obj:arraycontext.ScalarLike",
     # modepy
     "mp.Shape": "class:modepy.Shape",
     # meshmode
     "Discretization": "class:meshmode.discretization.Discretization",
     "DOFArray": "class:meshmode.dof_array.DOFArray",
+    # boxtree
+    # "TreeKind": "obj:boxtree.tree_build.TreeKind",
     # sumpy
+    "ExpansionBase": "class:sumpy.expansion.ExpansionBase",
+    "ExpansionFactoryBase": "class:sumpy.expansion.ExpansionFactoryBase",
     "Kernel": "class:sumpy.kernel.Kernel",
     "P2PBase": "class:sumpy.p2p.P2PBase",
-    "ExpansionFactoryBase": "class:sumpy.expansion.ExpansionFactoryBase",
     # pytential
-    "sym.IntG": "class:pytential.symbolic.primitives.IntG",
-    "sym.DOFDescriptor": "class:pytential.symbolic.dof_desc.DOFDescriptor",
+    "DOFDescriptorLike": "data:pytential.symbolic.dof_desc.DOFDescriptorLike",
+    "DOFGranularity": "data:pytential.symbolic.dof_desc.DOFGranularity",
+    "DiscretizationStage": "data:pytential.symbolic.dof_desc.DiscretizationStage",
+    "GeometryId": "data:pytential.symbolic.dof_desc.GeometryId",
     "Operand": "obj:pytential.symbolic.primitives.Operand",
     "QBXForcedLimit": "obj:pytential.symbolic.primitives.QBXForcedLimit",
     "TargetOrDiscretization": "obj:pytential.target.TargetOrDiscretization",
+    "pytential.symbolic.dof_desc.DOFDescriptorLike":
+        "data:pytential.symbolic.dof_desc.DOFDescriptorLike",
+    "sym.DOFDescriptor": "class:pytential.symbolic.dof_desc.DOFDescriptor",
+    "sym.IntG": "class:pytential.symbolic.primitives.IntG",
 }
 
 

pyproject.toml

@@ -164,11 +164,12 @@ pythonVersion = "3.10"
 pythonPlatform = "All"
 
 ignore = [
-    "pytential/symbolic/old_diffop_primitives.py",
-    "pytential/symbolic/pde/maxwell/generalized_debye.py",
     "build",
+    "doc/_build",
     "doc/conf.py",
     "experiments",
+    "pytential/symbolic/old_diffop_primitives.py",
+    "pytential/symbolic/pde/maxwell/generalized_debye.py",
 ]
 
 [[tool.basedpyright.executionEnvironments]]

pytential/collection.py

@@ -30,8 +30,10 @@
 from typing import Any
 
 from constantdict import constantdict
+from typing_extensions import override
 
 from meshmode.discretization import Discretization
+from meshmode.discretization.connection.direct import DiscretizationConnection
 
 import pytential.symbolic.primitives as sym
 from pytential.qbx import QBXLayerPotentialSource
@@ -40,7 +42,12 @@
     PointPotentialSource,
     PotentialSource,
 )
-from pytential.symbolic.dof_desc import DiscretizationStage, DOFDescriptorLike
+from pytential.symbolic.dof_desc import (
+    DiscretizationStage,
+    DOFDescriptor,
+    DOFDescriptorLike,
+    GeometryId,
+)
 from pytential.target import PointsTarget, TargetBase
 
 
@@ -121,11 +128,17 @@ class GeometryCollection:
 
     """
 
+    ambient_dim: int
+    places: Mapping[GeometryId, GeometryLike]
+    auto_where: tuple[DOFDescriptor, DOFDescriptor]
+
+    _caches: dict[Hashable, Any]
+
     def __init__(self,
             places: (
                 GeometryLike
                 | tuple[GeometryLike, GeometryLike]
-                | Mapping[Hashable, GeometryLike]
+                | Mapping[GeometryId, GeometryLike]
                 ),
             auto_where: AutoWhereLike | None = None) -> None:
         r"""
@@ -153,7 +166,7 @@ def __init__(self,
         from pytential.symbolic.execution import _prepare_auto_where
         auto_source, auto_target = _prepare_auto_where(auto_where)
 
-        places_dict: Mapping[Hashable, GeometryLike]
+        places_dict: Mapping[GeometryId, GeometryLike]
         if isinstance(places, QBXLayerPotentialSource):
             places_dict = {auto_source.geometry: places}
             auto_target = auto_source
@@ -215,7 +228,7 @@ def __init__(self,
         self.places = constantdict(places_dict)
         self.auto_where = (auto_source, auto_target)
 
-        self._caches: dict[str, Any] = {}
+        self._caches = {}
 
     @property
     def auto_source(self) -> sym.DOFDescriptor:
@@ -227,10 +240,12 @@ def auto_target(self) -> sym.DOFDescriptor:
 
     # {{{ cache handling
 
-    def _get_cache(self, name):
+    def _get_cache(self, name: Hashable) -> dict[Hashable, Any]:
         return self._caches.setdefault(name, {})
 
-    def _get_discr_from_cache(self, geometry, discr_stage):
+    def _get_discr_from_cache(self,
+                              geometry: GeometryId,
+                              discr_stage: DiscretizationStage) -> Discretization:
         cache = self._get_cache(_GeometryCollectionDiscretizationCacheKey)
         key = (geometry, discr_stage)
 
@@ -239,9 +254,15 @@ def _get_discr_from_cache(self, geometry, discr_stage):
                     f"cached discretization does not exist on '{geometry}' "
                     f"for stage '{discr_stage}'")
 
-        return cache[key]
+        result = cache[key]
+        assert isinstance(result, Discretization)
 
-    def _add_discr_to_cache(self, discr, geometry, discr_stage):
+        return result
+
+    def _add_discr_to_cache(self,
+                            discr: Discretization,
+                            geometry: GeometryId,
+                            discr_stage: DiscretizationStage) -> None:
         cache = self._get_cache(_GeometryCollectionDiscretizationCacheKey)
         key = (geometry, discr_stage)
 
@@ -251,17 +272,28 @@ def _add_discr_to_cache(self, discr, geometry, discr_stage):
 
         cache[key] = discr
 
-    def _get_conn_from_cache(self, geometry, from_stage, to_stage):
+    def _get_conn_from_cache(self,
+                             geometry: GeometryId,
+                             from_stage: DiscretizationStage | None,
+                             to_stage: DiscretizationStage | None
+                             ) -> DiscretizationConnection:
         cache = self._get_cache(_GeometryCollectionConnectionCacheKey)
         key = (geometry, from_stage, to_stage)
 
         if key not in cache:
             raise KeyError("cached connection does not exist on "
                     f"'{geometry}' from stage '{from_stage}' to '{to_stage}'")
 
-        return cache[key]
+        result = cache[key]
+        assert isinstance(result, DiscretizationConnection)
+
+        return result
 
-    def _add_conn_to_cache(self, conn, geometry, from_stage, to_stage):
+    def _add_conn_to_cache(self,
+                           conn: DiscretizationConnection,
+                           geometry: GeometryId,
+                           from_stage: DiscretizationStage,
+                           to_stage: DiscretizationStage) -> None:
         cache = self._get_cache(_GeometryCollectionConnectionCacheKey)
         key = (geometry, from_stage, to_stage)
 
@@ -271,8 +303,11 @@ def _add_conn_to_cache(self, conn, geometry, from_stage, to_stage):
 
         cache[key] = conn
 
-    def _get_qbx_discretization(self, geometry, discr_stage):
+    def _get_qbx_discretization(self,
+                                geometry: GeometryId,
+                                discr_stage: DiscretizationStage) -> Discretization:
         lpot_source = self.get_geometry(geometry)
+        assert isinstance(lpot_source, LayerPotentialSourceBase)
 
         try:
             discr = self._get_discr_from_cache(geometry, discr_stage)
@@ -291,7 +326,9 @@ def _get_qbx_discretization(self, geometry, discr_stage):
 
     # }}}
 
-    def get_connection(self, from_dd: DOFDescriptorLike, to_dd: DOFDescriptorLike):
+    def get_connection(self,
+                       from_dd: DOFDescriptorLike,
+                       to_dd: DOFDescriptorLike) -> DiscretizationConnection:
         """Construct a connection from *from_dd* to *to_dd* geometries.
 
         :returns: an object compatible with the
@@ -303,7 +340,7 @@ def get_connection(self, from_dd: DOFDescriptorLike, to_dd: DOFDescriptorLike):
         return connection_from_dds(self, from_dd, to_dd)
 
     def get_discretization(
-            self, geometry: Hashable,
+            self, geometry: GeometryId,
             discr_stage: DiscretizationStage | None = None
             ) -> GeometryLike:
         """Get the geometry or discretization in the collection.
@@ -334,7 +371,7 @@ def get_discretization(
         else:
             return discr
 
-    def get_geometry(self, geometry: Hashable) -> GeometryLike:
+    def get_geometry(self, geometry: GeometryId) -> GeometryLike:
         """
         :arg geometry: the identifier of the geometry in the collection.
         """
@@ -346,7 +383,7 @@ def get_geometry(self, geometry: Hashable) -> GeometryLike:
 
     def copy(
             self,
-            places: Mapping[Hashable, GeometryLike] | None = None,
+            places: Mapping[GeometryId, GeometryLike] | None = None,
             auto_where: AutoWhereLike | None = None,
             ) -> GeometryCollection:
         """Get a shallow copy of the geometry collection."""
@@ -356,7 +393,7 @@ def copy(
 
     def merge(
             self,
-            places: GeometryCollection | Mapping[Hashable, GeometryLike],
+            places: GeometryCollection | Mapping[GeometryId, GeometryLike],
             ) -> GeometryCollection:
         """Merges two geometry collections and returns the new collection.
 
@@ -373,10 +410,12 @@ def merge(
 
         return self.copy(places=new_places)
 
-    def __repr__(self):
+    @override
+    def __repr__(self) -> str:
         return f"{type(self).__name__}({self.places!r})"
 
-    def __str__(self):
+    @override
+    def __str__(self) -> str:
         return f"{type(self).__name__}({self.places!r})"
 
 
@@ -387,7 +426,7 @@ def __str__(self):
 
 def add_geometry_to_collection(
         places: GeometryCollection,
-        geometries: Mapping[Hashable, GeometryLike]) -> GeometryCollection:
+        geometries: Mapping[GeometryId, GeometryLike]) -> GeometryCollection:
     """Adds a mapping of geometries to an existing collection.
 
     This function is similar to :meth:`GeometryCollection.merge`, but it makes