spurious_reconstruction/analysis/1_case_study/feature-decoding/RR_sklearn.py

#!/usr/bin/env python
# coding: utf-8

# In[1]:


# Import packages and setup gpu configuration.
# This code block shouldnt need to be adjusted!
import os
import sys
import json
import yaml
import numpy as np
import copy
import math
import time
import random
from tqdm.auto import tqdm
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
from torchvision import transforms
import h5py
import utils
import pandas as pd
# here we import ridge regression from sklearn

from sklearn.linear_model import Ridge
import argparse

# tf32 data type is faster than standard float32
torch.backends.cuda.matmul.allow_tf32 = True
# following fixes a Conv3D CUDNN_NOT_SUPPORTED error
torch.backends.cudnn.benchmark = True



# outdir = os.path.abspath(f'checkpoints/{model_name}')


all_features = ['features[0]', 'features[2]', 'features[5]', 'features[7]', 'features[10]', 'features[12]', 'features[14]', 'features[16]', 'features[19]', 'features[21]', 'features[23]', 'features[25]', 'features[28]', 'features[30]', 'features[32]', 'features[34]', 'classifier[0]', 'classifier[3]', 'classifier[6]']

parser = argparse.ArgumentParser(description='Decoding features from a model')

parser.add_argument('--run_name', type=str, default='subj1_40_test', help='Name of the run')
parser.add_argument('--current_features', type=str, default='features[28]', help='Feature layer to decode')
parser.add_argument('--num_sessions', type=float, default=20, help='Number of sessions to use')
parser.add_argument('--subj', type=int, default=1, help='Subject number', choices=[1,2,5,7])

if utils.is_interactive():
    current_features = 'features[28]'
    num_sessions = 20
    subj = 2
    run_name = 'subj1_40_test'


else:
    args = parser.parse_args()
    for attribute_name in vars(args).keys():
        globals()[attribute_name] = getattr(args, attribute_name)

print(f"Configured run_name = {run_name}")
print(f"Configured current_features = {current_features}")
print(f"Configured num_sessions = {num_sessions}")
print(f"Configured subj = {subj}")


if utils.is_interactive():
    # Following allows you to change functions in models.py or utils.py and 
    # have this notebook automatically update with your revisions
    get_ipython().run_line_magic('load_ext', 'autoreload')
    get_ipython().run_line_magic('autoreload', '2')


save_ckpt = False


print("PID of this process =",os.getpid())
seed = 42
utils.seed_everything(seed)
data_type = torch.float32
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# this is the number of ridge regression splits to perform bcz of memory constraints
data_path = '/weka/proj-medarc/shared/mindeyev2_dataset/'



outdir = os.path.abspath(f'./decoded_features/{run_name}')
os.makedirs(outdir, exist_ok=True)


# In[2]:


# load the feature_encoder

from bdpy.dl.torch.models import VGG19, layer_map, model_factory
from bdpy.recon.torch.modules import build_encoder, build_generator, TargetNormalizedMSE
from bdpy.dl.torch.domain import Domain, image_domain, ComposedDomain



feature_network = VGG19()

feature_network.load_state_dict(torch.load('/weka/proj-fmri/ckadirt/spurious_reconstruction/analysis/VGG_ILSVRC_19_layers/VGG_ILSVRC_19_layers.pt'))
encoder = feature_network.to(device)

# encoder.eval()
# for param in encoder.parameters():
#     param.requires_grad = True

# # Define feature extractor
# class EncoderFeatureExtractor(nn.Module):
#     def __init__(self, encoder, target_layers):
#         super(EncoderFeatureExtractor, self).__init__()
#         self.encoder = encoder
#         self.target_layers = target_layers
#         self.features_layers = list(self.encoder.features.children())
#         self.classifier_layers = list(self.encoder.classifier.children())

#     def forward(self, x):
#         outputs = {}
        
#         for idx, layer in enumerate(self.features_layers):
#             x = layer(x)
#             layer_name = f'features[{idx}]'
#             if layer_name in self.target_layers:
#                 outputs[layer_name] = x

#         if 'avgpool' in self.target_layers:
#             x = self.encoder.avgpool(x)
#             outputs['avgpool'] = x
#         else:
#             x = self.encoder.avgpool(x)

#         x = torch.flatten(x, 1)

#         for idx, layer in enumerate(self.classifier_layers):
#             x = layer(x)
#             layer_name = f'classifier[{idx}]'
#             if layer_name in self.target_layers:
#                 outputs[layer_name] = x

#         return outputs

# if features == 'all':
#     feature_extractor = EncoderFeatureExtractor(encoder, target_layers=all_features)
# else:
#     feature_extractor = EncoderFeatureExtractor(encoder, target_layers=features)

if current_features == 'all':
    layer_names = all_features
else:
    layer_names = [current_features]


encoder = build_encoder(feature_network, layer_names,
                        domain= ComposedDomain([image_domain.BdPyVGGDomain(device=device,dtype=data_type), 
                                                image_domain.FixedResolutionDomain((224, 224))]),
                        )


# In[3]:


print("loading_betas")
betas = utils.create_snr_betas(subject=subj, data_type=torch.float16, data_path=data_path, threshold=-1.0)
print("betas_ loaded")

x_train, valid_nsd_ids_train, x_test, test_nsd_ids = utils.load_nsd(subject=subj, betas=betas, data_path=data_path)


# In[4]:


stim_descriptions = pd.read_csv(
        os.path.join(data_path, "nsd_stim_info_merged.csv"), index_col=0
    )
stim_descriptions.head()

rep_columns = [f"subject{subj}_rep{j}" for j in range(3)]

indexes_shared_1000 = torch.Tensor(stim_descriptions[
    (stim_descriptions[f'subject{subj}'] == 1) & (stim_descriptions['shared1000'] == 1)
][rep_columns].values.flatten()) - 1

nsd_ids = stim_descriptions[
    (stim_descriptions[f'subject{subj}'] == 1)
][rep_columns + ['nsdId']].values

valid_nsd_ids_full = torch.zeros(len(betas), dtype=torch.long)

for i, nsd_id in enumerate(nsd_ids):
    rep1, rep2, rep3, current_nsd_id = nsd_id
    valid_nsd_ids_full[rep1-1] = current_nsd_id
    valid_nsd_ids_full[rep2-1] = current_nsd_id
    valid_nsd_ids_full[rep3-1] = current_nsd_id

# check how many zeros are in valid_nsd_ids_full
print("Number of zeros in valid_nsd_ids_full", torch.sum(valid_nsd_ids_full == 0))


session_size = 750
num_examples_train = math.ceil(num_sessions * session_size)

x_train_subset = betas[:num_examples_train]
valid_nsd_ids_train_subset = valid_nsd_ids_full[:num_examples_train]

# filter the subset removing the indexes from the indexes_shared_1000
indexes_shared_1000_clap = indexes_shared_1000[indexes_shared_1000 < num_examples_train].to(torch.long)
mask = torch.ones(x_train_subset.size(0), dtype=torch.bool)
mask[indexes_shared_1000_clap] = False
# remove the examples from x_train_subset which are in indexes_shared_1000_clap
x_train_subset = x_train_subset[mask]
valid_nsd_ids_train_subset = valid_nsd_ids_train_subset[mask]


# In[5]:


x_train = x_train_subset
valid_nsd_ids_train = valid_nsd_ids_train_subset


# In[15]:


print('Num train examples', x_train.shape)


# In[6]:


f_images = h5py.File(f'{data_path}/coco_images_224_float16.hdf5', 'r')
images = f_images['images']


images = torch.Tensor(images[:])
print("Loaded all 73k possible NSD images to cpu!", images.shape)


# In[7]:


from torch.utils.data import Dataset, DataLoader

class RRDataset(Dataset):
    def __init__(self, x, valid_nsd_ids):
        self.x = x
        self.valid_nsd_ids = valid_nsd_ids

    def __len__(self):
        return len(self.x)

    def __getitem__(self, idx):
        betas = self.x[idx]
        nsd_id = self.valid_nsd_ids[idx]
        c_image = images[nsd_id]
        return betas, c_image, nsd_id
    
batch_size = 128

train_dataset = RRDataset(x_train, valid_nsd_ids_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

test_dataset = RRDataset(x_test, test_nsd_ids)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)


# In[8]:


# Define for the RR subdivision
splits_per_layer = {
    'features[0]': 32,
    'features[2]': 32,
    'features[5]': 16,
    'features[7]': 16,
    'features[10]': 8,
    'features[12]': 8,
    'features[14]': 8,
    'features[16]': 8,
    'features[19]': 4,
    'features[21]': 4,
    'features[23]': 4,
    'features[25]': 4,
    'features[28]': 2,
    'features[30]': 2,
    'features[32]': 2,
    'features[34]': 2,
    'classifier[0]': 1,
    'classifier[3]': 1,
    'classifier[6]': 1,
}

layer_sizes = {
    'features[0]': 64,
    'features[2]': 64,
    'features[5]': 128,
    'features[7]': 128,
    'features[10]': 256,
    'features[12]': 256,
    'features[14]': 256,
    'features[16]': 256,
    'features[19]': 512,
    'features[21]': 512,
    'features[23]': 512,
    'features[25]': 512,
    'features[28]': 512,
    'features[30]': 512,
    'features[32]': 512,
    'features[34]': 512,
    'classifier[0]': 4096,
    'classifier[3]': 4096,
    'classifier[6]': 1000,
}

features_shapes = {
    'features[0]': (64, 224, 224),
    'features[2]': (64, 224, 224),
    'features[5]': (128, 112, 112),
    'features[7]': (128, 112, 112),
    'features[10]': (256, 56, 56),
    'features[12]': (256, 56, 56),
    'features[14]': (256, 56, 56),
    'features[16]': (256, 56, 56),
    'features[19]': (512, 28, 28),
    'features[21]': (512, 28, 28),
    'features[23]': (512, 28, 28),
    'features[25]': (512, 28, 28),
    'features[28]': (512, 14, 14),
    'features[30]': (512, 14, 14),
    'features[32]': (512, 14, 14),
    'features[34]': (512, 14, 14),
    'classifier[0]': (4096,),
    'classifier[3]': (4096,),
    'classifier[6]': (1000,),
}

alpha_per_layer = {
    'features[0]': 30000,
    'features[2]': 30000,
    'features[5]': 30000,
    'features[7]': 30000,
    'features[10]': 30000,
    'features[12]': 30000,
    'features[14]': 30000,
    'features[16]': 25000,
    'features[19]': 25000,
    'features[21]': 25000,
    'features[23]': 25000,
    'features[25]': 25000,
    'features[28]': 25000,
    'features[30]': 25000,
    'features[32]': 25000,
    'features[34]': 25000,
    'classifier[0]': 20000,
    'classifier[3]': 20000,
    'classifier[6]': 20000,
}


# In[9]:


def get_numpy_subset_of_features(train_loader, test_loader, encoder, current_features, current_split):
    # check that the num_split is less than the splits_per_layer
    assert current_split <= splits_per_layer[current_features], "num_split is greater than splits_per_layer"
    
    size_of_features_for_split = math.ceil(layer_sizes[current_features] / splits_per_layer[current_features])
    start_feature_index = size_of_features_for_split * (current_split - 1)
    end_feature_index = size_of_features_for_split * current_split
    print(f"start_feature_index: {start_feature_index}, end_feature_index: {end_feature_index}")
    with torch.no_grad():
        if current_features in ['classifier[0]', 'classifier[3]', 'classifier[6]']:
            train_features = np.zeros(tuple([len(train_loader.dataset)] + [size_of_features_for_split])).astype(np.float32)
            test_features = np.zeros(tuple([len(test_loader.dataset)] + [size_of_features_for_split])).astype(np.float32)
        else:
            train_features = np.zeros(tuple([len(train_loader.dataset)] + [size_of_features_for_split] + list(features_shapes[current_features][1:]))).astype(np.float32)
            test_features = np.zeros(tuple([len(test_loader.dataset)] + [size_of_features_for_split] + list(features_shapes[current_features][1:]))).astype(np.float32)

        for i, (betas, c_image, nsd_id) in enumerate(tqdm(train_loader)):
            c_image = c_image.to(device)
            features = encoder(c_image)
            train_features[i * batch_size:features[current_features].shape[0] + i * batch_size] = features[current_features][:, start_feature_index:end_feature_index].cpu().numpy()

        for i, (betas, c_image, nsd_id) in enumerate(tqdm(test_loader)):
            c_image = c_image.to(device)
            features = encoder(c_image)
            test_features[i * batch_size:features[current_features].shape[0] + i * batch_size] = features[current_features][:, start_feature_index:end_feature_index].cpu().numpy()


    return train_features, test_features


# In[10]:


# train_features, test_features = get_numpy_subset_of_features(train_loader, test_loader, feature_extractor, current_features, 1)


# In[11]:


imagery_data_path = '/weka/proj-medarc/shared/umn-imagery'
# load nsd_imagery_data
voxels_vision, all_images_vision = utils.load_nsd_mental_imagery(subject=subj, mode='vision', stimtype="all", average=False, nest=True, data_root=imagery_data_path)
voxels_imagery, all_images_imagery = utils.load_nsd_mental_imagery(subject=subj, mode='imagery', stimtype="all", average=False, nest=True, data_root=imagery_data_path)


# In[12]:


def compute_mean_keepdims(train_features, feature_axis=1):
    axes_to_average = tuple(i for i in range(train_features.ndim) if i != feature_axis)
    y_mean = np.mean(train_features, axis=axes_to_average)
    return y_mean  # Shape: (1, features, 1, 1) or similar, depending on feature_axis


# In[ ]:


outdir_for_feature = os.path.join(outdir, current_features)
os.makedirs(outdir_for_feature, exist_ok=True)

# iterate over the splits
for calc_rn_split in tqdm(range(1,splits_per_layer[current_features]+1)):
    print(f"Calculating split {calc_rn_split} of {splits_per_layer[current_features]}")
    train_features, test_features = get_numpy_subset_of_features(train_loader, test_loader, encoder, current_features, calc_rn_split)
    size_of_features_for_split = math.ceil(layer_sizes[current_features] / splits_per_layer[current_features])
    
    print(f"Starting ridge regression for split {calc_rn_split} with alpha {alpha_per_layer[current_features]}")
    ridge = Ridge(alpha=alpha_per_layer[current_features])
    ridge.fit(x_train.reshape(x_train.shape[0], -1), train_features.reshape(train_features.shape[0], -1))
    print(f"Finished, now scoring")
    train_score = ridge.score(x_train.reshape(x_train.shape[0], -1), train_features.reshape(train_features.shape[0], -1))
    test_score = ridge.score(x_test.reshape(x_test.shape[0], -1), test_features.reshape(test_features.shape[0], -1))
    print(f"train_score: {train_score}, test_score: {test_score}")

    if current_features in ['classifier[0]', 'classifier[3]', 'classifier[6]']:
        target_feature_shape = (size_of_features_for_split,)
    else:
        target_feature_shape = (size_of_features_for_split,) + tuple(features_shapes[current_features][1:])

    y_mean = compute_mean_keepdims(train_features)
    # save the mean
    np.save(f'{outdir_for_feature}/ridge_y_mean_{current_features}_{calc_rn_split}.npy', y_mean.astype(np.float16))

    # save the scores
    with open(f'{outdir_for_feature}/ridge_scores_{current_features}_{calc_rn_split}.json', 'w') as f:
        json.dump({'train_score': train_score, 'test_score': test_score}, f)
    # save the weights
    if save_ckpt:
        np.save(f'{outdir_for_feature}/ridge_weights_{current_features}_{calc_rn_split}.npy', ridge.coef_.astype(np.float16))
    # save the intercept
    if save_ckpt:
        np.save(f'{outdir_for_feature}/ridge_intercept_{current_features}_{calc_rn_split}.npy', ridge.intercept_.astype(np.float16))

    # save the test predictions
    test_predictions = ridge.predict(x_test.reshape(x_test.shape[0], -1))
    np.save(f'{outdir_for_feature}/ridge_test_predictions_{current_features}_{calc_rn_split}.npy', test_predictions.reshape(tuple([test_predictions.shape[0]] + list(target_feature_shape))).astype(np.float16))

    # vision_preds = None
    # # get predictions for the imagery data: vision
    # for i, (voxel, image) in enumerate(tqdm(zip(voxels_vision, all_images_vision))):
    #     voxel = voxel # 8, 15724
    #     pred = ridge.predict(voxel.cpu().numpy())

    #     if vision_preds is None:
    #         vision_preds = np.expand_dims(pred, axis=0)
    #     else:
    #         vision_preds = np.concatenate((vision_preds, np.expand_dims(pred, axis=0)), axis=0)

    # np.save(f'{outdir_for_feature}/ridge_vision_preds_{current_features}_{calc_rn_split}.npy', vision_preds.reshape(tuple([vision_preds.shape[0]] + [vision_preds.shape[1]] + list(target_feature_shape))).astype(np.float16))

    # # get the predictions for the imagery data: imagery
    # imagery_preds = None
    # for i, (voxel, image) in enumerate(tqdm(zip(voxels_imagery, all_images_imagery))):
    #     voxel = voxel # 8, 15724
    #     pred = ridge.predict(voxel.cpu().numpy())

    #     if imagery_preds is None:
    #         imagery_preds = np.expand_dims(pred, axis=0)
    #     else:
    #         imagery_preds = np.concatenate((imagery_preds, np.expand_dims(pred, axis=0)), axis=0)

    # np.save(f'{outdir_for_feature}/ridge_imagery_preds_{current_features}_{calc_rn_split}.npy', imagery_preds.reshape(tuple([imagery_preds.shape[0]] + [imagery_preds.shape[1]] + list(target_feature_shape))).astype(np.float16))

    # get the predictions for averaged imagery data: vision
    vision_averaged_voxels = np.mean(np.array(voxels_vision), axis=1)
    vision_averaged_preds = ridge.predict(vision_averaged_voxels)
    np.save(f'{outdir_for_feature}/ridge_vision_averaged_preds_{current_features}_{calc_rn_split}.npy', vision_averaged_preds.reshape(tuple([vision_averaged_preds.shape[0]] + list(target_feature_shape))).astype(np.float16))

    # get the predictions for averaged imagery data: imagery
    imagery_averaged_voxels = np.mean(np.array(voxels_imagery), axis=1)
    imagery_averaged_preds = ridge.predict(imagery_averaged_voxels)
    np.save(f'{outdir_for_feature}/ridge_imagery_averaged_preds_{current_features}_{calc_rn_split}.npy', imagery_averaged_preds.reshape(tuple([imagery_averaged_preds.shape[0]] + list(target_feature_shape))).astype(np.float16))


# In[ ]:


# train_score: 0.36801715559650683, test_score: 0.1672737750357223
# train_score: 0.40319354674904023, test_score: 0.14771963878285566


# In[ ]:


# size_of_features_for_split = 2
# rsp = test_predictions.reshape(tuple([test_predictions.shape[0]] + [size_of_features_for_split] + list(features_shapes[current_features][1:])))


# In[ ]:


# # rsp.shape

# Error displaying widget: model not found
# Calculating split 1 of 2
# start_feature_index: 0, end_feature_index: 256
# Error displaying widget: model not found
# Error displaying widget: model not found
# Starting ridge regression for split 1
# Finished, now scoring
# train_score: 0.47192760353898633, test_score: 0.17910965480278365
# Error displaying widget: model not found
# Error displaying widget: model not found
# Calculating split 2 of 2
# start_feature_index: 256, end_feature_index: 512
# Error displaying widget: model not found
# Error displaying widget: model not found
# Starting ridge regression for split 2
# Finished, now scoring
# train_score: 0.47058061967151404, test_score: 0.17653868688099578
# Error displaying widget: model not found
# Error displaying widget: model not found


# In[ ]:


# # rsp.shape

# Error displaying widget: model not found
# Calculating split 1 of 2
# start_feature_index: 0, end_feature_index: 256
# Error displaying widget: model not found
# Error displaying widget: model not found
# Starting ridge regression for split 1
# Finished, now scoring
# train_score: 0.47192760353898633, test_score: 0.17910965480278365
# Error displaying widget: model not found
# Error displaying widget: model not found
# Calculating split 2 of 2
# start_feature_index: 256, end_feature_index: 512
# Error displaying widget: model not found
# Error displaying widget: model not found
# Starting ridge regression for split 2
# Finished, now scoring
# train_score: 0.47058061967151404, test_score: 0.17653868688099578
# Error displaying widget: model not found
# Error displaying widget: model not found


# In[ ]:


# 100.000
# # 100%
# #  2/2 [07:02<00:00, 205.79s/it]
# # Calculating split 1 of 2
# # start_feature_index: 0, end_feature_index: 256
# # 100%
# #  211/211 [00:28<00:00, 12.92it/s]
# # 100%
# #  8/8 [00:12<00:00,  1.23it/s]
# # Starting ridge regression for split 1
# # Finished, now scoring
# # train_score: 0.25576497027366507, test_score: 0.17944773415519444
# #  18/? [00:02<00:00,  9.75it/s]
# #  18/? [00:01<00:00, 10.16it/s]
# # Calculating split 2 of 2
# # start_feature_index: 256, end_feature_index: 512
# # 100%
# #  211/211 [00:25<00:00, 14.23it/s]
# # 100%
# #  8/8 [00:11<00:00,  1.35it/s]
# # Starting ridge regression for split 2
# # Finished, now scoring
# # train_score: 0.2538025531233917, test_score: 0.17658538319863515
# #  18/? [00:02<00:00,  7.65it/s]
# #  18/? [00:01<00:00, 10.60it/s]

# 60.000
# # Calculating split 1 of 2
# # start_feature_index: 0, end_feature_index: 256
# # 100%
# #  211/211 [00:28<00:00, 12.19it/s]
# # 100%
# #  8/8 [00:12<00:00,  1.24it/s]
# # Starting ridge regression for split 1
# # Finished, now scoring
# # train_score: 0.29882922368878595, test_score: 0.1865818620210305
# #  18/? [00:03<00:00,  6.38it/s]
# #  18/? [00:02<00:00,  7.63it/s]
# # Calculating split 2 of 2
# # start_feature_index: 256, end_feature_index: 512
# # 100%
# #  211/211 [00:28<00:00, 12.74it/s]
# # 100%
# #  8/8 [00:14<00:00,  1.07it/s]
# # Starting ridge regression for split 2
# # Finished, now scoring
# # train_score: 0.29697740748685114, test_score: 0.18380820114650484
# #  18/? [00:03<00:00,  6.08it/s]
# #  18/? [00:02<00:00,  7.28it/s]

# 3.000
# # 100%
# #  2/2 [05:16<00:00, 158.05s/it]
# # Calculating split 1 of 2
# # start_feature_index: 0, end_feature_index: 256
# # 100%
# #  211/211 [00:26<00:00, 13.16it/s]
# # 100%
# #  8/8 [00:11<00:00,  1.28it/s]
# # Starting ridge regression for split 1
# # Finished, now scoring
# # train_score: 0.5698654322488805, test_score: 0.13571111344383577
# #  18/? [00:02<00:00,  6.78it/s]
# #  18/? [00:02<00:00,  7.33it/s]
# # Calculating split 2 of 2
# # start_feature_index: 256, end_feature_index: 512
# # 100%
# #  211/211 [00:26<00:00, 13.34it/s]
# # 100%
# #  8/8 [00:11<00:00,  1.30it/s]
# # Starting ridge regression for split 2
# # Finished, now scoring
# # train_score: 0.5688183958383461, test_score: 0.13313861189424853
# #  18/? [00:03<00:00,  5.49it/s]
# #  18/? [00:02<00:00,  6.20it/s]

# 30.000
# # 100%
# #  2/2 [06:36<00:00, 191.34s/it]
# # Calculating split 1 of 2
# # start_feature_index: 0, end_feature_index: 256
# # 100%
# #  211/211 [00:27<00:00, 12.63it/s]
# # 100%
# #  8/8 [00:12<00:00,  1.29it/s]
# # Starting ridge regression for split 1
# # Finished, now scoring
# # train_score: 0.364419577220044, test_score: 0.19057156925390528
# #  18/? [00:02<00:00,  7.21it/s]
# #  18/? [00:01<00:00,  9.08it/s]
# # Calculating split 2 of 2
# # start_feature_index: 256, end_feature_index: 512
# # 100%
# #  211/211 [00:26<00:00, 13.14it/s]
# # 100%
# #  8/8 [00:11<00:00,  1.28it/s]
# # Starting ridge regression for split 2
# # Finished, now scoring
# # train_score: 0.3627511765041596, test_score: 0.18790273193089738
# #  18/? [00:03<00:00,  8.05it/s]
# #  18/? [00:02<00:00,  8.58it/s]

# 40.000
# # 100%
# #  8/8 [00:14<00:00,  1.05it/s]
# # Starting ridge regression for split 2
# # Finished, now scoring
# # train_score: 0.3347034806338601, test_score: 0.1871069173521844
# #  18/? [00:03<00:00,  5.37it/s]
# #  18/? [00:02<00:00,  7.88it/s]
# # 100%
# #  8/8 [00:14<00:00,  1.05it/s]
# # Starting ridge regression for split 2
# # Finished, now scoring
# # train_score: 0.3347034806338601, test_score: 0.1871069173521844
# #  18/? [00:03<00:00,  5.37it/s]
# #  18/? [00:02<00:00,  7.88it/s]

# 20.000
# # 100%
# #  2/2 [05:22<00:00, 161.87s/it]
# # Calculating split 1 of 2
# # start_feature_index: 0, end_feature_index: 256
# # 100%
# #  211/211 [00:27<00:00, 13.28it/s]
# # 100%
# #  8/8 [00:12<00:00,  1.27it/s]
# # Starting ridge regression for split 1
# # Finished, now scoring
# # train_score: 0.4046410458211125, test_score: 0.18916135958044172
# #  18/? [00:03<00:00,  5.01it/s]
# #  18/? [00:02<00:00,  6.18it/s]
# # Calculating split 2 of 2
# # start_feature_index: 256, end_feature_index: 512
# # 100%
# #  211/211 [00:26<00:00, 12.82it/s]
# # 100%
# #  8/8 [00:13<00:00,  1.13it/s]
# # Starting ridge regression for split 2
# # Finished, now scoring
# # train_score: 0.4030907987701695, test_score: 0.18653935361366922
# #  18/? [00:02<00:00,  7.51it/s]
# #  18/? [00:01<00:00,  8.18it/s]


# In[ ]:


# 30.000 0.33
# 1672 30