RK-382A Neural ODE Composition Invariance Demo

This code is a small “bridge” between two ideas:

Neural ODE (the model being solved)
The system has a hidden state called h that changes over time.
The rule for how it changes is a neural-style function called f_theta.
So at any moment, you plug the current time and the current state into f_theta, and it tells you the “instantaneous change” of h.
In this demo, f_theta is a tiny hard-coded neural network:

it mixes time and state using fixed weights,

passes that through a safe activation (softsign),

then outputs the change rate.

Runge–Kutta method (the numerical engine)
Because you usually can’t solve this kind of time-evolution exactly, you approximate it by stepping forward in small time increments.
A Runge–Kutta method does this by computing several stage evaluations of f_theta inside one step.
Each stage is like “peek at the slope using a specific intermediate time and intermediate state.”
Then it combines those stage slopes into the next state value.

What Theorem 382A is demonstrating here

Runge–Kutta methods have “internal stages.” You can relabel the order of stages as long as you also relabel all the method’s coefficients consistently (the A, b, and c tables).
When you do that correctly, you have not changed the actual numerical method—only the labeling of its internal work.

So Theorem 382A, in practical terms, says:

If you swap or reorder stages and you remap every coefficient consistently, the method’s output for a step stays the same (up to tiny floating-point rounding differences).

What the code does, in plain symbolic steps

It defines a tiny Neural ODE: “state h evolves according to f_theta(time, h).”

It builds two Runge–Kutta solvers: one is RK4 (four stages), one is Heun (two stages).

It creates permuted versions of both solvers by reordering their stages and remapping their coefficients.

It then performs a two-step composition: First it advances one step using RK4.

Then it advances one step using Heun, starting from the later time (important because f_theta depends on time).

It runs that same two-step procedure again, but using the permuted-but-equivalent RK4 and Heun.

Finally it prints both results and the absolute difference.

Code
// demo_neuralODE_theorem_382A_main.c (Zorro lite-C)
// Minimal Neural ODE example + Theorem 382A demo:
// - dh/dt = f_theta(t,h) where f_theta is a tiny "neural net" (hard-coded weights)
// - Integrate with composition: RK4 step then Heun RK2 step
// - Build stage-permuted equivalents of both methods
// - Show composed result is unchanged (Theorem 382A)
// Runs in main() (no bars/history). No malloc. No typedef unsigned.

#define MAXS 4
#define TRI_A 6   // for MAXS=4: 4*3/2 = 6

static var G_k[MAXS];

// ---- RK tableaus: Heun2 (M1) and RK4 (M2) + permuted versions (M1B, M2B)
static int   M1_s = 0;  static float M1_A[TRI_A];  static float M1_b[MAXS];  static float M1_c[MAXS];
static int   M2_s = 0;  static float M2_A[TRI_A];  static float M2_b[MAXS];  static float M2_c[MAXS];
static int   M1B_s = 0; static float M1B_A[TRI_A]; static float M1B_b[MAXS]; static float M1B_c[MAXS];
static int   M2B_s = 0; static float M2B_A[TRI_A]; static float M2B_b[MAXS]; static float M2B_c[MAXS];

// ----------------- Triangular indexing (explicit RK uses only i>j) -----------------
int Aidx(int i, int j)
{
  return (i*(i-1))/2 + j;
}

var getA(float* Atri, int i, int j)
{
  if(i <= j) return 0;
  return (var)Atri[Aidx(i,j)];
}

void setA(float* Atri, int i, int j, var v)
{
  if(i <= j) return;
  Atri[Aidx(i,j)] = (float)v;
}

void zeroTableau(int* s, float* Atri, float* b, float* c)
{
  int i;
  *s = 0;
  for(i=0;i<TRI_A;i++) Atri[i] = 0;
  for(i=0;i<MAXS;i++) { b[i]=0; c[i]=0; }
}

// ----------------- "Neural network" for the Neural ODE -----------------
// Hidden state: h(t) (scalar for simplicity)
// dh/dt = f_theta(t,h)
//
// We avoid fancy ops and use a safe activation: softsign(z) = z / (1 + abs(z))
var softsign(var z)
{
  return z / (1.0 + abs(z));
}

// Hard-coded weights (this stands in for a trained NN)
var f_theta(var t, var h)
{
  // "features" could be more inputs; here we just use t and h
  // layer1: z = w_h*h + w_t*t + b
  var z = 1.20*h + 0.70*t - 0.30;

  // activation
  var a = softsign(z);

  // output layer: dh/dt = w*a + b2
  return 0.90*a + 0.10;
}

// ----------------- One explicit RK step for Neural ODE -----------------
var rk_step(int s, float* Atri, float* b, float* c, var t, var y, var h)
{
  int i,j;

  if(s < 1) return y;
  if(s > MAXS) s = MAXS;

  for(i=0;i<s;i++) G_k[i] = 0;

  for(i=0;i<s;i++)
  {
    var yi = y;
    for(j=0;j<i;j++)
      yi += h * getA(Atri,i,j) * G_k[j];

    G_k[i] = f_theta(t + (var)c[i]*h, yi);
  }

  var y1 = y;
  for(i=0;i<s;i++)
    y1 += h * (var)b[i] * G_k[i];

  return y1;
}

// ----------------- Composition: do method2 step, then method1 step -----------------
// IMPORTANT for time-dependent ODE: the second step starts at t+h.
var compose_step(int s1,float* A1,float* b1,float* c1,
                 int s2,float* A2,float* b2,float* c2,
                 var t, var y, var h)
{
  var y_mid = rk_step(s2,A2,b2,c2,t,   y,     h);
  var y_out = rk_step(s1,A1,b1,c1,t+h, y_mid, h);
  return y_out;
}

// ----------------- Stage permutation (build equivalent tableau) -----------------
void permute_method_into(int s, float* A, float* b, float* c,
                         int* p,
                         int* sOut, float* Aout, float* bout, float* cout)
{
  int i,j;

  *sOut = s;
  if(*sOut < 0) *sOut = 0;
  if(*sOut > MAXS) *sOut = MAXS;

  for(i=0;i<TRI_A;i++) Aout[i] = 0;
  for(i=0;i<MAXS;i++)  { bout[i]=0; cout[i]=0; }

  // permute b,c
  for(i=0;i<*sOut;i++)
  {
    int pi = p[i];
    if(pi < 0) pi = 0;
    if(pi >= s) pi = s-1;
    bout[i] = b[pi];
    cout[i] = c[pi];
  }

  // permute A (lower triangle only)
  for(i=0;i<*sOut;i++)
    for(j=0;j<i;j++)
    {
      int pi = p[i];
      int pj = p[j];
      if(pi < 0) pi = 0;
      if(pj < 0) pj = 0;
      if(pi >= s) pi = s-1;
      if(pj >= s) pj = s-1;
      setA(Aout, i, j, getA(A, pi, pj));
    }
}

// ----------------- Build Heun RK2 and classic RK4 -----------------
void make_heun_rk2()
{
  zeroTableau(&M1_s, M1_A, M1_b, M1_c);
  M1_s = 2;
  M1_c[0] = 0.0; M1_c[1] = 1.0;
  setA(M1_A, 1, 0, 1.0);
  M1_b[0] = 0.5; M1_b[1] = 0.5;
}

void make_rk4()
{
  zeroTableau(&M2_s, M2_A, M2_b, M2_c);
  M2_s = 4;

  M2_c[0]=0.0; M2_c[1]=0.5; M2_c[2]=0.5; M2_c[3]=1.0;
  setA(M2_A, 1, 0, 0.5);
  setA(M2_A, 2, 1, 0.5);
  setA(M2_A, 3, 2, 1.0);

  M2_b[0]=(float)(1./6.);
  M2_b[1]=(float)(1./3.);
  M2_b[2]=(float)(1./3.);
  M2_b[3]=(float)(1./6.);
}

void buildAll()
{
  make_heun_rk2();
  make_rk4();

  // Example stage permutations:
  // Heun has 2 stages: swap them (1,0)
  int p1[MAXS]; p1[0]=1; p1[1]=0; p1[2]=0; p1[3]=0;

  // RK4 has 4 stages: swap middle two (0,2,1,3)
  int p2[MAXS]; p2[0]=0; p2[1]=2; p2[2]=1; p2[3]=3;

  permute_method_into(M1_s, M1_A, M1_b, M1_c, p1, &M1B_s, M1B_A, M1B_b, M1B_c);
  permute_method_into(M2_s, M2_A, M2_b, M2_c, p2, &M2B_s, M2B_A, M2B_b, M2B_c);
}

// ----------------- Entry point -----------------
function main()
{
  buildAll();

  // Neural ODE initial condition (hidden state)
  var t0 = 0.0;
  var h0 = 0.25;

  // One composed "macro step": RK4 step of size dt, then Heun step of size dt
  // Total time advanced = 2*dt because we advance t by dt between the two steps.
  var dt = 0.10;

  // Original composition
  var hA = compose_step(M1_s,  M1_A,  M1_b,  M1_c,
                        M2_s,  M2_A,  M2_b,  M2_c,
                        t0, h0, dt);

  // Permuted-but-equivalent composition (Theorem 382A says this should match)
  var hB = compose_step(M1B_s, M1B_A, M1B_b, M1B_c,
                        M2B_s, M2B_A, M2B_b, M2B_c,
                        t0, h0, dt);

  printf("\nNeural ODE demo: dh/dt = f_theta(t,h)");
  printf("\nStart: t=%.3f  h=%.17g", t0, h0);
  printf("\nStep:  RK4(dt) then Heun(dt), dt=%.3f (total time +%.3f)", dt, 2*dt);

  printf("\n\nhA = compose(original RK4, original Heun)   = %.17g", hA);
  printf("\nhB = compose(permuted RK4, permuted Heun)   = %.17g", hB);
  printf("\nabs diff = %.3e\n", abs(hA - hB));

  quit("done");
}

Attached Files
NeuralODE01.zip (3 downloads)
ZorroGPT - https://bit.ly/3Gbsm4S