{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "!pip install aiohttp matplotlib requests scipy \"xarray[io]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import requests\n",
    "import xarray as xr\n",
    "from matplotlib.colors import LogNorm\n",
    "from scipy.signal import stft\n",
    "\n",
    "plt.rcParams[\"font.family\"] = \"sans\"\n",
    "plt.rcParams[\"font.size\"] = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Level 1 Data\n",
    "\n",
    "This notebook contains example plots of data from different diagnostics across MAST without any preprocessing, interpolation, calibration, cropping, etc... applied to the dataset. Data in level 1 are supplied under the original names used during the time of MAST's operation. This dataset contains all shots that could be pulled from the MAST archive, including instrument calibration and testing shots.\n",
    "\n",
    "First we need to find the url to a particular shot. Here we are going to use shot 30421 as an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "shot_data = requests.get(\"https://mastapp.site/json/shots/30421\").json()\n",
    "endpoint, url = shot_data[\"endpoint_url\"], shot_data[\"url\"]\n",
    "shot_url = url.replace(\"s3:/\", endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plasma Current Data\n",
    "\n",
    "Data from the `amc` source contains\n",
    " - Plasma Current ($I_p$): Flows within the plasma, providing initial heating and contributing to the poloidal magnetic field for confinement and stability.\n",
    " - PF Coil Currents: Control the poloidal magnetic field, allowing for plasma shaping, vertical stability, and edge magnetic configuration control.\n",
    " - TF Coil Currents: Generate the strong toroidal magnetic field necessary for primary plasma confinement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(shot_url, group='amc')\n",
    "dataset = dataset.isel(time=(dataset.time > 0) & (dataset.time < .35))\n",
    "fig, axes = plt.subplots(3, 1, figsize=(10, 6))\n",
    "ax1, ax2, ax3 = axes.flatten()\n",
    "\n",
    "ax1.plot(dataset['time'], dataset['plasma_current'])\n",
    "ax1.set_xlabel('Time (s)')\n",
    "ax1.set_ylabel('Plasma Current $I_p$ (kA)')\n",
    "\n",
    "ax2.plot(dataset['time'], dataset['sol_current'])\n",
    "ax2.set_xlabel('Time (s)')\n",
    "ax2.set_ylabel('Solenoid Feed Current (kA)')\n",
    "\n",
    "ax3.plot(dataset['time'], dataset['tf_current'])\n",
    "ax3.set_xlabel('Time (s)')\n",
    "ax3.set_ylabel('TF Feed Current (kA)')\n",
    "\n",
    "for ax in axes:\n",
    "    ax.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Thompson Scattering Data\n",
    "\n",
    "`ayc` source holds the Thomspon Scattering data at the core. Thomson scattering diagnostics provide accurate measurements of electron temperature and density."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(shot_url, group='ayc')\n",
    "dataset = dataset[['te_core', 'ne_core']].dropna(dim='time')\n",
    "fig, axes = plt.subplots(2,1)\n",
    "ax1, ax2 = axes\n",
    "\n",
    "ax1.plot(dataset['time'], dataset['te_core'])\n",
    "ax1.set_xlabel('Time (s)')\n",
    "ax1.set_ylabel('Core Temperature (eV)')\n",
    "\n",
    "ax2.plot(dataset['time'], dataset['ne_core'])\n",
    "ax2.set_xlabel('Time (s)')\n",
    "ax2.set_ylabel('Peak Core Electron Density ($1 / m^3$)')\n",
    "\n",
    "for ax in axes:\n",
    "    ax.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CO2 Interferometers\n",
    "\n",
    "CO2 interferometers (`ane`) are used to measure the electron density in the plasma. By measuring the phase shift of the laser beam as it passes through the plasma, the electron density can be inferred with high precision."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(shot_url, group='ane')\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(shot_url, group='ane')\n",
    "plt.plot(dataset['time'], dataset['density'])\n",
    "ax.set_xlabel('Time (s)')\n",
    "ax.set_ylabel('Integrated Electron Density ($1 / m^2$)')\n",
    "ax.grid(alpha=0.3)\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Equillibrium Reconstruction Data\n",
    "\n",
    "The source `efm` contains data from EFIT. EFIT is a computational tool used to reconstruct the magnetic equilibrium configuration of the plasma in a tokamak. It calculates the shape and position of the plasma, as well as the distribution of the current and pressure within it, based on magnetic measurements.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(shot_url, group='efm')\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we show how to load and plot the plasma current denisty and with the last closed flux surface (LCFS)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d = dataset['plasma_current_rz'].dropna(dim='time')\n",
    "r = dataset['r']\n",
    "z = dataset['z']\n",
    "\n",
    "\n",
    "lcfs_R = dataset['lcfs_r'].sel(time=d.time)\n",
    "lcfs_Z = dataset['lcfs_z'].sel(time=d.time)\n",
    "\n",
    "R, Z = np.meshgrid(r, z)\n",
    "\n",
    "index = 50\n",
    "\n",
    "# Get the x-point\n",
    "xpoint_r = dataset['xpoint1_rc'][index]\n",
    "xpoint_z = dataset['xpoint1_zc'][index]\n",
    "\n",
    "# Get the current centre\n",
    "mag_axis_r = dataset['current_centrd_r'][index]\n",
    "mag_axis_z = dataset['current_centrd_z'][index]\n",
    "\n",
    "# Get the last closed flux surface (LCFS)\n",
    "lcfs_r = lcfs_R[index].values\n",
    "lcfs_r = lcfs_r[~np.isnan(lcfs_r)]\n",
    "lcfs_z = lcfs_Z[index].values\n",
    "lcfs_z = lcfs_z[~np.isnan(lcfs_z)]\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "ax.contourf(R, Z, d[index], cmap='magma', levels=20)\n",
    "ax.plot(lcfs_r, lcfs_z, c='red', linestyle='--', label='LCFS')\n",
    "ax.scatter(xpoint_r, xpoint_z, marker='x', color='green', label='X Point')\n",
    "ax.scatter(mag_axis_r, mag_axis_z, marker='o', color='purple', label='Current Centre')\n",
    "\n",
    "plt.title(f'EFIT Plasma Current & LCFS for Shot {d.attrs[\"shot_id\"]}')\n",
    "plt.ylabel('Z (m)')\n",
    "plt.xlabel('R (m)')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mirnov Coils\n",
    "Mirnov coils are primarily used to measure magnetic fluctuations in the plasma. These fluctuations can provide important information about various plasma instabilities. \n",
    "\n",
    "They are particularly useful for studying magnetohydrodynamic (MHD) phenomena. MHD activity includes various modes of instabilities, such as kink modes and tearing modes, which can affect plasma confinement and stability."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(\n",
    "    \"https://s3.echo.stfc.ac.uk/mast/level1/shots/29790.zarr\", group=\"xmo\"\n",
    ")\n",
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can first look at the line profile for one of the Mirnov coils:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n",
    "ax.plot(dataset['time'], dataset['omaha_3lz'])\n",
    "ax.grid()\n",
    "\n",
    "ax.grid(alpha=0.3)\n",
    "ax.set_ylabel('Voltage (V)')\n",
    "ax.set_xlabel('Time (s)')\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looking at the spectrogram of the dataset 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 = dataset['omaha_3lz']\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[1] - ds.time[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'XMO/OMAHA/3LZ - Shot {ds.attrs[\"shot_id\"]}')\n",
    "ax.set_ylabel('Frequency [Hz]')\n",
    "ax.set_xlabel('Time [sec]')\n",
    "plt.colorbar(cax, ax=ax)\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Photron Camera Data\n",
    "### RBA\n",
    "RBA contains the data from Photron bullet camera A. \n",
    "\n",
    "A Photron Bullet Camera provides high-speed, high-resolution imaging of fast transient events in the plasma. Its ability to capture detailed images of plasma instabilities, turbulence, and disruptions makes it essential for understanding and controlling plasma behavior, ultimately aiding in the pursuit of sustained nuclear fusion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(shot_url, group='rba')\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(dataset.data[50], cmap='gray')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### RBB\n",
    "\n",
    "RBB contains the data from Photron bullet camera B, which is looking at the central column. \n",
    "\n",
    "A Photron Bullet Camera provides high-speed, high-resolution imaging of fast transient events in the plasma. Its ability to capture detailed images of plasma instabilities, turbulence, and disruptions makes it essential for understanding and controlling plasma behavior, ultimately aiding in the pursuit of sustained nuclear fusion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = xr.open_zarr(shot_url, group='rbb')\n",
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(dataset.data[50], cmap='gray')\n",
    "plt.tight_layout()"
   ]
  }
 ],
 "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
}