<jupyter_start><jupyter_text>Load train set<jupyter_code>df_train_all = pd.read_csv(TRAIN_CSV)
df_train_all = df_train_all.set_index(['stock_id', 'time_id'])
print(df_train_all.shape)
#rows_for_stock_id_0 = df_train_all.query('stock_id == 0').shape[0]
#rows_for_stock_id_0
def load_book(ROOT, filename, stock_id):
    df_book_train_stock_X = pd.read_parquet(os.path.join(ROOT, f"{filename}/stock_id={stock_id}"))
    df_book_train_stock_X["stock_id"] = stock_id
    df_book_train_stock_X = df_book_train_stock_X.set_index(['stock_id', 'time_id'])
    return df_book_train_stock_X

df_book_train_stock_X = load_book(ROOT, 'book_train.parquet', 0)
def show_details(df):
    try:
        nbr_index_levels = len(df.index.levels)
    except AttributeError:
        nbr_index_levels = 1
    nbr_nulls = df.isnull().sum().sum()
    #nulls_msg = "Has no nulls"
    #if nbr_nulls==0:
    nulls_msg = f"{nbr_nulls} nulls"
    is_view_msg = f'is_view {df_train_all._data.is_view}'
    is_single_block_msg = f'is_single_block {df_train_all._data.is_single_block}'
    is_consolidated_msg = f'is_consolidated {df_train_all._data.is_consolidated()}'    
    print(f'[{nbr_index_levels}c] {df.shape[0]:,}x{df.shape[1]:,}, {nulls_msg}, {is_view_msg}, {is_single_block_msg}, {is_consolidated_msg}')

show_details(df_train_all)
# make feature columns
def make_features_stats(df_book, agg_type, cols):
    features_var1 = df_book.groupby(['stock_id', 'time_id'])[cols].agg(agg_type)
    #print(type(features_var1))
    if isinstance(features_var1, pd.Series):
        # .size yields a series not a df
        #features_var1.name = str(agg_type)
        features_var1 = pd.DataFrame(features_var1, columns=[agg_type])
        #pass
    else:
        features_var1_col_names = [f"{col}_{agg_type}" for col in cols]
        features_var1.columns = features_var1_col_names
    return features_var1

if True: # lightweight tests
    df_book_train_stock_XX = pd.read_parquet(os.path.join(ROOT, f"book_train.parquet/stock_id=0"))
    df_book_train_stock_XX["stock_id"] = 0
    df_book_train_stock_XX = df_book_train_stock_XX.set_index(['stock_id', 'time_id'])
    display(make_features_stats(df_book_train_stock_XX, 'nunique', ['ask_size1']).head())
%timeit make_features_stats(df_book_train_stock_XX, 'nunique', ['ask_size1'])
%lprun -f make_features_stats make_features_stats(df_book_train_stock_XX, 'nunique', ['ask_size1'])
def _realized_volatility_weighted_sub(ser, weights):
    ser_weighted = ser * weights
    return np.sqrt(np.sum(ser_weighted**2))

def realized_volatility_weighted(ser, weights_type):
    """Weighted volatility"""
    # as a numpy array
    # we drop from 12us to 3us by adding @njit to the _sub function
    # we can't make _sub a closure, it loses all compilation benefits
    # and we can't add njit(cache=True) in Jupyter as it can't
    # find a cache location    
    # as a Series we have 5us and 15us w/wo @njit respectively
    if isinstance(ser, pd.Series):
        ser = ser.to_numpy()
    nbr_items = ser.shape[0]
    if weights_type == 'uniform':
        weights = np.ones(nbr_items)
    elif weights_type == 'linear':
        weights = np.linspace(0.1, 1, nbr_items) # linear increasing weight
    elif weights_type == 'half0half1':
        half_way = int(ser.shape[0] / 2)
        weights = np.concatenate((np.zeros(half_way), np.ones(ser.shape[0] - half_way))) # 0s then 1s weight
    elif weights_type == 'geom':
        weights = np.geomspace(0.01, 1, nbr_items) # geometric increase
    #assert isinstance(weights_type, str) == False, f"Must not be a string like '{weights}' at this point"
    return _realized_volatility_weighted_sub(ser, weights)

#from joblib import Memory
#memory = Memory("joblib_cache", verbose=0)
<EMAIL>

from functools import lru_cache

@lru_cache
def get_weights(nbr_items, weights_type):
    if weights_type == 'uniform':
        weights = np.ones(nbr_items)
    elif weights_type == 'linear':
        weights = np.linspace(0.1, 1, nbr_items) # linear increasing weight
    elif weights_type == 'half0half1':
        half_way = int(ser.shape[0] / 2)
        weights = np.concatenate((np.zeros(half_way), np.ones(ser.shape[0] - half_way))) # 0s then 1s weight
    elif weights_type == 'geom':
        weights = np.geomspace(0.01, 1, nbr_items) # geometric increase
    return weights

def realized_volatility_weighted_v2(ser, weights_type):
    """Weighted volatility"""
    # as a numpy array
    # we drop from 12us to 3us by adding @njit to the _sub function
    # we can't make _sub a closure, it loses all compilation benefits
    # and we can't add njit(cache=True) in Jupyter as it can't
    # find a cache location    
    # as a Series we have 5us and 15us w/wo @njit respectively
    if isinstance(ser, pd.Series):
        ser = ser.to_numpy()
    nbr_items = ser.shape[0]
    weights = get_weights(nbr_items, weights_type)
    #assert isinstance(weights_type, str) == False, f"Must not be a string like '{weights}' at this point"
    return _realized_volatility_weighted_sub(ser, weights)

series_log_return = pd.Series(np.linspace(0, 10, 600))
print(realized_volatility_weighted(series_log_return, weights_type="uniform"))
assert realized_volatility_weighted(series_log_return, weights_type="uniform") == realized_volatility_weighted_v2(series_log_return, weights_type="uniform")
    #%timeit realized_volatility_weighted(series_log_return, weights_type="uniform")
%timeit realized_volatility_weighted(series_log_return, weights_type="uniform")
%timeit realized_volatility_weighted_v2(series_log_return, weights_type="uniform")
%lprun -f realized_volatility_weighted -f _realized_volatility_weighted_sub realized_volatility_weighted(series_log_return, weights_type="uniform")
def log_return(list_stock_prices):
    return np.log(list_stock_prices).diff()

#def realized_volatility(series_log_return):
#    return np.sqrt(np.sum(series_log_return**2))

def make_wap(df_book_data, num=1, wap_colname="wap"):
    """Modifies df_book_data"""
    assert num==1 or num==2
    wap_numerator = (df_book_data[f'bid_price{num}'] * df_book_data[f'ask_size{num}'] +
                                     df_book_data[f'ask_price{num}'] * df_book_data[f'bid_size{num}'])
    wap_denominator = df_book_data[f'bid_size{num}'] + df_book_data[f'ask_size{num}']
    df_book_data[wap_colname] = wap_numerator / wap_denominator

<EMAIL>
def make_realized_volatility(df_book_data, log_return_name='log_return', wap_colname='wap', weights=None):
    """Consume wap column"""
    df_book_data[log_return_name] = df_book_data.groupby(['stock_id', 'time_id'])[wap_colname].apply(log_return)
    df_book_data = df_book_data[~df_book_data[log_return_name].isnull()]
    df_realized_vol_per_stock =  pd.DataFrame(df_book_data.groupby(['stock_id', 'time_id'])[log_return_name].agg(realized_volatility_weighted, weights))
    return df_realized_vol_per_stock

make_wap(df_book_train_stock_X, 2) # adds 'wap' column
df_realized_vol_per_stockX = make_realized_volatility(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
display(df_realized_vol_per_stockX.head(2))
def make_volatility_ask_bid_diff(df_book_data, col='ask1_bid1_diff', weights='uniform'):
    new_name = col + "_log_return_" + weights
    df_book_data[new_name] = df_book_data.groupby(['stock_id', 'time_id'])[col].apply(log_return)
    df_book_data = df_book_data[~df_book_data[new_name].isnull()]
    # makes a new dataframe
    df_realized_vol_ask_bid_diff_per_stock =  pd.DataFrame(df_book_data.groupby(['stock_id', 'time_id'])[new_name].agg(realized_volatility_weighted, weights))
    return df_realized_vol_ask_bid_diff_per_stock

df_book_train_stock_X['ask1_bid1_diff'] = (df_book_train_stock_X['ask_price1'] / df_book_train_stock_X['bid_price1']) - 1
df_realized_vol_ask_bid_diff_per_stock = make_volatility_ask_bid_diff(df_book_train_stock_X, weights='linear')
display(df_realized_vol_ask_bid_diff_per_stock.head(2))
%lprof -f make_volatility_ask_bid_diff make_volatility_ask_bid_diff(df_book_train_stock_X, weights='linear')
def make_wap(df_book_data, num=1, wap_colname="wap"):
    """Modifies df_book_data"""
    assert num==1 or num==2
    wap_numerator = (df_book_data[f'bid_price{num}'] * df_book_data[f'ask_size{num}'] +
                                     df_book_data[f'ask_price{num}'] * df_book_data[f'bid_size{num}'])
    wap_denominator = df_book_data[f'bid_size{num}'] + df_book_data[f'ask_size{num}']
    df_book_data[wap_colname] = wap_numerator / wap_denominator

def make_realized_volatility(df_book_data, log_return_name='log_return', wap_colname='wap', weights=None):
    """Consume wap column"""
    df_book_data[log_return_name] = df_book_data.groupby(['stock_id', 'time_id'])[wap_colname].apply(log_return)
    df_book_data = df_book_data[~df_book_data[log_return_name].isnull()]
    df_realized_vol_per_stock =  pd.DataFrame(df_book_data.groupby(['stock_id', 'time_id'])[log_return_name].agg(realized_volatility_weighted, weights))
    return df_realized_vol_per_stock

make_wap(df_book_train_stock_X, 2) # adds 'wap' column
df_realized_vol_per_stockX = make_realized_volatility(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
display(df_realized_vol_per_stockX.head(2))
%lprun -f make_realized_volatility -f realized_volatility_weighted -f _realized_volatility_weighted_sub -f log_return make_realized_volatility(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
ser = df_book_train_stock_X.query('stock_id ==0 and time_id == 5')['wap']
log_return(ser)
from numba import njit

@njit(parallel=True)
def log_return_v2_inner(arr):
    arr = np.log(arr)
    arr2 = np.concatenate((np.array([np.NaN]), arr[1:] - arr[:-1]))
    return arr2

def log_return_v2(list_stock_prices_ser):
    # assume list_stock_prices is an array, not Series
    #print(type(list_stock_prices))
    list_stock_prices = list_stock_prices_ser.to_numpy()
    arr2 = log_return_v2_inner(list_stock_prices)
    return pd.Series(arr2, index=list_stock_prices_ser.index)


def make_realized_volatility_v2(df_book_data, log_return_name='log_return', wap_colname='wap', weights=None):
    """Consume wap column"""
    logged_diff = df_book_data.groupby(['stock_id', 'time_id'])[wap_colname].apply(log_return_v2)
    #print(logged_diff[:5])
    #df_book_data = df_book_data[~df_book_data[log_return_name].isnull()]
    logged_diff = logged_diff.dropna()
    df_realized_vol_per_stock =  pd.DataFrame(logged_diff.groupby(['stock_id', 'time_id']).agg(realized_volatility_weighted, weights))
    return df_realized_vol_per_stock

make_wap(df_book_train_stock_X, 1) # adds 'wap' column

df_realized_vol_per_stockX_orig = make_realized_volatility_v2(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
df_realized_vol_per_stockX = make_realized_volatility_v2(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
display(df_realized_vol_per_stockX[:5])
pd.testing.assert_frame_equal(df_realized_vol_per_stockX_orig, df_realized_vol_per_stockX)
#log_return_v2(ser.to_numpy()) # this looks like the original, sans the index
log_return_v2_inner.parallel_diagnostics(level=1)
%timeit make_realized_volatility_v2(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
%timeit make_realized_volatility(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
%lprun -f make_realized_volatility_v2 -f realized_volatility_weighted -f _realized_volatility_weighted_sub -f log_return_v2 make_realized_volatility_v2(df_book_train_stock_X, log_return_name="log_return2", weights='linear')
<EMAIL>
def load_data_build_features(stock_id, ROOT, book_filename, trade_filename, cols, df_target):
    # filename e.g. book_train.parquet
    assert isinstance(stock_id, int)
    df_book_stock_X = pd.read_parquet(
        os.path.join(ROOT, f"{book_filename}/stock_id={stock_id}")
    )
    df_book_stock_X["stock_id"] = stock_id
    df_book_stock_X = df_book_stock_X.set_index(['stock_id', 'time_id'])
    #assert df_book_train_stock_X.shape[0] > rows_for_stock_id_0, (df_book_train_stock_X.shape[0], rows_for_stock_id_0)
    
    df_trade_stock_X = pd.read_parquet(
        os.path.join(ROOT, f"{trade_filename}/stock_id={stock_id}")
    )
    df_trade_stock_X["stock_id"] = stock_id
    df_trade_stock_X = df_trade_stock_X.set_index(['stock_id', 'time_id'])
    
    #df_book_train_stock_X_gt500 = df_book_train_stock_X.query("seconds_in_bucket>500").copy()
    #df_realized_vol_per_stock_short500 = add_wap_make_realized_volatility(df_book_train_stock_X_gt500, log_return_name='log_return_gt500sec')
    #df_book_train_stock_X_gt300 = df_book_train_stock_X.query("seconds_in_bucket>300").copy()
    #df_realized_vol_per_stock_short300 = add_wap_make_realized_volatility(df_book_train_stock_X_gt300, log_return_name='log_return_gt300sec')
    if True:
        make_wap(df_book_stock_X, 2, "wap2") 
        df_realized_vol_per_stock_wap2_uniform = make_realized_volatility(df_book_stock_X, log_return_name="log_return2_uniform", wap_colname="wap2", weights='uniform')    
        df_realized_vol_per_stock_wap2_linear = make_realized_volatility(df_book_stock_X, log_return_name="log_return2_linear", wap_colname="wap2", weights='linear')
        df_realized_vol_per_stock_wap2_half0half1 = make_realized_volatility(df_book_stock_X, log_return_name="log_return2_half0half1", wap_colname="wap2", weights='half0half1')
        make_wap(df_book_stock_X, 1, "wap") # adds 'wap' column
        df_realized_vol_per_stock_wap1_uniform = make_realized_volatility(df_book_stock_X, log_return_name="log_return1_uniform", weights='uniform')
        df_realized_vol_per_stock_wap1_linear = make_realized_volatility(df_book_stock_X, log_return_name="log_return1_linear", weights='linear')
        df_realized_vol_per_stock_wap1_half0half1 = make_realized_volatility(df_book_stock_X, log_return_name="log_return1_half0half1", weights='half0half1')

        features_var1 = make_features_stats(df_book_stock_X, 'var', cols)
        features_mean1 = make_features_stats(df_book_stock_X, 'mean', cols)
        features_size1 = make_features_stats(df_book_stock_X, 'size', cols)
        features_min1 = make_features_stats(df_book_stock_X, 'min', cols)
        features_max1 = make_features_stats(df_book_stock_X, 'max', cols)
        features_nunique1 = make_features_stats(df_book_stock_X, 'nunique', cols)

        df_book_stock_X['ask1_bid1_diff'] = (df_book_stock_X['ask_price1'] / df_book_stock_X['bid_price1']) - 1
        df_realized_vol_ask1_bid1_diff_per_stock = make_volatility_ask_bid_diff(df_book_stock_X, col='ask1_bid1_diff')
        df_realized_vol_ask1_bid1_diff_per_stock_linear = make_volatility_ask_bid_diff(df_book_stock_X, col='ask1_bid1_diff', weights='linear')
        df_realized_vol_ask1_bid1_diff_per_stock_linear = make_volatility_ask_bid_diff(df_book_stock_X, col='ask1_bid1_diff', weights='half0half1')
        df_book_stock_X['ask2_bid2_diff'] = (df_book_stock_X['ask_price2'] / df_book_stock_X['bid_price2']) - 1
        df_realized_vol_ask2_bid2_diff_per_stock = make_volatility_ask_bid_diff(df_book_stock_X, col='ask2_bid2_diff')
        df_realized_vol_ask2_bid2_diff_per_stock = make_volatility_ask_bid_diff(df_book_stock_X, col='ask2_bid2_diff', weights='linear')
        df_realized_vol_ask2_bid2_diff_per_stock = make_volatility_ask_bid_diff(df_book_stock_X, col='ask2_bid2_diff', weights='half0half1')
        df_book_stock_X['ask1_bid2_diff'] = (df_book_stock_X['ask_price1'] / df_book_stock_X['bid_price2']) - 1
        df_realized_vol_ask1_bid2_diff_per_stock = make_volatility_ask_bid_diff(df_book_stock_X, col='ask1_bid2_diff')
        df_book_stock_X['ask2_bid1_diff'] = (df_book_stock_X['ask_price2'] / df_book_stock_X['bid_price1']) - 1
        df_realized_vol_ask2_bid1_diff_per_stock = make_volatility_ask_bid_diff(df_book_stock_X, col='ask2_bid1_diff')
    else:
        features_var1 = make_features_stats(df_book_stock_X, 'var', cols)
        
    #breakpoint()
    # trade stats
    df_trade_basic_stats = df_trade_stock_X.groupby(['stock_id', 'time_id']).agg(trade_size_count=pd.NamedAgg('size', 'count'), trade_order_count_sum=pd.NamedAgg('order_count', 'sum'))

    df_train_stock_X = df_target.query('stock_id == @stock_id')
    if True:
        to_merge_book = [df_train_stock_X, 
                    features_var1, features_mean1, features_size1, 
                    features_min1, features_max1, features_nunique1,
                    df_realized_vol_per_stock_wap1_uniform,
                    df_realized_vol_ask1_bid1_diff_per_stock_linear,
                    df_realized_vol_per_stock_wap2_uniform,
                    df_realized_vol_per_stock_wap1_linear,
                    df_realized_vol_per_stock_wap2_linear,
                    df_realized_vol_per_stock_wap1_half0half1,
                    df_realized_vol_per_stock_wap2_half0half1,
                    df_realized_vol_ask1_bid1_diff_per_stock,
                    df_realized_vol_ask2_bid2_diff_per_stock,
                    df_realized_vol_ask1_bid2_diff_per_stock,
                    df_realized_vol_ask2_bid1_diff_per_stock,]
        to_merge_trade = [df_trade_basic_stats]
    else:
        to_merge_book = [df_train_stock_X, features_var1]
        to_merge_trade = [df_trade_basic_stats]

    # some trade datasets are missing some time_ids, making the join a mess
    # we reindex and make the choice to fillna 0
    to_merge_trade = [to_merge_tr.reindex(to_merge_book[0].index, fill_value=0) for to_merge_tr in to_merge_trade]
    #to_merge_trade = [to_merge_tr.reindex(to_merge_book[0].index, fill_value=0).fillna(0) for to_merge_tr in to_merge_trade]
    
    row_lengths = [df.shape[0] for df in to_merge_book]
    assert len(set(row_lengths)) == 1, f"row_lengths are different for stock {stock_id}: {row_lengths}" # should all be same length
    to_merge = to_merge_book + to_merge_trade
    for idx, item_to_merge in enumerate(to_merge):
        assert item_to_merge.index.names == ['stock_id', 'time_id'], f"We must have the same index on idx {idx}"
    train_merged = pd.concat(to_merge, axis=1)
                             
    if 'target' in train_merged.columns:
        features = train_merged.drop(columns='target').columns
        #print(features)
        assert len(set(features)) == len(features), f"Feature duplication! {len(set(features))} vs {len(features)}"

    return train_merged

#if 'memory' in dir():
#    # only setup local cache if we're running locally in development
#    load_data_build_features = memory.cache(load_data_build_features)
    
cols = ['bid_price1', 'ask_price1', 'bid_price2', 'ask_price2',] 
cols += ['bid_size1', 'ask_size1', 'bid_size2', 'ask_size2']

# test...
train_mergedXX = load_data_build_features(31, ROOT, 'book_train.parquet', 'trade_train.parquet', cols, df_train_all)
display(train_mergedXX.head(5))

%lprun -f load_data_build_features load_data_build_features(31, ROOT, 'book_train.parquet', 'trade_train.parquet', cols, df_train_all)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10604787904241915, 'fraction_numerical': 0.01601217975640487, 'mean_word_length': 5.62113312478276, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '262938', 'license': 'MIT', 'n_tokens_mistral': 7301, 'n_tokens_neox': 7081, 'n_words': 1029, 'path': 'notebooks/20210825_light_profiling.ipynb', 'repo_name': 'ianozsvald/kaggle_optiver_volatility', 'size': '19069'}