code: from math import factorial
  import scipy.special
  import numpy as np
  
  def sloppy_spherical(y):
      r = np.linalg.norm(y)
      costheta = y[2] / r
      theta = np.arccos(costheta)
      phi = np.arccos(y[0] / r / np.sin(theta))
      return r, theta, phi
  
  def Rdirect(n_max, y):
      r, theta, phi = sloppy_spherical(y)
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
      Pmn = scipy.special.lpmn(n_max, n_max, np.cos(theta))[0]
      for i in range(n_max + 1):
          for j in range(-i, i + 1):
              if j < 0:
                  lp = (
                      ((-1) ** (-j)) * (factorial(i + j) / factorial(i - j))
                      * Pmn[-j, i] / ((-1) ** -j)
                  )
              else:
                  lp = Pmn[j, i] / ((-1) ** j)
              factor = (r ** i) * lp / factorial(i + j)
              real[i, n_max + j] = factor * np.cos(j * phi)
              imag[i, n_max + j] = factor * np.sin(j * phi)
      return real, imag
  
  def Sdirect(n_max, y):
      r, theta, phi = sloppy_spherical(y)
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
      Pmn = scipy.special.lpmn(n_max, n_max, np.cos(theta))[0]
      for i in range(n_max + 1):
          for j in range(-i, i + 1):
              if j < 0:
                  lp = (
                      ((-1) ** (-j)) * (factorial(i + j) / factorial(i - j))
                      * Pmn[-j, i] / ((-1) ** -j)
                  )
              else:
                  lp = Pmn[j, i] / ((-1) ** j)
              factor = factorial(i - j) * lp / (r ** (i + 1))
              real[i, n_max + j] = factor * np.cos(j * phi)
              imag[i, n_max + j] = factor * np.sin(j * phi)
      return real, imag
  
  def R(n_max, y):
      y1, y2, y3 = y
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
      real[0, n_max] = 1.0
      for i in range(0, n_max):
          real[i + 1, n_max + i + 1] = (
              (y1 * real[i, n_max + i] - y2 * imag[i, n_max + i])
              / (2 * (i + 1))
          )
          imag[i + 1, n_max + i + 1] = (
              (y1 * imag[i, n_max + i] + y2 * real[i, n_max + i])
              / (2 * (i + 1))
          )
  
      t2f = np.linalg.norm(y) ** 2
      for j in range(n_max + 1):
          for i in range(j, n_max):
              factor = 1.0 / ((i + 1) ** 2 - j ** 2)
              t1f = (2 * i + 1) * y3
              real[i + 1, n_max + j] = factor * (t1f * real[i, n_max + j] - t2f * real[i - 1, n_max + j])
              imag[i + 1, n_max + j] = factor * (t1f * imag[i, n_max + j] - t2f * imag[i - 1, n_max + j])
      for i in range(n_max + 1):
          for j in range(1, n_max + 1):
              real[i, n_max - j] = ((-1) ** j) * real[i, n_max + j]
              imag[i, n_max - j] = ((-1) ** (j + 1)) * imag[i, n_max + j]
      return real, imag
  
  def R_storagefree(n_max, y):
      def neg(real, imag, mi):
          return (
              ((-1) ** mi) * real,
              ((-1) ** (mi + 1)) * imag
          )
  
      y1, y2, y3 = y
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
  
      t2f = np.linalg.norm(y) ** 2
      Rsr = 1.0
      Rsi = 0.0
      for mi in range(0, n_max + 1):
          real[mi, n_max + mi] = Rsr
          imag[mi, n_max + mi] = Rsi
          real[mi, n_max - mi], imag[mi, n_max - mi] = neg(Rsr, Rsi, mi)
  
          Rm2r = 0.0
          Rm2i = 0.0
          Rm1r = Rsr
          Rm1i = Rsi
          for ni in range(mi, n_max):
              factor = 1.0 / ((ni + 1) ** 2 - mi ** 2)
              t1f = (2 * ni + 1) * y3
              Rvr = factor * (t1f * Rm1r - t2f * Rm2r)
              Rvi = factor * (t1f * Rm1i - t2f * Rm2i)
              real[ni + 1, n_max + mi] = Rvr
              imag[ni + 1, n_max + mi] = Rvi
              real[ni + 1, n_max - mi], imag[ni + 1, n_max - mi] = neg(Rvr, Rvi, mi)
              Rm2r = Rm1r
              Rm2i = Rm1i
              Rm1r = Rvr
              Rm1i = Rvi
          Rsrold = Rsr
          Rsiold = Rsi
          Rsr = (y1 * Rsrold - y2 * Rsiold) / (2 * (mi + 1))
          Rsi = (y1 * Rsiold + y2 * Rsrold) / (2 * (mi + 1))
      return real, imag
  
  def Rderivs(n_max, y, d):
      Rvr, Rvi = R(n_max + 1, y)
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
  
      if d == 0:
          for i in range(n_max):
              for j in range(-i, i + 1):
                  real[i, n_max + j] = 0.5 * (
                      Rvr[i - 1, (n_max + 1) + j - 1]
                      - Rvr[i - 1, (n_max + 1) + j + 1]
                  )
                  imag[i, n_max + j] = 0.5 * (
                      Rvi[i - 1, (n_max + 1) + j - 1]
                      - Rvi[i - 1, (n_max + 1) + j + 1]
                  )
      elif d == 1:
          for i in range(n_max + 1):
              for j in range(-i, i + 1):
                  real[i, n_max + j] = -0.5 * (
                      Rvi[i - 1, (n_max + 1) + j - 1]
                      + Rvi[i - 1, (n_max + 1) + j + 1]
                  )
                  imag[i, n_max + j] = 0.5 * (
                      Rvr[i - 1, (n_max + 1) + j - 1]
                      + Rvr[i - 1, (n_max + 1) + j + 1]
                  )
      else:
          for i in range(n_max + 1):
              for j in range(-i, i + 1):
                  real[i, n_max + j] = Rvr[i - 1, (n_max + 1) + j]
                  imag[i, n_max + j] = Rvi[i - 1, (n_max + 1) + j]
      return real, imag
  
  def S(n_max, y):
      y1, y2, y3 = y
      ynorm = np.linalg.norm(y)
      ynorm2 = ynorm ** 2
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
      real[0, n_max] = 1.0 / ynorm
      for i in range(0, n_max):
          factor = (2 * i + 1) / ynorm2
          real[i + 1, n_max + i + 1] = factor * (
              (y1 * real[i, n_max + i] - y2 * imag[i, n_max + i])
          )
          imag[i + 1, n_max + i + 1] = factor * (
              (y1 * imag[i, n_max + i] + y2 * real[i, n_max + i])
          )
  
      for j in range(n_max + 1):
          for i in range(j, n_max):
              factor = 1.0 / ynorm2
              t1f = (2 * i + 1) * y3
              t2f = i ** 2 - j ** 2
              real[i + 1, n_max + j] = factor * (
                  t1f * real[i, n_max + j] - t2f * real[i - 1, n_max + j]
              )
              imag[i + 1, n_max + j] = factor * (
                  t1f * imag[i, n_max + j] - t2f * imag[i - 1, n_max + j]
              )
      for i in range(n_max + 1):
          for j in range(1, n_max + 1):
              real[i, n_max - j] = ((-1) ** j) * real[i, n_max + j]
              imag[i, n_max - j] = ((-1) ** (j + 1)) * imag[i, n_max + j]
      return real, imag
  
  def S_storagefree(n_max, y):
      def neg(real, imag, mi):
          return (
              ((-1) ** mi) * real,
              ((-1) ** (mi + 1)) * imag
          )
  
      y1, y2, y3 = y
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
  
      ynorm = np.linalg.norm(y)
      ynorm2 = ynorm ** 2
      Ssr = 1.0 / ynorm
      Ssi = 0.0
      for mi in range(0, n_max + 1):
          real[mi, n_max + mi] = Ssr
          imag[mi, n_max + mi] = Ssi
          real[mi, n_max - mi], imag[mi, n_max - mi] = neg(Ssr, Ssi, mi)
  
          Sm2r = 0.0
          Sm2i = 0.0
          Sm1r = Ssr
          Sm1i = Ssi
          for ni in range(mi, n_max):
              factor = 1.0 / ynorm2
              t1f = (2 * ni + 1) * y3
              t2f = ni ** 2 - mi ** 2
              Svr = factor * (t1f * Sm1r - t2f * Sm2r)
              Svi = factor * (t1f * Sm1i - t2f * Sm2i)
              real[ni + 1, n_max + mi] = Svr
              imag[ni + 1, n_max + mi] = Svi
              real[ni + 1, n_max - mi], imag[ni + 1, n_max - mi] = neg(Svr, Svi, mi)
              Sm2r = Sm1r
              Sm2i = Sm1i
              Sm1r = Svr
              Sm1i = Svi
          Ssrold = Ssr
          Ssiold = Ssi
          factor = (2 * mi + 1) / ynorm2
          Ssr = factor * (y1 * Ssrold - y2 * Ssiold)
          Ssi = factor * (y1 * Ssiold + y2 * Ssrold)
      return real, imag
  
  def Sderivs(n_max, y, d):
      Svr, Svi = S(n_max + 1, y)
      real = np.zeros((n_max + 1, 2 * n_max + 1))
      imag = np.zeros((n_max + 1, 2 * n_max + 1))
  
      if d == 0:
          for i in range(n_max + 1):
              for j in range(-i, i + 1):
                  real[i, n_max + j] = 0.5 * (
                      Svr[i + 1, (n_max + 1) + j - 1]
                      - Svr[i + 1, (n_max + 1) + j + 1]
                  )
                  imag[i, n_max + j] = 0.5 * (
                      Svi[i + 1, (n_max + 1) + j - 1]
                      - Svi[i + 1, (n_max + 1) + j + 1]
                  )
      elif d == 1:
          for i in range(n_max + 1):
              for j in range(-i, i + 1):
                  real[i, n_max + j] = -0.5 * (
                      Svi[i + 1, (n_max + 1) + j - 1]
                      + Svi[i + 1, (n_max + 1) + j + 1]
                  )
                  imag[i, n_max + j] = 0.5 * (
                      Svr[i + 1, (n_max + 1) + j - 1]
                      + Svr[i + 1, (n_max + 1) + j + 1]
                  )
      else:
          for i in range(n_max + 1):
              for j in range(-i, i + 1):
                  real[i, n_max + j] = -Svr[i + 1, (n_max + 1) + j]
                  imag[i, n_max + j] = -Svi[i + 1, (n_max + 1) + j]
      return real, imag
  
   import time
  import numpy as np
  
  from tectosaur.farfield import farfield_pts_direct, get_gpu_module
  from tectosaur.ops.sparse_farfield_op import TriToTriDirectFarfieldOp
  from tectosaur.util.geometry import normalize
  from tectosaur.mesh.mesh_gen import make_rect
  from tectosaur.mesh.modify import concat
  
  def make_meshes(n_m = 8, sep = 2, w = 1, n_m2 = None):
      if n_m2 is None:
          n_m2 = n_m
  
      m1 = make_rect(n_m, n_m, [
          [-w, 0, w], [-w, 0, -w],
          [w, 0, -w], [w, 0, w]
      ])
      m2 = make_rect(n_m2, n_m2, [
          [-w, sep, w], [-w, sep, -w],
          [w, sep, -w], [w, sep, w]
      ])
      m = concat(m1, m2)
      surf1_idxs = np.arange(m1[1].shape[0])
      surf2_idxs = (surf1_idxs[-1] + 1) + surf1_idxs
      return m, surf1_idxs, surf2_idxs
  
  def test_tri_tri_farfield():
      m, surf1_idxs, surf2_idxs = make_meshes()
      T1, T2 = [
          C(
              2, 'elasticT3', [1.0,0.25], m[0], m[1],
              np.float32, obs_subset = surf1_idxs,
              src_subset = surf2_idxs
          ) for C in [PtToPtDirectFarfieldOp, TriToTriDirectFarfieldOp]
      ]
      in_vals = np.random.rand(T1.shape[1])
      out1 = T1.dot(in_vals)
      out2 = T2.dot(in_vals)
      np.testing.assert_almost_equal(out1, out2)
  
  def timing(n, runtime, name, flops):
      print("for " + name)
      cycles = runtime * 5e12
      entries = ((n * 3) ** 2)
      cycles_per_entry = cycles / entries
      print("total time: " + str(runtime))
      print("pts: " + str(n))
      print("interacts: " + str(entries))
      print("cycles/interact: " + str(cycles_per_entry))
      print("total flop count: " + str(flops * n ** 2))
      print("Tflop/s: " + str(flops * n ** 2 / runtime / 1e12))
  
  def run_kernel(n, k_name, flops, testit = False, timeit = False):
      block_size = 128
      np.random.seed(100)
      obs_pts = np.random.rand(n, 3)
      obs_ns = normalize(np.random.rand(n, 3))
      src_pts = np.random.rand(n, 3)
      src_ns = obs_ns
      weights = np.random.rand(n, 3).flatten()
  
      start = time.time()
      params = [1.0, 0.25]
      result = farfield_pts_direct(
          k_name, obs_pts, obs_ns, src_pts, src_ns, weights, params, np.float32
      ).reshape((n, 3))
      runtime = time.time() - start
      if timeit:
          timing(n, runtime, k_name, flops)
  
      # if testit:
      #     correct = fmm.direct_eval(
      #         "elastic" + k_name, obs_pts, obs_ns, src_pts, src_ns, [1.0, 0.25]
      #     )
      #     correct = correct.reshape((n * 3, n * 3))
      #     correct = correct.dot(weights.reshape(n * 3)).reshape((n, 3))
      #     np.testing.assert_almost_equal(
      #         np.abs((result - correct) / correct),
      #         np.zeros_like(result), 2
      #     )
  
  def test_U():
      run_kernel(1000, 'elasticU3', 28, testit = True)
  
  def test_T():
      run_kernel(1000, 'elasticT3', 63, testit = True)
  
  def test_A():
      run_kernel(1000, 'elasticA3', 63, testit = True)
  
  def test_H():
      run_kernel(1000, 'elasticH3', 102, testit = True)
  
  if __name__ == '__main__':
      n = 32 * 512
      run_kernel(n, 'elasticU3', 28, timeit = True)
      run_kernel(n, 'elasticA3', 63, timeit = True)
      run_kernel(n, 'elasticT3', 63, timeit = True)
      run_kernel(n, 'elasticH3', 102, timeit = True)
      run_kernel(n, 'laplaceS3', 3, timeit = True)
  
apis: numpy.zeros numpy.arccos numpy.cos numpy.sin numpy.linalg.norm
apis: numpy.arange numpy.random.seed numpy.random.rand numpy.testing.assert_almost_equal