Add continuous normalizing flow

1 files changed, 128 insertions, 0 deletions

diff --git a/src/cnf_train.c b/src/cnf_train.c
new file mode 100644
index 0000000..21b357f
--- /dev/null
+++ b/src/cnf_train.c
@@ -0,0 +1,128 @@
+#include "cnf_train.h"
+#include "cnf.h"
+#include "adam.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <math.h>
+
+static const double TARGET_MU[4][2] = {
+    { 1.5,  0.0}, {-1.5,  0.0},
+    { 0.0,  1.5}, { 0.0, -1.5}
+};
+static const double TARGET_SIGMA2 = 0.2;
+static const int    TARGET_K      = 4;
+
+static double log_p_target(const double *z) {
+    double lc[4], mx = -1e300;
+    for (int k = 0; k < TARGET_K; k++) {
+        double dx = z[0] - TARGET_MU[k][0], dy = z[1] - TARGET_MU[k][1];
+        lc[k] = -0.5 * (dx * dx + dy * dy) / TARGET_SIGMA2;
+        if (lc[k] > mx) mx = lc[k];
+    }
+    double s = 0.0;
+    for (int k = 0; k < TARGET_K; k++) s += exp(lc[k] - mx);
+    return mx + log(s / (double)TARGET_K)
+             - log(2.0 * M_PI * TARGET_SIGMA2);
+}
+
+static void grad_log_p_target(const double *z, double *g) {
+    double lc[4], mx = -1e300;
+    for (int k = 0; k < TARGET_K; k++) {
+        double dx = z[0] - TARGET_MU[k][0], dy = z[1] - TARGET_MU[k][1];
+        lc[k] = -0.5 * (dx * dx + dy * dy) / TARGET_SIGMA2;
+        if (lc[k] > mx) mx = lc[k];
+    }
+    double w[4], wsum = 0.0;
+    for (int k = 0; k < TARGET_K; k++) { w[k] = exp(lc[k] - mx); wsum += w[k]; }
+    g[0] = g[1] = 0.0;
+    for (int k = 0; k < TARGET_K; k++) {
+        double wk = w[k] / wsum;
+        g[0] += wk * (-(z[0] - TARGET_MU[k][0]) / TARGET_SIGMA2);
+        g[1] += wk * (-(z[1] - TARGET_MU[k][1]) / TARGET_SIGMA2);
+    }
+}
+
+/* ---- Training loop ---- */
+void cnf_train_demo(RNG *r) {
+    const int D = 2, H = 32;
+    const int ITERS = 200, BATCH = 16, LOG_EVERY = 20;
+    const double t0 = 0.0, t1 = 1.0;
+    const double atol = 1e-4, rtol = 1e-4;
+    const double LR = 1e-3;
+
+    int nparams = dynmlp_nparams(D, H);
+    double *theta = vec_alloc(nparams);
+    CNF cnf;
+    cnf_init(&cnf, D, H, theta, r);
+    Adam adam = adam_init(nparams, LR, 0.9, 0.999, 1e-8);
+
+    printf("\n--- CNF density matching (4-Gaussian target, D=2, H=%d) ---\n\n", H);
+    printf("%-6s  %-14s  %-10s  %-10s\n", "Iter", "Loss", "FwdNFE", "BwdNFE");
+    printf("--------------------------------------------------\n");
+
+    for (int iter = 1; iter <= ITERS; iter++) {
+        double *dL_dtheta = vec_zeros(nparams);
+        double loss = 0.0;
+        int total_nfe_fwd = 0, total_nfe_bwd = 0;
+
+        for (int b = 0; b < BATCH; b++) {
+            double z0[2];
+            for (int i = 0; i < D; i++) z0[i] = rng_normal(r);
+
+            CNFSampleResult sr = cnf_sample(&cnf, theta, z0, t0, t1, atol, rtol);
+            double *z1 = sr.z1;
+            total_nfe_fwd += sr.nfe;
+
+            double log_pb = -0.5 * (z0[0]*z0[0] + z0[1]*z0[1])
+                            - (double)D * 0.5 * log(2.0 * M_PI);
+
+            double log_pm = log_pb + sr.delta_logp;
+            double log_pt = log_p_target(z1);
+
+             *   = log_pb + delta_logp - log_pt
+             * (KL(p_model || p_target) estimator) */
+            loss += log_pm - log_pt;
+
+            double dL_dz1[2];
+            grad_log_p_target(z1, dL_dz1);
+            dL_dz1[0] = -dL_dz1[0];
+            dL_dz1[1] = -dL_dz1[1];
+
+            CNFBackwardResult br = cnf_backward(&cnf, theta, z1,
+                                                 1.0, dL_dz1,
+                                                 t0, t1, atol, rtol);
+            total_nfe_bwd += br.nfe;
+            vec_add_scaled(dL_dtheta, 1.0 / BATCH, br.dL_dtheta, nparams);
+            free(br.dL_dz0); free(br.dL_dtheta);
+            free(z1);
+        }
+
+        adam_update(&adam, theta, dL_dtheta);
+        free(dL_dtheta);
+
+        if (iter % LOG_EVERY == 0) {
+            printf("%-6d  %-14.4f  %-10d  %-10d\n",
+                   iter, loss / BATCH,
+                   total_nfe_fwd / BATCH, total_nfe_bwd / BATCH);
+            fflush(stdout);
+        }
+    }
+
+    printf("\n--------------------------------------------------\n");
+
+    const int EVAL_N = 200;
+    double eval_loss = 0.0;
+    for (int i = 0; i < EVAL_N; i++) {
+        double z0[2];
+        for (int j = 0; j < D; j++) z0[j] = rng_normal(r);
+        CNFSampleResult sr = cnf_sample(&cnf, theta, z0, t0, t1, atol, rtol);
+        eval_loss += log_p_target(sr.z1);
+        free(sr.z1);
+    }
+    printf("Eval mean log p_target of samples: %.4f\n", eval_loss / EVAL_N);
+    printf("Total parameters: %d\n", nparams);
+
+    adam_free(&adam);
+    free(theta);
+}