1r"""Physical property metadata for PSI's model quantities and scales.
2
3This module provides the :class:`ScaleProps` and :class:`ModelProps` dataclasses and
4three read-only mapping objects that describe every quantity PSI's modeling codes
5(*viz.* for MAS and POT3D) write to HDF files.
6
7For metadata access, one should use the provided getter functions rather than the mapping
8objects directly.
9
10
11.. list-table::
12 :header-rows: 1
13
14 * - Mapping
15 - Model
16 - Contents
17 * - :func:`get_mas_quantity_properties`
18 - `MAS <https://github.com/predsci/MAS>`_ (Magnetohydrodynamic Algorithm outside a Sphere)
19 - 19 3-D output fields: velocity, magnetic field, current density, temperature,
20 density, pressure, wave energy, Elsässer amplitudes, heating
21 * - :func:`get_pot3d_quantity_properties`
22 - `POT3D <https://github.com/predsci/POT3D>`_ (High Performance Potential Field Solver)
23 - 3 magnetic field components (br, bt, bp)
24 * - :func:`get_psi_scale_properties`
25 - PSI coordinate grids (shared by MAS, POT3D, and related codes)
26 - 3 1-D coordinate scale arrays (r, t, p)
27
28Coordinate scales
29-----------------
30Every PSI HDF file stores three 1-D coordinate arrays alongside its data. These
31arrays define the spherical grid on which model output is sampled and are standard
32across MAS, POT3D, and most other PSI codes:
33
34.. list-table::
35 :header-rows: 1
36 :widths: 10 12 28 24 16
37
38 * - Key
39 - Symbol
40 - Physical quantity
41 - Units
42 - Range
43 * - ``'r'``
44 - :math:`r`
45 - Radial distance
46 - solar radii (:data:`~psi_io.units.PSI_rsun`)
47 - :math:`r \geq 1\,R_\odot`
48 * - ``'t'``
49 - :math:`\theta`
50 - Co-latitude (pole = 0)
51 - radians (:data:`~psi_io.units.PSI_angle`)
52 - :math:`[0,\,\pi]`
53 * - ``'p'``
54 - :math:`\varphi`
55 - Longitude
56 - radians (:data:`~psi_io.units.PSI_angle`)
57 - :math:`[0,\,2\pi]`
58
59PSI HDF files are written in **Fortran order** (column-major), so the *first* Fortran
60dimension varies fastest in memory. When read into NumPy (row-major), the axis order
61is reversed: a 3-D array has shape :math:`(N_\varphi, N_\theta, N_r)` and the three
62scale arrays map to axes as follows:
63
64.. list-table::
65 :header-rows: 1
66
67 * - Key
68 - Fortran dimension
69 - NumPy axis
70 - ``arr.shape`` index
71 * - ``'r'``
72 - 1st (fastest-varying)
73 - last
74 - ``-1``
75 * - ``'t'``
76 - 2nd
77 - middle
78 - ``-2``
79 * - ``'p'``
80 - 3rd (slowest-varying)
81 - first
82 - ``0``
83
84
85MAS Model Quantities
86--------------------
87MAS outputs 19 3-D fields in spherical coordinates. Code-unit values are converted
88to physical (CGS/SI) unit by multiplying by the normalization constants defined in
89:mod:`psi_io.units`. Approximate physical scales are given in parentheses.
90
91.. list-table::
92 :header-rows: 1
93 :widths: 8 12 30 28 10 10
94
95 * - Key
96 - Symbol
97 - Physical quantity
98 - Physical unit (CGS)
99 - Type
100 - Mesh code
101 * - ``vr``
102 - :math:`v_r`
103 - Radial velocity
104 - km s⁻¹ (:data:`~psi_io.units.FN_V` ≈ 481 km s⁻¹)
105 - vector
106 - ``0b011``
107 * - ``vt``
108 - :math:`v_\theta`
109 - Co-latitude velocity
110 - km s⁻¹
111 - vector
112 - ``0b101``
113 * - ``vp``
114 - :math:`v_\varphi`
115 - Longitude velocity
116 - km s⁻¹
117 - vector
118 - ``0b110``
119 * - ``br``
120 - :math:`B_r`
121 - Radial magnetic field
122 - G (:data:`~psi_io.units.FN_B` ≈ 2.2 G)
123 - vector
124 - ``0b100``
125 * - ``bt``
126 - :math:`B_\theta`
127 - Co-latitude magnetic field
128 - G
129 - vector
130 - ``0b010``
131 * - ``bp``
132 - :math:`B_\varphi`
133 - Longitude magnetic field
134 - G
135 - vector
136 - ``0b001``
137 * - ``jr``
138 - :math:`J_r`
139 - Radial current density
140 - A m⁻² (:data:`~psi_io.units.FN_J`)
141 - vector
142 - ``0b011``
143 * - ``jt``
144 - :math:`J_\theta`
145 - Co-latitude current density
146 - A m⁻²
147 - vector
148 - ``0b101``
149 * - ``jp``
150 - :math:`J_\varphi`
151 - Longitude current density
152 - A m⁻²
153 - vector
154 - ``0b110``
155 * - ``t``
156 - :math:`T`
157 - Single-fluid temperature
158 - MK (:data:`~psi_io.units.FN_T` ≈ 28 MK)
159 - scalar
160 - ``0b111``
161 * - ``te``
162 - :math:`T_e`
163 - Electron temperature
164 - MK
165 - scalar
166 - ``0b111``
167 * - ``tp``
168 - :math:`T_p`
169 - Proton temperature
170 - MK
171 - scalar
172 - ``0b111``
173 * - ``rho``
174 - :math:`\rho`
175 - Plasma density
176 - cm⁻³ (:data:`~psi_io.units.FN_N` = 10⁸ cm⁻³)
177 - scalar
178 - ``0b111``
179 * - ``p``
180 - :math:`p`
181 - Plasma pressure
182 - erg cm⁻³ (:data:`~psi_io.units.FN_P`)
183 - scalar
184 - ``0b111``
185 * - ``ep``
186 - :math:`e^+`
187 - Forward Alfvén wave energy density
188 - erg cm⁻³
189 - scalar
190 - ``0b111``
191 * - ``em``
192 - :math:`e^-`
193 - Backward Alfvén wave energy density
194 - erg cm⁻³
195 - scalar
196 - ``0b111``
197 * - ``zp``
198 - :math:`z^+`
199 - Outward Elsässer amplitude
200 - km s⁻¹
201 - scalar
202 - ``0b111``
203 * - ``zm``
204 - :math:`z^-`
205 - Inward Elsässer amplitude
206 - km s⁻¹
207 - scalar
208 - ``0b111``
209 * - ``heat``
210 - :math:`Q`
211 - Volumetric heating rate
212 - erg cm⁻³ s⁻¹ (:data:`~psi_io.units.FN_HEAT`)
213 - scalar
214 - ``0b111``
215
216POT3D Model Quantities
217----------------------
218POT3D solves for the potential magnetic field :math:`\mathbf{B} = -\nabla\Psi` driven
219by a photospheric radial-field boundary condition. It outputs three spherical
220magnetic field components:
221
222.. list-table::
223 :header-rows: 1
224 :widths: 10 12 35 25 18
225
226 * - Key
227 - Symbol
228 - Physical quantity
229 - Physical unit
230 - Mesh code
231 * - ``br``
232 - :math:`B_r`
233 - Radial magnetic field
234 - input magnetogram unit (typically G)
235 - ``0b011``
236 * - ``bt``
237 - :math:`B_\theta`
238 - Co-latitude magnetic field
239 - input magnetogram unit
240 - ``0b101``
241 * - ``bp``
242 - :math:`B_\varphi`
243 - Longitude magnetic field
244 - input magnetogram unit
245 - ``0b110``
246
247.. note::
248
249 POT3D mesh codes are the bitwise complement of the corresponding MAS magnetic field
250 codes: ``POT3D_mesh = 0b111 ^ MAS_mesh``. Where MAS places each component on the
251 face through which it is the outward normal (face-centred), POT3D places the same
252 component on the opposite pair of edges (edge-centred).
253
254.. warning::
255
256 POT3D does not apply a normalization: values are in whatever physical unit the
257 input boundary magnetogram was provided in (typically Gauss, but run-dependent).
258 The code unit :data:`~psi_io.units.POT3D_b` has a scale factor of 1
259 (dimensionless-unscaled). Always supply the correct unit explicitly when
260 converting; ``read(unit='physical')`` without a ``unit`` override will return
261 dimensionless values.
262
263Staggered Grid Overview
264-----------------------
265MAS and POT3D solve their equations on a three-dimensional **staggered** (Yee-type)
266spherical grid :math:`(r, \theta, \varphi)`. Different physical quantities are
267located at different positions within each grid cell so that discrete differential
268operators (curl, divergence) are exactly satisfied.
269
270With the convention used in this module, numpy arrays loaded from PSI HDF files have
271shape :math:`(N_{\phi}, N_{\theta}, N_r)` — longitude index varies slowest, radial index
272fastest. Mesh codes therefore assign bits as follows:
273
274.. list-table::
275 :header-rows: 1
276
277 * - Bit position
278 - Numpy axis
279 - Coordinate
280 * - Most-significant bit (MSB)
281 - ``axis = -1`` (last)
282 - :math:`r` (radial)
283 * - Middle bit
284 - ``axis = -2``
285 - :math:`\theta` (co-latitude)
286 * - Least-significant bit (LSB)
287 - ``axis = -3`` (first)
288 - :math:`\varphi` (longitude)
289
290A bit value of ``1`` means the quantity is on the **half mesh** along that axis
291(displaced half a grid spacing), while ``0`` means it is on the **main mesh**.
292
293The resulting grid positions for MAS vector components are:
294
295.. list-table::
296 :header-rows: 1
297
298 * - Quantity type
299 - Mesh code
300 - :math:`r`
301 - :math:`\theta`
302 - :math:`\varphi`
303 - Grid position
304 * - ``br``
305 - ``0b100``
306 - half
307 - main
308 - main
309 - :math:`r`-face center
310 * - ``bt``
311 - ``0b010``
312 - main
313 - half
314 - main
315 - :math:`\theta`-face center
316 * - ``bp``
317 - ``0b001``
318 - main
319 - main
320 - half
321 - :math:`\varphi`-face center
322 * - ``vr``, ``jr``
323 - ``0b011``
324 - main
325 - half
326 - half
327 - :math:`r`-edge center
328 * - ``vt``, ``jt``
329 - ``0b101``
330 - half
331 - main
332 - half
333 - :math:`\theta`-edge center
334 * - ``vp``, ``jp``
335 - ``0b110``
336 - half
337 - half
338 - main
339 - :math:`\varphi`-edge center
340 * - Scalars (``t``, ``rho``, ``p``, …)
341 - ``0b111``
342 - half
343 - half
344 - half
345 - Cell corner
346
347This Yee-type arrangement ensures that :math:`\nabla \cdot \mathbf{B} = 0` is
348satisfied exactly at the discrete level: each :math:`B` component lives on the face
349through which it is the outward normal. The current density components
350:math:`\mathbf{J} = \nabla \times \mathbf{B}` then live on the corresponding cell
351edges. Scalar thermodynamic quantities (density, temperature, pressure) occupy the
352cell corners, which are equivalent to cell centers on the dual grid.
353
354See Also
355--------
356
357:mod:`psi_io.mesh` :
358 Defines :class:`~psi_io.mesh.Mesh` and the mesh normalization helpers.
359:mod:`psi_io.units` :
360 Provides the custom astropy unit referenced by each :class:`ScaleProps` and
361 :class:`ModelProps`.
362:mod:`psi_io.mhd_io` :
363 Uses these mappings to auto-configure lazy HDF readers.
364"""
365
366from __future__ import annotations
367
368__all__ = [
369 "get_mas_quantity_properties",
370 "get_pot3d_quantity_properties",
371 "get_psi_scale_properties",
372 "parse_psi_filename_schema",
373 "extract_quantity_from_filepath",
374 "extract_sequence_from_filepath"
375]
376
377import re
378from dataclasses import dataclass, asdict
379from pathlib import Path
380from types import MappingProxyType
381from typing import Literal, Optional, Callable
382
383import astropy.units as u
384
385from psi_io.mesh import ArrayOrdering, Mesh
386from psi_io.units import MAS_v, MAS_b, MAS_j, MAS_t, MAS_n, MAS_p, MAS_heat, POT3D_b, PSI_rsun, PSI_angle
387
388
[docs]
389@dataclass(frozen=True, repr=True)
390class ScaleProps:
391 """Immutable property bundle for a PSI coordinate scale array.
392
393 Associates a coordinate axis name with its human-readable description and
394 physical unit. Used for the three standard PSI spatial scales ``'r'``,
395 ``'t'``, and ``'p'``. Instances are frozen (immutable) dataclass instances.
396
397 Parameters
398 ----------
399 name : str
400 Canonical lower-case coordinate identifier (``'r'``, ``'t'``, or ``'p'``).
401 desc : str
402 Human-readable description (e.g. ``'PSI Radial Scale (Solar Radii)'``).
403 unit : Unit
404 Astropy unit for values stored along this coordinate axis. For example,
405 :data:`~psi_io.units.PSI_rsun` for the radial scale.
406
407 Notes
408 -----
409 The arithmetic dunder methods (``__mul__``, ``__rmul__``, ``__rtruediv__``)
410 delegate to :attr:`unit`, so ``value * props`` is equivalent to
411 ``value * props.unit`` and returns an :class:`~astropy.units.Quantity`.
412
413 See Also
414 --------
415 ModelProps : Extends :class:`ScaleProps` with scalar, mesh, and ordering fields
416 for 3-D model quantities.
417 get_psi_scale_properties : Retrieve a :class:`ScaleProps` by coordinate name.
418
419 Examples
420 --------
421 >>> from psi_io.models import ScaleProps
422 >>> import astropy.units as u
423 >>> p = ScaleProps('r', 'PSI Radial Scale (Solar Radii)', u.R_sun)
424 >>> str(p)
425 'r'
426 >>> p.name, p.desc
427 ('r', 'PSI Radial Scale (Solar Radii)')
428 >>> (1.0 * p).unit
429 Unit("solRad")
430 """
431
432 name: str
433 desc: str
434 unit: u.Unit
435
436 def __str__(self) -> str:
437 """Return the quantity name.
438
439 Returns
440 -------
441 out : str
442 The :attr:`name` field (e.g. ``'br'``, ``'vr'``).
443
444 Examples
445 --------
446 >>> from psi_io.models import _MAS_QUANTITY_PROPS_MAPPING
447 >>> str(_MAS_QUANTITY_PROPS_MAPPING['br'])
448 'br'
449 """
450 return self.name
451
452 def __mul__(self, other):
453 """Multiply *other* by this quantity's unit.
454
455 Allows ``props * value`` expressions that return an astropy
456 :class:`~astropy.units.Quantity`.
457
458 Parameters
459 ----------
460 other : numeric | Quantity
461 The value to multiply by :attr:`unit`.
462
463 Returns
464 -------
465 out : Quantity
466 ``other * self.unit``.
467
468 Examples
469 --------
470 >>> from psi_io.models import _MAS_QUANTITY_PROPS_MAPPING
471 >>> props = _MAS_QUANTITY_PROPS_MAPPING['vr']
472 >>> (props * 1.0).unit # doctest: +ELLIPSIS
473 Unit(...)
474 """
475 return other * self.unit
476
477 def __rmul__(self, other):
478 """Right-multiply *other* by this quantity's unit.
479
480 Allows ``value * props`` expressions that return an astropy
481 :class:`~astropy.units.Quantity`.
482
483 Parameters
484 ----------
485 other : numeric or Quantity
486 The value to multiply by :attr:`unit`.
487
488 Returns
489 -------
490 out : Quantity
491 ``other * self.unit``.
492
493 Examples
494 --------
495 >>> from psi_io.models import _MAS_QUANTITY_PROPS_MAPPING
496 >>> props = _MAS_QUANTITY_PROPS_MAPPING['vr']
497 >>> (1.0 * props).unit # doctest: +ELLIPSIS
498 Unit(...)
499 """
500 return other * self.unit
501
502 def __rtruediv__(self, other):
503 """Divide *other* by this quantity's unit.
504
505 Allows ``value / props`` expressions that return an astropy
506 :class:`~astropy.units.Quantity`.
507
508 Parameters
509 ----------
510 other : numeric | Quantity
511 The numerator; divided by :attr:`unit`.
512
513 Returns
514 -------
515 out : Quantity
516 ``other / self.unit``.
517
518 Examples
519 --------
520 >>> from psi_io.models import _MAS_QUANTITY_PROPS_MAPPING
521 >>> props = _MAS_QUANTITY_PROPS_MAPPING['br']
522 >>> (1.0 / props).unit # doctest: +ELLIPSIS
523 Unit(...)
524 """
525 return other / self.unit
526
[docs]
527 def _asdict(self):
528 """Return a plain dictionary of this instance's fields.
529
530 Returns
531 -------
532 out : dict[str, object]
533 Mapping of field names to their values.
534
535 Examples
536 --------
537 >>> from psi_io.models import ScaleProps
538 >>> import astropy.units as u
539 >>> p = ScaleProps('r', 'Radial', u.R_sun)
540 >>> list(p._asdict().keys())
541 ['name', 'desc', 'unit']
542 """
543 return asdict(self)
544
545
[docs]
546@dataclass(frozen=True, repr=True)
547class ModelProps(ScaleProps):
548 """Immutable property bundle for a single PSI model quantity.
549
550 Extends :class:`ScaleProps` with scalar/vector classification, staggered-grid mesh
551 code, array ordering, and scale axis labels. Used for MAS and POT3D output fields.
552 Instances are frozen (immutable) dataclass instances.
553
554 Parameters
555 ----------
556 name : str
557 Canonical lower-case quantity identifier (e.g. ``'br'``, ``'vr'``). Matches
558 the filename prefix used in MAS and POT3D HDF output.
559 desc : str
560 Human-readable description of the physical quantity.
561 unit : Unit
562 Astropy unit whose scale factor converts one code unit of this quantity to
563 physical unit. For example, :data:`~psi_io.units.MAS_b` ≈ 2.2 Gauss.
564 scalar : bool
565 ``True`` if the quantity is a scalar field (temperature, density, …);
566 ``False`` if it is a component of a vector field (velocity, magnetic field, …).
567 _mesh : int
568 Integer mesh code encoding the stagger position on the three-dimensional grid.
569 Each binary bit indicates whether the quantity is on the half mesh (``1``) or
570 main mesh (``0``) along one coordinate axis.
571 order : str, optional
572 Array memory layout: ``'F'`` for Fortran (column-major, PSI default) or
573 ``'C'`` for C (row-major). Default is ``'F'``.
574 scales : tuple[str, ...], optional
575 Ordered names of the coordinate scale axes associated with this quantity.
576 Default is ``('r', 't', 'p')``.
577
578 Attributes
579 ----------
580 mesh : Mesh
581 Normalized form of :attr:`_mesh` as a single :class:`~psi_io.mesh.Mesh`
582 instance built from the integer code and the number of coordinate scales.
583
584 Notes
585 -----
586 The arithmetic dunder methods (``__mul__``, ``__rmul__``, ``__rtruediv__``)
587 delegate to :attr:`unit`, so ``value * props`` is equivalent to
588 ``value * props.unit`` and returns an :class:`~astropy.units.Quantity`.
589
590 See Also
591 --------
592 ScaleProps : Base class for coordinate scale arrays without mesh or ordering.
593 get_mas_quantity_properties : Retrieve a :class:`ModelProps` for a MAS quantity.
594 get_pot3d_quantity_properties : Retrieve a :class:`ModelProps` for a POT3D quantity.
595
596 Examples
597 --------
598 >>> from psi_io.models import ModelProps
599 >>> import astropy.units as u
600 >>> p = ModelProps('br', 'Radial B field', u.Gauss, False, 0b100)
601 >>> str(p)
602 'br'
603 >>> p.scalar
604 False
605 >>> p.mesh
606 Mesh(HALF, MAIN, MAIN)
607 >>> (2.5 * p).unit
608 Unit("G")
609 """
610
611 name: str
612 desc: str
613 unit: u.Unit
614 scalar: bool
615 _mesh: int
616 order: ArrayOrdering = 'F'
617 scales: tuple = tuple('rtp')
618
619 @property
620 def mesh(self):
621 """Normalized mesh-stagger object for this quantity.
622
623 Wraps the integer :attr:`_mesh` code and the number of coordinate scales
624 in a single :class:`~psi_io.mesh.Mesh` instance via its constructor.
625
626 Returns
627 -------
628 out : Mesh
629 A :class:`~psi_io.mesh.Mesh` whose bits classify each spatial axis
630 ``(r, theta, phi)`` as half or main mesh, with the most-significant
631 bit mapping to :math:`r` (the last NumPy axis under PSI's Fortran
632 convention).
633
634 Examples
635 --------
636 >>> from psi_io.models import _MAS_QUANTITY_PROPS_MAPPING
637 >>> from psi_io.mesh import Mesh
638 >>> _MAS_QUANTITY_PROPS_MAPPING['br'].mesh
639 Mesh(HALF, MAIN, MAIN)
640 >>> _MAS_QUANTITY_PROPS_MAPPING['t'].mesh
641 Mesh(HALF, HALF, HALF)
642 >>> _MAS_QUANTITY_PROPS_MAPPING['vr'].mesh
643 Mesh(MAIN, HALF, HALF)
644 """
645 return Mesh(self._mesh, len(self.scales))
646
647 def _asdict(self):
648 """Return a plain dictionary of this instance's fields, with mesh resolved.
649
650 Replaces the raw ``_mesh`` integer with the normalized :attr:`mesh`
651 object before returning.
652
653 Returns
654 -------
655 out : dict[str, object]
656 Mapping of field names to their values; ``'_mesh'`` is replaced by
657 ``'mesh'`` containing the resolved :class:`~psi_io.mesh.Mesh` instance.
658
659 Examples
660 --------
661 >>> from psi_io.models import get_mas_quantity_properties
662 >>> d = get_mas_quantity_properties('br')._asdict()
663 >>> 'mesh' in d and '_mesh' not in d
664 True
665 """
666 dout = asdict(self)
667 dout.update(mesh=self.mesh)
668 del dout['_mesh']
669 return dout
670
671
672# ----------------------------------------------------------------
673# MAS Model
674# ----------------------------------------------------------------
675
676
677MasQuantities = Literal[
678 'vr', 'vt', 'vp', 'br', 'bt', 'bp', 'jr', 'jt', 'jp',
679 't', 'te', 'tp', 'rho', 'p', 'ep', 'em', 'zp', 'zm', 'heat'
680]
681"""Literal type alias for the 19 named MAS output quantities.
682
683Each string is the canonical lower-case identifier used as a key in
684:data:`_MAS_QUANTITY_PROPS_MAPPING` and as the filename prefix (e.g. ``br001001.h5``).
685
686Velocity components
687 ``'vr'``, ``'vt'``, ``'vp'`` — radial, co-latitude, and longitude velocity.
688
689Magnetic field components
690 ``'br'``, ``'bt'``, ``'bp'`` — radial, co-latitude, and longitude magnetic field.
691
692Current density components
693 ``'jr'``, ``'jt'``, ``'jp'`` — radial, co-latitude, and longitude current density.
694
695Temperature
696 ``'t'`` — single-fluid temperature; ``'te'`` — electron temperature;
697 ``'tp'`` — proton (ion) temperature (two-temperature model).
698
699Density and pressure
700 ``'rho'`` — mass density; ``'p'`` — plasma pressure.
701
702Alfvén wave quantities
703 ``'ep'``, ``'em'`` — wave energy density for forward- and backward-propagating
704 Alfvén waves (proportional to :math:`|z^+|^2` and :math:`|z^-|^2`);
705 ``'zp'``, ``'zm'`` — Elsässer variable amplitudes
706 :math:`z^\\pm = v \\pm v_A` for forward/backward propagating waves.
707
708Heating
709 ``'heat'`` — local coronal volumetric heating rate.
710"""
711
712_MAS_QUANTITY_PROPS_MAPPING = MappingProxyType({
713 'vr': ModelProps('vr', 'MAS Velocity (Radial Component)', MAS_v, False, 0b011),
714 'vt': ModelProps('vt', 'MAS Velocity (Theta Component)', MAS_v, False, 0b101),
715 'vp': ModelProps('vp', 'MAS Velocity (Phi Component)', MAS_v, False, 0b110),
716 'br': ModelProps('br', 'MAS Magnetic Field (Radial Component)', MAS_b, False, 0b100),
717 'bt': ModelProps('bt', 'MAS Magnetic Field (Theta Component)', MAS_b, False, 0b010),
718 'bp': ModelProps('bp', 'MAS Magnetic Field (Phi Component)', MAS_b, False, 0b001),
719 'jr': ModelProps('jr', 'MAS Current Density (Radial Component)', MAS_j, False, 0b011),
720 'jt': ModelProps('jt', 'MAS Current Density (Theta Component)', MAS_j, False, 0b101),
721 'jp': ModelProps('jp', 'MAS Current Density (Phi Component)', MAS_j, False, 0b110),
722 't': ModelProps('t', 'MAS Temperature', MAS_t, True, 0b111),
723 'te': ModelProps('te', 'MAS Electron Temperature', MAS_t, True, 0b111),
724 'tp': ModelProps('tp', 'MAS Proton Temperature', MAS_t, True, 0b111),
725 'rho': ModelProps('rho', 'MAS Density', MAS_n, True, 0b111),
726 'p': ModelProps('p', 'MAS Pressure', MAS_p, True, 0b111),
727 'ep': ModelProps('ep', 'MAS Wave Energy Density (Parallel to the Field)', MAS_p, True, 0b111),
728 'em': ModelProps('em', 'MAS Wave Energy Density (Anti-Parallel to the Field)', MAS_p, True, 0b111),
729 'zp': ModelProps('zp', 'MAS Outward Propagating Wave Amplitude', MAS_v, True, 0b111),
730 'zm': ModelProps('zm', 'MAS Inward Propagating Wave Amplitude', MAS_v, True, 0b111),
731 'heat': ModelProps('heat', 'MAS Local Coronal Heating Rate', MAS_heat, True, 0b111),
732})
733"""Read-only mapping from MAS quantity name to its :class:`ModelProps` descriptor."""
734
735
[docs]
736def get_mas_quantity_properties(variable: MasQuantities) -> ModelProps:
737 """Return the :class:`~psi_io.models.ModelProps` descriptor for a MAS quantity.
738
739 Parameters
740 ----------
741 variable : MasQuantities
742 Case-insensitive MAS quantity name. Valid values: ``'br'``, ``'bt'``, ``'bp'``,
743 ``'vr'``, ``'vt'``, ``'vp'``, ``'jr'``, ``'jt'``, ``'jp'``, ``'t'``, ``'te'``,
744 ``'tp'``, ``'rho'``, ``'p'``, ``'ep'``, ``'em'``, ``'zp'``, ``'zm'``, ``'heat'``.
745
746 Returns
747 -------
748 out : ModelProps
749 Immutable metadata descriptor for the requested MAS quantity.
750
751 Raises
752 ------
753 ValueError
754 If *variable* is not a recognized MAS quantity.
755
756 See Also
757 --------
758 get_pot3d_quantity_properties : Analogous lookup for POT3D quantities.
759 get_model_prop_caller : Return the getter function for a given model type.
760
761 Examples
762 --------
763 >>> from psi_io.models import get_mas_quantity_properties
764 >>> props = get_mas_quantity_properties('br')
765 >>> props.desc
766 'MAS Magnetic Field (Radial Component)'
767 >>> get_mas_quantity_properties('BR').name # case-insensitive
768 'br'
769 """
770 try:
771 return _MAS_QUANTITY_PROPS_MAPPING[variable.lower()]
772 except KeyError:
773 raise ValueError(f"Invalid variable {variable!r} for MAS model. "
774 f"Valid options are: {', '.join(_MAS_QUANTITY_PROPS_MAPPING.keys())}") from None
775
776
777# ----------------------------------------------------------------
778# POT3D Model
779# ----------------------------------------------------------------
780
781
782Pot3dQuantities = Literal['br', 'bt', 'bp',]
783"""Literal type alias for the 3 POT3D magnetic field output quantities.
784
785POT3D is a potential-field solver that outputs only the three spherical components of
786the magnetic field: ``'br'`` (radial), ``'bt'`` (co-latitude), and ``'bp'``
787(longitude).
788
789.. note::
790
791 POT3D mesh codes are the bitwise complement of the corresponding MAS magnetic field
792 codes: ``POT3D_mesh = 0b111 ^ MAS_mesh``. Where MAS places each component on the
793 face through which it is the outward normal (face-centred), POT3D places the same
794 component on the opposite pair of edges (edge-centred).
795
796.. warning::
797
798 POT3D output files carry **no intrinsic physical unit**. The values are stored as
799 dimensionless quantities (unit :data:`~psi_io.units.POT3D_b` = 1) whose physical
800 interpretation depends entirely on the unit of the photospheric boundary
801 magnetogram used to drive the simulation — typically Gauss, but this is not
802 guaranteed. Always supply the correct unit explicitly via the ``unit`` keyword
803 argument of :func:`~psi_io.mhd_io.PsiData`; calling ``read(unit='physical')``
804 on a reader opened without a ``unit`` override will **not** perform a meaningful
805 conversion.
806"""
807
808_POT3D_QUANTITY_PROPS_MAPPING = MappingProxyType({
809 'br': ModelProps('br', 'POT3D Magnetic Field (Radial Component)', POT3D_b, False, 0b011),
810 'bt': ModelProps('bt', 'POT3D Magnetic Field (Theta Component)', POT3D_b, False, 0b101),
811 'bp': ModelProps('bp', 'POT3D Magnetic Field (Phi Component)', POT3D_b, False, 0b110),
812})
813"""Read-only mapping from POT3D quantity name to its :class:`ModelProps` descriptor."""
814
815
[docs]
816def get_pot3d_quantity_properties(variable: Pot3dQuantities) -> ModelProps:
817 """Return the :class:`~psi_io.models.ModelProps` descriptor for a POT3D quantity.
818
819 Parameters
820 ----------
821 variable : Pot3dQuantities
822 Case-insensitive POT3D quantity name. Valid values: ``'br'``, ``'bt'``,
823 ``'bp'``.
824
825 Returns
826 -------
827 out : ModelProps
828 Immutable descriptor for the requested POT3D magnetic field component.
829
830 Raises
831 ------
832 ValueError
833 If *variable* is not a recognized POT3D quantity.
834
835 See Also
836 --------
837 get_mas_quantity_properties : Analogous lookup for MAS quantities.
838 get_model_prop_caller : Return the getter function for a given model type.
839
840 Examples
841 --------
842 >>> from psi_io.models import get_pot3d_quantity_properties
843 >>> get_pot3d_quantity_properties('bp').name
844 'bp'
845 """
846 try:
847 return _POT3D_QUANTITY_PROPS_MAPPING[variable.lower()]
848 except KeyError:
849 raise ValueError(f"Invalid variable {variable!r} for POT3D model. "
850 f"Valid options are: {', '.join(_POT3D_QUANTITY_PROPS_MAPPING.keys())}") from None
851
852
853# ----------------------------------------------------------------
854# PSI Coordinate Scales
855# ----------------------------------------------------------------
856
857
858PsiScales = Literal['r', 'radius', 't', 'theta', 'p', 'phi']
859"""Literal type alias for the three PSI coordinate scale identifiers.
860
861``'r'`` or ``'radius'``
862 Radial coordinate in unit of solar radii (:data:`~psi_io.units.PSI_rsun`).
863``'t'`` or ``'theta'``
864 Co-latitude :math:`\\theta` in radians (:data:`~psi_io.units.PSI_angle`).
865``'p'`` or ``'phi'``
866 Longitude :math:`\\varphi` in radians (:data:`~psi_io.units.PSI_angle`).
867"""
868
869
870_BASE_SCALE_PROPS_MAPPING = MappingProxyType({
871 'r': ScaleProps('r', 'PSI Radial Scale (Solar Radii)', PSI_rsun,),
872 't': ScaleProps('t', 'PSI Theta Scale (Co-Latitude)', PSI_angle,),
873 'p': ScaleProps('p', 'PSI Phi Scale (Longitude)', PSI_angle,),
874})
875"""Read-only mapping from coordinate scale label to its :class:`ScaleProps` descriptor."""
876
877_PSI_SCALE_PROPS_MAPPING = MappingProxyType({
878 **_BASE_SCALE_PROPS_MAPPING,
879 'radius': _BASE_SCALE_PROPS_MAPPING['r'],
880 'theta': _BASE_SCALE_PROPS_MAPPING['t'],
881 'phi': _BASE_SCALE_PROPS_MAPPING['p'],
882})
883"""Read-only mapping from coordinate scale label to its :class:`ScaleProps` descriptor.
884
885Includes both the canonical short keys (``'r'``, ``'t'``, ``'p'``) and their long-form
886aliases (``'radius'``, ``'theta'``, ``'phi'``). Alias keys share the same
887:class:`ScaleProps` instance as the corresponding canonical key.
888"""
889
890
[docs]
891def get_psi_scale_properties(variable: PsiScales) -> ScaleProps:
892 """Return the :class:`~psi_io.models.ScaleProps` descriptor for a PSI coordinate scale.
893
894 Parameters
895 ----------
896 variable : PsiScales
897 Coordinate label. The first character is used for lookup, so ``'r'``,
898 ``'radius'``, ``'t'``, ``'theta'``, ``'p'``, and ``'phi'`` are all accepted.
899
900 Returns
901 -------
902 out : ScaleProps
903 Descriptor for the requested coordinate axis.
904
905 Raises
906 ------
907 ValueError
908 If *variable* is not a recognized PSI coordinate label.
909
910 See Also
911 --------
912 get_mas_quantity_properties : Analogous lookup for MAS model quantities.
913 get_pot3d_quantity_properties : Analogous lookup for POT3D model quantities.
914
915 Examples
916 --------
917 >>> from psi_io.models import get_psi_scale_properties
918 >>> get_psi_scale_properties('r').desc
919 'PSI Radial Scale (Solar Radii)'
920 >>> get_psi_scale_properties('theta').name
921 't'
922 """
923 try:
924 return _PSI_SCALE_PROPS_MAPPING[variable.lower()]
925 except KeyError:
926 raise ValueError(f"Invalid variable {variable!r} for PSI coordinate scale. "
927 f"Valid options are: {', '.join(_PSI_SCALE_PROPS_MAPPING.keys())}") from None
928
929
930# ----------------------------------------------------------------
931# Model Collection Factory
932# ----------------------------------------------------------------
933
934
935ModelType = Literal['mas', 'pot3d']
936"""Literal type alias for the two recognized PSI model types.
937
938``'mas'``
939 MAS (Magnetohydrodynamic Algorithm outside a Sphere) plasma model output.
940``'pot3d'``
941 POT3D potential-field source-surface (PFSS) magnetic field output.
942"""
943
944
945_PROP_GETTER_MAPPING = MappingProxyType({
946 'mas': get_mas_quantity_properties,
947 'pot3d': get_pot3d_quantity_properties,})
948"""Read-only mapping from model label to its :class:`ModelProps` getter function."""
949
950
[docs]
951def get_model_prop_caller(model: ModelType) -> Callable:
952 """Return the :class:`ModelProps` getter function for the given model type.
953
954 Parameters
955 ----------
956 model : ModelType
957 Case-insensitive model label. Valid values: ``'mas'``, ``'pot3d'``.
958
959 Returns
960 -------
961 out : Callable
962 The getter function associated with *model* — either
963 :func:`get_mas_quantity_properties` or :func:`get_pot3d_quantity_properties`.
964
965 Raises
966 ------
967 ValueError
968 If *model* is not a recognized model label.
969
970 See Also
971 --------
972 get_mas_quantity_properties : Getter for MAS model quantities.
973 get_pot3d_quantity_properties : Getter for POT3D model quantities.
974
975 Examples
976 --------
977 >>> from psi_io.models import get_model_prop_caller
978 >>> caller = get_model_prop_caller('mas')
979 >>> caller('br').name
980 'br'
981 """
982 try:
983 return _PROP_GETTER_MAPPING[model.lower()]
984 except KeyError:
985 raise ValueError(f"Invalid model '{model}'. "
986 f"Valid options are: {', '.join(_PROP_GETTER_MAPPING.keys())}") from None
987
988
989MATCH_QUANTITIES = '|'.join(re.escape(q) for q in sorted(_MAS_QUANTITY_PROPS_MAPPING.keys(), key=len, reverse=True))
990"""Regex alternation string matching any valid MAS quantity name (case-insensitive).
991
992Quantities are sorted longest-first to avoid partial matches (e.g. ``'heat'`` must be
993tried before ``'h'``). Used in :data:`FILEPATH_SCHEMA` and
994:func:`extract_quantity_from_filepath`.
995
996Examples
997--------
998>>> import re
999>>> from psi_io.models import MATCH_QUANTITIES
1000>>> bool(re.match(MATCH_QUANTITIES, 'br', re.IGNORECASE))
1001True
1002>>> bool(re.match(MATCH_QUANTITIES, 'heat', re.IGNORECASE))
1003True
1004>>> bool(re.match(MATCH_QUANTITIES, 'unknown', re.IGNORECASE))
1005False
1006"""
1007
1008FILEPATH_SCHEMA = rf'^({MATCH_QUANTITIES})(\d{{3}}(?:\d{{3}})?)$'
1009"""Regex pattern for the strict MAS filename schema ``<quantity><sequence>``.
1010
1011The stem (filename without extension) must consist of exactly one recognized MAS
1012quantity name followed by a 3- or 6-digit decimal sequence number.
1013
1014Groups:
1015
10161. Quantity name (e.g. ``'br'``, ``'heat'``).
10172. Sequence number (e.g. ``'001'``, ``'001001'``).
1018
1019Used by :func:`parse_psi_filename_schema`.
1020
1021See Also
1022--------
1023extract_quantity_from_filepath :
1024 A lenient variant that does not require the sequence suffix.
1025
1026Examples
1027--------
1028>>> import re
1029>>> from psi_io.models import FILEPATH_SCHEMA
1030>>> bool(re.match(FILEPATH_SCHEMA, 'br001001', re.IGNORECASE))
1031True
1032>>> bool(re.match(FILEPATH_SCHEMA, 'heat001', re.IGNORECASE))
1033True
1034>>> bool(re.match(FILEPATH_SCHEMA, 'br_001', re.IGNORECASE))
1035False
1036"""
1037
1038
1078
1079
1117
1118
[docs]
1119def parse_psi_filename_schema(ifile: Path):
1120 """Parse a PSI HDF filename that follows the strict ``<quantity><sequence>`` schema.
1121
1122 The filename stem must consist of exactly one recognized MAS/POT3D quantity name
1123 followed immediately by a 3- or 6-digit sequence number, with no other characters.
1124 The match is case-insensitive.
1125
1126 Parameters
1127 ----------
1128 ifile : Path
1129 File path to parse. The stem is matched against :data:`FILEPATH_SCHEMA`.
1130
1131 Returns
1132 -------
1133 quantity : str
1134 Lower-case quantity name (e.g. ``'br'``).
1135 sequence : int
1136 Integer sequence number (e.g. ``1001``).
1137
1138 Raises
1139 ------
1140 ValueError
1141 If the filename stem does not match the expected schema.
1142
1143 See Also
1144 --------
1145 extract_quantity_from_filepath : Lenient quantity extraction without schema enforcement.
1146 extract_sequence_from_filepath : Lenient sequence extraction without schema enforcement.
1147 FILEPATH_SCHEMA : The regex pattern used for matching.
1148
1149 Examples
1150 --------
1151 >>> from pathlib import Path
1152 >>> from psi_io.models import parse_psi_filename_schema
1153 >>> parse_psi_filename_schema(Path('br001001.h5'))
1154 ('br', 1001)
1155 >>> parse_psi_filename_schema(Path('heat001.hdf'))
1156 ('heat', 1)
1157 >>> parse_psi_filename_schema(Path('notvalid.h5')) # doctest: +IGNORE_EXCEPTION_DETAIL
1158 Traceback (most recent call last):
1159 ...
1160 ValueError: Filename 'notvalid.h5' does not match expected MAS filename schema: ...
1161 """
1162 matches = re.match(FILEPATH_SCHEMA, ifile.stem, re.IGNORECASE)
1163 if not matches:
1164 raise ValueError(f"Filename '{ifile.name}' does not match expected MAS filename schema: "
1165 f"'<quantity><sequence>'. Valid quantities are: {', '.join(_MAS_QUANTITY_PROPS_MAPPING.keys())}. "
1166 f"Sequence should be a 3 or 6 digit number.")
1167 quantity, sequence = matches.groups()
1168 return quantity, int(sequence)
