Increase Performance of WhenEvent in NDSolve

Question

I'm trying to reset the value of functions $u_{i}$ when they reach a threshold $θ$. Below the threshold, they evolve according to a simple differential equation. The $r_{i}$ resets the derivative to zero for a while (refractory period of a neuron), but I doubt that's relevant to my question.

NDSolve is really slow when the number of events is increased (which can be done by decreasing recovery). I don't understand why an artificial discontinuity would make this code slower. Here is a minimal example where all $u_{i}$ are the same:

recovery = 0.01; θ = 0.5; n = 50; T = 2;
s = ConstantArray[{}, n];

AbsoluteTiming[
 sol = NDSolve[Table[With[{i = i}, {
  Subscript[u, i]'[t] == (-Subscript[u, i][t] + 10^3) Boole[Subscript[r, i][t] < t], 
  Subscript[r, i][0] == 0, 
  Subscript[u, i][0] == 0, 

  WhenEvent[
   Subscript[u, i][t] > θ, {Subscript[u, i][t] -> 0, 
   AppendTo[s[[i]], t], 
   Subscript[r, i][t] -> t + recovery}]}],
 {i, n}], 

  Table[{Subscript[u, i], Subscript[r, i]}, {i, n}] // Flatten, {t, 0, T}, 
  DiscreteVariables -> Table[Subscript[r, i], {i, n}]];
 ]

Length[s // First] (*number of events*)

(*{65.485, Null}*)
(*191*)

What could be causing this, and is there any way to make it faster?

Answer 1

That Boole worries me, so here's an approach that avoids it. I define another DiscreteVariable for each equation called Subscript[sw,i][t] that is either 0 or 1 and acts as a switch on the Subscript[u, i]'[t] equations. I also add another WhenEvent that turns the switch back off at time t==Subscript[r, i][t].

AbsoluteTiming[sol = NDSolve[Table[With[{i = i}, {
  Subscript[u, i]'[t] == Subscript[sw, i][t] (-Subscript[u, i][t] + 10^3),
  Subscript[r, i][0] == 0, Subscript[u, i][0] == 0, Subscript[sw, i][0] == 1, 
  WhenEvent[Subscript[u, i][t] > θ, 
    {Subscript[u, i][t] -> 0, AppendTo[s[[i]], t],
     Subscript[r, i][t] -> t + recovery, Subscript[sw, i][t] -> 0}],
  WhenEvent[t == Subscript[r, i][t], Subscript[sw, i][t] -> 1]
 }], {i, n}],
 Table[{Subscript[u, i], Subscript[r, i], Subscript[sw, i]}, {i, n}] // Flatten,
 {t, 0, T}, DiscreteVariables ->
   Evaluate[Table[{Subscript[r, i], Subscript[sw, i]}, {i, n}] // Flatten]
];]

Length[s // First]

(* {18.3326, Null} *)
(* 191 *)

I am hoping that there's a better way though, because this still seems too slow. I thought using Reap and Sow instead of AppendTo might help, but it really didn't.

Answer 2

This is in part some minor improvements in speed for ChrisK's good answer (+1), and in part an extended comment. To begin, using Subscript variables generally is not a good idea in computations. Also, building Tables inside NDSolve often is slow, although it does not matter much here. With these changes, the question's code as optimized by ChrisK becomes

recovery = 1/100; θ = 1/2; n = 50; T = 2;
s = ConstantArray[{}, n];
var = Table[ {u[ i], r[i], sw[i]}, {i, n}] // Flatten ;
dvar = Table[{r[i][t], sw[ i][t]}, {i, n}] // Flatten;
eq = Table[With[{i = i}, {u[ i]'[t] == sw[ i][t] (- u[i][t] + 10^3), 
    r[i][0] == 0, u[i][0] == 0, sw[i][0] == 1, 
    WhenEvent[u[i][t] > θ, {u[i][t] -> 0, AppendTo[s[[i]], t], 
        r[i][t] -> t + recovery, sw[ i][t] -> 0}],
    WhenEvent[t > r[i][t], sw[i][t] -> 1]}], {i, n}];
sol = NDSolve[eq, var, {t, 0, T}, DiscreteVariables -> dvar]; // AbsoluteTiming
Length[s // First]

(* {16.8441, Null} *)
(* 191 *)

which gives the desired answer in about seventeen seconds. As I noted in my earlier comment, this computation is very slow, because restarting the computation with WhenEvent 2*191 times results in many very small timesteps, 1791 in this case. In contrast, the computation requires only 49 timesteps and takes only about 0.03 seconds without WhenEvent.

The runtime can be reduced by about 30% by using the option, "LocationMethod" -> "LinearInterpolation" in WhenEvent and requesting only modest accuracy from NDSolve: AccuracyGoal -> 3, PrecisionGoal -> 3. (Nerve cells do not fire with high precision anyway.) Doing so produces essentially the same answers as before but in about twelve seconds. It seems likely that choosing an optimal integration Method will further reduce runtime. I say this based on the observation that the solution is almost always zero, punctuated by 191 evenly spaced spikes of width 10^-3 θ. I tried several possibilities but only succeeded in increasing runtime.

I also tried eliminating the second WhenEvent by using

u[i]'[t] == Piecewise[ {{(- u[i][t] + 10^3), t > r[i][t]}}, 0]

or

u[i]'[t] == If[t > r[i][t], (- u[i][t] + 10^3), 0]

but doing so increased runtime to a bit more than forty seconds. As noted in the question, Boole is slower yet, although it certainly yields the correct answer.

Finally I would remark that the OP probably intends this question as a prototype for many coupled neurons (those in the question are uncoupled) firing at different times. To see the effect of neurons firing at different times, I used the initial condition, u[i][0] == θ (i - 1)/n. Runtime increased to about 300 sec. All this suggests that computational models not involving NDSolve should be investigated.