diff --git a/.gitignore b/.gitignore index f1d900f..936d505 100644 --- a/.gitignore +++ b/.gitignore @@ -35,3 +35,7 @@ libtool .project .settings/ +build/ + +*.pyc +pvfmm.egg-info/ diff --git a/CMakeLists.txt b/CMakeLists.txt index b563ed4..68af91a 100644 --- a/CMakeLists.txt +++ b/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() diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index d5c3626..81bdb20 100644 --- a/examples/CMakeLists.txt +++ b/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) diff --git a/python/README.md b/python/README.md new file mode 100644 index 0000000..eb368b1 --- /dev/null +++ b/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/). diff --git a/python/examples/example1.py b/python/examples/example1.py new file mode 100644 index 0000000..e162a1b --- /dev/null +++ b/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 diff --git a/python/examples/example2.py b/python/examples/example2.py new file mode 100644 index 0000000..21ba191 --- /dev/null +++ b/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) diff --git a/python/pyproject.toml b/python/pyproject.toml new file mode 100644 index 0000000..12a791f --- /dev/null +++ b/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" = "bward@flatironinstitute.org" }] +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"] diff --git a/python/src/pvfmm/__init__.py b/python/src/pvfmm/__init__.py new file mode 100644 index 0000000..8e24482 --- /dev/null +++ b/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" diff --git a/python/src/pvfmm/ffi.py b/python/src/pvfmm/ffi.py new file mode 100644 index 0000000..5570d1a --- /dev/null +++ b/python/src/pvfmm/ffi.py @@ -0,0 +1,269 @@ +import ctypes +import ctypes.util +import os + +import numpy as np +from numpy.ctypeslib import ndpointer +from mpi4py import MPI + +# boilerplate from https://github.com/mpi4py/mpi4py/blob/master/demo/wrap-ctypes/helloworld.py +if MPI._sizeof(MPI.Comm) == ctypes.sizeof(ctypes.c_int): + MPI_Comm = ctypes.c_int +else: + MPI_Comm = ctypes.c_void_p + + +def get_MPI_COMM(comm) -> MPI_Comm: + comm_ptr = MPI._addressof(comm) + return MPI_Comm.from_address(comm_ptr) + + +def wrapped_ndptr(*args, **kwargs): + """ + A version of np.ctypeslib.ndpointer + which allows None (passed as NULL) + """ + base = ndpointer(*args, **kwargs) + + def from_param(cls, obj): + if obj is None: + return obj + return base.from_param(obj) + + return type(base.__name__, (base,), {"from_param": classmethod(from_param)}) + + +double_array = wrapped_ndptr(dtype=ctypes.c_double, flags=("C_CONTIGUOUS")) +float_array = wrapped_ndptr(dtype=ctypes.c_float, flags=("C_CONTIGUOUS")) + + +def volume_callback(type): + return ctypes.CFUNCTYPE( + None, ctypes.POINTER(type), ctypes.c_long, ctypes.POINTER(type), ctypes.c_void_p + ) + + +double_volume_callback = volume_callback(ctypes.c_double) +float_volume_callback = volume_callback(ctypes.c_float) + +PVFMMKernel = ctypes.c_uint # enum + +# try to load lib +_custom_location = os.getenv("PVFMM") +if _custom_location: + _dll = os.path.join(_custom_location, "libpvfmm.so") +else: + _dll = ctypes.util.find_library("pvfmm") + if _dll is None: + raise ImportError( + "Failed to find libpvfmm! \n" + "Set PVFMM environment variable or check your installation of pvfmm!" + ) + +SHARED_LIB = ctypes.CDLL(_dll) + +# somwhat automatically generated: +# clang2py --clang-args="-I/mnt/sw/nix/store/z5w5a7pr5cmdbds0pn9ajdgy0jg71sl6-gcc-11.3.0/lib/gcc/x86_64-pc-linux-gnu/11.3.0/include/" include/pvfmm.h -l ~/pvfmm/build/libpvfmm.so > python/auto.py + +PVFMMCreateVolumeFMMD = SHARED_LIB.PVFMMCreateVolumeFMMD +PVFMMCreateVolumeFMMD.restype = ctypes.c_void_p +PVFMMCreateVolumeFMMD.argtypes = [ctypes.c_int, ctypes.c_int, PVFMMKernel, MPI_Comm] + +PVFMMCreateVolumeFMMF = SHARED_LIB.PVFMMCreateVolumeFMMF +PVFMMCreateVolumeFMMF.restype = ctypes.c_void_p +PVFMMCreateVolumeFMMF.argtypes = [ctypes.c_int, ctypes.c_int, PVFMMKernel, MPI_Comm] + +PVFMMCreateVolumeTreeD = SHARED_LIB.PVFMMCreateVolumeTreeD +PVFMMCreateVolumeTreeD.restype = ctypes.c_void_p +PVFMMCreateVolumeTreeD.argtypes = [ + ctypes.c_int, + ctypes.c_int, + double_volume_callback, + ctypes.c_void_p, + double_array, + ctypes.c_long, + MPI_Comm, + ctypes.c_double, + ctypes.c_int, + ctypes.c_bool, + ctypes.c_int, +] +PVFMMCreateVolumeTreeF = SHARED_LIB.PVFMMCreateVolumeTreeF +PVFMMCreateVolumeTreeF.restype = ctypes.c_void_p +PVFMMCreateVolumeTreeF.argtypes = [ + ctypes.c_int, + ctypes.c_int, + float_volume_callback, + ctypes.c_void_p, + float_array, + ctypes.c_long, + MPI_Comm, + ctypes.c_float, + ctypes.c_int, + ctypes.c_bool, + ctypes.c_int, +] +PVFMMCreateVolumeTreeFromCoeffD = SHARED_LIB.PVFMMCreateVolumeTreeFromCoeffD +PVFMMCreateVolumeTreeFromCoeffD.restype = ctypes.c_void_p +PVFMMCreateVolumeTreeFromCoeffD.argtypes = [ + ctypes.c_long, + ctypes.c_int, + ctypes.c_int, + double_array, + double_array, + double_array, + ctypes.c_long, + MPI_Comm, + ctypes.c_bool, +] +PVFMMCreateVolumeTreeFromCoeffF = SHARED_LIB.PVFMMCreateVolumeTreeFromCoeffF +PVFMMCreateVolumeTreeFromCoeffF.restype = ctypes.c_void_p +PVFMMCreateVolumeTreeFromCoeffF.argtypes = [ + ctypes.c_long, + ctypes.c_int, + ctypes.c_int, + float_array, + float_array, + float_array, + ctypes.c_long, + MPI_Comm, + ctypes.c_bool, +] +PVFMMEvaluateVolumeFMMD = SHARED_LIB.PVFMMEvaluateVolumeFMMD +PVFMMEvaluateVolumeFMMD.restype = None +PVFMMEvaluateVolumeFMMD.argtypes = [ + double_array, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_long, +] +PVFMMEvaluateVolumeFMMF = SHARED_LIB.PVFMMEvaluateVolumeFMMF +PVFMMEvaluateVolumeFMMF.restype = None +PVFMMEvaluateVolumeFMMF.argtypes = [ + float_array, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_long, +] +PVFMMDestroyVolumeFMMD = SHARED_LIB.PVFMMDestroyVolumeFMMD +PVFMMDestroyVolumeFMMD.restype = None +PVFMMDestroyVolumeFMMD.argtypes = [ctypes.POINTER(ctypes.c_void_p)] +PVFMMDestroyVolumeFMMF = SHARED_LIB.PVFMMDestroyVolumeFMMF +PVFMMDestroyVolumeFMMF.restype = None +PVFMMDestroyVolumeFMMF.argtypes = [ctypes.POINTER(ctypes.c_void_p)] +PVFMMDestroyVolumeTreeD = SHARED_LIB.PVFMMDestroyVolumeTreeD +PVFMMDestroyVolumeTreeD.restype = None +PVFMMDestroyVolumeTreeD.argtypes = [ctypes.POINTER(ctypes.c_void_p)] +PVFMMDestroyVolumeTreeF = SHARED_LIB.PVFMMDestroyVolumeTreeF +PVFMMDestroyVolumeTreeF.restype = None +PVFMMDestroyVolumeTreeF.argtypes = [ctypes.POINTER(ctypes.c_void_p)] +PVFMMGetLeafCountD = SHARED_LIB.PVFMMGetLeafCountD +PVFMMGetLeafCountD.restype = ctypes.c_long +PVFMMGetLeafCountD.argtypes = [ctypes.c_void_p] +PVFMMGetLeafCountF = SHARED_LIB.PVFMMGetLeafCountF +PVFMMGetLeafCountF.restype = ctypes.c_long +PVFMMGetLeafCountF.argtypes = [ctypes.c_void_p] +PVFMMGetLeafCoordD = SHARED_LIB.PVFMMGetLeafCoordD +PVFMMGetLeafCoordD.restype = None +PVFMMGetLeafCoordD.argtypes = [double_array, ctypes.c_void_p] +PVFMMGetLeafCoordF = SHARED_LIB.PVFMMGetLeafCoordF +PVFMMGetLeafCoordF.restype = None +PVFMMGetLeafCoordF.argtypes = [float_array, ctypes.c_void_p] +PVFMMGetPotentialCoeffD = SHARED_LIB.PVFMMGetPotentialCoeffD +PVFMMGetPotentialCoeffD.restype = None +PVFMMGetPotentialCoeffD.argtypes = [ + double_array, + ctypes.c_void_p, +] +PVFMMGetPotentialCoeffF = SHARED_LIB.PVFMMGetPotentialCoeffF +PVFMMGetPotentialCoeffF.restype = None +PVFMMGetPotentialCoeffF.argtypes = [ + float_array, + ctypes.c_void_p, +] +PVFMMCoeff2NodesD = SHARED_LIB.PVFMMCoeff2NodesD +PVFMMCoeff2NodesD.restype = None +PVFMMCoeff2NodesD.argtypes = [ + double_array, + ctypes.c_long, + ctypes.c_int, + ctypes.c_int, + double_array, +] +PVFMMCoeff2NodesF = SHARED_LIB.PVFMMCoeff2NodesF +PVFMMCoeff2NodesF.restype = None +PVFMMCoeff2NodesF.argtypes = [ + float_array, + ctypes.c_long, + ctypes.c_int, + ctypes.c_int, + float_array, +] +PVFMMNodes2CoeffD = SHARED_LIB.PVFMMNodes2CoeffD +PVFMMNodes2CoeffD.restype = None +PVFMMNodes2CoeffD.argtypes = [ + double_array, + ctypes.c_long, + ctypes.c_int, + ctypes.c_int, + double_array, +] +PVFMMNodes2CoeffF = SHARED_LIB.PVFMMNodes2CoeffF +PVFMMNodes2CoeffF.restype = None +PVFMMNodes2CoeffF.argtypes = [ + float_array, + ctypes.c_long, + ctypes.c_int, + ctypes.c_int, + float_array, +] +PVFMMCreateContextD = SHARED_LIB.PVFMMCreateContextD +PVFMMCreateContextD.restype = ctypes.c_void_p +PVFMMCreateContextD.argtypes = [ + ctypes.c_double, + ctypes.c_int, + ctypes.c_int, + PVFMMKernel, + MPI_Comm, +] +PVFMMCreateContextF = SHARED_LIB.PVFMMCreateContextF +PVFMMCreateContextF.restype = ctypes.c_void_p +PVFMMCreateContextF.argtypes = [ + ctypes.c_float, + ctypes.c_int, + ctypes.c_int, + PVFMMKernel, + MPI_Comm, +] +PVFMMEvalD = SHARED_LIB.PVFMMEvalD +PVFMMEvalD.restype = None +PVFMMEvalD.argtypes = [ + double_array, + double_array, + double_array, + ctypes.c_long, + double_array, + double_array, + ctypes.c_long, + ctypes.c_void_p, + ctypes.c_int, +] +PVFMMEvalF = SHARED_LIB.PVFMMEvalF +PVFMMEvalF.restype = None +PVFMMEvalF.argtypes = [ + float_array, + float_array, + float_array, + ctypes.c_long, + float_array, + float_array, + ctypes.c_long, + ctypes.c_void_p, + ctypes.c_int, +] +PVFMMDestroyContextD = SHARED_LIB.PVFMMDestroyContextD +PVFMMDestroyContextD.restype = None +PVFMMDestroyContextD.argtypes = [ctypes.POINTER(ctypes.c_void_p)] +PVFMMDestroyContextF = SHARED_LIB.PVFMMDestroyContextF +PVFMMDestroyContextF.restype = None +PVFMMDestroyContextF.argtypes = [ctypes.POINTER(ctypes.c_void_p)] diff --git a/python/src/pvfmm/py.typed b/python/src/pvfmm/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/python/src/pvfmm/wrapper.py b/python/src/pvfmm/wrapper.py new file mode 100644 index 0000000..e54642e --- /dev/null +++ b/python/src/pvfmm/wrapper.py @@ -0,0 +1,333 @@ +import ctypes +import numpy as np +from enum import Enum +from typing import Optional, Callable, Union + +# calls MPI_Init_thread() and sets up MPI_Finalize() for you. +# see https://mpi4py.readthedocs.io/en/stable/mpi4py.run.html +from mpi4py import MPI + +from . import ffi + + +def nodes_to_coeff( + N_leaf: int, cheb_deg: int, dof: int, node_val: np.ndarray +) -> np.ndarray: + is_double = node_val.dtype == np.float64 + Ncoef = (cheb_deg + 1) * (cheb_deg + 2) * (cheb_deg + 3) // 6 + # TODO: is this the valid size of the output array? + coeff = np.empty(Ncoef * N_leaf * dof, dtype=node_val.dtype) + + get_function_dtype("PVFMMNodes2Coeff", node_val.dtype)( + coeff, N_leaf, cheb_deg, dof, node_val + ) + + return coeff + + +class FMMKernel(Enum): + """Mirroring PVFMMKernel in pvfmm.h""" + + LaplacePotential = 0 + LaplaceGradient = 1 + StokesPressure = 2 + StokesVelocity = 3 + StokesVelocityGrad = 4 + BiotSavartPotential = 5 + + +# read out of calls to BuildKernel in pvfmm/include/kernel.txx +KERNEL_DIMS = { + FMMKernel.LaplacePotential: (1, 1), + FMMKernel.LaplaceGradient: (1, 3), + FMMKernel.StokesPressure: (3, 1), + FMMKernel.StokesVelocity: (3, 3), + FMMKernel.StokesVelocityGrad: (3, 9), + FMMKernel.BiotSavartPotential: (3, 3), +} + + +def get_function_dtype(function_name: str, dtype: np.dtype) -> Callable: + """ + Helper function to switch between the double and float functions + from the FFI module + """ + if dtype == np.float64: + function_name += "D" + elif dtype == np.float32: + function_name += "F" + else: + raise ValueError("Invalid dtype, must be either float64 or float32") + + return getattr(ffi, function_name) + + +class FMMVolumeContext: + def __init__( + self, + multipole_order: int, + chebyshev_degree: int, + kernel: FMMKernel, + comm: MPI.Comm, + dtype=np.float64, + ): + self.kernel = kernel + self.dtype = np.dtype(dtype) + if multipole_order <= 0 or multipole_order % 2 != 0: + raise ValueError("multipole order must be even and postive") + self._ptr = get_function_dtype("PVFMMCreateVolumeFMM", dtype)( + multipole_order, + chebyshev_degree, + int(self.kernel.value), + ffi.get_MPI_COMM(comm), + ) + + def __del__(self): + if hasattr(self, "_ptr"): + get_function_dtype("PVFMMDestroyVolumeFMM", self.dtype)( + ctypes.byref(ctypes.c_void_p(self._ptr)) + ) + + +class FMMParticleContext: + def __init__( + self, + box_size: float, + max_points: int, + multipole_order: int, + kernel: FMMKernel, + comm: MPI.Comm, + dtype=np.float64, + ): + self.kernel = kernel + self.dtype = np.dtype(dtype) + if multipole_order <= 0 or multipole_order % 2 != 0: + raise ValueError("multipole order must be even and postive") + + self._ptr = get_function_dtype("PVFMMCreateContext", dtype)( + float(box_size), + max_points, + multipole_order, + int(self.kernel.value), + ffi.get_MPI_COMM(comm), + ) + + def __del__(self): + if hasattr(self, "_ptr"): + get_function_dtype("PVFMMDestroyContext", self.dtype)( + ctypes.byref(ctypes.c_void_p(self._ptr)) + ) + + def evaluate( + self, + src_pos: np.ndarray, + sl_den: np.ndarray, + dl_den: Optional[np.ndarray], + trg_pos: np.ndarray, + setup: bool = True, + ) -> np.ndarray: + if src_pos.dtype != self.dtype: + raise ValueError( + f"Source array had the wrong dtype: {src_pos.dtype}. " + f"This object was created with dtype {self.dtype}" + ) + + source_length = len(src_pos) + if source_length % 3 != 0: + raise ValueError( + "Source arrays must have a length which is a multiple of 3" + ) + + if sl_den is not None: + if sl_den.dtype != self.dtype: + raise ValueError( + f"Source array had the wrong dtype: {src_pos.dtype}. " + f"This object was created with dtype {self.dtype}" + ) + if len(sl_den) != source_length: + raise ValueError("Source arrays must all be of the same length!") + if dl_den is not None: + if dl_den.dtype != self.dtype: + raise ValueError( + f"Source array had the wrong dtype: {src_pos.dtype}. " + f"This object was created with dtype {self.dtype}" + ) + if len(dl_den) != source_length * 2: + raise ValueError("Source arrays must all be of the same length!") + + n_src = source_length // 3 + + target_length = len(trg_pos) + if target_length % 3 != 0: + raise ValueError( + "Target arrays must have a length which is a multiple of 3" + ) + n_trg = target_length // 3 + trg_val = np.empty(target_length, dtype=self.dtype) + + get_function_dtype("PVFMMEval", self.dtype)( + src_pos, + sl_den, + dl_den, + n_src, + trg_pos, + trg_val, + n_trg, + self._ptr, + int(setup), + ) + return trg_val + + +class FMMVolumeTree: + def __init__( + self, + ptr: ctypes.c_void_p, + cheb_deg: int, + data_dim: int, + n_trg: int, + dtype: np.dtype, + ): + self._ptr = ptr + self.cheb_deg = cheb_deg + self.n_cheb = (cheb_deg + 1) ** 3 + self.n_coeff = (cheb_deg + 1) * (cheb_deg + 2) * (cheb_deg + 3) // 6 + self.data_dim = data_dim + self.n_trg = n_trg + self.dtype = dtype + self._used_kernel = None + + @classmethod + def from_function( + cls, + cheb_deg: int, + data_dim: int, + fn: Union[ffi.double_volume_callback, ffi.float_volume_callback], + context: ctypes.c_void_p, + trg_coord: np.ndarray, + comm: MPI.Comm, + tol: float, + max_pts: int, + periodic: bool, + init_depth: int, + ) -> "FMMVolumeTree": + n_trg = len(trg_coord) // 3 + + dtype = trg_coord.dtype + ptr = get_function_dtype("PVFMMCreateVolumeTree", dtype)( + cheb_deg, + data_dim, + fn, + context, + trg_coord, + n_trg, + ffi.get_MPI_COMM(comm), + tol, + max_pts, + periodic, + init_depth, + ) + return cls(ptr, cheb_deg, data_dim, n_trg, dtype) + + @classmethod + def from_coefficients( + cls, + cheb_deg: int, + data_dim: int, + leaf_coord: np.ndarray, + fn_coeff: np.ndarray, + trg_coord: Optional[np.ndarray], + comm: MPI.Comm, + periodic: bool, + ) -> "FMMVolumeTree": + if len(leaf_coord) % 3 != 0: + raise ValueError( + "Leaf coordinates must have a length which is a multiple of 3" + ) + N_leaf = len(leaf_coord) // 3 + fn_coeff_size = ( + N_leaf * data_dim * (cheb_deg + 1) * (cheb_deg + 2) * (cheb_deg + 3) // 6 + ) + if len(fn_coeff) != fn_coeff_size: + raise ValueError( + "Function coefficients array has the wrong length, required length " + + fn_coeff_size + ) + dtype = leaf_coord.dtype + if fn_coeff.dtype != dtype: + raise ValueError( + f"Mismatching dtypes. Leaves had dtype {dtype}, " + f"but coefficients had dtype {fn_coeff.dtype}" + ) + if trg_coord is not None: + n_trg = len(trg_coord) // 3 + if trg_coord.dtype != dtype: + raise ValueError( + f"Mismatching dtypes. Leaves had dtype {dtype}, " + f"but targets had dtype {trg_coord.dtype}" + ) + else: + n_trg = 0 + + ptr = get_function_dtype("PVFMMCreateVolumeTreeFromCoeff", dtype)( + N_leaf, + cheb_deg, + data_dim, + leaf_coord, + fn_coeff, + trg_coord, + n_trg, + ffi.get_MPI_COMM(comm), + periodic, + ) + return cls(ptr, cheb_deg, data_dim, n_trg, dtype) + + def __del__(self): + if hasattr(self, "_ptr"): + get_function_dtype("PVFMMDestroyVolumeTree", self.dtype)( + ctypes.byref(ctypes.c_void_p(self._ptr)) + ) + + def evaluate(self, fmm: FMMVolumeContext, loc_size: int) -> np.ndarray: + if fmm.dtype != self.dtype: + raise ValueError( + f"Volume context has dtype {fmm.dtype}, " + f"but this tree has dtype {self.dtype}" + ) + (_kdim0, kdim1) = KERNEL_DIMS[fmm.kernel] + trg_val = np.empty(self.n_trg * kdim1, dtype=self.dtype) + get_function_dtype("PVFMMEvaluateVolumeFMM", self.dtype)( + trg_val, self._ptr, fmm._ptr, loc_size + ) + self._used_kernel = fmm.kernel + return trg_val + + def leaf_count(self) -> int: + return int(get_function_dtype("PVFMMGetLeafCount", self.dtype)(self._ptr)) + + def get_leaf_coordinates(self) -> np.ndarray: + Nleaf = self.leaf_count() + leaf_coord = np.empty(Nleaf * 3, dtype=self.dtype) + get_function_dtype("PVFMMGetLeafCoord", self.dtype)(leaf_coord, self._ptr) + return leaf_coord + + def get_coefficients(self) -> np.ndarray: + if self._used_kernel is None: + raise ValueError( + "Cannot get coefficients of an un-evaluated tree" + ) # TODO: is this true? what is the contract of this class + n_leaf = self.leaf_count() + (_kdim0, kdim1) = KERNEL_DIMS[self._used_kernel] + coeff = np.empty(n_leaf * self.n_coeff * kdim1, dtype=self.dtype) + get_function_dtype("PVFMMGetPotentialCoeff", self.dtype)(coeff, self._ptr) + return coeff + + def get_values(self) -> np.ndarray: + coeff = self.get_coefficients() + n_leaf = self.leaf_count() + (_kdim0, kdim1) = KERNEL_DIMS[self._used_kernel] + value = np.empty(n_leaf * self.n_cheb * kdim1, dtype=self.dtype) + get_function_dtype("PVFMMCoeff2Nodes", self.dtype)( + value, n_leaf, self.cheb_deg, self.data_dim, coeff + ) + return value