{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "98e6e23d-5ca7-4706-b1d9-dd57b54888ef",
   "metadata": {},
   "source": [
    "# Data preprocessing for further calculations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5324ceb9-24e7-454b-87b9-ba9a717078ae",
   "metadata": {},
   "source": [
    "### Import libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7b2a7f44-b0cb-4471-a0c6-e56da23caf86",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "36b9f49e-32e6-4544-a9d3-f6a8ba49d867",
   "metadata": {},
   "outputs": [],
   "source": [
    "# also available at https://eee.ipfran.ru/files/seasonal-variation-2024/\n",
    "# attention: the files are very large (~ 350 GB totally)\n",
    "src_path = \"../shared_files/eee_public_files/seasonal-variation-2024/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e16ee8e-f3b0-4251-9691-19d7dfd4aff7",
   "metadata": {},
   "source": [
    "### Preprocessing WRF T2m data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "78a4350c-59fb-479a-b7cd-e2bf9b996d36",
   "metadata": {},
   "outputs": [],
   "source": [
    "# available numbers of simulated days for analysis\n",
    "wrf_N_days = 4992\n",
    "inmcm_N_days = 3650"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "53cb9cc3-0e56-4da4-920b-2f071a0846fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dates corresponding to the indices (0 axis) of the data arrays\n",
    "# note: for WRF dates correspond to real dates\n",
    "\n",
    "wrf_dt_indicies = np.array(\n",
    "    [dt.date(1980, 1, 1) + dt.timedelta(i * 3) for i in range(wrf_N_days)]\n",
    ")\n",
    "inmcm_dt_indicies = np.array(\n",
    "    [dt.date(2022, 1, 1) + dt.timedelta(i % 365) for i in range(inmcm_N_days)]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "id": "0e4d0268-208f-45c7-bb0a-b5e8d3c7c1f7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4992, 180, 360)"
      ]
     },
     "execution_count": 147,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "wrf_T2_data = np.load(f\"{src_path}/T2-MAP-FULL.npy\")[:wrf_N_days]\n",
    "wrf_T2_data.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "id": "ec569ffd-93c2-4490-8ba1-69af4fab8f23",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mean surface air temperature values for different latitudes and months\n",
    "wrf_mon_T2 = np.zeros((180, 12))\n",
    "\n",
    "for month_idx in range(12):\n",
    "    monthly_indicies = [\n",
    "        i for i, date in enumerate(wrf_dt_indicies) if date.month == month_idx + 1\n",
    "    ]  # indicies of days available for `month_idx+1` month\n",
    "\n",
    "    wrf_mon_T2[:, month_idx] = wrf_mean_T2[monthly_indicies].mean(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "08302aa4-cb14-47f9-8216-a9db06ae53ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(f\"./data/WRF/WRF_T2_LATxMON.npy\",wrf_mon_T2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d4f093-a420-42c7-b885-a8409d9d8ee4",
   "metadata": {},
   "source": [
    "### INMCM and WRF IP: classic parametrization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "94a603c3-982d-4c78-be1c-bb6c74b86b5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "wrf_daily_latitudal_ip = {}\n",
    "inmcm_daily_latitudal_ip = {}\n",
    "\n",
    "wrf_hourly_total_ip = {}\n",
    "inmcm_hourly_total_ip = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "d8e43c4f-59af-483c-8979-535c696abb4e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "800\n",
      "1000\n",
      "1200\n"
     ]
    }
   ],
   "source": [
    "for cape_thres in [800, 1000, 1200]: # J/kg\n",
    "    print(cape_thres)\n",
    "\n",
    "    # grid cell contributions to the IP (not normalised) with the shape\n",
    "    # (number of days, number of hours, number of latitudes, number of longitudes)\n",
    "    wrf_raw_ip_data = np.load(f\"{src_path}/WRF-IP-MAP-{cape_thres}.npy\")[:wrf_N_days]\n",
    "    wrf_raw_ip_data = wrf_raw_ip_data[:, :24, :, :]\n",
    "    wrf_raw_ip_data /= (1/240e3) * wrf_raw_ip_data.sum(axis=(-2,-1)).mean()\n",
    "\n",
    "    wrf_daily_latitudal_ip[cape_thres] = wrf_raw_ip_data.mean(axis=1).sum(axis=-1)\n",
    "    wrf_hourly_total_ip[cape_thres] = wrf_raw_ip_data.sum(axis=(-2, -1))\n",
    "\n",
    "    inmcm_raw_ip_data = np.load(f\"{src_path}/INMCM-IP-MAP-{cape_thres}.npy\").reshape((inmcm_N_days, 24, 120, 180))\n",
    "    inmcm_raw_ip_data /= (1/240e3) * inmcm_raw_ip_data.sum(axis=(-2,-1)).mean()\n",
    "    \n",
    "    inmcm_daily_latitudal_ip[cape_thres] = inmcm_raw_ip_data.mean(axis=1).sum(axis=-1)\n",
    "    inmcm_hourly_total_ip[cape_thres] = inmcm_raw_ip_data.sum(axis=(-2, -1))\n",
    "\n",
    "    del wrf_raw_ip_data\n",
    "    del inmcm_raw_ip_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "id": "eb28cbc7-eb0a-49be-8cc1-734bba1d06f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "for cape_thres in [800, 1000, 1200]:  # J/kg\n",
    "    np.save(\n",
    "        f\"./data/INMCM/INMCM_HOURLY_TOTAL_IP_{cape_thres}.npy\",\n",
    "        inmcm_hourly_total_ip[cape_thres],\n",
    "    )\n",
    "    np.save(\n",
    "        f\"./data/WRF/WRF_HOURLY_TOTAL_IP_{cape_thres}.npy\",\n",
    "        wrf_hourly_total_ip[cape_thres],\n",
    "    )\n",
    "\n",
    "    wrf_data_LATxMON = np.zeros((180, 12))\n",
    "    inmcm_data_LATxMON = np.zeros((120, 12))\n",
    "\n",
    "    for month_idx in range(12):\n",
    "        monthly_indicies = [\n",
    "            i for i, date in enumerate(wrf_dt_indicies) if date.month == month_idx + 1\n",
    "        ]  # indicies of days available for `month_idx+1` month\n",
    "\n",
    "        wrf_data_MONxLAT[:, month_idx] = wrf_daily_latitudal_ip[cape_thres][monthly_indicies].mean(\n",
    "            axis=0\n",
    "        )\n",
    "    np.save(\n",
    "        f\"./data/WRF/WRF_IP_{cape_thres}_LATxMON.npy\",\n",
    "        wrf_data_MONxLAT,\n",
    "    )\n",
    "\n",
    "    for month_idx in range(12):\n",
    "        monthly_indicies = [\n",
    "            i for i, date in enumerate(inmcm_dt_indicies) if date.month == month_idx + 1\n",
    "        ]  # indicies of days available for `month_idx+1` month\n",
    "\n",
    "        inmcm_data_LATxMON[:, month_idx] = inmcm_daily_latitudal_ip[cape_thres][\n",
    "            monthly_indicies\n",
    "        ].mean(axis=0)\n",
    "    np.save(\n",
    "        f\"./data/INMCM/INMCM_IP_{cape_thres}_LATxMON.npy\",\n",
    "        inmcm_data_LATxMON,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91bc6d7a-393c-4078-9a6d-1955393d55f5",
   "metadata": {},
   "source": [
    "### WRF IP: parametrization based on T2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b6f987e-ba3c-4371-af7b-c9857a7d33d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "wrf_raw_ip_data = np.load(f\"{src_path}/WRF-IP-MAP-500-T2-25.npy\")[:wrf_N_days]\n",
    "wrf_raw_ip_data = wrf_raw_ip_data[:, :24, :, :]\n",
    "wrf_raw_ip_data /= (1/240e3) * wrf_raw_ip_data.sum(axis=(-2,-1)).mean()\n",
    "\n",
    "wrf_daily_latitudal_ip = wrf_raw_ip_data.mean(axis=1).sum(axis=-1)\n",
    "wrf_hourly_total_ip = wrf_raw_ip_data.sum(axis=(-2, -1))\n",
    "\n",
    "np.save(\n",
    "    f\"./data/WRF/WRF_HOURLY_TOTAL_IP_500_T2_25.npy\",\n",
    "    wrf_hourly_total_ip,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "17036c19-95f8-40df-a6c9-f8a23cf426f6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4992\n"
     ]
    }
   ],
   "source": [
    "wrf_data_LATxMON = np.zeros((180, 12))\n",
    "\n",
    "for month_idx in range(12):\n",
    "    monthly_indicies = [\n",
    "        i for i, date in enumerate(wrf_dt_indicies) \n",
    "        if date.month == month_idx + 1\n",
    "    ]\n",
    "\n",
    "    wrf_data_LATxMON[:, month_idx] = \\\n",
    "        wrf_daily_latitudal_ip[monthly_indicies].mean(axis=0)\n",
    "\n",
    "np.save(f\"./data/WRF/WRF_IP_500_T2_25_LATxMON.npy\", wrf_data_LATxMON)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e24297fc-cf81-4ea7-9a80-cdcaf277474a",
   "metadata": {},
   "source": [
    "### Saving number of days (used for monthly mean) for each month"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "6dfdae16-cb07-4580-97ba-414b5b2c1f2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "wrf_days = np.array([len([\n",
    "        i for i, date in enumerate(wrf_dt_indicies) \n",
    "        if date.month == m + 1\n",
    "    ]) for m in range(12)])\n",
    "\n",
    "np.save(f\"./data/WRF/WRF_NUMDAYS_MON.npy\", wrf_days)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "762dd08f-7d1e-4b68-b272-e7aa9b26d9f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "inm_days = np.array([len([\n",
    "        i for i, date in enumerate(inmcm_dt_indicies) \n",
    "        if date.month == m + 1\n",
    "    ]) for m in range(12)])\n",
    "\n",
    "np.save(f\"./data/INMCM/INMCM_NUMDAYS_MON.npy\", inm_days)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "78afb8e1-3975-4e19-993f-1b0c74c3e9bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for average over months use\n",
    "# `(wrf_data_LATxMON[:, :].sum(axis=0)*days).sum()/days.sum()`\n",
    "# unstead\n",
    "# `wrf_data_LATxMON[:, :].sum(axis=0).mean()`\n",
    "# because\n",
    "# ((a1+a2+a3)/3 + (b1+b2)/2)/2 != (a1+a2+a3+b1+b2)/5"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}