model/ContinuousFlowOperator.py

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional, Tuple

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F


def sinusoidal_time_embedding(t: torch.Tensor, embed_dim: int) -> torch.Tensor:
    half_dim = embed_dim // 2
    freqs = torch.exp(
        torch.linspace(0.0, math.log(10000.0), half_dim, device=t.device, dtype=t.dtype)
    )
    angles = t[:, None] * freqs[None, :]
    emb = torch.cat([torch.sin(angles), torch.cos(angles)], dim=-1)
    if emb.shape[-1] < embed_dim:
        emb = F.pad(emb, (0, embed_dim - emb.shape[-1]))
    return emb


class TimeEmbeddingMLP(nn.Module):
    def __init__(self, embed_dim: int = 64, hidden_dim: int = 128):
        super().__init__()
        self.embed_dim = embed_dim
        self.net = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, embed_dim),
        )

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        return self.net(sinusoidal_time_embedding(t, self.embed_dim))


class DownsampleBlock1D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, time_emb_dim: int):
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, padding=1)
        self.norm1 = nn.GroupNorm(32, out_channels)
        self.time_proj = nn.Linear(time_emb_dim, out_channels)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1)
        self.norm2 = nn.GroupNorm(32, out_channels)
        self.down = nn.Conv1d(out_channels, out_channels, kernel_size=4, stride=2, padding=1)

    def forward(self, x: torch.Tensor, t_emb: Optional[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
        x = F.silu(self.norm1(self.conv1(x)))
        if t_emb is not None:
            x = x + self.time_proj(t_emb)[:, :, None]
        x = F.silu(self.norm2(self.conv2(x)))
        skip = x
        return self.down(x), skip


class UpsampleBlock1D(nn.Module):
    def __init__(self, in_channels: int, skip_channels: int, out_channels: int, time_emb_dim: int):
        super().__init__()
        self.up = nn.ConvTranspose1d(in_channels, out_channels, kernel_size=4, stride=2, padding=1)
        self.conv1 = nn.Conv1d(out_channels + skip_channels, out_channels, kernel_size=3, padding=1)
        self.norm1 = nn.GroupNorm(32, out_channels)
        self.time_proj = nn.Linear(time_emb_dim, out_channels)
        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, padding=1)
        self.norm2 = nn.GroupNorm(32, out_channels)

    def forward(self, x: torch.Tensor, skip: torch.Tensor, t_emb: Optional[torch.Tensor]) -> torch.Tensor:
        x = self.up(x)
        if x.shape[-1] != skip.shape[-1]:
            x = F.interpolate(x, size=skip.shape[-1], mode="linear", align_corners=False)
        x = torch.cat([x, skip], dim=1)
        x = F.silu(self.norm1(self.conv1(x)))
        if t_emb is not None:
            x = x + self.time_proj(t_emb)[:, :, None]
        x = F.silu(self.norm2(self.conv2(x)))
        return x


class UNet1D(nn.Module):
    """1D U-Net backbone with additive sinusoidal time embeddings."""

    def __init__(
        self,
        in_channels: int = 1,
        channel_mult: Tuple[int, ...] = (64, 128),
        time_emb_dim: int = 64,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.time_mlp = TimeEmbeddingMLP(embed_dim=time_emb_dim)
        downs = []
        prev = in_channels
        for ch in channel_mult:
            downs.append(DownsampleBlock1D(prev, ch, time_emb_dim))
            prev = ch
        self.downs = nn.ModuleList(downs)
        self.mid1 = nn.Conv1d(prev, prev, kernel_size=3, padding=1)
        self.mid_norm1 = nn.GroupNorm(32, prev)
        self.mid_time = nn.Linear(time_emb_dim, prev)
        self.mid2 = nn.Conv1d(prev, prev, kernel_size=3, padding=1)
        self.mid_norm2 = nn.GroupNorm(32, prev)
        ups = []
        in_ch = prev
        for ch in reversed(channel_mult):
            ups.append(UpsampleBlock1D(in_ch, ch, ch, time_emb_dim))
            in_ch = ch
        self.ups = nn.ModuleList(ups)
        self.out = nn.Conv1d(in_ch, in_channels, kernel_size=1)

    def forward(self, x: torch.Tensor, t: Optional[torch.Tensor] = None) -> torch.Tensor:
        squeeze = False
        if x.ndim == 2:
            x = x[:, None, :]
            squeeze = True
        elif x.ndim == 3 and x.shape[-1] == self.in_channels:
            x = x.transpose(1, 2)
        if t is None:
            t = torch.zeros(x.shape[0], device=x.device, dtype=x.dtype)
        t_emb = self.time_mlp(t)
        skips = []
        h = x
        for down in self.downs:
            h, skip = down(h, t_emb)
            skips.append(skip)
        h = F.silu(self.mid_norm1(self.mid1(h)))
        h = h + self.mid_time(t_emb)[:, :, None]
        h = F.silu(self.mid_norm2(self.mid2(h)))
        for up in self.ups:
            h = up(h, skips.pop(), t_emb)
        h = self.out(h)
        return h[:, 0, :] if squeeze else h.transpose(1, 2)


def estimate_derivatives_np(x: np.ndarray, time: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray]:
    x = np.asarray(x, dtype=np.float32)
    t_steps = x.shape[0]
    if t_steps < 3:
        dx = np.gradient(x, axis=0).astype(np.float32)
        ddx = np.gradient(dx, axis=0).astype(np.float32)
        return dx, ddx
    if time is None:
        dt = np.float32(1.0 / (t_steps - 1))
        dx_left = (-3.0 * x[0] + 4.0 * x[1] - x[2]) / (2.0 * dt)
        dx_right = (3.0 * x[-1] - 4.0 * x[-2] + x[-3]) / (2.0 * dt)
        dx_c = (x[2:] - x[:-2]) / (2.0 * dt)
        dx = np.concatenate([dx_left[None], dx_c, dx_right[None]], axis=0)
        c2 = (x[2:] - 2.0 * x[1:-1] + x[:-2]) / (dt * dt)
        ddx = np.concatenate([c2[0:1], c2, c2[-1:]], axis=0)
        return dx.astype(np.float32), ddx.astype(np.float32)
    time = np.asarray(time, dtype=np.float32)
    h = np.diff(time).astype(np.float32)
    h_b = h.reshape(h.shape + (1,) * (x.ndim - 1))
    h0, h1 = h_b[:-1], h_b[1:]
    dx_c = (
        (-h1 / (h0 * (h0 + h1))) * x[:-2]
        + ((h1 - h0) / (h0 * h1)) * x[1:-1]
        + (h0 / (h1 * (h0 + h1))) * x[2:]
    )
    h0l, h1l = h_b[0], h_b[1]
    dx_left = (
        (-(2.0 * h0l + h1l) / (h0l * (h0l + h1l))) * x[0]
        + ((h0l + h1l) / (h0l * h1l)) * x[1]
        - (h0l / (h1l * (h0l + h1l))) * x[2]
    )
    g0, g1 = h_b[-2], h_b[-1]
    dx_right = (
        (g1 / (g0 * (g0 + g1))) * x[-3]
        - ((g0 + g1) / (g0 * g1)) * x[-2]
        + ((2.0 * g1 + g0) / (g1 * (g0 + g1))) * x[-1]
    )
    dx = np.concatenate([dx_left[None], dx_c, dx_right[None]], axis=0)
    ddx_c = 2.0 * (((x[2:] - x[1:-1]) / h1) - ((x[1:-1] - x[:-2]) / h0)) / (h0 + h1)
    ddx_left = 2.0 * (x[0] / (h0l * (h0l + h1l)) - x[1] / (h0l * h1l) + x[2] / (h1l * (h0l + h1l)))
    ddx_right = 2.0 * (x[-3] / (g0 * (g0 + g1)) - x[-2] / (g0 * g1) + x[-1] / (g1 * (g0 + g1)))
    ddx = np.concatenate([ddx_left[None], ddx_c, ddx_right[None]], axis=0)
    return dx.astype(np.float32), ddx.astype(np.float32)


def quintic_spline_coefficients_np(data: np.ndarray, time: Optional[np.ndarray] = None) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    data = np.asarray(data, dtype=np.float32)
    batch, steps = data.shape[:2]
    if time is None:
        time = np.linspace(0.0, 1.0, steps, dtype=np.float32)
        time = np.broadcast_to(time[None, :], (batch, steps))
    elif time.ndim == 1:
        time = np.broadcast_to(time[None, :].astype(np.float32), (batch, steps))
    else:
        time = np.asarray(time, dtype=np.float32)
    coeffs = []
    starts = []
    ends = []
    m_inv = np.asarray([[10.0, -4.0, 0.5], [-15.0, 7.0, -1.0], [6.0, -3.0, 0.5]], dtype=np.float32)
    for i in range(batch):
        x = data[i]
        ti = time[i]
        dx, ddx = estimate_derivatives_np(x, ti)
        dt = np.diff(ti).reshape((steps - 1,) + (1,) * (x.ndim - 1)).astype(np.float32)
        a0 = x[:-1]
        a1 = dx[:-1] * dt
        a2 = 0.5 * ddx[:-1] * (dt * dt)
        s1 = x[1:] - (a0 + a1 + a2)
        s2 = dx[1:] * dt - (a1 + 2.0 * a2)
        s3 = ddx[1:] * (dt * dt) - (2.0 * a2)
        b_vec = np.stack([s1, s2, s3], axis=-1)
        a345 = np.einsum("...k,lk->...l", b_vec, m_inv)
        ci = np.stack([a0, a1, a2, a345[..., 0], a345[..., 1], a345[..., 2]], axis=1)
        coeffs.append(ci.astype(np.float32))
        starts.append(ti[:-1].astype(np.float32))
        ends.append(ti[1:].astype(np.float32))
    return (
        np.concatenate(coeffs, axis=0),
        np.concatenate(starts, axis=0),
        np.concatenate(ends, axis=0),
    )


@dataclass
class CFOBatch:
    spline_coef: torch.Tensor
    t_start: torch.Tensor
    t_end: torch.Tensor
    delta_t: torch.Tensor
    eps: torch.Tensor


class ContinuousFlowOperator(nn.Module):
    def __init__(self, input_size: int = 100, gamma: float = 1e-5):
        super().__init__()
        self.input_size = int(input_size)
        self.gamma = float(gamma)
        self.net = UNet1D(in_channels=1, channel_mult=(64, 128), time_emb_dim=64)

    @staticmethod
    def _reshape_time_like(x: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
        return x.reshape((x.shape[0],) + (1,) * (ref.ndim - 1))

    def velocity(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        return self.net(x, t)

    def sample_conditional_path(self, tau: torch.Tensor, spline_coef: torch.Tensor, eps: torch.Tensor) -> torch.Tensor:
        x = (
            spline_coef[:, 0]
            + tau * spline_coef[:, 1]
            + tau.pow(2) * spline_coef[:, 2]
            + tau.pow(3) * spline_coef[:, 3]
            + tau.pow(4) * spline_coef[:, 4]
            + tau.pow(5) * spline_coef[:, 5]
        )
        gamma_t = self.gamma * tau.pow(3) * (1.0 - tau).pow(3)
        return x + gamma_t * eps

    def compute_targets(self, spline_coef: torch.Tensor, tau: torch.Tensor, dt: torch.Tensor, eps: torch.Tensor) -> torch.Tensor:
        poly_dtau = (
            spline_coef[:, 1]
            + 2.0 * spline_coef[:, 2] * tau
            + 3.0 * spline_coef[:, 3] * tau.pow(2)
            + 4.0 * spline_coef[:, 4] * tau.pow(3)
            + 5.0 * spline_coef[:, 5] * tau.pow(4)
        )
        gamma_prime = self.gamma * (3.0 / dt) * (tau.pow(2) * (1.0 - tau).pow(2) * (1.0 - 2.0 * tau))
        return (poly_dtau / dt) + gamma_prime * eps

    def loss_from_batch(self, batch: CFOBatch) -> torch.Tensor:
        dt = batch.t_end - batch.t_start
        tau_scalar = batch.delta_t / dt
        tau = self._reshape_time_like(tau_scalar, batch.spline_coef[:, 0])
        dt_view = self._reshape_time_like(dt, batch.spline_coef[:, 0])
        x = self.sample_conditional_path(tau, batch.spline_coef, batch.eps)
        pred = self.velocity(x, batch.t_start + batch.delta_t)
        target = self.compute_targets(batch.spline_coef, tau, dt_view, batch.eps)
        return F.mse_loss(pred, target)

    @torch.no_grad()
    def rk4_step(self, x: torch.Tensor, t: float, dt: float) -> torch.Tensor:
        tb = torch.full((x.shape[0],), float(t), device=x.device, dtype=x.dtype)
        dt_t = torch.as_tensor(dt, device=x.device, dtype=x.dtype)
        k1 = self.velocity(x, tb)
        k2 = self.velocity(x + 0.5 * dt_t * k1, tb + 0.5 * dt_t)
        k3 = self.velocity(x + 0.5 * dt_t * k2, tb + 0.5 * dt_t)
        k4 = self.velocity(x + dt_t * k3, tb + dt_t)
        return x + (dt_t / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)

    @torch.no_grad()
    def rollout(self, x0: torch.Tensor, trajectory_points: int, steps_per_segment: int = 2) -> torch.Tensor:
        self.eval()
        points = int(trajectory_points)
        segment_dt = 1.0 / max(points - 1, 1)
        x = x0
        preds = [x0.detach().cpu()]
        for idx in range(points - 1):
            start = idx * segment_dt
            sub_dt = segment_dt / float(max(steps_per_segment, 1))
            for inner in range(max(steps_per_segment, 1)):
                x = self.rk4_step(x, start + inner * sub_dt, sub_dt)
            preds.append(x.detach().cpu())
        return torch.stack(preds, dim=1)


def trajectory_relative_l2(gt: np.ndarray, pred: np.ndarray, eps: float = 1e-8) -> float:
    gt = np.asarray(gt, dtype=np.float64)
    pred = np.asarray(pred, dtype=np.float64)
    axes = tuple(range(1, gt.ndim))
    err = np.sqrt(np.mean((gt - pred) ** 2, axis=axes))
    den = np.sqrt(np.mean(gt ** 2, axis=axes)) + eps
    return float(np.mean(err / den))


def rmse(gt: np.ndarray, pred: np.ndarray) -> float:
    return float(np.sqrt(np.mean((np.asarray(gt) - np.asarray(pred)) ** 2)))