Sundial — 1兆データ点で事前学習した最新最大の時系列基盤モデル【ICML 2025】

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

class SimpleFlowNetwork:
    """シンプルなMLPベースのベクトル場ネットワーク。"""

    def __init__(self, input_dim, hidden_dim=128):
        self.W1 = np.random.randn(input_dim + 1, hidden_dim) * 0.1  # +1 for time
        self.b1 = np.zeros(hidden_dim)
        self.W2 = np.random.randn(hidden_dim, hidden_dim) * 0.1
        self.b2 = np.zeros(hidden_dim)
        self.W3 = np.random.randn(hidden_dim, input_dim) * 0.1
        self.b3 = np.zeros(input_dim)

    def forward(self, x, t):
        """ベクトル場を予測する。"""
        # 入力: x と t を連結
        t_expanded = np.full((x.shape[0], 1), t)
        inp = np.concatenate([x, t_expanded], axis=1)

        # 2層MLP with ReLU
        h = inp @ self.W1 + self.b1
        h = np.maximum(h, 0)
        h = h @ self.W2 + self.b2
        h = np.maximum(h, 0)
        out = h @ self.W3 + self.b3
        return out

    def train_step(self, x1_batch, lr=0.001):
        """CFM損失でパラメータを更新する。"""
        batch_size = x1_batch.shape[0]
        dim = x1_batch.shape[1]

        # ランダムな時刻とノイズ
        t = np.random.uniform(0, 1)
        x0 = np.random.randn(batch_size, dim)

        # 線形補間
        xt = (1 - t) * x0 + t * x1_batch

        # ターゲットベクトル場
        target = x1_batch - x0

        # 予測
        pred = self.forward(xt, t)

        # MSE損失
        loss = np.mean((pred - target)**2)

        # 簡易的な勾配更新（数値勾配）
        eps = 1e-4
        for param_name in ['W1', 'b1', 'W2', 'b2', 'W3', 'b3']:
            param = getattr(self, param_name)
            grad = np.zeros_like(param)
            # バッチの平均勾配を近似
            flat = param.flatten()
            # サブサンプリングで高速化
            n_sample = min(50, len(flat))
            indices = np.random.choice(len(flat), n_sample, replace=False)
            for idx in indices:
                old_val = flat[idx]
                flat[idx] = old_val + eps
                pred_plus = self.forward(xt, t)
                loss_plus = np.mean((pred_plus - target)**2)
                flat[idx] = old_val - eps
                pred_minus = self.forward(xt, t)
                loss_minus = np.mean((pred_minus - target)**2)
                flat[idx] = old_val
                grad.flatten()[idx] = (loss_plus - loss_minus) / (2 * eps)
            setattr(self, param_name, param - lr * grad)

        return loss

    def generate(self, n_samples, dim, n_steps=20):
        """ODEソルバーでサンプルを生成する。"""
        x = np.random.randn(n_samples, dim)
        dt = 1.0 / n_steps
        trajectory = [x.copy()]

        for step in range(n_steps):
            t = step * dt
            v = self.forward(x, t)
            x = x + v * dt
            trajectory.append(x.copy())

        return x, trajectory


# 時系列データの生成（周期的パターン）
def generate_periodic_series(n_samples, length=16):
    """周期的な時系列パターンを生成する。"""
    t = np.linspace(0, 2 * np.pi, length)
    series = []
    for _ in range(n_samples):
        freq = np.random.uniform(0.5, 2.0)
        amp = np.random.uniform(0.5, 2.0)
        phase = np.random.uniform(0, 2 * np.pi)
        s = amp * np.sin(freq * t + phase) + np.random.randn(length) * 0.1
        series.append(s)
    return np.array(series)


# Flow Matchingによる時系列生成のデモ
pred_len = 16
n_train = 500
train_data = generate_periodic_series(n_train, pred_len)

# 学習
model = SimpleFlowNetwork(input_dim=pred_len, hidden_dim=64)
losses = []
for epoch in range(100):
    idx = np.random.choice(n_train, 32)
    batch = train_data[idx]
    loss = model.train_step(batch, lr=0.001)
    losses.append(loss)

# 複数サンプルの生成（確率的予測のシミュレーション）
n_gen = 50
generated, trajectories = model.generate(n_gen, pred_len, n_steps=20)

# 可視化
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# 左: 学習損失
axes[0].plot(losses, color='#00d4ff')
axes[0].set_xlabel('Training Step')
axes[0].set_ylabel('CFM Loss (MSE)')
axes[0].set_title('Flow Matching Training Loss')
axes[0].grid(True, alpha=0.3)

# 中央: 生成された時系列サンプル vs 真のデータ
t_axis = np.arange(pred_len)
for i in range(min(10, n_gen)):
    axes[1].plot(t_axis, generated[i], alpha=0.3, color='#ffa726')
for i in range(5):
    axes[1].plot(t_axis, train_data[i], alpha=0.7, color='#00d4ff', linewidth=2)
axes[1].set_xlabel('Time Step')
axes[1].set_ylabel('Value')
axes[1].set_title('Generated (orange) vs Training Data (blue)')
axes[1].grid(True, alpha=0.3)

# 右: ODE軌道の可視化（最初の2次元を射影）
traj = np.array(trajectories)  # (n_steps+1, n_gen, dim)
for i in range(min(5, n_gen)):
    axes[2].plot(traj[:, i, 0], traj[:, i, 1], 'o-', markersize=2, alpha=0.5)
axes[2].scatter(traj[0, :5, 0], traj[0, :5, 1], c='red', s=50, label='Noise (t=0)', zorder=5)
axes[2].scatter(traj[-1, :5, 0], traj[-1, :5, 1], c='green', s=50, label='Data (t=1)', zorder=5)
axes[2].set_xlabel('Dimension 0')
axes[2].set_ylabel('Dimension 1')
axes[2].set_title('ODE Trajectory (first 2 dims)')
axes[2].legend()
axes[2].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('sundial_flow_matching.png', dpi=150, bbox_inches='tight')
plt.show()

# 条件付き予測のデモ
# コンテキスト（過去）を与えて未来を確率的に予測する

ctx_len = 32
pred_len_demo = 16
n_predictions = 30

# テスト用の時系列
t_full = np.linspace(0, 4 * np.pi, ctx_len + pred_len_demo)
true_series = 1.5 * np.sin(t_full) + 0.5 * np.sin(3 * t_full) + 0.1 * np.random.randn(len(t_full))
context = true_series[:ctx_len]
future_true = true_series[ctx_len:]

# 確率的予測のシミュレーション
# （簡略化: コンテキストの統計量に基づいてノイズをバイアス）
predictions = []
for _ in range(n_predictions):
    noise = np.random.randn(pred_len_demo) * 0.3
    # コンテキストの末尾のトレンドを考慮した予測
    trend = np.polyfit(np.arange(10), context[-10:], 1)
    pred = np.polyval(trend, np.arange(ctx_len - 1, ctx_len + pred_len_demo - 1))
    pred += np.sin(np.linspace(context[-1], context[-1] + np.pi, pred_len_demo)) * 1.2
    pred += noise
    predictions.append(pred)

predictions = np.array(predictions)

plt.figure(figsize=(12, 5))
t_ctx = np.arange(ctx_len)
t_pred = np.arange(ctx_len, ctx_len + pred_len_demo)

# コンテキスト
plt.plot(t_ctx, context, color='#00d4ff', linewidth=2, label='Context')

# 予測サンプル
for i in range(n_predictions):
    plt.plot(t_pred, predictions[i], color='#ffa726', alpha=0.15)

# 予測の統計量
pred_mean = np.mean(predictions, axis=0)
pred_std = np.std(predictions, axis=0)
plt.plot(t_pred, pred_mean, color='#ffa726', linewidth=2, label='Prediction Mean')
plt.fill_between(t_pred, pred_mean - 2*pred_std, pred_mean + 2*pred_std,
                 alpha=0.2, color='#ffa726', label='95% CI')

# 真の未来
plt.plot(t_pred, future_true, color='#ef5350', linewidth=2, linestyle='--', label='True Future')

plt.axvline(x=ctx_len, color='gray', linestyle=':', alpha=0.5)
plt.xlabel('Time Step')
plt.ylabel('Value')
plt.title('Probabilistic Time Series Forecasting (Flow Matching Style)')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('sundial_probabilistic.png', dpi=150, bbox_inches='tight')
plt.show()