.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "gallery/03_mhd_io/p04_writing_custom_fields.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_gallery_03_mhd_io_p04_writing_custom_fields.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_gallery_03_mhd_io_p04_writing_custom_fields.py:


Computing a Derived Field and Round-Tripping it as a Custom Model
=================================================================

Build a derived quantity from several MAS fields, write it to a new HDF file with
the metadata attributes that :func:`~psi_io.mhd_io.PsiData` needs, and read it back
as a ``model='custom'`` reader — relying entirely on the attributes stored in the
file to reconstruct the metadata.

This example demonstrates:

1. Reading the three magnetic-field components ``br``, ``bt``, ``bp`` onto a common
   mesh so they can be combined.
2. Computing the field magnitude :math:`|B| = \sqrt{B_r^2 + B_\theta^2 + B_\phi^2}`.
3. Writing the result with :func:`~psi_io.psi_io.write_hdf_data`, attaching the
   metadata attributes (``name``, ``desc``, ``unit``, ``scalar``, ``mesh``,
   ``order``, ``scales``) that describe a PSI-style dataset.
4. Reading it back with the default ``model='custom'`` and confirming the file's
   own attributes supply every piece of metadata — no model table required.

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

.. GENERATED FROM PYTHON SOURCE LINES 27-38

.. code-block:: Python


    import tempfile
    from pathlib import Path

    import numpy as np
    import warnings

    from psi_data import fetch_mas_data
    from psi_io import PsiData, write_hdf_data
    from psi_io.mhd_io import MetaDataWarning








.. GENERATED FROM PYTHON SOURCE LINES 39-48

Reading the components onto a common mesh
-----------------------------------------

Each magnetic-field component is stored on its own staggered (Yee) grid: ``br`` is
face-centred in :math:`r`, ``bt`` in :math:`\theta`, and ``bp`` in :math:`\phi`.
Their native shapes therefore differ by one point along the staggered axis. Passing
``mesh='main'`` to :meth:`~psi_io.mhd_io.PsiData.read` averages each half-mesh axis
onto the main mesh, so all three end up on the *same* grid and can be combined.
We read every component in Gauss.

.. GENERATED FROM PYTHON SOURCE LINES 48-63

.. code-block:: Python


    mas_files = fetch_mas_data(domains='cor', variables='br,bt,bp')

    components = {}
    scales = None
    for comp in ('br', 'bt', 'bp'):
        reader = PsiData(getattr(mas_files, f'cor_{comp}'), model='mas')
        native_mesh = reader.mesh
        data, r, t, p = reader.read(mesh='main', unit='Gauss')
        components[comp] = data
        scales = (r, t, p)  # main-mesh scales are shared across components
        print(f"{comp}: native mesh={native_mesh}  →  main-mesh shape={data.shape}")

    print(f"shapes aligned: {components['br'].shape == components['bt'].shape == components['bp'].shape}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      0%|                                              | 0.00/43.3M [00:00<?, ?B/s]
      0%|▏                                     | 147k/43.3M [00:00<00:32, 1.34MB/s]
      1%|▍                                     | 492k/43.3M [00:00<00:17, 2.38MB/s]
      3%|█                                    | 1.31M/43.3M [00:00<00:08, 4.67MB/s]
      8%|██▉                                  | 3.38M/43.3M [00:00<00:03, 10.2MB/s]
     21%|███████▋                             | 8.96M/43.3M [00:00<00:01, 24.8MB/s]
     40%|██████████████▉                      | 17.4M/43.3M [00:00<00:00, 43.8MB/s]
     59%|█████████████████████▊               | 25.5M/43.3M [00:00<00:00, 53.3MB/s]
     78%|████████████████████████████▉        | 33.9M/43.3M [00:00<00:00, 62.8MB/s]
     93%|██████████████████████████████████▍  | 40.3M/43.3M [00:00<00:00, 62.8MB/s]
      0%|                                              | 0.00/43.3M [00:00<?, ?B/s]
    100%|██████████████████████████████████████| 43.3M/43.3M [00:00<00:00, 239GB/s]

      0%|                                              | 0.00/43.5M [00:00<?, ?B/s]
      0%|▏                                     | 164k/43.5M [00:00<00:30, 1.44MB/s]
      1%|▍                                     | 475k/43.5M [00:00<00:19, 2.21MB/s]
      3%|█▏                                   | 1.39M/43.5M [00:00<00:08, 4.88MB/s]
      8%|███                                  | 3.54M/43.5M [00:00<00:03, 10.3MB/s]
     22%|████████                             | 9.43M/43.5M [00:00<00:01, 26.4MB/s]
     40%|██████████████▉                      | 17.5M/43.5M [00:00<00:00, 43.8MB/s]
     57%|█████████████████████▏               | 24.8M/43.5M [00:00<00:00, 52.5MB/s]
     76%|████████████████████████████▏        | 33.2M/43.5M [00:00<00:00, 61.8MB/s]
     95%|███████████████████████████████████▎ | 41.4M/43.5M [00:00<00:00, 68.1MB/s]
      0%|                                              | 0.00/43.5M [00:00<?, ?B/s]
    100%|██████████████████████████████████████| 43.5M/43.5M [00:00<00:00, 252GB/s]
    br: native mesh=HALF, MAIN, MAIN  →  main-mesh shape=(299, 142, 254)
    bt: native mesh=MAIN, HALF, MAIN  →  main-mesh shape=(299, 142, 254)
    bp: native mesh=MAIN, MAIN, HALF  →  main-mesh shape=(299, 142, 254)
    shapes aligned: True




.. GENERATED FROM PYTHON SOURCE LINES 64-70

Computing the field magnitude
-----------------------------

With every component on the common main mesh, the magnitude is a plain
element-wise reduction.  Because the inputs are :class:`~astropy.units.Quantity`
objects in Gauss, the result carries Gauss units automatically.

.. GENERATED FROM PYTHON SOURCE LINES 70-74

.. code-block:: Python


    bmag = np.sqrt(components['br'] ** 2 + components['bt'] ** 2 + components['bp'] ** 2)
    print(f"|B| : shape={bmag.shape}, unit={bmag.unit}, max={bmag.max():.3f}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    |B| : shape=(299, 142, 254), unit=G, max=86.892 G




.. GENERATED FROM PYTHON SOURCE LINES 75-103

Writing the derived field with metadata
---------------------------------------

:func:`~psi_io.psi_io.write_hdf_data` writes the data array and its coordinate
scales; any extra keyword arguments are attached to the dataset as HDF attributes.
To make the file self-describing for a later ``model='custom'`` read, we attach
the full metadata schema:

``name`` / ``desc``
    Quantity identifier and human-readable description.
``unit``
    The physical unit of the stored values (``'G'``).  A custom reader has no
    normalization table to fall back on, so the unit *must* be stored here.
``scalar``
    ``True`` — :math:`|B|` is a scalar field.
``mesh``
    Stagger code. We write the ``'main'`` shorthand string rather than an integer
    code: :meth:`~psi_io.mesh.Mesh.parse` accepts the ``'main'``/``'half'``
    shorthands and per-axis token sequences, but an integer attribute read back
    from HDF returns as a NumPy integer, which the parser does not accept.
``order``
    ``'F'`` — the data array follows PSI's Fortran (column-major) convention.
``scales``
    The ordered axis-name tuple ``('r', 't', 'p')``.  The scale readers use these
    names to resolve their own units from :mod:`psi_io.units`.

Note the data array is written in HDF storage order ``(Nφ, Nθ, Nr)`` with the
scales supplied in physical ``(r, t, p)`` order, exactly as the reader returned them.

.. GENERATED FROM PYTHON SOURCE LINES 103-120

.. code-block:: Python


    r, t, p = scales
    outfile = Path(tempfile.gettempdir()) / "bmag_custom.h5"
    write_hdf_data(
        outfile,
        bmag.value,
        r.value, t.value, p.value,
        name='bmag',
        desc='Magnetic field magnitude',
        unit='G',
        scalar=True,
        mesh='main',
        order='F',
        scales=('r', 't', 'p'),
    )
    print(f"wrote {outfile.name}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    wrote bmag_custom.h5




.. GENERATED FROM PYTHON SOURCE LINES 121-133

Reading it back as a custom model
---------------------------------

:func:`~psi_io.mhd_io.PsiData` defaults to ``model='custom'``, which infers *no*
metadata from any model table — it relies solely on the dataset attributes (and,
for anything still missing, explicit keyword arguments).  Because we stored the
complete schema above, the custom reader reconstructs the quantity name,
description, unit, mesh, ordering, and coordinate scales entirely from the file.

A single :class:`~psi_io.mhd_io.MetaDataWarning` is expected: ``'custom'`` is
deliberately not a recognized model, so the reader notes that it is trusting the
file-supplied metadata rather than a model table. This is informational, not an error.

.. GENERATED FROM PYTHON SOURCE LINES 133-148

.. code-block:: Python


    with warnings.catch_warnings(record=True) as caught:
        warnings.simplefilter("always")
        custom = PsiData(outfile)  # model='custom' by default

    print(repr(custom))
    print(f"model  : {custom.model}")
    print(f"unit   : {custom.unit}   (read from the file's 'unit' attribute)")
    print(f"mesh   : {custom.mesh}")
    print(f"order  : {custom.order}")
    print(f"scales : {[s.name for s in custom.scales]}  units={[str(s.unit) for s in custom.scales]}")
    for w in caught:
        if issubclass(w.category, MetaDataWarning):
            print(f"METADATA WARNING: {w.message}")





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    H5Data(name='bmag' [Magnetic field magnitude], order='F', shape=(254, 142, 299), unit=Unit("G"), mesh=Mesh(MAIN, MAIN, MAIN), cached=False)
    model  : custom
    unit   : G   (read from the file's 'unit' attribute)
    mesh   : MAIN, MAIN, MAIN
    order  : F
    scales : ['r', 't', 'p']  units=['PSI_rsun', 'PSI_angle', 'PSI_angle']
    METADATA WARNING: H5Data(bmag) has an unrecognized model 'custom'. Ensure metadata is declared at instantiation or written to the HDF dataset's attribute mapping.




.. GENERATED FROM PYTHON SOURCE LINES 149-155

Confirming the round trip
-------------------------

Reading the data back reproduces both the array and its coordinate scales exactly,
confirming that the attributes written to the file carried all the information the
custom reader needed.

.. GENERATED FROM PYTHON SOURCE LINES 155-160

.. code-block:: Python


    data_rt, r_rt, t_rt, p_rt = custom.read()
    print(f"data round-trip max |Δ| : {np.abs(data_rt.value - bmag.value).max():.3e}")
    print(f"r-scale round-trip max |Δ|: {np.abs(r_rt.value - r.value).max():.3e}")
    print(f"recovered unit / shape    : {data_rt.unit}, {data_rt.shape}")




.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    data round-trip max |Δ| : 0.000e+00
    r-scale round-trip max |Δ|: 0.000e+00
    recovered unit / shape    : G, (299, 142, 254)





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (0 minutes 2.884 seconds)


.. _sphx_glr_download_gallery_03_mhd_io_p04_writing_custom_fields.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: p04_writing_custom_fields.ipynb <p04_writing_custom_fields.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: p04_writing_custom_fields.py <p04_writing_custom_fields.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: p04_writing_custom_fields.zip <p04_writing_custom_fields.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_