From 07974279e12aa086e37490c610596cc3d0d41085 Mon Sep 17 00:00:00 2001 From: y-jan137 Date: Mon, 16 Mar 2026 16:07:35 +0300 Subject: Reorganize --- 1806.07366v5.pdf | Bin 0 -> 3990592 bytes README.md | 19 ++ neural_ode.c | 644 +++++++++++++++++++++++++++---------------------------- 3 files changed, 335 insertions(+), 328 deletions(-) create mode 100644 1806.07366v5.pdf diff --git a/1806.07366v5.pdf b/1806.07366v5.pdf new file mode 100644 index 0000000..ea7f407 Binary files /dev/null and b/1806.07366v5.pdf differ diff --git a/README.md b/README.md index 7d58cb3..52440a9 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,22 @@ # Neural ODE --- + A C implementation of [Neural Ordinary Differential Equations by Chen et al. (2018)](https://arxiv.org/abs/1806.07366) + +## Build + +```bash +gcc -O2 -Wall -Wextra -o neural_ode neural_ode.c +``` + +## Run +```bash +./neural_ode +``` + + + + + + + diff --git a/neural_ode.c b/neural_ode.c index 50fc098..a226906 100644 --- a/neural_ode.c +++ b/neural_ode.c @@ -4,11 +4,12 @@ #include #include #include -#include -typedef struct { - uint64_t state; -} RNG; +/* ============================================================ + § Utilities: RNG, memory, vector/matrix ops + ============================================================ */ + +typedef struct { uint64_t state; } RNG; static RNG rng_init(uint64_t seed) { RNG r; @@ -38,120 +39,60 @@ static double rng_normal(RNG *r) { static void *xmalloc(size_t n) { void *p = malloc(n); - if (!p) { - fprintf(stderr, "fatal: malloc(%zu) failed\n", n); - abort(); - } + if (!p) { fprintf(stderr, "fatal: malloc(%zu) failed\n", n); abort(); } return p; } static void *xcalloc(size_t count, size_t size) { void *p = calloc(count, size); - if (!p) { - fprintf(stderr, "fatal: calloc(%zu, %zu) failed\n", count, size); - abort(); - } + if (!p) { fprintf(stderr, "fatal: calloc(%zu, %zu) failed\n", count, size); abort(); } return p; } -static double *vec_alloc(int n) { - return (double *)xmalloc((size_t)n * sizeof(double)); -} - -static double *vec_zeros(int n) { - return (double *)xcalloc((size_t)n, sizeof(double)); -} - -static void vec_copy(const double *src, double *dst, int n) { - memcpy(dst, src, (size_t)n * sizeof(double)); -} - -static void vec_zero(double *v, int n) { - memset(v, 0, (size_t)n * sizeof(double)); -} - -static void vec_axpy(const double *a, double alpha, const double *b, - double *dst, int n) { - for (int i = 0; i < n; i++) - dst[i] = a[i] + alpha * b[i]; -} +static double *vec_alloc(int n) { return (double *)xmalloc((size_t)n * sizeof(double)); } +static double *vec_zeros(int n) { return (double *)xcalloc((size_t)n, sizeof(double)); } +static void vec_zero(double *v, int n) { memset(v, 0, (size_t)n * sizeof(double)); } +static void vec_copy(const double *src, double *dst, int n) { memcpy(dst, src, (size_t)n * sizeof(double)); } static void vec_add_scaled(double *dst, double alpha, const double *v, int n) { - for (int i = 0; i < n; i++) - dst[i] += alpha * v[i]; -} - -static void vec_scale(const double *v, double alpha, double *dst, int n) { - for (int i = 0; i < n; i++) - dst[i] = alpha * v[i]; + for (int i = 0; i < n; i++) dst[i] += alpha * v[i]; } static double vec_dot(const double *a, const double *b, int n) { double s = 0.0; - for (int i = 0; i < n; i++) - s += a[i] * b[i]; + for (int i = 0; i < n; i++) s += a[i] * b[i]; return s; } -static double vec_norm(const double *v, int n) { - return sqrt(vec_dot(v, v, n)); -} - -static void mat_vec(const double *M, const double *x, double *dst, - int rows, int cols) { +static void mat_vec(const double *M, const double *x, double *dst, int rows, int cols) { for (int i = 0; i < rows; i++) { double s = 0.0; - for (int j = 0; j < cols; j++) - s += M[i * cols + j] * x[j]; + for (int j = 0; j < cols; j++) s += M[i * cols + j] * x[j]; dst[i] = s; } } -static void mat_vec_T(const double *M, const double *v, double *dst, - int rows, int cols) { - for (int i = 0; i < rows; i++) { +static void mat_vec_T(const double *M, const double *v, double *dst, int rows, int cols) { + for (int i = 0; i < rows; i++) for (int j = 0; j < cols; j++) dst[j] += M[i * cols + j] * v[i]; - } } static void mat_outer_add(double *M, double alpha, - const double *a, const double *b, - int rows, int cols) { - for (int i = 0; i < rows; i++) { + const double *a, const double *b, int rows, int cols) { + for (int i = 0; i < rows; i++) for (int j = 0; j < cols; j++) M[i * cols + j] += alpha * a[i] * b[j]; - } -} - -static void xavier_init(double *w, int fan_in, int fan_out, RNG *r) { - double limit = sqrt(6.0 / (fan_in + fan_out)); - int n = fan_in * fan_out; - for (int i = 0; i < n; i++) - w[i] = (2.0 * rng_uniform(r) - 1.0) * limit; -} - -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]); -} - -static void act_dtanh(const double *y, double *dst, int n) { - for (int i = 0; i < n; i++) - dst[i] = 1.0 - y[i] * y[i]; -} - - -/* --- DynMLP: f(z, t, θ): R^(D+1) -> R^D via tanh hidden layer --- */ +/* ============================================================ + § DynMLP: f(z, t, θ): R^(D+1) -> R^D + ============================================================ */ typedef struct { - int D; /* state dimension */ - int H; /* hidden dimension */ - int nparams; /* total number of parameters */ + int D; + int H; + int nparams; } DynMLP; #define DYNMLP_W1(D, H) (0) @@ -160,26 +101,24 @@ typedef struct { #define DYNMLP_b2(D, H) ((D + 1) * (H) + (H) + (H) * (D)) static int dynmlp_nparams(int D, int H) { - return (D + 1) * H /* W1 */ - + H /* b1 */ - + H * D /* W2 */ - + D; /* b2 */ + return (D + 1) * H + H + H * D + D; +} + +static void xavier_init(double *w, int fan_in, int fan_out, RNG *r) { + double limit = sqrt(6.0 / (fan_in + fan_out)); + int n = fan_in * fan_out; + for (int i = 0; i < n; i++) + w[i] = (2.0 * rng_uniform(r) - 1.0) * limit; } static void dynmlp_init(DynMLP *net, int D, int H, double *theta, RNG *r) { - net->D = D; - net->H = H; + net->D = D; + net->H = H; net->nparams = dynmlp_nparams(D, H); - - double *W1 = theta + DYNMLP_W1(D, H); - double *b1 = theta + DYNMLP_b1(D, H); - double *W2 = theta + DYNMLP_W2(D, H); - double *b2 = theta + DYNMLP_b2(D, H); - - xavier_init(W1, D + 1, H, r); - bias_init(b1, H); - xavier_init(W2, H, D, r); - bias_init(b2, D); + xavier_init(theta + DYNMLP_W1(D, H), D + 1, H, r); + vec_zero(theta + DYNMLP_b1(D, H), H); + xavier_init(theta + DYNMLP_W2(D, H), H, D, r); + vec_zero(theta + DYNMLP_b2(D, H), D); } static void dynmlp_forward(const DynMLP *net, const double *theta, @@ -190,30 +129,26 @@ static void dynmlp_forward(const DynMLP *net, const double *theta, const double *W2 = theta + DYNMLP_W2(D, H); const double *b2 = theta + DYNMLP_b2(D, H); - double *x = vec_alloc(D + 1); + double *x = vec_alloc(D + 1); double *h_pre = vec_alloc(H); - double *h = vec_alloc(H); + double *h = vec_alloc(H); - /* x = [z; t], length D+1 */ vec_copy(z, x, D); x[D] = t; - /* h_pre = W1 * x + b1, length H */ mat_vec(W1, x, h_pre, H, D + 1); vec_add_scaled(h_pre, 1.0, b1, H); - /* h = tanh(h_pre), length H */ - act_tanh(h_pre, h, H); + for (int i = 0; i < H; i++) h[i] = tanh(h_pre[i]); - /* out = W2 * h + b2, length D */ mat_vec(W2, h, out, D, H); vec_add_scaled(out, 1.0, b2, D); - free(x); - free(h_pre); - free(h); + free(x); free(h_pre); free(h); } +/* Vector-Jacobian product: vjp_z = v^T (∂f/∂z), vjp_theta += v^T (∂f/∂θ) + Note: vjp_theta is accumulated into, not overwritten. */ static void dynmlp_vjp(const DynMLP *net, const double *theta, const double *z, double t, const double *v, double *vjp_z, double *vjp_theta) { @@ -221,108 +156,71 @@ static void dynmlp_vjp(const DynMLP *net, const double *theta, const double *W1 = theta + DYNMLP_W1(D, H); const double *b1 = theta + DYNMLP_b1(D, H); const double *W2 = theta + DYNMLP_W2(D, H); - double *dW1 = vjp_theta + DYNMLP_W1(D, H); double *db1 = vjp_theta + DYNMLP_b1(D, H); double *dW2 = vjp_theta + DYNMLP_W2(D, H); double *db2 = vjp_theta + DYNMLP_b2(D, H); - /* --- recover intermediates --- */ double *x = vec_alloc(D + 1); double *h_pre = vec_alloc(H); double *h = vec_alloc(H); - vec_copy(z, x, D); x[D] = t; mat_vec(W1, x, h_pre, H, D + 1); vec_add_scaled(h_pre, 1.0, b1, H); - act_tanh(h_pre, h, H); + for (int i = 0; i < H; i++) h[i] = tanh(h_pre[i]); - /* --- backward --- */ double *dh = vec_zeros(H); double *dh_pre = vec_alloc(H); double *dx = vec_zeros(D + 1); - mat_vec_T(W2, v, dh, D, H); /* W2: D×H */ - - /* dW2 += outer(v, h), db2 += v */ + mat_vec_T(W2, v, dh, D, H); mat_outer_add(dW2, 1.0, v, h, D, H); vec_add_scaled(db2, 1.0, v, D); - /* dh_pre = dh * (1 - h*h) element-wise */ - act_dtanh(h, dh_pre, H); - for (int i = 0; i < H; i++) - dh_pre[i] *= dh[i]; + for (int i = 0; i < H; i++) dh_pre[i] = (1.0 - h[i] * h[i]) * dh[i]; - /* dx = W1^T * dh_pre */ - mat_vec_T(W1, dh_pre, dx, H, D + 1); /* W1: H×(D+1) */ - - /* dW1 += outer(dh_pre, x), db1 += dh_pre */ + mat_vec_T(W1, dh_pre, dx, H, D + 1); mat_outer_add(dW1, 1.0, dh_pre, x, H, D + 1); vec_add_scaled(db1, 1.0, dh_pre, H); - /* vjp_z = dx[0..D-1] */ vec_copy(dx, vjp_z, D); - free(x); - free(h_pre); - free(h); - free(dh); - free(dh_pre); - free(dx); + free(x); free(h_pre); free(h); free(dh); free(dh_pre); free(dx); } -/* --- ODE Solver: Dormand-Prince RK45 --- */ +/* ============================================================ + § RK45 solver + ============================================================ */ 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 */ + double *y; + int nfe; } 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_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 }; - -/* 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, +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, + 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); @@ -330,20 +228,14 @@ ODEResult ode_solve(ode_rhs_fn f, const double *y0, double t0, double t1, double *err = vec_alloc(dim); double *stg = vec_alloc(dim); - ODEResult res; - res.y = vec_alloc(dim); - res.nfe = 0; - + ODEResult res = { vec_alloc(dim), 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; + double t = t0; + double h = 0.01 * (t1 - t0); + int k1_fresh = 0; - for (int step = 0; step < max_steps; step++) { - /* Check termination */ + for (int step = 0; step < 1000000; step++) { if (t1 > t0) { if (t >= t1) break; if (t + h > t1) h = t1 - t; @@ -352,64 +244,39 @@ ODEResult ode_solve(ode_rhs_fn f, const double *y0, double t0, double t1, 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; - } + 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++; + 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++; + 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++; + 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++; + 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++; + + 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]); + + 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++; - /* 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]); + 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])); @@ -418,206 +285,327 @@ ODEResult ode_solve(ode_rhs_fn f, const double *y0, double t0, double t1, } 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); + 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) { - /* 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; + 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); - + 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] */ - + ODEResult res = { vec_alloc(dim * ntimes), 0 }; + vec_copy(y0, res.y, dim); 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], + 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; } -/* --- ODE solver tests --- */ +/* ============================================================ + § Adjoint sensitivity method (Algorithm 1) + ============================================================ */ + +typedef struct { + DynMLP net; + const double *theta; + int state_dim; + int nparams; +} AdjointCtx; + +static void neural_ode_rhs(const double *state, double t, const double *params, + int dim, double *out, void *ctx) { + (void)params; (void)dim; + AdjointCtx *ac = (AdjointCtx *)ctx; + dynmlp_forward(&ac->net, ac->theta, state, t, out); +} + +static void adjoint_dynamics(const double *aug_state, double t, const double *params, + int aug_dim, double *aug_out, void *ctx) { + (void)params; (void)aug_dim; + AdjointCtx *ac = (AdjointCtx *)ctx; + int D = ac->state_dim; + int nparams = ac->nparams; + const double *z = aug_state; + const double *a = aug_state + D; + + dynmlp_forward(&ac->net, ac->theta, z, t, aug_out); + + double *neg_a = vec_alloc(D); + double *vjp_z = vec_alloc(D); + double *vjp_theta = vec_zeros(nparams); + for (int i = 0; i < D; i++) neg_a[i] = -a[i]; + + dynmlp_vjp(&ac->net, ac->theta, z, t, neg_a, vjp_z, vjp_theta); + + vec_copy(vjp_z, aug_out + D, D); + vec_copy(vjp_theta, aug_out + 2 * D, nparams); -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]; + free(neg_a); + free(vjp_z); + free(vjp_theta); } -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]; +typedef struct { + double *dL_dz0; + double *dL_dtheta; + int nfe; +} AdjointResult; + +AdjointResult adjoint_solve(const DynMLP *net, const double *theta, + const double *z0, double t0, double t1, + const double *z1, const double *dL_dz1, + double atol, double rtol) { + (void)z0; + int D = net->D; + int nparams = net->nparams; + int aug_dim = 2 * D + nparams; + + double *aug0 = vec_zeros(aug_dim); + vec_copy(z1, aug0, D); + vec_copy(dL_dz1, aug0 + D, D); + + AdjointCtx ac = { *net, theta, D, nparams }; + ODEResult res = ode_solve(adjoint_dynamics, aug0, t1, t0, + NULL, aug_dim, atol, rtol, &ac); + + AdjointResult ar; + ar.dL_dz0 = vec_alloc(D); + ar.dL_dtheta = vec_alloc(nparams); + ar.nfe = res.nfe; + vec_copy(res.y + D, ar.dL_dz0, D); + vec_copy(res.y + 2 * D, ar.dL_dtheta, nparams); + + free(res.y); free(aug0); + return ar; +} + +typedef struct { + double *z1; + double *dL_dz0; + double *dL_dtheta; + int nfe_forward; + int nfe_backward; +} NeuralODEOutput; + +NeuralODEOutput neural_ode_forward_backward(const DynMLP *net, const double *theta, + const double *z0, double t0, double t1, + const double *target, double atol, double rtol) { + int D = net->D; + AdjointCtx ac = { *net, theta, D, net->nparams }; + ODEResult fwd = ode_solve(neural_ode_rhs, z0, t0, t1, NULL, D, atol, rtol, &ac); + double *dL_dz1 = vec_alloc(D); + + for (int i = 0; i < D; i++) + dL_dz1[i] = fwd.y[i] - target[i]; + + AdjointResult ar = adjoint_solve(net, theta, z0, t0, t1, fwd.y, dL_dz1, atol, rtol); + + NeuralODEOutput out; + out.z1 = fwd.y; + out.dL_dz0 = ar.dL_dz0; + out.dL_dtheta = ar.dL_dtheta; + out.nfe_forward = fwd.nfe; + out.nfe_backward = ar.nfe; + + free(dL_dz1); + return out; +} + +/* ============================================================ + § Tests + ============================================================ */ + +static void rhs_decay(const double *y, double t, const double *p, int d, double *out, void *ctx) { + (void)t; (void)p; (void)d; (void)ctx; + out[0] = -y[0]; +} + +static void rhs_rotation(const double *y, double t, const double *p, int d, double *out, void *ctx) { + (void)t; (void)p; (void)d; (void)ctx; + out[0] = -y[1]; out[1] = y[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); - } + const double atol = 1e-8, rtol = 1e-8, tol = 1e-6; - /* 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); - } + { double y0 = 1.0; + ODEResult r = ode_solve(rhs_decay, &y0, 0.0, 1.0, NULL, 1, atol, rtol, NULL); + double err = fabs(r.y[0] - exp(-1.0)); + printf("ODE test 1 (decay): err=%.2e nfe=%d %s\n", err, r.nfe, err < 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); - } + { 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 < tol ? "PASS" : "FAIL"); + free(r.y); } + + { 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 < tol ? "PASS" : "FAIL"); + free(r.y); } } static void test_dynmlp_gradients(RNG *r) { - const int D = 3; - const int H = 8; - const double EPS = 1e-7; - const double TOL = 1e-5; + const int D = 3, H = 8; + const double EPS = 1e-7, TOL = 1e-5; int np = dynmlp_nparams(D, H); - double *theta = vec_alloc(np); - double *z = vec_alloc(D); - double *v = vec_alloc(D); - double *out_p = vec_alloc(D); - double *out_m = vec_alloc(D); + double *theta = vec_alloc(np); + double *z = vec_alloc(D); + double *v = vec_alloc(D); + double *out_p = vec_alloc(D); + double *out_m = vec_alloc(D); DynMLP net; dynmlp_init(&net, D, H, theta, r); - for (int i = 0; i < D; i++) z[i] = rng_normal(r); - for (int i = 0; i < D; i++) v[i] = rng_normal(r); + for (int i = 0; i < D; i++) z[i] = rng_normal(r); + for (int i = 0; i < D; i++) v[i] = rng_normal(r); double t = rng_normal(r); - /* --- analytical VJP --- */ double *vjp_z = vec_zeros(D); double *vjp_theta = vec_zeros(np); dynmlp_vjp(&net, theta, z, t, v, vjp_z, vjp_theta); - /* --- numerical VJP w.r.t. z --- */ double *num_vjp_z = vec_alloc(D); for (int i = 0; i < D; i++) { double zi = z[i]; - - z[i] = zi + EPS; - dynmlp_forward(&net, theta, z, t, out_p); - - z[i] = zi - EPS; - dynmlp_forward(&net, theta, z, t, out_m); - + z[i] = zi + EPS; dynmlp_forward(&net, theta, z, t, out_p); + z[i] = zi - EPS; dynmlp_forward(&net, theta, z, t, out_m); z[i] = zi; - num_vjp_z[i] = vec_dot(v, out_p, D) - vec_dot(v, out_m, D); - num_vjp_z[i] /= 2.0 * EPS; + num_vjp_z[i] = (vec_dot(v, out_p, D) - vec_dot(v, out_m, D)) / (2.0 * EPS); } - double max_err_z = 0.0; for (int i = 0; i < D; i++) { double e = fabs(vjp_z[i] - num_vjp_z[i]); if (e > max_err_z) max_err_z = e; } - /* --- numerical VJP w.r.t. theta --- */ double *num_vjp_theta = vec_alloc(np); for (int k = 0; k < np; k++) { double tk = theta[k]; + theta[k] = tk + EPS; dynmlp_forward(&net, theta, z, t, out_p); + theta[k] = tk - EPS; dynmlp_forward(&net, theta, z, t, out_m); + theta[k] = tk; + num_vjp_theta[k] = (vec_dot(v, out_p, D) - vec_dot(v, out_m, D)) / (2.0 * EPS); + } + double max_err_theta = 0.0; + for (int k = 0; k < np; k++) { + double e = fabs(vjp_theta[k] - num_vjp_theta[k]); + if (e > max_err_theta) max_err_theta = e; + } + + printf("MLP vjp dL/dz: max_err=%.2e %s\n", max_err_z, max_err_z < TOL ? "PASS" : "FAIL"); + printf("MLP vjp dL/dtheta: max_err=%.2e %s\n", max_err_theta, max_err_theta < TOL ? "PASS" : "FAIL"); + + free(theta); free(z); free(v); free(out_p); free(out_m); + free(vjp_z); free(vjp_theta); free(num_vjp_z); free(num_vjp_theta); +} + +static void test_adjoint_gradients(RNG *r) { + const int D = 2, H = 8; + const double EPS = 1e-5, atol = 1e-7, rtol = 1e-7; + const double t0 = 0.0, t1 = 1.0; + + int np = dynmlp_nparams(D, H); + double *theta = vec_alloc(np); + double *z0 = vec_alloc(D); + double *target = vec_alloc(D); - theta[k] = tk + EPS; - dynmlp_forward(&net, theta, z, t, out_p); + DynMLP net; + dynmlp_init(&net, D, H, theta, r); + for (int i = 0; i < D; i++) z0[i] = rng_normal(r); + for (int i = 0; i < D; i++) target[i] = rng_normal(r); - theta[k] = tk - EPS; - dynmlp_forward(&net, theta, z, t, out_m); + NeuralODEOutput out = neural_ode_forward_backward(&net, theta, z0, t0, t1, + target, atol, rtol); +#define FWD_LOSS(z0_, theta_) ({ \ + AdjointCtx ac_ = { net, (theta_), D, np }; \ + ODEResult r_ = ode_solve(neural_ode_rhs, (z0_), t0, t1, NULL, D, atol, rtol, &ac_); \ + double l_ = 0.0; \ + for (int _i = 0; _i < D; _i++) { double _d = r_.y[_i] - target[_i]; l_ += 0.5*_d*_d; } \ + free(r_.y); l_; \ +}) + + double *num_dL_dtheta = vec_alloc(np); + for (int k = 0; k < np; k++) { + double tk = theta[k]; + theta[k] = tk + EPS; double lp = FWD_LOSS(z0, theta); + theta[k] = tk - EPS; double lm = FWD_LOSS(z0, theta); theta[k] = tk; - num_vjp_theta[k] = vec_dot(v, out_p, D) - vec_dot(v, out_m, D); - num_vjp_theta[k] /= 2.0 * EPS; + num_dL_dtheta[k] = (lp - lm) / (2.0 * EPS); } - + double max_num_theta = 0.0; + for (int k = 0; k < np; k++) + if (fabs(num_dL_dtheta[k]) > max_num_theta) max_num_theta = fabs(num_dL_dtheta[k]); double max_err_theta = 0.0; for (int k = 0; k < np; k++) { - double e = fabs(vjp_theta[k] - num_vjp_theta[k]); + double e = fabs(out.dL_dtheta[k] - num_dL_dtheta[k]); if (e > max_err_theta) max_err_theta = e; } + double rel_theta = max_err_theta / (max_num_theta + 1e-8); + printf("adjoint dL/dtheta: max_rel_err=%.2e nfe_fwd=%d nfe_bwd=%d %s\n", + rel_theta, out.nfe_forward, out.nfe_backward, rel_theta < 1e-3 ? "PASS" : "FAIL"); - printf("max grad error z: %.2e %s\n", - max_err_z, max_err_z < TOL ? "PASS":"FAIL"); - printf("max grad error theta: %.2e %s\n", - max_err_theta, max_err_theta < TOL ? "PASS" : "FAIL"); + double *num_dL_dz0 = vec_alloc(D); + for (int i = 0; i < D; i++) { + double zi = z0[i]; + z0[i] = zi + EPS; double lp = FWD_LOSS(z0, theta); + z0[i] = zi - EPS; double lm = FWD_LOSS(z0, theta); + z0[i] = zi; + num_dL_dz0[i] = (lp - lm) / (2.0 * EPS); + } + double max_num_z0 = 0.0; + for (int i = 0; i < D; i++) + if (fabs(num_dL_dz0[i]) > max_num_z0) max_num_z0 = fabs(num_dL_dz0[i]); + double max_err_z0 = 0.0; + for (int i = 0; i < D; i++) { + double e = fabs(out.dL_dz0[i] - num_dL_dz0[i]); + if (e > max_err_z0) max_err_z0 = e; + } + double rel_z0 = max_err_z0 / (max_num_z0 + 1e-8); + printf("adjoint dL/dz0: max_rel_err=%.2e %s\n", rel_z0, rel_z0 < 1e-3 ? "PASS" : "FAIL"); - free(theta); - free(z); - free(v); - free(out_p); - free(out_m); - free(vjp_z); - free(vjp_theta); - free(num_vjp_z); - free(num_vjp_theta); +#undef FWD_LOSS + + free(out.z1); free(out.dL_dz0); free(out.dL_dtheta); + free(num_dL_dtheta); free(num_dL_dz0); + free(theta); free(z0); free(target); } int main(void) { RNG r = rng_init(42); - test_dynmlp_gradients(&r); test_ode_solver(); + test_dynmlp_gradients(&r); + test_adjoint_gradients(&r); return 0; } -- cgit v1.2.3