R/xarray_extract_station_SSR.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8e28b8d6-8c7d-462a-90e7-f842fe756094",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import datetime\n",
    "import gc\n",
    "import glob\n",
    "import os.path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pyarrow.parquet as pq\n",
    "import xarray as xr\n",
    "\n",
    "os.chdir(\n",
    "    \"/mnt/cloud/wwu1/ec_bronze/_nogroup/ae78a1ca-a0e8-4e4e-8992-69c34947db65/UseCase_AIRCON\"\n",
    ")\n",
    "\n",
    "new_names = [\"Longitude\", \"Latitude\", \"Start\", \"SSR\"]\n",
    "st = pq.read_table(\"AQ_stations/EEA_stations_meta_table.parquet\").to_pandas()\n",
    "countries = st.Countrycode.unique().sort_values()\n",
    "# countries = ['AD', 'AL']\n",
    "st = st[[\"Air.Quality.Station.EoI.Code\", \"Longitude\", \"Latitude\"]]\n",
    "years = range(2015, 2024)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "420231e2-ba32-421c-b707-304a3f8bd7df",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def open_ssr(y):\n",
    "\n",
    "    ssr_file = glob.glob(f\"supplementary/era5_download/*solar*{y}*\")\n",
    "    ssr = xr.open_mfdataset(ssr_file, parallel=True)\n",
    "\n",
    "    # Sort out expver dimension (for recent ERA5T data)\n",
    "    # https://confluence.ecmwf.int/display/CUSF/ERA5+CDS+requests+which+return+a+mixture+of+ERA5+and+ERA5T+data\n",
    "    if \"expver\" in ssr.dims:\n",
    "        ssr = ssr.reduce(np.nansum, \"expver\")\n",
    "        # print(\"\\nFound experimental version of SRR data (recent dates). Reducing dimension.\")\n",
    "    return ssr\n",
    "\n",
    "\n",
    "def extract_ssr(c, y, ssr):\n",
    "\n",
    "    out_pq = f\"AQ_data/02_hourly_SSR/{c}_hourly_{y}_gaps_SSR.parquet\"\n",
    "    in_pq = f\"AQ_data/01_hourly/{c}_hourly.parquet\"\n",
    "\n",
    "    if not os.path.isfile(out_pq) and os.path.isfile(in_pq):\n",
    "\n",
    "        sel = [\n",
    "            (\"Start\", \">\", pd.to_datetime(f\"{y}-01-01\", utc=False)),\n",
    "            (\"Start\", \"<\", pd.to_datetime(f\"{y+1}-01-01\", utc=False)),\n",
    "        ]\n",
    "        aq = pq.read_table(in_pq, filters=sel).to_pandas()\n",
    "        aq = pd.merge(aq, st, on=\"Air.Quality.Station.EoI.Code\")\n",
    "\n",
    "        # define + extract locations in 3D\n",
    "        target_lon = xr.DataArray(aq[\"Longitude\"].values, dims=\"points\")\n",
    "        target_lat = xr.DataArray(aq[\"Latitude\"].values, dims=\"points\")\n",
    "        target_time = xr.DataArray(\n",
    "            aq[\"Start\"].values.astype(\"datetime64[ns]\"), dims=\"points\"\n",
    "        )\n",
    "\n",
    "        # extract\n",
    "        ssr_ex = ssr.ssr.sel(\n",
    "            longitude=target_lon,\n",
    "            latitude=target_lat,\n",
    "            time=target_time,\n",
    "            method=\"nearest\",\n",
    "        )\n",
    "        ssr_df = ssr_ex.to_series().rename(\"SSR\")\n",
    "\n",
    "        result = aq.merge(ssr_df, left_index=True, right_on=\"points\")\n",
    "        result.reset_index(drop=True, inplace=True)\n",
    "        result.to_parquet(out_pq)   \n",
    "        del aq, ssr_ex, ssr_df, target_lon, target_lat, target_time\n",
    "\n",
    "        gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "13ccc022-c484-48ef-8932-9976e8a6f232",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#%%time\n",
    "\n",
    "# Test\n",
    "#for y in years:\n",
    "#    ssr = open_ssr(y)\n",
    "#    extract_ssr(\"AD\", y, ssr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fa723ab4-68c9-4af3-bf48-ce07fb60d834",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "2015: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2016: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2017: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2018: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2019: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2020: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2021: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2022: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK \n",
      "2023: AD AL AT BA BE BG CH CY CZ DE DK EE ES FI FR GB GE GR HR HU IE IS IT LT LU LV ME MK MT NL NO PL PT RO RS SE SI SK TR UA XK "
     ]
    }
   ],
   "source": [
    "for y in years:\n",
    "    ssr = open_ssr(y)\n",
    "    print(f\"\\n{y}:\", sep=\" \", end=\" \", flush=True)\n",
    "    for c in countries:\n",
    "        extract_ssr(c, y, ssr)\n",
    "        print(c, sep=\" \", end=\" \", flush=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (Xarray)",
   "language": "python",
   "name": "py-xarray"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}