{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Computing a Derived Field and Round-Tripping it as a Custom Model\n\nBuild a derived quantity from several MAS fields, write it to a new HDF file with\nthe metadata attributes that :func:`~psi_io.mhd_io.PsiData` needs, and read it back\nas a ``model='custom'`` reader \u2014 relying entirely on the attributes stored in the\nfile to reconstruct the metadata.\n\nThis example demonstrates:\n\n1. Reading the three magnetic-field components ``br``, ``bt``, ``bp`` onto a common\n mesh so they can be combined.\n2. Computing the field magnitude $|B| = \\sqrt{B_r^2 + B_\\theta^2 + B_\\phi^2}$.\n3. Writing the result with :func:`~psi_io.psi_io.write_hdf_data`, attaching the\n metadata attributes (``name``, ``desc``, ``unit``, ``scalar``, ``mesh``,\n ``order``, ``scales``) that describe a PSI-style dataset.\n4. Reading it back with the default ``model='custom'`` and confirming the file's\n own attributes supply every piece of metadata \u2014 no model table required.\n\n

Note

The three components live on *different* half-meshes (``br`` is half-mesh in\n $r$, ``bt`` in $\\theta$, ``bp`` in $\\phi$). They cannot be\n combined element-wise until they share a grid, so every component is first\n remeshed to the common **main** mesh.

\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import tempfile\nfrom pathlib import Path\n\nimport numpy as np\nimport warnings\n\nfrom psi_data import fetch_mas_data\nfrom psi_io import PsiData, write_hdf_data\nfrom psi_io.mhd_io import MetaDataWarning" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Reading the components onto a common mesh\n\nEach magnetic-field component is stored on its own staggered (Yee) grid: ``br`` is\nface-centred in $r$, ``bt`` in $\\theta$, and ``bp`` in $\\phi$.\nTheir native shapes therefore differ by one point along the staggered axis. Passing\n``mesh='main'`` to :meth:`~psi_io.mhd_io.PsiData.read` averages each half-mesh axis\nonto the main mesh, so all three end up on the *same* grid and can be combined.\nWe read every component in Gauss.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "mas_files = fetch_mas_data(domains='cor', variables='br,bt,bp')\n\ncomponents = {}\nscales = None\nfor comp in ('br', 'bt', 'bp'):\n reader = PsiData(getattr(mas_files, f'cor_{comp}'), model='mas')\n native_mesh = reader.mesh\n data, r, t, p = reader.read(mesh='main', unit='Gauss')\n components[comp] = data\n scales = (r, t, p) # main-mesh scales are shared across components\n print(f\"{comp}: native mesh={native_mesh} \u2192 main-mesh shape={data.shape}\")\n\nprint(f\"shapes aligned: {components['br'].shape == components['bt'].shape == components['bp'].shape}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Computing the field magnitude\n\nWith every component on the common main mesh, the magnitude is a plain\nelement-wise reduction. Because the inputs are :class:`~astropy.units.Quantity`\nobjects in Gauss, the result carries Gauss units automatically.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "bmag = np.sqrt(components['br'] ** 2 + components['bt'] ** 2 + components['bp'] ** 2)\nprint(f\"|B| : shape={bmag.shape}, unit={bmag.unit}, max={bmag.max():.3f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Writing the derived field with metadata\n\n:func:`~psi_io.psi_io.write_hdf_data` writes the data array and its coordinate\nscales; any extra keyword arguments are attached to the dataset as HDF attributes.\nTo make the file self-describing for a later ``model='custom'`` read, we attach\nthe full metadata schema:\n\n``name`` / ``desc``\n Quantity identifier and human-readable description.\n``unit``\n The physical unit of the stored values (``'G'``). A custom reader has no\n normalization table to fall back on, so the unit *must* be stored here.\n``scalar``\n ``True`` \u2014 $|B|$ is a scalar field.\n``mesh``\n Stagger code. We write the ``'main'`` shorthand string rather than an integer\n code: :meth:`~psi_io.mesh.Mesh.parse` accepts the ``'main'``/``'half'``\n shorthands and per-axis token sequences, but an integer attribute read back\n from HDF returns as a NumPy integer, which the parser does not accept.\n``order``\n ``'F'`` \u2014 the data array follows PSI's Fortran (column-major) convention.\n``scales``\n The ordered axis-name tuple ``('r', 't', 'p')``. The scale readers use these\n names to resolve their own units from :mod:`psi_io.units`.\n\nNote the data array is written in HDF storage order ``(N\u03c6, N\u03b8, Nr)`` with the\nscales supplied in physical ``(r, t, p)`` order, exactly as the reader returned them.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "r, t, p = scales\noutfile = Path(tempfile.gettempdir()) / \"bmag_custom.h5\"\nwrite_hdf_data(\n outfile,\n bmag.value,\n r.value, t.value, p.value,\n name='bmag',\n desc='Magnetic field magnitude',\n unit='G',\n scalar=True,\n mesh='main',\n order='F',\n scales=('r', 't', 'p'),\n)\nprint(f\"wrote {outfile.name}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Reading it back as a custom model\n\n:func:`~psi_io.mhd_io.PsiData` defaults to ``model='custom'``, which infers *no*\nmetadata from any model table \u2014 it relies solely on the dataset attributes (and,\nfor anything still missing, explicit keyword arguments). Because we stored the\ncomplete schema above, the custom reader reconstructs the quantity name,\ndescription, unit, mesh, ordering, and coordinate scales entirely from the file.\n\nA single :class:`~psi_io.mhd_io.MetaDataWarning` is expected: ``'custom'`` is\ndeliberately not a recognized model, so the reader notes that it is trusting the\nfile-supplied metadata rather than a model table. This is informational, not an error.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "with warnings.catch_warnings(record=True) as caught:\n warnings.simplefilter(\"always\")\n custom = PsiData(outfile) # model='custom' by default\n\nprint(repr(custom))\nprint(f\"model : {custom.model}\")\nprint(f\"unit : {custom.unit} (read from the file's 'unit' attribute)\")\nprint(f\"mesh : {custom.mesh}\")\nprint(f\"order : {custom.order}\")\nprint(f\"scales : {[s.name for s in custom.scales]} units={[str(s.unit) for s in custom.scales]}\")\nfor w in caught:\n if issubclass(w.category, MetaDataWarning):\n print(f\"METADATA WARNING: {w.message}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Confirming the round trip\n\nReading the data back reproduces both the array and its coordinate scales exactly,\nconfirming that the attributes written to the file carried all the information the\ncustom reader needed.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "data_rt, r_rt, t_rt, p_rt = custom.read()\nprint(f\"data round-trip max |\u0394| : {np.abs(data_rt.value - bmag.value).max():.3e}\")\nprint(f\"r-scale round-trip max |\u0394|: {np.abs(r_rt.value - r.value).max():.3e}\")\nprint(f\"recovered unit / shape : {data_rt.unit}, {data_rt.shape}\")" ] } ], "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.9" } }, "nbformat": 4, "nbformat_minor": 0 }