Reorganize

1 files changed, 315 insertions, 327 deletions

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 <string.h>
 #include <math.h>
 #include <time.h>
-#include <float.h>
 
-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];
-
-    /* dx = W1^T * dh_pre */
-    mat_vec_T(W1, dh_pre, dx, H, D + 1);  /* W1: H×(D+1) */
+    for (int i = 0; i < H; i++) dh_pre[i] = (1.0 - h[i] * h[i]) * dh[i];
 
-    /* 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]);
-
-        /* Stage 7 / FSAL: f at (t+h, y5) */
-        f(y5, t + h, params, dim, k[6], ctx);
-        res.nfe++;
+                               + 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++;
 
-        /* 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); 
+
+    free(neg_a); 
+    free(vjp_z);
+    free(vjp_theta);
+}
+
+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;
 
-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];
+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;
 }
 
-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];
+/* ============================================================
+   § 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;
+    const double atol = 1e-8, rtol = 1e-8, 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);
-    }
+    { 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 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[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); }
 
-    /* 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 = 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);
+
+    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_p);
+    NeuralODEOutput out = neural_ode_forward_backward(&net, theta, z0, t0, t1,
+                                                      target, atol, rtol);
 
-        theta[k] = tk - EPS;
-        dynmlp_forward(&net, theta, z, t, out_m);
+#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;
 }