Use GeometryCollection and fix Pylint

Author: gaohao95

examples/off-surface-eval.py

@@ -23,6 +23,7 @@
 import numpy as np
 import functools
 import pytential
+from pytential.collection import GeometryCollection
 import matplotlib.pyplot as plt
 from mpi4py import MPI
 
@@ -35,9 +36,10 @@ def main():
     queue = cl.CommandQueue(cl_context)
     actx = PyOpenCLArrayContext(queue)
 
-    sigma = None
     places = None
     op = None
+    sigma = None
+
     if mpi_rank == 0:
         nelements = 30
         target_order = 8
@@ -57,7 +59,6 @@ def main():
         discretization = Discretization(
             actx, mesh, InterpolatoryQuadratureSimplexGroupFactory(target_order))
 
-        # FIXME: Use GeometryCollection instead of DistributedQBXLayerPotentialSource
         from pytential.qbx.distributed import DistributedQBXLayerPotentialSource
         layer_pot_source = DistributedQBXLayerPotentialSource(
             comm,
@@ -78,7 +79,7 @@ def main():
         fplot = FieldPlotter(np.zeros(2), extent=0.54, npoints=30)
         targets = PointsTarget(fplot.points)
 
-        places = (layer_pot_source, targets)
+        places = GeometryCollection((layer_pot_source, targets))
 
     from pytential.symbolic.execution import bind_distributed
     bound_op = bind_distributed(comm, places, op)

pytential/qbx/__init__.py

@@ -548,11 +548,6 @@ def _tree_indep_data_for_wrangler(self, source_kernels, target_kernels):
         fmm_local_factory = partial(local_expn_class, base_kernel)
         qbx_local_factory = partial(qbx_local_expn_class, base_kernel)
 
-        from pytential.qbx.distributed import DistributedQBXLayerPotentialSource
-        use_distributed = False
-        if isinstance(self, DistributedQBXLayerPotentialSource):
-            use_distributed = True
-
         if self.fmm_backend == "sumpy":
             from pytential.qbx.fmm import \
                     QBXSumpyTreeIndependentDataForWrangler
@@ -575,8 +570,7 @@ def _tree_indep_data_for_wrangler(self, source_kernels, target_kernels):
                     local_expansion_factory=fmm_local_factory,
                     qbx_local_expansion_factory=qbx_local_factory,
                     target_kernels=target_kernels_new,
-                    _use_target_specific_qbx=self._use_target_specific_qbx,
-                    use_distributed=use_distributed)
+                    _use_target_specific_qbx=self._use_target_specific_qbx)
 
         else:
             raise ValueError(f"invalid FMM backend: {self.fmm_backend}")

pytential/qbx/distributed.py

@@ -742,15 +742,24 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext,
         output_and_expansion_dtype = self.comm.bcast(
             output_and_expansion_dtype, root=0)
 
-        wrangler = tree_indep.wrangler_cls(
+        from pytential.qbx.fmmlib import (
+            QBXFMMLibExpansionWrangler, DistributedQBXFMMLibExpansionWrangler)
+
+        wrangler_cls = None
+        if tree_indep.wrangler_cls == QBXFMMLibExpansionWrangler:
+            wrangler_cls = DistributedQBXFMMLibExpansionWrangler
+        else:
+            raise NotImplementedError
+
+        wrangler = wrangler_cls(
                         self._cl_context, self.comm, tree_indep,
                         local_geo_data, global_geo_data,
                         output_and_expansion_dtype,
                         self.qbx_order,
                         self.fmm_level_to_order,
                         source_extra_kwargs,
                         kernel_extra_kwargs,
-                        _use_target_specific_qbx=self._use_target_specific_qbx)
+                        self._use_target_specific_qbx)
 
         if self.comm.Get_rank() == 0:
             from pytential.qbx.geometry import target_state

pytential/qbx/fmmlib.py

@@ -48,13 +48,11 @@ class QBXFMMLibTreeIndependentDataForWrangler(FMMLibTreeIndependentDataForWrangl
     def __init__(self, cl_context, *,
             multipole_expansion_factory, local_expansion_factory,
             qbx_local_expansion_factory, target_kernels,
-            _use_target_specific_qbx,
-            use_distributed):
+            _use_target_specific_qbx):
         self.cl_context = cl_context
         self.multipole_expansion_factory = multipole_expansion_factory
         self.local_expansion_factory = local_expansion_factory
         self.qbx_local_expansion_factory = qbx_local_expansion_factory
-        self.use_distributed = use_distributed
 
         kernel = target_kernels[0].get_base_kernel()
         self.target_kernels = target_kernels
@@ -145,10 +143,7 @@ def is_supported_helmknl_for_tsqbx(knl):
 
     @property
     def wrangler_cls(self):
-        if self.use_distributed:
-            return DistributedQBXFMMLibExpansionWrangler
-        else:
-            return QBXFMMLibExpansionWrangler
+        return QBXFMMLibExpansionWrangler
 
 # }}}