diff options
| -rw-r--r-- | neural_ode.c | 156 |
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; } |