tutorials/T04-ModelTraining.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5686b0fa-1851-4ee7-9028-1f52dc6ef81a",
   "metadata": {},
   "source": [
    "# Tutorial 04 - Model Training\n",
    "\n",
    "This notebook shows you how to train models of various methods used in this study. These models (methods) include:\n",
    "\n",
    "- Dedicated supervised\n",
    "- Parameterised supervised\n",
    "- Ideal weakly\n",
    "- Semi weakly (PAWS)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c9de9cc-9de7-4ded-ba59-86e3e828e215",
   "metadata": {},
   "source": [
    "The easiest way will be to directly use the paws CLI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e30ce72a-02c0-44be-93b0-d39894228d24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage: paws train_dedicated_supervised [OPTIONS]\n",
      "\n",
      "  Train dedicated supervised models.\n",
      "\n",
      "Options:\n",
      "  -m, --mass-point TEXT           Signal mass point to use for training in the\n",
      "                                  form \"m1:m2\".  [required]\n",
      "  --high-level / --low-level      Whether to do training with low-evel or\n",
      "                                  high-level features.  [default: high-level]\n",
      "  --decay-modes [qq|qqq|qq,qqq]   Which decay mode should the signal undergo\n",
      "                                  (qq or qqq).Use \"qq,qqq\" to include both\n",
      "                                  decay modes.  [default: qq,qqq]\n",
      "  --variables TEXT                Select certain high-level jet features to\n",
      "                                  include in the trainingby the indices they\n",
      "                                  appear in the feature vector. For\n",
      "                                  example,\"3,5,6\" means select the 4th, 6th\n",
      "                                  and 7th feature from the jetfeature vector\n",
      "                                  to be used in the training.\n",
      "  --noise INTEGER                 Number of noise dimension to add to the\n",
      "                                  train features.  [default: 0]\n",
      "  --dataset-index-path TEXT       Path to the dataset split configuration\n",
      "                                  file. It determines theshard indices for the\n",
      "                                  train, validation, and test datasets in\n",
      "                                  eachrandom realization of data. If None, a\n",
      "                                  default configuration will be created.\n",
      "  -i, --split-index INTEGER       Index for dataset split.  [default: 0]\n",
      "  --seed INTEGER                  The default seed used for all random\n",
      "                                  processes.  [default: 2023]\n",
      "  --batchsize INTEGER             Batch size for training.\n",
      "  --interrupt-freq INTEGER        Frequency of training interruption for early\n",
      "                                  stopping.\n",
      "  --cache-dataset / --no-cache-dataset\n",
      "                                  Whether to cache the dataset during\n",
      "                                  training.\n",
      "  -d, --datadir TEXT              Input directory where the tfrecord datasets\n",
      "                                  are stored  [default: datasets]\n",
      "  -o, --outdir TEXT               Base output directory  [default: outputs]\n",
      "  --version TEXT                  Version of the model.  [default: v1]\n",
      "  --cache / --no-cache            Whether to cache the results.  [default:\n",
      "                                  cache]\n",
      "  --multi-gpu / --single-gpu      Whether to enable multi-GPU training.\n",
      "                                  [default: multi-gpu]\n",
      "  -v, --verbosity TEXT            Verbosity level (\"DEBUG\", \"INFO\", \"WARNING\"\n",
      "                                  or \"ERROR\").  [default: INFO]\n",
      "  --help                          Show this message and exit.\n"
     ]
    }
   ],
   "source": [
    "# dedicated supervised model training\n",
    "!paws train_dedicated_supervised --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef5797cf-9bc6-4b07-a682-a3448619c522",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example command \n",
    "!paws train_dedicated_supervised -d \"datasets\" -o \"outputs\" --mass-point 300:300 --decay-modes qq,qqq \\\n",
    "--variables 3,5,6 --split-index 0 --version v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4195510c-afcf-4223-9a7d-cf53717eb1c9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage: paws train_param_supervised [OPTIONS]\n",
      "\n",
      "  Train parameterised supervised models.\n",
      "\n",
      "Options:\n",
      "  --high-level / --low-level      Whether to do training with low-evel or\n",
      "                                  high-level features.  [default: high-level]\n",
      "  --decay-modes [qq|qqq|qq,qqq]   Which decay mode should the signal undergo\n",
      "                                  (qq or qqq).Use \"qq,qqq\" to include both\n",
      "                                  decay modes.  [default: qq]\n",
      "  --variables TEXT                Select certain high-level jet features to\n",
      "                                  include in the trainingby the indices they\n",
      "                                  appear in the feature vector. For\n",
      "                                  example,\"3,5,6\" means select the 4th, 6th\n",
      "                                  and 7th feature from the jetfeature vector\n",
      "                                  to be used in the training.\n",
      "  --noise INTEGER                 Number of noise dimension to add to the\n",
      "                                  train features.\n",
      "  --exclude-masses TEXT           Mass points to exclude (mass point separated\n",
      "                                  by commas, mass values separated by colon).\n",
      "  --include-masses TEXT           Mass points to include (mass point separated\n",
      "                                  by commas, mass values separated by colon).\n",
      "  --dataset-index-path TEXT       Path to the dataset split configuration\n",
      "                                  file. It determines theshard indices for the\n",
      "                                  train, validation, and test datasets in\n",
      "                                  eachrandom realization of data. If None, a\n",
      "                                  default configuration will be created.\n",
      "  -i, --split-index INTEGER       Index for dataset split.  [default: 0]\n",
      "  --seed INTEGER                  The default seed used for all random\n",
      "                                  processes.  [default: 2023]\n",
      "  --batchsize INTEGER             Batch size for training.\n",
      "  --interrupt-freq INTEGER        Frequency of training interruption for early\n",
      "                                  stopping.\n",
      "  --cache-dataset / --no-cache-dataset\n",
      "                                  Whether to cache the dataset during\n",
      "                                  training.\n",
      "  -d, --datadir TEXT              Input directory where the tfrecord datasets\n",
      "                                  are stored  [default: datasets]\n",
      "  -o, --outdir TEXT               Base output directory  [default: outputs]\n",
      "  --version TEXT                  Version of the model.  [default: v1]\n",
      "  --cache / --no-cache            Whether to cache the results.  [default:\n",
      "                                  cache]\n",
      "  --multi-gpu / --single-gpu      Whether to enable multi-GPU training.\n",
      "                                  [default: multi-gpu]\n",
      "  -v, --verbosity TEXT            Verbosity level (\"DEBUG\", \"INFO\", \"WARNING\"\n",
      "                                  or \"ERROR\").  [default: INFO]\n",
      "  --help                          Show this message and exit.\n"
     ]
    }
   ],
   "source": [
    "# parameterised supervised model training\n",
    "!paws train_param_supervised --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5508b6a7-6a24-4a82-b189-adfe6ca43597",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example command \n",
    "!paws train_param_supervised -d \"datasets\" -o \"outputs\" --decay-modes qq --variables 3,5,6 --split-index 0 --version v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cf065ed7-574c-4b03-819a-f87d264be90f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage: paws train_ideal_weakly [OPTIONS]\n",
      "\n",
      "  Train ideal weakly models.\n",
      "\n",
      "Options:\n",
      "  -m, --mass-point TEXT           Signal mass point to use for training in the\n",
      "                                  form \"m1:m2\".  [required]\n",
      "  --mu FLOAT                      Signal fraction in the training and\n",
      "                                  validation dataset.  [required]\n",
      "  --alpha FLOAT                   Signal branching fraction in the training\n",
      "                                  and validation dataset. Ignored when only\n",
      "                                  one signal decay mode is considered.\n",
      "  --high-level / --low-level      Whether to do training with low-evel or\n",
      "                                  high-level features.  [default: high-level]\n",
      "  --decay-modes [qq|qqq|qq,qqq]   Which decay mode should the signal undergo\n",
      "                                  (qq or qqq).Use \"qq,qqq\" to include both\n",
      "                                  decay modes.  [default: qq,qqq]\n",
      "  --variables TEXT                Select certain high-level jet features to\n",
      "                                  include in the trainingby the indices they\n",
      "                                  appear in the feature vector. For\n",
      "                                  example,\"3,5,6\" means select the 4th, 6th\n",
      "                                  and 7th feature from the jetfeature vector\n",
      "                                  to be used in the training.\n",
      "  --noise INTEGER                 Number of noise dimension to add to the\n",
      "                                  train features.\n",
      "  --dataset-index-path TEXT       Path to the dataset split configuration\n",
      "                                  file. It determines theshard indices for the\n",
      "                                  train, validation, and test datasets in\n",
      "                                  eachrandom realization of data. If None, a\n",
      "                                  default configuration will be created.\n",
      "  -i, --split-index INTEGER       Index for dataset split.  [default: 0]\n",
      "  --num-trials INTEGER            Number of trials (random model\n",
      "                                  initialization) to run.  [default: 10]\n",
      "  --seed INTEGER                  The default seed used for all random\n",
      "                                  processes.  [default: 2023]\n",
      "  --batchsize INTEGER             Batch size for training.\n",
      "  --interrupt-freq INTEGER        Frequency of training interruption for early\n",
      "                                  stopping.\n",
      "  --cache-dataset / --no-cache-dataset\n",
      "                                  Whether to cache the dataset during\n",
      "                                  training.\n",
      "  -d, --datadir TEXT              Input directory where the tfrecord datasets\n",
      "                                  are stored  [default: datasets]\n",
      "  -o, --outdir TEXT               Base output directory  [default: outputs]\n",
      "  --version TEXT                  Version of the model.  [default: v1]\n",
      "  --cache / --no-cache            Whether to cache the results.  [default:\n",
      "                                  cache]\n",
      "  --multi-gpu / --single-gpu      Whether to enable multi-GPU training.\n",
      "                                  [default: multi-gpu]\n",
      "  -v, --verbosity TEXT            Verbosity level (\"DEBUG\", \"INFO\", \"WARNING\"\n",
      "                                  or \"ERROR\").  [default: INFO]\n",
      "  --help                          Show this message and exit.\n"
     ]
    }
   ],
   "source": [
    "# ideal weakly model training\n",
    "!paws train_ideal_weakly --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cf773ef-53a1-468e-8d5f-42e59b30ddd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example command \n",
    "!paws train_ideal_weakly -d \"datasets\" -o \"outputs\" --mass-point 300:300 --decay-modes qq,qqq \\\n",
    "--variables 3,5,6 --mu 0.01 --alpha 0.5 --split-index 0 --version v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fd795a73-6fae-4fe2-8ac9-a8abb8034b8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Usage: paws train_semi_weakly [OPTIONS]\n",
      "\n",
      "  Train semi-weakly (PAWS) models.\n",
      "\n",
      "Options:\n",
      "  -m, --mass-point TEXT           Signal mass point to use for training in the\n",
      "                                  form \"m1:m2\".  [required]\n",
      "  --mu FLOAT                      Signal fraction in the training and\n",
      "                                  validation dataset.  [required]\n",
      "  --alpha FLOAT                   Signal branching fraction in the training\n",
      "                                  and validation dataset. Ignored when only\n",
      "                                  one signal decay mode is considered.\n",
      "  --high-level / --low-level      Whether to do training with low-evel or\n",
      "                                  high-level features.  [default: high-level]\n",
      "  --decay-modes [qq|qqq|qq,qqq]   Which decay mode should the signal undergo\n",
      "                                  (qq or qqq).Use \"qq,qqq\" to include both\n",
      "                                  decay modes.  [default: qq,qqq]\n",
      "  --variables TEXT                Select certain high-level jet features to\n",
      "                                  include in the trainingby the indices they\n",
      "                                  appear in the feature vector. For\n",
      "                                  example,\"3,5,6\" means select the 4th, 6th\n",
      "                                  and 7th feature from the jetfeature vector\n",
      "                                  to be used in the training.\n",
      "  --noise INTEGER                 Number of noise dimension to add to the\n",
      "                                  train features.\n",
      "  --dataset-index-path TEXT       Path to the dataset split configuration\n",
      "                                  file. It determines theshard indices for the\n",
      "                                  train, validation, and test datasets in\n",
      "                                  eachrandom realization of data. If None, a\n",
      "                                  default configuration will be created.\n",
      "  -i, --split-index INTEGER       Index for dataset split.  [default: 0]\n",
      "  --num-trials INTEGER            Number of trials (random model\n",
      "                                  initialization) to run.  [default: 10]\n",
      "  --fs-version TEXT               Version of the supervised model to use.\n",
      "                                  [default: v1]\n",
      "  --fs-version-2 TEXT             When signals of mixed decay modes are\n",
      "                                  considered, it corresponds to the version of\n",
      "                                  the three-prone supervised model. If None,\n",
      "                                  the same version as `fs_version` will be\n",
      "                                  used.\n",
      "  --retrain / --no-retrain        Retrain when m1 <-> m2 gives better\n",
      "                                  validation loss.  [default: no-retrain]\n",
      "  --seed INTEGER                  The default seed used for all random\n",
      "                                  processes.  [default: 2023]\n",
      "  --batchsize INTEGER             Batch size for training.\n",
      "  --interrupt-freq INTEGER        Frequency of training interruption for early\n",
      "                                  stopping.\n",
      "  --weight-clipping / --no-weight-clipping\n",
      "                                  Whether to apply weight clipping.  [default:\n",
      "                                  weight-clipping]\n",
      "  --cache-dataset / --no-cache-dataset\n",
      "                                  Whether to cache the dataset during\n",
      "                                  training.\n",
      "  -d, --datadir TEXT              Input directory where the tfrecord datasets\n",
      "                                  are stored  [default: datasets]\n",
      "  -o, --outdir TEXT               Base output directory  [default: outputs]\n",
      "  --version TEXT                  Version of the model.  [default: v1]\n",
      "  --cache / --no-cache            Whether to cache the results.  [default:\n",
      "                                  cache]\n",
      "  --multi-gpu / --single-gpu      Whether to enable multi-GPU training.\n",
      "                                  [default: multi-gpu]\n",
      "  -v, --verbosity TEXT            Verbosity level (\"DEBUG\", \"INFO\", \"WARNING\"\n",
      "                                  or \"ERROR\").  [default: INFO]\n",
      "  --help                          Show this message and exit.\n"
     ]
    }
   ],
   "source": [
    "# semi-weakly (PAWS) model training\n",
    "!paws train_semi_weakly --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b4d98b3-e89d-4766-8f95-63a611bf4696",
   "metadata": {},
   "outputs": [],
   "source": [
    "# example command \n",
    "!paws train_semi_weakly -d \"datasets\" -o \"outputs\" --mass-point 300:300 --decay-modes qq,qqq \\\n",
    "--variables 3,5,6 --mu 0.01 --alpha 0.5 --split-index 0 --version v1 --fs-version v1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f4c25df-0837-47ff-b1e3-f9e99e747061",
   "metadata": {},
   "source": [
    "Alternatively, you may use the paws API:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a487c25d-5255-4fe3-980a-b5735883e835",
   "metadata": {},
   "outputs": [],
   "source": [
    "from paws.components import ModelTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "98a8d580-4e15-4f03-8236-3021801c0862",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on function __init__ in module paws.components.model_trainer:\n",
      "\n",
      "__init__(self, model_type: Union[str, paws.settings.ModelType], model_options: Optional[Dict] = None, feature_level: str = 'high_level', decay_modes: str = 'qq,qqq', cache: bool = True, variables: Optional[str] = None, noise_dimension: Optional[int] = None, seed: int = 2023, split_index: int = 0, batchsize: Optional[int] = None, cache_dataset: Optional[bool] = None, version: str = 'v1', multi_gpu: bool = True, interrupt_freq: int = 0, datadir: str = 'datasets', outdir: str = 'outputs', index_path: Optional[str] = None, verbosity: str = 'INFO')\n",
      "    Initialize the ModelTrainer class.\n",
      "    \n",
      "    Parameters\n",
      "    ----------------------------------------------------\n",
      "    model_type : str or ModelType\n",
      "        The type of the model to train.\n",
      "    model_options : Dict, optional\n",
      "        Options specific to the model type.\n",
      "    feature_level : str or FeatureLevel, default \"high_level\"\n",
      "        Features to use for the training. It can be either\n",
      "        high-level (\"high_level\") or low-level (\"low_level\").\n",
      "    decay_modes : str, list of str or list of DecayMode, default \"qq,qqq\"\n",
      "        Decay modes of the signal to include in the training. Candidates are\n",
      "        two-prong decay (\"qq\") or three-prong decay (\"qqq\"). If it is a\n",
      "        string, it will be a comma delimited list of the decay modes.\n",
      "    cache : bool, default True\n",
      "        Whether to cache the results.\n",
      "    variables : str, optional\n",
      "        Select certain high-level jet features to include in the training\n",
      "        by the indices they appear in the feature vector. For example,\n",
      "        \"3,5,6\" means select the 4th, 6th and 7th feature from the jet\n",
      "        feature vector to be used in the training.\n",
      "    noise_dimension : int, optional\n",
      "        Number of noise dimension per jet to include in the training.\n",
      "    seed : int, optional, default 2023\n",
      "        The default seed used for all random processes.\n",
      "    split_index : int\n",
      "        Index for dataset split.\n",
      "    batchsize : int\n",
      "        Batch size for training.\n",
      "    cache_dataset : bool\n",
      "        Whether to cache the datasets.\n",
      "    version : str\n",
      "        Version of the model.\n",
      "    multi_gpu : bool, default True\n",
      "        Whether to enable multi-GPU training.\n",
      "    interrupt_freq : int, default 0\n",
      "        Frequency of training interruption for early stopping.\n",
      "    datadir : str, default \"datasets\"\n",
      "        Directory for datasets.\n",
      "    outdir : str, default \"outputs\"\n",
      "        Directory for outputs.\n",
      "    index_path : str, optional\n",
      "        Path to the dataset split configuration file. It determines the\n",
      "        shard indices for the train, validation, and test datasets in each\n",
      "        random realization of data.\n",
      "    verbosity : str, default \"INFO\"\n",
      "        Verbosity level (\"DEBUG\", \"INFO\", \"WARNING\" or \"ERROR\").\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(ModelTrainer.__init__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fad71a0b-e6f8-4831-92e3-eb353222369c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{<ModelType.DEDICATED_SUPERVISED: 0>: {'required': ['mass_point'],\n",
       "  'optional': []},\n",
       " <ModelType.PARAM_SUPERVISED: 1>: {'required': [],\n",
       "  'optional': ['include_masses', 'exclude_masses']},\n",
       " <ModelType.IDEAL_WEAKLY: 2>: {'required': ['mass_point', 'mu', 'alpha'],\n",
       "  'optional': ['num_trials']},\n",
       " <ModelType.SEMI_WEAKLY: 3>: {'required': ['mass_point', 'mu', 'alpha'],\n",
       "  'optional': ['num_trials',\n",
       "   'weight_clipping',\n",
       "   'retrain',\n",
       "   'fs_version',\n",
       "   'fs_version_2']}}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# options for various models:\n",
    "from paws.components.model_trainer import MODEL_OPTIONS\n",
    "MODEL_OPTIONS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c46a5a4-e3d7-4fbb-8102-be86813fab7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dedicated supervised model training\n",
    "model_options = {\n",
    "    'mass_point': [300, 300]\n",
    "}\n",
    "datadir = \"datasets\"\n",
    "outdir = \"outputs\"\n",
    "model_trainer = ModelTrainer(\"dedicated_supervised\", model_options=model_options, decay_modes='qq',\n",
    "                             variables=\"3,5,6\", version=\"v1\", datadir=datadir, outdir=outdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e10702e7-2ce5-4878-93c7-92a068f1c32d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "38c76440-199d-446d-b2d1-06f74c1aac8a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:CPU:0',)\n",
      "[INFO] Created MirroredStrategy for distributed training\n",
      "[INFO] Number of devices : 1\n",
      "[INFO]      aliad version : 0.1.0\n",
      "[INFO] tensorflow version : 2.15.0\n",
      "Sat Sep 28 01:26:37 2024       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 525.105.17   Driver Version: 525.105.17   CUDA Version: 12.2     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                               |                      |               MIG M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  NVIDIA A100-PCI...  On   | 00000000:C3:00.0 Off |                    0 |\n",
      "| N/A   34C    P0    36W / 250W |    834MiB / 40960MiB |      0%      Default |\n",
      "|                               |                      |             Disabled |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                                  |\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
      "|        ID   ID                                                   Usage      |\n",
      "|=============================================================================|\n",
      "|    0   N/A  N/A    578296      C   ...myenv_kernel37/bin/python      416MiB |\n",
      "|    0   N/A  N/A   1240275      C   ...myenv_kernel37/bin/python      416MiB |\n",
      "+-----------------------------------------------------------------------------+\n",
      "nvcc: NVIDIA (R) Cuda compiler driver\n",
      "Copyright (c) 2005-2023 NVIDIA Corporation\n",
      "Built on Tue_Jun_13_19:16:58_PDT_2023\n",
      "Cuda compilation tools, release 12.2, V12.2.91\n",
      "Build cuda_12.2.r12.2/compiler.32965470_0\n"
     ]
    }
   ],
   "source": [
    "# parameterised supervised model training\n",
    "model_options = {\n",
    "}\n",
    "datadir = \"/pscratch/sd/c/chlcheng/projects/paws/datasets\"\n",
    "outdir = \"/pscratch/sd/c/chlcheng/projects/paws/outputs\"\n",
    "model_trainer = ModelTrainer(\"param_supervised\", model_options=model_options, decay_modes='qq',\n",
    "                             variables=\"3,5,6\", version=\"v1\", datadir=datadir, outdir=outdir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48690650-176a-4cd8-8238-3d44a9ae1ecb",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trainer.train()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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": 5
}