6
$\begingroup$

I learned a long time ago that NDSolve can often be sped up by vectorizing systems of equations. I've got a problem now that seems ideal for this approach, but I can't get it to work efficiently. Any ideas on how to speed this up? This is a simplified version of my real problem, which is orders-of-magnitude larger, but has the same issues.

There are actually two related models. The first is the master equation of a birth-death-migration-extinction stochastic process. It's linear and governed by a large transition matrix TM, which is a function of the immigration rate im.

Setting up the transition matrix:

r = 1; (* birth rate *)
k = 200; (* carrying capacity *)
e = 0.01; (* extinction rate *)
m = 0.01; (* emigration rate *)
tmax = 10^6; (* maximum time *)

SetSystemOptions["SparseArrayOptions" -> "TreatRepeatedEntries" -> Total];
SetSystemOptions["NDSolveOptions" -> "DefaultSolveTimeConstraint" -> 1200];

diagonal[vec_] := Transpose[{vec + 1, vec + 1}];
superdiagonal[vec_] := Transpose[{vec + 2, vec + 1}];
subdiagonal[vec_] := Transpose[{vec, vec + 1}];
toprow[vec_] := Transpose[{ConstantArray[1, Length[vec]], vec + 1}];

from1[n_] := -r n - r n^2/k - m n - e; (* leaving n *)
dec1[n_] := r n^2/k + m n; (* going to n-1 *)

nmax = 400; (* maximum population size *)
nrange = Range[0., nmax];

transitions = Join[diagonal[nrange], subdiagonal[Rest@nrange], superdiagonal[Most@nrange], toprow[nrange]];

TM[im_] := Module[{rates},
  rates = Join[
    from1[nrange] - im,
    dec1[Rest@nrange],
    r*Most@nrange + im,
    ConstantArray[e, nmax + 1]
  ];
  SparseArray[transitions -> rates]
];

This linear master equation for the probability p of having n individuals can be quickly solved by NDSolve in vector form for a fixed immigration rate im:

(* initial probability distribution *)
p0 = Table[If[n == Round[k], 1, 0], {n, 0, nmax}];

(* 0.03 s *)
s = NDSolve`ProcessEquations[{p'[t] == TM[0.01].p[t], p[0] == p0}, p, t][[1]];

(* 0.62 s *)
NDSolve`Iterate[s, tmax];

sol = NDSolve`ProcessSolutions[s];
ListPlot[p[tmax] /. sol, PlotRange -> All]

Mathematica graphics

Good so far. Now the real problem: I want to make the immigration rate im a function of the probability distribution p[t] (a structured metapopulation similar to Metz & Gyllenberg 2001). This makes the system nonlinear and apparently impossible to solve efficiently in vectorized form.

First attempt:

s = NDSolve`ProcessEquations[{p'[t] == TM[m p[t].Range[0, nmax]].p[t],
  p[0] == p0}, p, t][[1]];
(* NDSolve`ProcessEquations::ndnum
  Encountered non-numerical value for a derivative at t == 0.`. *)

This seems wrong, since the initial immigration rate im=m p[0].Range[0, nmax] is in fact numerical:

m p0.Range[0, nmax]
(* 2. *)

Wrapping the SparseArray TM in Normal makes it work, but slowly:

(* 0.15 s *)
s = NDSolve`ProcessEquations[{p'[t] == Normal@TM[m p[t].Range[0, nmax]].p[t],
  p[0] == p0}, p, t][[1]];

(* 8.26 s *)
NDSolve`Iterate[s, tmax];

sol = NDSolve`ProcessSolutions[s];
ListPlot[p[tmax] /. sol, PlotRange -> All]

Mathematica graphics

Maybe that's as good as it's going to get? No -- here are two different implementations of an unvectorized version that are faster.

First, using built-in NDSolve`ProcessEquations:

rhs = TM[m*Sum[n p[n][t], {n, 0, nmax}]].Table[p[n][t], {n, 0, nmax}];
eqns = Thread[Table[p[n]'[t], {n, 0, nmax}] == rhs];
ics = Table[p[n][0] == If[n == k, 1, 0], {n, 0, nmax}];
unks = Table[p[n], {n, 0, nmax}];

(* 2.75 s *)
s = NDSolve`ProcessEquations[Flatten[{eqns, ics}], unks, t][[1]];

(* 2.73 s *)
NDSolve`Iterate[s, tmax];

sol = NDSolve`ProcessSolutions[s];
ListPlot[Table[p[n][tmax], {n, 0, nmax}] /. sol, PlotRange -> All]
(* same output as above *)

Second, using my homemade ProcessFirstOrderODEs:

rhs = TM[m*Sum[n p[n], {n, 0, nmax}]].Table[p[n], {n, 0, nmax}];
ics = Table[If[n == k, 1, 0], {n, 0, nmax}];
unks = Table[p[n], {n, 0, nmax}];

(* 0.55 s *)
s = ProcessFirstOrderODEs[unks, rhs, ics, 0];

(* 2.82 s *)
NDSolve`Iterate[s, tmax]

sol = NDSolve`ProcessSolutions[s];
ListPlot[Table[p[n][tmax], {n, 0, nmax}] /. sol, PlotRange -> All]
(* same output as above *)

Finally, both NDSolve`Iterate stages can be sped up to around 0.6 s by processing the equations with option Jacobian -> FiniteDifference (not shown).

To recap:

  • the vectorized linear problem works fine with NDSolve
  • the vectorized nonlinear problem does not work with NDSolve, except when using Normal, which results in a total timing of 8.41 s
  • the nonvectorized nonlinear problem with NDSolve is faster (5.48 s)
  • the nonvectorized nonlinear problem with my ProcessFirstOrderODEs is even faster (3.37 s)
  • the nonvectorized nonlinear problem with my ProcessFirstOrderODEs and Jacobian->FiniteDifference is fastest (1.1 s)

Questions:

  • can the vectorized nonlinear problem be made to work with NDSolve without using Normal?
  • is there any other way to speed up this NDSolve?
$\endgroup$
4
$\begingroup$

I think the problem is that in TM[m p[t].Range[0, nmax]], the NDSolve "variable" p[t] is incorporated into the SparseArray and protected from being assigned a value in whatever way NDSolve achieves that. This can be overcome by using _?NumericQ. This speeds up the ProcessEquations[] step and the Iterate[] step (from 9 to 4 sec.). Specifying Method -> "StiffnessSwitching" speeds up Iterate[] even more. (For some reason, it does not speed up the OP's fastest way, with ProcessFirstOrderODEs.)

ClearAll[TM];
TM[im_?NumericQ] := 
  Module[{rates}, 
   rates = Join[from1[nrange] - im, dec1[Rest@nrange], 
     r*Most@nrange + im, ConstantArray[e, nmax + 1]];
   SparseArray[transitions -> rates]];

s = NDSolve`ProcessEquations[{p'[t] == TM[m p[t].Range[0, nmax]].p[t],
       p[0] == p0}, p, t, 
     Method -> "StiffnessSwitching"][[1]]; // AbsoluteTiming

NDSolve`Iterate[s, tmax]; // AbsoluteTiming

sol = NDSolve`ProcessSolutions[s];
ListPlot[p[tmax] /. sol, PlotRange -> All]
(*
  {0.001504, Null}
  {1.58603, Null}
*)

enter image description here

$\endgroup$
  • $\begingroup$ Hmm, I thought I'd tried mucking around with ?NumericQ but evidently missed the right combination. That answers my first question about an elegant, compact vectorized implementation. Too bad it's not faster than doing it as separate first-order equations. But here's another Method that brings the NDSolve`Iterate down to 0.23s: Method -> {"IDA", "ImplicitSolver" -> {"Newton", "LinearSolveMethod" -> {"Band", "BandWidth" -> {2, 2}}}}. Decent! $\endgroup$ – Chris K Aug 4 at 13:55

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.