Merge pull request #293 from wgawmy/beta

Initial Attempt at Koopa Algorithm

Author: yisongfu

src/basicts/models/Koopa/__init__.py

@@ -0,0 +1,2 @@
+from .arch import Koopa
+from .config.koopa_config import KoopaConfig

src/basicts/models/Koopa/arch/__init__.py

@@ -0,0 +1 @@
+from .koopa_arch import Koopa

src/basicts/models/Koopa/arch/koopa_arch.py

@@ -0,0 +1,100 @@
+import torch
+from torch import nn
+from .layers import FourierFilter, MLP, TimeInvKP, TimeVarKP
+from ..config.koopa_config import KoopaConfig
+
+class Koopa(nn.Module):
+    """
+    Paper: Koopa: Learning Non-stationary Time Series Dynamics with Koopman Predictors
+    Official Code: https://github.com/thuml/Koopa
+    Link: https://arxiv.org/abs/2305.18803
+    Venue: NeurIPS 2024
+    Task: Long-term Time Series Forecasting
+    """
+    def __init__(self, config: KoopaConfig):
+        super().__init__()
+        self.mask_spectrum = None
+        self.amps = None
+        self.alpha = config.alpha
+        self.enc_in = config.enc_in
+        self.input_len = config.input_len
+        self.output_len = config.output_len
+        self.seg_len = config.seg_len
+        self.num_blocks = config.num_blocks
+        self.dynamic_dim = config.dynamic_dim
+        self.hidden_dim = config.hidden_dim
+        self.hidden_layers = config.hidden_layers
+        self.multistep = config.multistep
+        self.disentanglement = FourierFilter(self.mask_spectrum)
+        # shared encoder/decoder to make koopman embedding consistent
+        self.time_inv_encoder = MLP(f_in=self.input_len, f_out=self.dynamic_dim, activation='relu',
+                                    hidden_dim=self.hidden_dim, hidden_layers=self.hidden_layers)
+        # fix: use self.output_len instead of non-existent attribute
+        self.time_inv_decoder = MLP(f_in=self.dynamic_dim, f_out=self.output_len, activation='relu',
+                                    hidden_dim=self.hidden_dim, hidden_layers=self.hidden_layers)
+        # separate module lists for time-invariant and time-variant KPs
+        self.time_inv_kps = nn.ModuleList([
+            TimeInvKP(input_len=self.input_len,
+                      pred_len=self.output_len,
+                      dynamic_dim=self.dynamic_dim,
+                      encoder=self.time_inv_encoder,
+                      decoder=self.time_inv_decoder)
+            for _ in range(self.num_blocks)])
+
+        # shared encoder/decoder to make koopman embedding consistent
+        self.time_var_encoder = MLP(f_in=self.seg_len * self.enc_in, f_out=self.dynamic_dim, activation='tanh',
+                                   hidden_dim=self.hidden_dim, hidden_layers=self.hidden_layers)
+        self.time_var_decoder = MLP(f_in=self.dynamic_dim, f_out=self.seg_len * self.enc_in, activation='tanh',
+                                   hidden_dim=self.hidden_dim, hidden_layers=self.hidden_layers)
+        self.time_var_kps = nn.ModuleList([
+            TimeVarKP(enc_in=self.enc_in,
+                      input_len=self.input_len,
+                      pred_len=self.output_len,
+                      seg_len=self.seg_len,
+                      dynamic_dim=self.dynamic_dim,
+                      encoder=self.time_var_encoder,
+                      decoder=self.time_var_decoder,
+                      multistep=self.multistep)
+            for _ in range(self.num_blocks)])
+    def forward(self, inputs: torch.Tensor = None) -> torch.Tensor:
+        """
+        Single-`inputs` forward to match runner API.
+
+        Args:
+            inputs (torch.Tensor): history input with shape [B, L, C] or [B, L, C, 1]
+
+        Returns:
+            torch.Tensor: prediction tensor with shape [B, output_len, num_features] (may include trailing feature dim)
+        """
+        history_data = inputs
+        if history_data is None:
+            raise AssertionError('Model forward requires inputs(history data) as first argument.')
+
+        if history_data.dim() == 4:
+            x_enc = history_data[..., 0]
+        elif history_data.dim() == 3:
+            x_enc = history_data
+        else:
+            raise ValueError(f'Unsupported inputs shape: {tuple(history_data.shape)}')
+
+        mean_enc = x_enc.mean(1, keepdim=True).detach()
+        x_enc = x_enc - mean_enc
+        std_enc = torch.sqrt(torch.var(x_enc, dim=1, keepdim=True, unbiased=False) + 1e-5).detach()
+        x_enc = x_enc / std_enc
+        if self.disentanglement is None:
+            raise ValueError('Koopa mask_spectrum is not initialized.')
+
+        residual, forecast = x_enc, None
+        for i in range(self.num_blocks):
+            time_var_input, time_inv_input = self.disentanglement(residual)
+            time_inv_output = self.time_inv_kps[i](time_inv_input)
+            time_var_backcast, time_var_output = self.time_var_kps[i](time_var_input)
+            residual = residual - time_var_backcast
+            if forecast is None:
+                forecast = time_inv_output + time_var_output
+            else:
+                forecast += (time_inv_output + time_var_output)
+        res = forecast * std_enc + mean_enc
+        if history_data is not None and history_data.dim() == 4 and res.dim() == 3:
+            res = res.unsqueeze(-1)
+        return res

src/basicts/models/Koopa/arch/layers.py

@@ -0,0 +1,234 @@
+import math
+import torch
+from torch import nn
+class FourierFilter(nn.Module):
+    """
+    Fourier Filter: to time-variant and time-invariant term
+    """
+    def __init__(self, mask_spectrum):
+        super().__init__()
+        self.mask_spectrum = mask_spectrum
+
+    def forward(self, x):
+        xf = torch.fft.rfft(x, dim=1)
+        mask = torch.ones_like(xf)
+        mask[:, self.mask_spectrum, :] = 0
+        x_var = torch.fft.irfft(xf * mask, dim=1)
+        x_inv = x - x_var
+
+        return x_var, x_inv
+
+
+class MLP(nn.Module):
+    '''
+    Multilayer perceptron to encode/decode high dimension representation of sequential data
+    '''
+
+    def __init__(self,
+                 f_in,
+                 f_out,
+                 hidden_dim=128,
+                 hidden_layers=2,
+                 dropout=0.05,
+                 activation='tanh'):
+        super().__init__()
+        self.f_in = f_in
+        self.f_out = f_out
+        self.hidden_dim = hidden_dim
+        self.hidden_layers = hidden_layers
+        self.dropout = dropout
+        if activation == 'relu':
+            self.activation = nn.ReLU()
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        else:
+            raise NotImplementedError
+
+        layers = [nn.Linear(self.f_in, self.hidden_dim),
+                  self.activation, nn.Dropout(self.dropout)]
+        for _ in range(self.hidden_layers - 2):
+            layers += [nn.Linear(self.hidden_dim, self.hidden_dim),
+                       self.activation, nn.Dropout(dropout)]
+
+        layers += [nn.Linear(hidden_dim, f_out)]
+        self.layers = nn.Sequential(*layers)
+
+    def forward(self, x):
+        # x:     B x S x f_in
+        # y:     B x S x f_out
+        y = self.layers(x)
+        return y
+
+
+class KPLayer(nn.Module):
+    """
+    A demonstration of finding one step transition of linear system by DMD iteratively
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self.K = None  # B E E
+
+    def one_step_forward(self, z, return_rec=False):
+        B, input_len, _ = z.shape
+        assert input_len > 1, 'snapshots number should be larger than 1'
+        x, y = z[:, :-1], z[:, 1:]
+
+        # solve linear system
+        self.K = torch.linalg.lstsq(x, y).solution  # B E E
+        if torch.isnan(self.K).any():
+            print('Encounter K with nan, replace K by identity matrix')
+            self.K = torch.eye(self.K.shape[1]).to(self.K.device).unsqueeze(0).repeat(B, 1, 1)
+
+        z_pred = torch.bmm(z[:, -1:], self.K)
+        if return_rec:
+            z_rec = torch.cat((z[:, :1], torch.bmm(x, self.K)), dim=1)
+            return z_rec, z_pred
+
+        return z_pred
+
+    def forward(self, z, pred_len=1):
+        assert pred_len >= 1, 'prediction length should not be less than 1'
+        z_rec, z_pred = self.one_step_forward(z, return_rec=True)
+        z_preds = [z_pred]
+        for _ in range(1, pred_len):
+            z_pred = torch.bmm(z_pred, self.K)
+            z_preds.append(z_pred)
+        z_preds = torch.cat(z_preds, dim=1)
+        return z_rec, z_preds
+
+
+class KPLayerApprox(nn.Module):
+    """
+    Find koopman transition of linear system by DMD with multistep K approximation
+    """
+
+    def __init__(self):
+        super().__init__()
+
+        self.K = None  # B E E
+        self.K_step = None  # B E E
+
+    def forward(self, z, pred_len=1):
+        # z:       B L E, koopman invariance space representation
+        # z_rec:   B L E, reconstructed representation
+        # z_pred:  B S E, forecasting representation
+        B, input_len, _ = z.shape
+        assert input_len > 1, 'snapshots number should be larger than 1'
+        x, y = z[:, :-1], z[:, 1:]
+
+        # solve linear system
+        self.K = torch.linalg.lstsq(x, y).solution  # B E E
+
+        if torch.isnan(self.K).any():
+            print('Encounter K with nan, replace K by identity matrix')
+            self.K = torch.eye(self.K.shape[1]).to(self.K.device).unsqueeze(0).repeat(B, 1, 1)
+
+        z_rec = torch.cat((z[:, :1], torch.bmm(x, self.K)), dim=1)  # B L E
+
+        if pred_len <= input_len:
+            self.K_step = torch.linalg.matrix_power(self.K, pred_len)
+            if torch.isnan(self.K_step).any():
+                print('Encounter multistep K with nan, replace it by identity matrix')
+                self.K_step = torch.eye(self.K_step.shape[1]).to(self.K_step.device).unsqueeze(0).repeat(B, 1, 1)
+            z_pred = torch.bmm(z[:, -pred_len:, :], self.K_step)
+        else:
+            self.K_step = torch.linalg.matrix_power(self.K, input_len)
+            if torch.isnan(self.K_step).any():
+                print('Encounter multistep K with nan, replace it by identity matrix')
+                self.K_step = torch.eye(self.K_step.shape[1]).to(self.K_step.device).unsqueeze(0).repeat(B, 1, 1)
+            temp_z_pred, all_pred = z, []
+            for _ in range(math.ceil(pred_len / input_len)):
+                temp_z_pred = torch.bmm(temp_z_pred, self.K_step)
+                all_pred.append(temp_z_pred)
+            z_pred = torch.cat(all_pred, dim=1)[:, :pred_len, :]
+
+        return z_rec, z_pred
+
+
+class TimeVarKP(nn.Module):
+    """
+    Koopman Predictor with DMD (analysitical solution of Koopman operator)
+    Utilize local variations within individual sliding window to predict the future of time-variant term
+    """
+
+    def __init__(self,
+                 enc_in=8,
+                 input_len=96,
+                 pred_len=96,
+                 seg_len=24,
+                 dynamic_dim=128,
+                 encoder=None,
+                 decoder=None,
+                 multistep=False,
+                 ):
+        super().__init__()
+        self.input_len = input_len
+        self.pred_len = pred_len
+        self.enc_in = enc_in
+        self.seg_len = seg_len
+        self.dynamic_dim = dynamic_dim
+        self.multistep = multistep
+        self.encoder, self.decoder = encoder, decoder
+        self.freq = math.ceil(self.input_len / self.seg_len)  # segment number of input
+        self.step = math.ceil(self.pred_len / self.seg_len)  # segment number of output
+        self.padding_len = self.seg_len * self.freq - self.input_len
+        # Approximate mulitstep K by KPLayerApprox when pred_len is large
+        self.dynamics = KPLayerApprox() if self.multistep else KPLayer()
+
+    def forward(self, x):
+        B, L, _ = x.shape
+
+        res = torch.cat((x[:, L - self.padding_len:, :], x), dim=1)
+
+        res = res.chunk(self.freq, dim=1)  # F x B P C, P means seg_len
+        res = torch.stack(res, dim=1).reshape(B, self.freq, -1)  # B F PC
+
+        res = self.encoder(res)  # B F H
+        x_rec, x_pred = self.dynamics(res, self.step)  # B F H, B S H
+
+        x_rec = self.decoder(x_rec)  # B F PC
+        x_rec = x_rec.reshape(B, self.freq, self.seg_len, self.enc_in)
+        x_rec = x_rec.reshape(B, -1, self.enc_in)[:, :self.input_len, :]  # B L C
+
+        x_pred = self.decoder(x_pred)  # B S PC
+        x_pred = x_pred.reshape(B, self.step, self.seg_len, self.enc_in)
+        x_pred = x_pred.reshape(B, -1, self.enc_in)[:, :self.pred_len, :]  # B S C
+
+        return x_rec, x_pred
+
+
+class TimeInvKP(nn.Module):
+    """
+    Koopman Predictor with learnable Koopman operator
+    Utilize lookback and forecast window snapshots to predict the future of time-invariant term
+    """
+
+    def __init__(self,
+                 input_len=96,
+                 pred_len=96,
+                 dynamic_dim=128,
+                 encoder=None,
+                 decoder=None):
+        super().__init__()
+        self.dynamic_dim = dynamic_dim
+        self.input_len = input_len
+        self.pred_len = pred_len
+        self.encoder = encoder
+        self.decoder = decoder
+
+        K_init = torch.randn(self.dynamic_dim, self.dynamic_dim)
+        U, _, V = torch.svd(K_init)  # stable initialization
+        self.K = nn.Linear(self.dynamic_dim, self.dynamic_dim, bias=False)
+        self.K.weight.data = torch.mm(U, V.t())
+
+    def forward(self, x):
+        # x: B L C
+        res = x.transpose(1, 2)  # B C L
+        res = self.encoder(res)  # B C H
+        res = self.K(res)  # B C H
+        res = self.decoder(res)  # B C S
+        res = res.transpose(1, 2)  # B S C
+
+        return res

src/basicts/models/Koopa/config/koopa_config.py

@@ -0,0 +1,19 @@
+from dataclasses import dataclass, field
+
+from basicts.configs import BasicTSModelConfig
+
+@dataclass
+class KoopaConfig(BasicTSModelConfig):
+    """
+        Config class for Koopa model.
+    """
+    alpha: float = field(default=0.2, metadata={"help": "Scaling coefficient."})
+    enc_in: int = field(default=7, metadata={"help": "Input feature dimension."})
+    input_len: int = field(default=None, metadata={"help": "Input sequence length."})
+    output_len: int = field(default=None, metadata={"help": "Prediction length."})
+    seg_len: int = field(default=48, metadata={"help": "Segment length. Recommended: e.g., 24 for hourly data."})
+    num_blocks: int = field(default=3, metadata={"help": "Number of blocks."})
+    dynamic_dim: int = field(default=64, metadata={"help": "Dynamic feature dimension. Must be > 0."})
+    hidden_dim: int = field(default=64, metadata={"help": "Hidden dimension."})
+    hidden_layers: int = field(default=2, metadata={"help": "Number of hidden layers (>=2 recommended)."})
+    multistep: bool = field(default=False, metadata={"help": "Whether to use multistep forecasting."})