Official Implementation of  SimStock : Representation Model for Stock Similarities
import torch.nn.functional as F
import torch
from tqdm import tqdm
import numpy as np
import os
#
from torch.autograd import Variable

def make_noise(shape, type="Gaussian"):
    """
    Generate random noise.
    Parameters
    ----------
    shape: List or tuple indicating the shape of the noise
    type: str, "Gaussian" or "Uniform", default: "Gaussian".
    Returns
    -------
    noise tensor
    """

    if type == "Gaussian":
        noise = Variable(torch.randn(shape))
    elif type == 'Uniform':
        noise = Variable(torch.randn(shape).uniform_(-1, 1))
    else:
        raise Exception("ERROR: Noise type {} not supported".format(type))
    return noise

########Supervised###########

def train(dataloader, optimizer, model, args, log, device, input_E, input_hidden, task_id=0):
    E = input_E
    hidden = input_hidden
    loss_list = []
    
    log("Start Training on Domain {}...".format(task_id))
    for epoch in range(args.epoches):
        with tqdm(dataloader, unit="batch") as tepoch:
            for X, S in tepoch:
                X, S  = X.float().to(device), S.long().to(device)            
                initial_noise = make_noise((1, args.noise_dim), args.noise_type).to(device)
                model.train()
                optimizer.zero_grad()
                E, hidden, loss = model(X, initial_noise, S, E, hidden, False)
                E = E.detach()
                hidden = tuple([i.detach() for i in hidden])
                
                loss.backward()
                optimizer.step()    
                loss_list.append(loss.item())
                
        log("Task_ID: {}\tEpoch: {}\tAverage Training Loss: {}".format(task_id, epoch, np.mean(loss_list)))
    
    if not os.path.isdir('weight'):
        os.makedirs('weight')
        
    torch.save(model.state_dict(), './weight/model_weights_{}.pth'.format(args.save_name))
    
    return E, hidden, model

def test(dataloader, model, args, log, device, input_E, input_hidden, is_repre):
    model.eval()    
    E = input_E
    hidden = input_hidden
    loss_l = []
    representation_l = []
    attn_l = []
    
    with tqdm(dataloader, unit="batch") as tepoch:
        for X, S in tepoch:
            X, S  = X.float().to(device), S.long().to(device)
            initial_noise = make_noise((1, args.noise_dim), args.noise_type).to(device)
            
            with torch.no_grad():
                if not is_repre:
                    _, _, loss = models(X, initial_noise, S, E, hidden, False)
                    loss_l.append(loss.item())
                else:
                    representation, attn_ = models(X, initial_noise, S, E, hidden, True)
                    representation_l.append(representation)
                    attn_l.append(attn_)
                    
    if is_repre:
        return representation_l
    else:
        log("Average Testing Error is {}".format(np.mean(loss_l)))

    

def test_only_inference(dataloader, model, args, log, device, input_E, input_hidden, is_repre):
    model.eval()    
    E = input_E
    hidden = input_hidden
    loss_l = []
    representation_l = []
    attn_l = []
    
    model.load_state_dict(torch.load(os.path.join('./weight/', 'model_weights_{}.pth'.format(args.save_name))))
    # TODO : Must save the E and hiddens
    
    with tqdm(dataloader, unit="batch") as tepoch:
        for X, S in tepoch:
            X, S  = X.float().to(device), S.long().to(device)
            initial_noise = make_noise((1, args.noise_dim), args.noise_type).to(device)
            
            with torch.no_grad():
                if not is_repre:
                    _, _, loss = models(X, initial_noise, S, E, hidden, False)
                    loss_l.append(loss.item())
                else:
                    representation, attn_ = models(X, initial_noise, S, E, hidden, True)
                    representation_l.append(representation)
                    attn_l.append(attn_)
                    
    if is_repre:
        return representation_l
    else:
        log("Average Testing Error is {}".format(np.mean(loss_l)))

exp/training.py
Yoontae6719/SimStock-Representation-Model-for-Stock-Similarities
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils import weight_norm
import math


class PositionalEmbedding(nn.Module):
    def __init__(self, d_model, max_len=5000):
        super(PositionalEmbedding, self).__init__()
        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model).float()
        pe.require_grad = False

        position = torch.arange(0, max_len).float().unsqueeze(1)
        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        return self.pe[:, :x.size(1)]


class TokenEmbedding(nn.Module):
    def __init__(self, c_in, d_model):
        super(TokenEmbedding, self).__init__()
        padding = 1 if torch.__version__ >= '1.5.0' else 2
        self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model, kernel_size=3, padding=padding, padding_mode='circular', bias=False)
        for m in self.modules():
            if isinstance(m, nn.Conv1d):
                nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')

    def forward(self, x):
        x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)
        return x


class FixedEmbedding(nn.Module):
    def __init__(self, c_in, d_model):
        super(FixedEmbedding, self).__init__()

        w = torch.zeros(c_in, d_model).float()
        w.require_grad = False

        position = torch.arange(0, c_in).float().unsqueeze(1)
        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()

        w[:, 0::2] = torch.sin(position * div_term)
        w[:, 1::2] = torch.cos(position * div_term)

        self.emb = nn.Embedding(c_in, d_model)
        self.emb.weight = nn.Parameter(w, requires_grad=False)

    def forward(self, x):
        return self.emb(x).detach()


class TemporalEmbedding(nn.Module):
    def __init__(self, d_model, embed_type='fixed', freq='h'):
        super(TemporalEmbedding, self).__init__()

        minute_size = 4
        hour_size = 24
        weekday_size = 7
        day_size = 32
        month_size = 13

        Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding
        if freq == 't':
            self.minute_embed = Embed(minute_size, d_model)
        self.hour_embed       = Embed(hour_size, d_model)
        self.weekday_embed    = Embed(weekday_size, d_model)
        self.day_embed        = Embed(day_size, d_model)
        self.month_embed      = Embed(month_size, d_model)

    def forward(self, x):
        x = x.long()

        minute_x = self.minute_embed(x[:, :, 4]) if hasattr(self, 'minute_embed') else 0.
        hour_x = self.hour_embed(x[:, :, 3])
        weekday_x = self.weekday_embed(x[:, :, 2])
        day_x = self.day_embed(x[:, :, 1])
        month_x = self.month_embed(x[:, :, 0])

        return hour_x + weekday_x + day_x + month_x + minute_x


class TimeFeatureEmbedding(nn.Module):
    def __init__(self, d_model, embed_type='timeF', freq='h'):
        super(TimeFeatureEmbedding, self).__init__()

        freq_map = {'h': 4, 't': 5, 's': 6, 'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}
        d_inp = freq_map[freq]
        self.embed = nn.Linear(d_inp, d_model, bias=False)

    def forward(self, x):
        return self.embed(x)


class DataEmbedding(nn.Module):
    def __init__(self, c_in, d_model,len_sector, dropout=0.1):
        super(DataEmbedding, self).__init__()

        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model) # OHLCV embedding
        self.sector_embedding = nn.Embedding(len_sector, d_model) # Sector embedding
        self.position_embedding = PositionalEmbedding(d_model=d_model)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, x_mark):
        x = x + self.sector_embedding(x_mark) #+ self.sector_embedding(sector)
        return self.dropout(x)


class DataEmbedding_wo_pos(nn.Module):
    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0.1):
        super(DataEmbedding_wo_pos, self).__init__()

        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
        self.position_embedding = PositionalEmbedding(d_model=d_model)
        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
            d_model=d_model, embed_type=embed_type, freq=freq)
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, x, x_mark):
        x = self.value_embedding(x) + self.temporal_embedding(x_mark)
        return self.dropout(x)

class Sector_embedding(nn.Module):
    def __init__(self, d_model, len_sector, dropout=0.1):
        super(Sector_embedding, self).__init__()  # Call parent class constructor
        self.sector_embedding = nn.Embedding(len_sector, d_model)
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, sector):
        x = self.sector_embedding(sector)
        return self.dropout(x)

class Sector_and_exchange_embedding(nn.Module):
    def __init__(self, d_model, len_sector, len_exchange, dropout=0.1):
        super(Sector_and_exchange_embedding, self).__init__()  # Call parent class constructor
        self.sector_embedding = nn.Embedding(len_sector, d_model)
        self.ex_embedding = nn.Embedding(len_exchange, d_model)

        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, sector, exchange):
        x = self.sector_embedding(sector) +  self.sector_embedding(exchange)
        return self.dropout(x)

    
    
class StockEmbedding_w_pos(nn.Module):
    def __init__(self, origin_input, d_model, len_sector, dropout = 0.1):
        super(StockEmbedding_w_pos, self).__init__()
        
        self.value_embedding = TokenEmbedding(c_in=origin_input, d_model=d_model) # Value embedding
        self.position_embedding = PositionalEmbedding(d_model=d_model)            # tickers
        self.sector_embedding = nn.Embedding(len_sector, d_model)                 # Sector
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, value, sector):
        x = self.value_embedding(value)  + self.sector_embedding(sector) # + self.position_embedding(tickers)
        return self.dropout(value)
        
class PV_embedding(nn.Module):
    def __init__(self, seq_len, d_model):
        super(PV_embedding, self).__init__()
        self.seq_len = seq_len
        self.linears = nn.ModuleList([nn.Linear(int(seq_len/(2**i)), seq_len) for i in range(4)])
        self.out_linear = nn.Linear(d_model * 4, d_model)
        
    def forward(self, x):
        batch, seq, dim = x.shape
        
        origins = []
        for i in range(4):
            origin_i = self.linears[i](x[:, 0::2**i, :].transpose(1, 2))
            origins.append(origin_i)
        
        out = torch.cat(origins, dim=1)
        #out = out.view(batch, 4, seq, dim)
        out = self.out_linear(out.permute(0, 2, 1)) # Noted that dim * 4
        
        return out #out.permute(0, 2, 1) # Noted that dim * 4
    
    
class StockEmbedding(nn.Module):
    def __init__(self, c_in, len_sector, d_model, dropout = 0.1):
        super(StockEmbedding, self).__init__()
        
        self.temporal_embedding = TimeFeatureEmbedding(d_model)
    
        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
        self.sector_embedding = nn.Embedding(len_sector, d_model)                 # Sector
        self.dropout = nn.Dropout(p=dropout)
        
    def forward(self, value, time):
        x =  self.value_embedding(value) + self.temporal_embedding(time) 
        return self.dropout(x)
    
layers/Embed.py
Yoontae6719/SimStock-Representation-Model-for-Stock-Similarities
# -*- coding: utf-8 -*-

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

import os
import logging
import time
import datetime
from tqdm import tqdm
import argparse

import os
import logging
import time
import pandas as pd

from models.Simstock import model
from utils.helper import make_noise
from utils.prepro import dataset_for_modeling
from exp.training import train, test


# create logger
logger = logging.getLogger('main')
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

# create console handler and set level to debug
ch = logging.StreamHandler()
ch.setLevel(logging.INFO)
ch.setFormatter(formatter)
logger.addHandler(ch)


def log(str): logger.info(str)


parser = argparse.ArgumentParser(description="SimStock")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

datasets = ['train_nasdaq', 'train_sse', 'train_szse', 'train_tse']
test_datasets = ['test_nasdaq', 'tset_sse', 'test_szse', 'test_tse', "ARKK", "LIT", "BOTZ", "SKYY"]

# dataset param
parser.add_argument("--train_dataset", default="train_nasdaq", type=str, help="one of: {}".format(", ".join(sorted(datasets))))
parser.add_argument("--test_dataset", default="test_nasdaq", type=str, help="one of: {}".format(", ".join(sorted(test_datasets))))
parser.add_argument("--batch_size", default=512, type=int,      help="the number of epoches for each task.")
parser.add_argument("--data_size", default=25, type=int,      help="the number of input features.")

# model param
parser.add_argument("--noise_dim", default=25, type=float,     help="the dimension of the LSTM input noise.")
parser.add_argument("--latent_dim", default=25, type=float,     help="the latent dimension of RNN variables.")
parser.add_argument("--hidden_dim", default=128, type=float,     help="the latent dimension of RNN variables.")
parser.add_argument("--noise_type", choices=["Gaussian", "Uniform"], default="Gaussian", help="The noise type to feed into the generator.")
parser.add_argument("--num_rnn_layer", default=1, type=float,   help="the number of RNN hierarchical layers.")
parser.add_argument("--sector_size", default=138, type=int,help="the number of sector size. WARNING : total + 1")
parser.add_argument("--sector_emb", default=256, type=int,help="the number of sector embedding size")
parser.add_argument("--lambda_values", default=0.7, type=float,help="the number of sector argument")

# training param
parser.add_argument("--learning_rate", default=1e-3, type=float,help="the unified learning rate for each single task.")
parser.add_argument("--epoches", default=5, type=int, help="the number of epoches for each task.")
parser.add_argument("--save_name", default="test", type=str,help="model save weight")

args = parser.parse_args([])


def main(arsgs):
    train_out = dataset_for_modeling(args, train_type = False)
    test_out = dataset_for_modeling(args, train_type = True)
    models =  model(args, device).to(device)
    optimizer = torch.optim.Adam(models.parameters(), lr=args.learning_rate)

    starting_time = time.time()


    Es, hiddens = [None], [None]
    for task_id, dataloader in enumerate(train_out[:-1]):
        E, hidden, rnn_unit = train(dataloader, optimizer, models, args, log, device, Es[-1], hiddens[-1], task_id)
        Es.append(E)
        hiddens.append(hidden)    
    ending_time = time.time()

    print("Training time:", ending_time - starting_time)

    # Testing
    representation_l = test(train_out[-1], model, args, log, device, Es[-1], hiddens[-1], is_repre = True)
    test(train_out[-1], model, args, log, device, Es[-1], hiddens[-1], is_repre = False)

    # Testing 2
    representation_ll = test(test_out[-1], model, args, log, device, Es[-1], hiddens[-1], is_repre = True)
    
    return representation_l, representation_ll


if __name__ == "__main__":
    print("Start Training and get embeddings")
    r1, r2 = main(args)

main.py
Yoontae6719/SimStock-Representation-Model-for-Stock-Similarities
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from functools import wraps
from math import sqrt
from einops import rearrange, repeat
import copy
from layers.Embed import Sector_embedding


def loss_fn(x, y):
    x = torch.nn.functional.normalize(x, dim=-1, p=2)
    y = torch.nn.functional.normalize(y, dim=-1, p=2)
    return 2 - 2 * (x * y).sum(dim=-1)

triplet_loss = nn.TripletMarginLoss(margin=1.0, p=2)


# helper
def default(val, def_val):
    return def_val if val is None else val

def flatten(t):
    return t.reshape(t.shape[0], -1)

def singleton(cache_key):
    def inner_fn(fn):
        @wraps(fn)
        def wrapper(self, *args, **kwargs):
            instance = getattr(self, cache_key)
            if instance is not None:
                return instance

            instance = fn(self, *args, **kwargs)
            setattr(self, cache_key, instance)
            return instance
        return wrapper
    return inner_fn


def set_requires_grad(model, val):
    for p in model.parameters():
        p.requires_grad = val
        
        
# Notice that we donot tune our model.


class GEGLU(nn.Module):
    def forward(self, x):
        x, gates = x.chunk(2, dim = -1)
        return x * F.gelu(gates)

    
class NumericalEmbedder(nn.Module):
    def __init__(self, dim, num_numerical_types):
        super().__init__()
        self.weights = nn.Parameter(torch.randn(num_numerical_types, dim))
        self.biases = nn.Parameter(torch.randn(num_numerical_types, dim))

    def forward(self, x):
        x = rearrange(x, 'b n -> b n 1')
        return x * self.weights + self.biases
    
def FeedForward(dim, mult = 4, dropout = 0.):
    return nn.Sequential(
        nn.LayerNorm(dim),
        nn.Linear(dim, dim * mult * 2),
        GEGLU(),
        nn.Dropout(dropout),
        nn.Linear(dim * mult, dim)
    )

class model(nn.Module):
    def __init__(self, args, device):
        super(model, self).__init__()
        
        self.lambda_values =  args.lambda_values
        
        self.init_lin_h = nn.Linear(args.noise_dim, args.latent_dim)
        self.init_lin_c = nn.Linear(args.noise_dim, args.latent_dim)
        self.init_input = nn.Linear(args.noise_dim, args.latent_dim)

        self.rnn = nn.LSTM(args.latent_dim, args.latent_dim, args.num_rnn_layer)
        
        self.NumericalEmbedder = NumericalEmbedder(args.sector_emb,args.latent_dim)
        
        # Transforming LSTM output to vector shape
        self.lin_transform_down = nn.Sequential(
                            nn.Linear(args.latent_dim, args.hidden_dim),
                            nn.ReLU(),
                            nn.Linear(args.hidden_dim, (args.data_size+1)*256 + 256*256 + 256 + 256+ 1))
        
        # Transforming vector to LSTM input shape
        self.lin_transform_up = nn.Sequential(
                            nn.Linear( (args.data_size+1)*256 + 256*256 + 256 + 256+ 1, args.hidden_dim),
                            nn.ReLU(),
                            nn.Linear(args.hidden_dim, args.latent_dim))
        
        self.num_rnn_layer = args.num_rnn_layer
        self.data_size = args.data_size
        self.device = device
        
        # Sector embedding
        self.sector_emb = Sector_embedding(args.sector_emb, args.sector_size)
        self.sector_projection = nn.Linear(args.sector_emb, args.noise_dim) # input size
        
        # Projection term
        
        self.online_encoder = nn.Linear(args.noise_dim, args.noise_dim)
        
        self.target_encoder = None
        
        ### Attn params
        self.cls_token = nn.Parameter(torch.randn(1, 1, args.sector_emb))
        self.scale = (256 // 4) ** -0.5
        self.dropout = nn.Dropout(0.05)
        self.depth = 2
        
        self.mlp = FeedForward(256) #FeedForward(256)
        self.norm = nn.LayerNorm(256)
    
    def nn_construction_attns(self, E):
        # First layers
        attn_qw  = E[:, :(self.data_size + 1)* 256 ] # dim * heads
        attn_kw  = E[:, :(self.data_size + 1)* 256 ] 
        attn_vw  = E[:, :(self.data_size + 1)* 256 ]
        
        attn_out_w = E[:, (self.data_size + 1)* 256 : (self.data_size + 1)* 256 + 256*256 ]
        attn_out_b = E[:, (self.data_size + 1)* 256 + 256*256 : (self.data_size + 1)* 256 + 256*256 + 256]

        out_m = E[:, (self.data_size + 1)* 256 + 256*256 + 256 : (self.data_size + 1)* 256 + 256*256 + 256 +  256]
        out_b = E[:, -1]

        return  [attn_qw.view((self.data_size +1), 256), 
                 attn_kw.view((self.data_size +1), 256), 
                 attn_vw.view((self.data_size +1), 256),
                 
                 attn_out_w.view(-1, 256),
                 attn_out_b.view(-1, 256),
                 
                 out_m.view(-1, 1),
                 out_b
                ] 
    
    def forward(self, X, z, sector, E=None, hidden=None, return_embedding = False):
        s = self.sector_emb(sector.squeeze(1))
        s = self.sector_projection(s)
        
        if hidden == None and E == None:
            init_c, init_h = [], []
            for _ in range(self.num_rnn_layer):
                init_c.append(torch.tanh(self.init_lin_h(z)))
                init_h.append(torch.tanh(self.init_lin_c(z)))
            # Initialize hidden inputs for the LSTM
            hidden = (torch.stack(init_c, dim=0), torch.stack(init_h, dim=0))
        
            # Initialize an input for the LSTM
            inputs = torch.tanh(self.init_input(z))
        else:
            inputs = self.lin_transform_up(E)

        out, hidden = self.rnn(inputs.unsqueeze(0), hidden)

        E = self.lin_transform_down(out.squeeze(0))
                
        #### add cls toekn ####
        X = torch.add(X, s)
        X_out = X # Anchor param
        
        X = self.NumericalEmbedder(X)
        cls_tokens = repeat(self.cls_token, '1 1 d -> b 1 d', b = X.shape[0])
        X = torch.cat((cls_tokens, X), dim = 1) # X shape : torch.Size([256, 26, 256])
        
        
        ############ Attn (1-depth)####################
        m_list = self.nn_construction_attns(E)        
        
        if return_embedding == True:

            q = torch.einsum('bnd,nd->bnd', X, m_list[0]) # attn_q
            k = torch.einsum('bnd,nd->bnd', X, m_list[1]) # attn_k
            v = torch.einsum('bnd,nd->bnd', X, m_list[2]) # attn_v

            q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = 4), (q, k, v))
            q = q * self.scale
            sim = torch.einsum('b h i d, b h j d -> b h i j', q, k)
            attn = sim.softmax(dim = -1)
            dropped_attn = self.dropout(attn)

            out = torch.einsum('b h i j, b h j d -> b h i d', dropped_attn, v)
            out = rearrange(out, 'b h n d -> b n (h d)', h = 4)


            out = torch.einsum('b f i, i j-> b f j', out, m_list[3]) + m_list[4]
            output = out + X

            output = torch.einsum('b f i, i j-> b f j', output, m_list[5]) + m_list[6]
            output = output[:, 0]
            return output, attn
        #############################################
        
        positive, negative = self.augment(X, self.lambda_values)
        # f_\theta (weight sharing)

        pos_q = torch.einsum('bnd,nd->bnd', positive, m_list[0]) # attn_q
        pos_k = torch.einsum('bnd,nd->bnd', positive, m_list[1]) # attn_k
        pos_v = torch.einsum('bnd,nd->bnd', positive, m_list[2]) # attn_v
        pos_q, pos_k, pos_v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = 4), (pos_q, pos_k, pos_v))
        pos_q = pos_q * self.scale
        pos_sim = torch.einsum('b h i d, b h j d -> b h i j', pos_q, pos_k)
        pos_attn = pos_sim.softmax(dim = -1)
        pos_dropped_attn = self.dropout(pos_attn)
        pos_out = torch.einsum('b h i j, b h j d -> b h i d', pos_dropped_attn, pos_v)
        pos_out = rearrange(pos_out, 'b h n d -> b n (h d)', h = 4)
        pos_out = torch.einsum('b f i, i j-> b f j', pos_out, m_list[3]) + m_list[4]
        pos_output = pos_out + positive
        pos_output = torch.einsum('b f i, i j-> b f j', pos_output, m_list[5]) + m_list[6]
        pos_output = pos_output[:, 0]
        
        neg_q = torch.einsum('bnd,nd->bnd', negative, m_list[0]) # attn_q
        neg_k = torch.einsum('bnd,nd->bnd', negative, m_list[1]) # attn_k
        neg_v = torch.einsum('bnd,nd->bnd', negative, m_list[2]) # attn_v
        neg_q, neg_k, neg_v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = 4), (neg_q, neg_k, neg_v))
        neg_q = neg_q * self.scale
        neg_sim = torch.einsum('b h i d, b h j d -> b h i j', neg_q, neg_k)
        neg_attn = neg_sim.softmax(dim = -1)
        neg_dropped_attn = self.dropout(neg_attn)
        neg_out = torch.einsum('b h i j, b h j d -> b h i d', neg_dropped_attn, neg_v)
        neg_out = rearrange(neg_out, 'b h n d -> b n (h d)', h = 4)
        neg_out = torch.einsum('b f i, i j-> b f j', neg_out, m_list[3]) + m_list[4]
        neg_output = neg_out + negative
        neg_output = torch.einsum('b f i, i j-> b f j', neg_output, m_list[5]) + m_list[6]
        neg_output = neg_output[:, 0]

        loss = triplet_loss(X_out, pos_output, neg_output)
        
        
        return E, hidden, loss
    
    @singleton('target_encoder')
    def _get_target_encoder(self):
        target_encoder = copy.deepcopy(self.online_encoder)
        set_requires_grad(target_encoder, False)
        return target_encoder    
    
    def augment(self, x, lambda_value, use_cuda=True):
        batch_size, token, feature_dim = x.shape
        if use_cuda:
            index1 = torch.randperm(feature_dim).cuda()
            index2 = torch.randperm(feature_dim).cuda()
        else:
            index1 = torch.randperm(feature_dim)
            index2 = torch.randperm(feature_dim)

        mixed_x1 = lambda_value * x + (1 - lambda_value) * x[:,  :, index1]
        mixed_x2 = (1 - lambda_value) * x + lambda_value * x[:,  :, index2]

        return mixed_x1, mixed_x2

models/Simstock.py
Yoontae6719/SimStock-Representation-Model-for-Stock-Similarities
import torch
import numpy as np

from torch.autograd import Variable
from torch.utils.data import DataLoader, Dataset


def make_noise(shape, type="Gaussian"):
    """
    Generate random noise.
    Parameters
    ----------
    shape: List or tuple indicating the shape of the noise
    type: str, "Gaussian" or "Uniform", default: "Gaussian".
    Returns
    -------
    noise tensor
    """

    if type == "Gaussian":
        noise = Variable(torch.randn(shape))
    elif type == 'Uniform':
        noise = Variable(torch.randn(shape).uniform_(-1, 1))
    else:
        raise Exception("ERROR: Noise type {} not supported".format(type))
    return noise

utils/helper.py
Yoontae6719/SimStock-Representation-Model-for-Stock-Similarities
import pandas as pd
import numpy as np
from torch.utils.data import DataLoader, Dataset

def preprocess_stock(df, is_daily = True):
    '''
    Preprocess stock dataframe by generating rolling metrics for Open, High, Low, Close, and Volume columns.
    
    Parameters:
    - df (pd.DataFrame): Input dataframe containing stock data with columns: Open, High, Low, Close, Volume.
    - is_daily (bool): Flag to determine if the stock data is daily. If set to False, it assumes hourly data.
    
    Returns:
    - df_pre (pd.DataFrame): Preprocessed dataframe with new rolling metrics columns added.
    '''
    
    if is_daily:
        rolling_size = [5, 10, 15, 20, 25]
    else:
        rolling_size = [6, 12, 18, 24, 30]
    
    df["Open06"] =  df['Open'].rolling(rolling_size[0]).sum() /   (rolling_size[0] * df['Open']) - 1
    df["Open12"] =  df['Open'].rolling(rolling_size[1]).sum() /  (rolling_size[1] * df['Open']) - 1
    df["Open18"] =  df['Open'].rolling(rolling_size[2]).sum() /  (rolling_size[2] * df['Open']) - 1
    df["Open24"] =  df['Open'].rolling(rolling_size[3]).sum() /  (rolling_size[3] * df['Open']) - 1
    df["Open30"] =  df['Open'].rolling(rolling_size[4]).sum() /  (rolling_size[4] * df['Open']) - 1

    df["High06"] =  df['High'].rolling(rolling_size[0]).sum() /   (rolling_size[0] * df['High']) - 1
    df["High12"] =  df['High'].rolling(rolling_size[1]).sum() /   (rolling_size[1] * df['High']) - 1
    df["High18"] =  df['High'].rolling(rolling_size[2]).sum() /   (rolling_size[2] * df['High']) - 1
    df["High24"] =  df['High'].rolling(rolling_size[3]).sum() /   (rolling_size[3] * df['High']) - 1
    df["High30"] =  df['High'].rolling(rolling_size[4]).sum() /   (rolling_size[4] * df['High']) - 1
    
    df["Low06"] =  df['Low'].rolling(rolling_size[0]).sum() /   (rolling_size[0] * df['Low']) - 1
    df["Low12"] =  df['Low'].rolling(rolling_size[1]).sum() /   (rolling_size[1] * df['Low']) - 1
    df["Low18"] =  df['Low'].rolling(rolling_size[2]).sum() /   (rolling_size[2] * df['Low']) - 1
    df["Low24"] =  df['Low'].rolling(rolling_size[3]).sum() /   (rolling_size[3] * df['Low']) - 1
    df["Low30"] =  df['Low'].rolling(rolling_size[4]).sum() /   (rolling_size[4] * df['Low']) - 1
    
    df["Close06"] =  df['Close'].rolling(rolling_size[0]).sum() /  (rolling_size[0] * df['Close']) - 1
    df["Close12"] =  df['Close'].rolling(rolling_size[1]).sum() /  (rolling_size[1] * df['Close']) - 1
    df["Close18"] =  df['Close'].rolling(rolling_size[2]).sum() /  (rolling_size[2] * df['Close']) - 1
    df["Close24"] =  df['Close'].rolling(rolling_size[3]).sum() /  (rolling_size[3] * df['Close']) - 1
    df["Close30"] =  df['Close'].rolling(rolling_size[4]).sum() /  (rolling_size[4] * df['Close']) - 1
    
    df["Volume06"] =  df['Volume'].rolling(rolling_size[0]).sum() /   (rolling_size[0] * df['Volume']) - 1
    df["Volume12"] =  df['Volume'].rolling(rolling_size[1]).sum() / (rolling_size[1] * df['Volume']) - 1
    df["Volume18"] =  df['Volume'].rolling(rolling_size[2]).sum() / (rolling_size[2] * df['Volume']) - 1
    df["Volume24"] =  df['Volume'].rolling(rolling_size[3]).sum() / (rolling_size[3] * df['Volume']) - 1
    df["Volume30"] =  df['Volume'].rolling(rolling_size[4]).sum() / (rolling_size[4] * df['Volume']) - 1

    df_pre = df.dropna()
    df_pre = df_pre.replace(np.inf, np.nan).dropna().reset_index(drop = True)
    df_pre = df_pre.replace(-np.inf, np.nan).dropna().reset_index(drop = True)
    
    return df_pre.reset_index(drop = True)




class DomainDataset(Dataset):
    """ Customized dataset for each domain"""
    def __init__(self,X, S):
        self.X = X                           # set data
        self.S = S                           # set Stock Sector

    def __len__(self):
        return len(self.X)                   # return length

    def __getitem__(self, idx):
        return [self.X[idx], self.S[idx]]    # return list of batch data [data, labels]
    
    

def dataset_for_modeling(args, train_type = False):
    data = pd.read_csv("./data/{}.csv".format(args.train_dataset))
    test = pd.read_csv("./data/{}.csv".format(args.test_dataset))
    dataloaders = []
    input_var = ['Open06', 'Open12', 'Open18', 'Open24', 'Open30',
           'High06', 'High12', 'High18', 'High24', 'High30', 'Low06', 'Low12',
           'Low18', 'Low24', 'Low30', 'Close06', 'Close12', 'Close18', 'Close24',
           'Close30', 'Volume06', 'Volume12', 'Volume18', 'Volume24', 'Volume30']
    sector_var = ["IndustryCode_"]

        
    if not train_type:
        index_range = [126, 126*2, 126*3, 126*4, 126*5, 126*6, 126*7, 126*8, 126*9, 126*10]
        data_0 = data[data["Date_index"] <= index_range[0]]
        data_1 = data[(data["Date_index"] > index_range[0]) & (data["Date_index"] <= index_range[1]) ]
        data_2 = data[(data["Date_index"] > index_range[1]) & (data["Date_index"] <= index_range[2]) ]
        data_3 = data[(data["Date_index"] > index_range[2]) & (data["Date_index"] <= index_range[3]) ]
        data_4 = data[(data["Date_index"] > index_range[3]) & (data["Date_index"] <= index_range[4]) ]
        data_5 = data[(data["Date_index"] > index_range[4]) & (data["Date_index"] <= index_range[5]) ]
        data_6 = data[(data["Date_index"] > index_range[5]) & (data["Date_index"] <= index_range[6]) ]
        data_7 = data[(data["Date_index"] > index_range[6]) & (data["Date_index"] <= index_range[7]) ]
        data_8 = data[(data["Date_index"] > index_range[7]) & (data["Date_index"] <= index_range[8]) ]
        data_9 = data[(data["Date_index"] > index_range[8]) & (data["Date_index"] <= index_range[9]) ]
        data_a = data[(data["Date_index"] > index_range[9])]


        data_0_in= data_0[input_var].values
        data_1_in= data_1[input_var].values
        data_2_in= data_2[input_var].values
        data_3_in= data_3[input_var].values
        data_4_in= data_4[input_var].values
        data_5_in= data_5[input_var].values
        data_6_in= data_6[input_var].values
        data_7_in= data_7[input_var].values
        data_8_in= data_8[input_var].values
        data_9_in= data_9[input_var].values
        data_a_in= data_a[input_var].values

        data_0_se= data_0[sector_var].values
        data_1_se= data_1[sector_var].values
        data_2_se= data_2[sector_var].values
        data_3_se= data_3[sector_var].values
        data_4_se= data_4[sector_var].values
        data_5_se= data_5[sector_var].values
        data_6_se= data_6[sector_var].values
        data_7_se= data_7[sector_var].values
        data_8_se= data_8[sector_var].values
        data_9_se= data_9[sector_var].values
        data_a_se= data_a[sector_var].values

        data_in = [data_0_in
                  ,data_1_in
                  ,data_2_in
                  ,data_3_in
                  ,data_4_in
                  ,data_5_in
                  ,data_6_in
                  ,data_7_in
                  ,data_8_in
                  ,data_9_in
                  ,data_a_in]

        data_se = [data_0_se
                  ,data_1_se
                  ,data_2_se
                  ,data_3_se
                  ,data_4_se
                  ,data_5_se
                  ,data_6_se
                  ,data_7_se
                  ,data_8_se
                  ,data_9_se
                  ,data_a_se]

        dataloaders=[]
        for i in range(len(data_in)-1):
            temp_X = data_in[i]
            temp_S = data_se[i]
                        
            domain_dataset = DomainDataset(temp_X, temp_S) # create dataset for each domain
            temp_dataloader = DataLoader(domain_dataset, batch_size= args.batch_size, shuffle=True, num_workers=0)
            dataloaders.append(temp_dataloader)
            
        return dataloaders
    else:
        
        data_test_in= test[input_var].values
        data_test_se= test[sector_var].values
        domain_data_testset = DomainDataset(data_test_in, data_test_se) 
        test_dataloader = DataLoader(domain_data_testset, batch_size= args.batch_size, shuffle=False, num_workers=0)
    
    return test_dataloader
utils/prepro.py
Yoontae6719/SimStock-Representation-Model-for-Stock-Similarities