1 files changed, 73 insertions, 0 deletions
diff --git a/src/spiral.c b/src/spiral.c
new file mode 100644
index 0000000..5d5cc00
--- /dev/null
+++ b/src/spiral.c
@@ -0,0 +1,73 @@
+#include "spiral.h"
+#include "adjoint.h"
+#include "ode_solver.h"
+
+#include <math.h>
+#include <stdlib.h>
+
+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];
+}
+
+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);
+
+        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;
+}
+
+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);
+}
+
+double evaluate(const DynMLP *net, const double *theta,
+                const Dataset *ds, double t0, double t1,
+                double atol, double rtol) {
+    int D = net->D;
+    Workspace ws = workspace_alloc(D, net->H, net->nparams);
+    AdjointCtx ac = { *net, theta, D, net->nparams, &ws };
+    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);
+    }
+    workspace_free(&ws);
+    return total_loss / (double)ds->num_samples;
+}