#!/usr/bin/env python3 """ benchmark_table.py Core baseline utilities for held-out interior-day population imputation. The functions here can be imported by 02_bracket_heldout.py. """ import math import numpy as np import torch import torch.nn as nn import ot # ============================================================ # Neural velocity field # ============================================================ def fourier(x, n=8): f = ( 2.0 ** torch.arange(n, device=x.device, dtype=x.dtype) ) * math.pi a = x * f return torch.cat( [ torch.sin(a), torch.cos(a), ], dim=-1, ) class Field(nn.Module): def __init__(self, d, h=256, nt=8): super().__init__() self.nt = nt self.net = nn.Sequential( nn.Linear(d + 2 * nt, h), nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, h), nn.SiLU(), nn.Linear(h, d), ) def forward(self, z, tau): if tau.ndim == 1: tau = tau[:, None] return self.net( torch.cat( [ z, fourier(tau, self.nt), ], dim=-1, ) ) # ============================================================ # Optimal transport # ============================================================ def ot_plan(X0, X1): """ Balanced quadratic-cost OT plan. Exact EMD is attempted first. If it fails, fall back to a simple Sinkhorn-like scaling procedure. """ X0 = np.asarray(X0, dtype=np.float64) X1 = np.asarray(X1, dtype=np.float64) C = ( ( X0[:, None, :] - X1[None, :, :] ) ** 2 ).sum(-1) a = np.ones(len(X0), dtype=np.float64) / len(X0) b = np.ones(len(X1), dtype=np.float64) / len(X1) try: return ot.emd( a, b, C, numItermax=2_000_000, ) except Exception: Cn = C / (C.max() + 1e-12) K = np.exp( -Cn / 0.05 ) u = np.ones_like(a) v = np.ones_like(b) for _ in range(300): u = a / ( K @ v + 1e-12 ) v = b / ( K.T @ u + 1e-12 ) return ( u[:, None] * K * v[None, :] ) def align_all(clouds, rng, batch=128): """ Sample approximately OT-aligned tuples across all retained clouds. Returns ------- Z : ndarray Shape (K, B, d) """ B = min( batch, min( len(c) for c in clouds ), ) first_idx = rng.choice( len(clouds[0]), B, replace=False, ) z = [ clouds[0][first_idx] ] for k in range( 1, len(clouds), ): nxt_idx = rng.choice( len(clouds[k]), B, replace=False, ) nxt = clouds[k][nxt_idx] P = ot_plan( z[-1], nxt, ) Pn = P / ( P.sum( axis=1, keepdims=True, ) + 1e-12 ) selected = np.array( [ rng.choice( B, p=Pn[i], ) for i in range(B) ] ) z.append( nxt[selected] ) return np.stack( z, axis=0, ) # ============================================================ # Spline interpolation # ============================================================ def pchip(x, Y, t): """ Vectorized shape-preserving cubic Hermite interpolation. Parameters ---------- x : (K,) Knot times. Y : (K, B, d) Aligned tuples. t : (B,) Evaluation times. Returns ------- mu : (B, d) Interpolated value. mup : (B, d) Time derivative. """ K = len(x) h = np.diff(x) d = ( Y[1:] - Y[:-1] ) / h[:, None, None] m = np.zeros_like(Y) # Interior PCHIP slopes for k in range( 1, K - 1, ): d0 = d[k - 1] d1 = d[k] h0 = h[k - 1] h1 = h[k] w0 = 2 * h1 + h0 w1 = h1 + 2 * h0 mono = ( d0 * d1 ) > 0 with np.errstate( divide="ignore", invalid="ignore", ): mk = ( w0 + w1 ) / ( w0 / d0 + w1 / d1 ) m[k] = np.where( mono, mk, 0.0, ) def endpoint_slope( h0, h1, d0, d1, ): me = ( ( 2 * h0 + h1 ) * d0 - h0 * d1 ) / ( h0 + h1 ) me = np.where( np.sign(me) != np.sign(d0), 0.0, me, ) me = np.where( ( np.sign(d0) != np.sign(d1) ) & ( np.abs(me) > 3 * np.abs(d0) ), 3 * d0, me, ) return me m[0] = endpoint_slope( h[0], h[1] if K > 2 else h[0], d[0], d[1] if K > 2 else d[0], ) m[-1] = endpoint_slope( h[-1], h[-2] if K > 2 else h[-1], d[-1], d[-2] if K > 2 else d[-1], ) B = Y.shape[1] idx = np.clip( np.searchsorted( x, t, side="right", ) - 1, 0, K - 2, ) xi = x[idx] xi1 = x[idx + 1] hi = xi1 - xi s = ( ( t - xi ) / hi )[:, None] yi = Y[ idx, np.arange(B), ] yi1 = Y[ idx + 1, np.arange(B), ] mi = m[ idx, np.arange(B), ] mi1 = m[ idx + 1, np.arange(B), ] hb = hi[:, None] mu = ( ( 2 * s**3 - 3 * s**2 + 1 ) * yi + ( s**3 - 2 * s**2 + s ) * hb * mi + ( -2 * s**3 + 3 * s**2 ) * yi1 + ( s**3 - s**2 ) * hb * mi1 ) mup = ( ( 6 * s**2 - 6 * s ) * yi + ( 3 * s**2 - 4 * s + 1 ) * hb * mi + ( -6 * s**2 + 6 * s ) * yi1 + ( 3 * s**2 - 2 * s ) * hb * mi1 ) / hb return mu, mup def natural(x, Y, t): """ Natural cubic spline interpolation. """ K = len(x) h = np.diff(x) n = K - 2 if n <= 0: M = np.zeros_like(Y) else: A = np.zeros( (n, n), dtype=np.float64, ) rhs = np.zeros( (n,) + Y.shape[1:], dtype=np.float64, ) for i in range( 1, K - 1, ): r = i - 1 A[r, r] = 2 * ( h[i - 1] + h[i] ) if r > 0: A[ r, r - 1, ] = h[i - 1] if r < n - 1: A[ r, r + 1, ] = h[i] rhs[r] = 6 * ( ( Y[i + 1] - Y[i] ) / h[i] - ( Y[i] - Y[i - 1] ) / h[i - 1] ) M = np.zeros_like(Y) # Solve rather than explicitly invert A. flat_rhs = rhs.reshape( n, -1, ) flat_solution = np.linalg.solve( A, flat_rhs, ) M[ 1:K - 1 ] = flat_solution.reshape( rhs.shape ) B = Y.shape[1] idx = np.clip( np.searchsorted( x, t, side="right", ) - 1, 0, K - 2, ) xi = x[idx] xi1 = x[idx + 1] hh = xi1 - xi a = ( ( xi1 - t ) / hh )[:, None] b = ( ( t - xi ) / hh )[:, None] hb = hh[:, None] yi = Y[ idx, np.arange(B), ] yi1 = Y[ idx + 1, np.arange(B), ] mi = M[ idx, np.arange(B), ] mi1 = M[ idx + 1, np.arange(B), ] mu = ( a * yi + b * yi1 + ( ( a**3 - a ) * mi + ( b**3 - b ) * mi1 ) * ( hb**2 ) / 6 ) mup = ( ( yi1 - yi ) / hb + hb / 6 * ( ( 1 - 3 * a**2 ) * mi + ( 3 * b**2 - 1 ) * mi1 ) ) return mu, mup def linear(x, Y, t): """ Piecewise-linear interpolation. """ B = Y.shape[1] idx = np.clip( np.searchsorted( x, t, side="right", ) - 1, 0, len(x) - 2, ) xi = x[idx] xi1 = x[idx + 1] hb = ( xi1 - xi )[:, None] s = ( ( t - xi ) / ( xi1 - xi ) )[:, None] yi = Y[ idx, np.arange(B), ] yi1 = Y[ idx + 1, np.arange(B), ] mu = ( 1 - s ) * yi + s * yi1 mup = ( yi1 - yi ) / hb return mu, mup SPL = { "linear": linear, "natural": natural, "pchip": pchip, } # ============================================================ # Conditional flow matching # ============================================================ def fit_flow( clouds_by_day, dn, spline, dev, steps, seed, batch=128, lr=1e-3, ): """ Fit a conditional velocity field to one of the reference interpolants. """ torch.manual_seed(seed) rng = np.random.default_rng( seed ) d = next( iter( clouds_by_day.values() ) ).shape[1] field = Field( d ).to(dev) opt = torch.optim.Adam( field.parameters(), lr=lr, ) x = np.array( sorted(dn), dtype=float, ) clouds = [ clouds_by_day[t] for t in x ] ni = len(x) - 1 for step in range(steps): Z = align_all( clouds, rng, batch, ) B = Z.shape[1] per = B // ni ts = [] for i in range(ni): m = per + ( 1 if i < B - per * ni else 0 ) ts.append( rng.uniform( x[i], x[i + 1], size=m, ) ) t = np.concatenate( ts )[:B] rng.shuffle( t ) mu, mup = SPL[ spline ]( x, Z, t, ) tt = torch.tensor( t[:, None], dtype=torch.float32, device=dev, ) mu_t = torch.tensor( mu, dtype=torch.float32, device=dev, ) mup_t = torch.tensor( mup, dtype=torch.float32, device=dev, ) pred = field( mu_t, tt, ) loss = ( ( pred - mup_t ) ** 2 ).mean() opt.zero_grad() loss.backward() nn.utils.clip_grad_norm_( field.parameters(), 1.0, ) opt.step() field.eval() return field def ode( field, z0, ts, te, dev, N=50, ): """ Forward Euler integration of the learned field. """ z = z0.clone() dt = ( te - ts ) / N with torch.no_grad(): for k in range(N): tcur = ( ts + ( te - ts ) * k / N ) tau = torch.full( ( len(z), 1, ), float(tcur), device=dev, ) z = ( z + dt * field( z, tau, ) ) return z.cpu().numpy() # ============================================================ # Simple held-out baselines # ============================================================ def bracketing_days( retained_days, heldout_day, ): """ Return nearest observed retained day before and after heldout_day. Example ------- Gupta d7: retained = [2.5, 14, 21] -> left=2.5, right=14 Gupta d14: retained = [2.5, 7, 21] -> left=7, right=21 """ retained = np.array( sorted( float(x) for x in retained_days ), dtype=float, ) left = retained[ retained < heldout_day ] right = retained[ retained > heldout_day ] if len(left) == 0: raise ValueError( f"No retained day before heldout={heldout_day}" ) if len(right) == 0: raise ValueError( f"No retained day after heldout={heldout_day}" ) return ( float(left.max()), float(right.min()), ) def no_move( X_source, rng=None, n_out=None, ): """ No-move baseline. Predict that the held-out population is simply the previous observed source population. If n_out is supplied, resample source cells to n_out. """ X_source = np.asarray( X_source ) if n_out is None: return X_source.copy() if rng is None: rng = np.random.default_rng( 0 ) idx = rng.choice( len(X_source), size=n_out, replace=( n_out > len(X_source) ), ) return X_source[ idx ].copy() def mccann( Xa, Xb, alpha, rng, n_out=None, ): """ McCann / displacement interpolation under an OT coupling. Parameters ---------- Xa, Xb Bracketing populations. alpha Fractional position in time: alpha = (t_h - t_left) / (t_right - t_left) n_out Number of samples drawn from the interpolated distribution. Defaults to len(Xa). """ Xa = np.asarray( Xa ) Xb = np.asarray( Xb ) P = ot_plan( Xa, Xb, ) total = P.sum() if total <= 0: raise RuntimeError( "OT plan has zero total mass." ) if n_out is None: n_out = len(Xa) idx = rng.choice( P.size, size=n_out, replace=True, p=( P / total ).ravel(), ) pi, pj = np.unravel_index( idx, P.shape, ) return ( ( 1.0 - alpha ) * Xa[pi] + alpha * Xb[pj] ) def plain_ot_impute( clouds_by_day, retained_days, heldout_day, seed=0, n_fit=1500, n_out=1500, ): """ Plain OT / McCann interpolation baseline. Uses only the two retained observed populations immediately bracketing the held-out time. To keep exact EMD computationally manageable, each bracketing cloud is subsampled to at most n_fit cells before constructing the OT plan. The same n_fit should be used for every method/dataset in the benchmark. """ rng = np.random.default_rng( seed ) t0, t1 = bracketing_days( retained_days, heldout_day, ) Xa = np.asarray( clouds_by_day[t0] ) Xb = np.asarray( clouds_by_day[t1] ) if len(Xa) > n_fit: ia = rng.choice( len(Xa), n_fit, replace=False, ) Xa_fit = Xa[ ia ] else: Xa_fit = Xa if len(Xb) > n_fit: ib = rng.choice( len(Xb), n_fit, replace=False, ) Xb_fit = Xb[ ib ] else: Xb_fit = Xb alpha = ( heldout_day - t0 ) / ( t1 - t0 ) pred = mccann( Xa_fit, Xb_fit, alpha, rng, n_out=n_out, ) return pred def no_move_impute( clouds_by_day, retained_days, heldout_day, seed=0, n_out=1500, ): """ Held-out no-move baseline using the nearest earlier retained population. """ rng = np.random.default_rng( seed ) t0, _ = bracketing_days( retained_days, heldout_day, ) return no_move( clouds_by_day[t0], rng=rng, n_out=n_out, ) # ============================================================ # FSI / spline image helper # ============================================================ def fsi( clouds_by_day, dn, t_h, rng, ): """ Push the first observed population through a sequence of barycentric OT maps, then fit a natural spline through the resulting trajectories. Retained here for compatibility with the existing benchmark code. """ x = np.array( sorted(dn), dtype=float, ) clouds = [ clouds_by_day[t] for t in x ] cur = clouds[0] imgs = [ cur ] for k in range( 1, len(clouds), ): P = ot_plan( cur, clouds[k], ) Pn = P / ( P.sum( axis=1, keepdims=True, ) + 1e-12 ) cur = ( Pn @ clouds[k] ) imgs.append( cur ) Y = np.stack( imgs, axis=0, ) mu, _ = natural( x, Y, np.full( Y.shape[1], t_h, ), ) return mu # ============================================================ # Evaluation # ============================================================ def w2( A, B, mx=1500, seed=0, ): """ Empirical W2 using exact balanced OT after paired-size subsampling. """ r = np.random.default_rng( seed ) A = np.asarray( A ) B = np.asarray( B ) if len(A) > mx: A = A[ r.choice( len(A), mx, replace=False, ) ] if len(B) > mx: B = B[ r.choice( len(B), mx, replace=False, ) ] C = ( ( A[:, None, :] - B[None, :, :] ) ** 2 ).sum(-1) cost = ot.emd2( ot.unif( len(A) ), ot.unif( len(B) ), C, numItermax=2_000_000, ) return float( np.sqrt( max( cost, 0.0, ) ) ) def mse_means( A, B, ): """ Squared error between population means. """ A = np.asarray( A ) B = np.asarray( B ) return float( ( ( A.mean(0) - B.mean(0) ) ** 2 ).mean() ) # ============================================================ # Convenience function for held-out baselines # ============================================================ def basic_baseline_predictions( clouds_by_day, retained_days, heldout_day, seed=0, n_fit=1500, n_out=1500, ): """ Return the two simple reference baselines needed for the first TRACE comparison. Returns ------- dict { "No_move": ..., "OT_McCann": ... } """ pred_no_move = no_move_impute( clouds_by_day, retained_days, heldout_day, seed=seed, n_out=n_out, ) pred_ot = plain_ot_impute( clouds_by_day, retained_days, heldout_day, seed=seed, n_fit=n_fit, n_out=n_out, ) return { "No_move": pred_no_move, "OT_McCann": pred_ot, } # ============================================================ # Import check # ============================================================ if __name__ == "__main__": print( "benchmark_table.py loaded successfully." ) print( "This file is intended to be imported by " "02_bracket_heldout.py." ) print( "Available simple baselines: " "No_move, OT_McCann." ) print( "Available CFM splines: " "linear, natural, pchip." )