Skip to content

Commit

Permalink
Add some uniformity, rsvd, and test blkrcmp for helmholtz
Browse files Browse the repository at this point in the history
  • Loading branch information
@skailasa
skailasa committed Dec 3, 2024
1 parent abdb62f commit e7eb5ba
Show file tree
Hide file tree
Showing 9 changed files with 937 additions and 60 deletions.
14 changes: 11 additions & 3 deletions kifmm/examples/single_node_helmholtz.rs
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use green_kernels::{helmholtz_3d::Helmholtz3dKernel, types::GreenKernelEvalType};
use kifmm::fmm::types::BlasFieldTranslationIa;
use kifmm::{Evaluate, FftFieldTranslation, SingleNodeBuilder};
use kifmm::{Evaluate, FftFieldTranslation, FmmSvdMode, SingleNodeBuilder};

use kifmm::tree::helpers::points_fixture;
use num::{FromPrimitive, One};
Expand Down Expand Up @@ -70,7 +70,11 @@ fn main() {
&expansion_order,
Helmholtz3dKernel::new(wavenumber),
GreenKernelEvalType::Value,
BlasFieldTranslationIa::new(singular_value_threshold, None),
BlasFieldTranslationIa::new(
singular_value_threshold,
None,
FmmSvdMode::Deterministic,
),
)
.unwrap()
.build()
Expand Down Expand Up @@ -99,7 +103,11 @@ fn main() {
&expansion_order,
Helmholtz3dKernel::new(wavenumber),
GreenKernelEvalType::Value,
BlasFieldTranslationIa::new(singular_value_threshold, None),
BlasFieldTranslationIa::new(
singular_value_threshold,
None,
FmmSvdMode::Deterministic,
),
)
.unwrap()
.build()
Expand Down
108 changes: 108 additions & 0 deletions kifmm/include/kifmm_rs.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -613,6 +613,114 @@ struct FmmEvaluator *laplace_fft_f64_alloc(bool timed,
uint64_t depth,
uintptr_t block_size);

/**
* Constructor for F32 Helmholtz FMM with BLAS based M2L translations compressed
* with deterministic SVD.
*
* Note that either `n_crit` or `depth` must be specified. If `n_crit` is specified, depth
* must be set to 0 and vice versa. If `depth` is specified, `expansion_order` must be specified
* at each level, and stored as a buffer of length `depth` + 1.
*
*
* # Parameters
* - `timed`: Modulates whether operators and metadata are timed.
* - `expansion_order`: A pointer to an array of expansion orders.
* - `n_expansion_order`: The number of expansion orders.
* - `eval_type`: true corresponds to evaluating potentials, false corresponds to evaluating potentials and potential derivatives
* - `wavenumber`: The wavenumber.
* - `sources`: A pointer to the source points.
* - `n_sources`: The length of the source points buffer
* - `targets`: A pointer to the target points.
* - `n_targets`: The length of the target points buffer.
* - `charges`: A pointer to the charges associated with the source points.
* - `n_charges`: The length of the charges buffer.
* - `prune_empty`: A boolean flag indicating whether to prune empty leaf nodes, and their ancestors.
* - `n_crit`: Threshold for tree refinement, if set to 0 ignored. Otherwise will refine until threshold is
* reached based on a uniform particle distribution.
* - `depth`: The maximum depth of the tree, max supported depth is 16.
* - `singular_value_threshold`: Threshold for singular values used in compressing M2L matrices.
* - `surface_diff`: Set to 0 to disable, otherwise uses surface_diff+equivalent_surface_expansion_order = check_surface_expansion_order
* - `n_components`: If known, can specify the rank of the M2L matrix for randomised range finding, otherwise set to 0.
* - `n_oversamples`: Optionally choose the number of oversamples for randomised range finding, otherwise set to 10.
*
* # Safety
* This function is intended to be called from C. The caller must ensure that:
* - Input data corresponds to valid pointers
* - That they remain valid for the duration of the function call
*/
struct FmmEvaluator *helmholtz_blas_rsvd_f32_alloc(bool timed,
const uintptr_t *expansion_order,
uintptr_t n_expansion_order,
bool eval_type,
float wavenumber,
const void *sources,
uintptr_t n_sources,
const void *targets,
uintptr_t n_targets,
const void *charges,
uintptr_t n_charges,
bool prune_empty,
uint64_t n_crit,
uint64_t depth,
float singular_value_threshold,
uintptr_t surface_diff,
uintptr_t n_components,
uintptr_t n_oversamples);

/**
* Constructor for F64 Helmholtz FMM with BLAS based M2L translations compressed
* with deterministic SVD.
*
* Note that either `n_crit` or `depth` must be specified. If `n_crit` is specified, depth
* must be set to 0 and vice versa. If `depth` is specified, `expansion_order` must be specified
* at each level, and stored as a buffer of length `depth` + 1.
*
*
* # Parameters
* - `timed`: Modulates whether operators and metadata are timed.
* - `expansion_order`: A pointer to an array of expansion orders.
* - `n_expansion_order`: The number of expansion orders.
* - `eval_type`: true corresponds to evaluating potentials, false corresponds to evaluating potentials and potential derivatives
* - `wavenumber`: The wavenumber.
* - `sources`: A pointer to the source points.
* - `n_sources`: The length of the source points buffer
* - `targets`: A pointer to the target points.
* - `n_targets`: The length of the target points buffer.
* - `charges`: A pointer to the charges associated with the source points.
* - `n_charges`: The length of the charges buffer.
* - `prune_empty`: A boolean flag indicating whether to prune empty leaf nodes, and their ancestors.
* - `n_crit`: Threshold for tree refinement, if set to 0 ignored. Otherwise will refine until threshold is
* reached based on a uniform particle distribution.
* - `depth`: The maximum depth of the tree, max supported depth is 16.
* - `singular_value_threshold`: Threshold for singular values used in compressing M2L matrices.
* - `surface_diff`: Set to 0 to disable, otherwise uses surface_diff+equivalent_surface_expansion_order = check_surface_expansion_order
* - `n_components`: If known, can specify the rank of the M2L matrix for randomised range finding, otherwise set to 0.
* - `n_oversamples`: Optionally choose the number of oversamples for randomised range finding, otherwise set to 10.
*
* # Safety
* This function is intended to be called from C. The caller must ensure that:
* - Input data corresponds to valid pointers
* - That they remain valid for the duration of the function call
*/
struct FmmEvaluator *helmholtz_blas_rsvd_f64_alloc(bool timed,
const uintptr_t *expansion_order,
uintptr_t n_expansion_order,
bool eval_type,
double wavenumber,
const void *sources,
uintptr_t n_sources,
const void *targets,
uintptr_t n_targets,
const void *charges,
uintptr_t n_charges,
bool prune_empty,
uint64_t n_crit,
uint64_t depth,
double singular_value_threshold,
uintptr_t surface_diff,
uintptr_t n_components,
uintptr_t n_oversamples);

/**
* Constructor for F32 Helmholtz FMM with BLAS based M2L translations compressed
* with deterministic SVD.
Expand Down
81 changes: 57 additions & 24 deletions kifmm/python/examples/API Usage.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@
{
"data": {
"text/plain": [
"{'source_data': 4, 'target_data': 3, 'source_to_target_data': 256}"
"{'source_data': 2, 'target_data': 2, 'source_to_target_data': 252}"
]
},
"execution_count": 3,
Expand Down Expand Up @@ -266,7 +266,7 @@
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": 9,
"id": "89147a31-182c-4bd0-920e-000fb5ff5f9f",
"metadata": {},
"outputs": [],
Expand All @@ -288,7 +288,7 @@
},
{
"cell_type": "code",
"execution_count": 9,
"execution_count": 10,
"id": "333eaf61-cb42-45d2-ae37-382cf866e9d6",
"metadata": {},
"outputs": [],
Expand All @@ -313,7 +313,7 @@
},
{
"cell_type": "code",
"execution_count": 10,
"execution_count": 11,
"id": "363d8d4b-f967-44dd-b93e-c97192295ff4",
"metadata": {},
"outputs": [],
Expand Down Expand Up @@ -351,7 +351,7 @@
},
{
"cell_type": "code",
"execution_count": 11,
"execution_count": 12,
"id": "25d4692c-d7ab-4f03-b2aa-ca92da29b961",
"metadata": {},
"outputs": [],
Expand Down Expand Up @@ -390,7 +390,7 @@
},
{
"cell_type": "code",
"execution_count": 12,
"execution_count": 5,
"id": "cecccd0e-f229-4d1a-9849-7c73c725b58f",
"metadata": {},
"outputs": [
Expand Down Expand Up @@ -519,33 +519,66 @@
"direct = np.zeros(n_targets * eval_type.value).astype(ctype)\n",
"fmm.evaluate_kernel(EvalType.Value, sources, targets, charges, direct)\n",
"\n",
"# assert np.allclose(fmm.all_potentials_u, direct)\n",
"assert np.allclose(np.abs(direct), np.abs(fmm.all_potentials_u[0].T), 1e-3)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "24a5eb2e-b718-4b14-9a6b-e270666df62d",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "f3c61f7e-8efa-45da-a754-d02ef29341cb",
"cell_type": "markdown",
"id": "6191bb15-81c4-4516-92e8-491ef6cc765a",
"metadata": {},
"outputs": [],
"source": []
"source": [
"## Randomised SVD"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b1e1d1c6-8bf0-4452-b62f-0bf6262fe475",
"execution_count": 14,
"id": "24a5eb2e-b718-4b14-9a6b-e270666df62d",
"metadata": {},
"outputs": [],
"source": []
"source": [
"np.random.seed(0)\n",
"\n",
"dim = 3\n",
"dtype = np.float32\n",
"ctype = np.complex64\n",
"\n",
"# Set FMM Parameters\n",
"expansion_order = np.array([6], np.uint64) # Single expansion order as using n_crit\n",
"n_vec = 1\n",
"n_crit = 150\n",
"n_sources = 1000\n",
"n_targets = 2000\n",
"wavenumber = dtype(10)\n",
"prune_empty = True\n",
"svd_threshold=1e-5\n",
"\n",
"# Setup source/target/charge data in Fortran order\n",
"sources = np.random.rand(n_sources * dim).astype(dtype)\n",
"targets = np.random.rand(n_targets * dim).astype(dtype)\n",
"charges = np.random.rand(n_sources * n_vec).astype(ctype)\n",
"\n",
"eval_type = EvalType.Value\n",
"\n",
"# EvalType computes either potentials (EvalType.Value) or potentials + derivatives (EvalType.ValueDeriv)\n",
"kernel = HelmholtzKernel(dtype, wavenumber, eval_type)\n",
"\n",
"tree = SingleNodeTree(sources, targets, charges, n_crit=n_crit, prune_empty=prune_empty)\n",
"\n",
"field_translation = BlasFieldTranslation(kernel, svd_threshold, random=True)\n",
"\n",
"# Create FMM runtime object\n",
"fmm = KiFmm(expansion_order, tree, field_translation, timed=True)\n",
"\n",
"# Evaluate FMM\n",
"fmm.evaluate()\n",
"\n",
"# Direct evaluation check\n",
"direct = np.zeros(n_targets * eval_type.value).astype(ctype)\n",
"fmm.evaluate_kernel(EvalType.Value, sources, targets, charges, direct)\n",
"\n",
"assert np.allclose(np.abs(direct), np.abs(fmm.all_potentials_u[0].T), 1e-3)"
]
},
{
"cell_type": "code",
Expand All @@ -572,7 +605,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.10.14"
"version": "3.10.15"
}
},
"nbformat": 4,
Expand Down
59 changes: 51 additions & 8 deletions kifmm/python/kifmm_py/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ def __init__(self, times_p):
def __repr__(self):
return str(self._times)


class MetadataTimes:
"""
Helper class to destructure metadata runtimes.
Expand Down Expand Up @@ -98,6 +99,7 @@ def __init__(self, times_p):
def __repr__(self):
return str(self._times)


class CommunicationTimes:
"""
Helper class to destructure communication runtimes.
Expand Down Expand Up @@ -210,11 +212,9 @@ def __init__(
self.surface_diff = surface_diff
self.random = random

if isinstance(self.kernel, HelmholtzKernel):
if self.random:
raise TypeError("Randomised compression unimplemented for this kernel")

elif isinstance(self.kernel, LaplaceKernel):
if isinstance(self.kernel, HelmholtzKernel) or isinstance(
self.kernel, LaplaceKernel
):
if self.random:
self.rsvd = RandomSvdSettings(
n_components,
Expand Down Expand Up @@ -628,9 +628,52 @@ def _construct(self):
elif isinstance(self._field_translation.kernel, HelmholtzKernel):

if self._field_translation.random:
raise TypeError(
"Randomised SVD not yet supported for Helmholtz Kernel"
)
if self._field_translation.kernel.dtype == np.float32:
self._fmm = lib.helmholtz_blas_rsvd_f32_alloc(
self._timed,
self._expansion_order_c,
self._nexpansion_order,
self._field_translation.kernel.eval_type,
self._field_translation.kernel.wavenumber,
self._tree.sources_c,
self._tree.n_sources,
self._tree.targets_c,
self._tree.n_targets,
self._tree.charges_c,
self._tree.ncharges,
self._tree.prune_empty,
self._tree.n_crit,
self._tree.depth,
self._field_translation.svd_threshold,
self._field_translation.surface_diff,
self._field_translation.rsvd.n_components,
self._field_translation.rsvd.n_oversamples,
)
self.potential_dtype = np.complex64

elif self._field_translation.kernel.dtype == np.float64:
self._fmm = lib.helmholtz_blas_rsvd_f64_alloc(
self._timed,
self._expansion_order_c,
self._nexpansion_order,
self._field_translation.kernel.eval_type,
self._field_translation.kernel.wavenumber,
self._tree.sources_c,
self._tree.n_sources,
self._tree.targets_c,
self._tree.n_targets,
self._tree.charges_c,
self._tree.ncharges,
self._tree.prune_empty,
self._tree.n_crit,
self._tree.depth,
self._field_translation.svd_threshold,
self._field_translation.surface_diff,
self._field_translation.rsvd.n_components,
self._field_translation.rsvd.n_oversamples,
)
self.potential_dtype = np.complex128

else:
if self._field_translation.kernel.dtype == np.float32:
self._fmm = lib.helmholtz_blas_svd_f32_alloc(
Expand Down
Loading

0 comments on commit e7eb5ba

Please sign in to comment.