Neural ODE

A implementation of Neural Ordinary Differential Equations by Chen et al. (2018) in C.

Build

make

Run

make run

Usage

DynNet

A composable network used as the ODE right-hand side: dz/dt = f_θ(z, t). Build a network by stacking layers, then call dynnet_finalize.

#include "dynnet.h"

DynNet *dynnet_create(int D);
void dynnet_add_linear(DynNet *net, int out_dim);
void dynnet_add_tanh(DynNet *net);
void dynnet_add_softplus(DynNet *net);
void dynnet_add_swish(DynNet *net);
void dynnet_add_layernorm(DynNet *net);
void dynnet_add_residual_begin(DynNet *net);
void dynnet_add_residual_end(DynNet *net);
void dynnet_add_time_concat(DynNet *net);
void dynnet_finalize(DynNet *net);
void dynnet_init_params(DynNet *net, double *theta, RNG *r);
void dynnet_free(DynNet *net);

void dynnet_forward(DynNet *net, const double *theta,
                    const double *z, double t, double *out, double *workspace);
void dynnet_vjp(DynNet *net, const double *theta,
                const double *z, double t, const double *v,
                double *vjp_z, double *vjp_theta, double *workspace);

D: state dimension (input/output).
theta must point to at least net->total_params doubles (caller-owned).
dynnet_init_params fills theta with random Xavier weights.
workspace must be at least net->total_workspace doubles; allocate once and reuse.
time_concat appends t to the current feature vector (increases dim by 1).
Residual blocks: wrap layers between residual_begin / residual_end (dims must match).

ODE solver

Adaptive-step Dormand–Prince (RK45) integrator.

#include "ode_solver.h"

// Integrate from t0 to t1, return state at t1.
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);

// Integrate and record state at each entry of times[0..ntimes-1].
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.y is a heap-allocated array of dim doubles (single solve) or ntimes * dim doubles (multi-time solve). The caller must free(result.y). ODEResult.nfe is the number of function evaluations used.

Neural ODE forward + adjoint backward

#include "adjoint.h"

// Single target: integrate z0 → z1, compute MSE loss, backprop via adjoint.
NeuralODEOutput neural_ode_forward_backward(
    DynNet *net, const double *theta,
    const double *z0, double t0, double t1,
    const double *target, double atol, double rtol,
    int num_checkpoints);

// Multiple observation times: integrate z0 → z(times[0..ntimes-1]).
MultiObsNeuralODEOutput neural_ode_forward_backward_multi(
    DynNet *net, const double *theta,
    const double *z0, const double *times,
    const double *targets, int ntimes,
    double atol, double rtol);

Both outputs carry heap-allocated dL_dz0 (D doubles) and dL_dtheta (nparams doubles). The caller must free each field.

For inference only, set up an AdjointCtx and call ode_solve directly:

double *ws = vec_alloc(net->total_workspace);
double *neg_a = vec_alloc(D);
AdjointCtx ctx = { net, theta, ws, neg_a };
ODEResult fwd = ode_solve(neural_ode_rhs, z0, t0, t1, NULL, D, atol, rtol, &ctx);
// fwd.y holds the predicted state
free(fwd.y);
free(ws);
free(neg_a);

Adam optimizer

#include "adam.h"

Adam adam_init(int nparams, double lr, double beta1, double beta2, double eps);
void adam_update(Adam *a, double *theta, const double *grad);
void adam_free(Adam *a);

Training step

#include "train.h"

TrainStepResult train_step(DynNet *net, double *theta,
                           const double **z0s, const double **targets,
                           double t0, double t1, int batch_size,
                           Adam *adam, double atol, double rtol,
                           int num_checkpoints);

Runs a forward+backward pass over a mini-batch, averages gradients, and calls adam_update. Returns the mean MSE loss and forward/backward NFE counts.

Continuous Normalizing Flow (CNF)

#include "cnf.h"

void cnf_init(CNF *cnf, int D, int H, double *theta, RNG *r);

// Sample noise z0 forward to data space z1.
CNFSampleResult  cnf_sample  (const CNF *cnf, const double *theta,
                               const double *z0, double t0, double t1,
                               double atol, double rtol);

// Evaluate log-probability of a data point z1.
CNFLogProbResult cnf_log_prob(const CNF *cnf, const double *theta,
                               const double *z1, double t0, double t1,
                               double atol, double rtol);

// Backpropagate through the CNF.
CNFBackwardResult cnf_backward(const CNF *cnf, const double *theta,
                                const double *z1,
                                double dL_dlogp, const double *dL_dz1_in,
                                double t0, double t1,
                                double atol, double rtol);

CNFSampleResult.z1 and CNFLogProbResult.z0 are heap-allocated; the caller must free them. CNFBackwardResult.dL_dz0 and .dL_dtheta must also be freed.

Utilities

#include "utils.h"

RNG rng_init(uint64_t seed);   // seed the PRNG
double rng_normal(RNG *r);     // standard normal sample

double *vec_alloc(int n);      // malloc n doubles
double *vec_zeros(int n);      // calloc n doubles

End-to-end example

#include "utils.h"
#include "dynnet.h"
#include "adjoint.h"
#include "adam.h"
#include "train.h"
#include <stdio.h>
#include <stdlib.h>

int main(void) {
    const int D = 2, H = 32;
    const double t0 = 0.0, t1 = 1.5;

    RNG r = rng_init(42);

    DynNet *net = dynnet_create(D);
    dynnet_add_time_concat(net);
    dynnet_add_linear(net, H);
    dynnet_add_tanh(net);
    dynnet_add_linear(net, D);
    dynnet_finalize(net);

    int nparams = net->total_params;
    double *theta = vec_alloc(nparams);
    dynnet_init_params(net, theta, &r);

    Adam adam = adam_init(nparams, 1e-3, 0.9, 0.999, 1e-8);

    // training data
    double z0[2] = {1.0, 0.0};
    double target[2] = {0.0, 1.0};
    const double *z0s[]  = {z0};
    const double *tgts[] = {target};

    // training loop
    for (int iter = 1; iter <= 500; iter++) {
        TrainStepResult res = train_step(net, theta, z0s, tgts,
                                         t0, t1, 1, &adam,
                                         1e-3, 1e-3, 10);
        if (iter % 100 == 0)
            printf("iter %d  loss=%.4f\n", iter, res.loss);
    }

    // inference
    double *ws    = vec_alloc(net->total_workspace);
    double *neg_a = vec_alloc(D);
    AdjointCtx ctx = { net, theta, ws, neg_a };
    ODEResult fwd = ode_solve(neural_ode_rhs, z0, t0, t1,
                              NULL, D, 1e-5, 1e-5, &ctx);
    free(fwd.y);
    free(ws);
    free(neg_a);

    free(theta);
    adam_free(&adam);
    dynnet_free(net);
    return 0;
}

TODOs

[ ] Data and training visualization (matplotlib)