TRACE / benchmark_table.py

Initial TRACE implementation

4c4d79d verified 2 days ago

23 kB

#!/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."
    )

	#!/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."
	)