{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 09 \u2014 Reproducing Zhu, Kollias, Yang (2023): particle-inertia spectra\n",
    "\n",
    "Reference: Zhu, Z., Kollias, P., and Yang, F., 2023. *Particle Inertia\n",
    "Effects on Radar Doppler Spectra Simulation*, Atmos. Meas. Tech.\n",
    "Discuss. MATLAB simulator source: https://zenodo.org/records/7897981\n",
    "\n",
    "The conventional way to get a turbulent radar Doppler spectrum from a\n",
    "drop-size distribution is to convolve the drop reflectivity spectrum\n",
    "with a single Gaussian of width \u03c3_air. Zhu 2023 argues \u2014 and shows\n",
    "with a per-drop drag-ODE simulator \u2014 that this is wrong: small drops\n",
    "are passive tracers, but large drops are ballistic and under-respond\n",
    "to the small-scale eddies of the wind field. The broadening kernel is\n",
    "*diameter-dependent*, \u03c3_t(D), and narrows toward the large-drop end.\n",
    "\n",
    "rustmatrix ships an analytical Stokes-number low-pass response,\n",
    "`InertialZeng2023`, implementing this physics family (the Zeng et al.\n",
    "2023 companion paper):\n",
    "\n",
    "$$\\sigma_t(D) = \\frac{\\sigma_{\\rm air}}{\\sqrt{1 + \\mathrm{St}(D)^2}}, \\qquad \\mathrm{St} = \\tau_p(D) / \\tau_{\\rm eddy}.$$\n",
    "\n",
    "This notebook reproduces Zhu 2023's configuration \u2014 W-band, \u00b112 m/s\n",
    "Nyquist, 1024 bins, SNR = 40 dB, exponential warm-rain PSD \u2014 and\n",
    "compares conventional Gaussian broadening to the inertia-aware kernel.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from rustmatrix import Scatterer, SpectralIntegrator, spectra\n",
    "from rustmatrix.psd import ExponentialPSD, PSDIntegrator\n",
    "from rustmatrix.refractive import m_w_10C\n",
    "from rustmatrix.tmatrix_aux import K_w_sqr, geom_vert_back, wl_W\n",
    "\n",
    "# --- Zhu 2023 configuration ---\n",
    "V_MIN, V_MAX = -12.0, 12.0\n",
    "NFFT = 1024\n",
    "SNR_DB = 40.0\n",
    "Z_DBZ = 20.0\n",
    "SIGMA_AIR = 0.289        # std of uniform [-0.5, 0.5] wind field\n",
    "EPSILON = 1e-2           # dissipation rate (m\u00b2/s\u00b3)\n",
    "L_O = 0.1                # integral length scale (m)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Build the W-band rain scatterer and PSD\n",
    "\n",
    "Zhu 2023's MATLAB code uses N(D) = N\u2080\u00b7exp(shape\u00b7D_cm)\u00b71e4 m\u207b\u00b3 \u00b5m\u207b\u00b9\n",
    "with N\u2080 = 0.08 and shape = -20. In rustmatrix's mm-based units this\n",
    "is N(D_mm) = 8\u00d710\u2075 \u00b7 exp(-2\u00b7D_mm) m\u207b\u00b3 mm\u207b\u00b9, truncated at D_max=4 mm.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "s = Scatterer(wavelength=wl_W, m=m_w_10C[wl_W],\n",
    "              Kw_sqr=K_w_sqr[wl_W], ddelt=1e-4, ndgs=2)\n",
    "integ = PSDIntegrator()\n",
    "integ.D_max = 4.0; integ.num_points = 128\n",
    "integ.geometries = (geom_vert_back,)\n",
    "s.psd_integrator = integ\n",
    "s.psd_integrator.init_scatter_table(s)\n",
    "s.psd = ExponentialPSD(N0=8e5, Lambda=2.0, D_max=4.0)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The inertial reduction factor \u03c3_t(D) / \u03c3_air\n",
    "\n",
    "At fixed \u03c3_air and \u03b5, the reduction factor depends only on D through\n",
    "the particle relaxation time \u03c4_p(D) = v_t(D)/g. This plot is the\n",
    "physical crux of Zhu 2023 \u2014 large drops experience less-effective\n",
    "turbulence broadening than small drops.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "zeng = spectra.InertialZeng2023(sigma_air=SIGMA_AIR,\n",
    "                                 epsilon=EPSILON, L_o=L_O)\n",
    "D_plot = np.linspace(0.1, 4.0, 200)\n",
    "fig, ax = plt.subplots(figsize=(7, 4))\n",
    "ax.plot(D_plot, zeng(D_plot) / SIGMA_AIR, 'C0-', lw=2)\n",
    "ax.axhline(1.0, color='k', ls='--', alpha=0.5,\n",
    "           label='conventional: \u03c3_t / \u03c3_air = 1')\n",
    "ax.set_xlabel('drop diameter D [mm]')\n",
    "ax.set_ylabel('\u03c3_t(D) / \u03c3_air')\n",
    "ax.set_title(f'Inertial reduction (\u03c3_air={SIGMA_AIR}, '\n",
    "             f'\u03b5={EPSILON}, L_o={L_O})')\n",
    "ax.legend(); ax.grid(True, alpha=0.3)\n",
    "ax.set_ylim(0.0, 1.05)\n",
    "fig.tight_layout();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run two spectra: conventional vs. inertia-aware\n",
    "\n",
    "Both use the same fall-speed model, velocity grid, and noise floor.\n",
    "The only difference is the turbulence kernel \u2014 constant \u03c3_air vs.\n",
    "\u03c3_t(D) from Zeng 2023.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Receiver noise: SNR = 40 dB over 20 dBZ reference.\n",
    "noise = 10 ** (Z_DBZ / 10.0) / 10 ** (SNR_DB / 10.0)\n",
    "\n",
    "common = dict(\n",
    "    fall_speed=spectra.fall_speed.atlas_srivastava_sekhon_1973,\n",
    "    v_min=V_MIN, v_max=V_MAX, n_bins=NFFT,\n",
    "    geometry_backscatter=geom_vert_back,\n",
    "    noise=noise,\n",
    ")\n",
    "\n",
    "conv = SpectralIntegrator(s,\n",
    "    turbulence=spectra.GaussianTurbulence(SIGMA_AIR),\n",
    "    **common,\n",
    ").run()\n",
    "\n",
    "iner = SpectralIntegrator(s,\n",
    "    turbulence=spectra.InertialZeng2023(\n",
    "        sigma_air=SIGMA_AIR, epsilon=EPSILON, L_o=L_O),\n",
    "    **common,\n",
    ").run()\n",
    "\n",
    "print(f'Conventional integrated Z_h = '\n",
    "      f'{10*np.log10(np.trapezoid(conv.sZ_h, conv.v)):.2f} dBZ')\n",
    "print(f'Inertia-aware integrated Z_h = '\n",
    "      f'{10*np.log10(np.trapezoid(iner.sZ_h, iner.v)):.2f} dBZ')\n",
    "print(f'Noise floor              = '\n",
    "      f'{noise:.4f} mm\u2076/m\u00b3 ('\n",
    "      f'{10*np.log10(noise / (V_MAX - V_MIN)):.2f} '\n",
    "      f'dBZ / (m/s))')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Spectrum comparison in dBZ\n",
    "\n",
    "On a log scale the inertia correction visibly narrows the spectrum\n",
    "near the fall-speed of the largest drops (the right-hand tail). The\n",
    "total power integrates to the same reflectivity \u2014 this is a\n",
    "redistribution of spectral shape, not a gain change.\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "def dBZ(x):\n",
    "    return 10 * np.log10(np.maximum(x, 1e-12))\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.plot(conv.v, dBZ(conv.sZ_h), 'C0-',\n",
    "        label='conventional Gaussian')\n",
    "ax.plot(iner.v, dBZ(iner.sZ_h), 'C3-',\n",
    "        label='inertia-aware (Zeng 2023)')\n",
    "ax.axhline(dBZ(noise / (V_MAX - V_MIN)), color='k',\n",
    "           ls='--', alpha=0.5, label='noise floor')\n",
    "peak = max(dBZ(conv.sZ_h).max(), dBZ(iner.sZ_h).max())\n",
    "ax.set_xlim(0, 10)\n",
    "ax.set_ylim(-50, np.ceil(peak / 5.0) * 5.0 + 5)\n",
    "ax.set_xlabel('Doppler velocity v [m/s]')\n",
    "ax.set_ylabel('sZ_h [dBZ / (m/s)]')\n",
    "ax.set_title('Zhu 2023 setup \u2014 W-band warm-rain Doppler spectrum')\n",
    "ax.legend(); ax.grid(True, alpha=0.3)\n",
    "fig.tight_layout();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Zoom on the large-drop tail\n",
    "\n",
    "Narrowing is most visible between v \u2248 7\u20139 m/s, where the fastest\n",
    "drops land. Zhu 2023's Lagrangian drop-tracking simulator produces\n",
    "the same qualitative signature (their Fig. 3\u20135).\n"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "ax.plot(conv.v, dBZ(conv.sZ_h), 'C0-', label='conventional')\n",
    "ax.plot(iner.v, dBZ(iner.sZ_h), 'C3-', label='inertia-aware')\n",
    "ax.set_xlim(6, 10)\n",
    "tail_peak = max(dBZ(conv.sZ_h[(conv.v > 6) & (conv.v < 10)]).max(),\n",
    "                dBZ(iner.sZ_h[(iner.v > 6) & (iner.v < 10)]).max())\n",
    "ax.set_ylim(-25, np.ceil(tail_peak / 5.0) * 5.0 + 5)\n",
    "ax.set_xlabel('Doppler velocity v [m/s]')\n",
    "ax.set_ylabel('sZ_h [dBZ / (m/s)]')\n",
    "ax.set_title('Large-drop tail: inertia narrows the spectrum')\n",
    "ax.legend(); ax.grid(True, alpha=0.3)\n",
    "fig.tight_layout();\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}