Bokeh JS callback when plotting differential equations

Question

I'm relatively new to Python and Stackoverflow so apologies if this question has already been asked but I've searched for quite a while now and can't really find a solution to what I'm trying to do.

Problem:

I've been trying to create a very basic model of the COVID-19 epidemic to familiarise myself with Python. My intention is to produce a basic SIR model that calculates susceptible, infected and removed individuals in a population. So far so good.

My problem is that I would like the plot to have an interactive slider that alters one of the constants in the differential equation.

I am using Bokeh and am trying to use Javascript Callbacks for this however I am having difficulty with the Javascript. All examples I have seen so far use line equations where y is a function of x, and which are relatively straightforward to code. In my case, since its a system of differential equations I'm not sure how I should be going about this.

I've also tried using scipy but I still encounter the same problem.

Code below. Any help/feedback/suggestions would be greatly appreciated.

Thanks!

from bokeh.layouts import column, row
from bokeh.models import CustomJS, Slider
from bokeh.plotting import ColumnDataSource, figure, output_file, show

t = []

for i in range(200):
    t.append(i)

slst = []
ilst = []
rlst = []

s = 489599/489609
i = 10/489609
r = 0/489609
bet = 0.28
gam = 0.189

for f in range(200):
    ds = (-bet * (s * i))
    di = ((bet * (s * i)) - (gam * i))
    dr = gam * i
    s = s + (ds)
    slst.append(s)
    i = i + (di)
    ilst.append(i)
    r = r + (dr)
    rlst.append(r)

source = ColumnDataSource(data=dict(x=t, y=[], s=slst, i=ilst))

plot = figure(plot_width=400, plot_height=400)
plot.line('x', 'y', source=source, line_width=3, line_alpha=0.6)

callback = CustomJS(args=dict(source=source), code="""

                    ????????????

    """)

slider = Slider(start=0.1, end=4, value=1, step=.1, title="Beta ")
slider.js_on_change('value', callback)

layout = column(slider, plot)

show(layout)

Answer 1

Here's what I've come up with. It's interesting to see that with high values of beta, the susceptible line goes below 0. Maybe I've made a mistake while porting your code to JavaScript - please correct me if so.

from bokeh.core.property.instance import Instance
from bokeh.io import save
from bokeh.layouts import column
from bokeh.model import Model
from bokeh.models import CustomJS, Slider, Callback
from bokeh.plotting import ColumnDataSource, figure

source = ColumnDataSource(data=dict(t=[], s=[], i=[], r=[]))

plot = figure(plot_width=400, plot_height=400)
plot.line('t', 's', source=source, line_width=3, line_alpha=0.6)
plot.line('t', 'i', source=source, line_width=3, line_alpha=0.6, color='orange')
plot.line('t', 'r', source=source, line_width=3, line_alpha=0.6, color='green')

callback = CustomJS(args=dict(source=source), code="""\
    const N = 200;
    let s = 489599 / 489609;
    let i = 10 / 489609;
    let r = 0 / 489609;
    const bet = cb_obj.value;
    const gam = 0.189;
    const tlst = source.data.t = [];
    const slst = source.data.s = [];
    const ilst = source.data.i = [];
    const rlst = source.data.r = [];
    for (let t = 0; t < N; ++t) {
        s -= bet * s * i;
        i += bet * s * i - gam * i;
        r += gam * i;
        tlst.push(t);
        slst.push(s);
        ilst.push(i);
        rlst.push(r);
    }
    source.change.emit();
""")

slider = Slider(start=0.1, end=4, value=1, step=.1, title="Beta ")
slider.js_on_change('value', callback)


class IdleDocObserver(Model):
    """Work around https://github.com/bokeh/bokeh/issues/4272."""
    on_idle = Instance(Callback)

    # language=TypeScript
    __implementation__ = """\
        import {View} from "core/view"
        import {Model} from "model"
        import * as p from "core/properties"

        export class IdleDocObserverView extends View {}

        export namespace IdleDocObserver {
            export type Attrs = p.AttrsOf<Props>
            export type Props = Model.Props & {on_idle: p.Property<any>}
        }

        export interface IdleDocObserver extends IdleDocObserver.Attrs {}

        export class IdleDocObserver extends Model {
            static init_IdleDocObserver(): void {
                this.prototype.default_view = IdleDocObserverView
                this.define<IdleDocObserver.Props>({on_idle: [p.Any]})
            }

            _doc_attached(): void {
                super._doc_attached()
                const execute = () => this.on_idle!.execute(this)
                if (this.document!.is_idle)
                    execute();
                else
                    this.document!.idle.connect(execute);
            }
        }
    """


idle_doc_observer = IdleDocObserver(on_idle=CustomJS(args=dict(callback=callback, slider=slider),
                                                     code="callback.execute(slider);"))

layout = column(slider, plot)
save([idle_doc_observer, layout])

Answer 2

As per my comments, one can use RK4 in a javascript implementation to integrate the ODE in the html document. There appears no way in bokeh to implement javascript functions outside of callbacks, utility functions and common computations. So to avoid code duplication one has to make the one callback universal enough that it can serve for all slider change events. (Alternatively, one could implement a "recompute" button.)

To look more professional, make 2 plots, one for all components, and one for I alone.

# Set up the plots and their data source
source = ColumnDataSource(data=dict(T=[], S=[], I=[], R=[]))

SIR_plot = figure(plot_width=400, plot_height=400)
SIR_plot.line('T', 'S', source=source, legend_label="S", line_width=3, line_alpha=0.6, color='blue')
SIR_plot.line('T', 'I', source=source, legend_label="I", line_width=3, line_alpha=0.6, color='orange')
SIR_plot.line('T', 'R', source=source, legend_label="R", line_width=3, line_alpha=0.6, color='green')

I_plot = figure(plot_width=400, plot_height=400)
I_plot.line('T', 'I', source=source, line_width=3, line_alpha=0.6, color='orange')

Next set up 4 sliders for parameters one is likely to want to influence

# declare the interactive interface elements
trans_rate = Slider(start=0.01, end=0.4, value=0.3, step=.01, title="transmission rate ")
recov_rate = Slider(start=0.01, end=0.4, value=0.1, step=.01, title="recovery rate")

I_init = Slider(start=0.01, end=0.1, value=0.05, step=.002, title="initial infected [proportion] ")
max_time = Slider(start=10, end=200, value=50, step=1, title="time range [days] ")

Now as the main change to the answer of Eugene Pakhomov, make one call-back for all sliders (see bokeh gallery slider demo) and use a vectorized RK4 method to do the ODE integration

callback = CustomJS(args=dict(source=source, I_init=I_init, max_time=max_time, 
                              trans_rate=trans_rate, recov_rate=recov_rate), 
                    code="""\
    let i = I_init.value;
    let s = 1-i;
    let r = 0;
    const bet = trans_rate.value;
    const gam = recov_rate.value;
    let tf = max_time.value;
    const dt = 0.1;
    const tlst = source.data.T = [0];
    const slst = source.data.S = [s];
    const ilst = source.data.I = [i];
    const rlst = source.data.R = [r];

    function odefunc(t,sir) {
        let tr = bet*sir[0]*sir[1];
        let rc = gam*sir[1];
        return [-tr, tr-rc, rc];
    }
    let sir = [s,i,r];
    for (let t = 0; t < tf; t+=dt) {
        sir = RK4Step(t,sir,dt);
        tlst.push(t+dt);
        slst.push(sir[0]);
        ilst.push(sir[1]);
        rlst.push(sir[2]);
    }
    source.change.emit();

    function axpy(a,x,y) { 
        // returns a*x+y for arrays x,y of the same length
        var k = y.length >>> 0;
        var res = new Array(k);
        while(k-->0) { res[k] = y[k] + a*x[k]; }
        return res;
    }

    function RK4Step(t,y,h) {
        var k0 = odefunc(t      ,               y );
        var k1 = odefunc(t+0.5*h, axpy(0.5*h,k0,y));
        var k2 = odefunc(t+0.5*h, axpy(0.5*h,k1,y));
        var k3 = odefunc(t+    h, axpy(    h,k2,y));
        // ynext = y+h/6*(k0+2*k1+2*k2+k3);
        return axpy(h/6,axpy(1,k0,axpy(2,k1,axpy(2,k2,k3))),y);
    }

""")
trans_rate.js_on_change('value', callback)
recov_rate.js_on_change('value', callback)

I_init.js_on_change('value', callback)
max_time.js_on_change('value', callback)

Finally, string all together in some layout

# generate the layout

parameters_panel = column(trans_rate, recov_rate)
initials_panel = column(I_init,max_time)

plots = row(SIR_plot, I_plot)
inputs = row(parameters_panel, initials_panel)

simulation = column(plots, inputs)

show(simulation)

To avoid the initial empty plots I refer to the answer of Eugene Pakhomov, as it is the plots appear after the first slider is moved.