diff --git a/sumpy/expansion/__init__.py b/sumpy/expansion/__init__.py index 995f522d5603d5daaf2c3927901e89626d985170..cc49c6ca2e712ecbf43a0a6d460d0e0dd2667152 100644 --- a/sumpy/expansion/__init__.py +++ b/sumpy/expansion/__init__.py @@ -1,5 +1,6 @@ from __future__ import division from __future__ import absolute_import +from six.moves import range __copyright__ = "Copyright (C) 2012 Andreas Kloeckner" @@ -25,12 +26,16 @@ THE SOFTWARE. import numpy as np import sympy as sp +import logging from pytools import memoize_method +from sumpy.tools import MiDerivativeTaker __doc__ = """ .. autoclass:: ExpansionBase """ +logger = logging.getLogger(__name__) + # {{{ base class @@ -112,24 +117,25 @@ class ExpansionBase(object): # }}} -# {{{ volume taylor +# {{{ derivative wrangler -class VolumeTaylorExpansionBase(object): +class DerivativeWrangler(object): + + def __init__(self, order, dim): + self.order = order + self.dim = dim def get_coefficient_identifiers(self): - """ - Returns the identifiers of the coefficients that actually get stored. - """ raise NotImplementedError - @property - @memoize_method - def _storage_loc_dict(self): - return dict((i, idx) for idx, i in - enumerate(self.get_coefficient_identifiers())) + def get_full_kernel_derivatives_from_stored(self, stored_kernel_derivatives): + raise NotImplementedError - def get_storage_index(self, i): - return self._storage_loc_dict[i] + def get_stored_mpole_coefficients_from_full(self, full_mpole_coefficients): + raise NotImplementedError + + def get_derivative_taker(self, expr, var_list): + raise NotImplementedError @memoize_method def get_full_coefficient_identifiers(self): @@ -140,75 +146,162 @@ class VolumeTaylorExpansionBase(object): generate_nonnegative_integer_tuples_summing_to_at_most as gnitstam) - return sorted(gnitstam(self.order, self.kernel.dim), key=sum) - - def stored_to_full(self, coeff_idx, stored_coeffs): - raise NotImplementedError + res = sorted(gnitstam(self.order, self.dim), key=sum) + return res - def full_to_stored(self, coeff_idx, full_coeffs): - raise NotImplementedError +class FullDerivativeWrangler(DerivativeWrangler): -class VolumeTaylorExpansion(VolumeTaylorExpansionBase): + def get_derivative_taker(self, expr, dvec): + return MiDerivativeTaker(expr, dvec) get_coefficient_identifiers = ( - VolumeTaylorExpansionBase.get_full_coefficient_identifiers) + DerivativeWrangler.get_full_coefficient_identifiers) + + def get_full_kernel_derivatives_from_stored(self, stored_kernel_derivatives): + return stored_kernel_derivatives - def stored_to_full(self, stored_coeffs): - return stored_coeffs + get_stored_mpole_coefficients_from_full = ( + get_full_kernel_derivatives_from_stored) - full_to_stored = stored_to_full +# {{{ sparse matrix-vector multiplication -class LinearRecurrenceBasedVolumeTaylorExpansion(VolumeTaylorExpansionBase): +def _spmv(spmat, x, sparse_vectors): + """ + :param spmat: maps row indices to list of (col idx, value) + :param x: maps vector indices to values + :param sparse_vectors: If True, treat vectors as dict-like, otherwise list-like + """ + if sparse_vectors: + result = {} + else: + result = [] - def __init__(self): - self.precompute_coeff_matrix() + for row in range(len(spmat)): + acc = 0 + + for j, coeff in spmat[row]: + if sparse_vectors: + # Check if the index exists in the vector. + if x.get(j, 0) == 0: + continue + + acc += coeff * x[j] + + if sparse_vectors: + if acc != 0: + result[row] = acc + else: + result.append(acc) + + return result + +# }}} + + +class LinearRecurrenceBasedDerivativeWrangler(DerivativeWrangler): + + def __init__(self, order, dim): + DerivativeWrangler.__init__(self, order, dim) def get_coefficient_identifiers(self): return self.stored_identifiers - def stored_to_full(self, stored_coeffs): - return self.coeff_matrix.dot(stored_coeffs) + def get_full_kernel_derivatives_from_stored(self, stored_kernel_derivatives): + return _spmv(self.coeff_matrix, stored_kernel_derivatives, + sparse_vectors=False) + + def get_stored_mpole_coefficients_from_full(self, full_mpole_coefficients): + # = M^T x, where M = coeff matrix + result = [0] * len(self.stored_identifiers) + for row, coeff in enumerate(full_mpole_coefficients): + for col, val in self.coeff_matrix[row]: + result[col] += coeff * val + return result - def full_to_stored(self, full_coeffs): - return self.coeff_matrix.T.dot(full_coeffs) + def precompute_recurrences(self): + """ + Build up a matrix that expresses every derivative in terms of a + set of "stored" derivatives. - def precompute_coeff_matrix(self): + For example, for the recurrence + + u_xx + u_yy + u_zz = 0 + + the coefficient matrix features the following entries: + + ... u_xx u_yy ... <= cols = only stored derivatives + ================== + ...| ... ... ... ... + | + u_zz| ... -1 -1 ... + + ^ rows = one for every derivative + """ stored_identifiers = [] - identifiers_so_far = [] + identifiers_so_far = {} + + import time + start_time = time.time() + logger.debug("computing recurrence for Taylor coefficients: start") - ncoeffs = len(self.get_full_coefficient_identifiers()) - coeff_matrix = [] + # Sparse matrix, indexed by row + from collections import defaultdict + coeff_matrix_transpose = defaultdict(lambda: []) - # Build up the matrix by row. - for identifier in self.get_full_coefficient_identifiers(): + # Build up the matrix transpose by row. + from six import iteritems + for i, identifier in enumerate(self.get_full_coefficient_identifiers()): expr = self.try_get_recurrence_for_derivative( - identifier, identifiers_so_far, ncoeffs) + identifier, identifiers_so_far) if expr is None: # Identifier should be stored - row = np.zeros(ncoeffs, dtype=object) - row[len(stored_identifiers)] = 1 + coeff_matrix_transpose[len(stored_identifiers)] = [(i, 1)] stored_identifiers.append(identifier) else: - # Rewrite in terms of the stored identifiers - ncoeffs_so_far = len(coeff_matrix) - row = np.dot(np.transpose(coeff_matrix), expr[:ncoeffs_so_far]) + expr = dict((identifiers_so_far[ident], val) for ident, val in + iteritems(expr)) + result = _spmv(coeff_matrix_transpose, expr, sparse_vectors=True) + for j, item in iteritems(result): + coeff_matrix_transpose[j].append((i, item)) - coeff_matrix.append(row) - identifiers_so_far.append(identifier) + identifiers_so_far[identifier] = i self.stored_identifiers = stored_identifiers - ncols = len(stored_identifiers) - self.coeff_matrix = np.vstack(coeff_matrix)[:, :ncols] + coeff_matrix = defaultdict(lambda: []) + for i, row in iteritems(coeff_matrix_transpose): + for j, val in row: + coeff_matrix[j].append((i, val)) + + logger.debug("computing recurrence for Taylor coefficients: " + "done after {dur:.2f} seconds" + .format(dur=time.time() - start_time)) + + logger.debug("number of Taylor coefficients was reduced from {orig} to {red}" + .format(orig=len(self.get_full_coefficient_identifiers()), + red=len(self.stored_identifiers))) + + self.coeff_matrix = coeff_matrix + + def try_get_recurrence_for_derivative(self, deriv, in_terms_of): + raise NotImplementedError + + def get_derivative_taker(self, expr, var_list): + from sumpy.tools import LinearRecurrenceBasedMiDerivativeTaker + return LinearRecurrenceBasedMiDerivativeTaker(expr, var_list, self) + + +class LaplaceDerivativeWrangler(LinearRecurrenceBasedDerivativeWrangler): -class LaplaceConformingVolumeTaylorExpansion( - LinearRecurrenceBasedVolumeTaylorExpansion): + def __init__(self, order, dim): + LinearRecurrenceBasedDerivativeWrangler.__init__(self, order, dim) + self.precompute_recurrences() - def try_get_recurrence_for_derivative(self, deriv, in_terms_of, ncoeffs): - deriv = np.array(deriv) + def try_get_recurrence_for_derivative(self, deriv, in_terms_of): + deriv = np.array(deriv, dtype=int) for dim in np.nonzero(2 <= deriv)[0]: # Check if we can reduce this dimension in terms of the other @@ -216,32 +309,32 @@ class LaplaceConformingVolumeTaylorExpansion( reduced_deriv = deriv.copy() reduced_deriv[dim] -= 2 + coeffs = {} - needed_derivs = [] - for other_dim in range(self.kernel.dim): + for other_dim in range(self.dim): if other_dim == dim: continue needed_deriv = reduced_deriv.copy() needed_deriv[other_dim] += 2 + needed_deriv = tuple(needed_deriv) - needed_derivs.append(tuple(needed_deriv)) + if needed_deriv not in in_terms_of: + break - expr = np.zeros(ncoeffs, dtype=object) - try: - for needed_deriv in needed_derivs: - deriv_idx = in_terms_of.index(needed_deriv) - expr[deriv_idx] = -1 - except ValueError: - continue + coeffs[needed_deriv] = -1 + else: + return coeffs - return expr +class HelmholtzDerivativeWrangler(LinearRecurrenceBasedDerivativeWrangler): -class HelmholtzConformingVolumeTaylorExpansion( - LinearRecurrenceBasedVolumeTaylorExpansion): + def __init__(self, order, dim, helmholtz_k_name): + LinearRecurrenceBasedDerivativeWrangler.__init__(self, order, dim) + self.helmholtz_k_name = helmholtz_k_name + self.precompute_recurrences() - def try_get_recurrence_for_derivative(self, deriv, in_terms_of, ncoeffs): - deriv = np.array(deriv) + def try_get_recurrence_for_derivative(self, deriv, in_terms_of): + deriv = np.array(deriv, dtype=int) for dim in np.nonzero(2 <= deriv)[0]: # Check if we can reduce this dimension in terms of the other @@ -249,29 +342,94 @@ class HelmholtzConformingVolumeTaylorExpansion( reduced_deriv = deriv.copy() reduced_deriv[dim] -= 2 + coeffs = {} - needed_derivs = [] - for other_dim in range(self.kernel.dim): + for other_dim in range(self.dim): if other_dim == dim: continue needed_deriv = reduced_deriv.copy() needed_deriv[other_dim] += 2 + needed_deriv = tuple(needed_deriv) + + if needed_deriv not in in_terms_of: + break + + coeffs[needed_deriv] = -1 + else: + k = sp.Symbol(self.helmholtz_k_name) + coeffs[tuple(reduced_deriv)] = -k*k + return coeffs + +# }}} + + +# {{{ volume taylor - needed_derivs.append((-1, tuple(needed_deriv))) +class VolumeTaylorExpansionBase(object): + + @classmethod + def get_or_make_derivative_wrangler(cls, *key): + """ + This stores the derivative wrangler at the class attribute level because + precomputing the derivative wrangler may become expensive. + """ + try: + wrangler = cls.derivative_wrangler_cache[key] + except KeyError: + wrangler = cls.derivative_wrangler_class(*key) + cls.derivative_wrangler_cache[key] = wrangler + + return wrangler - k = sp.Symbol(self.kernel.get_base_kernel().helmholtz_k_name) + @property + def derivative_wrangler(self): + return self.get_or_make_derivative_wrangler(*self.derivative_wrangler_key) - needed_derivs.append((-k*k, tuple(reduced_deriv))) + def get_coefficient_identifiers(self): + """ + Returns the identifiers of the coefficients that actually get stored. + """ + return self.derivative_wrangler.get_coefficient_identifiers() - expr = np.zeros(ncoeffs, dtype=object) - try: - for coeff, needed_deriv in needed_derivs: - deriv_idx = in_terms_of.index(needed_deriv) - expr[deriv_idx] = coeff - except ValueError: - continue + def get_full_coefficient_identifiers(self): + return self.derivative_wrangler.get_full_coefficient_identifiers() - return expr + @property + @memoize_method + def _storage_loc_dict(self): + return dict((i, idx) for idx, i in + enumerate(self.get_coefficient_identifiers())) + + def get_storage_index(self, i): + return self._storage_loc_dict[i] + + +class VolumeTaylorExpansion(VolumeTaylorExpansionBase): + + derivative_wrangler_class = FullDerivativeWrangler + derivative_wrangler_cache = {} + + def __init__(self, kernel, order): + self.derivative_wrangler_key = (order, kernel.dim) + + +class LaplaceConformingVolumeTaylorExpansion(VolumeTaylorExpansionBase): + + derivative_wrangler_class = LaplaceDerivativeWrangler + derivative_wrangler_cache = {} + + def __init__(self, kernel, order): + self.derivative_wrangler_key = (order, kernel.dim) + + +class HelmholtzConformingVolumeTaylorExpansion(VolumeTaylorExpansionBase): + + derivative_wrangler_class = HelmholtzDerivativeWrangler + derivative_wrangler_cache = {} + + def __init__(self, kernel, order): + helmholtz_k_name = kernel.get_base_kernel().helmholtz_k_name + self.derivative_wrangler_key = (order, kernel.dim, helmholtz_k_name) # }}} diff --git a/sumpy/expansion/local.py b/sumpy/expansion/local.py index c20eaf20fe1223cfc7704c3fea878ffa3ba69f26..3ba643cd2e1bdf6af214f069b278f01d650b26d3 100644 --- a/sumpy/expansion/local.py +++ b/sumpy/expansion/local.py @@ -94,12 +94,14 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): from sumpy.tools import MiDerivativeTaker ppkernel = self.kernel.postprocess_at_source( self.kernel.get_expression(avec), avec) + taker = MiDerivativeTaker(ppkernel, avec) return [taker.diff(mi) for mi in self.get_coefficient_identifiers()] def evaluate(self, coeffs, bvec): from sumpy.tools import mi_power, mi_factorial - evaluated_coeffs = self.stored_to_full(coeffs) + evaluated_coeffs = ( + self.derivative_wrangler.get_full_kernel_derivatives_from_stored(coeffs)) result = sum( coeff * self.kernel.postprocess_at_target(mi_power(bvec, mi), bvec) @@ -128,8 +130,7 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): # derivatives. taker = src_expansion.get_kernel_derivative_taker(dvec) - def mi_sum(a, b): - return tuple(aval + bval for aval, bval in zip(a, b)) + from sumpy.tools import add_mi result = [] for deriv in self.get_coefficient_identifiers(): @@ -137,7 +138,7 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): for coeff, term in zip( src_coeff_exprs, src_expansion.get_coefficient_identifiers()): - kernel_deriv = taker.diff(mi_sum(deriv, term)) + kernel_deriv = taker.diff(add_mi(deriv, term)) local_result.append(coeff * kernel_deriv) result.append(sp.Add(*local_result)) else: @@ -156,7 +157,7 @@ class VolumeTaylorLocalExpansion( def __init__(self, kernel, order): VolumeTaylorLocalExpansionBase.__init__(self, kernel, order) - VolumeTaylorExpansion.__init__(self) + VolumeTaylorExpansion.__init__(self, kernel, order) class LaplaceConformingVolumeTaylorLocalExpansion( @@ -165,7 +166,7 @@ class LaplaceConformingVolumeTaylorLocalExpansion( def __init__(self, kernel, order): VolumeTaylorLocalExpansionBase.__init__(self, kernel, order) - LaplaceConformingVolumeTaylorExpansion.__init__(self) + LaplaceConformingVolumeTaylorExpansion.__init__(self, kernel, order) class HelmholtzConformingVolumeTaylorLocalExpansion( @@ -174,7 +175,7 @@ class HelmholtzConformingVolumeTaylorLocalExpansion( def __init__(self, kernel, order): VolumeTaylorLocalExpansionBase.__init__(self, kernel, order) - HelmholtzConformingVolumeTaylorExpansion.__init__(self) + HelmholtzConformingVolumeTaylorExpansion.__init__(self, kernel, order) # }}} diff --git a/sumpy/expansion/multipole.py b/sumpy/expansion/multipole.py index 0444179588b49c987a5b604edf57ad9cc3a20303..fa227e89ab3a0290dd64d9420340ceeec8d93843 100644 --- a/sumpy/expansion/multipole.py +++ b/sumpy/expansion/multipole.py @@ -86,7 +86,8 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): result = [ mi_power(avec, mi) / mi_factorial(mi) for mi in self.get_full_coefficient_identifiers()] - return self.full_to_stored(result) + return ( + self.derivative_wrangler.get_stored_mpole_coefficients_from_full(result)) def evaluate(self, coeffs, bvec): taker = self.get_kernel_derivative_taker(bvec) @@ -99,8 +100,7 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): ppkernel = self.kernel.postprocess_at_target( self.kernel.get_expression(bvec), bvec) - from sumpy.tools import MiDerivativeTaker - return MiDerivativeTaker(ppkernel, bvec) + return self.derivative_wrangler.get_derivative_taker(ppkernel, bvec) def translate_from(self, src_expansion, src_coeff_exprs, dvec): if not isinstance(src_expansion, type(self)): @@ -151,7 +151,8 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase): result[i] /= mi_factorial(tgt_mi) logger.info("building translation operator: done") - return self.full_to_stored(result) + return ( + self.derivative_wrangler.get_stored_mpole_coefficients_from_full(result)) class VolumeTaylorMultipoleExpansion( @@ -160,7 +161,7 @@ class VolumeTaylorMultipoleExpansion( def __init__(self, kernel, order): VolumeTaylorMultipoleExpansionBase.__init__(self, kernel, order) - VolumeTaylorExpansion.__init__(self) + VolumeTaylorExpansion.__init__(self, kernel, order) class LaplaceConformingVolumeTaylorMultipoleExpansion( @@ -169,7 +170,7 @@ class LaplaceConformingVolumeTaylorMultipoleExpansion( def __init__(self, kernel, order): VolumeTaylorMultipoleExpansionBase.__init__(self, kernel, order) - LaplaceConformingVolumeTaylorExpansion.__init__(self) + LaplaceConformingVolumeTaylorExpansion.__init__(self, kernel, order) class HelmholtzConformingVolumeTaylorMultipoleExpansion( @@ -178,7 +179,7 @@ class HelmholtzConformingVolumeTaylorMultipoleExpansion( def __init__(self, kernel, order): VolumeTaylorMultipoleExpansionBase.__init__(self, kernel, order) - HelmholtzConformingVolumeTaylorExpansion.__init__(self) + HelmholtzConformingVolumeTaylorExpansion.__init__(self, kernel, order) # }}} diff --git a/sumpy/tools.py b/sumpy/tools.py index a18358c939e9932340058487c6af20efc4951fe3..5e0d4341c136c1594194e84e4478f613e576dbb5 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -68,21 +68,67 @@ class MiDerivativeTaker(object): try: expr = self.cache_by_mi[mi] except KeyError: - closest_mi = min( - (other_mi - for other_mi in self.cache_by_mi.keys() - if (np.array(mi) >= np.array(other_mi)).all()), - key=lambda other_mi: sum(self.mi_dist(mi, other_mi))) + current_mi = self.get_closest_cached_mi(mi) + expr = self.cache_by_mi[current_mi] + for next_deriv, next_mi in self.get_derivative_taking_sequence( + current_mi, mi): + expr = expr.diff(next_deriv) + self.cache_by_mi[next_mi] = expr + + return expr + + def get_derivative_taking_sequence(self, start_mi, end_mi): + current_mi = np.array(start_mi, dtype=int) + for idx, (mi_i, vec_i) in enumerate( + zip(self.mi_dist(end_mi, start_mi), self.var_list)): + for i in range(1, 1 + mi_i): + current_mi[idx] += 1 + yield vec_i, tuple(current_mi) + + def get_closest_cached_mi(self, mi): + return min((other_mi + for other_mi in self.cache_by_mi.keys() + if (np.array(mi) >= np.array(other_mi)).all()), + key=lambda other_mi: sum(self.mi_dist(mi, other_mi))) + + +class LinearRecurrenceBasedMiDerivativeTaker(MiDerivativeTaker): + """ + See :class:`sumpy.expansion.LinearRecurrenceBasedDerivativeWrangler` + """ + + def __init__(self, expr, var_list, wrangler): + MiDerivativeTaker.__init__(self, expr, var_list) + self.wrangler = wrangler + + def diff(self, mi): + try: + expr = self.cache_by_mi[mi] + except KeyError: + from six import iteritems + from sympy import Add + + closest_mi = self.get_closest_cached_mi(mi) expr = self.cache_by_mi[closest_mi] - current_mi = np.array(closest_mi, dtype=int) - for idx, (mi_i, vec_i) in enumerate( - zip(self.mi_dist(mi, closest_mi), self.var_list)): - for i in range(1, 1 + mi_i): - current_mi[idx] += 1 - expr = expr.diff(vec_i) - self.cache_by_mi[tuple(current_mi)] = expr + # Try to reduce the derivative using recurrences first, and if that + # fails fall back to derivative taking. + for next_deriv, next_mi in ( + self.get_derivative_taking_sequence(closest_mi, mi)): + + recurrence = ( + self.wrangler.try_get_recurrence_for_derivative( + next_mi, self.cache_by_mi)) + + if recurrence is not None: + expr = Add(*tuple( + coeff * self.cache_by_mi[ident] + for ident, coeff in iteritems(recurrence))) + else: + expr = expr.diff(next_deriv) + + self.cache_by_mi[next_mi] = expr return expr # }}}