From 7d5bb22193983e2c24c6c92cb3ad08dd57007701 Mon Sep 17 00:00:00 2001
From: Alexandru Fikl <alexfikl@gmail.com>
Date: Thu, 22 Sep 2022 19:30:02 +0300
Subject: [PATCH] use the full double-layer in source-proxy skeletonization

---
 pytential/linalg/direct_solver_symbolic.py | 8 ++++++--
 pytential/linalg/skeletonization.py        | 6 ++++--
 pytential/symbolic/matrix.py               | 9 +++------
 3 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/pytential/linalg/direct_solver_symbolic.py b/pytential/linalg/direct_solver_symbolic.py
index c57752840..5c11eee1d 100644
--- a/pytential/linalg/direct_solver_symbolic.py
+++ b/pytential/linalg/direct_solver_symbolic.py
@@ -54,12 +54,16 @@ def prepare_expr(places, exprs, auto_where=None):
     return _prepare_expr(places, exprs, auto_where=auto_where)
 
 
-def prepare_proxy_expr(places, exprs, auto_where=None):
+def prepare_proxy_expr(
+        places, exprs,
+        auto_where=None,
+        remove_transforms: bool = True):
     def _prepare_expr(expr):
         # remove all diagonal / non-operator terms in the expression
         expr = IntGTermCollector()(expr)
         # ensure all IntGs remove all the kernel derivatives
-        expr = KernelTransformationRemover()(expr)
+        if remove_transforms:
+            expr = KernelTransformationRemover()(expr)
         # ensure all IntGs have their source and targets set
         expr = DOFDescriptorReplacer(auto_where[0], auto_where[1])(expr)
 
diff --git a/pytential/linalg/skeletonization.py b/pytential/linalg/skeletonization.py
index 30f72abdb..50ce51c5b 100644
--- a/pytential/linalg/skeletonization.py
+++ b/pytential/linalg/skeletonization.py
@@ -446,9 +446,11 @@ def make_skeletonization_wrangler(
 
     exprs = prepare_expr(places, exprs, auto_where)
     source_proxy_exprs = prepare_proxy_expr(
-            places, exprs, (auto_where[0], PROXY_SKELETONIZATION_TARGET))
+            places, exprs, (auto_where[0], PROXY_SKELETONIZATION_TARGET),
+            remove_transforms=True)
     target_proxy_exprs = prepare_proxy_expr(
-            places, exprs, (PROXY_SKELETONIZATION_SOURCE, auto_where[1]))
+            places, exprs, (PROXY_SKELETONIZATION_SOURCE, auto_where[1]),
+            remove_transforms=True)
 
     # }}}
 
diff --git a/pytential/symbolic/matrix.py b/pytential/symbolic/matrix.py
index 24fa84226..c52e94017 100644
--- a/pytential/symbolic/matrix.py
+++ b/pytential/symbolic/matrix.py
@@ -651,7 +651,6 @@ def map_int_g(self, expr):
                 expr.target.geometry, expr.target.discr_stage)
 
         actx = self.array_context
-        target_base_kernel = expr.target_kernel.get_base_kernel()
 
         result = 0
         for kernel, density in zip(expr.source_kernels, expr.densities):
@@ -672,12 +671,10 @@ def map_int_g(self, expr):
 
             # {{{ generator
 
-            base_kernel = kernel.get_base_kernel()
-
             from sumpy.p2p import P2PMatrixSubsetGenerator
             mat_gen = P2PMatrixSubsetGenerator(actx.context,
-                    source_kernels=(base_kernel,),
-                    target_kernels=(target_base_kernel,),
+                    source_kernels=(kernel,),
+                    target_kernels=(expr.target_kernel,),
                     exclude_self=self.exclude_self)
 
             # }}}
@@ -687,7 +684,7 @@ def map_int_g(self, expr):
             # {{{ kernel args
 
             # NOTE: copied from pytential.symbolic.primitives.IntG
-            kernel_args = base_kernel.get_args() + base_kernel.get_source_args()
+            kernel_args = kernel.get_args() + kernel.get_source_args()
             kernel_args = {arg.loopy_arg.name for arg in kernel_args}
 
             kernel_args = _get_layer_potential_args(