SIR Model in D3

At long last, I overcame my fear and made the leap into the world of javascript and D3. My first task whenever learning a new language is to implement a simple version of the Susceptible-Infected-Recovered (SIR) model of infectious disease transmission.

When you move the slider marked “Number initially susceptible,” the simulation runs again with 1 initial infectious individual and the specified number of susceptible individuals. Similarly, when you change the value of the basic reproductive number, R0, the simulation also runs again for the new value of R0 and current value of S.

Right now, the model only runs for 20 steps and then stops. But I’d like to make a version that grows and shrinks the width of the bars to accommodate longer runs. But this one’s good enough for now.

This is the final result:

Not too shabby.

This is the code for the simulation model itself. It’s a pretty simple affair:

  function outbreak(N,b) {
      var max_t = 20;
      var I = 1;
      var S = N-1;
      var incidence = [{t:0, I:I}];
      for (t = 1; t < max_t; t++) {
          var p_inf = 1.0-Math.exp(-b*I/N);
          var new_I = 0;
          for (i = 0; i < S; i++) {
              if (Math.random() < p_inf) {
                  new_I++;
              }
          }
          if (new_I == 0) {
              break;
          }
          incidence.push({t:t, I:new_I});
          S -= new_I;
          I = new_I;
      }
      return(incidence);
  }

This is a simple version of a frequency dependent SIR model where the force of infection on each susceptible individual is the product of the number of infectious individuals (I) and the reproductive rate (in this case, b), divided by the size of the total population. We’re assuming here that everyone has a constant infectious period lasting one step of the model, so in this case b is the same as the mythical basic reproductive number R0. Note that the probability of infection p_inf is the exponential cumulative distribution function with rate equal to bI. This accounts for the discretization of time in our model, and translates our continuous-time rate b to a probability.

You’ll see that we loop over every susceptible individual and determine whether they’ll be infected by drawing a random float on [0,1). There are decidedly better ways of doing this, but I wasn’t able to find a good binomial random number generator floating around for javascript, so I went for the brute force approach. This works fine for this toy example, but obviously wouldn’t be so great for a large population.

Most of the guts of the bar chart are pirated from Mike Bostock’s great series of Let’s Make a Bar Chart tutorials, so I won’t go over the definition of the bars, etc in any detail. One thing to note is that the bars highlighted in green are incoming bars added to the DOM after the last simulation, following the general update pattern for d3. This allows for smooth transitions between bars.

Aside from implementing a vanilla SIR model, one of my goals with this was to understand how to effectively use transitions to re-scale the bar heights and axes between simulations, and to use a slider to set and update the input parameters.

To do this, I put an “onchange” listener on the slider for the number of susceptibles as follows:

 d3.select("#S").on("change", function() { update(+this.value, d3.select("#R").node().value/100.0)});

This grabs the value of the slider when the mouse is released and sends it to the update function. Note that we also grab the value for the reproductive number R to pass it to the update function.

But to continuously update the value the variable S next to the slider, we use an “oninput” event as follows:

  d3.select("#S").on("input", function() {
      d3.select("#S-value").text(+this.value);
  })

I did the same thing for R0 but because the slider gives inputs only in integer values, we’ll set the range from 0 to 200 and divide by 100 to get a range from 0 to 2.0 by 0.1 increments.

The update function is itself pretty simple, just generating a new simulation and then re-calling the render function we used to build the plot in the first instance:

  function update(numS,r) {
      o = outbreak(numS,r);
      render(o);
  }

And that’s it! A small first step into D3, but an exciting one as far as I’m concerned. I’d been intimidated by the relative complexity of javascript-based visualization for awhile. But once I took the time to get comfortable with the idea of manipulating the DOM and wrapped my head around CSS selectors, it didn’t seem quite as different from the kind of graphing I’m used to doing in ggplot, especially when it comes to the kind of low-level control you need to make good-looking visualizations.

My next step is going to be to try to make some interactive disease maps in D3, so watch this space…