diff --git a/pytential/qbx/__init__.py b/pytential/qbx/__init__.py index 3e4e1026b..a6261ba59 100644 --- a/pytential/qbx/__init__.py +++ b/pytential/qbx/__init__.py @@ -603,6 +603,9 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, geo_data.tree().user_source_ids, insn.kernel_arguments, evaluate)) + from sumpy.fmm import SumpyTranslationClassesData + translation_classes_data = SumpyTranslationClassesData(actx.queue, + geo_data.traversal()) tree_indep = self._tree_indep_data_for_wrangler( target_kernels=insn.target_kernels, source_kernels=insn.source_kernels) @@ -612,7 +615,8 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, self.qbx_order, self.fmm_level_to_order, source_extra_kwargs=source_extra_kwargs, - kernel_extra_kwargs=kernel_extra_kwargs) + kernel_extra_kwargs=kernel_extra_kwargs, + translation_classes_data=translation_classes_data) from pytential.qbx.geometry import target_state if actx.to_numpy(actx.np.any( diff --git a/pytential/qbx/fmm.py b/pytential/qbx/fmm.py index 91cce641a..5bf14c535 100644 --- a/pytential/qbx/fmm.py +++ b/pytential/qbx/fmm.py @@ -112,13 +112,14 @@ class QBXExpansionWrangler(SumpyExpansionWrangler): def __init__(self, tree_indep, geo_data, dtype, qbx_order, fmm_level_to_order, source_extra_kwargs, kernel_extra_kwargs, - _use_target_specific_qbx=None): + translation_classes_data=None, *, _use_target_specific_qbx=None): if _use_target_specific_qbx: raise ValueError("TSQBX is not implemented in sumpy") super().__init__( tree_indep, geo_data.traversal(), - dtype, fmm_level_to_order, source_extra_kwargs, kernel_extra_kwargs) + dtype, fmm_level_to_order, source_extra_kwargs, kernel_extra_kwargs, + translation_classes_data=translation_classes_data) self.qbx_order = qbx_order self.geo_data = geo_data diff --git a/pytential/qbx/fmmlib.py b/pytential/qbx/fmmlib.py index c2029ad1b..e91a9fe62 100644 --- a/pytential/qbx/fmmlib.py +++ b/pytential/qbx/fmmlib.py @@ -153,7 +153,8 @@ class QBXFMMLibExpansionWrangler(FMMLibExpansionWrangler): def __init__(self, tree_indep, geo_data, dtype, qbx_order, fmm_level_to_order, source_extra_kwargs, - kernel_extra_kwargs, + kernel_extra_kwargs, *, + translation_classes_data=None, _use_target_specific_qbx=None): # FMMLib is CPU-only. This wrapper gets the geometry out of # OpenCL-land. diff --git a/pytential/unregularized.py b/pytential/unregularized.py index 8502b8a8f..00bba1c7b 100644 --- a/pytential/unregularized.py +++ b/pytential/unregularized.py @@ -256,12 +256,16 @@ def exec_compute_potential_insn_fmm(self, actx: PyOpenCLArrayContext, tree_indep = self._tree_indep_data_for_wrangler( fmm_kernel, target_kernels=insn.target_kernels, source_kernels=insn.source_kernels) + from sumpy.fmm import SumpyTranslationClassesData + translation_classes_data = SumpyTranslationClassesData(actx.queue, + geo_data.traversal()) from sumpy.fmm import SumpyExpansionWrangler wrangler = SumpyExpansionWrangler( tree_indep, geo_data.traversal(), output_and_expansion_dtype, self.fmm_level_to_order, + translation_classes_data=translation_classes_data, source_extra_kwargs=source_extra_kwargs, kernel_extra_kwargs=kernel_extra_kwargs)