Speed up Expectation-Maximization algorithm

Question

I would like to implement the EM algorithm for parameter estimation of a bivariate poisson model from Karlis 2003. The code is able to run, but I'm not very satisfied with the speed. I would like to ask for some ideas to speed up my code.

 fitBivariatePoissonModel[x_?ListQ, init3_: 1, maxIteration_: 300, 
  prec_: 10^-8] := 
 Block[{n = Length[x], s, like, zeroQ, lambda1 = 0., lambda2 = 0., 
   lambda3 = init3, var1 = x[[All, 1]], var2 = x[[All, 2]], 
   difllike = 1000., loglike0 = 1000., i = 0, lbp1, lbp2, x1, x2, 
   loglike, loglikeli},
  s = ConstantArray[0., n];
  like = ConstantArray[0., n];
  zeroQ = If[MemberQ[#, 0], 1., .0] & /@ x;
  lambda1 = Max[0.1, Total[var1]/n*1.0 - lambda3];
  lambda2 = Max[0.1, Total[var2]/n*1.0 - lambda3];
  loglikeli = Reap[While[difllike > prec && i <= maxIteration,
      i++;
      lbp1 = 
       MapThread[
        LogLikelihood[
          MultivariatePoissonDistribution[
           lambda3, {lambda1, lambda2}], {{#1 - 1, #2 - 1}}] &, {var1,
          var2}];
      lbp2 = 
       MapThread[
        LogLikelihood[
          MultivariatePoissonDistribution[
           lambda3, {lambda1, lambda2}], {{#1, #2}}] &, {var1, 
         var2}];
      s = 
       MapThread[
        If[MemberQ[{#1, #2}, 0], 0, 
          Exp[Log[lambda3] + 
            LogLikelihood[

             MultivariatePoissonDistribution[
              lambda3, {lambda1, lambda2}], {{#1 - 1, #2 - 1}}] - 
            LogLikelihood[
             MultivariatePoissonDistribution[
              lambda3, {lambda1, lambda2}], {{#1, #2}}]]] &, {var1, 
         var2}];
      like = MapThread[If[MemberQ[{#1, #2}, 0],

          LogLikelihood[PoissonDistribution[lambda1], {#1}] + 
           LogLikelihood[PoissonDistribution[lambda2], {#2}] - 
           lambda3, 
          LogLikelihood[
           MultivariatePoissonDistribution[
            lambda3, {lambda1, lambda2}], {{#1, #2}}]] &, {var1, 
         var2}];
      x1 = var1 - s;
      x2 = var2 - s;

      Sow[loglike = Total[like]];
      difllike = Abs[(loglike0 - loglike)/loglike0];
      loglike0 = loglike;

      lambda1 = Mean[x1];
      lambda2 = Mean[x2];
      lambda3 = Mean[s]]][[2, 1]];
  If[i == maxIteration + 1, Print["Maximum iterations reached"]];
  Print[ListLinePlot[loglikeli]];
  {lambda1, lambda2, lambda3, loglike0}]

A test evaluation of 100 samples took 30 seconds on my computer. Apparently, for large application, this evaluation speed is quite unsatisfied. Any ideas to accelerate would be appreciated.

l1 = Table[RandomInteger[{0, 10}, 2], {i, 100}];
fitBivariatePoissonModel[l1, 0.01, 500, 10^-8] // AbsoluteTiming
(* {38.5076, {4.84859, 5.34859, 0.00140634, -550.099}} *)

Answer 1

Discussion

This is maybe not the main problem but the frequent calls to LogLikelihood cause a lot of symbolic computation that tends to be slow. Indeed, we are in a situation in which the outcome of LogLikelihood has a simple symbolic expression. Thus, we may exploit that to perform the symbolic computation only once by defining:

Block[{u, v, λ1, λ2, λ3},
  f[λ1_, λ2_, λ3_][u_, v_] = LogLikelihood[
    MultivariatePoissonDistribution[λ3, {λ1, λ2}], {{u, v}}
    ]
  ];

Then one may use it as follows:

{var1, var2} = RandomInteger[{0, 10}, {2, 1000}];
lambda1 = 1.;
lambda2 = 1.;
lambda3 = 1.;
aa = MapThread[f[lambda1, lambda2, lambda3], {var1 - 1, var2 - 1}]; // AbsoluteTiming // First
bb = MapThread[
    LogLikelihood[
      MultivariatePoissonDistribution[lambda3, {lambda1, lambda2}], {{#1 - 1, #2 - 1}}] &, 
      {var1, var2}
    ]; // AbsoluteTiming // First
aa == bb

0.018541

0.197101

True

As you can see, this becomes 10 times as fast.

Similar tweaks should accelerate also the other calls to LogLikelihood.

Another issue might be that -∞ may occur quite often. Since this is symbolic, every array in which it occurs must be unpacked. That might slow down things. But I guess the biggest bottleneck is HypergeometricU (appearing in LogLikelihood); it requires 65 % of the computation time of the above code.

Next thing I found is that you perform many computations redundantly. For example, setting

s = Unitize[var1 var2] Exp[Log[lambda3] + lbp1 - lbp2];

instead of

s = MapThread[ If[MemberQ[{#1, #2}, 0], 0, <<...>>] &, {var1, var2}];

cuts the computation time in half.

Moreover, by setting

g[λ1_][u_] = Simplify[LogLikelihood[PoissonDistribution[λ1], {u}], u >= 0];

which exploits that g is never fed a negative second argument, we obtain a vectorized functions g[λ1] which can be applied to lists without Map or MapThread.

Updated Implementation

With some further minor changes, this is the improved code.

fitBivariatePoissonModel2[x_?ListQ, init3_: 1, maxIteration_: 300, 
  prec_: 10^-8] := 
 Block[{n, s, like, zeroQ, lambda1, lambda2, lambda3, var1, var2, 
   difllike = 1000., loglike0 = 1000., i, lbp1, lbp2, loglike, 
   loglikeli, f, g, χ, meanvar1, meanvar2},
  i = 0;
  lambda1 = 0.;
  lambda2 = 0.;
  lambda3 = N[init3];
  
  {var1, var2} = Transpose[x];
  meanvar1 = Mean[N[var1]];
  meanvar2 = Mean[N[var2]];
  n = Length[x];
  s = ConstantArray[0., n];
  like = ConstantArray[0., n];
  zeroQ = 1. - Unitize[var1 var2];
  lambda1 = Max[0.1, Mean[N[var1]] - lambda3];
  lambda2 = Max[0.1, Mean[N[var2]] - lambda3];
  
  Block[{u, v, λ1, λ2, λ3},
   f[λ1_, λ2_, λ3_][u_, v_] = LogLikelihood[ MultivariatePoissonDistribution[λ3, {λ1, λ2}], {{u, v}}];
   g[λ1_][u_] = Simplify[LogLikelihood[PoissonDistribution[λ1], {u}], u >= 0];
   ];
  
  loglikeli = Reap[
     While[difllike > prec && i <= maxIteration, i++;
      lbp1 = MapThread[f[lambda1, lambda2, lambda3], {var1 - 1, var2 - 1}];
      lbp2 = MapThread[f[lambda1, lambda2, lambda3], {var1, var2}];
      χ = N[Unitize[var1 var2]];
      s = χ Exp[Log[lambda3] + lbp1 - lbp2];
      like = χ lbp2 + (1. - χ) (g[lambda1][var1] + g[lambda2][var2] - lambda3);
      Sow[loglike = Total[like]];
      difllike = Abs[1. - loglike/loglike0];
      loglike0 = loglike;
      lambda3 = Mean[s];
      lambda1 = meanvar1 - lambda3;
      lambda2 = meanvar2 - lambda3;
      ]
     ][[2, 1]];
  If[i == maxIteration + 1, Print["Maximum iterations reached"]];
  Print[ListLinePlot[loglikeli]];
  {lambda1, lambda2, lambda3, loglike0}
  ]

Timing and accuracy comparison:

SeedRandom[1234];
l1 = RandomInteger[{0, 10}, {100, 2}];
aa = fitBivariatePoissonModel[l1, 0.01, 500, 10^-8]; // 
  AbsoluteTiming // First
bb = fitBivariatePoissonModel2[l1, 0.01, 500, 10^-8]; // 
  AbsoluteTiming // First
Max[Abs[aa - bb]]

18.7093

0.863204

2.27374*10^-13

There is certainly more potential for improvement. Probably the algorithm is not implemented correctly; this comes to my mind because I observed that the algorithm takes rather many interations to "converge" and that the result is very sensitive to decreasing prec.