Clear code written to combine off/on-axis recurrence

Author: hirish99

sumpy/recurrence.py

@@ -425,7 +425,7 @@ def get_reindexed_and_center_origin_on_axis_recurrence(pde) -> tuple[int, int,
     return n_initial, order, r_s
 
 # ================ OFF-AXIS RECURRENCE =================
-def move_center_origin_source_arbitrary_expression(pde: LinearPDESystemOperator) -> sp.Expr:
+def _move_center_origin_source_arbitrary_expression(pde: LinearPDESystemOperator) -> sp.Expr:
     r"""
     A function that "shifts" the recurrence so it's center is placed
     at the origin and source is the input for the recurrence generated.
@@ -447,14 +447,14 @@ def move_center_origin_source_arbitrary_expression(pde: LinearPDESystemOperator)
     return r_ret
 
 
-def get_reindexed_and_center_origin_off_axis_recurrence(pde: LinearPDESystemOperator) -> sp.Expr:
+def get_reindexed_and_center_origin_off_axis_recurrence(pde: LinearPDESystemOperator) -> [int, int, sp.Expr]:
     r"""
     A function that takes in as input a pde and outputs a off-axis recurrence
     for derivatives taken at the origin with an arbitrary source location.
     The recurrence is reindexed so it gives s(n) in terms of s(n-1), ..., etc.
     """
     var = _make_sympy_vec("x", 1)
-    r_exp = move_center_origin_source_arbitrary_expression(pde).subs(var[0], 0)
+    r_exp = _move_center_origin_source_arbitrary_expression(pde).subs(var[0], 0)
 
     var = _make_sympy_vec("x", 2)
     var_t = _make_sympy_vec("t", 2)
@@ -469,7 +469,7 @@ def get_reindexed_and_center_origin_off_axis_recurrence(pde: LinearPDESystemOper
     return start_order, recur_order, recur
 
 
-def get_off_axis_expression(pde, taylor_order=4):
+def get_off_axis_expression(pde, taylor_order=4) -> [sp.Expr, int]:
     r"""
     A function that takes in as input a pde, and outputs
     the Taylor expression that gives the n th derivative

sumpy/recurrence_qbx.py

@@ -39,7 +39,12 @@
 import numpy as np
 import sympy as sp
 
-from sumpy.recurrence import _make_sympy_vec, get_reindexed_and_center_origin_on_axis_recurrence
+from sumpy.recurrence import (
+    _make_sympy_vec,
+    get_off_axis_expression,
+    get_reindexed_and_center_origin_off_axis_recurrence,
+    get_reindexed_and_center_origin_on_axis_recurrence,
+)
 
 
 # ================ Transform/Rotate =================
@@ -114,11 +119,10 @@ def generate_lamb_expr(i, n_initial):
         for j in range(ndim):
             arg_list.append(var[j])
 
-        lamb_expr_symb_deriv = sp.diff(g_x_y, var_t[0], i)
-        for j in range(ndim):
-            lamb_expr_symb_deriv = lamb_expr_symb_deriv.subs(var_t[j], 0)
-        
         if i < n_initial:
+            lamb_expr_symb_deriv = sp.diff(g_x_y, var_t[0], i)
+            for j in range(ndim):
+                lamb_expr_symb_deriv = lamb_expr_symb_deriv.subs(var_t[j], 0)
             lamb_expr_symb = lamb_expr_symb_deriv
         else:
             lamb_expr_symb = recurrence.subs(n, i)
@@ -148,7 +152,61 @@ def generate_lamb_expr(i, n_initial):
 
 
     ### NEW CODE - COMPUTE OFF AXIS INTERACTIONS
+    start_order, t_recur_order, t_recur = get_reindexed_and_center_origin_off_axis_recurrence(pde)
+    t_exp, t_exp_order = get_off_axis_expression(pde)
+    storage_taylor_order = max(t_recur_order, t_exp_order+1)
+
+    storage = [np.zeros((n_p, n_p))] * storage_taylor_order
+
+    def gen_lamb_expr_t_recur(i, start_order):
+        arg_list = []
+        for j in range(t_recur_order, 0, -1):
+            # pylint: disable-next=not-callable
+            arg_list.append(s(i-j))
+        for j in range(ndim):
+            arg_list.append(var[j])
+
+        if i < start_order:
+            lamb_expr_symb_deriv = sp.diff(g_x_y, var_t[0], i)
+            for j in range(ndim):
+                lamb_expr_symb_deriv = lamb_expr_symb_deriv.subs(var_t[j], 0)
+            lamb_expr_symb = lamb_expr_symb_deriv
+        else:
+            lamb_expr_symb = t_recur.subs(n, i)
+
+        return lamb_expr_symb
+
+
+    def gen_lamb_expr_t_exp(i, t_exp_order):
+        arg_list = []
+        for j in range(t_exp_order, -1, -1):
+            # pylint: disable-next=not-callable
+            arg_list.append(s(i-j))
+        for j in range(ndim):
+            arg_list.append(var[j])
+
+        if i < t_exp_order:
+            lamb_expr_symb_deriv = sp.diff(g_x_y, var_t[0], i)
+            for j in range(ndim):
+                lamb_expr_symb_deriv = lamb_expr_symb_deriv.subs(var_t[j], 0)
+            lamb_expr_symb = lamb_expr_symb_deriv.subs(var[0], 0)
+        else:
+            lamb_expr_symb = t_exp.subs(n, i)
+
+        return lamb_expr_symb
+
+    interactions_off_axis = 0
+    for i in range(p+1):
+        lamb_expr_t_recur = gen_lamb_expr_t_recur(i, start_order)
+        a1 = [*storage[:-t_recur_order], *coord]
+
+        storage.pop(0)
+        storage.append(lamb_expr_t_recur(a1))
+
+        lamb_expr_t_exp = gen_lamb_expr_t_exp(i, t_exp_order)
+        a2 = [*storage[:-(t_exp_order+1)], *coord]
 
+        interactions_off_axis += lamb_expr_t_exp(a2) * radius**i/math.factorial(i)
 
     ################
 
@@ -159,6 +217,7 @@ def generate_lamb_expr(i, n_initial):
 
     interactions_total = np.zeros(coord[0].shape)
     interactions_total[mask_on_axis] = interactions_on_axis[mask_on_axis]
+    interactions_total[mask_off_axis] = interactions_off_axis[mask_off_axis]
 
     exp_res = (interactions_total * strengths[None, :]).sum(axis=1)
 

test/test_recurrence_qbx.py

@@ -303,7 +303,4 @@ def _construct_laplace_axis_2d(orders, resolutions):
 plt.ylabel("Error")
 plt.title("2D Ellipse LP Eval Error")
 plt.legend()
-plt.show()
-
-
-
+plt.show()