{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Plotting the Squashing Factor\n\nThis example demonstrates how to use and plot the :func:`~mapflpy.scripts.compute_q_on_surface`\nto visualize key topology and morphology metrics of the magnetic field.\nThe :func:`~mapflpy.scripts.compute_q_on_surface` combines :func:`~mapflpy.scripts.expansion_factor` with\n :func:`~mapflpy.utils.calc_jacobian` and :func:`~mapflpy.utils.calc_q` using\n :func:`~mapflpy.scripts.map_pt_forward` or :func:`~mapflpy.scripts.map_pt_backward`\n to calculate the squashing factor.\n\nFor a more complete description of the squashing factor, see\n[Titov et al. 2007](https://ui.adsabs.harvard.edu/abs/2007ApJ...660..863T).\nGenerally, the squashing factor indicates how much a given flux tube distorts.\nOther ways to imagine the squashing factor, Q, includes places where current sheets\nare likely (but not guaranteed) to form as high Q lines indicate quasi-separatrix layers (QSLs).\nAlternatively, high Q lines indicate different magnetic flux domains.\n\nAdditionally, this plots :func:`~mapflpy.scripts.expansion_factor`, critical to models such as\nWSA (see [Wang & Sheeley 1990](https://ui.adsabs.harvard.edu/abs/1990ApJ...355..726W/abstract),\n[Arge & Pizzo 2000](https://ui.adsabs.harvard.edu/abs/2000JGR...10510465A/abstract), and\n[Arge et al. 2004](https://ui.adsabs.harvard.edu/abs/2004JASTP..66.1295A/abstract)).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nfrom psi_data import fetch_mas_data\nimport numpy as np\nimport matplotlib.pylab as plt\nfrom mapflpy.scripts import map_pt_forward, expansion_factor, compute_q_on_surface"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The squashing factor, Q, is a measure of the topology of the magnetic field.\nSo, let's read in magnetic field files. We're loading in from\na CORHEL-MAS thermodynamic MHD calculation for CR2282. These aren't\ncurrently standard datasets in mapflpy or psi-io, so we're fetching them\nmanually and placing them in the default cache location.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "files = fetch_mas_data(domains=\"cor\", variables=\"br,bt,bp,t\")\nmagnetic_field_files = files.cor_br, files.cor_bt, files.cor_bp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "As a quick look, we can immediately calculate Q\nspecifying only the magnetic field files and visualize the output.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# compute Q\nsquashing_factor_default = compute_q_on_surface(magnetic_field_files)\n# plot Q\nax = plt.figure().add_subplot()\nq_map = ax.pcolormesh(np.rad2deg(squashing_factor_default.p), 90 - np.rad2deg(squashing_factor_default.t),\n                      np.log10(squashing_factor_default.q),\n                      cmap='Grays')\nax.set_aspect(\"equal\", adjustable=\"box\")\nax.set_title('Log$_{10}$ Q')\nplt.colorbar(q_map)\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "That should run fairly quickly, as the resolution is a little bit more than 1 degree.\nHowever, we can customize the map. We can, for example,\nchange the direction of mapping to \"bwd\" and choose a trace_radius of 3.\nWe can either specify a [start, end] for theta and phi with t_range and p_range,\nor specify our own array with t_arr and p_arr.\nWe are intentionally picking a low resolution so this runs fast, you should use more points!\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "squashing_factor_3_bwd = compute_q_on_surface(magnetic_field_files, direction='bwd', nproc=4, trace_radius=3,\n                                              p_arr=np.linspace(0, 2 * np.pi, 80), t_arr=np.linspace(0, np.pi, 40))\n# and visualizing:\nax = plt.figure().add_subplot()\nq_map = ax.pcolormesh(np.rad2deg(squashing_factor_3_bwd.p), 90 - np.rad2deg(squashing_factor_3_bwd.t),\n                      np.log10(squashing_factor_3_bwd.q),\n                      cmap='Grays')\nax.set_aspect(\"equal\", adjustable=\"box\")\nax.set_title('Log$_{10}$ Q')\nplt.colorbar(q_map)\nplt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "This wrapper makes it easy to get the squashing factor. If we're interested\nin just say, the expansion factor, we can plot that.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# We first need to calculate a mapping on a set of given points\np_to_trace = np.linspace(0, 2 * np.pi, 100)\nt_to_trace = np.linspace(0, np.pi, 50)\n\n# let's  map and get the expansion factor\nmapping = map_pt_forward(*magnetic_field_files, p_to_trace, t_to_trace)\nef, p_ef, t_ef = expansion_factor(magnetic_field_files, mapping, 3, p_to_trace, t_to_trace)\n\n# and now we can visualize\nax = plt.figure().add_subplot()\nef_map = ax.pcolormesh(np.rad2deg(p_to_trace), 90 - np.rad2deg(t_to_trace), np.log10(ef).T,\n                       cmap='plasma')\nax.set_aspect(\"equal\", adjustable=\"box\")\nax.set_title('expansion factor')\nplt.colorbar(ef_map)\nplt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.14"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}