{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "b8e236f1-385a-4d93-bb39-bea3ee384d76", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "outdir /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/checkpoints/HCPflat_large_gsrFalse_\n", "Loaded config.yaml from ckpt folder /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/checkpoints/HCPflat_large_gsrFalse_\n", "\n", "__CONFIG__\n", "base_lr = 0.001\n", "batch_size = 32\n", "ckpt_interval = 5\n", "ckpt_saving = True\n", "cls_embed = True\n", "contrastive_loss_weight = 1.0\n", "datasets_to_include = HCP\n", "decoder_embed_dim = 512\n", "grad_accumulation_steps = 1\n", "grad_clip = 1.0\n", "gsr = False\n", "hcp_flat_path = /weka/proj-medarc/shared/HCP-Flat\n", "mask_ratio = 0.75\n", "model_name = HCPflat_large_gsrFalse_\n", "no_qkv_bias = False\n", "norm_pix_loss = False\n", "nsd_flat_path = /weka/proj-medarc/shared/NSD-Flat\n", "num_epochs = 100\n", "num_frames = 16\n", "num_samples_per_epoch = 200000\n", "num_workers = 10\n", "patch_size = 16\n", "pct_masks_to_decode = 1\n", "plotting = True\n", "pred_t_dim = 8\n", "print_interval = 20\n", "probe_base_lr = 0.0003\n", "probe_batch_size = 8\n", "probe_num_epochs = 30\n", "probe_num_samples_per_epoch = 100000\n", "resume_from_ckpt = True\n", "seed = 42\n", "sep_pos_embed = True\n", "t_patch_size = 2\n", "test_num_samples_per_epoch = 50000\n", "test_set = False\n", "trunc_init = False\n", "use_contrastive_loss = False\n", "wandb_log = True\n", "\n", "\n", "WORLD_SIZE=1\n", "PID of this process = 1830098\n" ] } ], "source": [ "# Import packages and setup gpu configuration.\n", "# This code block shouldnt need to be adjusted!\n", "import os\n", "import sys\n", "import json\n", "import yaml\n", "import numpy as np\n", "import copy\n", "import math\n", "import time\n", "import random\n", "from tqdm import tqdm\n", "import webdataset as wds\n", "import matplotlib.pyplot as plt\n", "\n", "import torch\n", "import torch.nn as nn\n", "from torchvision import transforms\n", "import utils\n", "from mae_utils.flat_models import *\n", "import pandas as pd\n", "from typing import List, Dict, Any, Tuple\n", "from sklearn.preprocessing import StandardScaler\n", "\n", "# tf32 data type is faster than standard float32\n", "torch.backends.cuda.matmul.allow_tf32 = True\n", "# following fixes a Conv3D CUDNN_NOT_SUPPORTED error\n", "torch.backends.cudnn.benchmark = True\n", "\n", "## MODEL TO LOAD ##\n", "# YOU WILL NEED TO PRECOMPUTE AND SAVE THE FEATURES SEE: prep_HCP_downstream.ipynb\n", "model_name = \"HCPflat_large_gsrFalse_\"\n", "parquet_folder = \"epoch99\"\n", "\n", "## DEFINE TARGET AND GLOBAL POOLING ##\n", "global_pooling = True\n", "target = \"trial_type\" # \"trial_type\" or \"sex\" or \"age\"\n", "\n", "\n", "# outdir = os.path.abspath(f'checkpoints/{model_name}')\n", "outdir = os.path.abspath(f'checkpoints/{model_name}')\n", "\n", "print(\"outdir\", outdir)\n", "# Load previous config.yaml if available\n", "if os.path.exists(f\"{outdir}/config.yaml\"):\n", " config = yaml.load(open(f\"{outdir}/config.yaml\", 'r'), Loader=yaml.FullLoader)\n", " print(f\"Loaded config.yaml from ckpt folder {outdir}\")\n", " # create global variables from the config\n", " print(\"\\n__CONFIG__\")\n", " for attribute_name in config.keys():\n", " print(f\"{attribute_name} = {config[attribute_name]}\")\n", " globals()[attribute_name] = config[f'{attribute_name}']\n", " print(\"\\n\")\n", "\n", "world_size = os.getenv('WORLD_SIZE')\n", "if world_size is None: \n", " world_size = 1\n", "else:\n", " world_size = int(world_size)\n", "print(f\"WORLD_SIZE={world_size}\")\n", "\n", "if utils.is_interactive():\n", " # Following allows you to change functions in models.py or utils.py and \n", " # have this notebook automatically update with your revisions\n", " %load_ext autoreload\n", " %autoreload 2\n", "\n", "batch_size = probe_batch_size\n", "num_epochs = probe_num_epochs\n", "\n", "data_type = torch.float32 # change depending on your mixed_precision\n", "global_batch_size = batch_size * world_size\n", "\n", "device = torch.device('cuda')\n", "\n", "print(\"PID of this process =\",os.getpid())\n", "\n", "utils.seed_everything(seed)" ] }, { "cell_type": "markdown", "id": "ab15aca0-148e-435f-b8f2-7a708b61a6d9", "metadata": {}, "source": [ "### prepare hcp_flat dataset" ] }, { "cell_type": "code", "execution_count": 2, "id": "de1a5b87-fa69-44e1-bdb9-bd9e257485c6", "metadata": { "tags": [] }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/tmp/ipykernel_1830098/4239417916.py:33: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " state = torch.load(checkpoint_path)\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "img_size (144, 320) patch_size (16, 16) frames 16 t_patch_size 2\n", "model initialized\n", "\n", "Loaded checkpoint epoch99.pth from /weka/proj-fmri/ckadirt/fMRI-foundation-model/src/checkpoints/HCPflat_large_gsrFalse_\n", "\n" ] }, { "data": { "text/plain": [ "MaskedAutoencoderViT(\n", " (patch_embed): PatchEmbed(\n", " (proj): Conv3d(1, 1024, kernel_size=(2, 16, 16), stride=(2, 16, 16))\n", " )\n", " (blocks): ModuleList(\n", " (0-23): 24 x Block(\n", " (norm1): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n", " (attn): Attention(\n", " (q): Linear(in_features=1024, out_features=1024, bias=True)\n", " (k): Linear(in_features=1024, out_features=1024, bias=True)\n", " (v): Linear(in_features=1024, out_features=1024, bias=True)\n", " (proj): Linear(in_features=1024, out_features=1024, bias=True)\n", " (proj_drop): Dropout(p=0.0, inplace=False)\n", " )\n", " (drop_path): Identity()\n", " (norm2): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n", " (mlp): Mlp(\n", " (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n", " (act): GELU(approximate='none')\n", " (drop1): Dropout(p=0.0, inplace=False)\n", " (norm): Identity()\n", " (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n", " (drop2): Dropout(p=0.0, inplace=False)\n", " )\n", " )\n", " )\n", " (norm): LayerNorm((1024,), eps=1e-06, elementwise_affine=True)\n", " (decoder_embed): Linear(in_features=1024, out_features=512, bias=True)\n", " (decoder_blocks): ModuleList(\n", " (0-3): 4 x Block(\n", " (norm1): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n", " (attn): Attention(\n", " (q): Linear(in_features=512, out_features=512, bias=True)\n", " (k): Linear(in_features=512, out_features=512, bias=True)\n", " (v): Linear(in_features=512, out_features=512, bias=True)\n", " (proj): Linear(in_features=512, out_features=512, bias=True)\n", " (proj_drop): Dropout(p=0.0, inplace=False)\n", " )\n", " (drop_path): Identity()\n", " (norm2): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n", " (mlp): Mlp(\n", " (fc1): Linear(in_features=512, out_features=2048, bias=True)\n", " (act): GELU(approximate='none')\n", " (drop1): Dropout(p=0.0, inplace=False)\n", " (norm): Identity()\n", " (fc2): Linear(in_features=2048, out_features=512, bias=True)\n", " (drop2): Dropout(p=0.0, inplace=False)\n", " )\n", " )\n", " )\n", " (decoder_norm): LayerNorm((512,), eps=1e-06, elementwise_affine=True)\n", " (decoder_pred): Linear(in_features=512, out_features=256, bias=True)\n", ")" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from mae_utils.flat import load_hcp_flat_mask\n", "from mae_utils.flat import create_hcp_flat\n", "import mae_utils.visualize as vis\n", "\n", "if utils.is_interactive(): # Use less samples per epoch for debugging\n", " probe_num_samples_per_epoch = 100000\n", " test_num_samples_per_epoch = 100000\n", " num_epochs = 10\n", "\n", "\n", "# Load ckpt\n", "if not os.path.exists(outdir) or not os.path.isdir(outdir):\n", " assert True, (f\"\\nCheckpoint folder {outdir} does not exist.\\n\")\n", "else:\n", " checkpoint_files = [f for f in os.listdir(outdir) if f.endswith('.pth')]\n", "\n", " # Find the latest ckpt to load\n", " epoch_numbers = []\n", " for file in checkpoint_files:\n", " try:\n", " epoch_number = int(file.split('epoch')[-1].split('.')[0])\n", " epoch_numbers.append(epoch_number)\n", " except ValueError:\n", " continue\n", " latest_epoch = max(epoch_numbers)\n", " checkpoint_name = f\"epoch{latest_epoch}.pth\"\n", " \n", " ### Or provide the specific checkpoint you want to load\n", " # checkpoint_name = \"epoch10.pth\" #\"epoch15.pth\"\n", "\n", " # Load the checkpoint\n", " checkpoint_path = os.path.join(outdir, checkpoint_name)\n", " state = torch.load(checkpoint_path)\n", "\n", "model = mae_vit_large_fmri(\n", " patch_size=16,\n", " decoder_embed_dim=decoder_embed_dim,\n", " t_patch_size=t_patch_size,\n", " pred_t_dim=pred_t_dim,\n", " decoder_depth=4,\n", " cls_embed=cls_embed,\n", " norm_pix_loss=norm_pix_loss,\n", " no_qkv_bias=no_qkv_bias,\n", " sep_pos_embed=sep_pos_embed,\n", " trunc_init=trunc_init,\n", " img_mask=state[\"model_state_dict\"]['img_mask']\n", ")\n", "\n", "model.load_state_dict(state[\"model_state_dict\"], strict=True) #model_state_dict\n", "print(f\"\\nLoaded checkpoint {checkpoint_name} from {outdir}\\n\")\n", "\n", "model.eval()\n", "model.requires_grad_(False)\n", "model.to(device)" ] }, { "cell_type": "markdown", "id": "e4ab9f5d-ce67-430a-8f1f-d4e97f217fd2", "metadata": {}, "source": [ "### Prepare subjects information" ] }, { "cell_type": "code", "execution_count": 3, "id": "29672348-687e-44b5-916e-4dd4779cf560", "metadata": {}, "outputs": [], "source": [ "# open the file containing subject information\n", "if target == \"age\" or target == \"sex\":\n", " subject_information_HCP_path = os.path.join(hcp_flat_path, \"subjects_data_restricted.csv\")\n", " try:\n", " subject_information_HCP = pd.read_csv(subject_information_HCP_path)\n", " except:\n", " try:\n", " subject_information_HCP = pd.read_csv('./unrestricted_clane9_4_23_2024_13_28_14.csv') \n", " except:\n", " assert False, \"Subject information file not found\"\n", "\n", " ###### This is for unrestricted\n", " # age_related_columns = [\n", " # 'Age', 'PicSeq_AgeAdj', 'CardSort_AgeAdj', 'Flanker_AgeAdj',\n", " # 'ReadEng_AgeAdj', 'PicVocab_AgeAdj', 'ProcSpeed_AgeAdj',\n", " # 'CogFluidComp_AgeAdj', 'CogEarlyComp_AgeAdj', 'CogTotalComp_AgeAdj',\n", " # 'CogCrystalComp_AgeAdj', 'Endurance_AgeAdj', 'Dexterity_AgeAdj',\n", " # 'Strength_AgeAdj', 'Odor_AgeAdj', 'Taste_AgeAdj'\n", " # ]\n", " \n", " # sex_related_columns = [\n", " # 'Gender'\n", " # ]\n", "\n", " ###### This is for restricted\n", " gender_related_columns = [\n", " 'Gender'\n", " ]\n", "\n", " age_related_columns = [\n", " 'Age_in_Yrs',\n", " 'Menstrual_AgeBegan',\n", " 'Menstrual_AgeIrreg',\n", " 'Menstrual_AgeStop',\n", " 'SSAGA_Alc_Age_1st_Use',\n", " 'SSAGA_TB_Age_1st_Cig',\n", " 'SSAGA_Mj_Age_1st_Use',\n", " 'Endurance_AgeAdj',\n", " 'Dexterity_AgeAdj',\n", " 'Strength_AgeAdj',\n", " 'PicSeq_AgeAdj',\n", " 'CardSort_AgeAdj',\n", " 'Flanker_AgeAdj',\n", " 'ReadEng_AgeAdj',\n", " 'PicVocab_AgeAdj',\n", " 'ProcSpeed_AgeAdj',\n", " 'Odor_AgeAdj',\n", " 'Taste_AgeAdj'\n", " ]\n", "\n", " # # show the first few rows of the subject information\n", " # subject_information_HCP[age_related_columns + sex_related_columns].head()\n", "\n", " # Handle missing values (e.g., impute with mean)\n", " mean_age = subject_information_HCP['Age_in_Yrs'].mean()\n", " \n", " # Initialize the scaler\n", " scaler = StandardScaler()\n", " \n", " # Perform z-score normalization\n", " subject_information_HCP['Age_in_Yrs_z'] = scaler.fit_transform(subject_information_HCP[['Age_in_Yrs']])\n", "\n", "\n", " \n", "def get_label_unrestricted(subject_id: List[str], target: str, method_for_age: str = 'mean') -> List:\n", " \"\"\"\n", " Get the label for the given subject id and target.\n", "\n", " For sex 0 is F and 1 is M\n", " \"\"\"\n", "\n", " # convert to list of ints\n", " subject_id = [int(x) for x in subject_id]\n", "\n", " if target == \"age\":\n", " age_array = []\n", " for subject in subject_id:\n", " c_age = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Age'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_age) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_age) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", "\n", " c_age = c_age[0].split('-')\n", " if len(c_age) < 2:\n", " c_age = c_age[0].split('+')\n", " age_array.append(int(c_age[0]))\n", " else:\n", " if method_for_age == 'mean':\n", " age_array.append(np.mean([int(x) for x in c_age]))\n", " elif method_for_age == 'min':\n", " age_array.append(np.min([int(x) for x in c_age]))\n", " elif method_for_age == 'max':\n", " age_array.append(np.max([int(x) for x in c_age]))\n", " else:\n", " assert False, f\"Method {method_for_age} not recognized\"\n", "\n", " return np.array(age_array) \n", " \n", " elif target == 'sex':\n", " sex_array = []\n", " for subject in subject_id:\n", " c_sex = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Gender'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_sex) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_sex) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", " sex_array.append(int(c_sex[0] == 'M'))\n", " return sex_array\n", "\n", "def get_label_restricted(subject_id: List[str], target: str, normalized: bool = True) -> List:\n", " \"\"\"\n", " Get the label for the given subject id and target.\n", "\n", " For sex 0 is F and 1 is M\n", " \"\"\"\n", "\n", " # convert to list of ints\n", " subject_id = [int(x) for x in subject_id]\n", "\n", " if target == \"age\":\n", " age_array = []\n", " for subject in subject_id:\n", " c_age = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Age_in_Yrs' if not normalized else 'Age_in_Yrs_z'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_age) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_age) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", "\n", " age_array.append(np.int8(c_age[0]))\n", "\n", " return np.array(age_array) \n", " \n", " elif target == 'sex':\n", " sex_array = []\n", " for subject in subject_id:\n", " c_sex = subject_information_HCP[subject_information_HCP['Subject'] == subject]['Gender'].values\n", " # if the subject is not in the subject information file trigger an error\n", " if len(c_sex) == 0:\n", " assert False, f\"Subject {subject} not found in subject information file\"\n", " if len(c_sex) > 1:\n", " print(f\"Warning: Multiple entries for subject {subject}\")\n", " sex_array.append(int(c_sex[0] == 'M'))\n", " return sex_array" ] }, { "cell_type": "markdown", "id": "59ece3e1-3b04-4f92-8417-148fc933757d", "metadata": {}, "source": [ "### Fit sklearn model" ] }, { "cell_type": "code", "execution_count": 4, "id": "c3461199-e805-4e9c-8c91-894e83cf8bc3", "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Target: trial_type\n", "train: (111302, 9), test: (12082, 9)\n", "test: (12082, 9)\n", "X_train: (111302, 1024), X_test: (12082, 1024)\n", "X_test: (12082, 1024)\n", "classes (21): ['0bk_body' '0bk_faces' '0bk_places' '0bk_tools' '2bk_body' '2bk_faces'\n", " '2bk_places' '2bk_tools' 'fear' 'lf' 'lh' 'match' 'math' 'mental' 'neut'\n", " 'relation' 'rf' 'rh' 'rnd' 'story' 't']\n", "\n", "y_train: (111302,) [10 13 13 15 14 17 13 17 14 14 16 13 18 11 8 9 13 9 8 8]\n", "y_test: (12082,) [ 7 4 10 4 14 13 4 14 15 13 7 4 16 4 8 18 4 8 11 18]\n", "\n", "train_ind: 100171 [73280 30515 54498 3785 49428 41996 11968 37870 49360 80177]\n", "val_ind: 11131 [ 52061 23724 105636 84928 69162 94892 2339 55435 106582 63659]\n", "Fitting PCA projection\n", "Fitting logistic regression\n", "Done:\n", "{\"target\": \"trial_type\", \"train_acc\": 0.9657385870162023, \"val_acc\": 0.9542718533824455, \"test_acc\": 0.9557192517795067}\n" ] } ], "source": [ "import argparse\n", "import json\n", "import os\n", "import pickle\n", "from pathlib import Path\n", "\n", "import pandas as pd\n", "import numpy as np\n", "from sklearn.decomposition import PCA\n", "from sklearn.linear_model import LogisticRegressionCV, Ridge\n", "from sklearn.metrics import mean_squared_error\n", "from sklearn.model_selection import train_test_split\n", "from sklearn.preprocessing import LabelEncoder\n", "\n", "\n", "print(f\"Target: {target}\")\n", "\n", "train_features = pd.read_parquet(f\"{outdir}_gp{str(global_pooling)}/{parquet_folder}/HCP/train.parquet\")\n", "test_features = pd.read_parquet(f\"{outdir}_gp{str(global_pooling)}/{parquet_folder}/HCP/test.parquet\")\n", "\n", "print(f\"train: {train_features.shape}, test: {test_features.shape}\")\n", "print(f\"test: {test_features.shape}\")\n", "\n", "X_train = np.stack(train_features[\"feature\"])\n", "X_test = np.stack(test_features[\"feature\"])\n", "print(f\"X_train: {X_train.shape}, X_test: {X_test.shape}\")\n", "print(f\"X_test: {X_test.shape}\")\n", "\n", "\n", "if target == \"task\":\n", " labels_train = train_features[\"task\"].str.rstrip(\"1234\").values\n", " labels_test = test_features[\"task\"].str.rstrip(\"1234\").values\n", "elif target == \"trial_type\":\n", " labels_train = train_features[\"trial_type\"].values\n", " labels_test = test_features[\"trial_type\"].values\n", "elif target == \"sex\":\n", " labels_train = [str(sample[0]) for sample in train_features['sub'].values.tolist()]\n", " labels_test = [str(sample[0]) for sample in test_features['sub'].values.tolist()]\n", " labels_train = get_label_restricted(labels_train, target=\"sex\")\n", " labels_test = get_label_restricted(labels_test, target=\"sex\")\n", "elif target == \"age\":\n", " labels_train = [str(sample[0]) for sample in train_features['sub'].values.tolist()]\n", " labels_test = [str(sample[0]) for sample in test_features['sub'].values.tolist()]\n", " labels_train = get_label_restricted(labels_train, target=\"age\")\n", " labels_test = get_label_restricted(labels_test, target=\"age\")\n", "\n", "labels_train = np.array([label[0] if isinstance(label, np.ndarray) else label for label in labels_train])\n", "labels_test = np.array([label[0] if isinstance(label, np.ndarray) else label for label in labels_test])\n", "\n", "if target == \"trial_type\" or target == \"sex\":\n", " label_enc = LabelEncoder()\n", " y_train = label_enc.fit_transform(labels_train)\n", " y_test = label_enc.transform(labels_test)\n", " print(f\"classes ({len(label_enc.classes_)}): {label_enc.classes_}\")\n", "elif target == \"age\":\n", " y_train = np.array(labels_train)\n", " y_test = np.array(labels_test)\n", "\n", "\n", "print(\n", " f\"\\ny_train: {y_train.shape} {y_train[:20]}\\n\"\n", " f\"y_test: {y_test.shape} {y_test[:20]}\"\n", ")\n", "del train_features, test_features\n", "\n", "train_ind, val_ind = train_test_split(\n", " np.arange(len(X_train)), train_size=0.9, random_state=42\n", ")\n", "print(\n", " f\"\\ntrain_ind: {len(train_ind)} {train_ind[:10]}\\n\"\n", " f\"val_ind: {len(val_ind)} {val_ind[:10]}\"\n", ")\n", "X_train, X_val = X_train[train_ind], X_train[val_ind]\n", "y_train, y_val = y_train[train_ind], y_train[val_ind]\n", "\n", "print(\"Fitting PCA projection\")\n", "pca = PCA(n_components=384, whiten=True, svd_solver=\"randomized\")\n", "pca.fit(X_train)\n", "\n", "X_train = pca.transform(X_train)\n", "X_val = pca.transform(X_val)\n", "X_test = pca.transform(X_test)\n", "\n", "if target == \"trial_type\" or target == \"sex\":\n", " print(\"Fitting logistic regression\")\n", " clf = LogisticRegressionCV()\n", " clf.fit(X_train, y_train)\n", " \n", " train_acc = clf.score(X_train, y_train)\n", " val_acc = clf.score(X_val, y_val)\n", " test_acc = clf.score(X_test, y_test)\n", " \n", " result = {\n", " \"target\": target,\n", " \"train_acc\": train_acc,\n", " \"val_acc\": val_acc,\n", " \"test_acc\": test_acc,\n", " }\n", "\n", "elif target == \"age\":\n", " alpha = 10000\n", " print(\"Fitting ridge regression\")\n", " clf = Ridge()\n", " clf.fit(X_train, y_train)\n", " \n", " # Calculate R² scores\n", " train_r2 = clf.score(X_train, y_train)\n", " val_r2 = clf.score(X_val, y_val)\n", " test_r2 = clf.score(X_test, y_test)\n", " \n", " # Make predictions\n", " train_pred = clf.predict(X_train)\n", " val_pred = clf.predict(X_val)\n", " test_pred = clf.predict(X_test)\n", " \n", " # Calculate MSE\n", " train_mse = mean_squared_error(y_train, train_pred)\n", " val_mse = mean_squared_error(y_val, val_pred)\n", " test_mse = mean_squared_error(y_test, test_pred)\n", " \n", " # Compile results\n", " result = {\n", " \"target\": target,\n", " \"train_r2\": train_r2,\n", " \"val_r2\": val_r2,\n", " \"test_r2\": test_r2,\n", " \"train_mse\": float(train_mse),\n", " \"val_mse\": float(val_mse),\n", " \"test_mse\": float(test_mse),\n", " }\n", "\n", "with open(f\"{outdir}_gp{str(global_pooling)}/{parquet_folder}/HCP/downstream_{target}.json\", 'w') as out_json:\n", " json.dump(result, out_json)\n", "\n", "print(f\"Done:\\n{json.dumps(result)}\")" ] } ], "metadata": { "kernelspec": { "display_name": "foundation_env", "language": "python", "name": "foundation_env" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.10" } }, "nbformat": 4, "nbformat_minor": 5 }