{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d8002a9b-edff-4381-8dfd-1568156ae3fa",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "# Example: TROPOMI NO$_2$ plume\n",
    "An example demonstrating the application of the CSF and Gaussian plume inversion for estimating NOx emissions from the Matimba/Medupi power plant using TROPOMI NO2 images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0a070ea-f8aa-431e-8d0d-2c73dcb6137a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import cartopy.crs as ccrs\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import ucat\n",
    "import xarray as xr\n",
    "\n",
    "import ddeq\n",
    "\n",
    "sources = ddeq.misc.read_point_sources()\n",
    "sources = sources.sel(source=['Matimba'])\n",
    "\n",
    "filename = os.path.join(ddeq.DATA_PATH, 'Matimba_S5P_RPRO_L2__NO2____20210725T110715.nc')\n",
    "\n",
    "DOMAIN = ddeq.misc.Domain('Matimba', 25, -25.2, 29, -22.9)\n",
    "CRS = ddeq.misc.get_opt_crs(DOMAIN)\n",
    "\n",
    "CURVE_STYLE = \"BEZIER\"     # \"POLY\" or \"BEZIER\"\n",
    "DETECTION = \"THR\"         # \"THR\" or \"WIND\"\n",
    "BACKGROUND = \"LINEAR\"      # \"LINEAR\" or \"SMOOTH\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bb085d6-e593-4b5e-8805-b41ee48f8ce4",
   "metadata": {},
   "source": [
    "## Sentinel-5P/TROPMI NO2 dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2422f026-bbc8-45d3-a653-0cfdcb6c6799",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = xr.open_dataset(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "918df6dc-c7d3-4362-91c7-faaf75859c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ddeq.vis.show_level2(\n",
    "    data, 1e6 * data['NO2'],\n",
    "    sources=sources,\n",
    "    vmin=0,\n",
    "    vmax=300,\n",
    "    label=\"NO$_2$ vertical column density [µmol m$^{-2}$]\", do_zoom=True,\n",
    "    units='umol m-2'\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a15f5120-72bf-4985-926b-df397167839d",
   "metadata": {},
   "source": [
    "## ERA5 wind files\n",
    "The code below would download ERA5 single level and pressure level. Since the files are already in `ddeq.DATA_PATH`, no files will be downloaded. If you want to download own ERA5 data, please the change `data_path` to the location on your system, where you like to save your files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "064f00ed-6e6e-48bc-b650-d9d69e068374",
   "metadata": {},
   "outputs": [],
   "source": [
    "sng_filename = ddeq.era5dl.download_single_lvl(\n",
    "    time=pd.Timestamp(data.time.values),\n",
    "    area=DOMAIN.extent,\n",
    "    timesteps=24,\n",
    "    data_path=ddeq.DATA_PATH,\n",
    "    prefix=\"Matimba\"\n",
    ")\n",
    "\n",
    "lvl_filename = ddeq.era5dl.download_pressure_lvl(\n",
    "    time=pd.Timestamp(data.time.values),\n",
    "    area=DOMAIN.extent,\n",
    "    timesteps=24,\n",
    "    data_path=ddeq.DATA_PATH,\n",
    "    prefix=\"Matimba\"\n",
    ")\n",
    "\n",
    "sng_filename, lvl_filename\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e47c234b-2bd6-48c0-b1f3-b6ebe5ddfe29",
   "metadata": {},
   "source": [
    "## Cross-sectional flux method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52a03f2f-5337-474f-84e6-a221b6c9dff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Effective wind at source locations using GNFR-A weighted winds:\n",
    "winds = ddeq.era5.read(\n",
    "    sng_filename,\n",
    "    lvl_filename,\n",
    "    method=\"gnfra\",\n",
    "    sources=sources,\n",
    "    times=data.time\n",
    ")\n",
    "winds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92cfce6c-e759-40ae-8b9a-2dec972bab85",
   "metadata": {},
   "outputs": [],
   "source": [
    "if DETECTION == \"THR\":\n",
    "\n",
    "    var_sys = (ucat.convert_columns(0.5e15, 'cm-2', 'mol m-2',\n",
    "                                    molar_mass='NO2'))**2\n",
    "\n",
    "    data = ddeq.dplume.detect_plumes(data, sources,\n",
    "                                     variable='NO2',\n",
    "                                     variable_std='NO2_std',\n",
    "                                     var_sys=var_sys,\n",
    "                                     filter_type='gaussian',\n",
    "                                     filter_size=0.5, crs=CRS)\n",
    "elif DETECTION == \"WIND\":\n",
    "    data = ddeq.dplume.detect_from_wind(data, sources, winds, dmax=30e3, crs=CRS)\n",
    "else:\n",
    "    raise ValueError"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4113cfbb-ac6c-4f40-88be-3e8991b7c56e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ddeq.vis.show_level2(\n",
    "    data, 1e6 * data['NO2'], gas='NO2',\n",
    "    vmin=0,\n",
    "    vmax=300,\n",
    "    label=\"NO$_2$ columns [µmol m$^{-2}$]\",\n",
    "    winds=winds,\n",
    "    domain=DOMAIN,\n",
    "    do_zoom=False, show_clouds=False, crs=CRS,\n",
    "    figwidth=4,\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7288d802-a2f5-4fcd-9308-5c649ada3596",
   "metadata": {},
   "outputs": [],
   "source": [
    "variable = \"NO2_mass\"\n",
    "background = \"linear\"\n",
    "ddeq.emissions.convert_units(data, \"NO2\", \"NO2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cdb5563-c641-41b7-a769-95ff70bd57a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "if DETECTION == \"THR\":\n",
    "    data = ddeq.curves.fit_to_detections(data, n_nodes=3, force_origin=True, use_weights=True)\n",
    "\n",
    "data = ddeq.curves.compute_natural_coords(data)\n",
    "data = ddeq.curves.compute_plume_areas(data, plume_width=120e3)\n",
    "\n",
    "# estimate NO2 enhancement\n",
    "ddeq.background.estimate(data, \"NO2\")\n",
    "ddeq.emissions.compute_plume_signal(data, \"NO2\")\n",
    "\n",
    "ddeq.emissions.convert_units(data, \"NO2\", \"NO2_estimated_background\")\n",
    "ddeq.emissions.convert_units(data, \"NO2\", \"NO2_minus_estimated_background\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f86c5ce4-9c91-491f-a5cd-b64db01b3a9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = None if DETECTION == \"WIND\" else 10e3\n",
    "\n",
    "results_csf = ddeq.csf.estimate_emissions(data, winds, sources, 'NO2',\n",
    "                                          xmin=10e3, dx=dx,\n",
    "                                          f_model=2.24, crs=CRS,\n",
    "                                          variable=variable, background=background\n",
    "                                         )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a72a300-eb97-4649-9a8d-08aca159d46b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddeq.vis.plot_csf_result(['NO2'], data, winds, results_csf, 'Matimba',\n",
    "                         vmins=[0], vmaxs=[200e-6], crs=CRS);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59154b37-ac87-49cd-a2ed-f7d33b151d09",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ddeq.vis.show_level2(\n",
    "    data,\n",
    "    1e6 * data['NO2'],\n",
    "    gas='NO2',\n",
    "    vmin=0,\n",
    "    vmax=300,\n",
    "    label=\"NO$_2$ columns [µmol m$^{-2}$]\",\n",
    "    winds=winds,\n",
    "    domain=DOMAIN,\n",
    "    do_zoom=False,\n",
    "    show_clouds=False,\n",
    "    crs=CRS,\n",
    "    figwidth=4,\n",
    ");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63188d98-434b-4c3b-bb68-70f0e436cb8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax  = plt.subplots(1, figsize=(8,2))\n",
    "ddeq.vis.plot_along_plume(ax, 'NO2', results_csf.sel(source='Matimba'))\n",
    "ax.set_xlim(right=207)\n",
    "ax.set_ylim(0,150)\n",
    "plt.tight_layout()\n",
    "\n",
    "ax.legend().remove()\n",
    "ax.legend(loc=1, ncol=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4250c3c8-ebb8-4062-b28e-6d1799fc514c",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes  = plt.subplots(1,3, figsize=(8,2), sharey=True)\n",
    "\n",
    "\n",
    "for i, ax in zip([1,9,-1], axes):\n",
    "    try:\n",
    "        r = results_csf.sel(source='Matimba').isel(polygon=i)\n",
    "    except IndexError:\n",
    "        continue\n",
    "    ddeq.vis.plot_across_section(r, gases=['NO2'], method='gauss', ax=ax,\n",
    "                                 legend='simple')\n",
    "\n",
    "    ax.text(-36, 19.5, '%d-%d km' % (r.xa/1e3, r.xb/1e3), ha='left', va='top')\n",
    "\n",
    "for ax in axes[1:]:\n",
    "    ax.set_ylim(-1, 20)\n",
    "    ax.set_ylabel('')\n",
    "\n",
    "\n",
    "plt.tight_layout()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "449cd7bf-d4c4-4425-acc8-e324e72addfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddeq.vis.plot_csf_result(['NO2'], data, winds, results_csf, 'Matimba',\n",
    "                         vmins=[0], vmaxs=[200e-6], crs=CRS);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "549bdeff-9320-4e62-8417-81860a753751",
   "metadata": {},
   "source": [
    "## Gaussian plume inversion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22bff546-5254-4c3e-a9ff-fd60a014f08e",
   "metadata": {},
   "outputs": [],
   "source": [
    "priors = {'Matimba': {\n",
    "    'NO2': {'Q': 3.0,       # kg/s\n",
    "            'tau': 4*60**2  # seconds\n",
    "           }\n",
    "}}\n",
    "data, results_gauss = ddeq.gauss.estimate_emissions(data, winds, sources,\n",
    "                                                    ['NO2'], priors=priors,\n",
    "                                                    fit_decay_times=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "159549c0-ae6d-4c4a-8318-3815809e4d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_gauss = ddeq.emissions.convert_NO2_to_NOx_emissions(results_gauss, f=2.38)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d63089c9-04e5-4e76-8d2f-968835e4e576",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ddeq.vis.plot_gauss_result(data, results_gauss, ['Matimba'],\n",
    "                                 'NO2', crs=CRS, vmin=0, vmax=10e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc0cc2d3-b5f4-4879-a682-aae3317a5c12",
   "metadata": {},
   "source": [
    "## IME method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fadf60de-b4af-4855-9e7c-c84101a09d03",
   "metadata": {},
   "outputs": [],
   "source": [
    "results_ime = ddeq.ime.estimate_emissions(\n",
    "    data,\n",
    "    winds,\n",
    "    sources,\n",
    "    \"NO2\",\n",
    "    decay_time=4.0*60**2\n",
    ")\n",
    "results_ime = ddeq.emissions.convert_NO2_to_NOx_emissions(results_ime, f=1.32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c8c8605-5756-48ac-8e7c-1a7d10117c29",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddeq.vis.plot_ime_result(\n",
    "    \"NO2\",\n",
    "    data,\n",
    "    winds,\n",
    "    results_ime,\n",
    "    \"Matimba\",\n",
    "    crs=CRS,\n",
    "    vmin=0,\n",
    "    vmax=300e-6,\n",
    ");"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44396af2-b405-4557-bcd1-5674fa91841d",
   "metadata": {},
   "source": [
    "## LCSF method\n",
    "Quick example applying the LCSF method to TROPOMI NO2 data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7cda1e8c-62fe-4c05-98fa-176c8fe06526",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert from mol/m2 to molec/cm2 for LCSF method\n",
    "data[\"NO2\"] = ucat.convert_columns(data[\"NO2\"], \"mol/m2\", \"cm-2\")\n",
    "data[\"NO2_std\"] = ucat.convert_columns(data[\"NO2_std\"], \"mol/m2\", \"cm-2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f62fd9b-1767-4fb3-bb5c-71044809e2b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "wind_field = ddeq.era5.read(\n",
    "    sng_filename,\n",
    "    lvl_filename,\n",
    "    method=\"gnfra\",\n",
    "    times=data.time\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2209b2b9-ab88-4457-b0b1-43f75f2e731a",
   "metadata": {},
   "outputs": [],
   "source": [
    "lcs_params = {}\n",
    "lcs_params['use_prior'] = True\n",
    "lcs_params[\"verbose\"] = True\n",
    "\n",
    "lcs_params['n_min_fit_pts'] = 10\n",
    "lcs_params['fit_pt_slice_width'] = 15\n",
    "\n",
    "lcs_params['f_NOx_NO2'] = 3.5\n",
    "lcs_params['NOx_NO2_tau_depletion'] = 4.0 * 60**2\n",
    "\n",
    "lcsf_results = ddeq.lcsf.estimate_emissions(data, wind_field, sources, ['NO2'], priors=priors, lcs_params=lcs_params, all_diags=True)\n",
    "print(f\"NOx emissions: {np.nanmedian(lcsf_results['NO2_emissions']):.1f} kt NO2 / a\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44c86a66-3591-4fa9-a063-8743127043b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = ddeq.vis.plot_lcsf_result(\"Matimba\", lcsf_results, data, sources,\n",
    "                                gases=['NO2'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd12113e-5059-4f37-94f8-82ba2e4f26f4",
   "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.9.21"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}