online_testing/model_postprocessing/v5_nn_wrapper.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3a186fd1-49ff-41d8-85d8-91e52ea3c4f3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/IPython/core/magics/osm.py:417: UserWarning: using dhist requires you to install the `pickleshare` library.\n",
      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/global/u2/z/zeyuanhu/public_codes/ClimSim\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-07-23 03:40:28.985421: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-07-23 03:40:28.985454: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-07-23 03:40:28.986968: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-07-23 03:40:28.994981: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "%cd /global/u2/z/zeyuanhu/public_codes/ClimSim/\n",
    "from climsim_utils.data_utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9443bcdf-de3a-456e-821a-7cf2e0dc7ab8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import torch.optim as optim\n",
    "import torch.nn as nn\n",
    "import modulus"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "027f9d6a-056a-4126-b2f1-df0339c3470c",
   "metadata": {},
   "source": [
    "# Create a wrapper model to include normalization and de-normalization inside model's forward method"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66f9d868-967e-4f7e-ae1c-00b172b7a942",
   "metadata": {},
   "source": [
    "We define below a new class \"NewModel\" that takes the trained U-Net model (v5, i.e., applied microphysics constraints) and include all the preprocessing and post-processing steps inside the forward method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ae7e628f-e7ec-4534-9c8f-d0ac0acc2ff9",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/global/u2/z/zeyuanhu/public_codes/ClimSim/online_testing/baseline_models/Unet_v5/training\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/IPython/core/magics/osm.py:417: UserWarning: using dhist requires you to install the `pickleshare` library.\n",
      "  self.shell.db['dhist'] = compress_dhist(dhist)[-100:]\n"
     ]
    }
   ],
   "source": [
    "%cd /global/u2/z/zeyuanhu/public_codes/ClimSim/online_testing/baseline_models/Unet_v5/training\n",
    "from climsim_unet import ClimsimUnet\n",
    "import climsim_unet as climsim_unet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4ac47ff5-7cdd-4d2f-8c20-3e0c3d5d4e91",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NewModel(nn.Module):\n",
    "    def __init__(self, original_model, input_sub, input_div, out_scale, lbd_qn):\n",
    "        super(NewModel, self).__init__()\n",
    "        self.original_model = original_model\n",
    "        self.input_sub = torch.tensor(input_sub, dtype=torch.float32)\n",
    "        self.input_div = torch.tensor(input_div, dtype=torch.float32)\n",
    "        self.out_scale = torch.tensor(out_scale, dtype=torch.float32)\n",
    "        self.lbd_qn = torch.tensor(lbd_qn, dtype=torch.float32)\n",
    "    \n",
    "    def apply_temperature_rules(self, T):\n",
    "        # Create an output tensor, initialized to zero\n",
    "        output = torch.zeros_like(T)\n",
    "\n",
    "        # Apply the linear transition within the range 253.16 to 273.16\n",
    "        mask = (T >= 253.16) & (T <= 273.16)\n",
    "        output[mask] = (T[mask] - 253.16) / 20.0  # 20.0 is the range (273.16 - 253.16)\n",
    "\n",
    "        # Values where T > 273.16 set to 1\n",
    "        output[T > 273.16] = 1\n",
    "\n",
    "        # Values where T < 253.16 are already set to 0 by the initialization\n",
    "        return output\n",
    "\n",
    "    def preprocessing(self, x):\n",
    "        \n",
    "        # convert v4 input array to v5 input array:\n",
    "        xout = x\n",
    "        xout_new = torch.zeros((xout.shape[0], 1405), dtype=xout.dtype)\n",
    "        xout_new[:,0:120] = xout[:,0:120]\n",
    "        xout_new[:,120:180] = xout[:,120:180] + xout[:,180:240]\n",
    "        xout_new[:,180:240] = self.apply_temperature_rules(xout[:,0:60])\n",
    "        xout_new[:,240:840] = xout[:,240:840] #60*14\n",
    "        xout_new[:,840:900] = xout[:,840:900]+ xout[:,900:960] #dqc+dqi\n",
    "        xout_new[:,900:1080] = xout[:,960:1140]\n",
    "        xout_new[:,1080:1140] = xout[:,1140:1200]+ xout[:,1200:1260]\n",
    "        xout_new[:,1140:1405] = xout[:,1260:1525]\n",
    "        x = xout_new\n",
    "        \n",
    "        #do input normalization\n",
    "        x[:,120:180] = 1 - torch.exp(-x[:,120:180] * self.lbd_qn)\n",
    "        x= (x - self.input_sub) / self.input_div\n",
    "        x = torch.where(torch.isnan(x), torch.tensor(0.0, device=x.device), x)\n",
    "        x = torch.where(torch.isinf(x), torch.tensor(0.0, device=x.device), x)\n",
    "        \n",
    "        #prune top 15 levels in qn input\n",
    "        x[:,120:120+15] = 0\n",
    "        #clip rh input\n",
    "        x[:, 60:120] = torch.clamp(x[:, 60:120], 0, 1.2)\n",
    "        return x\n",
    "\n",
    "    def postprocessing(self, x):\n",
    "        x[:,60:75] = 0\n",
    "        x[:,120:135] = 0\n",
    "        x[:,180:195] = 0\n",
    "        x[:,240:255] = 0\n",
    "        x = x/self.out_scale\n",
    "        return x\n",
    "\n",
    "    def forward(self, x):\n",
    "        t_before = x[:,0:60].clone()\n",
    "        qc_before = x[:,120:180].clone()\n",
    "        qi_before = x[:,180:240].clone()\n",
    "        qn_before = qc_before + qi_before\n",
    "        \n",
    "        x = self.preprocessing(x)\n",
    "        x = self.original_model(x)\n",
    "        x = self.postprocessing(x)\n",
    "        \n",
    "        t_new = t_before + x[:,0:60]*1200.\n",
    "        qn_new = qn_before + x[:,120:180]*1200.\n",
    "        liq_frac = self.apply_temperature_rules(t_new)\n",
    "        qc_new = liq_frac*qn_new\n",
    "        qi_new = (1-liq_frac)*qn_new\n",
    "        xout = torch.zeros((x.shape[0],368))\n",
    "        xout[:,0:120] = x[:,0:120]\n",
    "        xout[:,240:] = x[:,180:]\n",
    "        xout[:,120:180] = (qc_new - qc_before)/1200.\n",
    "        xout[:,180:240] = (qi_new - qi_before)/1200.\n",
    "    \n",
    "        return xout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "58c18e23-cea9-4d72-9c06-8e64e5c1e02b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def save_wrapper(casename):\n",
    "    # casename = 'v5_noclassifier_huber_1y_noaggressive'\n",
    "    f_torch_model = f'/global/homes/z/zeyuanhu/scratch/hugging/E3SM-MMF_ne4/saved_models/{casename}/model.mdlus'\n",
    "    f_inp_sub     = f'/global/homes/z/zeyuanhu/scratch/hugging/E3SM-MMF_ne4/saved_models/{casename}/inp_sub.txt'\n",
    "    f_inp_div     = f'/global/homes/z/zeyuanhu/scratch/hugging/E3SM-MMF_ne4/saved_models/{casename}/inp_div.txt'\n",
    "    f_out_scale   = f'/global/homes/z/zeyuanhu/scratch/hugging/E3SM-MMF_ne4/saved_models/{casename}/out_scale.txt'\n",
    "    f_qn_lbd = '/global/u2/z/zeyuanhu/nvidia_codes/Climsim_private/preprocessing/normalizations/inputs/qn_exp_lambda_large.txt'\n",
    "    lbd_qn = np.loadtxt(f_qn_lbd, delimiter=',')\n",
    "    input_sub = np.loadtxt(f_inp_sub, delimiter=',')\n",
    "    input_div = np.loadtxt(f_inp_div, delimiter=',')\n",
    "    out_scale = np.loadtxt(f_out_scale, delimiter=',')\n",
    "    model_inf = modulus.Module.from_checkpoint(f_torch_model).to('cpu')\n",
    "\n",
    "    new_model = NewModel(model_inf, input_sub, input_div, out_scale, lbd_qc, lbd_qi)\n",
    "\n",
    "    NewModel.device = \"cpu\"\n",
    "    device = torch.device(\"cpu\")\n",
    "    scripted_model = torch.jit.script(new_model)\n",
    "    scripted_model = scripted_model.eval()\n",
    "    save_file_torch = os.path.join('/global/homes/z/zeyuanhu/scratch/hugging/E3SM-MMF_ne4/saved_models_wrapper_tmp/', f'{casename}.pt')\n",
    "    scripted_model.save(save_file_torch)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "33d6ba33-2c89-49fa-acde-0445b6de85ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_wrapper('v5_unet_nonaggressive_cliprh_huber_rop2_r2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "77e372c9-5151-4f82-805f-4cd358b7d762",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_wrapper('v5_unet_nonaggressive_cliprh_huber')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "316a370d-b9b9-405f-a74a-2aa84f758ee2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "save_wrapper('v5_unet_nonaggressive_cliprh_mae')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d2a89f5-16d2-4138-a411-a571d9535130",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MyEnvironment",
   "language": "python",
   "name": "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}