{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "!pip install matplotlib s3fs \"xarray[io]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import zarr\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm\n",
    "from scipy.signal import stft"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Level 2 Data\n",
    "\n",
    "In this notebook we demonstrate the variety of different data available in the FAIR MAST dataset. In this example we are using level 2 MAST data, which includes cropping, interpolation, calibration, etc. of each signal, as well as mapping each diagnostic group try and follow [IMAS](https://imas-data-dictionary.readthedocs.io/en/latest/) naming convetions. The level 2 data is well-indexed data and follows the FAIR principles. Shots are also filtered using the plasma current to remove shots which were only used for testing, comissioning, machine calibration etc."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First we need to connect to the remote S3 storage bucket to access the data. Each shot from MAST is stored as a seperate [Zarr](https://zarr.readthedocs.io/en/stable/) file.\n",
    "\n",
    "Using `fsspec` and `xarray` we can remotely read data directly over the web. In the example below we also turn on local file caching, allowing us to avoid reading over the network multiple times."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shot_id = 30421\n",
    "\n",
    "endpoint_url = 'https://s3.echo.stfc.ac.uk'\n",
    "url = f's3://mast/level2/shots/{shot_id}.zarr'\n",
    "\n",
    "# Get a handle to the remote file\n",
    "store = zarr.storage.FsspecStore.from_url(\n",
    "    url,\n",
    "    storage_options=dict(\n",
    "        protocol='simplecache',\n",
    "        target_protocol=\"s3\",\n",
    "        cache_storage='.cache',\n",
    "        target_options=dict(\n",
    "            anon=True, endpoint_url=endpoint_url, asynchronous=True\n",
    "        )\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary Profiles\n",
    "\n",
    "The summary group provides a collection of general physics quantities for an experiment. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-input"
    ]
   },
   "outputs": [],
   "source": [
    "def plot_1d_profiles(profiles: xr.Dataset):\n",
    "    \"\"\"Helper function for plotting 1D profiles\"\"\"\n",
    "    n = int(np.ceil(len(profiles.data_vars) / 2))\n",
    "    fig, axes = plt.subplots(n, 2, figsize=(10, 2*n))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    for i, name in enumerate(profiles.data_vars.keys()):\n",
    "        profiles[name].plot(x='time', ax=axes[i])\n",
    "\n",
    "    for ax in axes:\n",
    "        ax.grid('on', alpha=0.5)\n",
    "        ax.set_xlim(profiles.time.min(), profiles.time.max())\n",
    "\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='summary')\n",
    "\n",
    "plot_1d_profiles(profiles)\n",
    "profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pulse Schedule"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='pulse_schedule')\n",
    "\n",
    "fig, axes = plt.subplots(2, 1, figsize=(10, 5))\n",
    "axes = axes.flatten()\n",
    "profiles['i_plasma'].plot(x='time', ax=axes[0])\n",
    "profiles['n_e_line'].plot(x='time', ax=axes[1])\n",
    "\n",
    "\n",
    "for ax in axes:\n",
    "    ax.grid('on', alpha=0.5)\n",
    "plt.tight_layout()\n",
    "\n",
    "profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Magnetics\n",
    "\n",
    "Magnetic diagnostics for equilibrium identification and plasma shape control."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='magnetics')\n",
    "\n",
    "fig, axes = plt.subplots(4, 3, figsize=(8, 10))\n",
    "axes = axes.flatten()\n",
    "\n",
    "profiles['b_field_pol_probe_ccbv_field'].plot.line(x='time', ax=axes[0], add_legend=False)\n",
    "profiles['b_field_pol_probe_obv_field'].plot.line(x='time', ax=axes[1], add_legend=False)\n",
    "profiles['b_field_pol_probe_obr_field'].plot.line(x='time', ax=axes[2], add_legend=False)\n",
    "\n",
    "\n",
    "profiles['b_field_pol_probe_omv_voltage'].plot.line(x='time_mirnov', ax=axes[3], add_legend=False)\n",
    "profiles['b_field_pol_probe_cc_field'].plot.line(x='time_mirnov', ax=axes[4], add_legend=False)\n",
    "profiles['b_field_tor_probe_cc_field'].plot.line(x='time_mirnov', ax=axes[5], add_legend=False)\n",
    "\n",
    "profiles['b_field_tor_probe_saddle_field'].plot.line(x='time_saddle', ax=axes[6], add_legend=False)\n",
    "profiles['b_field_tor_probe_saddle_voltage'].plot.line(x='time_saddle', ax=axes[7], add_legend=False)\n",
    "profiles['b_field_tor_probe_omaha_voltage'].plot.line(x='time_omaha', ax=axes[8], add_legend=False)\n",
    "\n",
    "profiles['flux_loop_flux'].plot.line(x='time', ax=axes[9], add_legend=False)\n",
    "profiles['ip'].plot.line(x='time', ax=axes[10], add_legend=False)\n",
    "\n",
    "for ax in axes:\n",
    "    ax.grid('on', alpha=0.5)\n",
    "plt.tight_layout()\n",
    "\n",
    "profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looking at the spectrogram of one of the mirnov coils can show us information about the MHD modes. Here we see several mode instabilities occuring before the plasma is lost."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = profiles['b_field_pol_probe_omv_voltage'].isel(b_field_pol_probe_omv_channel=1)\n",
    "# Parameters to limit the number of frequencies\n",
    "nperseg = 2000  # Number of points per segment\n",
    "nfft = 2000  # Number of FFT points\n",
    "\n",
    "# Compute the Short-Time Fourier Transform (STFT)\n",
    "sample_rate = 1/(ds.time_mirnov[1] - ds.time_mirnov[0])\n",
    "f, t, Zxx = stft(ds, fs=int(sample_rate), nperseg=nperseg, nfft=nfft)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 5))\n",
    "cax = ax.pcolormesh(t, f/1000, np.abs(Zxx), shading='nearest', cmap='jet', norm=LogNorm(vmin=1e-5))\n",
    "ax.set_ylim(0, 50)\n",
    "ax.set_title(f'Shot {shot_id}')\n",
    "ax.set_ylabel('Frequency [Hz]')\n",
    "ax.set_xlabel('Time [sec]')\n",
    "plt.colorbar(cax, ax=ax)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spectrometer Visible\n",
    "\n",
    "Spectrometer in visible light range diagnostic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='spectrometer_visible')\n",
    "profiles['filter_spectrometer_dalpha_voltage'].plot.line(x='time')\n",
    "profiles['filter_spectrometer_bes_voltage'].isel(bes_channel=0).plot.line(x='time_bes')\n",
    "profiles\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### PF Active"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='pf_active')\n",
    "fig, axes = plt.subplots(3, 1)\n",
    "axes = axes.flatten()\n",
    "\n",
    "profiles['coil_current'].plot.line(x='time', ax=axes[0], add_legend=False)\n",
    "profiles['coil_voltage'].plot.line(x='time', ax=axes[1], add_legend=False)\n",
    "profiles['solenoid_current'].plot.line(x='time', ax=axes[2], add_legend=False)\n",
    "\n",
    "plt.tight_layout()\n",
    "profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Soft X-rays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='soft_x_rays')\n",
    "fig, axes = plt.subplots(3, 1)\n",
    "\n",
    "\n",
    "profiles['horizontal_cam_lower'].plot.line(x='time', ax=axes[1], add_legend=False)\n",
    "axes[1].set_ylim(0, 0.02)\n",
    "\n",
    "profiles['horizontal_cam_upper'].plot.line(x='time', ax=axes[2], add_legend=False)\n",
    "axes[2].set_ylim(0, 0.02)\n",
    "\n",
    "if \"tangential_cam\" in profiles:\n",
    "    profiles['tangential_cam'].plot.line(x='time', ax=axes[0], add_legend=False)\n",
    "    axes[0].set_ylim(0, 0.2)\n",
    "\n",
    "plt.tight_layout()\n",
    "profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Thomson Profiles\n",
    "\n",
    "Thomson scattering measurements in a tokamak provide information about the plasma's electron temperature and density profiles. The diagnostic analyses the scattering of laser light off free electrons in the plasma from a number of radial channels.\n",
    "\n",
    "Below we plot the following profiles measured by the Thomson diagnostic\n",
    "\n",
    "- $T_e$ - Electron temperature\n",
    "- $N_e$ - Electron density\n",
    "- $P_e$ - Electron pressure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='thomson_scattering')\n",
    "profiles\n",
    "\n",
    "fig, axes = plt.subplots(3, 1)\n",
    "axes = axes.flatten()\n",
    "profiles.t_e.plot(x='time', y='major_radius', ax=axes[0])\n",
    "profiles.n_e.plot(x='time', y='major_radius', ax=axes[1])\n",
    "profiles.p_e.plot(x='time', y='major_radius', ax=axes[2])\n",
    "plt.tight_layout()\n",
    "\n",
    "profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 1)\n",
    "profiles['t_e_core'].plot(x='time', ax=axes[0])\n",
    "profiles['n_e_core'].plot(x='time', ax=axes[1])\n",
    "for ax in axes:\n",
    "    ax.grid('on', alpha=0.5)\n",
    "plt.tight_layout()\n",
    "profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CXRS Profiles\n",
    "\n",
    "Charge Exchange Recombination Spectroscopy (CXRS) measurements provide information about ion temperature and plasma rotation. This diagnostic analyses the light emitted from charge exchange reactions between injected neutral beams and plasma ions.\n",
    "\n",
    "Below we plot the following profiles measured by the CXRS diagnostic\n",
    "\n",
    " - $T_i$ - Ion temperature\n",
    " - $V_i$ - Ion velocity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='charge_exchange')\n",
    "\n",
    "fig, axes = plt.subplots(2, 1)\n",
    "profiles['t_i'].plot(x='time', y='major_radius', ax=axes[0], vmax=1000)\n",
    "profiles['v_i'].plot(x='time', y='major_radius', ax=axes[1], vmax=1000)\n",
    "plt.tight_layout()\n",
    "profiles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Equilibrium"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='equilibrium')\n",
    "\n",
    "profile_1d = profiles.drop_vars(['j_tor', 'psi', 'q'])\n",
    "plot_1d_profiles(profile_1d)\n",
    "\n",
    "profiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(10, 5))\n",
    "\n",
    "profiles['j_tor'].isel(time=50).plot(ax=axes[0], x='major_radius')\n",
    "profiles['psi'].isel(time=50).plot(ax=axes[1], x='major_radius')\n",
    "profiles['q'].isel(time=50).plot(ax=axes[2])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Gas Injection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "profiles = xr.open_zarr(store, group='gas_injection')\n",
    "\n",
    "plot_1d_profiles(profiles)\n",
    "plt.tight_layout()\n",
    "profiles"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.13.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}