{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tuning the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although our model is quite general for different time span, observational strategies, and colors. It can imporve to better performance on a specific mission or to strengthen the model's ability for dealing with some special cases.\n",
    "\n",
    "Thus, we supply an interface for you to fine-tune the model on given datasets. In this tutorial, we will show how to use the fine-tunning interface. Here we choose some ZTF archive data as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "make sure you have GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.set_device(0) # change this if you have more than one GPU\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "make sure the training is deterministic and reproducible"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "seed = 42\n",
    "\n",
    "random.seed(seed)  # python random generator\n",
    "np.random.seed(seed)  # numpy random generator\n",
    "\n",
    "torch.manual_seed(seed)\n",
    "torch.cuda.manual_seed_all(seed)\n",
    "torch.backends.cudnn.deterministic = True\n",
    "\n",
    "\n",
    "def seed_worker(worker_id):\n",
    "    worker_seed = torch.initial_seed() % 2**32\n",
    "    np.random.seed(worker_seed)\n",
    "    import random\n",
    "\n",
    "    random.seed(worker_seed)\n",
    "\n",
    "\n",
    "g = torch.Generator()\n",
    "g.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load your own dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Your own dataset should located at a directory as the following structure:\n",
    "\n",
    "```\n",
    "dataset\n",
    "├── train\n",
    "│   ├── 0\n",
    "│   ├── 1\n",
    "│   └── ...\n",
    "└── test\n",
    "    ├── 0\n",
    "    ├── 1\n",
    "    └── ...\n",
    "└── train_labels.csv\n",
    "└── test_labels.csv\n",
    "```\n",
    "\n",
    "All the data inside the train/test folder can be any format that can be constructed as numpy ndarrays.\n",
    "\n",
    "In our work, we only use light curve as input, so the data are all light curve ndarrays. But you can also include parameters since we have parameter component in our combined network. Other components like spectrum/SED are not implenmented in our model, but can be easily added."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You might need to modify the following code to load your own dataset properly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First, load your dataset to lists of numpy ndarrays:\n",
    "\n",
    "train_lcs = []\n",
    "train_params = []\n",
    "train_labels = []\n",
    "\n",
    "test_lcs = []\n",
    "test_params = []\n",
    "test_labels = []\n",
    "\n",
    "train_csv = np.loadtxt('dataset/train_labels.csv', delimiter=',')\n",
    "test_csv = np.loadtxt('dataset/test_labels.csv', delimiter=',')\n",
    "\n",
    "# If light curves are not normalized to relative variability, please apply our \n",
    "# `light_curve_preparation` function to them first\n",
    "\n",
    "for i in range(len(train_csv)):\n",
    "    lc_id = int(train_csv[i][0])\n",
    "    train_lcs.append(np.load(f'dataset/train/{lc_id}.npy'))\n",
    "    train_params.append(train_csv[i][2:])\n",
    "    train_labels.append(train_csv[i][1])\n",
    "    \n",
    "    test_lcs.append(np.load(f'dataset/test/{lc_id}.npy'))\n",
    "    test_params.append(test_csv[i][2:])\n",
    "    test_labels.append(test_csv[i][1])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, send the data to our dataset object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deep_lc import LocalDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_set = LocalDataset(train_lcs, train_params, train_labels)\n",
    "test_set = LocalDataset(test_lcs, test_params, test_labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Modify hyperparameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deep_lc.config import PROPOSAL_NUM, BATCH_SIZE, LR, LABELS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Change above imported hyperparameters to your own hyperparameters. For example, you can change the number of epochs, the learning rate, the batch size, and labels etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hyper_parameters={\n",
    "    \"batch_size\": BATCH_SIZE,\n",
    "    \"lr\": LR/10, # lower learning rate for fine-tuning\n",
    "    \"weight_decay\": 0,\n",
    "    \"labels\": LABELS,\n",
    "    \"proposal_num\": PROPOSAL_NUM,\n",
    "},\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load baseline model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here I only show how to fine-tune the combined model, since it's the most powerful model. For other models, you may need to implement the fine-tuning interface by yourself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dict = torch.load(\"combined_17.pth\", map_location=device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tune the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are lots of strageties for fine-tuning. Here we use the simplest one: fine-tune the whole model with a small learning rate. In pratices, you may need to try different learning rate, different optimizer, different loss function, and different fine-tuning strageties for different datasets. For example, smaller but similar datasets may only need to fine-tune the last few layers while larger or more different datasets may need to fine-tune the whole model with a smaller learning rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from deep_lc import finetune\n",
    "\n",
    "finetune(\n",
    "    training_set=training_set,\n",
    "    test_set=test_set,\n",
    "    hyper_params=hyper_parameters,\n",
    "    base_model=model_dict,\n",
    "    save_dir=\".\",\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calibration\n",
    "\n",
    "We provide a conformal prediction algorithm to calibrate the model"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}