{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configure the model architecture and training parameters, and start training\n",
    "\n",
    "This section shows how to configure the model architectures and training hyperparameters using in the UNAGI. \n",
    "\n",
    "The model architecutures includes \n",
    "-   The number of fully connected layers in the VAE of UNAGI.\n",
    "-   The number of neurons for each fully connected layer.\n",
    "\n",
    "The training hyperparameters includes\n",
    "-   The learning rate of graph VAE.\n",
    "-   The learning rate of the discriminator.\n",
    "-   The batch size.\n",
    "-   The number of iterations to run the UNAGI tool.\n",
    "-   The number of epochs for the inital iteration.\n",
    "-   The number of eacphs for the iteratively training iterations.\n",
    "-   The gene expression distribution of the input single-cell dataset.\n",
    "-   Using GPU or not.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Register the data used for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "from UNAGI import UNAGI\n",
    "unagi = UNAGI()\n",
    "unagi.setup_data('PATH_TO_YOUR_DATA',total_stage=4,stage_key='stage') # total_stage is the number of time-points of your data, stage_key is the column name of the stage in your data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure the model architecture."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_dim = 128 # neurons of the hidden dimension of the VAE\n",
    "latent_dim = 32 # neurons of the latent dimension of the VAE\n",
    "graph_dim = 512 # neurons of the output of graph convolutional network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure the training hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task_name = 'example'\n",
    "gene_expression_distribution = 'ziln' # using zero-inflated lognormal distribution\n",
    "GPU=True # enable GPU\n",
    "device='cuda:0' # GPU device\n",
    "training_epochs_initial = 10 # training epochs for initial stage\n",
    "training_epochs_iterative = 8 # training epochs for iterative stages\n",
    "maximum_iteration = 3 # maximum number of iterations\n",
    "batch_size = 256 # batch size\n",
    "learning_rate_vae = 0.0001 # learning rate for VAE\n",
    "learning_rate_discriminator = 0.0005 # learning rate for discriminator (GAN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unagi.setup_training(task=task_name,dist=gene_expression_distribution,hidden_dim=hidden_dim,latent_dim=latent_dim,graph_dim=graph_dim,GPU=GPU,device=device,epoch_initial=training_epochs_initial,epoch_iter=training_epochs_iterative,max_iter=maximum_iteration,BATCHSIZE=batch_size,lr=learning_rate_vae,lr_dis=learning_rate_discriminator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start training\n",
    "\n",
    "Running the UNAGI tool requires to specify the directory of iDREM software."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "iDREM_Path = 'iDREM_Dir' # path to iDREM directory\n",
    "unagi.run_UNAGI(iDREM_Path)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}