Rewriting For loop to FoldList or similar

Question

In CrossValidated, the answer to a question on prediction of number of viewers of Gangnam Style is given in R. I have mapped the answer directly to Mathematica:

S0 = 10000000; a = 0.00000005; b = 0.01; k = 1.2;
views = {{0, 0}}; S = S0; inf = 1;
For[t = 1, t < 365, t++, dS = -a*inf*S;
  dinf = a*inf*S - b*inf;
  S = S + dS;
  inf = inf + dinf;
  noviews = Last[Last[views]] + k inf; 
  views = Append[views, {t, noviews}]];
ListPlot[views]

However, this doesn't really feel or look like good Mathematica code. I tried to follow the advice in Alternatives to procedural loops and iterating over lists in Mathematica. But I get stuck in trying to create the function that the FoldList should use. It seems like the solution would be to have FoldList create a number of lists, one for each "state" variable.

So my question is: how do I rewrite the For loop that has a number of "state" variables and where the output depends on both these variables and earlier values of the output?

Answer 1

I wouldn't try to force FoldList, if there are better suited built-ins like RecurrenceTable

We could try to simplify your equations first

neweqs = {
     views[n + 1] == views[n] + k inf[n + 1], 
     inf[n + 1] == inf[n] + dinf, 
     dinf == a inf[n] S[n] - b inf[n], 
     S[n + 1] == S[n] + dS, 
     dS == -a inf[n] S[n]
     }~Eliminate~{dS, dinf} /. And -> List // Simplify

to get

(* {b inf[n] + inf[1 + n] == inf[n] (1 + a S[n]), 
   a inf[n] S[n] + S[1 + n] == S[n], 
   k inf[1 + n] + views[n] == views[1 + n]} *)

So now we can write

Block[{S0 = 10000000, a = 0.00000005, b = 0.01, k = 1.2,
  view0 = {{0, 0}}, inf0 = 1, max = 365, views, inf, S, n},

 view2 = RecurrenceTable[neweqs~Join~{
       views[0] == Last@Last@view0, inf[0] == inf0, S[0] == S0},
     {views[n]}, {n, 1, max},
     DependentVariables -> {inf, S, views}] /. 
    vals_ :> Transpose@{Range[max], Flatten@vals};
 ]
ListPlot @ view2

and get the same output as your example

For a maximum of 30000 I get in my machine an absolute timing of

your version: 5.6 s

this version: 0.017 s

So, at least for that size, thats a 330× speedup.

Answer 2

UPDATE: With FoldList and NestList one can eliminate the need for Join at every iteration and get much simpler and faster forms:

 foldListVws = Transpose[{Range[0, 364],
  (FoldList[{#[[1]] + k #[[3]] (1 + a #[[2]] - b),
     #[[2]] (1 - a #[[3]]), #[[3]] (1 + a #[[2]] - b)} &,
   {0, S0, 1}, Range[0, 364]][[All, 1]] //  Most)}]; 
 foldListVws == views
 (* True *)

 nestListVws =  NestList[{#[[1]] + 1, #[[2]] + k #[[4]] (1 + a #[[3]] - b),
     #[[3]] (1 - a #[[4]]), #[[4]] (1 + a #[[3]] - b)} &,
    {0, 0, S0, 1}, 364][[All, ;; 2]]; 
 nestListVws == views
 (* True *)

Using max= 30 000 as in @Rojo's post, the timings I get for the five methods (on a Windows Vista 64bit laptop with Intel Core2 Duo 2.80GHz CPU and 8G memory) are:

More readable update functions for FoldList and NestList:

ClearAll[nstLstUF, fldLstUF];
nstLstUF =Function[{t,x,y,z}, {t+1, x+k z (1 + a y - b), y (1 - a z), z (1 + a y - b)}];
fldLstUF = Function[{x,y,z}, {x + k z (1 + a y - b), y (1 - a z), z (1 + a y - b)}]; 

... and almost-one liners:

fldLstVws2 = Transpose[{Range[0, 364],
  (FoldList[fldLstUF[Sequence@@#]&, {0, S0, 1}, Range[0, 364]][[All, 1]] // Most)}];
nstLstVws2 = NestList[nstLstUF[Sequence @@ #]&,{0, 0, S0, 1}, 364][[All, ;; 2]];
fldLstVws2 == nstLstVws2 == views
(* True *)

---

Original post:

 foldVws =First@Fold[{Join[#1[[1]],
    {{#2, Last@Last[#1[[1]]] + k #1[[3]] (1 + a #1[[2]] - b)}}],
   #1[[2]] (1 - a #1[[3]]), 
   #1[[3]] (1 + a #1[[2]] - b)} &,
 {{{0, 0}}, S0, 1},
 Range[364]];
 foldVws==views
 (* True *)

Alternatively, you can use Nest:

 nestVws = Nest[{Join[#1[[1]],
  {{#1[[4]], Last@Last[#1[[1]]] + k #1[[3]] (1 + a #1[[2]] - b)}}],
  #1[[2]] (1 - a #1[[3]]), 
  #1[[3]] (1 + a #1[[2]] - b), #1[[4]] + 1} &,
  {{{0, 0}}, S0, 1, 1}, 364] // First;
 nestVws==views
 (* True *)

Answer 3

I suggest using Fold instead of FoldList. The conversion from For loop to Fold loop is algorithmic, could be done even programmatically. The definiton of Fold is:

Fold[f, x, {a, b, ...}] returns the last of the list {x, f[x,a], f[f[x, a], b], ...} (that is, it returns the last element of FoldList[f, x, {a, b, ...}]).

First, identify the function f that you want to use in Fold.

(* parameters *)
S0 = 10000000; a = 0.00000005; b = 0.01; k = 1.2;
(* initial values *)
{views, S, inf} = {{{0, 0}}, S0, 1};
update[{S_, inf_, views_}, t_] := Module[{dS, dinf, noviews},
   dS = -a*inf*S;
   dinf = a*inf*S - b*inf;
   noviews = Last[Last[views]] + k (inf + dinf);
   (* update will return the updated values in a list, as it accepts the same first argument *)
   {S + dS, inf + dinf, Append[views, {t, noviews}]}
   ];
For[t = 1, t < 365, t++,
  {S, inf, views} = update[{S, inf, views}, t]
  ];

Now it is easy to convert it to Fold. Define t as a range of values that the function will iterate through.

{views, S, inf} = {{{0, 0}}, S0, 1}; (* these must be reset *)
ListPlot@Last@Fold[update, {S, inf, views}, Range[1, 364]]