Time Series Search

!pip install -qU convertdate kaggle matplotlib==3.1.3 
!pip install -q git+https://github.com/facebookresearch/Kats.git
!pip install -qU pinecone-client

import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import pprint
from sklearn.preprocessing import MinMaxScaler
from kats.consts import TimeSeriesData
from kats.tsfeatures.tsfeatures import TsFeatures
import itertools
from decimal import Decimal
from IPython.display import clear_output

import warnings
warnings.simplefilter(action='ignore')

import pinecone

# Load Pinecone API key
api_key = os.getenv('PINECONE_API_KEY') or 'YOUR_API_KEY'
# Set Pinecone environment. Default environment is YOUR_ENVIRONMENT
env = os.getenv('PINECONE_ENVIRONMENT') or 'YOUR_ENVIRONMENT'
pinecone.init(api_key=api_key, environment=env)

# Pick a name for the new index
simple_index_name = 'stocks-trends'

# Check whether the index with the same name already exists
if simple_index_name in pinecone.list_indexes():
    pinecone.delete_index(simple_index_name)

# Create a new index
pinecone.create_index(name=simple_index_name, dimension=128)

# Establish a connection
simple_index = pinecone.Index(index_name=simple_index_name)

%%writefile kaggle.json
{"username":"KAGGLE_USERNAME","key":"KAGGLE_KEY"}

#Check Kaggle username and key
! cat ./kaggle.json

!mkdir ~/.kaggle
!cp kaggle.json ~/.kaggle/
!chmod 600 ~/.kaggle/kaggle.json
!kaggle datasets download -d jacksoncrow/stock-market-dataset
!unzip -q stock-market-dataset.zip -d data

# Define sliding window and step size
SLIDING_WINDOW_SIZE = 64
STEP = 10

def windows(data, window_size, step):
    r = np.arange(len(data))
    s = r[::step]
    z = list(zip(s, s + window_size))
    f = '{0[0]}:{0[1]}'.format
    g = lambda t: data.iloc[t[0]:t[1]]
    return pd.concat(map(g, z), keys=map(f, z))

def get_feature_embedding_for_window(df, stock):
    ts_name = f"{stock.strip('.csv')}_{str(df.Date.min())}_{str(df.Date.max())}"
    scaler=MinMaxScaler()
    df[['Open', 'Close']] = scaler.fit_transform(df[['Open', 'Close']])
    prices = df[['Open', 'Close']].values.tolist()
    flat_values = [item for sublist in prices for item in sublist]
    df = df.rename(columns={"Date":"time"}) 
    ts_df = pd.DataFrame({'time':df.time.repeat(2), 
                          'price':flat_values})
    ts_df.drop_duplicates(keep='first', inplace=True)  

    # Use Kats to extract features for the time window
    try:
        if not (len(np.unique(ts_df.price.tolist())) == 1 \
            or len(np.unique(ts_df.price.tolist())) == 0):
            timeseries = TimeSeriesData(ts_df)
            features = TsFeatures().transform(timeseries)
            feature_list = [float(v) if not pd.isnull(v) else float(0) for _, v in features.items()]
            if Decimal('Infinity') in feature_list or Decimal('-Infinity') in feature_list:
                return None
            return (ts_name, feature_list)
    except np.linalg.LinAlgError as e:
        print(f"Can't process {ts_name}:{e}")
    return None

def get_simple_pair_for_window(df, stock):
    ts_name = f"{stock.strip('.csv')}_{str(df.Date.min())}_{str(df.Date.max())}"
    prices = df[['Open', 'Close']].values.tolist()
    flat_values = [item for sublist in prices for item in sublist]
    return (ts_name, flat_values)

def chunks(iterable, batch_size=100):
    it = iter(iterable)
    chunk = tuple(itertools.islice(it, batch_size))
    while chunk:
        yield chunk
        chunk = tuple(itertools.islice(it, batch_size))

def upload_data_to_index(index, create_pair_func, verbose=False):
    # Define path to the folder
    stocks = sorted(os.listdir('./data/stocks'))
    
    # Iterate over files, create vectors and upload data
    for stock in stocks[::50]:
        print(stock.strip('.csv'))
        data = pd.read_csv(os.path.join('./data/stocks', stock))
        data = data.sort_index(axis=0, ascending=True)
        data["Date"] = pd.to_datetime(data["Date"]).dt.date

        # Interpolate data for missing dates
        data.set_index('Date', inplace=True)
        data = data.reindex(pd.date_range(start=data.index.min(),
                                          end=data.index.max(),
                                          freq='1D'))
        data = data.interpolate(method='linear')
        data = data.reset_index().rename(columns={'index': 'Date'})
        data["Date"] = pd.to_datetime(data["Date"]).dt.date
        
        # Create sliding windows dataset
        wdf = windows(data, SLIDING_WINDOW_SIZE, STEP)
        
        # Prepare sequences for upload 
        items_to_upload = []
        for window, new_df in wdf.groupby(level=0):
            if new_df.shape[0] == SLIDING_WINDOW_SIZE:
                pair = create_pair_func(new_df, stock)
                if pair:
                    items_to_upload.append(pair)

        # Upload data for the symbol
        for batch in chunks(items_to_upload, 500):
            index.upsert(vectors=batch)

upload_data_to_index(simple_index, get_simple_pair_for_window)
clear_output()

# Check the index size
simple_index.describe_index_stats()

def prepare_items_for_graph(data):    
    scaler = MinMaxScaler()
    result_list = []
    
    for _, row in data.iterrows():
        id = row['id']
        vec = row['values']
        scaled_vec = scaler.fit_transform(np.array(vec).reshape(-1,1))
        result_list.append((id, (vec, scaled_vec)))
    return result_list

def show_query_results(query_item, data):
    data_prepared = prepare_items_for_graph(data)
    graph_index = pd.Float64Index(np.arange(start=0, stop=SLIDING_WINDOW_SIZE, step=0.5))

    print('\n The most similar items from the vector index:')
    data.reset_index(inplace=True, drop=True)
    display(data)
      
    fig = plt.figure(figsize=(20,7))
    for item in data_prepared:
        _id, vectors = item
        ax1 = plt.subplot(1, 2, 1)
        graph = plt.plot(graph_index, vectors[0], label = _id, marker='o' if _id == query_item else None)
        ax2 = plt.subplot(1, 2, 2)
        graph = plt.plot(graph_index, vectors[1], label = _id, marker='o' if _id == query_item else None)    
    ax1.set_xlabel("Days in time window")
    ax2.set_xlabel("Days in time window")
    ax1.set_ylabel("Stock values")
    ax2.set_ylabel("Normalized Stock Values")
    ax1.title.set_text(f'Similar stock patterns and their market values')
    ax2.title.set_text(f'Similar stock patterns and their normalized market values')
    plt.legend(loc='center left', bbox_to_anchor=(1.0, 0.5))
    plt.show()

def filter_results(query_item, data, historical_only=False):
    already_present = []
    
    # Remove symbol that is already included
    for i, row in data.iterrows():
        check_name = row.id.split('_')[0]
        if check_name not in already_present:
            already_present.append(check_name)
        else:
            data.drop(i,axis=0,inplace=True)
            
    # Include only data prior to query interval
    if historical_only:
        _, start_dt, end_dt = query_item.split('_')
        start_dt = pd.to_datetime(start_dt).date()
        data['final_date'] = data.id.apply(lambda x: x.split('_')[2])
        data['final_date'] =  data.final_date.apply(lambda x: pd.to_datetime(x).date())
        data = data[data.final_date <= start_dt]
        del data['final_date']
       
    return data

# Define query examples
items_to_query = ['BORR_2019-10-18_2019-12-20', 'HCCO_2020-01-28_2020-03-31', 'PUMP_2019-11-22_2020-01-24']

# Fetch vectors from the index
fetch_res = simple_index.fetch(ids=items_to_query)

# Create a list of ids and vectors for the fetched items
query_ids = [res.id for res in fetch_res.vectors.values()]
query_vectors = [res.values for res in fetch_res.vectors.values()]

# Query the pinecone index
query_results = []
for xq in query_vectors:
    res = simple_index.query(xq, top_k=100, include_values=True)
    query_results.append(res)

# Iterate and show query results for each query item
for query_item, q_res in zip(query_ids, query_results):
    print(f'\nQueried: {query_item}')
    res_df = pd.DataFrame(
        {
            'id': [res.id for res in q_res.matches], 
            'score': [res.score for res in q_res.matches],
            'values': [res.values for res in q_res.matches]
         }
    )
    res_df = filter_results(query_item, res_df, historical_only=False)
    show_query_results(query_item, res_df.head(6))

# Pick a name for the new index
kats_index_name = 'stocks-trends-with-features'

# Check whether the index with the same name already exists
if kats_index_name in pinecone.list_indexes():
    pinecone.delete_index(kats_index_name)

# Create a new index
pinecone.create_index(name=kats_index_name, dimension=40)

# Establish a connection
kats_index = pinecone.Index(index_name=kats_index_name)

upload_data_to_index(kats_index, get_feature_embedding_for_window)
clear_output()

kats_index.describe_index_stats()

# Fetch vectors from the index
fetch_res = kats_index.fetch(ids=items_to_query)

# Create a list of ids and vectors for the fetched items
query_ids = [res.id for res in fetch_res.vectors.values()]
query_vectors = [res.values for res in fetch_res.vectors.values()]

# Query the pinecone index
query_results = []
for xq in query_vectors:
    res = kats_index.query(xq, top_k=100, include_values=True)
    query_results.append(res)

# Iterate and show query results for each query item
for query_item, q_res in zip(query_ids, query_results):
    print(f'\nQueried: {query_item}')
    res_df = pd.DataFrame(
        {
            'id': [res.id for res in q_res.matches], 
            'score': [res.score for res in q_res.matches]
        }
    )

    # Use simple index to retrieve historical prices for query result items
    res_df['values'] = res_df.id.apply(lambda x: [res.values for res in simple_index.fetch(ids=[x]).vectors.values()][0])
    res_df = filter_results(query_item, res_df, historical_only=False)
    show_query_results(query_item, res_df.head(6))

pinecone.delete_index(simple_index_name)
pinecone.delete_index(kats_index_name)