{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Reading MHD Data: Units, Meshes, and Value-Space Slicing\n\nDemonstrate the full read and vslice API of :func:`~psi_io.mhd_io.PsiData`:\nloading data in different unit systems, remeshing from the half-mesh to the\nmain mesh, and slicing by physical coordinate value using\n:class:`~astropy.units.Quantity`.\n\nThis example demonstrates:\n\n1. Reading the full dataset in code units, physical CGS units, and a specific unit.\n2. Remeshing a half-mesh field to the main mesh via the *mesh* keyword.\n3. Reading a radial subset by index.\n4. Slicing by physical coordinate value (``vslice``) with a bare scalar and with\n   an :class:`~astropy.units.Quantity` object.\n5. Combining index-space and value-space arguments in a single ``vslice`` call.\n\n<div class=\"alert alert-info\"><h4>Note</h4><p>All positional arguments to :meth:`~psi_io.mhd_io.PsiData.read` and\n   :meth:`~psi_io.mhd_io.PsiData.vslice` are interpreted in physical\n   ``(r, \u03b8, \u03c6)`` order, even though the HDF storage order is ``(N\u03c6, N\u03b8, Nr)``.\n   The returned data array shape always reflects the HDF storage order.</p></div>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import astropy.units as u\nimport numpy as np\nfrom psi_io import data\nfrom psi_io._units import PSI_rsun\nfrom psi_io.mhd_io import PsiData"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Create a reader**\n\nOpen the example radial magnetic field file.  Metadata \u2014 quantity name, physical\nunit, and mesh stagger \u2014 is resolved from the filename and HDF attributes.\nNo array data is transferred from disk at this step.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "br_filepath = data.get_3d_data()\nreader = PsiData(br_filepath)\nprint(f\"quantity : {reader.quantity!r}\")\nprint(f\"shape    : {reader.shape}  (N\u03c6 \u00d7 N\u03b8 \u00d7 Nr)\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Reading in code (native) units**\n\nBy default, :meth:`~psi_io.mhd_io.PsiData.read` returns the dataset in MAS\ncode units.  The values stored on disk are wrapped in the normalization unit\nobject so the result is a :class:`~astropy.units.Quantity`.  Passing\n``unit='native'`` (aliases: ``'code'``, ``'model'``, ``'psi'``) is equivalent.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_code, r, t, p = reader.read()\nprint(f\"unit (code) : {data_code.unit}\")\nprint(f\"shape       : {data_code.shape}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Reading in physical (CGS) units**\n\nPassing ``unit='physical'`` (aliases: ``'real'``, ``'phys'``, ``'cgs'``)\ndecomposes the data into CGS base units.  Any :class:`~astropy.units.Unit`-compatible\nstring is also accepted \u2014 the example below converts directly to Gauss:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_phys, r, t, p = reader.read(unit='physical')\nprint(f\"unit (phys) : {data_phys.unit}\")\n\ndata_gauss, r, t, p = reader.read(unit='Gauss')\nprint(f\"unit (Gauss): {data_gauss.unit}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Remeshing to the main mesh**\n\n``br`` is stored on the *half*-mesh in the radial direction (Yee face-centred).\nPassing ``mesh='main'`` shifts every half-mesh axis to the main mesh by\naveraging adjacent array elements; the radial axis shrinks by one element.\nThe mesh stagger for each quantity is encoded in its\n:class:`~psi_io._models.Props` descriptor (see :mod:`psi_io._models`), and the\naveraging itself is performed by :func:`~psi_io._mesh.remesh_array` from\n:mod:`psi_io._mesh`.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_main, r_main, t_main, p_main = reader.read(mesh='main', unit='Gauss')\nprint(f\"original r points  : {r.shape[0]}\")\nprint(f\"main-mesh r points : {r_main.shape[0]}\")\nprint(f\"data shape (main)  : {data_main.shape}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Reading a subset by index**\n\nPositional arguments to :meth:`~psi_io.mhd_io.PsiData.read` are index-space\nslice specifiers in physical ``(r, \u03b8, \u03c6)`` order.  Here we read only the\nfirst ten radial grid points (all \u03b8 and \u03c6):\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_rslice, r_rslice, t_rslice, p_rslice = reader.read(slice(0, 10), unit='Gauss')\nprint(f\"r-slice shape : {data_rslice.shape}\")\nprint(f\"r scale range : [{r_rslice[0]:.3f}, {r_rslice[-1]:.3f}]\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Value-space slicing with** ``vslice``\n\n:meth:`~psi_io.mhd_io.PsiData.vslice` accepts physical coordinate values as\npositional arguments alongside the usual index-space arguments.  A bare\nscalar is interpreted in the native coordinate unit (solar radii for r,\nradians for \u03b8 and \u03c6).\n\nWe first inspect the radial scale to pick a coordinate safely within the domain:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "r_scale = reader.scales.r.read()\nprint(f\"r domain : [{r_scale[0]:.5f}, {r_scale[-1]:.5f}]\")\n\nr_target = float(r_scale.value[len(r_scale) // 2])\nprint(f\"r target : {r_target:.5f} {r_scale.unit}\")\n\ndata_vs, r_vs, t_vs, p_vs = reader.vslice(r_target, unit='Gauss')\nprint(f\"vslice shape : {data_vs.shape}  (r axis collapsed to 1)\")\nprint(f\"r value      : {r_vs[0]:.5f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Passing an :class:`~astropy.units.Quantity` allows specifying the coordinate\nin any compatible unit.  Here we use :data:`~psi_io._units.PSI_rsun`, the\nsolar radius constant used internally by MAS coordinate scales:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "r_qty = r_scale[len(r_scale) // 2].to(u.cm)\ndata_vsq, r_vsq, t_vsq, p_vsq = reader.vslice(r_qty, unit='Gauss')\nprint(f\"input (Qty) r  : {r_qty:.4e}\")\nprint(f\"vslice (Qty) r : {r_vsq[0]:.4f}\")\nnp.testing.assert_allclose(data_vs.value, data_vsq.value, rtol=1e-5)\nprint(\"Scalar and Quantity results match.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Mixing value-space and index-space arguments**\n\n``vslice`` accepts any combination of value-space (Quantity / scalar) and\nindex-space (``slice``, ``int``, ``None``) arguments, one per spatial axis\nin ``(r, \u03b8, \u03c6)`` order.  The example below fixes the radial coordinate at\nthe chosen target value and reads only the first five co-latitude grid points:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "data_mixed, r_mixed, t_mixed, p_mixed = reader.vslice(\n    r_qty,           # r: value-space (PSI_rsun) \u2192 collapses to 1\n    slice(0, 5),     # \u03b8: index-space \u2192 first 5 points\n    None,            # \u03c6: all points\n    unit='Gauss',\n)\nprint(f\"mixed shape : {data_mixed.shape}  (r=1; \u03b8=5; \u03c6=full)\")\nprint(f\"r           : {r_mixed[0]:.4f}\")\nprint(f\"\u03b8 range     : [{t_mixed[0]:.4f}, {t_mixed[-1]:.4f}]\")"
      ]
    }
  ],
  "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.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}