diff --git a/sumpy/expansion/local.py b/sumpy/expansion/local.py index 2f494a66bced904e7f8425597b23094ce84c8edf..ad29ea056ac84a898aa39c08ad816b4b64fd6f81 100644 --- a/sumpy/expansion/local.py +++ b/sumpy/expansion/local.py @@ -157,64 +157,44 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase): # # coeff0 * diff(kernel, mi0) + coeff1 * diff(kernel, mi1) + ... # - # To get the local expansion coefficients, we take derivatives of - # the multipole expansion. + # The local is formed by differentiating the mpole, i.e. the + # target variable is in the mpole basis--so that is what the + # derivatives for the local will hit. # - # This code speeds up derivative taking by caching all kernel - # derivatives. + # We group the calculation here so that we take all derivatives + # in one big gulp, which ends up being more efficient than + # evaluating the mpole (as a giant sum) and then differentiating + # that. + # + # In particular, this speeds up derivative taking by allowing all + # kernel derivatives to be cached. taker = src_expansion.get_kernel_derivative_taker(dvec) from sumpy.tools import add_mi - # Create a expansion terms wrangler for derivatives up to order - # (tgt order)+(src order). - tgtplusderiv_exp_terms_wrangler = \ - src_expansion.expansion_terms_wrangler.copy( - order=self.order + src_expansion.order) - tgtplusderiv_coeff_ids = \ - tgtplusderiv_exp_terms_wrangler.get_coefficient_identifiers() - tgtplusderiv_full_coeff_ids = \ - tgtplusderiv_exp_terms_wrangler.get_full_coefficient_identifiers() - - tgtplusderiv_ident_to_index = dict((ident, i) for i, ident in - enumerate(tgtplusderiv_full_coeff_ids)) + src_exp_terms_wrangler = src_expansion.expansion_terms_wrangler + src_exp_stored_coeff_ids = \ + src_exp_terms_wrangler.get_coefficient_identifiers() result = [] - for lexp_mi in self.get_coefficient_identifiers(): - lexp_mi_terms = [] - - # Embed the source coefficients in the coefficient array - # for the (tgt order)+(src order) wrangler, offset by lexp_mi. - embedded_coeffs = [0] * len(tgtplusderiv_full_coeff_ids) - for coeff, term in zip( - src_coeff_exprs, - src_expansion.get_coefficient_identifiers()): - embedded_coeffs[ - tgtplusderiv_ident_to_index[add_mi(lexp_mi, term)]] \ - = coeff - - # Compress the embedded coefficient set - stored_coeffs = tgtplusderiv_exp_terms_wrangler \ - .get_stored_mpole_coefficients_from_full( - embedded_coeffs, src_rscale) - - # Sum the embeeded coefficient set - for i, coeff in enumerate(stored_coeffs): - if coeff == 0: - continue - nderivatives_for_scaling = \ - sum(tgtplusderiv_coeff_ids[i])-sum(lexp_mi) + for tgt_coeff_id in self.get_coefficient_identifiers(): + tgt_coeff_terms = [] + + for src_coeff_id, coeff in zip( + src_exp_stored_coeff_ids, src_coeff_exprs): + nderivatives_for_scaling = sum(src_coeff_id) + term_mi = add_mi(src_coeff_id, tgt_coeff_id) kernel_deriv = ( src_expansion.get_scaled_multipole( - taker.diff(tgtplusderiv_coeff_ids[i]), + taker.diff(term_mi), dvec, src_rscale, - nderivatives=sum(tgtplusderiv_coeff_ids[i]), + nderivatives=sum(term_mi), nderivatives_for_scaling=nderivatives_for_scaling)) - lexp_mi_terms.append( - coeff * kernel_deriv * tgt_rscale**sum(lexp_mi)) - result.append(sym.Add(*lexp_mi_terms)) + tgt_coeff_terms.append( + coeff * kernel_deriv * tgt_rscale**sum(tgt_coeff_id)) + result.append(sym.Add(*tgt_coeff_terms)) else: from sumpy.tools import MiDerivativeTaker