jaxFMM 0.0.1__tar.gz

jaxfmm-0.0.1/.gitignore

@@ -0,0 +1,22 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+*.cpp
+
+# build
+.pytest_cache/
+build/
+*.egg-info/
+
+# vim
+*.swp
+
+# docs build files
+docs/_build
+
+# local dev folder
+dev

jaxfmm-0.0.1/.gitlab-ci.yml

@@ -0,0 +1,61 @@
+default:
+  interruptible: true
+
+.test_template:
+  stage: test
+  variables:
+    PYTEST_ADDOPTS: "--color=yes"
+  script: python -m pytest --verbose
+
+.cpu_template:
+  extends: .test_template
+  image: mambaorg/micromamba:jammy
+  before_script:
+    - micromamba install -y -n base -c conda-forge "python=$PYTHON_VERSION"
+    - python --version
+    - pip install .[dev]
+
+.gpu_template:
+  extends: .test_template
+  image: mambaorg/micromamba:jammy-cuda-12.3.1
+  before_script:
+    - micromamba install -y -n base -c conda-forge "python=$PYTHON_VERSION"
+    - python --version
+    - nvidia-smi
+    - pip install .[dev,cuda]
+  tags:
+    - saas-linux-medium-amd64-gpu-standard
+
+# test_py3_11_gpu:
+#   extends: .gpu_template
+#   variables:
+#     PYTHON_VERSION: "3.11"
+
+test_py3_11_cpu:
+  extends: .cpu_template
+  variables:
+    PYTHON_VERSION: "3.11"
+
+deploy_pypi:
+  stage: deploy
+  image: !reference [.cpu_template, image]
+  before_script:
+    - micromamba install -y -n base -c conda-forge 'python=3.11' make pandoc
+    - python -m pip install --upgrade twine
+    - python -m pip install --upgrade build
+  variables:
+    TWINE_USERNAME: __token__
+    TWINE_PASSWORD: ${PYPI_TOKEN}
+  script:
+    - python -m build
+    - python -m twine upload dist/*
+  rules:
+    - if: $CI_COMMIT_TAG
+  interruptible: false
+
+workflow:
+  rules:
+    - if: $CI_COMMIT_TAG
+    - if: $CI_COMMIT_REF_NAME == $CI_DEFAULT_BRANCH
+    - if: $CI_PIPELINE_SOURCE == "merge_request_event"
+    - if: $CI_PIPELINE_SOURCE == "web"

jaxfmm-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Robert Kraft
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

jaxfmm-0.0.1/PKG-INFO

@@ -0,0 +1,98 @@
+Metadata-Version: 2.4
+Name: jaxFMM
+Version: 0.0.1
+Summary: Adaptive Fast Multipole Method with Laplace kernel in JAX.
+Project-URL: Homepage, https://gitlab.com/jaxfmm/jaxfmm
+Author-email: Robert Kraft <robert.kraft@univie.ac.at>
+License: MIT License
+        
+        Copyright (c) 2025 Robert Kraft
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: FMM,N-body,jax,potential,treecode
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.9
+Requires-Dist: jax
+Provides-Extra: cuda
+Requires-Dist: jax[cuda]; extra == 'cuda'
+Provides-Extra: dev
+Requires-Dist: matplotlib; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# jaxFMM
+
+jaxFMM is an open source implementation of the Fast Multipole Method in JAX. The goal is to offer an easily readable/maintainable FMM implementation with good performance that runs on CPU/GPU and supports autodiff. This is enabled through JAX's just-in-time compiler.
+
+## Installation and Usage
+
+jaxFMM depends only on JAX and can be installed from pypi or by downloading the source as follows:
+
+    pip install jaxfmm
+
+If you want to run jaxFMM on GPUs, the easiest way is to use NVIDIA CUDA and cuDNN from pip wheels by instead typing:
+
+    pip install jaxfmm[cuda]
+
+Using a custom, self-installed CUDA with jax is [described in the JAX documentation](https://docs.jax.dev/en/latest/installation.html).
+
+The [unitcube demo](/demos/unitcube.py) is a short and simple example demonstrating how to use jaxFMM.
+
+## Features
+
+There are many flavors of FMM implementations. In short, jaxFMM currently:
+
+- only supports the Laplacian kernel.
+- only supports point charges.
+- only supports evaluation positions that are the same as the source positions.
+- uses real basis functions computed via recurrence relations.
+- uses "nested sum" O(p^4) M2M/M2L/L2L transformations.
+- uses a non-uniform 2^N-ary tree hierarchy (directly inspired by [this work of A. Goude and S. Engblom](https://link.springer.com/article/10.1007/s11227-012-0836-0)), allowing arbitrary shape of the boxes in the hierarchy and guaranteeing balanced trees but requiring storage of interaction lists.
+- has jit-compiled functions and autodiff for every substep of the algorithm except for the generation of interaction lists.
+
+In summary, jaxFMM in its current state can do adaptive point charge FMM for Laplace kernels with good performance for lower expansion orders (p <= 3) and reasonably homogenous distributions. Autodiff only works if the particle positions remain constant. A first benchmark of uniformly distributed charges in the unit cube, computed on Google Cloud [g2-standard-8](https://cloud.google.com/compute/docs/gpus#l4-gpus) (GPU timings) and [c3d-highmem-16](https://cloud.google.com/compute/docs/general-purpose-machines#c3d_series) (CPU timings) machines can be found below:
+
+<img src="docs/images/jax_unitcube_benchmark_p3.png" alt="unitcube benchmark" width="500"/>
+
+## TODOs
+
+jaxFMM is primarily developed for my PhD project, where I am working on a GPU parallel FMM stray field evaluation routine with autodiff for finite-element micromagnetics. Alongside the very early state that it is in, this explains the currently limited feature set and design decisions mentioned above.
+
+Contributions are always welcome however, and I plan to still improve jaxFMM. Topics that come to mind here are:
+
+- volume FMM.
+- other kernels or even a kernel-independent formulation.
+- independent evaluation and source positions.
+- distributed parallelism via jax.sharding.
+- faster M2M/M2L/L2L transformations.
+- faster jit-compilation, especially for large systems, higher orders and gradient computations.
+- jit-compilable (and differentiable) interaction list generation, if possible.
+- various other performance improvements.
+- some degree of autotuning for the parameters.

jaxfmm-0.0.1/README.md

@@ -0,0 +1,49 @@
+# jaxFMM
+
+jaxFMM is an open source implementation of the Fast Multipole Method in JAX. The goal is to offer an easily readable/maintainable FMM implementation with good performance that runs on CPU/GPU and supports autodiff. This is enabled through JAX's just-in-time compiler.
+
+## Installation and Usage
+
+jaxFMM depends only on JAX and can be installed from pypi or by downloading the source as follows:
+
+    pip install jaxfmm
+
+If you want to run jaxFMM on GPUs, the easiest way is to use NVIDIA CUDA and cuDNN from pip wheels by instead typing:
+
+    pip install jaxfmm[cuda]
+
+Using a custom, self-installed CUDA with jax is [described in the JAX documentation](https://docs.jax.dev/en/latest/installation.html).
+
+The [unitcube demo](/demos/unitcube.py) is a short and simple example demonstrating how to use jaxFMM.
+
+## Features
+
+There are many flavors of FMM implementations. In short, jaxFMM currently:
+
+- only supports the Laplacian kernel.
+- only supports point charges.
+- only supports evaluation positions that are the same as the source positions.
+- uses real basis functions computed via recurrence relations.
+- uses "nested sum" O(p^4) M2M/M2L/L2L transformations.
+- uses a non-uniform 2^N-ary tree hierarchy (directly inspired by [this work of A. Goude and S. Engblom](https://link.springer.com/article/10.1007/s11227-012-0836-0)), allowing arbitrary shape of the boxes in the hierarchy and guaranteeing balanced trees but requiring storage of interaction lists.
+- has jit-compiled functions and autodiff for every substep of the algorithm except for the generation of interaction lists.
+
+In summary, jaxFMM in its current state can do adaptive point charge FMM for Laplace kernels with good performance for lower expansion orders (p <= 3) and reasonably homogenous distributions. Autodiff only works if the particle positions remain constant. A first benchmark of uniformly distributed charges in the unit cube, computed on Google Cloud [g2-standard-8](https://cloud.google.com/compute/docs/gpus#l4-gpus) (GPU timings) and [c3d-highmem-16](https://cloud.google.com/compute/docs/general-purpose-machines#c3d_series) (CPU timings) machines can be found below:
+
+<img src="docs/images/jax_unitcube_benchmark_p3.png" alt="unitcube benchmark" width="500"/>
+
+## TODOs
+
+jaxFMM is primarily developed for my PhD project, where I am working on a GPU parallel FMM stray field evaluation routine with autodiff for finite-element micromagnetics. Alongside the very early state that it is in, this explains the currently limited feature set and design decisions mentioned above.
+
+Contributions are always welcome however, and I plan to still improve jaxFMM. Topics that come to mind here are:
+
+- volume FMM.
+- other kernels or even a kernel-independent formulation.
+- independent evaluation and source positions.
+- distributed parallelism via jax.sharding.
+- faster M2M/M2L/L2L transformations.
+- faster jit-compilation, especially for large systems, higher orders and gradient computations.
+- jit-compilable (and differentiable) interaction list generation, if possible.
+- various other performance improvements.
+- some degree of autotuning for the parameters.

jaxfmm-0.0.1/demos/inverse.py

@@ -0,0 +1,23 @@
+import jax.numpy as jnp
+from jax.scipy.optimize import minimize
+from jaxfmm import *
+
+print("Assembling grid of charges.")
+nside, sidelen = 15, 4*jnp.pi
+pts = (jnp.mgrid[:nside,:nside,:nside].T/(nside-1) * sidelen - sidelen/2).reshape((-1,3))
+chrgs = 0.01*jnp.ones(pts.shape[0])
+
+print("Generating hierarchy.")
+tree_info = gen_hierarchy(pts)
+
+print("Computing desired potential.")
+desired_pot = jnp.sin(jnp.linalg.norm(pts,axis=-1))
+norm = jnp.linalg.norm(desired_pot)
+
+def loss(chrgs):
+    return jnp.linalg.norm(desired_pot-eval_potential(*tree_info,chrgs))/norm
+
+print("Initial loss: %.2e"%loss(chrgs))
+print("Compiling + running minimizer.")
+res = minimize(loss,chrgs,method="BFGS",options={"maxiter": 1000})
+print("Final loss (%i iterations): %.2e"%(res.nit,res.fun))

jaxfmm-0.0.1/demos/unitcube.py

@@ -0,0 +1,20 @@
+import jax.numpy as jnp
+from jax import random
+from jaxfmm import *
+
+print("Generating sample points and charges.")
+N = 2**15
+key = random.key(124)
+pts = random.uniform(key,(N,3))
+chrgs = random.uniform(key,N)
+
+print("Generating FMM hierarchy.")
+tree_info = gen_hierarchy(pts)
+
+print("Compiling and computing FMM potential.")
+pot_FMM = eval_potential(*tree_info,chrgs)
+
+print("Compiling and computing analytic potential.")
+pot_dir = eval_potential_direct(pts,chrgs)
+
+print("FMM normwise relative error: %.2e"%(jnp.linalg.norm(pot_dir-pot_FMM)/jnp.linalg.norm(pot_dir)))

jaxfmm-0.0.1/jaxfmm/__init__.py

@@ -0,0 +1,5 @@
+from jaxfmm.fmm import *
+from jaxfmm.debug_helpers import *
+
+__all__ = (fmm.__all__ +
+           debug_helpers.__all__)

jaxfmm-0.0.1/jaxfmm/debug_helpers.py

@@ -0,0 +1,134 @@
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import jax.numpy as jnp
+import jax.scipy as jsp
+import jaxfmm.fmm as fmm
+from math import isqrt
+
+__all__ = ["print_stats", "plot_fmm_boxes"]
+
+def print_stats(pts, idcs, rev_idcs, boxcenters, mpl_cnct, dir_cnct, n_split):
+    r"""
+    Print some basic information about the hierarchy.
+    """
+    print("----------------------------------------FMM Hierarchy Stats----------------------------------------")
+    print("%i points, %i levels, %i children per box, %i charges per box"%(pts.shape[0],len(mpl_cnct)-1, 2**n_split, idcs.shape[1]))
+    for i in range(len(mpl_cnct)):
+        print("Total number of mpl interactions on level %i (with padding, fraction: %.2f): %i"%(i, mpl_cnct[i].size / (mpl_cnct[i]<mpl_cnct[i].shape[0]).sum(),mpl_cnct[i].size))
+    print("Total number of dir interactions on max lvl (with padding, fraction: %.2f): %i"%(dir_cnct.size / (dir_cnct < dir_cnct.shape[0]).sum(),dir_cnct.size))
+    print("Near field compression ratio (without padding):", (((dir_cnct<dir_cnct.shape[0]).sum()) * idcs.shape[1]**2) / (float(pts.shape[0])**2))
+    memory_usage = idcs.nbytes + rev_idcs.nbytes + dir_cnct.nbytes
+    for i in range(len(boxcenters)):
+        memory_usage += boxcenters[i].nbytes + mpl_cnct[i].nbytes
+    print("Total memory consumed by the hierarchy: %.2e Bytes"%memory_usage)
+    print("---------------------------------------------------------------------------------------------------")
+
+def plot_box(center, L, ax, facecolor='cyan', edgecolor='k', alpha=0.25):
+    r"""
+    Plot a single box in the given ax.
+    """
+    faces = []
+    for perm in range(3):   # for the 3 coordinate axes
+        faces.append([])
+        faces.append([])
+        for id in [(-1,-1), (1,-1), (1,1), (-1,1)]: # we generate 4 points on each side
+            shift = jnp.array([-1,id[0],id[1]])
+            shift2 = jnp.array([1,id[0],id[1]])
+
+            start = center + L/2 * shift[(jnp.arange(3,dtype=int)+perm)%3]
+            end = center + L/2 * shift2[(jnp.arange(3,dtype=int)+perm)%3]
+
+            faces[-2].append(start)
+            faces[-1].append(end)
+        faces[-2].append(faces[-2][0])
+        faces[-1].append(faces[-1][0])
+    faces = jnp.array(faces)
+    ax.add_collection3d(Poly3DCollection(faces, facecolors=facecolor, linewidths=0.5, edgecolors=edgecolor, alpha=alpha))
+
+def plot_fmm_boxes(pts, boxcenters, boxlens, level, show_wellsep=None, mpl_cnct=None, dir_cnct=None, n_split=3, elev=45, azim=45, roll=0, fname=None, plot_all=False):
+    r"""
+    Plot the hierarchy.
+    """
+    fig = plt.figure()
+    ax = fig.add_subplot(projection='3d')
+    ax.set_box_aspect((jnp.ptp(pts[:,0]), jnp.ptp(pts[:,1]), jnp.ptp(pts[:,2])))    # setting equal aspect ratio so we see the true shape of objects
+
+    nboxs = (2**n_split)**level
+    colors = plt.cm.jet(jnp.linspace(0,1,nboxs))
+    plot_all = plot_all or show_wellsep is None # if there is no show_wellsep set, we always plot everything
+    for i in range(nboxs):
+        if(show_wellsep is not None):
+            if(i==show_wellsep):
+                color = 'red'
+            elif(show_wellsep in mpl_cnct[level][i]):
+                color = 'green'
+            elif(level==(len(boxcenters)-1) and dir_cnct is not None and show_wellsep in dir_cnct[i]):
+                color = 'blue'
+            else:
+                color = 'gray'
+        else:
+            color = colors[i]
+        if(plot_all or color == 'blue' or color == 'red' or color=="green"):
+            plot_box(boxcenters[level][i],boxlens[level][i],ax,facecolor=color,alpha=0.5)
+    
+    plt.axis('off') 
+    ax.view_init(elev=elev, azim=azim, roll=roll)
+    plt.tight_layout()
+    plt.show()
+    if(fname is not None):
+        fig.savefig(fname,transparent=True,bbox_inches='tight',dpi=300)
+    return fig, ax
+
+def eval_multipole(coeff, boxcenter, eval_pts):
+    r"""
+    Evaluate a single multipole expansion.
+    """
+    p = get_deg(coeff.shape[-1])
+    sing = fmm.eval_singular_basis(eval_pts - boxcenter,p)
+    res = jnp.zeros(eval_pts.shape[0])
+    for n in range(p+1):
+        for m in range(-n,n+1):
+            if(m!=0):
+                res += (-1)**n * 2 * coeff[...,fmm.mpl_idx(m,n)] * sing[...,fmm.mpl_idx(-m,n)]
+            else:
+                res += (-1)**n * coeff[...,fmm.mpl_idx(m,n)] * sing[...,fmm.mpl_idx(-m,n)] 
+    res /= (4*jnp.pi)
+    return res
+
+def get_local_expansions(pts, chrgs, exp_centers, p):
+    r"""
+    Generate local expansions.
+    """
+    dist = pts[None,...] - exp_centers[:,None,:]
+    coeff = (fmm.eval_singular_basis(dist,p) * chrgs[None,:,None]).sum(axis=1)
+    return coeff
+
+def binom(x, y):
+  return jnp.exp(jsp.special.gammaln(x + 1) - jsp.special.gammaln(y + 1) - jsp.special.gammaln(x - y + 1))
+
+def gen_multipole_dist(m, n, eps = 0.5):
+    r"""
+    Generate a point charge distribution corresponding to a specific multipole moment (Majic, Matt. (2022). Point charge representations of multipoles. European Journal of Physics. 43. 10.1088/1361-6404/ac578b.)
+    """
+    if(m == 0):   # axial
+        k = jnp.arange(-n, n+1, 2)
+        chrgs = (-1)**((n-k)/2) * binom(n, (n-k)/2.0) / (jsp.special.factorial(n) * (2*eps)**n)
+        pts = jnp.zeros((k.shape[0],3))
+        pts = pts.at[:,2].set(k*eps)
+    else:         # (stacked) bracelet
+        rotate = m < 0
+        m = abs(m)      # we work with the real basis and rotate later
+        knum = n-m+1
+        jnum = 2*m
+        j = jnp.tile(jnp.arange(jnum),knum)
+        k = jnp.repeat(jnp.arange(-n+m,n-m+1,2),jnum)
+        phi = (j-0.5) * jnp.pi/m if rotate else j * jnp.pi/m
+        pts = jnp.array([eps*jnp.cos(phi), eps*jnp.sin(phi), k*eps]).T
+        chrgs = 4**(m-1) * jsp.special.factorial(m-1) / ((2*eps)**n * jsp.special.factorial(n-m)) * (-1)**((n-m-k)/2 + j) * binom(n-m,(n-m-k)/2)
+    return pts, chrgs
+
+def get_deg(N_coeff):
+    r"""
+    Get the degree p of a multipole expansion from the number of coefficients.
+    """
+    return isqrt(N_coeff) - 1

jaxfmm-0.0.1/jaxfmm/fmm.py

@@ -0,0 +1,542 @@
+import jax
+import jax.numpy as jnp
+from functools import partial
+from math import log2, ceil, sqrt
+
+__all__ = ["gen_hierarchy", "eval_potential", "eval_potential_direct"]
+
+def get_max_l(N_tot, N_max, n_split = 3): # need to get this outside of the jit compiled function - if the particle number changes too much, we must recompile...
+    r"""
+    Compute number of levels in the hierarchy.
+
+    :param N_tot: Total number of point charges.
+    :type N_tot: int
+    :param N_max: Maximum allowed number of point charges per box.
+    :type N_max: int
+    :param n_split: How many splits per level and box will be performed. Each box will have 2^n_split children.
+    :type n_split: int, optional
+
+    :return: The maximum level in the hierarchy.
+    :rtype: int
+    """
+    max_l = int(ceil(log2(N_tot/N_max)/n_split))
+    return 0 if max_l < 0 else max_l
+
+@partial(jax.jit, static_argnames = ["max_l", "n_split"])
+def balanced_tree(pts, max_l, n_split = 3):
+    r"""
+    Generate a balanced 2^n-tree hierarchy.
+
+    :param pts: Array of shape (N_tot,3) containing the positions of N_tot point charges.
+    :type pts: jnp.array
+    :param max_l: Maximum level, computed with get_max_l().
+    :type max_l: int
+    :param n_split: How many splits per level and box will be performed. Each box will have 2^n_split children.
+    :type n_split: int, optional
+
+    :return: Indices to resort pts into the highest level of the hierarchy (includes padding), indices to reverse sorting, centers of the boxes on all levels, sidelengths of the boxes on all levels.
+    :rtype: (jnp.array, jnp.array, list(jnp.array), list(jnp.array))
+    """
+    n_chi = 2**n_split
+    idcs = jnp.arange(pts.shape[0], dtype = jnp.int32)[None,:]
+
+    for l in range(max_l*n_split):    # carry out n_split splits on max_l levels in total (we cannot make this a for_i loop as the shape constantly changes)
+        splitpos = idcs.shape[1]//2   # split position in the middle
+        needpad = idcs.shape[1]%2     # modulo tells us if padding must be inserted
+
+        pts_sorted = pts.at[idcs].get(mode="fill",fill_value=-jnp.nan**2)  # padded values get converted to NaNs - NOTE: due to the argpartition implementation, we need to make sure that all the NaNs have the same sign
+        axis_to_split = jnp.argmax(jnp.nanmax(pts_sorted,axis=1) - jnp.nanmin(pts_sorted,axis=1),axis=1)    # nanmax and -min to ignore NaNs
+        idcs = idcs[jnp.arange(idcs.shape[0], dtype = jnp.int32)[:,None],jnp.argpartition(pts_sorted[jnp.arange(axis_to_split.shape[0], dtype = jnp.int32),:,axis_to_split],splitpos,axis=1)] # splitting - the NaNs introduced below get transported to the beginning
+
+        # padding so the next array has the correct shape - this does not break JIT compiling because it can be computed from only the input shape
+        idcs = jax.lax.pad(idcs,jnp.int32(pts.shape[0]),[(0,0,0),(0,needpad,0)])   # pad at the end with out of range values
+        idcs = idcs.reshape((-1,idcs.shape[1]//2))                      # now that we padded, we can safely reshape this
+    
+    idcs = jnp.sort(idcs,axis=1)    # sorting is good for locality but might be overkill TODO: swap only first and last positions instead of full sort?
+    rev_idcs = jnp.argsort(idcs.flatten())[:pts.shape[0]]     # reverse sorting indices, to undo the sorting
+    pts_sorted = pts.at[idcs].get(mode="fill",fill_value=jnp.nan)  
+    boxcenters, boxlens = [jnp.zeros((n_chi**l,3)) for l in range(max_l+1)], [jnp.zeros((n_chi**l,3)) for l in range(max_l+1)]
+
+    for l in range(max_l,0,-1):
+        minc, maxc = jnp.nanmin(pts_sorted,axis=1), jnp.nanmax(pts_sorted,axis=1)    # nanmax and -min to ignore NaNs
+        boxlens[l] = maxc - minc    # TODO: in principle we only need to save the norm of this, but it is nice to have for visualizations
+        boxcenters[l] = minc + boxlens[l]/2
+        pts_sorted = pts_sorted.reshape((-1,pts_sorted.shape[1]*n_chi,3))
+
+    minc, maxc = jnp.nanmin(pts_sorted,axis=1), jnp.nanmax(pts_sorted,axis=1)
+    boxlens[0] = maxc - minc             # box sidelengths
+    boxcenters[0] = minc + boxlens[0]/2  # box center
+    return idcs, rev_idcs, boxcenters, boxlens
+
+#@partial(jax.jit, static_argnames = ["theta", "n_split"])
+def gen_connectivity(boxcenters, boxlens, theta = 0.75, n_split = 3): # TODO: is there a way to make this JIT compilable?
+    r"""
+    Compute connectivity information for a given hierarchy. To determine well-separatedness, we check for
+
+    R + theta * r <= theta * d
+
+    where R = max(r1,r2), r = min(r1,r2) and d is the distance between the centers of two boxes with radii r1 and r2. This should give a FMM error that scales as
+
+    theta^(p+1)
+    
+    with the expansion order p.
+
+    :param boxcenters: List of boxcenters computed with balanced_tree().
+    :type boxcenters: list(jnp.array)
+    :param boxlens: List of box sidelengths computed with balanced_tree().
+    :type boxlens: list(jnp.array)
+    :param theta: Well-separatedness parameter, determines accuracy.
+    :type theta: float
+    :param n_split: How many splits per level and box have been performed. Each box has 2^n_split children.
+    :type n_split: int, optional
+
+    :return: List of (padded) M2L interaction partner index arrays for each box on each level, array of (padded) P2P interaction partner indices on the highest level.
+    :rtype: (list(jnp.array), jnp.array)
+    """
+    n_l = len(boxcenters)  # number of levels
+    n_chi = 2**n_split     # number of child boxes per split
+
+    mpl_cnct = []
+    non_wellseps = jnp.array([[0]], dtype = jnp.int32)     # initial value
+
+    for l in range(n_l):   # for every level
+        r1 = jnp.linalg.norm(boxlens[l],axis=1)/2   # L/2 is the radius of the box TODO: see comment on only storing the norm above
+        d = jnp.linalg.norm(boxcenters[l][:,None,:] - boxcenters[l].at[non_wellseps].get(mode="fill",fill_value=jnp.nan),axis=-1)  # find the distances between boxcenters
+        R = r1[non_wellseps]    # r2
+        r = R.copy()
+        tmp = (r1[:,None] > R)
+        R = jnp.where(tmp,R,r1[:,None])     # wherever r2 < r1, we replace r2 with r1 => R = max(r1,r2)
+        r = jnp.where(~tmp,r,r1[:,None])    # wherever r2 > r1, we replace r2 with r1 => r = min(r1,r2)
+
+        wellsep = (R + theta*r <= theta*d)      # NOTE: NaNs always return false here - this is how we rid ourselves of the padding
+        non_wellsep = (R + theta*r > theta*d)
+        wellsep_nums = wellsep.sum(axis=1, dtype = jnp.int32)          # number of well-separated boxes for each box
+        non_wellsep_nums = non_wellsep.sum(axis=1, dtype = jnp.int32)  # number of non-well-separated boxes for each box
+        
+        to_pad_wellsep = jnp.max(wellsep_nums) - wellsep_nums    # how much padding per box must be inserted
+        wellsep_padding = jnp.repeat(jnp.cumsum(wellsep_nums),to_pad_wellsep)   # the correct indices for the insert below NOTE: this has a dynamic size and therefore breaks JIT compilation
+        mpl_cnct.append(jnp.insert(non_wellseps[wellsep],wellsep_padding,n_chi**l).reshape(n_chi**l,-1))    # save result to the interaction lists
+
+        to_pad_non_wellsep = jnp.max(non_wellsep_nums) - non_wellsep_nums    # how much padding per box must be inserted for non-well-separated boxes
+        non_wellsep_padding = jnp.repeat(jnp.cumsum(non_wellsep_nums),to_pad_non_wellsep)   # the correct indices for the insert below
+        non_wellseps = jnp.insert(non_wellseps[non_wellsep],non_wellsep_padding,n_chi**l).reshape((n_chi**l,-1))    # we overwrite the old values here
+
+        if(l<n_l-1):    # prepare for the computation on the next level by transforming to child indices
+            non_wellseps = jnp.repeat(jnp.repeat(non_wellseps,n_chi,axis=1)*n_chi + jnp.tile(jnp.arange(n_chi, dtype = jnp.int32),non_wellseps.shape[1]),n_chi,axis=0)
+
+    return mpl_cnct, non_wellseps
+
+def mpl_idx(m,n):
+    r"""
+    Compute "flattened" array position of multipole coefficient C^m_n with order m and degree n.
+    """
+    return n**2 + (m+n)
+
+def inv_mpl_idx(idx):
+    r"""
+    Compute order m and degree n of "flattened" array position idx.
+    """
+    n = int(sqrt(idx))
+    m = idx - n*(n+1)
+    return m, n
+
+@partial(jax.jit, static_argnames=['p'])
+def eval_regular_basis(rvec, p):
+    r"""
+    Evaluate real regular basis functions (Laplace kernel) with a recursion relation [Gumerov, N. A. et al. Fast multipole methods on graphics processors. J. Comp. Phys., B 227, 8290 (2008)].
+
+    :param rvec: Array of positions in cartesian coordinates to evaluate the basis at.
+    :type rvec: jnp.array
+    :param p: Maximum degree (degree of multipole expansion) for the evaluation.
+    :type p: int
+
+    :return: Regular basis evaluated until the given degree at the given locations.
+    :rtype: jnp.array
+    """
+    x, y, z = rvec[..., 0], rvec[...,1], rvec[...,2]
+    coeff = jnp.zeros((*rvec.shape[:-1], (p+1)**2))
+    coeff = coeff.at[...,mpl_idx(0,0)].set(1)
+
+    if(p>0):
+        coeff = coeff.at[...,mpl_idx(1,1)].set(-0.5*x)
+        coeff = coeff.at[...,mpl_idx(-1,1)].set(0.5*y)
+
+    for n in range(2,p+1):    # first recursion: extreme values
+        coeff = coeff.at[...,mpl_idx(n,n)].set(-(x*coeff[...,mpl_idx(n-1,n-1)] + y*coeff[...,mpl_idx(1-n,n-1)])/(2*n))
+        coeff = coeff.at[...,mpl_idx(-n,n)].set((y*coeff[...,mpl_idx(n-1,n-1)] - x*coeff[...,mpl_idx(1-n,n-1)])/(2*n))
+
+    for n in range(0,p):      # second recursion: neighbors of extreme values
+        coeff = coeff.at[...,mpl_idx(n,n+1)].set(-z*coeff[...,mpl_idx(n,n)])
+        coeff = coeff.at[...,mpl_idx(-n,n+1)].set(-z*coeff[...,mpl_idx(-n,n)])
+
+    for n in range(2,p+1):    # third recursion: all values inbetween
+        for m in range(-n+2,n-1):
+            coeff = coeff.at[...,mpl_idx(m,n)].set(-((2*n-1)*z*coeff[...,mpl_idx(m,n-1)] + (x**2+y**2+z**2)*coeff[...,mpl_idx(m,n-2)])/((n-abs(m))*(n+abs(m))))
+    
+    return coeff
+
+@partial(jax.jit, static_argnames=['p'])
+def eval_singular_basis(rvec, p):   # NOTE: might have to include a factor (-1)^n, also use S_n^-m for evaluating
+    r"""
+    Evaluate real singular basis functions (Laplace kernel) with a recursion relation.
+
+    :param rvec: Array of positions in cartesian coordinates to evaluate the basis at.
+    :type rvec: jnp.array
+    :param p: Maximum degree (degree of multipole expansion) for the evaluation.
+    :type p: int
+
+    :return: Singular basis evaluated until the given degree at the given locations.
+    :rtype: jnp.array
+    """
+    x, y, z = rvec[..., 0], rvec[...,1], rvec[...,2]
+    r2 = x**2 + y**2 + z**2
+    coeff = jnp.zeros((*rvec.shape[:-1], (p+1)**2))
+    coeff = coeff.at[...,mpl_idx(0,0)].set(1/jnp.sqrt(r2))
+
+    if(p>0):
+        coeff = coeff.at[...,mpl_idx(1,1)].set(-coeff[...,mpl_idx(0,0)]*y/r2)
+        coeff = coeff.at[...,mpl_idx(-1,1)].set(coeff[...,mpl_idx(0,0)]*x/r2)
+
+    for n in range(2,p+1):    # first recursion: extreme values
+        coeff = coeff.at[...,mpl_idx(n,n)].set((2*n-1)*(x*coeff[...,mpl_idx(n-1,n-1)] - y*coeff[...,mpl_idx(1-n,n-1)])/r2)
+        coeff = coeff.at[...,mpl_idx(-n,n)].set((2*n-1)*(y*coeff[...,mpl_idx(n-1,n-1)] + x*coeff[...,mpl_idx(1-n,n-1)])/r2)
+
+    for n in range(0,p):      # second recursion: neighbors of extreme values
+        coeff = coeff.at[...,mpl_idx(n,n+1)].set((2*n+1)*z*coeff[...,mpl_idx(n,n)]/r2)
+        coeff = coeff.at[...,mpl_idx(-n,n+1)].set((2*n+1)*z*coeff[...,mpl_idx(-n,n)]/r2)
+
+    for n in range(2,p+1):    # third recursion: all values inbetween
+        for m in range(-n+2,n-1):
+            coeff = coeff.at[...,mpl_idx(m,n)].set(((2*n-1)*z*coeff[...,mpl_idx(m,n-1)] - (n-1-m)*(n-1+m)*coeff[...,mpl_idx(m,n-2)])/r2)
+    
+    return coeff
+
+@partial(jax.jit, static_argnames=['p'])
+def get_initial_mpls(padded_pts, padded_chrgs, boxcenters, p):
+    r"""
+    Get initial multipole expansions for each box on the highest level.
+
+    :param padded_pts: Array of point positions, sorted into the highest level of the hierarchy via balanced_tree().
+    :type padded_pts: jnp.array
+    :param padded_chrgs: Array of point charges, sorted into the highest level of the hierarchy via balanced_tree().
+    :type padded_chrgs: jnp.array
+    :param boxcenters: Array of box centers, generated via balanced_tree().
+    :type boxcenters: jnp.array
+    :param p: Multipole expansion order.
+    :type p: int
+
+    :return: Array of multipole coefficients of the boxes on the highest level.
+    :rtype: jnp.array
+    """
+    dist = padded_pts - boxcenters[:,None]
+    return (eval_regular_basis(dist,p) * padded_chrgs[...,None]).sum(axis=1)
+
+@partial(jax.jit, static_argnames=['p', 'n_split'])
+def M2M(coeff, oldboxdims, newboxdims, p, n_split):
+    r"""
+    Multipole-to-multipole transformation, merging "small" multipole expansions on higher levels into "large" multipole expansions on lower levels. This is the O(p^4) algorithm proposed in the original 3D FMM paper.
+
+    :param coeff: Array of multipole coefficients on level l.
+    :type coeff: jnp.array
+    :param oldboxdims: Array of box centers on level l.
+    :type oldboxdims: jnp.array
+    :param newboxdims: Array of box centers on level l-1.
+    :type newboxdims: jnp.array
+    :param p: Multipole expansion order.
+    :type p: int
+    :param n_split: How many splits per level and box have been performed. Each box has 2^n_split children which are merged into one in this function.
+    :type n_split: int
+
+    :return: Array of multipole coefficients of the boxes on level l-1.
+    :rtype: jnp.array
+    """
+    n_chi = 2**n_split
+    mpls = coeff.reshape((coeff.shape[0]//n_chi,n_chi,coeff.shape[1]))
+    new_mpls = jnp.zeros((mpls.shape[0],mpls.shape[2]))
+    oldboxdims = oldboxdims.reshape((-1,n_chi,3))
+    reg = eval_regular_basis(oldboxdims - newboxdims[:,None,:],p)   # shift direction points from target to source
+
+    for j in range(p+1):
+        for k in range(0,j+1):  # Real coeffs
+            for n in range(j+1):
+                for m in range(max(k+n-j,-n),min(k+j-n,n)+1):
+                    new_mpls = new_mpls.at[...,mpl_idx(k,j)].set(new_mpls[...,mpl_idx(k,j)] + (-1)**((abs(k)-abs(m)-abs(k-m))//2) * (
+                                                  reg[...,mpl_idx(abs(m),n)]*mpls[...,mpl_idx(abs(k-m),j-n)] - 
+                                                  jnp.sign(m)*jnp.sign(k-m)*reg[...,mpl_idx(-abs(m),n)]*mpls[...,mpl_idx(-abs(k-m),j-n)]).sum(axis=-1))
+        for k in range(-j,0):   # Imag coeffs
+            for n in range(j+1):
+                for m in range(max(k+n-j,-n),min(k+j-n,n)+1):
+                    new_mpls = new_mpls.at[...,mpl_idx(k,j)].set( new_mpls[...,mpl_idx(k,j)] - (-1)**((abs(k)-abs(m)-abs(k-m))//2) * (
+                                                  jnp.sign(k-m)*reg[...,mpl_idx(abs(m),n)]*mpls[...,mpl_idx(-abs(k-m),j-n)] + 
+                                                  jnp.sign(m)*reg[...,mpl_idx(-abs(m),n)]*mpls[...,mpl_idx(abs(k-m),j-n)]).sum(axis=-1))
+    return new_mpls
+
+@partial(jax.jit, static_argnames=['p', 'n_split'])
+def L2L(locs, oldboxdims, newboxdims, p, n_split):
+    r"""
+    Local-to-local transformation, distributing "large" local expansions on lower levels to "small" local expansions on higher levels. This is the O(p^4) algorithm proposed in the original 3D FMM paper.
+
+    :param locs: Array of local coefficients on level l.
+    :type locs: jnp.array
+    :param oldboxdims: Array of box centers on level l.
+    :type oldboxdims: jnp.array
+    :param newboxdims: Array of box centers on level l+1.
+    :type newboxdims: jnp.array
+    :param p: Multipole expansion order.
+    :type p: int
+    :param n_split: How many splits per level and box have been performed. Each box has 2^n_split children which all receive a local expansion of their parent in this function.
+    :type n_split: int
+
+    :return: Array of local coefficients of the boxes on level l+1.
+    :rtype: jnp.array
+    """
+    n_chi = 2**n_split
+    new_locs = jnp.zeros((locs.shape[0], n_chi, locs.shape[1]))
+    newboxdims = newboxdims.reshape((-1,n_chi,3))
+    reg = eval_regular_basis(oldboxdims[:,None,:]-newboxdims,p)   # shift direction points from target to source
+
+    for j in range(p+1):
+        for k in range(1,j+1):  # -Imag coeffs!
+            for n in range(j,p+1):
+                for m in range(k+j-n,k+n-j+1):
+                    new_locs = new_locs.at[...,mpl_idx(k,j)].set(new_locs[...,mpl_idx(k,j)] - (-1)**((abs(m)-abs(m-k)-abs(k))//2) * (
+                                                  -jnp.sign(m)*reg[...,mpl_idx(abs(m-k),n-j)]*locs[...,None,mpl_idx(abs(m),n)] + 
+                                                  jnp.sign(m-k)*reg[...,mpl_idx(-abs(m-k),n-j)]*locs[...,None,mpl_idx(-abs(m),n)]))
+        for k in range(-j,1):   # Real coeffs!
+            for n in range(j,p+1):
+                for m in range(k+j-n,k+n-j+1):
+                    new_locs = new_locs.at[...,mpl_idx(k,j)].set(new_locs[...,mpl_idx(k,j)] + (-1)**((abs(m)-abs(m-k)-abs(k))//2) * (
+                                                  reg[...,mpl_idx(abs(m-k),n-j)]*locs[...,None,mpl_idx(-abs(m),n)] + 
+                                                  jnp.sign(m-k)*jnp.sign(m)*reg[...,mpl_idx(-abs(m-k),n-j)]*locs[...,None,mpl_idx(abs(m),n)]))
+    return new_locs.reshape((-1,new_locs.shape[-1]))
+
+@partial(jax.jit, static_argnames=['p'])
+def M2L(mpls, locs, boxcenters, mpl_cnct, p):  # TODO: try to write this as a jax.pallas kernel?
+    r"""
+    Multipole-to-local transformation, turning multipole expansions into local expansions on the same level. This is the O(p^4) algorithm proposed in the original 3D FMM paper.
+
+    :param mpls: Array of multipole coefficients.
+    :type mpls: jnp.array
+    :param locs: Array of local coefficients.
+    :type locs: jnp.array
+    :param boxcenters: Array of box centers.
+    :type boxcenters: jnp.array
+    :param mpl_cnct: Array containing interaction partner indices for each box.
+    :type mpl_cnct: jnp.array
+    :param p: Multipole expansion order.
+    :type p: int
+
+    :return: Updated array of local coefficients, now includes all the well-separated multipole expansions on this level.
+    :rtype: jnp.array
+    """
+    sing = eval_singular_basis(boxcenters.at[mpl_cnct].get(mode="fill",fill_value=123456789)-boxcenters[:,None,:],p)    # shift direction points from target to source TODO: find a proper way of getting non-nan values here
+
+    for j in range(p+1):
+        for k in range(1,j+1):  # -Imag coeffs!
+            for n in range(p+1-j):
+                for m in range(max(k-j-n,-n),min(j+n+k,n)+1):
+                    locs = locs.at[:,mpl_idx(k,j)].set( locs[:,mpl_idx(k,j)] + (-1)**((abs(k-m)-abs(k)-abs(m))//2) * (
+                                                -jnp.sign(m-k)*sing[...,mpl_idx(abs(m-k),j+n)]*mpls.at[mpl_cnct,mpl_idx(abs(m),n)].get(mode="fill",fill_value=0) + 
+                                                jnp.sign(m)*sing[...,mpl_idx(-abs(m-k),j+n)]*mpls.at[mpl_cnct,mpl_idx(-abs(m),n)].get(mode="fill",fill_value=0)).sum(axis=1))
+        for k in range(-j,1):  # Real coeffs!
+            for n in range(p+1-j):
+                for m in range(max(k-j-n,-n),min(j+n+k,n)+1):
+                    locs = locs.at[:,mpl_idx(k,j)].set(locs[:,mpl_idx(k,j)] + (-1)**((abs(k-m)-abs(k)-abs(m))//2) * (
+                                                sing[...,mpl_idx(-abs(m-k),j+n)]*mpls.at[mpl_cnct,mpl_idx(abs(m),n)].get(mode="fill",fill_value=0) + 
+                                                jnp.sign(m)* jnp.sign(m-k) * sing[...,mpl_idx(abs(m-k),j+n)]*mpls.at[mpl_cnct,mpl_idx(-abs(m),n)].get(mode="fill",fill_value=0)).sum(axis=1))
+    return locs
+
+@partial(jax.jit, static_argnames=['p', 'n_split'])
+def go_up(coeff, boxcenters, p, n_split):
+    r"""
+    Using multipole-to-multipole transformation, descend the hierarchy.
+
+    :param coeff: Array of multipole coefficients on the highest level.
+    :type coeff: jnp.array
+    :param boxcenters: List containing box center arrays for every level.
+    :type boxcenters: list(jnp.array)
+    :param p: Multipole expansion order.
+    :type p: int
+    :param n_split: How many splits per level and box have been performed. Each box has 2^n_split children which are merged on every level for this step.
+    :type n_split: int
+
+    :return: List of multipole coefficient arrays, containing multipole expansion coefficients for every box on every level.
+    :rtype: list(jnp.array)
+    """
+    mpls = [coeff]
+    n_l = len(boxcenters)
+    for i in range(n_l-1):
+        mpls.append(M2M(mpls[-1],boxcenters[-(i+1)],boxcenters[-(i+2)], p, n_split))
+    return mpls
+
+@partial(jax.jit, static_argnames=['p', 'n_split'])
+def go_down(mpls, boxcenters, mpl_cnct, p, n_split):
+    r"""
+    Using multipole-to-local and local-to-local transformation, ascend the hierarchy.
+
+    :param mpls: List of multipole coefficient arrays on every level.
+    :type mpls: list(jnp.array)
+    :param boxcenters: List containing box center arrays for every level.
+    :type boxcenters: list(jnp.array)
+    :param mpl_cnct: List of interaction partner index arrays on every level.
+    :type mpl_cnct: list(jnp.array)
+    :param p: Multipole expansion order.
+    :type p: int
+    :param n_split: How many splits per level and box have been performed. Each box has 2^n_split children which receive local expansions from their parent in this step.
+    :type n_split: int
+
+    :return: List of local coefficient arrays, containing local expansion coefficients for every box on every level.
+    :rtype: list(jnp.array)
+    """
+    locs = [jnp.zeros((mpls[-1].shape))]
+    n_l = len(boxcenters)
+    for i in range(n_l-1):  # for every level TODO: level 1 can be skipped
+        locs.append(L2L(locs[-1],boxcenters[i],boxcenters[i+1], p, n_split)) # get local expansions from parent boxes
+        locs[-1] = M2L(mpls[-(i+2)], locs[-1], boxcenters[i+1], mpl_cnct[i+1], p) # next, it is time to do the M2L shifts for this level
+    return locs
+
+@partial(jax.jit, static_argnames=['p'])
+def eval_local(locs, padded_pts, rev_idcs, boxcenters, p):
+    r"""
+    Evaluate local expansions on the highest level.
+
+    :param locs: Array of local coefficients on the highest level.
+    :type locs: jnp.array
+    :param padded_pts: Array containing point positions, sorted into the highest level of the hierarchy.
+    :type padded_pts: jnp.array
+    :param rev_idcs: Index array to remove padding and return to original sorting.
+    :type rev_idcs: jnp.array
+    :param boxcenters: Array of box center positions on the highest level.
+    :type boxcenters: jnp.array
+    :param p: Multipole expansion order.
+    :type p: int
+
+    :return: Far-field potential in the original sorting positions.
+    :rtype: jnp.array
+    """
+    padded_res = jnp.zeros(padded_pts.shape[:2])
+    reg = eval_regular_basis(padded_pts-boxcenters[:,None],p)   # this evaluation needs to be relative to the boxcenter
+
+    for n in range(p+1):
+        for m in range(-n,n+1):
+            if(m!=0):
+                padded_res += (-1)**n * 2*locs[:,None,mpl_idx(-m,n)] * reg[...,mpl_idx(m,n)]
+            else:
+                padded_res += (-1)**n * locs[:,None,mpl_idx(-m,n)] * reg[...,mpl_idx(m,n)] 
+    return padded_res.flatten()[rev_idcs] / (4*jnp.pi)  # TODO: flatten and resort only once, after adding?
+
+@jax.jit
+def eval_direct(padded_pts, padded_chrgs, rev_idcs, direct_cnct): # TODO: replace some of these evaluations with multipole evaluations for better performance
+    r"""
+    Evaluate the near-field potential directly (P2P).
+
+    :param padded_pts: Array containing point positions, sorted into the highest level of the hierarchy.
+    :type padded_pts: jnp.array
+    :param padded_chrgs: Array containing point charges, sorted into the highest level of the hierarchy.
+    :type padded_chrgs: jnp.array
+    :param rev_idcs: Index array to remove padding and return to original sorting.
+    :type rev_idcs: jnp.array
+    :param direct_cnct: Index array of interaction partners for each box on the highest level.
+    :type direct_cnct: jnp.array
+
+    :return: Near-field potential in the original sorting positions.
+    :rtype: jnp.array
+    """
+    n = padded_pts.shape[1]
+    nrows = padded_pts.shape[0]
+
+    def i_body(i):
+        acc = jnp.zeros((n,))
+        xblk = padded_pts[i]
+        def k_body(k, acc):
+            partner = direct_cnct[i,k]
+            dists = jnp.linalg.norm(xblk[:,None] - padded_pts.at[partner].get(mode="fill",fill_value=jnp.inf),axis=-1)
+            dists = 1/jnp.where(dists==0,jnp.inf,dists)
+            chunk = padded_chrgs[partner]
+            return acc + dists.dot(chunk)
+        acc = jax.lax.fori_loop(0,direct_cnct.shape[1], k_body, acc)
+        return acc
+
+    accs = jax.vmap(i_body)(jnp.arange(nrows))
+    return accs.flatten()[rev_idcs]/(4*jnp.pi) # TODO: flatten and resort only once, after adding?
+
+def gen_hierarchy(pts, N_max = 128, theta = 0.75, n_split = 3):
+    r"""
+    Generate the balanced tree and connectivity for the FMM.
+
+    :param pts: Array of shape (N,3) containing the positions of N point charges.
+    :type pts: jnp.array
+    :param N_max: Maximum allowed number of point charges per box.
+    :type N_max: int, optional
+    :param theta: Well-separatedness parameter, determines accuracy.
+    :type theta: float, optional
+    :param n_split: How many splits per level and box are been performed. Each box has 2^n_split children.
+    :type n_split: int, optional
+
+    :return: Tuple containing full hierarchy information.
+    :rtype: tuple
+    """
+    max_l = get_max_l(pts.shape[0], N_max, n_split)
+    idcs, rev_idcs, boxcenters, boxlens = balanced_tree(pts, max_l, n_split)
+    mpl_cnct, dir_cnct = gen_connectivity(boxcenters, boxlens, theta, n_split)
+    return pts, idcs, rev_idcs, boxcenters, mpl_cnct, dir_cnct, n_split
+
+@partial(jax.jit, static_argnames=['p', 'n_split'])#,backend="cpu")
+def eval_potential(pts, idcs, rev_idcs, boxcenters, mpl_cnct, direct_cnct, n_split, chrgs, p = 3):
+    r"""
+    Full FMM potential evaluation, does not include creation of the hierarchy (see gen_hierarchy(), which generates the first 7 parameters). Therefore, only chrgs and p can be safely changed without recomputing the hierarchy.
+
+    :param pts: Array containing point positions.
+    :type pts: jnp.array
+    :param idcs: Index array to sort points/charges into highest level of the hierarchy.
+    :type idcs: jnp.array
+    :param rev_idcs: Index array to remove padding and return to original sorting.
+    :type rev_idcs: jnp.array
+    :param boxcenters: List containing box center arrays for every level.
+    :type boxcenters: list(jnp.array)
+    :param mpl_cnct: List of interaction partner index arrays on every level.
+    :type mpl_cnct: list(jnp.array)
+    :param direct_cnct: Index array of interaction partners for each box on the highest level.
+    :type direct_cnct: jnp.array
+    :param n_split: How many splits per level and box have been performed. Each box has 2^n_split children.
+    :type n_split: int
+    :param chrgs: Array containing point charges.
+    :type chrgs: jnp.array
+    :param p: Multipole expansion order.
+    :type p: int, optional
+
+    :return: Electrostatic potential of the points and corresponding charges.
+    :rtype: jnp.array
+    """
+    padded_pts = pts.at[idcs].get(mode="fill",fill_value=0.0)   # TODO: we could buffer this alternatively
+    padded_chrgs = chrgs.at[idcs].get(mode="fill",fill_value=0.0)
+    coeff = get_initial_mpls(padded_pts, padded_chrgs, boxcenters[-1], p)
+    mpls = go_up(coeff, boxcenters, p, n_split)
+    locs = go_down(mpls, boxcenters, mpl_cnct, p, n_split)
+    pot = eval_direct(padded_pts, padded_chrgs, rev_idcs, direct_cnct) + eval_local(locs[-1], padded_pts, rev_idcs, boxcenters[-1], p)
+    return pot
+
+@jax.jit
+def eval_potential_direct(pts, chrgs, eval_pts = None):
+    r"""
+    Evaluate the potential directly via pairwise sums.
+
+    :param pts: Array containing point positions.
+    :type padded_pts: jnp.array
+    :param chrgs: Array containing point charges.
+    :type chrgs: jnp.array
+    :param eval_pts: Array containing points to evaluate the potential at. Defaults to pts.
+    :type eval_pts: jnp.array, optional
+
+    :return: Electrostatic potential of the points and corresponding charges.
+    :rtype: jnp.array
+    """
+    if(eval_pts is None):
+        eval_pts = pts
+    res = jnp.zeros(eval_pts.shape[0])
+    def eval_direct_body(i, val):
+        inv_dists = jnp.linalg.norm(pts[:,:] - eval_pts[i,None,:],axis=-1)
+        inv_dists = 1/jnp.where(inv_dists==0,jnp.inf,inv_dists) # take out self-interaction
+        val = val.at[i].set((chrgs * inv_dists).sum()) 
+        return val
+    return jax.lax.fori_loop(0,eval_pts.shape[0],eval_direct_body,res)/(4*jnp.pi)

jaxfmm-0.0.1/pyproject.toml

@@ -0,0 +1,44 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "jaxFMM"
+version = "0.0.1"
+description = "Adaptive Fast Multipole Method with Laplace kernel in JAX."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { file = "LICENSE" }
+keywords = ["jax", "FMM", "potential", "treecode", "N-body"]
+authors = [{ name = "Robert Kraft", email = "robert.kraft@univie.ac.at" }]
+# maintainers = [
+#   { name = "A. Great Maintainer", email = "maintainer@example.com" },
+# ]
+classifiers = [
+  "Development Status :: 3 - Alpha",    # 4 - Beta, 5 - Production
+  "Intended Audience :: Science/Research",
+  "Topic :: Scientific/Engineering :: Physics",
+  "License :: OSI Approved :: GNU General Public License v3 (GPLv3)",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Programming Language :: Python :: 3.12",
+  "Programming Language :: Python :: 3.13",
+  "Programming Language :: Python :: 3 :: Only",
+]
+dependencies = ["jax"]
+
+[project.optional-dependencies]
+cuda = ["jax[cuda]"]
+dev = ["pytest","matplotlib"]
+
+[project.urls]
+"Homepage" = "https://gitlab.com/jaxfmm/jaxfmm"
+
+[tool.pytest.ini_options]
+minversion = "6.0"
+# addopts = "-ra -q"
+testpaths = [
+    "tests"
+]

jaxfmm-0.0.1/tests/test_expansion.py

@@ -0,0 +1,34 @@
+import jax.numpy as jnp
+import jaxfmm.fmm as fmm
+import jaxfmm.debug_helpers as debug
+from jax import random
+
+### TODO:
+#   - find similar tests for local coeffs and expansions
+
+def test_mpl_coeffs():  # TODO: also check for correct normalization
+    for n in range(9):
+        for m in range(-n,n+1):
+            pts, chrgs = debug.gen_multipole_dist(m,n,eps=10.0)    # special point charge distribution corresponding to multipole moments - set eps large to minimize error
+            tree_info = fmm.gen_hierarchy(pts)
+            coeff = fmm.get_initial_mpls(pts,chrgs,tree_info[3][0],n)[0,:]
+            test = jnp.where(jnp.abs(coeff)>1e-5)[0]   # only the (m,n) coefficient should be nonzero
+            assert test.shape[0] == 1
+            assert test[0] == fmm.mpl_idx(m,n)
+
+def test_mpl_eval():  # TODO: do a better test, maybe independent from computing coeffs (related to the above todo)
+    key = random.key(743)
+    pts = random.uniform(key,(128,3),minval=-1,maxval=1)
+    chrgs = random.uniform(key,128,minval=-1,maxval=1)
+
+    tree_info = fmm.gen_hierarchy(pts)
+    coeff = fmm.get_initial_mpls(pts,chrgs,tree_info[3][0],10)[0,:]
+
+    nside, sidelen = 100, 10
+    eval_pts = (jnp.mgrid[:nside+1,:nside+1,:nside+1].T/nside * sidelen - sidelen/2).reshape((-1,3))
+    eval_pts = eval_pts[jnp.linalg.norm(eval_pts,axis=-1)>2*jnp.sqrt(3)]
+
+    pot_fmm = debug.eval_multipole(coeff, tree_info[3][0], eval_pts)
+    pot_dir = fmm.eval_potential_direct(pts,chrgs,eval_pts)
+    max_err = jnp.max(jnp.abs(pot_fmm-pot_dir))
+    assert max_err < 5.75e-6

jaxfmm-0.0.1/tests/test_potential.py

@@ -0,0 +1,20 @@
+import jax.numpy as jnp
+from jaxfmm import *
+from jax import random
+
+### TODO:
+# - check if the FMM error scaling roughly works out
+
+def test_potential_unitcube():
+    N = 2**15
+    key = random.key(856)
+    pts = random.uniform(key,(N,3))
+    chrgs = random.uniform(key,N,minval=-1,maxval=1)
+
+    tree_info = gen_hierarchy(pts)
+    pot_FMM = eval_potential(*tree_info,chrgs)
+
+    pot_dir = eval_potential_direct(pts,chrgs)
+
+    err = jnp.linalg.norm(pot_dir-pot_FMM)/jnp.linalg.norm(pot_dir)
+    assert err < 3.75e-3

jaxfmm-0.0.1/tests/test_transformation.py

@@ -0,0 +1,21 @@
+import jax.numpy as jnp
+import jaxfmm.fmm as fmm
+from jax import random
+
+### TODO:
+#   - find similar tests for L2L, M2L
+ 
+def test_M2M():
+    key = random.key(825)
+    pts = random.uniform(key,(8*128,3),minval=-1,maxval=1)
+    chrgs = random.uniform(key,8*128,minval=-1,maxval=1)
+
+    tree_info = fmm.gen_hierarchy(pts)
+    padded_pts = pts.at[tree_info[1]].get(mode="fill",fill_value=0.0)
+    padded_chrgs = chrgs.at[tree_info[1]].get(mode="fill",fill_value=0.0)
+
+    coeff = fmm.get_initial_mpls(padded_pts,padded_chrgs,tree_info[3][1],10)
+    coeff_merged = fmm.M2M(coeff,tree_info[3][1], tree_info[3][0], 10, 3)[0,:]
+    coeff_dir = fmm.get_initial_mpls(pts,chrgs,tree_info[3][0],10)[0,:]
+
+    assert jnp.allclose(coeff_merged, coeff_dir, rtol=1e-3, atol=1e-6)

jaxfmm-0.0.1/tests/test_trees.py

@@ -0,0 +1,24 @@
+import jax.numpy as jnp
+import jaxfmm.fmm as fmm
+
+### TODO:
+#   - find better tests
+
+def test_tree_connectivity_unitcube():
+    nside, sidelen = 64, 10.0
+    pts = (jnp.mgrid[:nside,:nside,:nside].T/(nside-1) * sidelen - sidelen/2).reshape((-1,3))
+
+    max_l = fmm.get_max_l(pts.shape[0], 128)
+    assert max_l==4
+
+    idcs, rev_idcs, boxcenters, boxlens = fmm.balanced_tree(pts,max_l)
+    assert idcs.shape[1]==64
+    assert jnp.all(idcs < pts.shape[0])
+    assert jnp.allclose(pts, pts[idcs].reshape((-1,3))[rev_idcs])
+    for i in range(len(boxlens)):
+        assert jnp.allclose(boxlens[i], pts[nside//(2**i)-1,0] - pts[0,0])
+
+    mpl_cnct, dir_cnct = fmm.gen_connectivity(boxcenters, boxlens)
+
+    assert (mpl_cnct[-1]<mpl_cnct[-1].shape[0]).sum() == 611136
+    assert (dir_cnct<dir_cnct.shape[0]).sum() == 70336