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