Placate pylint and mypy

Author: gaohao95

pytential/qbx/distributed.py

@@ -22,7 +22,7 @@
 
 from pytential.qbx import QBXLayerPotentialSource
 from arraycontext import PyOpenCLArrayContext, unflatten
-from typing import Optional, Any
+from typing import Any, Dict
 from dataclasses import dataclass
 import numpy as np
 import pyopencl as cl
@@ -135,7 +135,7 @@ def __call__(self, queue, tree, **kwargs):
 
 
 def broadcast_global_geometry_data(
-        comm, actx: PyOpenCLArrayContext, traversal_builder, global_geometry_data):
+        comm, actx, traversal_builder, global_geometry_data):
     """Broadcasts useful fields of global geometry data from the root rank to the
     worker ranks, so that each rank can form local geometry data independently.
 
@@ -212,7 +212,7 @@ def broadcast_global_geometry_data(
 
 
 def compute_local_geometry_data(
-        actx: PyOpenCLArrayContext, comm, global_geometry_data, boxes_time,
+        actx, comm, global_geometry_data, boxes_time,
         traversal_builder):
     """Compute the local geometry data of the current rank from the global geometry
     data.
@@ -572,10 +572,7 @@ def distribute_geo_data(comm, actx, insn, bound_expr, evaluate,
 
 
 class DistributedQBXLayerPotentialSource(QBXLayerPotentialSource):
-    def __init__(self, comm, cl_context, *args,
-                 _use_target_specific_qbx: Optional[bool] = None,
-                 fmm_backend: str = "fmmlib",
-                 **kwargs):
+    def __init__(self, comm, cl_context, *args, **kwargs):
         """
         :arg comm: MPI communicator.
         :arg cl_context: This argument is necessary because although the root rank
@@ -584,6 +581,8 @@ def __init__(self, comm, cl_context, *args,
 
         `*args` and `**kwargs` will be forwarded to
         `QBXLayerPotentialSource.__init__` on the root rank.
+
+        Currently, `fmm_backend` has to be set to `"fmmlib"`.
         """
         self.comm = comm
         self._cl_context = cl_context
@@ -600,17 +599,14 @@ def __init__(self, comm, cl_context, *args,
                 "distributed implementation")
 
         # Only fmmlib is supported
-        assert fmm_backend == "fmmlib"
+        assert kwargs["fmm_backend"] == "fmmlib"
 
         if self.comm.Get_rank() == 0:
-            super().__init__(
-                *args,
-                _use_target_specific_qbx=_use_target_specific_qbx,
-                fmm_backend=fmm_backend,
-                **kwargs)
+            super().__init__(*args, **kwargs)
         else:
-            self._use_target_specific_qbx = _use_target_specific_qbx
-            self.fmm_backend = fmm_backend
+            self.fmm_backend = "fmmlib"
+            self._use_target_specific_qbx = kwargs.get(
+                "_use_target_specific_qbx", None)
             self.qbx_order = kwargs.get("qbx_order", None)
             self.fmm_level_to_order = kwargs.get("fmm_level_to_order", None)
             self.expansion_factory = kwargs.get("expansion_factory", None)
@@ -727,6 +723,7 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext,
 
         user_source_ids = None
         if self.comm.Get_rank() == 0:
+            assert global_geo_data_device is not None
             user_source_ids = global_geo_data_device.tree().user_source_ids
 
         kernel_extra_kwargs, source_extra_kwargs = (
@@ -767,6 +764,7 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext,
                         self._use_target_specific_qbx)
 
         if self.comm.Get_rank() == 0:
+            assert global_geo_data_device is not None
             from pytential.qbx.geometry import target_state
             if actx.to_numpy(actx.np.any(
                     actx.thaw(global_geo_data_device.user_target_to_center())
@@ -784,11 +782,12 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext,
             # }}}
 
         # Execute global QBX.
-        timing_data = {}
+        timing_data: Dict[str, Any] = {}
         all_potentials_on_every_target = drive_dfmm(
             self.comm, flat_strengths, wrangler, timing_data)
 
         if self.comm.Get_rank() == 0:
+            assert global_geo_data_device is not None
             results = []
 
             for o in insn.outputs:
@@ -830,6 +829,7 @@ def drive_dfmm(comm, src_weight_vecs, wrangler, timing_data=None):
 
     # {{{ Distribute source weights
 
+    template_ary = None
     if current_rank == 0:
         template_ary = src_weight_vecs[0]