plotting in scratch.py and trying TCN

.gitignore

@@ -0,0 +1,87 @@
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# DotEnv configuration
+.env
+
+# Database
+*.db
+*.rdb
+
+# Pycharm
+.idea
+
+# VS Code
+.vscode/
+
+# Spyder
+.spyproject/
+
+# Jupyter NB Checkpoints
+.ipynb_checkpoints/
+
+# exclude data from source control by default
+/data/
+
+# Mac OS-specific storage files
+.DS_Store
+
+# vim
+*.swp
+*.swo

experiments/forecast.py

@@ -223,4 +223,4 @@ def validate(model: nn.Module,
 
 if __name__ == '__main__':
     logging.root.setLevel(logging.INFO)
-    Fire(ForecastExperiment)
+    Fire(ForecastExperiment)

mjc_notes.md

@@ -0,0 +1,14 @@
+```sh
+# try with pip torch WORKS!
+export PROJ=deeptime
+conda create -n $PROJ python=3.8 -y
+conda activate $PROJ
+mamba install -y ipykernel pip ipywidgets
+pip install torch==1.10.0+cu113 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
+# 117 does not exist yet
+python -m ipykernel install --user --name $PROJ
+pip install gin-config fire pandas matplotlib numpy scikit-learn einops tensorboard
+
+
+python -m experiments.forecast --config_path=storage/experiments/Exchange/192S/repeat=0/config.gin run >> storage/experiments/Exchange/192S/repeat=0/instance.log 2>&1%
+```

models/DeepTIMe2.py

@@ -0,0 +1,60 @@
+ # Copyright (c) 2022, salesforce.com, inc.
+ # All rights reserved.
+ # SPDX-License-Identifier: BSD-3-Clause
+ # For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+
+from typing import Optional
+
+import gin
+import torch
+import torch.nn as nn
+from torch import Tensor
+from einops import rearrange, repeat, reduce
+
+from models.modules.inrplus2 import INRPlus2
+from models.modules.regressors import RidgeRegressor
+
+
+@gin.configurable()
+def deeptime2(datetime_feats: int, layer_size: int, inr_layers: int, n_fourier_feats: int, scales: float):
+    return DeepTIMe2(datetime_feats, layer_size, inr_layers, n_fourier_feats, scales)
+
+
+class DeepTIMe2(nn.Module):
+    def __init__(self, datetime_feats: int, layer_size: int, inr_layers: int, n_fourier_feats: int, scales: float):
+        super().__init__()
+        self.inr = INRPlus2(in_feats=datetime_feats + 1, layers=inr_layers, layer_size=layer_size,
+                       n_fourier_feats=n_fourier_feats, scales=scales)
+        self.adaptive_weights = RidgeRegressor()
+
+        self.datetime_feats = datetime_feats
+        self.inr_layers = inr_layers
+        self.layer_size = layer_size
+        self.n_fourier_feats = n_fourier_feats
+        self.scales = scales
+
+    def forward(self, x: Tensor, x_time: Tensor, y_time: Tensor) -> Tensor:
+        tgt_horizon_len = y_time.shape[1]
+        batch_size, lookback_len, _ = x.shape
+        coords = self.get_coords(lookback_len, tgt_horizon_len).to(x.device)
+
+        if y_time.shape[-1] != 0:
+            time = torch.cat([x_time, y_time], dim=1)
+            coords = repeat(coords, '1 t 1 -> b t 1', b=time.shape[0])
+            coords = torch.cat([coords, time], dim=-1)
+            time_reprs = self.inr(coords)
+        else:
+            time_reprs = repeat(self.inr(coords), '1 t d -> b t d', b=batch_size)
+
+        lookback_reprs = time_reprs[:, :-tgt_horizon_len]
+        horizon_reprs = time_reprs[:, -tgt_horizon_len:]
+        w, b = self.adaptive_weights(lookback_reprs, x)
+        preds = self.forecast(horizon_reprs, w, b)
+        return preds
+
+    def forecast(self, inp: Tensor, w: Tensor, b: Tensor) -> Tensor:
+        return torch.einsum('... d o, ... t d -> ... t o', [w, inp]) + b
+
+    def get_coords(self, lookback_len: int, horizon_len: int) -> Tensor:
+        coords = torch.linspace(0, 1, lookback_len + horizon_len)
+        return rearrange(coords, 't -> 1 t 1')

models/__init__.py

@@ -3,11 +3,14 @@ from typing import Union
 import torch
 
 from .DeepTIMe import deeptime
+from .DeepTIMe2 import deeptime2
 
 
 def get_model(model_type: str, **kwargs: Union[int, float]) -> torch.nn.Module:
     if model_type == 'deeptime':
         model = deeptime(datetime_feats=kwargs['datetime_feats'])
+    elif model_type=="deeptime2":
+        model = deeptime2(datetime_feats=kwargs['datetime_feats'])
     else:
         raise ValueError(f"Unknown model type {model_type}")
     return model

models/modules/causalinception.py

@@ -0,0 +1,140 @@
+"""
+Modifie from https://github.com/timeseriesAI/tsai/blob/main/tsai/models/InceptionTimePlus.py
+"""
+from tsai.models.InceptionTimePlus import Conv, Module, noop, Integral, nn, is_listy, SimpleSelfAttention, Concat, SqueezeExciteBlock, Norm, BN1d, delegates, ConvBlock, Add, np, random, ifnone, OrderedDict, Flatten, SigmoidRange, LinBnDrop, GACP1d, GAP1d, named_partial, F, torch, CausalConv1d, Noop
+
+Conv = named_partial('Conv', ConvBlock, norm=None, act=None, padding='causal')
+# CausalConvBlock = named_partial('CausalConv', ConvBlock, padding='causal')
+
+class CausalMaxPool1d(torch.nn.MaxPool1d):
+    def __init__(self, ks, stride=1, padding=0, dilation=1):
+        super().__init__(kernel_size=ks, stride=stride, padding=0, dilation=dilation)
+        self.__padding = (ks - 1) * dilation
+    def forward(self, input):
+        return super().forward(F.pad(input, (self.__padding, 0)))
+
+class InceptionModulePlus(Module):
+    def __init__(self, ni, nf, ks=40, bottleneck=True, padding='causal', coord=False, separable=False, dilation=1, stride=1, conv_dropout=0., sa=False, se=None,
+                 norm='Batch', zero_norm=False, bn_1st=True, act=nn.ReLU, act_kwargs={}):
+
+        dilation = max(1, dilation)
+
+        if not (is_listy(ks) and len(ks) == 3):
+            if isinstance(ks, Integral): ks = [ks // (2**i) for i in range(3)]
+            ks = [ksi if ksi % 2 != 0 else ksi - 1 for ksi in ks]  # ensure odd ks for padding='same'
+
+        bottleneck = False if ni == nf else bottleneck
+        self.bottleneck = Conv(ni, nf, 1, coord=coord, bias=False) if bottleneck else noop #
+        self.convs = nn.ModuleList()
+        for i in range(len(ks)): self.convs.append(Conv(nf if bottleneck else ni, nf, ks[i], padding=padding, coord=coord, separable=separable,
+                                                         dilation=dilation**i, stride=stride, bias=False))
+        self.mp_conv = nn.Sequential(*[Conv(ni, nf, 1, coord=coord, bias=False)])
+        self.concat = Concat()
+        if norm is not None:
+            self.norm = Norm(nf * 4, norm=norm, zero_norm=zero_norm)
+        else:
+            self.norm = noop
+        self.conv_dropout = nn.Dropout(conv_dropout) if conv_dropout else noop
+        self.sa = SimpleSelfAttention(nf * 4) if sa else noop
+        self.act = act(**act_kwargs) if act else noop
+        self.se = nn.Sequential(SqueezeExciteBlock(nf * 4, reduction=se), BN1d(nf * 4)) if se else noop
+
+        self._init_cnn(self)
+    
+    def _init_cnn(self, m):
+        if getattr(self, 'bias', None) is not None: nn.init.constant_(self.bias, 0)
+        if isinstance(self, (nn.Conv1d,nn.Conv2d,nn.Conv3d,nn.Linear)): nn.init.kaiming_normal_(self.weight)
+        for l in m.children(): self._init_cnn(l)
+
+    def forward(self, x):
+        input_tensor = x
+        x = self.bottleneck(x)
+        x = self.concat([l(x) for l in self.convs] + [self.mp_conv(input_tensor)])
+        x = self.norm(x)
+        x = self.conv_dropout(x)
+        x = self.sa(x)
+        x = self.act(x)
+        x = self.se(x)
+        return x
+
+
+@delegates(InceptionModulePlus.__init__)
+class InceptionBlockPlus(Module):
+    def __init__(self, ni, nf, residual=True, depth=6, coord=False, norm=None, zero_norm=False, act=nn.ReLU, act_kwargs={}, sa=False, se=None, dilation=1,
+                 stoch_depth=1., **kwargs):
+        self.residual, self.depth = residual, depth
+        self.inception, self.shortcut, self.act = nn.ModuleList(), nn.ModuleList(), nn.ModuleList()
+        for d in range(depth):
+            self.inception.append(InceptionModulePlus(ni if d == 0 else nf * 4, nf, coord=coord, norm=norm,
+                                                      zero_norm=zero_norm if d % 3 == 2 else False,
+                                                      act=act if d % 3 != 2 else None, act_kwargs=act_kwargs,
+                                                      sa=sa if d % 3 == 2 else False,
+                                                      se=se if d % 3 != 2 else None,
+                                                      dilation=dilation*d*(dilation>1),
+                                                      **kwargs))
+            if self.residual and d % 3 == 2:
+                n_in, n_out = ni if d == 2 else nf * 4, nf * 4
+                if norm is not None:
+                    n = Norm(n_in, norm=norm)
+                else:
+                    n = Noop
+                self.shortcut.append(n if n_in == n_out else ConvBlock(n_in, n_out, 1, coord=coord, bias=False, norm=norm, padding='causal', act=None))
+                self.act.append(act(**act_kwargs))
+        self.add = Add()
+        if stoch_depth != 0: keep_prob = np.linspace(1, stoch_depth, depth)
+        else: keep_prob = np.array([1] * depth)
+        self.keep_prob = keep_prob
+
+    def forward(self, x):
+        res = x
+        for i in range(self.depth):
+            if self.keep_prob[i] > random.random() or not self.training:
+                x = self.inception[i](x)
+            if self.residual and i % 3 == 2:
+                res = x = self.act[i//3](self.add(x, self.shortcut[i//3](res)))
+        return x
+
+# Cell
+@delegates(InceptionModulePlus.__init__)
+class CausalInceptionTimePlus(nn.Sequential):
+    def __init__(self, c_in, c_out, seq_len=None, nf=32, nb_filters=None,
+                 flatten=False, concat_pool=False, fc_dropout=0., bn=False, y_range=None, custom_head=None, **kwargs):
+
+        if nb_filters is not None: nf = nb_filters
+        else: nf = ifnone(nf, nb_filters) # for compatibility
+        backbone = InceptionBlockPlus(c_in, nf, **kwargs)
+
+        #head
+        self.head_nf = nf * 4
+        self.c_out = c_out
+        self.seq_len = seq_len
+        if custom_head: head = custom_head(self.head_nf, c_out, seq_len)
+        else: head = self.create_head(self.head_nf, c_out, seq_len, flatten=flatten, concat_pool=concat_pool,
+                                      fc_dropout=fc_dropout, bn=bn, y_range=y_range)
+
+        layers = OrderedDict([('backbone', nn.Sequential(backbone)), ('head', nn.Sequential(head))])
+        super().__init__(layers)
+
+        self.calc_receptive_field(kwargs.get('ks'), kwargs.get('depth'), kwargs.get('dilation', 1))
+
+    def calc_receptive_field(self, ks, depth, dilation):       
+        # receptive fields vs R
+        ks=np.array(ks)
+        d=np.array([dilation**i for i in range(3)])
+        rf = (ks-1)*d*depth
+
+        dilations = np.array([max(1, d*dilation) for d in range(depth)])
+        d=np.array([dilations**i for i in range(3)]).T
+        rf = ((ks-1)*d).sum(0)
+        print(f"receptive field {rf}={ks-1}*{d}")
+
+    def create_head(self, nf, c_out, seq_len, flatten=False, concat_pool=False, fc_dropout=0., bn=False, y_range=None):
+        if flatten:
+            nf *= seq_len
+            layers = [Flatten()]
+        else:
+            if concat_pool: nf *= 2
+            layers = [GACP1d(1) if concat_pool else GAP1d(1)]
+        layers += [LinBnDrop(nf, c_out, bn=bn, p=fc_dropout)]
+        if y_range: layers += [SigmoidRange(*y_range)]
+        return nn.Sequential(*layers)

models/modules/inrplus2.py

@@ -0,0 +1,43 @@
+# Copyright (c) 2022, salesforce.com, inc.
+# All rights reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+# For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from models.modules.feature_transforms import GaussianFourierFeatureTransform
+
+from tsai.models.InceptionTimePlus import InceptionTimePlus
+from .causalinception import CausalInceptionTimePlus, CausalConv1d
+
+def custom_head(head_nf, c_out, seq_len):
+    return nn.Sequential(
+        CausalConv1d(head_nf, c_out, 1, bias=False)
+        
+    )
+
+class INRPlus2(nn.Module):
+    def __init__(self, in_feats: int, layers: int, layer_size: int, n_fourier_feats: int, scales: float,
+                 dropout: Optional[float] = 0.5, bn=False, *args, **kwargs):
+        super().__init__()
+        self.features = nn.Linear(in_feats, layer_size) if n_fourier_feats == 0 \
+            else GaussianFourierFeatureTransform(in_feats, n_fourier_feats, scales)
+        in_size = layer_size if n_fourier_feats == 0 \
+            else n_fourier_feats+in_feats
+        # import pdb; pdb.set_trace()
+        self.layers = CausalInceptionTimePlus(
+            in_size-1, layer_size, seq_len=None, nf=layer_size, depth=layers,
+                 flatten=False, concat_pool=False, fc_dropout=dropout, conv_dropout=0.05, bn=bn, y_range=None, custom_head=custom_head, ks=[139, 19, 3], dilation=2, *args, **kwargs
+        )
+        # layers = [INRPlusLayer(in_size, layer_size, dropout=dropout)] + \
+        #          [INRPlusLayer(layer_size, layer_size, dropout=dropout) for _ in range(layers - 1)]
+        # self.layers = nn.Sequential(*layers)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.features(x)
+        # import pdb; pdb.set_trace()
+        return self.layers(x.permute((0, 2, 1))).permute((0, 2, 1))

models/modules/regressors.py

@@ -29,15 +29,17 @@ class RidgeRegressor(nn.Module):
 
         if n_samples >= n_dim:
             # standard
-            A = torch.bmm(X.mT, X)
+            A = torch.bmm(X.transpose(-2, -1), X)
             A.diagonal(dim1=-2, dim2=-1).add_(reg_coeff)
-            B = torch.bmm(X.mT, Y)
+            B = torch.bmm(X.transpose(-2, -1), Y)
             weights = torch.linalg.solve(A, B)
         else:
             # Woodbury
-            A = torch.bmm(X, X.mT)
+            # A = torch.bmm(X, X.mT)
+            A = torch.bmm(X, X.transpose(-2, -1))
             A.diagonal(dim1=-2, dim2=-1).add_(reg_coeff)
-            weights = torch.bmm(X.mT, torch.linalg.solve(A, Y))
+            # weights = torch.bmm(X.mT, torch.linalg.solve(A, Y))
+            weights = torch.bmm(X.transpose(-2, -1), torch.linalg.solve(A, Y))
 
         return weights[:, :-1], weights[:, -1:]
 

scratch.py

@@ -0,0 +1,194 @@
+# %%
+
+import os
+from os.path import join
+import math
+import logging
+from typing import Callable, Optional, Union, Dict, Tuple
+
+from matplotlib import pyplot as plt
+
+import gin
+from fire import Fire
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+from torch import optim
+from torch import nn
+
+from experiments.base import Experiment
+from data.datasets import ForecastDataset
+from models import get_model
+from utils.checkpoint import Checkpoint
+from utils.ops import default_device, to_tensor
+from utils.losses import get_loss_fn
+from utils.metrics import calc_metrics
+
+from experiments.forecast import get_data
+gin.enter_interactive_mode()
+# %%
+
+
+gin.clear_config()
+# gin.parse_config(open("storage/experiments/Exchange/96M/repeat=0/config.gin"))
+gin.parse_config(open("storage/experiments/Exchange/96Mplus/repeat=0/config.gin"))
+
+# %%
+train_set, train_loader = get_data(flag='train', batch_size=16)
+# x, _, _, _ =train_set[0]
+# x = x * 1.0
+
+# %%
+# x -= x[0]
+# x /= x.std()
+# plt.plot(x)
+# # %%
+
+
+# %%
+model = get_model("deeptime2",
+                    dim_size=train_set.data_x.shape[1],
+                    datetime_feats=train_set.timestamps.shape[-1]).to(default_device())
+model.load_state_dict(torch.load('storage/experiments/Exchange/96Mplus/repeat=0/model.pth'))
+model = model.eval()
+
+# %%
+b = train_set[1]
+b = [bb[None, :] for bb in b]
+x, y, x_time, y_time = map(to_tensor, b)
+with torch.no_grad():
+    forecast = model(x, x_time, y_time)
+# %%
+
+
+# %%
+plt.title('inception inr')
+import matplotlib.colors as mcolors
+colors = list(mcolors.BASE_COLORS.keys())
+l = x.shape[1]
+forecast2 = forecast[0].detach().cpu().numpy()
+x2 = x[0].cpu()
+y2 = y[0].cpu()
+i_past = list(range(l))
+i_future = list(range(l, l*2))
+for i in range(x.shape[-1]):
+    plt.plot(range(l), x2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(range(l, l*2), y2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(range(l, l*2), forecast2[:, i], c=colors[i], linestyle='--')
+
+# %%
+gin.clear_config()
+gin.parse_config(open("storage/experiments/Exchange/96M/repeat=0/config.gin"))
+
+train_set, train_loader = get_data(flag='train', batch_size=16)
+
+model = get_model("deeptime",
+                    dim_size=train_set.data_x.shape[1],
+                    datetime_feats=train_set.timestamps.shape[-1]).to(default_device())
+model.load_state_dict(torch.load('storage/experiments/Exchange/96M/repeat=0/model.pth'))
+model = model.eval()
+
+
+b = train_set[1]
+b = [bb[None, :] for bb in b]
+x, y, x_time, y_time = map(to_tensor, b)
+with torch.no_grad():
+    forecast = model(x, x_time, y_time)
+
+
+plt.title('mlp inr')
+import matplotlib.colors as mcolors
+colors = list(mcolors.BASE_COLORS.keys())
+l = x.shape[1]
+forecast2 = forecast[0].detach().cpu().numpy()
+x2 = x[0].cpu()
+y2 = y[0].cpu()
+i_past = list(range(l))
+i_future = list(range(l, l*2))
+for i in range(x.shape[-1]):
+    plt.plot(range(l), x2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(range(l, l*2), y2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(range(l, l*2), forecast2[:, i], c=colors[i], linestyle='--')
+# %%
+
+
+gin.clear_config()
+gin.parse_config(open("storage/experiments/Exchange/96Mplus2/repeat=0/config.gin"))
+
+train_set, train_loader = get_data(flag='train', batch_size=16)
+
+model = get_model("deeptime2",
+                    dim_size=train_set.data_x.shape[1],
+                    datetime_feats=train_set.timestamps.shape[-1]).to(default_device())
+model.load_state_dict(torch.load('storage/experiments/Exchange/96Mplus2/repeat=0/model.pth'))
+model = model.eval()
+
+
+b = train_set[1]
+b = [bb[None, :] for bb in b]
+x, y, x_time, y_time = map(to_tensor, b)
+with torch.no_grad():
+    forecast = model(x, x_time, y_time)
+
+
+plt.title('inception inr')
+import matplotlib.colors as mcolors
+colors = list(mcolors.BASE_COLORS.keys())
+l = x.shape[1]
+forecast2 = forecast[0].detach().cpu().numpy()
+x2 = x[0].cpu()
+y2 = y[0].cpu()
+l2 = y.shape[1]
+i_past = list(range(l))
+i_future = list(range(l, l+l2))
+for i in range(x.shape[-1]):
+    plt.plot(i_past, x2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(i_future, y2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(i_future, forecast2[:, i], c=colors[i], linestyle='--')
+
+# %%
+
+
+gin.clear_config()
+gin.parse_config(open("storage/experiments/Exchange/96M2/repeat=0/config.gin"))
+
+train_set, train_loader = get_data(flag='train', batch_size=16)
+
+model = get_model("deeptime",
+                    dim_size=train_set.data_x.shape[1],
+                    datetime_feats=train_set.timestamps.shape[-1]).to(default_device())
+model.load_state_dict(torch.load('storage/experiments/Exchange/96M2/repeat=0/model.pth'))
+model = model.eval()
+
+
+b = train_set[1]
+b = [bb[None, :] for bb in b]
+x, y, x_time, y_time = map(to_tensor, b)
+with torch.no_grad():
+    forecast = model(x, x_time, y_time)
+
+
+plt.title('mlp inr2')
+import matplotlib.colors as mcolors
+colors = list(mcolors.BASE_COLORS.keys())
+l = x.shape[1]
+forecast2 = forecast[0].detach().cpu().numpy()
+x2 = x[0].cpu()
+y2 = y[0].cpu()
+l2 = y.shape[1]
+i_past = list(range(l))
+i_future = list(range(l, l+l2))
+for i in range(x.shape[-1]):
+    plt.plot(i_past, x2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(i_future, y2[:, i], c=colors[i])
+for i in range(x.shape[-1]):
+    plt.plot(i_future, forecast2[:, i], c=colors[i], linestyle='--')
+
+# %%

storage/datasets/.gitignore

@@ -1,4 +1,5 @@
 # Ignore everything in this directory
 *
 # Except this file
-!.gitignore
+!.gitignore
+