EventFlow — Minimal Reproduction

import math
import numpy as np
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

torch.manual_seed(0)
np.random.seed(0)
device = "cpu"

T = 10.0                       # observation window [0, T]
SIGMA = 0.05                   # interpolant noise std
SIGMA_DATA = 0.35              # data cluster width
COUNT_PROBS = {1: 0.2, 2: 0.5, 3: 0.3}
CENTERS = {1: [5.0], 2: [3.0, 7.0], 3: [2.0, 5.0, 8.0]}
N_MAX = max(COUNT_PROBS)

TRAIN_SIZE, VAL_SIZE = 2000, 500
BATCH_SIZE = 128
EPOCHS = 60
LR = 1e-3
ODE_STEPS = 50
N_GEN = 2000

counts = np.array(list(COUNT_PROBS.keys()))
probs = np.array(list(COUNT_PROBS.values()))

def sample_data_sequence():
    n = int(np.random.choice(counts, p=probs))
    c = np.array(CENTERS[n])
    t = np.clip(c + SIGMA_DATA * np.random.randn(n), 0.0, T)
    return np.sort(t).astype(np.float32)

train_data = [sample_data_sequence() for _ in range(TRAIN_SIZE)]
val_data = [sample_data_sequence() for _ in range(VAL_SIZE)]
print("example sequences:", [s.round(2).tolist() for s in train_data[:4]])

def sample_reference_like(gamma_1):
    n = gamma_1.shape[0]
    return np.sort(np.random.rand(n) * T).astype(np.float32)

def collate_balanced_batch(seqs):
    B = len(seqs)
    gamma_hat = np.zeros((B, N_MAX), np.float32)
    target_v = np.zeros((B, N_MAX), np.float32)
    mask = np.zeros((B, N_MAX), np.float32)
    s_all = np.zeros((B, 1), np.float32)
    for i, g1 in enumerate(seqs):
        n = g1.shape[0]
        g0 = sample_reference_like(g1)
        s = np.random.rand()
        gs = (1.0 - s) * g0 + s * g1
        ghat = np.clip(gs + SIGMA * np.random.randn(n), 0.0, T)
        gamma_hat[i, :n] = ghat
        target_v[i, :n] = g1 - g0
        mask[i, :n] = 1.0
        s_all[i, 0] = s
    to = lambda a: torch.from_numpy(a).to(device)
    return to(gamma_hat), to(s_all), to(target_v), to(mask)

def sinusoidal_embedding(s, dim):
    half = dim // 2
    freqs = torch.exp(-math.log(10000.0) * torch.arange(half, device=s.device) / max(half - 1, 1))
    a = s * freqs
    return torch.cat([torch.sin(a), torch.cos(a)], dim=-1)

class VectorFieldMLP(nn.Module):
    def __init__(self, hidden=128, layers=3, time_emb=32):
        super().__init__()
        self.time_emb = time_emb
        in_dim = 5 + time_emb
        net = [nn.Linear(in_dim, hidden), nn.SiLU()]
        for _ in range(layers - 1):
            net += [nn.Linear(hidden, hidden), nn.SiLU()]
        net += [nn.Linear(hidden, 1)]
        self.net = nn.Sequential(*net)

    def forward(self, gamma, s, mask):
        B, Nm = gamma.shape
        count = mask.sum(-1, keepdim=True)
        masked_mean = (gamma * mask).sum(-1, keepdim=True) / count.clamp_min(1.0)
        s_exp = s.expand(B, Nm)
        feats = torch.stack([
            gamma, s_exp, gamma / T,
            (count / N_MAX).expand(B, Nm),
            (masked_mean / T).expand(B, Nm),
        ], dim=-1)
        emb = sinusoidal_embedding(s, self.time_emb).unsqueeze(1).expand(B, Nm, self.time_emb)
        v = self.net(torch.cat([feats, emb], dim=-1)).squeeze(-1)
        return v * mask

def masked_mse(pred, target, mask):
    se = ((pred - target) ** 2) * mask
    return se.sum() / mask.sum().clamp_min(1.0)

model = VectorFieldMLP().to(device)
print("parameters:", sum(p.numel() for p in model.parameters()))

opt = torch.optim.Adam(model.parameters(), lr=LR)

def epoch_pass(data, train):
    model.train(train)
    order = np.random.permutation(len(data)) if train else np.arange(len(data))
    total, nb = 0.0, 0
    for i in range(0, len(data), BATCH_SIZE):
        seqs = [data[j] for j in order[i:i + BATCH_SIZE]]
        gamma, s, tv, mask = collate_balanced_batch(seqs)
        pred = model(gamma, s, mask)
        loss = masked_mse(pred, tv, mask)
        if train:
            opt.zero_grad(); loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step()
        total += loss.item(); nb += 1
    return total / nb

for ep in range(1, EPOCHS + 1):
    tr = epoch_pass(train_data, True)
    with torch.no_grad():
        va = epoch_pass(val_data, False)
    if ep == 1 or ep % 10 == 0 or ep == EPOCHS:
        print(f"epoch {ep:3d}  train_mse={tr:.4f}  val_mse={va:.4f}")

def build_reference(n_list):
    B = len(n_list)
    gamma = np.zeros((B, N_MAX), np.float32)
    mask = np.zeros((B, N_MAX), np.float32)
    for i, n in enumerate(n_list):
        gamma[i, :n] = np.sort(np.random.rand(n) * T)
        mask[i, :n] = 1.0
    return torch.from_numpy(gamma).to(device), torch.from_numpy(mask).to(device)

@torch.no_grad()
def sample_with_ode(n_list, num_steps=ODE_STEPS, return_traj=False):
    gamma, mask = build_reference(n_list)
    B = gamma.shape[0]
    dt = 1.0 / num_steps
    traj = [gamma.clone()]
    for j in range(num_steps):
        s = torch.full((B, 1), j / num_steps, device=device)
        gamma = (gamma + dt * model(gamma, s, mask)).clamp(0.0, T) * mask
        if return_traj:
            traj.append(gamma.clone())
    out = gamma.clone()
    for i in range(B):
        n = int(mask[i].sum())
        if n > 0:
            out[i, :n] = torch.sort(gamma[i, :n]).values
    return (out, mask, torch.stack(traj)) if return_traj else (out, mask)

gen_counts = list(np.random.choice(counts, size=N_GEN, p=probs))
gen, gen_mask = sample_with_ode(gen_counts)

print(f"count-distribution L1 error : {count_l1:.3f}")
print(f"mean-time error  generated  : {mt_gen:.3f}   reference: {mt_ref:.3f}")
print(f"wasserstein      generated  : {w_gen:.3f}   reference: {w_ref:.3f}")

# count distribution (data vs generated)
# pooled event-time histograms (reference vs data vs generated)
# flow trajectories t^k(s) for a few n=3 reference sequences
#   ... see the .ipynb for the full plotting code ...
plt.tight_layout(); plt.show()