{
 "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": [
    "### Importing libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": "markdown",
   "id": "6b2b903f-fa30-4e35-97e1-74bc0ee6b944",
   "metadata": {},
   "source": [
    "### Helper variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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": "code",
   "execution_count": 3,
   "id": "78a4350c-59fb-479a-b7cd-e2bf9b996d36",
   "metadata": {},
   "outputs": [],
   "source": [
    "# the number of simulated days used for analysis\n",
    "wrf_N_days = 4992\n",
    "inm_N_days = 3650"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53cb9cc3-0e56-4da4-920b-2f071a0846fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dates corresponding to the indices (0 axis) of the data arrays\n",
    "# note: in the case of the WRF the dates correspond to real dates\n",
    "# note: in the case of the INMCM the are 10 365-day years\n",
    "wrf_dt_indices = np.array(\n",
    "    [dt.date(1980, 1, 1) + dt.timedelta(i * 3) for i in range(wrf_N_days)]\n",
    ")\n",
    "inm_dt_indices = np.array(\n",
    "    [dt.date(2022, 1, 1) + dt.timedelta(i % 365) for i in range(inm_N_days)]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e16ee8e-f3b0-4251-9691-19d7dfd4aff7",
   "metadata": {},
   "source": [
    "### Preprocessing T2m data from the WRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ccff6ad9-dc83-440f-8642-93cfc150b03f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# air temperature values (in K) at the height of 2 m with the shape\n",
    "# (number of days, number of hours, number of latitudes, number of longitudes)\n",
    "# contains temperature values depending on (d, h, lat, lon)\n",
    "#   d (axis 0) is the number of a day starting with 0 and ending with 5113\n",
    "#              every third day is taken\n",
    "#              d = 0 corresponds to 1 Jan 1980\n",
    "#              d = 5113 corresponds to 30 Dec 2021\n",
    "#              d = 4991 corresponds to 29 Dec 2020\n",
    "#              (we will restrict our attention to 1980–2020)\n",
    "#   h (axis 1) is the hour of the day (an integer in [0, 25])\n",
    "#              the values corresponding to h = 0 and h = 24 are the same\n",
    "#              (we delete the 25th value)\n",
    "# lat (axis 2) describes the latitude (an integer in [0, 179])\n",
    "# lon (axis 3) describes the longitude (an integer in [0, 359])\n",
    "wrf_T2_data = np.load(f\"{src_path}/WRF-T2-MAP.npy\")[:wrf_N_days, :24]\n",
    "wrf_T2_data_DAYxLAT = wrf_T2_data.mean(axis=(1, 3))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a7148097-51d8-492b-80c4-7b4174bba4b4",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# initialising an array to store monthly averaged values\n",
    "# for different latitudes\n",
    "wrf_T2_LATxMON = np.zeros((180, 12))\n",
    "\n",
    "# iterating over month numbers (starting with 0)\n",
    "for month_idx in range(12):\n",
    "    # filtering indices by the month number\n",
    "    wrf_monthly_indices = [\n",
    "        i for i, date in enumerate(wrf_dt_indices)\n",
    "        if date.month == month_idx + 1\n",
    "    ]\n",
    "\n",
    "    # putting the values for the specific month into the array\n",
    "    # (also converting the values from K to °C)\n",
    "    wrf_T2_LATxMON[:, month_idx] = (\n",
    "        wrf_T2_data_DAYxLAT[wrf_monthly_indices].mean(axis=0) - 273.15\n",
    "    )\n",
    "\n",
    "np.save(\"./data/WRF/WRF_T2_LATxMON.npy\", wrf_T2_LATxMON)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d4f093-a420-42c7-b885-a8409d9d8ee4",
   "metadata": {},
   "source": [
    "### Preprocessing IP data from the INMCM and WRF: the usual parameterisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "052374df-a505-420a-aa69-2902ce5fc23b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dictionaries where the processed data are saved\n",
    "# dictionary keys represent CAPE threshold values\n",
    "\n",
    "# dictionaries to store diurnal average IP values summed over longitudes\n",
    "# the dimensions are (4992, 180) for the WRF and (3650, 120) for the INMCM\n",
    "wrf_daily_lat_ip = {}\n",
    "inm_daily_lat_ip = {}\n",
    "\n",
    "# dictionaries to store hourly IP values summed over longitudes and latitudes\n",
    "# the dimensions are (4992, 24) for the WRF and (3650, 24) for the INMCM\n",
    "wrf_hourly_total_ip = {}\n",
    "inm_hourly_total_ip = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d8e43c4f-59af-483c-8979-535c696abb4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# iterating over CAPE threshold values (in J/kg) used in modeling \n",
    "# for each threshold, there are corresponding model data sets\n",
    "for cape_thres in [800, 1000, 1200]:\n",
    "    # grid cell contributions to the IP (not normalised) with the shape\n",
    "    # (number of days, number of hours,\n",
    "    # number of latitudes, number of longitudes)\n",
    "    # simulated using the WRF with CAPE threshold = `cape_thres` J/kg\n",
    "    # contains values of contributions to the IP depending on (d, h, lat, lon)\n",
    "    #   d (axis 0) is the number of a day starting with 0 and ending with 5113\n",
    "    #              every third day is taken\n",
    "    #              d = 0 corresponds to 1 Jan 1980\n",
    "    #              d = 5113 corresponds to 30 Dec 2021\n",
    "    #              d = 4991 corresponds to 29 Dec 2020\n",
    "    #              (we will restrict our attention to 1980–2020)\n",
    "    #   h (axis 1) is the hour of the day (an integer in [0, 24])\n",
    "    #              the values corresponding to h = 0 and h = 24 are the same\n",
    "    #              (we delete the 25th value)\n",
    "    # lat (axis 2) describes the latitude (an integer in [0, 179])\n",
    "    # lon (axis 3) describes the longitude (an integer in [0, 359])\n",
    "    wrf_raw_ip_data = np.load(\n",
    "        f\"{src_path}/WRF-IP-MAP-{cape_thres}.npy\"\n",
    "    )[:wrf_N_days, :24]\n",
    "\n",
    "    # normalising contributions to the IP to the global mean of 240 kV\n",
    "    wrf_raw_ip_data /= (1/240e3) * wrf_raw_ip_data.sum(axis=(-2,-1)).mean()\n",
    "\n",
    "    # filling the dictionaries with averaged values\n",
    "    wrf_daily_lat_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",
    "    np.save(\n",
    "        f\"./data/WRF/WRF_HOURLY_TOTAL_IP_{cape_thres}.npy\",\n",
    "        wrf_hourly_total_ip[cape_thres]\n",
    "    )\n",
    "\n",
    "    # grid cell contributions to the IP (not normalised) reshaped to\n",
    "    # (number of days, number of hours,\n",
    "    # number of latitudes, number of longitudes)\n",
    "    # simulated using the INMCM with CAPE threshold = `cape_thres` J/kg\n",
    "    # contains values of contributions to the IP depending on (d, h, lat, lon)\n",
    "    #   d (axis 0) is the number of a day\n",
    "    #              (10 consecutive 365-day years have been simulated)\n",
    "    #   h (axis 1) is the hour of the day (an integer in [0, 23])\n",
    "    # lat (axis 2) describes the latitude (an integer in [0, 179])\n",
    "    # lon (axis 3) describes the longitude (an integer in [0, 359])\n",
    "    inm_raw_ip_data = np.load(\n",
    "        f\"{src_path}/INMCM-IP-MAP-{cape_thres}.npy\"\n",
    "    ).reshape((inm_N_days, 24, 120, 180))\n",
    "\n",
    "    # normalising contributions to the IP to the global mean of 240 kV\n",
    "    inm_raw_ip_data /= (1/240e3) * inm_raw_ip_data.sum(axis=(-2,-1)).mean()\n",
    "\n",
    "    # filling the dictionaries with averaged values\n",
    "    inm_daily_lat_ip[cape_thres] = inm_raw_ip_data.mean(axis=1).sum(axis=-1)\n",
    "    inm_hourly_total_ip[cape_thres] = inm_raw_ip_data.sum(axis=(-2, -1))\n",
    "\n",
    "    np.save(\n",
    "        f\"./data/INMCM/INMCM_HOURLY_TOTAL_IP_{cape_thres}.npy\",\n",
    "        inm_hourly_total_ip[cape_thres]\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "eb28cbc7-eb0a-49be-8cc1-734bba1d06f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# iterating over CAPE threshold values (in J/kg) used in modeling \n",
    "# for each threshold, there are corresponding model data sets\n",
    "for cape_thres in [800, 1000, 1200]:\n",
    "\n",
    "    # initialising arrays to store monthly averaged values\n",
    "    # for different latitudes\n",
    "    wrf_data_LATxMON = np.zeros((180, 12))\n",
    "    inm_data_LATxMON = np.zeros((120, 12))\n",
    "\n",
    "    # iterating over month numbers (starting with 0)\n",
    "    for month_idx in range(12):\n",
    "        # filtering indices by the month number\n",
    "        wrf_monthly_indices = [i for i, date in enumerate(wrf_dt_indices)\n",
    "                               if date.month == month_idx + 1]\n",
    "        inm_monthly_indices = [i for i, date in enumerate(inm_dt_indices)\n",
    "                               if date.month == month_idx + 1]\n",
    "\n",
    "        # putting the values for the specific month into the array\n",
    "        wrf_data_LATxMON[:, month_idx] = \\\n",
    "            wrf_daily_lat_ip[cape_thres][wrf_monthly_indices].mean(axis=0)\n",
    "        inm_data_LATxMON[:, month_idx] = \\\n",
    "            inm_daily_lat_ip[cape_thres][inm_monthly_indices].mean(axis=0)\n",
    "\n",
    "    np.save(\n",
    "        f\"./data/WRF/WRF_IP_{cape_thres}_LATxMON.npy\",\n",
    "        wrf_data_LATxMON\n",
    "    )\n",
    "    np.save(\n",
    "        f\"./data/INMCM/INMCM_IP_{cape_thres}_LATxMON.npy\",\n",
    "        inm_data_LATxMON\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91bc6d7a-393c-4078-9a6d-1955393d55f5",
   "metadata": {},
   "source": [
    "### Preprocessing IP data from the WRF: the new parameterisation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5bc66a8f-aa8f-4681-91a9-60c9fbdbf8f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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",
    "# simulated using the WRF with CAPE threshold = 500 J/kg\n",
    "# and temperature threshold = 25 °C\n",
    "# contains values of contributions to the IP depending on (d, h, lat, lon)\n",
    "#   d (axis 0) is the number of a day starting with 0 and ending with 5113\n",
    "#              every third day is taken\n",
    "#              d = 0 corresponds to 1 Jan 1980\n",
    "#              d = 5113 corresponds to 30 Dec 2021\n",
    "#              d = 4991 corresponds to 29 Dec 2020\n",
    "#              (we will restrict our attention to 1980–2020)\n",
    "#   h (axis 1) is the hour of the day (an integer in [0, 24])\n",
    "#              the values corresponding to h = 0 and h = 24 are the same\n",
    "#              (we delete the 25th value)\n",
    "# lat (axis 2) describes the latitude (an integer in [0, 179])\n",
    "# lon (axis 3) describes the longitude (an integer in [0, 359])\n",
    "wrf_raw_ip_data = np.load(\n",
    "    f\"{src_path}/WRF-IP-MAP-500-T2-25.npy\"\n",
    ")[:wrf_N_days, :24]\n",
    "\n",
    "# normalising contributions to the IP to the global mean of 240 kV\n",
    "wrf_raw_ip_data /= (1/240e3) * wrf_raw_ip_data.sum(axis=(-2,-1)).mean()\n",
    "\n",
    "# filling the dictionaries with averaged values\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",
    "    \"./data/WRF/WRF_HOURLY_TOTAL_IP_500_T2_25.npy\",\n",
    "    wrf_hourly_total_ip,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "17036c19-95f8-40df-a6c9-f8a23cf426f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialising an array to store monthly averaged values\n",
    "# for different latitudes\n",
    "wrf_data_LATxMON = np.zeros((180, 12))\n",
    "\n",
    "# iterating over month numbers (starting with 0)\n",
    "for month_idx in range(12):\n",
    "    # filtering indices by the month number\n",
    "    wrf_monthly_indices = [i for i, date in enumerate(wrf_dt_indices)\n",
    "                           if date.month == month_idx + 1]\n",
    "\n",
    "    # putting the values for the specific month into the array\n",
    "    wrf_data_LATxMON[:, month_idx] = \\\n",
    "        wrf_daily_latitudal_ip[wrf_monthly_indices].mean(axis=0)\n",
    "\n",
    "np.save(\"./data/WRF/WRF_IP_500_T2_25_LATxMON.npy\", wrf_data_LATxMON)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e24297fc-cf81-4ea7-9a80-cdcaf277474a",
   "metadata": {},
   "source": [
    "### Saving the number of days for each month (used to compute mean values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "894ad630-17a5-4744-907e-a07768ff7848",
   "metadata": {},
   "outputs": [],
   "source": [
    "# saving the number of days for each month\n",
    "# necessary for correct averaging due to \n",
    "# different numbers of days in different months\n",
    "\n",
    "wrf_days = np.array([len([i for i, date in enumerate(wrf_dt_indices) \n",
    "                          if date.month == m + 1])\n",
    "                     for m in range(12)])\n",
    "inm_days = np.array([len([i for i, date in enumerate(inm_dt_indices) \n",
    "                          if date.month == m + 1])\n",
    "                     for m in range(12)])\n",
    "\n",
    "np.save(\"./data/WRF/WRF_NUMDAYS_MON.npy\", wrf_days)\n",
    "np.save(\"./data/INMCM/INMCM_NUMDAYS_MON.npy\", inm_days)\n",
    "\n",
    "# to calculate the annual mean value, use\n",
    "# `(wrf_data_LATxMON[:, :].sum(axis=0) * days).sum() / days.sum()`\n",
    "# rather than\n",
    "# `wrf_data_LATxMON[:, :].sum(axis=0).mean()`,\n",
    "# since\n",
    "# `((a1+a2+a3)/3 + (b1+b2)/2)/2 != (a1+a2+a3+b1+b2)/5`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04edcb46-b3f9-491a-ba88-509d0fceaac5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}