Optimizing code containing an infinite sum

Question

I have the following expression. It's based on the negative-log-likelihood of a Poisson-Gaussian noise model, i.e. a convolution of a Poisson & Normal distribution. The value $x_{i}$ is the original value, and $z_{i}$ is the measured value. $\sigma^{2}$ is the variance of the additive noise component.

$$ \sum_{i} \left [ x_{i} - \ln \left ( \sum_{n=0}^{\infty } \frac{\left ( x_{i} \right )^{n}}{n!}\textrm{e}^{-x_{i}} \times \textrm{e}^{-\frac{\left ( z_{i} - n \right )^{2}}{2\sigma ^{2}}} \right ) \right ] $$

The infinite sum inside the log makes things nice and complicated :-) It typically converges after about 50-60 terms, although it varies depending on parameters.

If I create a list of data $z_{i}$ and evaluate this sum to 500 terms (to be on the safe side), this is the performance, initially with no noise.

func[x_, a_, b_, c_] := a*Exp[-((x - b)^2/(2*c^2))];

data = Chop@Table[func[x, 5, 50, 3], {x, 1, 100}];

First@AbsoluteTiming@Do[
    With[{a = 5., b = 50., c = 3., sig = 2.},
     result1 = 
      Total@Map[
        func[#, a, b, c] - 
          Log[
            Sum[func[#, a, b, c]^n/n! * Exp[-((data[[#]] - n)^2/(2. * sig^2))],
              {n, 0, 500}]] &, Range[100]]
     ], {10}]/10
(* takes 0.872 seconds to evaluate the result *)

An alternative is to replace Sum with NSum, namely:

First@AbsoluteTiming@Do[
    With[{a = 5., b = 50., c = 3., sig = 2.},
     result2 = 
      Total@Map[
        func[#, a, b, c] - 
          Log[
            NSum[func[#, a, b, c]^n/n! * Exp[-((data[[#]] - n)^2/(2. * sig^2))],
              {n, 0, Infinity}]] &, Range[100]]
     ], {10}]/10
(* takes 0.197 seconds to evaluate the result *)

NumberForm[result1, 15]
(* 3.38562655668556 *)

NumberForm[result2, 15]
(* 3.38562655668556 *)

result2 - result1
(* 4.44*10^-16 *)

So I can speed up by about 4x using NSum, for a negligible difference in the end result.

However, here is an alternative approach, based on some literature I've been reading, to approximate the infinite series in a better/more rigorous way, rather than just saying "it converges after about 60 terms".

It leads to this expression:

$$ \sum_{i} \left [ x_{i} - \ln \left ( \textrm{e}^{-\frac{z_{i}^{2}}{\sigma^{2}}} + \sum_{n=\textrm{max}\left [ 1, \,n^{*} - \Delta \sigma \right ]}^{n^{*} + \Delta \sigma } \frac{\left ( x_{i} \right )^{n}}{n!}\textrm{e}^{-x_{i}} \times \textrm{e}^{-\frac{\left ( z_{i} - n \right )^{2}}{2\sigma ^{2}}} \right ) \right ] $$

where

$$ n^{*}=\sigma^{2} \: \textrm{W}\left ( \frac{x_{i}}{\sigma^{2}} \textrm{e}^{\frac{z_{i}}{\sigma^{2}}} \right ) $$

and $\textrm{W}$ is the Lambert function, implemented in Mathematica as ProductLog. Here $\Delta$ is a number, typically about 5 or 6.

Implemented in Mathematica

First@AbsoluteTiming@Do[
 With[{a = 5., b = 50., c = 3., sig = 2., delta = 6},
  result3 = 
   Total@Map[Block[{nstar, exp1, var},
     var = data[[#]];
     exp1 = a*Exp[-((# - b)^2/(2*c^2))];
     nstar = sig^2 * ProductLog[(exp1/sig^2) * Exp[(var/sig^2)]];
     exp1 - Log[Exp[-(var^2/(2. * sig^2))] + Sum[exp1^n/n!*Exp[-((var - n)^2/(2. * sig^2))], 
     {n, Max[1., nstar - (delta * sig)], nstar + (delta * sig)}]]
  ] &, Range[100]]
], {10}]/10
(* takes 0.0073 seconds to evaluate the result *)

Which is a much, much bigger speed-up. One can also check the convergence of this approach.

NumberForm[result1, 15]
(* 3.38562655668556 *)

NumberForm[result3, 15]
(* 3.38562660516702 *)

result3 - result1
(* 4.85*10^-8 *)

So while the convergence isn't as good, it's still acceptable. In fact, if you increase the delta parameter to 9...

(* Takes 0.0097 seconds *)

NumberForm[result4, 15]
(* 3.38562655668556 *)

result4 - result1
(* 4.44*10^-16 *)

So it gives the same answer as Sum and NSum, much faster.

My eventual aim is fit a model $m(\theta)$ such that $x = m(\theta)$ to some measured data $z$, which will involve using these functions as the minimizers inside FindMinimum/NMinimize.

func[x_, a_, b_, c_] := a*Exp[-((x - b)^2/(2*c^2))];
addMixedNoise[x_, sigma_] := Block[{temp},
   temp = Map[If[# == 0., 0., RandomVariate[PoissonDistribution[#]]] &, x];
   Map[# + RandomVariate[NormalDistribution[0., sigma]] &, temp]
 ];

cleandata = Chop@Table[func[x, 5, 50, 3], {x, 1, 100}];
data = addMixedNoise[cleandata, 2.];

minimumfunction[inputdata_, a_?NumberQ, b_?NumberQ, c_?NumberQ, 
  sig_?NumberQ, delta_?NumberQ] :=
 Total@Map[Block[{nstar, exp1, var},
     var = inputdata[[#]];
     exp1 = a*Exp[-((# - b)^2/(2*c^2))];
     nstar = sig^2 * ProductLog[(exp1/sig^2) * Exp[(var/sig^2)]];
     exp1 - Log[Exp[-(var^2/(2. * sig^2))] + Sum[exp1^n/n!*Exp[-((var - n)^2/(2. * sig^2))], 
         {n, Max[1., nstar - (delta * sig)], nstar + (delta * sig)}]]
     ] &, Range[100]]

fitval = FindMinimum[minimumfunction[data, a, b, c, 2., 6.], 
           {{a, 5.}, {b, 50.}, {c, 3.}},
           Method -> "QuasiNewton"] // AbsoluteTiming
(* Takes: 0.78 seconds *)
(* {51.7574, {a -> 5.97587,b -> 49.7221, c -> 2.62345}} *)

So my question is:

Are there any further ways to speed up my function, given that ProductLog cannot, to my knowledge, be compiled?

I've got much bigger models than a single 1D Gaussian to fit (with more than 100 data points too).

Given that ProductLog doesn't compile, and that NSum proved slower than the approximation using a Lambert function I've used here, this may well be the fastest I can get it to work.

If anyone does have more suggestions, I'd be happy to hear them!

Answer 1

One can rewrite first OP's method a bit and obtain up to 500x speedup by using PackedArrays

SetSystemOptions["CatchMachineUnderflow" -> False];

x = N@Range[100];
data = func[x, 5., 50., 3.];
n = N@Range[0, 500];
a = 5.; b = 50.; c = 3.; sig = 2.;

Total[func[x, a, b, c] - Log@Total@Exp[
      Outer[Times, n, Log@func[x, a, b, c]] -
       Outer[Times, LogGamma[n + 1], ConstantArray[1., Length[x]]] -
       Outer[Plus, n, -data]^2/(2.*sig^2)
      ]] // AbsoluteTiming
(* {0.002632, 3.38563} *)

First option here allows us to not catch machine underflow when evaluating Exp of big negative numbers (see here).

Answer 2

A small speed-up can be gained by computing the ProductLog only once, and also by accessing the input data list just once, namely:

minimimumfunction2[inputdata_, a_?NumberQ, b_?NumberQ, c_?NumberQ, 
  sig_?NumberQ, delta_?NumberQ] :=
 Total@Map[Block[{nstar, exp1, var},
     var = inputdata[[#]];
     nstar = sig^2 * ProductLog[(a*Exp[-((# - b)^2/(2 * c^2))])/sig^2 * Exp[-(var/sig^2)]];
     exp1 = a*Exp[-((# - b)^2/(2*c^2))];
     exp1 - Log[Exp[-(var^2/(2. * sig^2))] + Sum[exp1^n/n!*Exp[-((var - n)^2/(2. * sig^2))], 
         {n, Max[1., nstar - (delta * sig)], nstar + (delta * sig)}]]
     ] &, Range[100]]

fitval = FindMinimum[minimimumfunction2[data, a, b, c, 2., 6.], 
           {{a, 5.}, {b, 50.}, {c, 3.}},
           Method -> "QuasiNewton"] // AbsoluteTiming
(* Takes: 0.78 seconds *)

Which is about 50% faster than in the question above.

Given that ProductLog doesn't compile, and that NSum proved slower than the approximation using a Lambert function I've used here, this may well be the fastest I can get it to work.

If anyone does have more suggestions, I'd be happy to hear them!