Renamed functions in recurrence.py, added off-axis section

Author: hirish99

sumpy/recurrence.py

@@ -379,34 +379,6 @@ def _get_initial_order_off_axis(recurrence):
         r_c = recurrence.subs(n, i)
     return i
 
-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.
-    Outputs an expression that evaluates to 0 rather than s(n) in terms
-    of s(n-1), etc. This is different from move_center_origin_source_arbitrary,
-    because we are "shifting" an EXPRESSION, not s(n) in terms of s(n-1), etc.
-
-    :arg recurrence: a recurrence relation in :math:`s(n)`
-    """
-    r = recurrence_from_pde(pde)
-
-    idx_l, terms = _extract_idx_terms_from_recurrence(r)
-
-    # How much do we need to shift the recurrence relation
-    shift_idx = max(idx_l)
-
-    n = sp.symbols("n")
-    r = r.subs(n, n-shift_idx)
-
-    idx_l, terms = _extract_idx_terms_from_recurrence(r)
-
-    r_ret = r
-    for i in range(len(idx_l)):
-        r_ret = r_ret.subs(terms[i], (-1)**(n+idx_l[i])*terms[i])
-
-    return r_ret, (max(idx_l)+1-min(idx_l))
-
 
 def move_center_origin_source_arbitrary(r: sp.Expr) -> sp.Expr:
     r"""
@@ -415,8 +387,8 @@ def move_center_origin_source_arbitrary(r: sp.Expr) -> sp.Expr:
     Assuming the recurrence is formulated so that evaluating it gives
     s(n) in terms of s(n-1), .., etc. We do NOT want a recurrence
     EXPRESSION as input, i.e. an expression containing s(n), s(n-1),
-    ..., that evaluates to 0. 
-    Use move_center_origin_source_arbitrary_expression for this.
+    ..., that evaluates to 0.
+    Use move_center_origin_source_arbitrary_expression for EXPRESSIONS.
 
     :arg recurrence: a recurrence relation in :math:`s(n)`
     """
@@ -431,7 +403,7 @@ def move_center_origin_source_arbitrary(r: sp.Expr) -> sp.Expr:
     return r_ret*((-1)**(n+1))
 
 
-def get_reindexed_and_center_origin_recurrence(pde) -> tuple[int, int,
+def get_reindexed_and_center_origin_on_axis_recurrence(pde) -> tuple[int, int,
                                                        sp.Expr]:
     r"""
     A function that "shifts" the recurrence so the expansion center is placed
@@ -453,24 +425,70 @@ def get_reindexed_and_center_origin_recurrence(pde) -> tuple[int, int,
     return n_initial, order, r_s
 
 # ================ OFF-AXIS RECURRENCE =================
-def get_off_axis_recurrence(pde: LinearPDESystemOperator) -> sp.Expr:
+def move_center_origin_source_arbitrary_expression(pde: LinearPDESystemOperator) -> sp.Expr:
     r"""
-    A function that takes in as input a pde and outputs a SHIFTED
-    + REINDEXED off-axis recurrence
+    A function that "shifts" the recurrence so it's center is placed
+    at the origin and source is the input for the recurrence generated.
+    Outputs an expression that evaluates to 0 rather than s(n) in terms
+    of s(n-1), etc. This is different from move_center_origin_source_arbitrary,
+    because we are "shifting" an EXPRESSION, not s(n) in terms of s(n-1), etc.
+
+    :arg recurrence: a recurrence relation in :math:`s(n)`
+    """
+    r = recurrence_from_pde(pde)
+
+    idx_l, terms = _extract_idx_terms_from_recurrence(r)
+    n = sp.symbols("n")
+
+    r_ret = r
+    for i in range(len(idx_l)):
+        r_ret = r_ret.subs(terms[i], (-1)**(n+idx_l[i])*terms[i])
+
+    return r_ret
+
+def get_reindexed_and_center_origin_off_axis_recurrence(pde: LinearPDESystemOperator) -> 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)[0].subs(var[0], 0)
+    r_exp = move_center_origin_source_arbitrary_expression(pde).subs(var[0], 0)
     recur_order, recur = reindex_recurrence_relation(r_exp)
-    print(recur)
     start_order = _get_initial_order_off_axis(recur)
     return start_order, recur_order, recur
 
-def get_taylor_expression(pde, deriv_order, taylor_order=4):
-    #recursively substitute, and output the "order" of the taylor expression
-    t_recurrence = get_off_axis_recurrence(pde)[2]
-    var = _make_sympy_vec("x", 2)
+
+def get_off_axis_expression(pde, taylor_order=4):
+    r"""
+    A function that takes in as input a pde, and outputs
+    the Taylor expression that gives the deriv_order th derivative
+    to a taylor_order order Taylor series with respect to x_1 and
+    s(i) where s(i) comes from the off-axis recurrence. See
+    get_reindexed_and_center_origin_off_axis_recurrence.
+    """
+    s = sp.Function("s")
     n = sp.symbols("n")
+    deriv_order = n
+
+    t_recurrence = get_reindexed_and_center_origin_off_axis_recurrence(pde)[2]
+    var = _make_sympy_vec("x", 2)
     exp = 0
     for i in range(taylor_order):
         exp += t_recurrence.subs(n, deriv_order+i)/math.factorial(i) * var[0]**i
-    return exp
+
+    #While derivatives w/order larger than the deriv_order exist in our taylor expression
+    #replace them with smaller order derivatives
+    idx_l, _ = _extract_idx_terms_from_recurrence(exp)
+    max_idx = max(idx_l)
+
+    while max_idx > 0:
+        for ind in idx_l:
+            if ind > 0:
+                exp = exp.subs(s(n+ind), t_recurrence.subs(n, n+ind))
+        
+        idx_l, _ = _extract_idx_terms_from_recurrence(exp)
+        max_idx = max(idx_l)
+    exp_range = max(idx_l) - min(idx_l)
+
+    return exp, exp_range

sumpy/recurrence_qbx.py

@@ -39,7 +39,7 @@
 import numpy as np
 import sympy as sp
 
-from sumpy.recurrence import _make_sympy_vec, get_reindexed_and_center_origin_recurrence, get_off_axis_recurrence, eval_taylor_recurrence_laplace_processed
+from sumpy.recurrence import _make_sympy_vec, get_reindexed_and_center_origin_on_axis_recurrence, get_reindexed_and_center_origin_off_axis_recurrence, eval_taylor_recurrence_laplace_processed
 
 
 # ================ Transform/Rotate =================
@@ -97,8 +97,8 @@ def recurrence_qbx_lp(sources, centers, normals, strengths, radius, pde, g_x_y,
     var_t = _make_sympy_vec("t", ndim)
 
     # ------------ 5. Compute recurrence
-    n_initial, order, recurrence = get_reindexed_and_center_origin_recurrence(pde)
-    t_order, t_recurrence = get_off_axis_recurrence(pde)
+    n_initial, order, recurrence = get_reindexed_and_center_origin_on_axis_recurrence(pde)
+    t_order, t_recurrence = get_reindexed_and_center_origin_off_axis_recurrence(pde)
     t_order += 2
 
     # ------------ 6. Set order p = 5

test/gridfree_taylor_recurrence.ipynb

@@ -20,7 +20,7 @@
     "    make_identity_diff_op,\n",
     ")\n",
     "\n",
-    "from sumpy.recurrence import get_taylor_expression, get_off_axis_recurrence\n",
+    "from sumpy.recurrence import get_off_axis_expression, get_reindexed_and_center_origin_off_axis_recurrence, _extract_idx_terms_from_recurrence\n",
     "\n",
     "import sympy as sp\n",
     "from sympy import hankel1\n",
@@ -53,21 +53,15 @@
    "source": [
     "from sumpy.expansion.diff_op import laplacian,make_identity_diff_op\n",
     "w = make_identity_diff_op(2)\n",
-    "laplace2d = laplacian(w)"
+    "laplace2d = laplacian(w)\n",
+    "helmholtz2d = laplacian(w) + w"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 4,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(-5*n + (n + 1)**2 + 1)*s(n - 2)/x1**2\n"
-     ]
-    },
     {
      "data": {
       "text/latex": [
@@ -83,98 +77,52 @@
     }
    ],
    "source": [
-    "get_off_axis_recurrence(laplace2d)[2].subs(n, 3)"
+    "get_reindexed_and_center_origin_off_axis_recurrence(laplace2d)[2].subs(n, 3)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(-6*n + (n + 1)**2 + 2)*s(n - 4)/x1**2 + (-5*n + (n + 1)**2 + 1)*s(n - 2)/x1**2\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "(3,\n",
-       " 4,\n",
-       " (-6*n + (n + 1)**2 + 2)*s(n - 4)/x1**2 + (-5*n + (n + 1)**2 + 1)*s(n - 2)/x1**2)"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "w = make_identity_diff_op(2)\n",
-    "helmholtz2d = laplacian(w) + w\n",
-    "get_off_axis_recurrence(helmholtz2d)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 9,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(-6*n + (n + 1)**2 + 2)*s(n - 4)/x1**2 + (-5*n + (n + 1)**2 + 1)*s(n - 2)/x1**2\n"
-     ]
-    },
     {
      "data": {
-      "text/latex": [
-       "$\\displaystyle \\frac{3 s{\\left(0 \\right)}}{x_{1}^{2}} + \\frac{6 s{\\left(2 \\right)}}{x_{1}^{2}}$"
-      ],
       "text/plain": [
-       "3*s(0)/x1**2 + 6*s(2)/x1**2"
+       "(x0**3*(((-6*n + (n + 2)**2 - 4)*s(n - 3)/x1**2 + (-5*n + (n + 2)**2 - 4)*s(n - 1)/x1**2)*(-5*n + (n + 4)**2 - 14)/(6*x1**2) + (-6*n + (n + 4)**2 - 16)*s(n - 1)/(6*x1**2)) + x0**2*((-6*n + (n + 3)**2 - 10)*s(n - 2)/(2*x1**2) + (-5*n + (n + 3)**2 - 9)*s(n)/(2*x1**2)) + x0*((-6*n + (n + 2)**2 - 4)*s(n - 3)/x1**2 + (-5*n + (n + 2)**2 - 4)*s(n - 1)/x1**2) + (-6*n + (n + 1)**2 + 2)*s(n - 4)/x1**2 + (-5*n + (n + 1)**2 + 1)*s(n - 2)/x1**2,\n",
+       " 4)"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 9,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "get_off_axis_recurrence(helmholtz2d)[2].subs(n, 4)"
+    "exp = get_off_axis_expression(helmholtz2d)\n",
+    "exp"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 11,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(-5*n + (n + 1)**2 + 1)*s(n - 2)/x1**2\n"
-     ]
-    },
     {
      "data": {
       "text/latex": [
-       "$\\displaystyle \\frac{x_{0}^{3} s{\\left(1 \\right)}}{3 x_{1}^{2}} + \\frac{2 s{\\left(-2 \\right)}}{x_{1}^{2}}$"
+       "$\\displaystyle \\frac{x_{0}^{3} \\left(- x_{1}^{2} \\left(6 n - \\left(n + 4\\right)^{2} + 16\\right) s{\\left(n - 1 \\right)} + \\left(\\left(5 n - \\left(n + 2\\right)^{2} + 4\\right) s{\\left(n - 1 \\right)} + \\left(6 n - \\left(n + 2\\right)^{2} + 4\\right) s{\\left(n - 3 \\right)}\\right) \\left(5 n - \\left(n + 4\\right)^{2} + 14\\right)\\right) + 3 x_{1}^{2} \\left(- x_{0}^{2} \\left(\\left(5 n - \\left(n + 3\\right)^{2} + 9\\right) s{\\left(n \\right)} + \\left(6 n - \\left(n + 3\\right)^{2} + 10\\right) s{\\left(n - 2 \\right)}\\right) - 2 x_{0} \\left(\\left(5 n - \\left(n + 2\\right)^{2} + 4\\right) s{\\left(n - 1 \\right)} + \\left(6 n - \\left(n + 2\\right)^{2} + 4\\right) s{\\left(n - 3 \\right)}\\right) + 2 \\left(- 6 n + \\left(n + 1\\right)^{2} + 2\\right) s{\\left(n - 4 \\right)} + 2 \\left(- 5 n + \\left(n + 1\\right)^{2} + 1\\right) s{\\left(n - 2 \\right)}\\right)}{6 x_{1}^{4}}$"
       ],
       "text/plain": [
-       "x0**3*s(1)/(3*x1**2) + 2*s(-2)/x1**2"
+       "(x0**3*(-x1**2*(6*n - (n + 4)**2 + 16)*s(n - 1) + ((5*n - (n + 2)**2 + 4)*s(n - 1) + (6*n - (n + 2)**2 + 4)*s(n - 3))*(5*n - (n + 4)**2 + 14)) + 3*x1**2*(-x0**2*((5*n - (n + 3)**2 + 9)*s(n) + (6*n - (n + 3)**2 + 10)*s(n - 2)) - 2*x0*((5*n - (n + 2)**2 + 4)*s(n - 1) + (6*n - (n + 2)**2 + 4)*s(n - 3)) + 2*(-6*n + (n + 1)**2 + 2)*s(n - 4) + 2*(-5*n + (n + 1)**2 + 1)*s(n - 2)))/(6*x1**4)"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "get_taylor_expression(laplace2d, 0)"
+    "exp[0].simplify()"
    ]
   },
   {