diff --git a/sumpy/expansion/__init__.py b/sumpy/expansion/__init__.py
index 647ebb01589a502c128a30aa597377c8fed1d0e4..f9ffc4f54dfd5427e8d47088872ce2a967a62032 100644
--- a/sumpy/expansion/__init__.py
+++ b/sumpy/expansion/__init__.py
@@ -29,6 +29,7 @@ import logging
from pytools import memoize_method
import sumpy.symbolic as sym
from collections import defaultdict
+from sumpy.tools import CoeffIdentifier, add_mi
__doc__ = """
@@ -364,24 +365,23 @@ class LinearRecurrenceBasedExpansionTermsWrangler(ExpansionTermsWrangler):
from pytools import ProcessLogger
plog = ProcessLogger(logger, "compute recurrence for Taylor coefficients")
- pde_dict = self.get_pde_dict()
+ pdes, iexpr, nexpr = self.get_pdes()
pde_mat = []
mis = self.get_full_coefficient_identifiers()
coeff_ident_enumerate_dict = dict((tuple(mi), i) for
(i, mi) in enumerate(mis))
+ offset = len(mis)
for mi in mis:
- for pde_mi, coeff in iteritems(pde_dict):
- diff = np.array(mi, dtype=int) - pde_mi
- if (diff >= 0).all():
- eq = [0]*len(mis)
- for pde_mi2, coeff2 in iteritems(pde_dict):
- c = tuple(pde_mi2 + diff)
- if c not in coeff_ident_enumerate_dict:
- break
- eq[coeff_ident_enumerate_dict[c]] = coeff
- else:
- pde_mat.append(eq)
+ for pde_dict in pdes:
+ eq = [0]*(len(mis)*nexpr)
+ for ident, coeff in iteritems(pde_dict):
+ c = tuple(add_mi(ident.mi, mi))
+ if c not in coeff_ident_enumerate_dict:
+ break
+ eq[offset*ident.iexpr + coeff_ident_enumerate_dict[c]] = coeff
+ else:
+ pde_mat.append(eq)
if len(pde_mat) > 0:
r"""
@@ -398,8 +398,9 @@ class LinearRecurrenceBasedExpansionTermsWrangler(ExpansionTermsWrangler):
:math:`S = N_{[r, :]}^{-T} N^T = N_{[r, :]}^{-T} N^T`.
"""
n = nullspace(pde_mat)
+ n = n[offset*iexpr:offset*(iexpr+1), :]
idx = self.get_reduced_coeffs(n)
- assert len(idx) >= n.shape[1]
+ n = n[:,:len(idx)]
s = solve_symbolic(n.T[:, idx], n.T)
stored_identifiers = [mis[i] for i in idx]
else:
@@ -422,10 +423,12 @@ class LinearRecurrenceBasedExpansionTermsWrangler(ExpansionTermsWrangler):
return stored_identifiers, coeff_matrix
- def get_pde_dict(self):
+ def get_pdes(self):
r"""
- Returns the PDE as a dictionary of (mi, coeff) such that mi
- is the multi-index of the derivative and the PDE is represented by,
+ Returns a list of PDEs, expression number, number of expressions.
+ A PDE is a dictionary of (ident, coeff) such that ident is a
+ namedtuple of (mi, iexpr) where mi is the multi-index of the
+ derivative, iexpr is the expression number and the PDE is represented by,
.. math::
@@ -487,15 +490,16 @@ class LaplaceExpansionTermsWrangler(LinearRecurrenceBasedExpansionTermsWrangler)
init_arg_names = ("order", "dim", "max_mi")
def __init__(self, order, dim, max_mi=None):
- super(LaplaceExpansionTermsWrangler, self).__init__(order, dim, 1, max_mi)
+ super(LaplaceExpansionTermsWrangler, self).__init__(order=order, dim=dim,
+ deriv_multiplier=1, max_mi=max_mi)
- def get_pde_dict(self):
+ def get_pdes(self):
pde_dict = {}
for d in range(self.dim):
mi = [0]*self.dim
mi[d] = 2
- pde_dict[tuple(mi)] = 1
- return pde_dict
+ pde_dict[CoeffIdentifier(tuple(mi), 0)] = 1
+ return [pde_dict], 0, 1
def _get_reduced_coeffs(self, nullspace):
idx = []
@@ -515,17 +519,17 @@ class HelmholtzExpansionTermsWrangler(LinearRecurrenceBasedExpansionTermsWrangle
self.helmholtz_k_name = helmholtz_k_name
multiplier = sym.Symbol(helmholtz_k_name)
- super(HelmholtzExpansionTermsWrangler, self).__init__(order, dim, multiplier,
- max_mi)
+ super(HelmholtzExpansionTermsWrangler, self).__init__(order=order, dim=dim,
+ deriv_multiplier=multiplier, max_mi=max_mi)
- def get_pde_dict(self, **kwargs):
+ def get_pdes(self, **kwargs):
pde_dict = {}
for d in range(self.dim):
mi = [0]*self.dim
mi[d] = 2
- pde_dict[tuple(mi)] = 1
- pde_dict[tuple([0]*self.dim)] = 1
- return pde_dict
+ pde_dict[CoeffIdentifier(tuple(mi), 0)] = 1
+ pde_dict[CoeffIdentifier(tuple([0]*self.dim), 0)] = 1
+ return [pde_dict], 0, 1
def _get_reduced_coeffs(self, nullspace):
idx = []
@@ -536,6 +540,46 @@ class HelmholtzExpansionTermsWrangler(LinearRecurrenceBasedExpansionTermsWrangle
idx.append(i)
return idx
+
+class StokesExpansionTermsWrangler(LinearRecurrenceBasedExpansionTermsWrangler):
+
+ init_arg_names = ("order", "dim", "icomp", "viscosity_mu_name", "max_mi")
+
+ def __init__(self, order, dim, icomp, viscosity_mu_name, max_mi=None):
+ self.viscosity_mu_name = viscosity_mu_name
+ self.icomp = icomp
+ multiplier = 1/sym.Symbol(viscosity_mu_name)
+ super(StokesExpansionTermsWrangler, self).__init__(order=order, dim=dim,
+ deriv_multiplier=multiplier, max_mi=max_mi)
+
+ def get_pdes(self, **kwargs):
+ pdes = []
+ # velocity
+ for eq in range(self.dim):
+ pde_dict = {}
+ for d in range(self.dim):
+ mi = [0]*self.dim
+ mi[d] = 2
+ pde_dict[CoeffIdentifier(tuple(mi), eq)] = 1
+ mi = [0]*self.dim
+ mi[eq] = 1
+ pde_dict[CoeffIdentifier(tuple(mi), self.dim)] = -1
+ pdes.append(pde_dict)
+ # pressure
+ pde_dict = {}
+ for d in range(self.dim):
+ mi = [0]*self.dim
+ mi[d] = 2
+ pde_dict[CoeffIdentifier(tuple(mi), self.dim)] = 1
+ pdes.append(pde_dict)
+ # divergence
+ pde_dict = {}
+ for d in range(self.dim):
+ mi = [0]*self.dim
+ mi[d] = 1
+ pde_dict[CoeffIdentifier(tuple(mi), d)] = 1
+ pdes.append(pde_dict)
+ return pdes, self.icomp, self.dim+1
# }}}
@@ -612,6 +656,19 @@ class HelmholtzConformingVolumeTaylorExpansion(VolumeTaylorExpansionBase):
helmholtz_k_name = kernel.get_base_kernel().helmholtz_k_name
self.expansion_terms_wrangler_key = (order, kernel.dim, helmholtz_k_name)
+
+class StokesConformingVolumeTaylorExpansion(VolumeTaylorExpansionBase):
+
+ expansion_terms_wrangler_class = StokesExpansionTermsWrangler
+ expansion_terms_wrangler_cache = {}
+
+ # not user-facing, be strict about having to pass use_rscale
+ def __init__(self, kernel, order, use_rscale):
+ icomp = kernel.get_base_kernel().icomp
+ viscosity_mu_name = kernel.get_base_kernel().viscosity_mu_name
+ self.expansion_terms_wrangler_key = (order, kernel.dim,
+ icomp, viscosity_mu_name)
+
# }}}
diff --git a/sumpy/expansion/local.py b/sumpy/expansion/local.py
index cf94ebd674201178047cfee86acca0d74512c1ec..f080b7460029ce2f5325e0ed58f3d97611aa791c 100644
--- a/sumpy/expansion/local.py
+++ b/sumpy/expansion/local.py
@@ -27,7 +27,7 @@ import sumpy.symbolic as sym
from sumpy.expansion import (
ExpansionBase, VolumeTaylorExpansion, LaplaceConformingVolumeTaylorExpansion,
- HelmholtzConformingVolumeTaylorExpansion)
+ HelmholtzConformingVolumeTaylorExpansion, StokesConformingVolumeTaylorExpansion)
class LocalExpansionBase(ExpansionBase):
@@ -269,6 +269,16 @@ class HelmholtzConformingVolumeTaylorLocalExpansion(
HelmholtzConformingVolumeTaylorExpansion.__init__(
self, kernel, order, use_rscale)
+
+class StokesConformingVolumeTaylorLocalExpansion(
+ StokesConformingVolumeTaylorExpansion,
+ VolumeTaylorLocalExpansionBase):
+
+ def __init__(self, kernel, order, use_rscale=None):
+ VolumeTaylorLocalExpansionBase.__init__(self, kernel, order, use_rscale)
+ StokesConformingVolumeTaylorExpansion.__init__(
+ self, kernel, order, use_rscale)
+
# }}}
diff --git a/sumpy/expansion/multipole.py b/sumpy/expansion/multipole.py
index 25488952e4b7425371ecf23aefa261c0c70900f5..07f2eeb6929c47b0437e954819ebc449fc7c0c27 100644
--- a/sumpy/expansion/multipole.py
+++ b/sumpy/expansion/multipole.py
@@ -28,7 +28,7 @@ import sumpy.symbolic as sym # noqa
from sumpy.symbolic import vector_xreplace
from sumpy.expansion import (
ExpansionBase, VolumeTaylorExpansion, LaplaceConformingVolumeTaylorExpansion,
- HelmholtzConformingVolumeTaylorExpansion)
+ HelmholtzConformingVolumeTaylorExpansion, StokesConformingVolumeTaylorExpansion)
import logging
logger = logging.getLogger(__name__)
@@ -220,6 +220,16 @@ class HelmholtzConformingVolumeTaylorMultipoleExpansion(
HelmholtzConformingVolumeTaylorExpansion.__init__(
self, kernel, order, use_rscale)
+
+class StokesConformingVolumeTaylorMultipoleExpansion(
+ StokesConformingVolumeTaylorExpansion,
+ VolumeTaylorMultipoleExpansionBase):
+
+ def __init__(self, kernel, order, use_rscale=None):
+ VolumeTaylorMultipoleExpansionBase.__init__(self, kernel, order, use_rscale)
+ StokesConformingVolumeTaylorExpansion.__init__(
+ self, kernel, order, use_rscale)
+
# }}}
diff --git a/sumpy/tools.py b/sumpy/tools.py
index 32d14bcbfcf09f48322ff285d37ba5461b12560d..c13a04b6a73f858650e4235c9cb753d6b863c119 100644
--- a/sumpy/tools.py
+++ b/sumpy/tools.py
@@ -27,6 +27,7 @@ THE SOFTWARE.
import six
from six.moves import range, zip
+from collections import namedtuple
from pytools import memoize_method, memoize_in
import numpy as np
import sumpy.symbolic as sym
@@ -729,4 +730,7 @@ def solve_symbolic(a, b):
red = rref(big)[0]
return red[:, big.shape[0]:]
+
+CoeffIdentifier = namedtuple('CoeffIdentifier', ['mi', 'iexpr'])
+
# vim: fdm=marker
diff --git a/test/test_kernels.py b/test/test_kernels.py
index 3fcb195bba0be9cb0b67a4822a45dec616d370aa..6a96d1ad8f44133b8b4bff7ba87078a5cded528a 100644
--- a/test/test_kernels.py
+++ b/test/test_kernels.py
@@ -37,13 +37,15 @@ from sumpy.expansion.multipole import (
VolumeTaylorMultipoleExpansion, H2DMultipoleExpansion,
VolumeTaylorMultipoleExpansionBase,
LaplaceConformingVolumeTaylorMultipoleExpansion,
- HelmholtzConformingVolumeTaylorMultipoleExpansion)
+ HelmholtzConformingVolumeTaylorMultipoleExpansion,
+ StokesConformingVolumeTaylorMultipoleExpansion)
from sumpy.expansion.local import (
VolumeTaylorLocalExpansion, H2DLocalExpansion,
LaplaceConformingVolumeTaylorLocalExpansion,
- HelmholtzConformingVolumeTaylorLocalExpansion)
+ HelmholtzConformingVolumeTaylorLocalExpansion,
+ StokesConformingVolumeTaylorLocalExpansion)
from sumpy.kernel import (LaplaceKernel, HelmholtzKernel, AxisTargetDerivative,
- DirectionalSourceDerivative)
+ DirectionalSourceDerivative, StokesletKernel)
from pytools.convergence import PConvergenceVerifier
import logging
@@ -155,6 +157,8 @@ def test_p2e2p(ctx_getter, base_knl, expn_class, order, with_source_derivative):
extra_kwargs["k"] = 0.2 * (0.707 + 0.707j)
else:
extra_kwargs["k"] = 0.2
+ if isinstance(base_knl, StokesletKernel):
+ extra_kwargs["mu"] = 0.2
if with_source_derivative:
knl = DirectionalSourceDerivative(base_knl, "dir_vec")
@@ -342,7 +346,10 @@ def test_p2e2p(ctx_getter, base_knl, expn_class, order, with_source_derivative):
(HelmholtzKernel(2), VolumeTaylorLocalExpansion, VolumeTaylorMultipoleExpansion),
(HelmholtzKernel(2), HelmholtzConformingVolumeTaylorLocalExpansion,
HelmholtzConformingVolumeTaylorMultipoleExpansion),
- (HelmholtzKernel(2), H2DLocalExpansion, H2DMultipoleExpansion)
+ (HelmholtzKernel(2), H2DLocalExpansion, H2DMultipoleExpansion),
+ (StokesletKernel(2, 0, 0), VolumeTaylorLocalExpansion, VolumeTaylorMultipoleExpansion),
+ (StokesletKernel(2, 0, 0), StokesConformingVolumeTaylorLocalExpansion,
+ StokesConformingVolumeTaylorMultipoleExpansion),
])
def test_translations(ctx_getter, knl, local_expn_class, mpole_expn_class):
logging.basicConfig(level=logging.INFO)
@@ -363,6 +370,8 @@ def test_translations(ctx_getter, knl, local_expn_class, mpole_expn_class):
extra_kwargs = {}
if isinstance(knl, HelmholtzKernel):
extra_kwargs["k"] = 0.05
+ if isinstance(knl, StokesletKernel):
+ extra_kwargs["mu"] = 0.05
# Just to make sure things also work away from the origin
origin = np.array([2, 1], np.float64)
@@ -404,7 +413,7 @@ def test_translations(ctx_getter, knl, local_expn_class, mpole_expn_class):
# FIXME: Embarrassing--but we run out of memory for higher orders.
orders = [2, 3]
else:
- orders = [2, 3, 4]
+ orders = [3, 4, 5]
nboxes = centers.shape[-1]
def eval_at(e2p, source_box_nr, rscale):