📌 JAX로 Prophet 스타일 + Neural Basis Expansion Hybrid 시계열 예측 모델 구현

📌 JAX로 Prophet 스타일 + Neural Basis Expansion Hybrid 시계열 예측 모델 구현

---

🌟 왜 Hybrid 모델인가?

• Prophet은 계절성, 추세, 휴일 효과 등 해석 가능한 예측이 강점
• **Neural Basis Expansion (N-BEATS)**는 딥러닝 기반의 고정된 블록을 쌓아 높은 예측 정확도를 달성
• JAX에서는 정형적 계절성 + 비선형 딥러닝 구조를 동시에 반영한 Hybrid 모델을 유연하게 구현할 수 있음

---

🎯 이 글에서 구현하는 Hybrid 구성

구성요소 역할

Trend Component (Linear/Logistic)	Prophet의 추세 모델링
Seasonal Component (Fourier basis)	주기성 정보 반영
Residual Neural Component (MLP or Transformer)	오차 보정, 비선형 패턴 처리

---

💾 1. 데이터 준비

import pandas as pd
import jax.numpy as jnp
from sklearn.preprocessing import MinMaxScaler

df = pd.read_csv("time_series_data.csv", parse_dates=["ds"])
y = df['y'].values.reshape(-1, 1)

scaler = MinMaxScaler()
y_scaled = scaler.fit_transform(y)
t = jnp.arange(len(y_scaled)).reshape(-1, 1)  # 시점 정보

---

🧠 2. Fourier 기반 계절성 함수 구현

def fourier_features(t, period=365.25, order=3):
    t = t / period
    features = [jnp.sin(2 * jnp.pi * t * i) for i in range(1, order + 1)]
    features += [jnp.cos(2 * jnp.pi * t * i) for i in range(1, order + 1)]
    return jnp.concatenate(features, axis=1)  # shape: (N, order*2)

---

🧱 3. Hybrid 모델 정의 (Flax)

from flax import linen as nn

class ProphetLikeHybridModel(nn.Module):
    fourier_order: int = 3

    @nn.compact
    def __call__(self, t):
        # Trend: 선형 추세
        trend = nn.Dense(1)(t)  # (batch, 1)

        # Seasonality: Fourier
        fourier = fourier_features(t, order=self.fourier_order)
        seasonal = nn.Dense(1)(fourier)  # (batch, 1)

        # Residual (non-linear part)
        x = jnp.concatenate([t, fourier], axis=1)
        x = nn.Dense(32)(x)
        x = nn.relu(x)
        x = nn.Dense(1)(x)  # Residual output

        return trend + seasonal + x  # Hybrid 합산

---

⚙️ 4. 학습 초기화

반응형

import optax
from flax.training import train_state
import jax

model = ProphetLikeHybridModel()
params = model.init(jax.random.PRNGKey(0), t)
tx = optax.adam(1e-3)
state = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)

---

📉 5. 손실 및 학습 Step

@jax.jit
def loss_fn(params, t, y_true):
    preds = model.apply(params, t)
    return jnp.mean((preds - y_true) ** 2)

@jax.jit
def train_step(state, t, y_true):
    loss, grads = jax.value_and_grad(loss_fn)(state.params, t, y_true)
    state = state.apply_gradients(grads=grads)
    return state, loss

---

🏃 6. 학습 루프

epochs = 500
for epoch in range(epochs):
    state, loss = train_step(state, t, y_scaled)
    if epoch % 50 == 0:
        print(f"Epoch {epoch} - Loss: {loss:.6f}")

---

🔮 7. 예측 및 시각화

import matplotlib.pyplot as plt

preds = model.apply(state.params, t)
preds_inv = scaler.inverse_transform(preds)

plt.plot(df['ds'], y, label="True")
plt.plot(df['ds'], preds_inv, label="Hybrid Prediction")
plt.title("Prophet + NN Hybrid Forecast")
plt.legend()
plt.show()

---

✅ 8. 활용 시나리오

분야 설명

에너지 예측	주간/연간 주기 + 이상치
금융	트렌드 + 비정형적 급등락
수요 예측	명절/계절 효과 + 머신러닝 보정

---

📌 다음 글 예고: JAX 기반의 Graph Neural Network(GNN)로 시계열 + 관계 데이터 예측 모델 구현하기

 

JAX, 시계열 예측, Hybrid 모델, Prophet, Fourier Transform, Neural Basis Expansion, Flax, 딥러닝 시계열, 시계열 해석, 시계열 트렌드, 시계열 계절성, AI 예측, 비선형 예측, Multi-component Forecasting, Python ML