diff options
| -rw-r--r-- | neural_ode.c | 262 |
1 files changed, 259 insertions, 3 deletions
diff --git a/neural_ode.c b/neural_ode.c index ac6e805..50fc098 100644 --- a/neural_ode.c +++ b/neural_ode.c @@ -36,7 +36,6 @@ static double rng_normal(RNG *r) { return sqrt(-2.0 * log(u1)) * cos(2.0 * M_PI * u2); } - static void *xmalloc(size_t n) { void *p = malloc(n); if (!p) { @@ -55,7 +54,6 @@ static void *xcalloc(size_t count, size_t size) { return p; } - static double *vec_alloc(int n) { return (double *)xmalloc((size_t)n * sizeof(double)); } @@ -137,7 +135,6 @@ static void bias_init(double *b, int n) { vec_zero(b, n); } - static void act_tanh(const double *x, double *dst, int n) { for (int i = 0; i < n; i++) dst[i] = tanh(x[i]); @@ -275,6 +272,262 @@ static void dynmlp_vjp(const DynMLP *net, const double *theta, free(dx); } +/* --- ODE Solver: Dormand-Prince RK45 --- */ + +typedef void (*ode_rhs_fn)(const double *state, double t, const double *params, + int dim, double *out, void *ctx); + +typedef struct { + double *y; /* final state, length dim — caller must free */ + int nfe; /* total number of f evaluations */ +} ODEResult; + +/* Butcher tableau nodes */ +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 +}; + +/* a coefficients, row by row (lower triangular) */ +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 }; + +/* 5th-order weights (b); b7=0 so FSAL: stage-7 input == y5 */ +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 +}; + +/* 4th-order weights (b*) — stored for reference */ +static const double dp_bs[7] = { + 5179.0/57600.0, 0.0, 7571.0/16695.0, 393.0/640.0, + -92097.0/339200.0, 187.0/2100.0, 1.0/40.0 +}; + +/* Error coefficients: e = b - b* */ +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) { + (void)dp_bs; /* stored for reference; dp_e encodes b-b* */ + + double **k = (double **)xmalloc(7 * sizeof(double *)); + for (int i = 0; i < 7; i++) k[i] = vec_alloc(dim); + double *y = vec_alloc(dim); + double *y5 = vec_alloc(dim); + double *err = vec_alloc(dim); + double *stg = vec_alloc(dim); + + ODEResult res; + res.y = vec_alloc(dim); + res.nfe = 0; + + vec_copy(y0, y, dim); + double t = t0; + double h = 0.01 * (t1 - t0); /* sign encodes direction */ + int k1_fresh = 0; + + const double safety = 0.9; + const int max_steps = 1000000; + + for (int step = 0; step < max_steps; step++) { + /* Check termination */ + 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; + } + + /* Stage 1 (skip on FSAL reuse) */ + if (!k1_fresh) { + f(y, t, params, dim, k[0], ctx); + res.nfe++; + k1_fresh = 1; + } + + /* Stage 2 */ + 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++; + + /* Stage 3 */ + 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++; + + /* Stage 4 */ + 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++; + + /* Stage 5 */ + 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++; + + /* Stage 6 */ + 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++; + + /* 5th-order solution y5 (b2=0, b7=0) */ + 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]); + + /* Stage 7 / FSAL: f at (t+h, y5) */ + f(y5, t + h, params, dim, k[6], ctx); + res.nfe++; + + /* Error estimate: e2=0, k[1] not used */ + 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]); + + /* RMS error norm with mixed tolerance scaling */ + 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); + + /* Compute new step size factor */ + double factor; + if (err_norm == 0.0) { + factor = 5.0; + } else { + factor = safety * pow(err_norm, -0.2); + if (factor < 0.2) factor = 0.2; + if (factor > 5.0) factor = 5.0; + } + + if (err_norm <= 1.0) { + /* Accept: advance state, FSAL swap k[0] <-> k[6] */ + vec_copy(y5, y, dim); + t += h; + double *tmp = k[0]; k[0] = k[6]; k[6] = tmp; + h *= factor; + } else { + /* Reject: shrink only */ + if (factor > 1.0) factor = 1.0; + h *= factor; + /* k1_fresh stays 1 — y and t unchanged */ + } + } + + vec_copy(y, res.y, dim); + + for (int i = 0; i < 7; i++) free(k[i]); + free(k); + free(y); + free(y5); + free(err); + free(stg); + + return res; +} + +/* Solve and record state at each time in times[0..ntimes-1]. + times[0] is the start; result->y is dim*ntimes doubles. */ +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; + res.y = vec_alloc(dim * ntimes); + res.nfe = 0; + + vec_copy(y0, res.y, dim); /* state at times[0] */ + + for (int i = 1; i < ntimes; i++) { + const double *cur = res.y + (i - 1) * dim; + ODEResult seg = ode_solve(f, cur, 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; +} + +/* --- ODE solver tests --- */ + +static void rhs_decay(const double *state, double t, const double *params, + int dim, double *out, void *ctx) { + (void)t; (void)params; (void)dim; (void)ctx; + out[0] = -state[0]; +} + +static void rhs_rotation(const double *state, double t, const double *params, + int dim, double *out, void *ctx) { + (void)t; (void)params; (void)dim; (void)ctx; + out[0] = -state[1]; + out[1] = state[0]; +} + +static void test_ode_solver(void) { + const double atol = 1e-8, rtol = 1e-8; + const double check_tol = 1e-6; + + /* Test 1: scalar decay dy/dt = -y, y(0)=1 -> y(1) = e^{-1} */ + { + double y0 = 1.0; + ODEResult r = ode_solve(rhs_decay, &y0, 0.0, 1.0, NULL, 1, atol, rtol, NULL); + double exact = exp(-1.0); + double err = fabs(r.y[0] - exact); + printf("ODE test 1 (decay): err=%.2e nfe=%d %s\n", + err, r.nfe, err < check_tol ? "PASS" : "FAIL"); + free(r.y); + } + + /* Test 2: 2D rotation, one full period -> back to [1, 0] */ + { + double y0[2] = {1.0, 0.0}; + ODEResult r = ode_solve(rhs_rotation, y0, 0.0, 2.0 * M_PI, + NULL, 2, atol, rtol, NULL); + double err = sqrt((r.y[0]-1.0)*(r.y[0]-1.0) + r.y[1]*r.y[1]); + printf("ODE test 2 (rotation): err=%.2e nfe=%d %s\n", + err, r.nfe, err < check_tol ? "PASS" : "FAIL"); + free(r.y); + } + + /* Test 3: backward integration, decay from t=1 to t=0 */ + { + double y0 = exp(-1.0); + ODEResult r = ode_solve(rhs_decay, &y0, 1.0, 0.0, NULL, 1, atol, rtol, NULL); + double err = fabs(r.y[0] - 1.0); + printf("ODE test 3 (backward): err=%.2e nfe=%d %s\n", + err, r.nfe, err < check_tol ? "PASS" : "FAIL"); + free(r.y); + } +} static void test_dynmlp_gradients(RNG *r) { const int D = 3; @@ -360,8 +613,11 @@ static void test_dynmlp_gradients(RNG *r) { free(num_vjp_theta); } + + int main(void) { RNG r = rng_init(42); test_dynmlp_gradients(&r); + test_ode_solver(); return 0; } |