From ea6c161de5caa917c77cf73ae28dffceff0d064b Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Tue, 5 Jul 2011 21:11:23 -0400
Subject: [PATCH] Add M2M and M2L computer algebra (untested), plus rudiments
of a test.
---
sumpy/expansion.py | 93 +++++++++++++++++++++++++++++-------
test/test_kernels.py | 109 ++++++++++++++++++++++++++++++++++++++++++-
2 files changed, 184 insertions(+), 18 deletions(-)
diff --git a/sumpy/expansion.py b/sumpy/expansion.py
index 5528c2f0..97de6c22 100644
--- a/sumpy/expansion.py
+++ b/sumpy/expansion.py
@@ -26,7 +26,7 @@ class Expansion(object):
-class TaylorExpansion(Expansion):
+class TaylorMultipoleExpansion(Expansion):
def __init__(self, kernel, order, dimensions):
"""
:arg: kernel in terms of 'b' variable
@@ -41,36 +41,97 @@ class TaylorExpansion(Expansion):
as gnitstam)
self.multi_indices = sorted(gnitstam(self.order, self.dimensions), key=sum)
+ self.mi_to_index = dict(
+ (mi, i) for i, mi in enumerate(self.multi_indices))
# given in terms of b variable
+ from sumpy.symbolic import diff_multi_index
self.basis = [
- self.diff_kernel(mi)
+ diff_multi_index(kernel, mi, "b")
for mi in self.multi_indices]
# given in terms of a variable
from sumpy.symbolic import make_sym_vector
- a = make_sym_vector("a", 3)
+ a = make_sym_vector("a", dimensions)
from sumpy.tools import mi_power, mi_factorial
self.coefficients = [
mi_power(a, mi)/mi_factorial(mi)
for mi in self.multi_indices]
- @memoize_method
- def diff_kernel(self, multi_index):
- if sum(multi_index) == 0:
- return self.kernel
+ def m2m_exprs(self, get_coefficient_expr):
+ """Expressions for coefficients of shifted expansion, in terms of s
+ (the shift from the old center to the new center), as well as the
+ coefficients returned by *get_coefficient_expr* for each multi-index.
+ """
+
+ from sumpy.symbolic import make_sym_vector
+ a = make_sym_vector("a", self.dimensions)
+ s = make_sym_vector("s", self.dimensions)
+ # new center = c+s
+
+ from sumpy.symbolic import IdentityMapper, find_power_of
+ class ToCoefficientMapper(IdentityMapper):
+ def map_Mul(subself, expr):
+ a_powers = tuple(int(find_power_of(ai, expr)) for ai in a)
+
+ return (
+ expr/mi_power(a, a_powers)
+ * get_coefficient_expr(self.mi_to_index[a_powers]))
+
+ map_Symbol = map_Mul
+
+ tcm = ToCoefficientMapper()
+
+ from sumpy.tools import mi_power
+ return [tcm(mi_power(a+s, mi).expand()) for mi in self.multi_indices]
+
+ def m2l_exprs(self, loc_exp, get_coefficient_expr):
+ """Expressions for coefficients of the local expansion *loc_exp* of the
+ multipole expansion *self*, whose coefficients are obtained by
+ *get_coefficient_expr* for each coefficient index. The expressions are
+ given in terms of *s* (the shift from the multipole center to the local
+ center).
+ """
+
+ b2s_map = dict(
+ (sp.Symbol("b%d" % i), sp.Symbol("s%d" % i))
+ for i in range(self.dimensions))
+
+ expansion = sum(
+ get_coefficient_expr(i)
+ * basis_func.subs(b2s_map)
+ for i, basis_func in enumerate(self.basis))
+
+ from sumpy.symbolic import diff_multi_index
+ from sumpy.tools import mi_factorial
+ return [diff_multi_index(expansion, loc_mi, "s") / mi_factorial(loc_mi)
+ for loc_mi in loc_exp.multi_indices]
+
+
- first_nonzero_axis = min(
- i for i in range(self.dimensions)
- if multi_index[i] > 0)
- lowered_mi = list(multi_index)
- lowered_mi[first_nonzero_axis] -= 1
- lowered_mi = tuple(lowered_mi)
- lower_diff_kernel = self.diff_kernel(lowered_mi)
- return sp.diff(lower_diff_kernel,
- sp.Symbol("b%d" % first_nonzero_axis))
+class TaylorLocalExpansion(Expansion):
+ def __init__(self, order, dimensions):
+ self.order = order
+ self.dimensions = dimensions
+
+ from pytools import (
+ generate_nonnegative_integer_tuples_summing_to_at_most
+ as gnitstam)
+
+ self.multi_indices = sorted(gnitstam(self.order, self.dimensions), key=sum)
+ self.mi_to_index = dict(
+ (mi, i) for i, mi in enumerate(self.multi_indices))
+ from sumpy.symbolic import make_sym_vector
+
+ from sumpy.tools import mi_power, mi_factorial
+
+ # given in terms of b variable
+ b = make_sym_vector("b", dimensions)
+ self.basis = [
+ mi_power(b, mi)/mi_factorial(mi)
+ for mi in self.multi_indices]
diff --git a/test/test_kernels.py b/test/test_kernels.py
index 755b1968..eb466aea 100644
--- a/test/test_kernels.py
+++ b/test/test_kernels.py
@@ -99,8 +99,8 @@ def test_p2m2p(ctx_getter):
cell_idx_to_particle_cnt_tgt_dev = to_device(queue, cell_idx_to_particle_cnt_tgt)
from sumpy.symbolic import make_coulomb_kernel_in
- from sumpy.expansion import TaylorExpansion
- texp = TaylorExpansion(
+ from sumpy.expansion import TaylorMultipoleExpansion
+ texp = TaylorMultipoleExpansion(
make_coulomb_kernel_in("b", dimensions),
order=3, dimensions=dimensions)
@@ -161,6 +161,111 @@ def test_p2m2p(ctx_getter):
+@pytools.test.mark_test.opencl
+def test_p2m2m2p(ctx_getter):
+ ctx = ctx_getter()
+ queue = cl.CommandQueue(ctx)
+
+ res = 100
+
+ dimensions = 3
+ sources = np.random.rand(dimensions, 5).astype(np.float32)
+ #targets = np.random.rand(dimensions, 10).astype(np.float32)
+ targets = np.mgrid[-2:2:res*1j, -2:2:res*1j, 2:2:1j].reshape(3, -1).astype(np.float32)
+ centers = np.array([[0.5]]*dimensions).astype(np.float32)
+
+ from sumpy.tools import vector_to_device
+ targets_dev = vector_to_device(queue, targets)
+ sources_dev = vector_to_device(queue, sources)
+ centers_dev = vector_to_device(queue, centers)
+ strengths_dev = cl_array.empty(queue, sources.shape[1], dtype=np.float32)
+ strengths_dev.fill(1)
+
+ cell_idx_to_particle_offset = np.array([0], dtype=np.uint32)
+ cell_idx_to_particle_cnt_src = np.array([sources.shape[1]], dtype=np.uint32)
+ cell_idx_to_particle_cnt_tgt = np.array([targets.shape[1]], dtype=np.uint32)
+
+ from pyopencl.array import to_device
+ cell_idx_to_particle_offset_dev = to_device(queue, cell_idx_to_particle_offset)
+ cell_idx_to_particle_cnt_src_dev = to_device(queue, cell_idx_to_particle_cnt_src)
+ cell_idx_to_particle_cnt_tgt_dev = to_device(queue, cell_idx_to_particle_cnt_tgt)
+
+ from sumpy.symbolic import make_coulomb_kernel_in
+ from sumpy.expansion import TaylorMultipoleExpansion
+ texp = TaylorMultipoleExpansion(
+ make_coulomb_kernel_in("b", dimensions),
+ order=3, dimensions=dimensions)
+
+ def get_coefficient_expr(idx):
+ return sp.Symbol("coeff%d" % idx)
+
+ for i in texp.m2m_exprs(get_coefficient_expr):
+ print i
+
+ print "-------------------------------------"
+ from sumpy.expansion import TaylorLocalExpansion
+ locexp = TaylorLocalExpansion(order=2, dimensions=3)
+
+ for i in texp.m2l_exprs(locexp, get_coefficient_expr):
+ print i
+
+ 1/0
+
+ coeff_dtype = np.float32
+
+ # {{{ apply P2M
+
+ from sumpy.p2m import P2MKernel
+ p2m = P2MKernel(ctx, texp)
+ mpole_coeff = p2m(cell_idx_to_particle_offset_dev,
+ cell_idx_to_particle_cnt_src_dev,
+ sources_dev, strengths_dev, centers_dev, coeff_dtype)
+
+ # }}}
+
+ # {{{ apply M2P
+
+ output_maps = [
+ lambda expr: expr,
+ lambda expr: sp.diff(expr, sp.Symbol("t0"))
+ ]
+
+ m2p_ilist_starts = np.array([0, 1], dtype=np.uint32)
+ m2p_ilist_mpole_offsets = np.array([0], dtype=np.uint32)
+
+ m2p_ilist_starts_dev = to_device(queue, m2p_ilist_starts)
+ m2p_ilist_mpole_offsets_dev = to_device(queue, m2p_ilist_mpole_offsets)
+
+ from sumpy.m2p import M2PKernel
+ m2p = M2PKernel(ctx, texp, output_maps=output_maps)
+ potential_dev, x_derivative = m2p(targets_dev, m2p_ilist_starts_dev, m2p_ilist_mpole_offsets_dev, mpole_coeff,
+ cell_idx_to_particle_offset_dev,
+ cell_idx_to_particle_cnt_tgt_dev)
+
+ # }}}
+
+ # {{{ compute (direct) reference solution
+
+ from sumpy.p2p import P2PKernel
+ from sumpy.symbolic import make_coulomb_kernel_ts
+ coulomb_knl = make_coulomb_kernel_ts(dimensions)
+
+ knl = P2PKernel(ctx, dimensions,
+ exprs=[f(coulomb_knl) for f in output_maps], exclude_self=False)
+
+ potential_dev_direct, x_derivative_dir = knl(targets_dev, sources_dev, strengths_dev)
+
+ if 0:
+ import matplotlib.pyplot as pt
+ pt.imshow((potential_dev-potential_dev_direct).get().reshape(res, res))
+ pt.colorbar()
+ pt.show()
+
+ assert la.norm((potential_dev-potential_dev_direct).get())/res**2 < 1e-3
+
+
+
+
# You can test individual routines by typing
# $ python test_kernels.py 'test_p2p(cl.create_some_context)'
--
GitLab