main.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Pretrain using the Barlow Twin method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from modules.pretraining import main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the config file to set the hyperparameters\n",
    "config = yaml.load(open(\"config.yaml\", \"r\"), Loader=yaml.FullLoader)\n",
    "print(config)\n",
    "\n",
    "main(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Finetune the model by loading in the pretrained model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "from utils.plots import plot_srcc_MAE\n",
    "from modules.finetune import main_finetune\n",
    "from modules.finetune import run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def run(task, main_folder, tune_from, runs, targets):\n",
    "    '''Run the fine-tuning experiments for the given task.\n",
    "    task: str, the name of the task\n",
    "    main_folder: str, the folder to save the results\n",
    "    tune_from: str, the name of the pre-trained model\n",
    "    runs: int, the number of runs\n",
    "    targets: int, the number of train splits to fine-tune on'''\n",
    "\n",
    "    config = yaml.load(open(\"config_finetune.yaml\", \"r\"), Loader=yaml.FullLoader)\n",
    "\n",
    "    config['task_name'] = task\n",
    "\n",
    "    if config[\"task_name\"] == 'visc':\n",
    "        config['dataset']['task'] = 'regression'\n",
    "        config['dataset']['data_path'] = 'GNN_BT_Data/visc_data.csv'\n",
    "        target_list = [['Viscosity_1', 'Viscosity_2', 'Viscosity_3', 'Viscosity_4', 'Viscosity_5', 'T1', 'T2', 'T3', 'T4', 'T5', 'Ln(A)', 'Ea/R', 'smiles']]\n",
    "    elif config[\"task_name\"] == 'cond':\n",
    "        config['dataset']['task'] = 'regression'\n",
    "        config['dataset']['data_path'] = 'GNN_BT_Data/cond_data.csv'\n",
    "        target_list = [['K1', 'K2', 'K3', 'K4', 'K5', 'T1', 'T2', 'T3', 'T4', 'T5', 'Intercept', 'Coefficients', 'SMILES']]\n",
    "    elif config[\"task_name\"] == 'visc_hc':\n",
    "        config['dataset']['task'] = 'regression'\n",
    "        config['dataset']['data_path'] = 'GNN_BT_Data/visc_hc_data.csv'\n",
    "        target_list = [['Viscosity_1', 'Viscosity_2', 'Viscosity_3', 'Viscosity_4', 'Viscosity_5', 'T1', 'T2', 'T3', 'T4', 'T5', 'Intercept (A)', 'Coefficient (B)', 'smiles']]\n",
    "    else:\n",
    "        raise ValueError('Undefined downstream task!')\n",
    "\n",
    "    print(config)\n",
    "\n",
    "    # save all the srcc and mae values in a dictionary\n",
    "    # save results in a yaml file\n",
    "    s_avg = np.zeros((targets, 5))\n",
    "    s_std = np.zeros((targets, 5))\n",
    "    m_avg = np.zeros((targets, 5))\n",
    "    m_std = np.zeros((targets, 5))\n",
    "\n",
    "    BT_s_avg = np.zeros((targets, 5))\n",
    "    BT_s_std = np.zeros((targets, 5))\n",
    "    BT_m_avg = np.zeros((targets, 5))\n",
    "    BT_m_std = np.zeros((targets, 5))\n",
    "\n",
    "    results = {\n",
    "        'scratch': np.zeros((runs, targets, 5)).tolist(),\n",
    "        'BT': np.zeros((runs, targets, 5)).tolist(),\n",
    "        'scratch_mae': np.zeros((runs, targets, 5)).tolist(),\n",
    "        'BT_mae': np.zeros((runs, targets, 5)).tolist(),\n",
    "        'axis': np.zeros(targets).tolist()\n",
    "    }\n",
    "\n",
    "    for j in range(runs):\n",
    "        folder = os.path.join(main_folder, f'test_{j+1}')\n",
    "        os.makedirs(folder, exist_ok=False)\n",
    "        \n",
    "        scratch = np.zeros((targets, 5))\n",
    "        BT = np.zeros((targets, 5))\n",
    "        scratch_mae = np.zeros((targets, 5))\n",
    "        BT_mae = np.zeros((targets, 5))\n",
    "        axis = np.zeros(targets)\n",
    "\n",
    "        a = 0\n",
    "        for i in range(6, 6-targets, -1):\n",
    "            print(\"start\")\n",
    "            config['dataset']['target'] = target_list[0]\n",
    "            config['dataset']['train_size'] = (i+1)*0.1\n",
    "            config['save_folder'] = folder\n",
    "            config['fine_tune_from'] = tune_from\n",
    "            config['name'] = f'BT_{(i+1)*0.1:.1f}'\n",
    "            \n",
    "            srcc_BT, mae_BT = main_finetune(config)\n",
    "            BT[a] = srcc_BT\n",
    "            BT_mae[a] = mae_BT\n",
    "\n",
    "            config['fine_tune_from'] = 'None'\n",
    "            config['name'] = f'Scratch_{(i+1)*0.1:.1f}'\n",
    "            srcc_s, mae_s = main_finetune(config)\n",
    "            scratch[a] = srcc_s\n",
    "            scratch_mae[a] = mae_s\n",
    "            \n",
    "\n",
    "            if config[\"task_name\"] == 'visc':\n",
    "                axis[a] = int((i+1)*0.1*477)\n",
    "            elif config[\"task_name\"] == 'cond':\n",
    "                axis[a] = int((i+1)*0.1*1222)\n",
    "            elif config[\"task_name\"] == 'visc_hc':\n",
    "                axis[a] = int((i+1)*0.1*182)\n",
    "\n",
    "            \n",
    "            print(\"done\")\n",
    "        \n",
    "        for i in range(5):\n",
    "            plot_srcc_MAE(scratch[:, i], BT[:, i], s_std[:, i], BT_s_std[:, i], axis, f'{i+1}', folder, tag2='srcc')\n",
    "            plot_srcc_MAE(scratch_mae[:, i], BT_mae[:, i], m_std[:, i], BT_m_std[:, i], axis, f'{i+1}', folder, tag2='mae')\n",
    "\n",
    "        results['scratch'][j] = scratch.tolist()\n",
    "        results['BT'][j] = BT.tolist()\n",
    "        results['scratch_mae'][j] = scratch_mae.tolist()\n",
    "        results['BT_mae'][j] = BT_mae.tolist()\n",
    "        results['axis'] = axis.tolist()\n",
    "\n",
    "        s_avg += scratch\n",
    "        m_avg += scratch_mae\n",
    "        BT_s_avg += BT\n",
    "        BT_m_avg += BT_mae\n",
    "\n",
    "    s_avg /= runs\n",
    "    m_avg /= runs\n",
    "    BT_s_avg /= runs\n",
    "    BT_m_avg /= runs\n",
    "\n",
    "    if runs > 1:\n",
    "        s_std = np.std(np.array(results['scratch']), axis=0)\n",
    "        m_std = np.std(np.array(results['scratch_mae']), axis=0)\n",
    "        BT_s_std = np.std(np.array(results['BT']), axis=0)\n",
    "        BT_m_std = np.std(np.array(results['BT_mae']), axis=0)\n",
    "\n",
    "        for i in range(5):\n",
    "            plot_srcc_MAE(s_avg[:, i], BT_s_avg[:, i], s_std[:, i], BT_s_std[:, i], axis, f'{i+1}', main_folder, tag2='srcc', tag3 = 'avg')\n",
    "            plot_srcc_MAE(m_avg[:, i], BT_m_avg[:, i], m_std[:, i], BT_m_std[:, i], axis, f'{i+1}', main_folder, tag2='mae', tag3 = 'avg')\n",
    "\n",
    "    # Save results to YAML file\n",
    "    with open(f\"{main_folder}/results.yaml\", \"w\") as file:\n",
    "        yaml.dump(results, file)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use the finetune config_finetune file to change the rest of the hyperparameters\n",
    "# changes are made to the config file based on the infomation below. \n",
    "\n",
    "# lists used to automate the process of finetuning multiple models\n",
    "task = ['visc']   # 'cond', 'visc', 'visc_hc'\n",
    "tune_from = ['BT1']  # BT1: 15% subgraph, BT2: 15% subgraph, 20% node/edge\n",
    "runs = [1]     # number of tests to run with the same hyperparameters\n",
    "targets = [1]      # number of train set size splits to use (max is 7) - 0.1, 0.2, 0.3...0.7\n",
    "main_folder = ['results']  # folder to save the results\n",
    "\n",
    "for i in range(len(main_folder)):\n",
    "    run(task[i], f\"finetune/{main_folder[i]}\", tune_from[i], runs[i], targets[i])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "research",
   "language": "python",
   "name": "python3"
  },
  "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}