aboutsummaryrefslogtreecommitdiff
path: root/src/ode_solver.c
blob: 6c6963daac6478f24b4c286148418448b8bf42d9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
#include "ode_solver.h"
#include "utils.h"

#include <math.h>
#include <stdlib.h>

static const double dp_c[7]  = { 0.0, 1.0/5.0, 3.0/10.0, 4.0/5.0, 8.0/9.0, 1.0, 1.0 };
static const double dp_a2[1] = { 1.0/5.0 };
static const double dp_a3[2] = { 3.0/40.0, 9.0/40.0 };
static const double dp_a4[3] = { 44.0/45.0, -56.0/15.0, 32.0/9.0 };
static const double dp_a5[4] = { 19372.0/6561.0, -25360.0/2187.0, 64448.0/6561.0, -212.0/729.0 };
static const double dp_a6[5] = { 9017.0/3168.0, -355.0/33.0, 46732.0/5247.0, 49.0/176.0, -5103.0/18656.0 };
static const double dp_b[7]  = { 35.0/384.0, 0.0, 500.0/1113.0, 125.0/192.0, -2187.0/6784.0, 11.0/84.0, 0.0 };
static const double dp_e[7]  = {
     35.0/384.0   - 5179.0/57600.0,
     0.0,
     500.0/1113.0 - 7571.0/16695.0,
     125.0/192.0  - 393.0/640.0,
     -2187.0/6784.0 + 92097.0/339200.0,
     11.0/84.0    - 187.0/2100.0,
     -1.0/40.0
};

ODEResult ode_solve(ode_rhs_fn f, const double *y0, double t0, double t1,
                    const double *params, int dim, double atol, double rtol,
                    void *ctx) {
    double *buf = vec_alloc(11 * dim);
    double *k[7];
    for (int i = 0; i < 7; i++) k[i] = buf + i * dim;
    double *y   = buf + 7 * dim;
    double *y5  = buf + 8 * dim;
    double *err = buf + 9 * dim;
    double *stg = buf + 10 * dim;

    ODEResult res = { vec_alloc(dim), 0 };
    vec_copy(y0, y, dim);

    double t = t0;
    double h = 0.01 * (t1 - t0);
    int k1_fresh = 0;

    for (int step = 0; step < 1000000; step++) {
        if (t1 > t0) {
            if (t >= t1) break;
            if (t + h > t1) h = t1 - t;
        } else {
            if (t <= t1) break;
            if (t + h < t1) h = t1 - t;
        }

        if (!k1_fresh) { f(y, t, params, dim, k[0], ctx); res.nfe++; k1_fresh = 1; }

        for (int i = 0; i < dim; i++)
            stg[i] = y[i] + h * dp_a2[0]*k[0][i];
        f(stg, t + dp_c[1]*h, params, dim, k[1], ctx); res.nfe++;

        for (int i = 0; i < dim; i++)
            stg[i] = y[i] + h * (dp_a3[0]*k[0][i] + dp_a3[1]*k[1][i]);
        f(stg, t + dp_c[2]*h, params, dim, k[2], ctx); res.nfe++;

        for (int i = 0; i < dim; i++)
            stg[i] = y[i] + h * (dp_a4[0]*k[0][i] + dp_a4[1]*k[1][i] + dp_a4[2]*k[2][i]);
        f(stg, t + dp_c[3]*h, params, dim, k[3], ctx); res.nfe++;

        for (int i = 0; i < dim; i++)
            stg[i] = y[i] + h * (dp_a5[0]*k[0][i] + dp_a5[1]*k[1][i]
                                + dp_a5[2]*k[2][i] + dp_a5[3]*k[3][i]);
        f(stg, t + dp_c[4]*h, params, dim, k[4], ctx); res.nfe++;

        for (int i = 0; i < dim; i++)
            stg[i] = y[i] + h * (dp_a6[0]*k[0][i] + dp_a6[1]*k[1][i]
                                + dp_a6[2]*k[2][i] + dp_a6[3]*k[3][i] + dp_a6[4]*k[4][i]);
        f(stg, t + dp_c[5]*h, params, dim, k[5], ctx); res.nfe++;

        for (int i = 0; i < dim; i++)
            y5[i] = y[i] + h * (dp_b[0]*k[0][i] + dp_b[2]*k[2][i]
                               + dp_b[3]*k[3][i] + dp_b[4]*k[4][i] + dp_b[5]*k[5][i]);
        f(y5, t + h, params, dim, k[6], ctx); res.nfe++;

        for (int i = 0; i < dim; i++)
            err[i] = h * (dp_e[0]*k[0][i] + dp_e[2]*k[2][i] + dp_e[3]*k[3][i]
                        + dp_e[4]*k[4][i] + dp_e[5]*k[5][i] + dp_e[6]*k[6][i]);

        double err_sq = 0.0;
        for (int i = 0; i < dim; i++) {
            double sc = atol + rtol * fmax(fabs(y[i]), fabs(y5[i]));
            double e  = err[i] / sc;
            err_sq   += e * e;
        }
        double err_norm = sqrt(err_sq / (double)dim);

        double factor;
        if (err_norm == 0.0) {
            factor = 5.0;
        } else {
            factor = 0.9 * pow(err_norm, -0.2);
            if (factor < 0.2) factor = 0.2;
            if (factor > 5.0) factor = 5.0;
        }

        if (err_norm <= 1.0) {
            vec_copy(y5, y, dim);
            t += h;
            double *tmp = k[0]; k[0] = k[6]; k[6] = tmp;
            h *= factor;
        } else {
            if (factor > 1.0) factor = 1.0;
            h *= factor;
        }
    }

    vec_copy(y, res.y, dim);
    free(buf);

    return res;
}

ODEResult ode_solve_times(ode_rhs_fn f, const double *y0, const double *times,
                          int ntimes, const double *params, int dim,
                          double atol, double rtol, void *ctx) {
    ODEResult res = { vec_alloc(dim * ntimes), 0 };
    vec_copy(y0, res.y, dim);
    for (int i = 1; i < ntimes; i++) {
        ODEResult seg = ode_solve(f, res.y + (i-1)*dim, times[i-1], times[i],
                                  params, dim, atol, rtol, ctx);
        vec_copy(seg.y, res.y + i * dim, dim);
        res.nfe += seg.nfe;

        free(seg.y);
    }
    return res;
}