Source: chemistry/RungeKutta4.js

'use strict'

/**
 * The fourth order Runge-Kutta integration method
 *
 * @class RungeKutta4
 * @param {Function} derives Differential equations which needed to integrate
 * @param {Number}   t0  Initial value problem: starting time
 * @param {Array}    yStart  Initial value problem: y0
 * @param {Number}   h       Each step-size value
 */
var RungeKutta4 = function(derives, t0, yStart, h) {
    this.derives    = derives
    this.t          = t0
    this.y          = yStart || []
    this.dimension  = this.y.length
    this.h          = h || 0.01

    // cache the k1, k2, k3, k4 of each step
    this._k1
    this._k2
    this._k3
    this._k4
}

RungeKutta4.prototype = {
    /**
     * Calculate each step according to step-size h
     *
     * @return {Array} calculated result at this.t
     */
    step: function() {
        var derives     = this.derives,
            t           = this.t,
            dimension   = this.dimension,
            h           = this.h

        var i, _y = []

        // Alias: f() <=> this.derives()
        //        Xn  <=> this.t
        //        Yn  <=> this.y
        //        H   <=> this.h
        //        K1  <=> this._k1

        // calculate _k1: K1 = f(t, Yn)
        this._k1 = derives(t, this.y)

        // calculate _k2: K2 = f(Xn + 0.5 * H, Yn + 0.5 * H * K1)
        for (i = 0; i < dimension; i++) {
            _y[i] = this.y[i] + h * 0.5 * this._k1[i]
        }
        this._k2 = derives(t + h * 0.5, _y)

        // calculate _k3: K3 = f(Xn + 0.5 * H, Yn + 0.5 * H * K2)
        for (i = 0; i < dimension; i++) {
            _y[i] = this.y[i] + h * 0.5 * this._k2[i]
        }
        this._k3 = derives(t + h * 0.5, _y)

        // calculate _k4: K4 = f(Xn + H, Yn + H * K3)
        for (i = 0; i < dimension; i++) {
            _y[i] = this.y[i] + h * this._k3[i]
        }
        this._k4 = derives(t + h, _y)

        // calculate yNext: Y(n + 1) = Yn + H / 6 * (K1 + 2 * K2 + 2 * K3 + K4)
        for (i = 0; i < dimension; i++) {
            this.y[i] += h / 6 * (this._k1[i] + 2 * this._k2[i] + 2 * this._k3[i] + this._k4[i])
        }

        // calculate xNext: X(n + 1) = Xn + H
        this.t += h

        return this.y
    },

    /**
     * Calculate according step times
     *
     * @param  {Number} steps Iterative times
     * @return {Array}        Calculated result at this.t
     */
    steps: function(steps) {
        for (var i = 0; i < steps; i++) {
            this.step()
        }

        return this.y
    },

    /**
     * Calculate according to xEnd
     *
     * @param  {Number} tEnd final t
     * @return {Array}       Calculated result at xEnd
     */
    end: function(tEnd) {
        const wholeSteps = Math.floor((tEnd - this.t) / this.h)
        this.steps(wholeSteps);
        if (this.t < tEnd) {
            const oldH = this.h;
            const newH = tEnd - this.t;
            this.h = newH;
            this.step();
            this.h = oldH;
        }
        return this.y;
    },

    deltaT: function(dT) {
        const t0 = this.t;
        const wholeSteps = Math.floor((dT) / this.h);
        this.steps(wholeSteps);
        if (this.t < (t0 + dT)) {
            const oldH = this.h;
            const newH = t0 + dT - this.t;
            this.h = newH;
            this.step();
            this.h = oldH;
        }
        return this.y;
    }

}

module.exports = RungeKutta4;