Merge pull request #15 from WardBrian/python-bindings

Add Python bindings

.gitignore

@@ -35,3 +35,7 @@ libtool
 .project
 .settings/
 
+build/
+
+*.pyc
+pvfmm.egg-info/

CMakeLists.txt

@@ -2,7 +2,7 @@ cmake_minimum_required(VERSION 3.10)
 project(
   PVFMM
   VERSION 1.2
-  LANGUAGES CXX)
+  LANGUAGES CXX C)
 
 set(CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake" ${CMAKE_MODULE_PATH})
 set(CMAKE_CXX_STANDARD 14)
@@ -28,7 +28,7 @@ set(MPI_CXX_SKIP_MPICXX
     true
     CACHE BOOL
           "If true, the MPI-2 C++ bindings are disabled using definitions.")
-          
+
 # required compiler features
 find_package(MPI REQUIRED)
 find_package(OpenMP REQUIRED)
@@ -44,7 +44,7 @@ else()
   find_package(FFTW REQUIRED)
   set(DEP_INC ${FFTW_INCLUDE_DIRS})
   set(DEP_LIB ${FFTW_FLOAT_OPENMP_LIB} ${FFTW_FLOAT_LIB}
-              ${FFTW_DOUBLE_OPENMP_LIB} ${FFTW_DOUBLE_LIB} 
+              ${FFTW_DOUBLE_OPENMP_LIB} ${FFTW_DOUBLE_LIB}
               ${BLAS_LIBRARIES}
               ${LAPACK_LIBRARIES})
 endif()

examples/CMakeLists.txt

@@ -5,6 +5,16 @@ else()
   unset(example_DEPS)
 endif()
 
+
+add_executable(example-c ./src/example-c.c)
+target_link_libraries(example-c PRIVATE pvfmmStatic "${example_DEPS}")
+target_include_directories(example-c PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)
+
+add_executable(example2-c ./src/example2-c.c)
+target_link_libraries(example2-c PRIVATE pvfmmStatic "${example_DEPS}")
+target_include_directories(example2-c PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)
+
+
 add_executable(example1 ./src/example1.cpp)
 target_link_libraries(example1 PRIVATE pvfmmStatic "${example_DEPS}")
 target_include_directories(example1 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/include)

python/README.md

@@ -0,0 +1,26 @@
+# Python bindings to PVFMM
+
+This package provides Python bindings to the [C API](../include/pvfmm.h) using
+[`ctypes`](https://docs.python.org/3/library/ctypes.html) and
+[`mpi4py`](https://mpi4py.readthedocs.io/en/stable/index.html).
+
+## Installation
+
+Build and install the PVFMM library as usual, and then run `pip install .`
+inside this directory.
+
+## Usage
+
+If you did not install PVFMM (see [INSTALL](../INSTALL)), you will need
+to tell the Python library where to find the `libpvfmm.so` file with
+`export PVFMM=/path/to/build/folder`.
+
+Then, you should be able to run code which uses it by using
+```
+mpirun $other_args_here python -m mpi4py path/to/program.py
+```
+
+## Examples
+
+See the `examples/` subdirectory for Python implementations of a few example programs.
+The examples rely on [numba](https://numba.pydata.org/).

python/examples/example1.py

@@ -0,0 +1,87 @@
+# based on examples/src/example-c.c
+
+import time
+
+from mpi4py import MPI
+import numpy as np
+import numba
+
+import pvfmm
+
+
+@numba.njit(parallel=True)
+def BiotSavart(src_X, src_V, trg_X):
+    Ns = len(src_X) // 3
+    Nt = len(trg_X) // 3
+    trg_V = np.zeros(Nt * 3)
+    oofp = 1 / (4 * np.pi)
+    for t in numba.prange(Nt):
+        for s in range(Ns):
+            X0 = trg_X[t * 3 + 0] - src_X[s * 3 + 0]
+            X1 = trg_X[t * 3 + 1] - src_X[s * 3 + 1]
+            X2 = trg_X[t * 3 + 2] - src_X[s * 3 + 2]
+            r2 = X0 * X0 + X1 * X1 + X2 * X2
+            rinv = 1 / np.sqrt(r2) if r2 > 0 else 0
+            rinv3 = rinv * rinv * rinv
+
+            trg_V[t * 3 + 0] += (
+                (src_V[s * 3 + 1] * X2 - src_V[s * 3 + 2] * X1) * rinv3 * oofp
+            )
+            trg_V[t * 3 + 1] += (
+                (src_V[s * 3 + 2] * X0 - src_V[s * 3 + 0] * X2) * rinv3 * oofp
+            )
+            trg_V[t * 3 + 2] += (
+                (src_V[s * 3 + 0] * X1 - src_V[s * 3 + 1] * X0) * rinv3 * oofp
+            )
+    return trg_V
+
+
+def test_FMM(ctx):
+    kdim = (3, 3)
+    Ns = 20000
+    Nt = 20000
+
+    # Initialize target coordinates
+    trg_X = np.random.rand(Nt * 3)
+
+    # Initialize source coordinates and density
+    src_X = np.random.rand(Ns * 3)
+    src_V = np.random.rand(Ns * kdim[0]) - 0.5
+
+    # FMM evaluation
+    tick = time.perf_counter_ns()
+    trg_V = ctx.evaluate(src_X, src_V, None, trg_X)
+    tock = time.perf_counter_ns()
+    print("FMM evaluation time (with setup) :", (tock - tick) / 1e9)
+
+    # FMM evaluation (without setup)
+    tick = time.perf_counter_ns()
+    trg_V2 = ctx.evaluate(src_X, src_V, None, trg_X, setup=False)
+    tock = time.perf_counter_ns()
+    print("FMM evaluation time (without setup) :", (tock - tick) / 1e9)
+
+    # Direct evaluation
+    tick = time.perf_counter_ns()
+    trg_V0 = BiotSavart(src_X, src_V, trg_X)
+    tock = time.perf_counter_ns()
+    print("Direct evaluation time :", (tock - tick) / 1e9)
+
+    max_val = np.max(np.abs(trg_V0))
+    max_err = np.max(np.abs(trg_V - trg_V0))
+    print("Max relative error :", max_err / max_val)
+
+
+if __name__ == "__main__":
+    # MPI init handled by mpi4py import
+    box_size = -1
+    points_per_box = 1000
+    multipole_order = 10
+    kernel = pvfmm.FMMKernel.BiotSavartPotential
+
+    print("Loaded.")
+    ctx = pvfmm.FMMParticleContext(
+        box_size, points_per_box, multipole_order, kernel, MPI.COMM_WORLD
+    )
+    print("Running!")
+    test_FMM(ctx)
+    # MPI finalize handled by mpi4py atexit handler

python/examples/example2.py

@@ -0,0 +1,168 @@
+# based on examples/src/example2-c.c
+
+import numpy as np
+import numba
+from numba import types
+import ctypes
+from mpi4py import MPI
+
+import pvfmm
+
+
+@numba.njit
+def fn_input(coord, out):
+    n = len(coord) // 3
+    L = 125
+    for i in range(n):
+        idx = i * 3
+        c = coord[idx : idx + 3]
+        r_2 = (c[0] - 0.5) ** 2 + (c[1] - 0.5) ** 2 + (c[2] - 0.5) ** 2
+        out[idx + 0] = 0 + 2 * L * np.exp(-L * r_2) * (c[0] - 0.5)
+        out[idx + 1] = 4 * L**2 * (c[2] - 0.5) * (5 - 2 * L * r_2) * np.exp(
+            -L * r_2
+        ) + 2 * L * np.exp(-L * r_2) * (c[1] - 0.5)
+        out[idx + 2] = -4 * L**2 * (c[1] - 0.5) * (5 - 2 * L * r_2) * np.exp(
+            -L * r_2
+        ) + 2 * L * np.exp(-L * r_2) * (c[2] - 0.5)
+
+
+c_sig = types.void(
+    types.CPointer(types.double),
+    types.int64,
+    types.CPointer(types.double),
+    types.voidptr,
+)
+
+
+# see https://numba.readthedocs.io/en/stable/user/cfunc.html
+@numba.cfunc(c_sig, nopython=True)
+def fn_input_C(coord_, n, out_, _ctx):
+    coord = numba.carray(coord_, n * 3)
+    out = numba.carray(out_, n * 3)
+    fn_input(coord, out)
+
+
+@numba.njit
+def fn_poten(coord):
+    n = len(coord) // 3
+    dof = 3
+    out = np.zeros(n * dof, dtype=np.float64)
+    L = 125
+
+    for i in range(n):
+        idx = i * dof
+        c = coord[idx : idx + dof]
+        r_2 = (c[0] - 0.5) ** 2 + (c[1] - 0.5) ** 2 + (c[2] - 0.5) ** 2
+        out[idx + 0] = 0
+        out[idx + 1] = 2 * L * (c[2] - 0.5) * np.exp(-L * r_2)
+        out[idx + 2] = -2 * L * (c[1] - 0.5) * np.exp(-L * r_2)
+
+    return out
+
+
+@numba.njit
+def GetChebNodes(Nleaf, cheb_deg, depth, leaf_coord):
+    leaf_length = 1.0 / (1 << depth)
+    Ncheb = (cheb_deg + 1) ** 3
+    cheb_coord = np.zeros(Nleaf * Ncheb * 3, dtype=np.float64)
+
+    for leaf_idx in range(Nleaf):
+        for j2 in range(cheb_deg + 1):
+            for j1 in range(cheb_deg + 1):
+                for j0 in range(cheb_deg + 1):
+                    node_idx = (
+                        leaf_idx * Ncheb
+                        + (j2 * (cheb_deg + 1) + j1) * (cheb_deg + 1)
+                        + j0
+                    )
+                    cheb_coord[node_idx * 3 + 0] = (
+                        leaf_coord[leaf_idx * 3 + 0]
+                        + (1 - np.cos(np.pi * (j0 * 2 + 1) / (cheb_deg * 2 + 2)))
+                        * leaf_length
+                        * 0.5
+                    )
+                    cheb_coord[node_idx * 3 + 1] = (
+                        leaf_coord[leaf_idx * 3 + 1]
+                        + (1 - np.cos(np.pi * (j1 * 2 + 1) / (cheb_deg * 2 + 2)))
+                        * leaf_length
+                        * 0.5
+                    )
+                    cheb_coord[node_idx * 3 + 2] = (
+                        leaf_coord[leaf_idx * 3 + 2]
+                        + (1 - np.cos(np.pi * (j2 * 2 + 1) / (cheb_deg * 2 + 2)))
+                        * leaf_length
+                        * 0.5
+                    )
+
+    return cheb_coord
+
+
+def test1(fmm, kdim0, kdim1, cheb_deg, comm):
+    Nt = 100
+    trg_coord = np.random.rand(Nt * 3)
+    trg_value_ref = fn_poten(trg_coord)
+    tree = pvfmm.FMMVolumeTree.from_function(
+        cheb_deg, kdim0, fn_input_C.ctypes, None, trg_coord, comm, 1e-6, 100, False, 0
+    )
+    trg_value = tree.evaluate(fmm, Nt)
+
+    max_val = np.max(np.abs(trg_value_ref))
+    max_err = np.max(np.abs(trg_value - trg_value_ref))
+    print("Maximum relative error = ", max_err / max_val)
+
+
+def test2(fmm, kdim0, kdim1, cheb_deg, comm):
+    Ncheb = (cheb_deg + 1) * (cheb_deg + 1) * (cheb_deg + 1)
+    Ncoef = (cheb_deg + 1) * (cheb_deg + 2) * (cheb_deg + 3) // 6
+
+    depth = 3
+    Nleaf = 1 << (3 * depth)
+    leaf_length = 1 / (1 << depth)
+
+    leaf_coord = np.empty(Nleaf * 3, dtype=np.float64)
+    dens_value = np.empty(Nleaf * Ncheb * kdim0, dtype=np.float64)
+
+    for leaf_idx in range(Nleaf):
+        leaf_coord[leaf_idx * 3 + 0] = (
+            (leaf_idx // (1 << (depth * 0))) % (1 << depth) * leaf_length
+        )
+        leaf_coord[leaf_idx * 3 + 1] = (
+            (leaf_idx // (1 << (depth * 1))) % (1 << depth) * leaf_length
+        )
+        leaf_coord[leaf_idx * 3 + 2] = (
+            (leaf_idx // (1 << (depth * 2))) % (1 << depth) * leaf_length
+        )
+    print("Getting nodes")
+    cheb_coord = GetChebNodes(Nleaf, cheb_deg, depth, leaf_coord)
+
+    fn_input(cheb_coord, dens_value)
+
+    dense_coeff = pvfmm.nodes_to_coeff(Nleaf, cheb_deg, kdim0, dens_value)
+    tree = pvfmm.FMMVolumeTree.from_coefficients(
+        cheb_deg, kdim0, leaf_coord, dense_coeff, None, comm, False
+    )
+    tree.evaluate(fmm, 0)
+    potn_value = tree.get_values()
+
+    Nleaf = tree.leaf_count()  # TODO: does this change from above?
+    leaf_coord = tree.get_leaf_coordinates()
+    cheb_coord = GetChebNodes(Nleaf, cheb_deg, depth, leaf_coord)
+    potn_value_ref = fn_poten(cheb_coord)
+
+    max_val = np.max(np.abs(potn_value_ref))
+    max_err = np.max(np.abs(potn_value - potn_value_ref))
+    print("Maximum relative error = ", max_err / max_val)
+
+
+if __name__ == "__main__":
+    mult_order = 10
+    cheb_deg = 14
+    kdim0 = 3
+    kdim1 = 3
+
+    comm = MPI.COMM_WORLD
+    fmm = pvfmm.FMMVolumeContext(
+        mult_order, cheb_deg, pvfmm.FMMKernel.StokesVelocity, comm
+    )
+    test1(fmm, kdim0, kdim1, cheb_deg, comm)
+    test2(fmm, kdim0, kdim1, cheb_deg, comm)

python/pyproject.toml

@@ -0,0 +1,31 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "pvfmm"
+readme = "README.md"
+authors = [{ "name" = "Brian Ward", "email" = "[email protected]" }]
+dependencies = ["numpy", "mpi4py"]
+requires-python = ">=3.8"
+license = { text = "BSD-3-Clause" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Development Status :: 4 - Beta",
+]
+dynamic = ["version"]
+
+[project.optional-dependencies]
+all = ["numba"]
+
+[tool.setuptools.dynamic]
+version = { attr = "pvfmm.__version__" }
+
+[tool.isort]
+profile = "black"
+
+[tool.setuptools.package-data]
+"pvfmm" = ["py.typed"]
+
+[tool.setuptools.packages.find]
+where = ["src"]

python/src/pvfmm/__init__.py

@@ -0,0 +1,18 @@
+from .wrapper import (
+    FMMKernel,
+    FMMVolumeContext,
+    FMMParticleContext,
+    FMMVolumeTree,
+    nodes_to_coeff,
+)
+
+
+__all__ = [
+    "FMMKernel",
+    "FMMVolumeContext",
+    "FMMParticleContext",
+    "FMMVolumeTree",
+    "nodes_to_coeff",
+]
+
+__version__ = "0.1.0"