differential equations - Return partial result when MemoryConstrained aborts NDSolve

I use NDSolve to solve a large set (~400) of coupled ODEs. Sometimes, the memory (~4GB) gets filled up, and my computer becomes impossible to work with, because it spends too much time writing to swap and the process can only be killed violently by the OS.

I circumvent this by using MemoryConstrained, but when the solver reaches the memory limit it is simply aborted and does not return the solution it obtained so far. Is there a way to obtain this solution (much like what happens when the solver encounters a singularity or reaches MaxSteps)?

Note: using a hack of the form

 StepMonitor :> If[MemoryInUse[]>...,...]

results in serious computational overhead.

Answer

Borrowing from an example of WhenEvent from the documentation in which a Button is used to stop the integration, I came up with this.

ClearAll[ndsolveMemConstrained];
SetAttributes[ndsolveMemConstrained, HoldFirst];
ndsolveMemConstrained::mlim = "Memory used `` exceeded limit ``.";

ndsolveMemConstrained[(nd_: NDSolve | NDSolveValue)[eqns_, rest___], bytes_] :=
 Module[{sol, stop, task, mstart},
  mstart = MemoryInUse[];
  stop = False;
  task = RunScheduledTask[
    stop = (ndsolve`mem = MemoryInUse[] - mstart) > bytes,
    0.2];
  sol = nd[Append[eqns,
           WhenEvent[stop,
            Message[ndsolveMemConstrained::mlim, ndsolve`mem, bytes];

            "StopIntegration"]],
     rest];
  RemoveScheduledTask[task];
  sol]

As a baseline, here is an example DE from the documentation:

NDSolveValue[{D[u[t, x], t, t] == D[u[t, x], x, x], 
   u[0, x] == Exp[-10 x^2], Derivative[1, 0][u][0, x] == 0, 
   u[t, -10] == u[t, 10]}, u, {t, 0, 100}, {x, -10, 10}, 
  Method -> "StiffnessSwitching"] // AbsoluteTiming


(* {10.969550,InterpolatingFunction[{{0.,100.},{\[Ellipsis],-10.,10.,\[Ellipsis]}},<>]} *)

When the memory is not exceeded, it takes about the same amount of time:

ndsolveMemConstrained[
  NDSolveValue[{
    D[u[t, x], t, t] == D[u[t, x], x, x], u[0, x] == Exp[-10 x^2], 
    Derivative[1, 0][u][0, x] == 0, u[t, -10] == u[t, 10]},
   u, {t, 0, 100}, {x, -10, 10}, Method -> "StiffnessSwitching"],
  8000000] // AbsoluteTiming

ndsolve`mem

(* {10.962278, InterpolatingFunction[{{0., 100.}, {..., -10., 10.,...}}, <>]} *)
(* 6992160 *)

When the memory limit is exceeded, there is frequently an extra warning message generated. I assume it has to do with where the solver is when stop is checked. (It's odd that it doesn't always produce the convergence warning.)

ndsolveMemConstrained[
  NDSolveValue[{
    D[u[t, x], t, t] == D[u[t, x], x, x], u[0, x] == Exp[-10 x^2], 
    Derivative[1, 0][u][0, x] == 0, u[t, -10] == u[t, 10]},

   u, {t, 0, 100}, {x, -10, 10}, Method -> "StiffnessSwitching"],
  4000000] // AbsoluteTiming
ndsolve`mem

NDSolveValue::evcvmit: Event location failed to converge to the requested accuracy or precision within 100 iterations between t = 56.32617731294334and t = 56.50060870314276. >>

ndsolveMemConstrained::mlim: Memory used 4158544 exceeded limit 4000000.

(* {5.595887, InterpolatingFunction[{{0., 56.3262}, {..., -10., 10.,...}}, <>]} *)
(* 4047584 *)

You can also monitor memory usage if the following is executed before ndsolveMemConstrained.

Dynamic @ ndsolve`mem

(* 6992160 *)