Add spiral demo

1 files changed, 156 insertions, 0 deletions

diff --git a/neural_ode.c b/neural_ode.c
index 5308ec5..3489c82 100644
--- a/neural_ode.c
+++ b/neural_ode.c
@@ -547,6 +547,81 @@ static TrainStepResult train_step(const DynMLP *net, double *theta,
 }
 
 /* ============================================================
+   § Spiral dataset
+   ============================================================ */
+
+static void spiral_rhs(const double *state, double t, const double *params,
+                       int dim, double *out, void *ctx) {
+    (void)t; (void)dim; (void)ctx;
+    double alpha = params[0];
+    out[0] = alpha * state[1];
+    out[1] = -alpha * state[0];
+}
+
+typedef struct {
+    double **z0;
+    double **target;
+    int num_samples;
+} Dataset;
+
+static Dataset generate_spiral_dataset(int num_samples, double t0, double t1,
+                                       double noise_std, RNG *r) {
+    Dataset ds;
+    ds.num_samples = num_samples;
+    ds.z0     = (double **)xmalloc(num_samples * sizeof(double *));
+    ds.target = (double **)xmalloc(num_samples * sizeof(double *));
+
+    for (int i = 0; i < num_samples; i++) {
+        double alpha = 1.0;
+        double angle  = 2.0 * M_PI * rng_uniform(r);
+        double radius = 0.5 + 1.0 * rng_uniform(r);   /* uniform in [0.5, 1.5] */
+
+        ds.z0[i]     = vec_alloc(2);
+        ds.target[i] = vec_alloc(2);
+
+        ds.z0[i][0] = radius * cos(angle);
+        ds.z0[i][1] = radius * sin(angle);
+
+        ODEResult res = ode_solve(spiral_rhs, ds.z0[i], t0, t1,
+                                  &alpha, 2, 1e-8, 1e-8, NULL);
+        vec_copy(res.y, ds.target[i], 2);
+        free(res.y);
+
+        /* add noise to the initial observation */
+        ds.z0[i][0] += noise_std * rng_normal(r);
+        ds.z0[i][1] += noise_std * rng_normal(r);
+    }
+    return ds;
+}
+
+static void dataset_free(Dataset *ds) {
+    for (int i = 0; i < ds->num_samples; i++) {
+        free(ds->z0[i]);
+        free(ds->target[i]);
+    }
+    free(ds->z0);
+    free(ds->target);
+}
+
+static double evaluate(const DynMLP *net, const double *theta,
+                       const Dataset *ds, double t0, double t1,
+                       double atol, double rtol) {
+    int D = net->D;
+    AdjointCtx ac = { *net, theta, D, net->nparams };
+    double total_loss = 0.0;
+    for (int i = 0; i < ds->num_samples; i++) {
+        ODEResult fwd = ode_solve(neural_ode_rhs, ds->z0[i], t0, t1,
+                                  NULL, D, atol, rtol, &ac);
+        for (int j = 0; j < D; j++) {
+            double d = fwd.y[j] - ds->target[i][j];
+            total_loss += 0.5 * d * d;
+        }
+        free(fwd.y);
+    }
+    return total_loss / (double)ds->num_samples;
+}
+
+/* ============================================================
    § Tests
    ============================================================ */
 
@@ -774,9 +849,90 @@ static void test_training(RNG *r) {
 
 int main(void) {
     RNG r = rng_init(42);
+
+    /* --- sanity checks --- */
     test_ode_solver();
     test_dynmlp_gradients(&r);
     test_adjoint_gradients(&r);
     test_training(&r);
+
+    printf("\n--- Training demo (spiral) ---\n\n");
+
+    r = rng_init((uint64_t)time(NULL));
+
+    const double t0 = 0.0, t1 = 1.5;
+    const double noise_std = 0.1;
+    const double atol_train = 1e-3, rtol_train = 1e-3;
+    const double atol_eval  = 1e-5, rtol_eval  = 1e-5;
+    const int TRAIN_N = 200, TEST_N = 50;
+    const int ITERS = 500, BATCH = 16, LOG_EVERY = 25;
+
+    Dataset train_ds = generate_spiral_dataset(TRAIN_N, t0, t1, noise_std, &r);
+    Dataset test_ds  = generate_spiral_dataset(TEST_N,  t0, t1, noise_std, &r);
+
+    const int D = 2, H = 32;
+    DynMLP net;
+    int nparams = dynmlp_nparams(D, H);
+    double *theta = vec_alloc(nparams);
+    dynmlp_init(&net, D, H, theta, &r);
+
+    Adam adam = adam_init(nparams, 1e-3, 0.9, 0.999, 1e-8);
+
+    printf("%-6s  %-12s  %-12s  %-10s  %-10s",
+           "Iter", "Train Loss", "Test Loss", "Fwd NFE", "Bwd NFE");
+    printf("\n--------------------------------------------------\n");
+
+    const double **batch_z0  = (const double **)xmalloc(BATCH * sizeof(double *));
+    const double **batch_tgt = (const double **)xmalloc(BATCH * sizeof(double *));
+
+    for (int iter = 1; iter <= ITERS; iter++) {
+        for (int b = 0; b < BATCH; b++) {
+            int idx = (int)(rng_next(&r) % (uint64_t)TRAIN_N);
+            batch_z0[b]  = train_ds.z0[idx];
+            batch_tgt[b] = train_ds.target[idx];
+        }
+
+        TrainStepResult res = train_step(&net, theta, batch_z0, batch_tgt,
+                                         t0, t1, BATCH, &adam,
+                                         atol_train, rtol_train);
+
+        if (iter % LOG_EVERY == 0) {
+            double test_loss = evaluate(&net, theta, &test_ds,
+                                        t0, t1, atol_eval, rtol_eval);
+            printf("%-6d  %-12.6f  %-12.6f  %-10d  %-10d\n",
+                   iter, res.loss, test_loss,
+                   res.nfe_fwd, res.nfe_bwd);
+            fflush(stdout);
+        }
+    }
+
+    free(batch_z0);
+    free(batch_tgt);
+
+    printf("\n--------------------------------------------------\n");
+    double final_test_loss = evaluate(&net, theta, &test_ds,
+                                      t0, t1, atol_eval, rtol_eval);
+    printf("Final test loss : %.6f\n", final_test_loss);
+    printf("Total parameters: %d\n", nparams);
+
+    printf("\nSample predictions:\n");
+    AdjointCtx ac = { net, theta, D, nparams };
+    for (int s = 0; s < 5; s++) {
+        int idx = (int)(rng_next(&r) % (uint64_t)TEST_N);
+        ODEResult fwd = ode_solve(neural_ode_rhs,
+                                  test_ds.z0[idx], t0, t1,
+                                  NULL, D, atol_eval, rtol_eval, &ac);
+        printf("  z0=(%.4f, %.4f)  predicted=(%.4f, %.4f)  target=(%.4f, %.4f)\n",
+               test_ds.z0[idx][0], test_ds.z0[idx][1],
+               fwd.y[0], fwd.y[1],
+               test_ds.target[idx][0], test_ds.target[idx][1]);
+        free(fwd.y);
+    }
+
+    free(theta);
+    adam_free(&adam);
+    dataset_free(&train_ds);
+    dataset_free(&test_ds);
+
     return 0;
 }