1"""
2Routines for reading/writing PSI style HDF5 and HDF4 data files.
3
4Written by Ronald M. Caplan, Ryder Davidson, & Cooper Downs.
5
62023/09: Start with SVN version r454, 2023/09/12 by RC, Predictive Science Inc.
7
82024/06: CD: add the get_scales subroutines.
9
102024/11: RD: Major Update: Add several generic data loading capabilites for faster IO.
11 - Read only the portions of data required (`read_hdf_by_index`, `read_hdf_by_value`).
12 - Interpolate to slices along a given axes (`np_interpolate_slice_from_hdf`) or
13 generic positions (`interpolate_positions_from_hdf`).
14
152025/06: CD: Prep for integration into psi-io package, HDF4 is now optional.
16
172026/01: RD: Refactor legacy routines to use new generic routines where possible.
18"""
19
20from __future__ import annotations
21
22__all__ = [
23 "read_hdf_meta",
24 "read_rtp_meta",
25
26 "get_scales_1d",
27 "get_scales_2d",
28 "get_scales_3d",
29
30 "read_hdf_by_index",
31 "read_hdf_by_value",
32 "read_hdf_by_ivalue",
33
34 "np_interpolate_slice_from_hdf",
35 "sp_interpolate_slice_from_hdf",
36 "interpolate_positions_from_hdf",
37
38 "instantiate_linear_interpolator",
39 "interpolate_point_from_1d_slice",
40 "interpolate_point_from_2d_slice",
41
42 "read_hdf_data",
43 "rdhdf_1d",
44 "rdhdf_2d",
45 "rdhdf_3d",
46
47 "write_hdf_data",
48 "wrhdf_1d",
49 "wrhdf_2d",
50 "wrhdf_3d",
51]
52
53import math
54from collections import namedtuple
55from pathlib import Path
56from types import MappingProxyType
57from typing import Optional, Literal, Tuple, Iterable, List, Dict, Union, Callable, Any
58
59import numpy as np
60import h5py as h5
61
62# -----------------------------------------------------------------------------
63# Optional Imports and Import Checking
64# -----------------------------------------------------------------------------
65# These packages are needed by several functions and must be imported in the
66# module namespace.
67try:
68 import pyhdf.SD as h4
69 H4_AVAILABLE = True
70 NPTYPES_TO_SDCTYPES = MappingProxyType({
71 "int8": h4.SDC.INT8,
72 "uint8": h4.SDC.UINT8,
73 "int16": h4.SDC.INT16,
74 "uint16": h4.SDC.UINT16,
75 "int32": h4.SDC.INT32,
76 # "int64": h4.SDC.INT32, # Not supported by pyhdf
77 "uint32": h4.SDC.UINT32,
78 # "float16": h4.SDC.FLOAT32, # Not supported by pyhdf
79 "float32": h4.SDC.FLOAT32,
80 "float64": h4.SDC.FLOAT64,
81 })
82except ImportError:
83 H4_AVAILABLE = False
84 NPTYPES_TO_SDCTYPES = {}
85
86try:
87 from scipy.interpolate import RegularGridInterpolator
88 SCIPY_AVAILABLE = True
89except ImportError:
90 SCIPY_AVAILABLE = False
91
92
93# Functions to stop execution if a package doesn't exist.
94def _except_no_pyhdf():
95 if not H4_AVAILABLE:
96 raise ImportError('The pyhdf package is required to read/write HDF4 .hdf files!')
97 return
98
99
100def _except_no_scipy():
101 if not SCIPY_AVAILABLE:
102 raise ImportError('The scipy package is required for the interpolation routines!')
103 return
104
105
106SDC_TYPE_CONVERSIONS = MappingProxyType({
107 3: np.dtype("ubyte"),
108 4: np.dtype("byte"),
109 5: np.dtype("float32"),
110 6: np.dtype("float64"),
111 20: np.dtype("int8"),
112 21: np.dtype("uint8"),
113 22: np.dtype("int16"),
114 23: np.dtype("uint16"),
115 24: np.dtype("int32"),
116 25: np.dtype("uint32")
117})
118"""Helper dictionary for mapping HDF4 types to numpy dtypes"""
119
120
121PSI_DATA_ID = MappingProxyType({
122 'h4': 'Data-Set-2',
123 'h5': 'Data'
124})
125"""Mapping of PSI standard dataset names for HDF4 and HDF5 files"""
126
127
128PSI_SCALE_ID = MappingProxyType({
129 'h4': ('fakeDim0', 'fakeDim1', 'fakeDim2'),
130 'h5': ('dim1', 'dim2', 'dim3')
131})
132"""Mapping of PSI standard scale names for HDF4 and HDF5 files"""
133
134
135HDFEXT = {'.hdf', '.h5'}
136"""Set of possible HDF file extensions"""
137
138
139HdfExtType = Literal['.hdf', '.h5']
140"""Type alias for possible HDF file extensions"""
141
142
143HdfScaleMeta = namedtuple('HdfScaleMeta', ['name', 'type', 'shape', 'imin', 'imax'])
144"""
145 Named tuples for HDF metadata
146
147 Parameters
148 ----------
149 name : str
150 The name of the scale.
151 type : str
152 The data type of the scale.
153 shape : Tuple[int, ...]
154 The shape of the scale.
155 imin : float
156 The minimum value of the scale.
157 This assumes the scale is monotonically increasing.
158 imax : float
159 The maximum value of the scale.
160 This assumes the scale is monotonically increasing.
161"""
162
163
164HdfDataMeta = namedtuple('HdfDataMeta', ['name', 'type', 'shape', 'scales'])
165"""
166 Named tuple for HDF dataset metadata
167
168 Parameters
169 ----------
170 name : str
171 The name of the dataset.
172 type : str
173 The data type of the dataset.
174 shape : tuple of int
175 The shape of the dataset.
176 scales : list of HdfScaleMeta
177 A list of scale metadata objects corresponding to each dimension of the dataset.
178 If the dataset has no scales, this list will be empty.
179"""
180
181
182def _dispatch_by_ext(ifile: Union[Path, str],
183 hdf4_func: Callable,
184 hdf5_func: Callable,
185 *args: Any, **kwargs: Any
186 ):
187 """
188 Dispatch function to call HDF4 or HDF5 specific functions based on file extension.
189
190 Parameters
191 ----------
192 ifile : Path | str
193 The path to the HDF file.
194 hdf4_func : Callable
195 The function to call for HDF4 files.
196 hdf5_func : Callable
197 The function to call for HDF5 files.
198 *args : Any
199 Positional arguments to pass to the selected function.
200
201 Raises
202 ------
203 ValueError
204 If the file does not have a `.hdf` or `.h5` extension.
205 ImportError
206 If the file is HDF4 and the `pyhdf` package is not available
207 """
208 ipath = Path(ifile)
209 if ipath.suffix == '.h5':
210 return hdf5_func(ifile, *args, **kwargs)
211 if ipath.suffix == '.hdf':
212 _except_no_pyhdf()
213 return hdf4_func(ifile, *args, **kwargs)
214 raise ValueError("File must be HDF4 (.hdf) or HDF5 (.h5)")
215
216
217# -----------------------------------------------------------------------------
218# "Classic" HDF reading and writing routines adapted from psihdf.py or psi_io.py.
219# -----------------------------------------------------------------------------
220
221
[docs]
222def rdhdf_1d(hdf_filename: str
223 ) -> Tuple[np.ndarray, np.ndarray]:
224 """Read a 1D PSI-style HDF5 or HDF4 file.
225
226 Parameters
227 ----------
228 hdf_filename : str
229 The path to the 1D HDF5 (.h5) or HDF4 (.hdf) file to read.
230
231 Returns
232 -------
233 x : np.ndarray
234 1D array of scales.
235 f : np.ndarray
236 1D array of data.
237
238 See Also
239 --------
240 read_hdf_data : Generic HDF data reading routine.
241 """
242 return _rdhdf_nd(hdf_filename, dimensionality=1)
243
244
[docs]
245def rdhdf_2d(hdf_filename: str
246 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
247 """Read a 2D PSI-style HDF5 or HDF4 file.
248
249 The data in the HDF file is assumed to be ordered X,Y in Fortran order.
250
251 Each dimension is assumed to have a 1D "scale" associated with it that
252 describes the rectilinear grid coordinates in each dimension.
253
254 Parameters
255 ----------
256 hdf_filename : str
257 The path to the 1D HDF5 (.h5) or HDF4 (.hdf) file to read.
258
259 Returns
260 -------
261 x : np.ndarray
262 1D array of scales in the X dimension.
263 y : np.ndarray
264 1D array of scales in the Y dimension.
265 f : np.ndarray
266 2D array of data, C-ordered as shape(ny,nx) for Python (see note 1).
267
268 See Also
269 --------
270 read_hdf_data : Generic HDF data reading routine.
271 """
272 return _rdhdf_nd(hdf_filename, dimensionality=2)
273
274
[docs]
275def rdhdf_3d(hdf_filename: str
276 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
277 """Read a 3D PSI-style HDF5 or HDF4 file.
278
279 The data in the HDF file is assumed to be ordered X,Y,Z in Fortran order.
280
281 Each dimension is assumed to have a 1D "scale" associated with it that
282 describes the rectilinear grid coordinates in each dimension.
283
284 Parameters
285 ----------
286 hdf_filename : str
287 The path to the 1D HDF5 (.h5) or HDF4 (.hdf) file to read.
288
289 Returns
290 -------
291 x : np.ndarray
292 1D array of scales in the X dimension.
293 y : np.ndarray
294 1D array of scales in the Y dimension.
295 z : np.ndarray
296 1D array of scales in the Z dimension.
297 f : np.ndarray
298 3D array of data, C-ordered as shape(nz,ny,nx) for Python (see note 1).
299
300 See Also
301 --------
302 read_hdf_data : Generic HDF data reading routine.
303 """
304 return _rdhdf_nd(hdf_filename, dimensionality=3)
305
306
[docs]
307def wrhdf_1d(hdf_filename: str,
308 x: np.ndarray,
309 f: np.ndarray,
310 **kwargs) -> None:
311 """Write a 1D PSI-style HDF5 or HDF4 file.
312
313 Parameters
314 ----------
315 hdf_filename : str
316 The path to the 1D HDF5 (.h5) or HDF4 (.hdf) file to write.
317 x : np.ndarray
318 1D array of scales.
319 f : np.ndarray
320 1D array of data.
321 **kwargs
322 Additional keyword arguments passed through to the underlying
323 :func:`~psi_io.psi_io.write_hdf_data` routine, specifically:
324
325 - ``dataset_id`` : str | None
326 The identifier of the dataset to write.
327 If None, a default dataset is used ('Data-Set-2' for HDF4 and 'Data' for HDF5).
328 - ``sync_dtype``: bool
329 If True, the data type of the scales will be matched to that of the data array.
330
331 Omitting these will yield the same behavior as the legacy routines, *i.e.* writing to
332 the default PSI dataset IDs for HDF4/HDF5 files and synchronizing datatypes between
333 the dataset and scales.
334
335 Raises
336 ------
337 ValueError
338 If the file does not have a `.hdf` or `.h5` extension.
339 KeyError
340 If, for HDF4 files, the provided data is of a type not supported by :py:mod:`pyhdf`.
341 *viz.* float16 or int64.
342
343 Notes
344 -----
345 This routine is provided for backwards compatibility with existing PSI codes that
346 expect this API signature. While the underlying implementation has been refactored to
347 dispatch to the newer :func:`~psi_io.psi_io.write_hdf_data` routine – along with a
348 dimensionality check – the API signature and behavior has been largely preserved.
349
350 .. warning::
351 It is worth reiterating that this routine – when called with its default arguments –
352 will **write to the default PSI dataset IDs for HDF4/HDF5 files and synchronize datatypes
353 between the dataset and scales.** This behavior is crucial for interacting with certain
354 Fortran-based tools developed by PSI.
355
356 See Also
357 --------
358 write_hdf_data : Generic HDF data writing routine.
359 """
360 return _wrhdf_nd(hdf_filename, f, x, dimensionality=1, **kwargs)
361
362
[docs]
363def wrhdf_2d(hdf_filename: str,
364 x: np.ndarray,
365 y: np.ndarray,
366 f: np.ndarray,
367 **kwargs) -> None:
368 """Write a 2D PSI-style HDF5 or HDF4 file.
369
370 The data in the HDF file will appear as X,Y in Fortran order.
371
372 Each dimension requires a 1D "scale" associated with it that
373 describes the rectilinear grid coordinates in each dimension.
374
375 Parameters
376 ----------
377 hdf_filename : str
378 The path to the 2D HDF5 (.h5) or HDF4 (.hdf) file to write.
379 x : np.ndarray
380 1D array of scales in the X dimension.
381 y : np.ndarray
382 1D array of scales in the Y dimension.
383 f : np.ndarray
384 2D array of data, C-ordered as shape(ny,nx) for Python (see note 1).
385 **kwargs
386 Additional keyword arguments passed through to the underlying
387 :func:`~psi_io.psi_io.write_hdf_data` routine, specifically:
388
389 - ``dataset_id`` : str | None
390 The identifier of the dataset to write.
391 If None, a default dataset is used ('Data-Set-2' for HDF4 and 'Data' for HDF5).
392 - ``sync_dtype``: bool
393 If True, the data type of the scales will be matched to that of the data array.
394
395 Omitting these will yield the same behavior as the legacy routines, *i.e.* writing to
396 the default PSI dataset IDs for HDF4/HDF5 files and synchronizing datatypes between
397 the dataset and scales.
398
399 Raises
400 ------
401 ValueError
402 If the file does not have a `.hdf` or `.h5` extension.
403 KeyError
404 If, for HDF4 files, the provided data is of a type not supported by :py:mod:`pyhdf`.
405 *viz.* float16 or int64.
406
407 Notes
408 -----
409 This routine is provided for backwards compatibility with existing PSI codes that
410 expect this API signature. While the underlying implementation has been refactored to
411 dispatch to the newer :func:`~psi_io.psi_io.write_hdf_data` routine – along with a
412 dimensionality check – the API signature and behavior has been largely preserved.
413
414 .. warning::
415 It is worth reiterating that this routine – when called with its default arguments –
416 will **write to the default PSI dataset IDs for HDF4/HDF5 files and synchronize datatypes
417 between the dataset and scales.** This behavior is crucial for interacting with certain
418 Fortran-based tools developed by PSI.
419
420 See Also
421 --------
422 write_hdf_data : Generic HDF data writing routine.
423 """
424 return _wrhdf_nd(hdf_filename, f, x, y, dimensionality=2, **kwargs)
425
426
[docs]
427def wrhdf_3d(hdf_filename: str,
428 x: np.ndarray,
429 y: np.ndarray,
430 z: np.ndarray,
431 f: np.ndarray,
432 **kwargs) -> None:
433 """Write a 3D PSI-style HDF5 or HDF4 file.
434
435 The data in the HDF file will appear as X,Y,Z in Fortran order.
436
437 Each dimension requires a 1D "scale" associated with it that
438 describes the rectilinear grid coordinates in each dimension.
439
440 Parameters
441 ----------
442 hdf_filename : str
443 The path to the 3D HDF5 (.h5) or HDF4 (.hdf) file to write.
444 x : np.ndarray
445 1D array of scales in the X dimension.
446 y : np.ndarray
447 1D array of scales in the Y dimension.
448 z : np.ndarray
449 1D array of scales in the Z dimension.
450 f : np.ndarray
451 3D array of data, C-ordered as shape(nz,ny,nx) for Python (see note 1).
452 **kwargs
453 Additional keyword arguments passed through to the underlying
454 :func:`~psi_io.psi_io.write_hdf_data` routine, specifically:
455
456 - ``dataset_id`` : str | None
457 The identifier of the dataset to write.
458 If None, a default dataset is used ('Data-Set-2' for HDF4 and 'Data' for HDF5).
459 - ``sync_dtype``: bool
460 If True, the data type of the scales will be matched to that of the data array.
461
462 Omitting these will yield the same behavior as the legacy routines, *i.e.* writing to
463 the default PSI dataset IDs for HDF4/HDF5 files and synchronizing datatypes between
464 the dataset and scales.
465
466 Raises
467 ------
468 ValueError
469 If the file does not have a `.hdf` or `.h5` extension.
470 KeyError
471 If, for HDF4 files, the provided data is of a type not supported by :py:mod:`pyhdf`.
472 *viz.* float16 or int64.
473
474 Notes
475 -----
476 This routine is provided for backwards compatibility with existing PSI codes that
477 expect this API signature. While the underlying implementation has been refactored to
478 dispatch to the newer :func:`~psi_io.psi_io.write_hdf_data` routine – along with a
479 dimensionality check – the API signature and behavior has been largely preserved.
480
481 .. warning::
482 It is worth reiterating that this routine – when called with its default arguments –
483 will **write to the default PSI dataset IDs for HDF4/HDF5 files and synchronize datatypes
484 between the dataset and scales.** This behavior is crucial for interacting with certain
485 Fortran-based tools developed by PSI.
486
487 See Also
488 --------
489 write_hdf_data : Generic HDF data writing routine.
490 """
491 return _wrhdf_nd(hdf_filename, f, x, y, z, dimensionality=3, **kwargs)
492
493
[docs]
494def get_scales_1d(filename: str
495 ) -> np.ndarray:
496 """Wrapper to return the scales of a 1D PSI style HDF5 or HDF4 dataset.
497
498 This routine does not load the data array so it can be quite fast for large files.
499
500 Parameters
501 ----------
502 filename : str
503 The path to the 1D HDF5 (.h5) or HDF4 (.hdf) file to read.
504
505 Returns
506 -------
507 x : np.ndarray
508 1D array of scales in the X dimension.
509 """
510 return _dispatch_by_ext(filename, _get_scales_nd_h4, _get_scales_nd_h5,
511 dimensionality=1)
512
513
[docs]
514def get_scales_2d(filename: str
515 ) -> Tuple[np.ndarray, np.ndarray]:
516 """Wrapper to return the scales of a 2D PSI style HDF5 or HDF4 dataset.
517
518 This routine does not load the data array so it can be quite fast for large files.
519
520 The data in the HDF file is assumed to be ordered X,Y in Fortran order.
521
522 Parameters
523 ----------
524 filename : str
525 The path to the 1D HDF5 (.h5) or HDF4 (.hdf) file to read.
526
527 Returns
528 -------
529 x : np.ndarray
530 1D array of scales in the X dimension.
531 y : np.ndarray
532 1D array of scales in the Y dimension.
533 """
534 return _dispatch_by_ext(filename, _get_scales_nd_h4, _get_scales_nd_h5,
535 dimensionality=2)
536
537
[docs]
538def get_scales_3d(filename: str
539 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
540 """Wrapper to return the scales of a 3D PSI style HDF5 or HDF4 dataset.
541
542 This routine does not load the data array so it can be quite fast for large files.
543
544 The data in the HDF file is assumed to be ordered X,Y,Z in Fortran order.
545
546 Parameters
547 ----------
548 filename : str
549 The path to the 1D HDF5 (.h5) or HDF4 (.hdf) file to read.
550
551 Returns
552 -------
553 x : np.ndarray
554 1D array of scales in the X dimension.
555 y : np.ndarray
556 1D array of scales in the Y dimension.
557 z : np.ndarray
558 1D array of scales in the Z dimension.
559 """
560 return _dispatch_by_ext(filename, _get_scales_nd_h4, _get_scales_nd_h5,
561 dimensionality=3)
562
563
564# -----------------------------------------------------------------------------
565# "Updated" HDF reading and slicing routines for Hdf4 and Hdf5 datasets.
566# -----------------------------------------------------------------------------
567
568
616
617
649
650
[docs]
651def read_hdf_data(ifile: Union[Path, str], /,
652 dataset_id: Optional[str] = None,
653 return_scales: bool = True,
654 ) -> Tuple[np.ndarray]:
655 """
656 Read data from an HDF4 (.hdf) or HDF5 (.h5) file.
657
658 Parameters
659 ----------
660 ifile : Path | str
661 The path to the HDF file to read.
662 dataset_id : str | None
663 The identifier of the dataset to read.
664 If None, a default dataset is used ('Data-Set-2' for HDF4 and 'Data' for HDF5).
665 return_scales : bool
666 If True, the scales (coordinate arrays) for each dimension are also returned.
667
668 Returns
669 -------
670 out : np.ndarray | tuple[np.ndarray]
671 The data array.
672 If ``return_scales`` is True, returns a tuple containing the data array
673 and the scales for each dimension.
674
675 Raises
676 ------
677 ValueError
678 If the file does not have a `.hdf` or `.h5` extension.
679
680 See Also
681 --------
682 read_hdf_by_index
683 Read HDF datasets by index.
684 read_hdf_by_value
685 Read HDF datasets by value ranges.
686 read_hdf_by_ivalue
687 Read HDF datasets by subindex values.
688
689 Notes
690 -----
691 This function delegates to :func:`_read_h5_data` for HDF5 files
692 and :func:`_read_h4_data` for HDF4 files based on the file extension.
693 """
694 return _dispatch_by_ext(ifile, _read_h4_data, _read_h5_data,
695 dataset_id=dataset_id, return_scales=return_scales)
696
697
[docs]
698def read_hdf_by_index(ifile: Union[Path, str], /,
699 *xi: Union[int, Tuple[Union[int, None], Union[int, None]], None],
700 dataset_id: Optional[str] = None,
701 return_scales: bool = True,
702 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
703 """
704 Read data from an HDF4 (.hdf) or HDF5 (.h5) file by index.
705
706 .. attention::
707 For each dimension, the *minimum* number of elements returned is 1 *e.g.*
708 if 3 ints are passed (as positional `*xi` arguments) for a 3D dataset,
709 the resulting subset will have a shape of (1, 1, 1,) with scales of length 1.
710
711 Parameters
712 ----------
713 ifile : Path | str
714 The path to the HDF file to read.
715 *xi : int | tuple[int | None, int | None] | None
716 Indices or ranges for each dimension of the `n`-dimensional dataset.
717 Use None for a dimension to select all indices. If no arguments are passed,
718 the entire dataset (and its scales) will be returned – see
719 :func:`~psi_io.psi_io.read_hdf_data`.
720 dataset_id : str | None
721 The identifier of the dataset to read.
722 If None, a default dataset is used ('Data-Set-2' for HDF4 and 'Data' for HDF5).
723 return_scales : bool
724 If True, the scales (coordinate arrays) for each dimension are also returned.
725
726 Returns
727 -------
728 out : np.ndarray | tuple[np.ndarray]
729 The selected data array.
730 If ``return_scales`` is True, returns a tuple containing the data array
731 and the scales for each dimension.
732
733 Raises
734 ------
735 ValueError
736 If the file does not have a `.hdf` or `.h5` extension.
737
738 See Also
739 --------
740 read_hdf_by_value
741 Read HDF datasets by value ranges.
742 read_hdf_by_ivalue
743 Read HDF datasets by subindex values.
744 read_hdf_data
745 Read entire HDF datasets.
746
747 Notes
748 -----
749 This function delegates to :func:`_read_h5_by_index` for HDF5 files and
750 :func:`_read_h4_by_index` for HDF4 files based on the file extension.
751
752 This function assumes that the dataset is Fortran (or column-major) ordered *viz.* for
753 compatibility with PSI's data ecosystem; as such, a given :math:`n`-dimensional array,
754 of shape :math:`(D_0, D_1, ..., D_{n-1})`, has scales :math:`(x_0, x_1, ..., x_{n-1})`,
755 such that :math:`| x_i | = | D_{(n-1)-i} |`. For example, a 3D dataset with shape
756 :math:`(D_p, D_t, D_r)` has scales :math:`r, t, p` corresponding to the radial, theta,
757 and phi dimensions respectively.
758
759 This function extracts a subset of the given dataset/scales without reading the
760 entire data into memory. For a given scale :math:`x_j`, an index, index range, or ``None``
761 can be provided; these inputs are forwarded through to Python's builtin :class:`slice` to
762 extract the desired subset of the scale(s) / dataset.
763
764 Examples
765 --------
766 Import a 3D HDF5 cube.
767
768 >>> from psi_io.data import get_3d_data
769 >>> from psi_io import read_hdf_by_index
770 >>> filepath = get_3d_data()
771
772 Extract a radial slice (at the first radial-scale index) from a 3D cube:
773
774 >>> f, r, t, p = read_hdf_by_value(filepath, 0, None, None)
775 >>> f.shape, r.shape, t.shape, p.shape
776 ((181, 100, 1), (1,), (100,), (181,))
777
778 Extract a phi slice at the 90th index from a 3D cube:
779
780 >>> f, r, t, p = read_hdf_by_value(filepath, None, None, 90)
781 >>> f.shape, r.shape, t.shape, p.shape
782 ((1, 100, 151), (151,), (100,), (1,))
783
784 Extract the values up to the 20th index (in the radial dimension) and with
785 phi indices from 10 to 25:
786
787 >>> f, r, t, p = read_hdf_by_value(filepath, (None, 20), None, (10, 25))
788 >>> f.shape, r.shape, t.shape, p.shape
789 ((15, 100, 20), (20,), (100,), (15,))
790 """
791 if not xi:
792 return read_hdf_data(ifile, dataset_id=dataset_id, return_scales=return_scales)
793 return _dispatch_by_ext(ifile, _read_h4_by_index, _read_h5_by_index,
794 *xi, dataset_id=dataset_id, return_scales=return_scales)
795
796
[docs]
797def read_hdf_by_value(ifile: Union[Path, str], /,
798 *xi: Union[float, Tuple[float, float], None],
799 dataset_id: Optional[str] = None,
800 return_scales: bool = True,
801 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
802 """
803 Read data from an HDF4 (.hdf) or HDF5 (.h5) file by value.
804
805 .. note::
806 For each dimension, the minimum number of elements returned is 2 *e.g.*
807 if 3 floats are passed (as positional `*xi` arguments) for a 3D dataset,
808 the resulting subset will have a shape of (2, 2, 2,) with scales of length 2.
809
810 Parameters
811 ----------
812 ifile : Path | str
813 The path to the HDF file to read.
814 *xi : float | tuple[float, float] | None
815 Values or value ranges corresponding to each dimension of the `n`-dimensional
816 dataset specified by the ``dataset_id``. If no arguments are passed, the
817 entire dataset (and its scales) will be returned.
818 dataset_id : str | None
819 The identifier of the dataset to read.
820 If None, a default dataset is used ('Data-Set-2' for HDF4 and 'Data' for HDF5).
821 return_scales : bool
822 If True, the scales for the specified dataset are also returned.
823
824 Returns
825 -------
826 out : np.ndarray | tuple[np.ndarray]
827 The selected data array.
828 If ``return_scales`` is True, returns a tuple containing the data array
829 and the scales for each dimension.
830
831 Raises
832 ------
833 ValueError
834 If the file does not have a `.hdf` or `.h5` extension.
835
836 See Also
837 --------
838 read_hdf_by_index
839 Read HDF datasets by index.
840 read_hdf_by_ivalue
841 Read HDF datasets by subindex values.
842 read_hdf_data
843 Read entire HDF datasets.
844 sp_interpolate_slice_from_hdf
845 Interpolate slices from HDF datasets using SciPy's
846 :class:`~scipy.interpolate.RegularGridInterpolator`
847 np_interpolate_slice_from_hdf
848 Perform linear, bilinear, or trilinear interpolation using vectorized numpy-based
849 routines.
850
851 Notes
852 -----
853 This function delegates to :func:`_read_h5_by_value` for HDF5 files and
854 :func:`_read_h4_by_value` for HDF4 files based on the file extension.
855
856 This function assumes that the dataset is Fortran (or column-major) ordered *viz.* for
857 compatibility with PSI's data ecosystem; as such, a given :math:`n`-dimensional array,
858 of shape :math:`(D_0, D_1, ..., D_{n-1})`, has scales :math:`(x_0, x_1, ..., x_{n-1})`,
859 such that :math:`| x_i | = | D_{(n-1)-i} |`. For example, a 3D dataset with shape
860 :math:`(D_p, D_t, D_r)` has scales :math:`r, t, p` corresponding to the radial, theta,
861 and phi dimensions respectively.
862
863 This function extracts a subset of the given dataset/scales without reading the
864 entire data into memory. For a given scale :math:`x_j`, if:
865
866 - *i)* a single float is provided (:math:`a`), the function will return a 2-element
867 subset of the scale (:math:`xʹ_j`) such that :math:`xʹ_j[0] <= a < xʹ_j[1]`.
868 - *ii)* a (float, float) tuple is provided (:math:`a_0, a_1`), the function will return an
869 *m*-element subset of the scale (:math:`xʹ_j`) where
870 :math:`xʹ_j[0] <= a_0` and :math:`xʹ_j[m-1] > a_1`.
871 - *iii)* a **None** value is provided, the function will return the entire scale :math:`x_j`
872
873 The returned subset can then be passed to a linear interpolation routine to extract the
874 "slice" at the desired fixed dimensions.
875
876 Examples
877 --------
878 Import a 3D HDF5 cube.
879
880 >>> from psi_io.data import get_3d_data
881 >>> from psi_io import read_hdf_by_value
882 >>> filepath = get_3d_data()
883
884 Extract a radial slice at r=15 from a 3D cube:
885
886 >>> f, r, t, p = read_hdf_by_value(filepath, 15, None, None)
887 >>> f.shape, r.shape, t.shape, p.shape
888 ((181, 100, 2), (2,), (100,), (181,))
889
890 Extract a phi slice at p=1.57 from a 3D cube:
891
892 >>> f, r, t, p = read_hdf_by_value(filepath, None, None, 1.57)
893 >>> f.shape, r.shape, t.shape, p.shape
894 ((2, 100, 151), (151,), (100,), (2,))
895
896 Extract the values between 3.2 and 6.4 (in the radial dimension) and with
897 phi equal to 4.5
898
899 >>> f, r, t, p = read_hdf_by_value(filepath, (3.2, 6.4), None, 4.5)
900 >>> f.shape, r.shape, t.shape, p.shape
901 ((2, 100, 15), (15,), (100,), (2,))
902 """
903 if not xi:
904 return read_hdf_data(ifile, dataset_id=dataset_id, return_scales=return_scales)
905 return _dispatch_by_ext(ifile, _read_h4_by_value, _read_h5_by_value,
906 *xi, dataset_id=dataset_id, return_scales=return_scales)
907
908
[docs]
909def read_hdf_by_ivalue(ifile: Union[Path, str], /,
910 *xi: Union[float, Tuple[float, float], None],
911 dataset_id: Optional[str] = None,
912 return_scales: bool = True,
913 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
914 """
915 Read data from an HDF4 (.hdf) or HDF5 (.h5) file by value.
916
917 .. note::
918 For each dimension, the minimum number of elements returned is 2 *e.g.*
919 if 3 floats are passed (as positional `*xi` arguments) for a 3D dataset,
920 the resulting subset will have a shape of (2, 2, 2,) with scales of length 2.
921
922 Parameters
923 ----------
924 ifile : Path | str
925 The path to the HDF file to read.
926 *xi : float | tuple[float, float] | None
927 Values or value ranges corresponding to each dimension of the `n`-dimensional
928 dataset specified by the ``dataset_id``. If no arguments are passed, the
929 entire dataset (and its scales) will be returned.
930 dataset_id : str | None
931 The identifier of the dataset to read.
932 If None, a default dataset is used ('Data-Set-2' for HDF4 and 'Data' for HDF5).
933 return_scales : bool
934 If True, arrays of indices for the specified dataset are also returned.
935 Note, regardless of whether the provided dataset has coordinate variables
936 (scales), the returned index arrays are always 0-based indices generated
937 through :func:`~numpy.arange`.
938
939 Returns
940 -------
941 out : np.ndarray | tuple[np.ndarray]
942 The selected data array.
943 If ``return_scales`` is True, returns a tuple containing the data array
944 and the scales for each dimension.
945
946 Raises
947 ------
948 ValueError
949 If the file does not have a `.hdf` or `.h5` extension.
950
951 See Also
952 --------
953 read_hdf_by_index
954 Read HDF datasets by index.
955 read_hdf_data
956 Read entire HDF datasets.
957 sp_interpolate_slice_from_hdf
958 Interpolate slices from HDF datasets using SciPy's
959 :class:`~scipy.interpolate.RegularGridInterpolator`
960 np_interpolate_slice_from_hdf
961 Perform linear, bilinear, or trilinear interpolation using vectorized numpy-based
962 routines.
963
964 Notes
965 -----
966 This function delegates to :func:`_read_h5_by_ivalue` for HDF5 files and
967 :func:`_read_h4_by_ivalue` for HDF4 files based on the file extension.
968
969 This function assumes that the dataset is Fortran (or column-major) ordered *viz.* for
970 compatibility with PSI's data ecosystem; as such, a given :math:`n`-dimensional array,
971 of shape :math:`(D_0, D_1, ..., D_{n-1})`, has scales :math:`(x_0, x_1, ..., x_{n-1})`,
972 such that :math:`| x_i | = | D_{(n-1)-i} |`. For example, a 3D dataset with shape
973 :math:`(D_p, D_t, D_r)` has scales :math:`r, t, p` corresponding to the radial, theta,
974 and phi dimensions respectively.
975
976 This function extracts a subset of the given dataset/scales without reading the
977 entire data into memory. For a given scale :math:`x_j`, if:
978
979 - *i)* a single float is provided (:math:`a`), the function will return a 2-element
980 subset of the scale (:math:`xʹ_j`): :math:`xʹ_j[floor(a)], xʹ_j[ceil(a)]`.
981 - *ii)* a (float, float) tuple is provided (:math:`a_0, a_1`), the function will return an
982 *m*-element subset of the scale (:math:`xʹ_j`): :math:`xʹ_j[floor(a_0)], xʹ_j[ceil(a_1)]`
983 - *iii)* a **None** value is provided, the function will return the entire scale :math:`x_j`
984
985 The returned subset can then be passed to a linear interpolation routine to extract the
986 "slice" at the desired fixed dimensions.
987 """
988 if not xi:
989 return read_hdf_data(ifile, dataset_id=dataset_id, return_scales=return_scales)
990 return _dispatch_by_ext(ifile, _read_h4_by_ivalue, _read_h5_by_ivalue,
991 *xi, dataset_id=dataset_id, return_scales=return_scales)
992
993
[docs]
994def write_hdf_data(ifile: Union[Path, str], /,
995 data: np.ndarray,
996 *scales: Iterable[Union[np.ndarray, None]],
997 dataset_id: Optional[str] = None,
998 sync_dtype: bool = False,
999 ) -> Path:
1000 """
1001 Write data to an HDF4 (.hdf) or HDF5 (.h5) file.
1002
1003 Following PSI conventions, the data array is assumed to be Fortran-ordered,
1004 with the scales provided in the order corresponding to each dimension *e.g.* a
1005 3D dataset with shape (Dp, Dt, Dr) has scales r, t, p corresponding to the
1006 radial, theta, and phi dimensions respectively.
1007
1008 Parameters
1009 ----------
1010 ifile : Path | str
1011 The path to the HDF file to write.
1012 data : np.ndarray
1013 The data array to write.
1014 *scales : Iterable[np.ndarray | None]
1015 The scales (coordinate arrays) for each dimension.
1016 dataset_id : str | None
1017 The identifier of the dataset to write.
1018 If None, a default dataset is used ('Data-Set-2' for HDF4
1019 and 'Data' for HDF5).
1020 sync_dtype : bool
1021 If True, the data type of the scales will be matched to that of the data array.
1022 This argument is included to mimic the behavior of PSI's legacy HDF writing routines
1023 and, more importantly, to ensure compatibility with certain Fortran tools within
1024 the PSI ecosystem that require uniform precision between datasets and their scales.
1025
1026 Returns
1027 -------
1028 out : Path
1029 The path to the written HDF file.
1030
1031 Raises
1032 ------
1033 ValueError
1034 If the file does not have a `.hdf` or `.h5` extension.
1035 KeyError
1036 If, for HDF4 files, the provided data is of a type not supported by :py:mod:`pyhdf`.
1037 *viz.* float16 or int64.
1038
1039 Notes
1040 -----
1041 This function delegates to :func:`_write_h5_data` for HDF5 files
1042 and :func:`_write_h4_data` for HDF4 files based on the file extension.
1043
1044 If no scales are provided, the dataset will be written without coordinate variables.
1045 If scales are provided, the number of scales must be less than or equal to the number
1046 of dimensions in the data array. To attach scales to specific dimensions only, provide
1047 ``None`` for the dimensions without scales.
1048
1049 See Also
1050 --------
1051 wrhdf_1d
1052 Write 1D HDF files.
1053 wrhdf_2d
1054 Write 2D HDF files.
1055 wrhdf_3d
1056 Write 3D HDF files.
1057 """
1058 return _dispatch_by_ext(ifile, _write_h4_data, _write_h5_data, data,
1059 *scales, dataset_id=dataset_id, sync_dtype=sync_dtype)
1060
1061
[docs]
1062def instantiate_linear_interpolator(*args, **kwargs):
1063 """
1064 Instantiate a linear interpolator using the provided data and scales.
1065
1066 Parameters
1067 ----------
1068 *args : sequence[array_like]
1069 The first argument is the data array.
1070 Subsequent arguments are the scales (coordinate arrays) for each dimension.
1071 **kwargs : dict
1072 Additional keyword arguments to pass to
1073 :class:`~scipy.interpolate.RegularGridInterpolator`.
1074
1075 Returns
1076 -------
1077 out : RegularGridInterpolator
1078 An instance of RegularGridInterpolator initialized
1079 with the provided data and scales.
1080
1081 Notes
1082 -----
1083 This function transposes the data array and passes it along with the scales
1084 to RegularGridInterpolator. Given a PSI-style Fortran-ordered 3D dataset,
1085 the resulting interpolator can be queried using (r, theta, phi) coordinates.
1086
1087 Examples
1088 --------
1089 Import a 3D HDF5 cube.
1090
1091 >>> from psi_io.data import get_3d_data
1092 >>> from psi_io import read_hdf_by_value
1093 >>> from numpy import pi
1094 >>> filepath = get_3d_data()
1095
1096 Read the dataset by value (at 15 R_sun in the radial dimension).
1097
1098 >>> data_and_scales = read_hdf_by_value(filepath, 15, None, None)
1099 >>> interpolator = instantiate_linear_interpolator(*data_and_scales)
1100
1101 Interpolate at a specific position.
1102
1103 >>> interpolator((15, pi/2, pi))
1104 0.0012864485109423877
1105 """
1106 _except_no_scipy()
1107 return RegularGridInterpolator(
1108 values=args[0].T,
1109 points=args[1:],
1110 **kwargs)
1111
1112
[docs]
1113def sp_interpolate_slice_from_hdf(*xi, **kwargs):
1114 """
1115 Interpolate a slice from HDF data using SciPy's `RegularGridInterpolator`.
1116
1117 .. note::
1118 Slicing routines result in a dimensional reduction. The dimensions
1119 that are fixed (i.e. provided as float values in `*xi`) are removed
1120 from the output slice, while the dimensions that are not fixed
1121 (*i.e.* provided as None in `*xi`) are retained.
1122
1123 Parameters
1124 ----------
1125 *xi : sequence
1126 Positional arguments passed-through to :func:`read_hdf_by_value`.
1127 **kwargs : dict
1128 Keyword arguments passed-through to :func:`read_hdf_by_value`.
1129 **NOTE: Instantiating a linear interpolator requires the** ``return_scales``
1130 **keyword argument to be set to True; this function overrides
1131 any provided value for** ``return_scales`` **to ensure this behavior.**
1132
1133 Returns
1134 -------
1135 slice : np.ndarray
1136 The interpolated data slice.
1137 scales : list
1138 A list of scales for the dimensions that were not fixed.
1139
1140 Notes
1141 -----
1142 This function reads data from an HDF file, creates a linear interpolator,
1143 and interpolates a slice based on the provided values.
1144
1145 .. note::
1146 The returned slice is Fortran-ordered *e.g.* radial slices will have shape
1147 (phi, theta), phi slices will have shape (r, theta), etc.
1148
1149 .. note::
1150 SciPy's `RegularGridInterpolator` casts all input data to `float64` internally.
1151 Therefore, PSI HDF datasets with single-precision (`float32`) data will be upcast
1152 during interpolation.
1153
1154 Examples
1155 --------
1156 >>> from psi_io.data import get_3d_data
1157 >>> from psi_io import sp_interpolate_slice_from_hdf
1158 >>> from numpy import pi
1159 >>> filepath = get_3d_data()
1160
1161 Fetch a 2D slice at r=15 from 3D map
1162
1163 >>> slice_, theta_scale, phi_scale = sp_interpolate_slice_from_hdf(filepath, 15, None, None)
1164 >>> slice_.shape, theta_scale.shape, phi_scale.shape
1165 ((181, 100), (100,), (181,))
1166
1167 Fetch a single point from 3D map
1168
1169 >>> point_value, *_ = sp_interpolate_slice_from_hdf(filepath, 1, pi/2, pi)
1170 >>> point_value
1171 6.084495480971823
1172 """
1173 filepath, *args = xi
1174 kwargs.pop('return_scales', None)
1175 result = read_hdf_by_value(filepath, *args, **kwargs)
1176 interpolator = instantiate_linear_interpolator(*result)
1177 grid = [yi[0] if yi[0] is not None else yi[1] for yi in zip(args, result[1:])]
1178 slice_ = interpolator(tuple(np.meshgrid(*grid, indexing='ij')))
1179 indices = [0 if yi is not None else slice(None, None) for yi in args]
1180 return slice_[tuple(indices)].T, *[yi[1] for yi in zip(args, result[1:]) if yi[0] is None]
1181
1182
[docs]
1183def np_interpolate_slice_from_hdf(ifile: Union[Path, str], /,
1184 *xi: Union[float, Tuple[float, float], None],
1185 dataset_id: Optional[str] = None,
1186 by_index: bool = False,
1187 ):
1188 """
1189 Interpolate a slice from HDF data using linear interpolation.
1190
1191 .. note::
1192 Slicing routines result in a dimensional reduction. The dimensions
1193 that are fixed (i.e. provided as float values in `*xi`) are removed
1194 from the output slice, while the dimensions that are not fixed
1195 (*i.e.* provided as `None` in `*xi`) are retained.
1196
1197 Parameters
1198 ----------
1199 ifile : Path | str
1200 The path to the HDF file to read.
1201 *xi : sequence
1202 Positional arguments passed-through to reader function.
1203 dataset_id : str | None
1204 The identifier of the dataset to read.
1205 If None, a default dataset is used ('Data-Set-2' for HDF
1206 and 'Data' for HDF5).
1207 by_index : bool
1208 If True, use :func:`read_hdf_by_ivalue` to read data by subindex values.
1209 If False, use :func:`read_hdf_by_value` to read data by value ranges.
1210
1211 Returns
1212 -------
1213 slice : np.ndarray
1214 The interpolated data slice.
1215 scales : list
1216 A list of scales for the dimensions that were not fixed.
1217
1218 Raises
1219 ------
1220 ValueError
1221 If the number of dimensions to interpolate over is not supported.
1222
1223 Notes
1224 -----
1225 This function supports linear, bilinear, and trilinear interpolation
1226 depending on the number of dimensions fixed in `xi`.
1227
1228 Examples
1229 --------
1230 >>> from psi_io.data import get_3d_data
1231 >>> from psi_io import np_interpolate_slice_from_hdf
1232 >>> from numpy import pi
1233 >>> filepath = get_3d_data()
1234
1235 Fetch a 2D slice at r=15 from 3D map
1236
1237 >>> slice_, theta_scale, phi_scale = np_interpolate_slice_from_hdf(filepath, 15, None, None)
1238 >>> slice_.shape, theta_scale.shape, phi_scale.shape
1239 ((181, 100), (100,), (181,))
1240
1241 Fetch a single point from 3D map
1242
1243 >>> point_value, *_ = np_interpolate_slice_from_hdf(filepath, 1, pi/2, pi)
1244 >>> point_value
1245 6.084496
1246
1247 """
1248 reader = read_hdf_by_value if not by_index else read_hdf_by_ivalue
1249 data, *scales = reader(ifile, *xi, dataset_id=dataset_id, return_scales=True)
1250 f_ = np.transpose(data)
1251 slice_type = sum([yi is not None for yi in xi])
1252 if slice_type == 1:
1253 return _np_linear_interpolation(xi, scales, f_).T, *[yi[1] for yi in zip(xi, scales) if yi[0] is None]
1254 elif slice_type == 2:
1255 return _np_bilinear_interpolation(xi, scales, f_).T, *[yi[1] for yi in zip(xi, scales) if yi[0] is None]
1256 elif slice_type == 3:
1257 return _np_trilinear_interpolation(xi, scales, f_).T, *[yi[1] for yi in zip(xi, scales) if yi[0] is None]
1258 else:
1259 raise ValueError("Not a valid number of dimensions for supported linear interpolation methods")
1260
1261
[docs]
1262def interpolate_positions_from_hdf(ifile, *xi, **kwargs):
1263 """
1264 Interpolate at a list of scale positions using SciPy's `RegularGridInterpolator`.
1265
1266 Parameters
1267 ----------
1268 ifile : Path | str
1269 The path to the HDF file to read.
1270 *xi : sequence[np.ndarray]
1271 Iterable scale values for each dimension of the `n`-dimensional dataset.
1272 Each array should have the same shape *i.e.* :math:`(m, )` – the function composes
1273 these array into a :math:`m x n` column stack for interpolation.
1274 **kwargs : dict
1275 Keyword arguments to pass to :func:`read_hdf_by_value`.
1276
1277 Returns
1278 -------
1279 out : np.ndarray
1280 The interpolated values at the provided positions.
1281
1282 Notes
1283 -----
1284 This function reads data from an HDF file, creates a linear interpolator,
1285 and interpolates at the provided scale values. For each dimension, the
1286 minimum and maximum values from the provided arrays are used to read
1287 the necessary subset of data from the HDF file *viz.* to avoid loading
1288 the entire dataset into memory.
1289
1290 Examples
1291 --------
1292 Import a 3D HDF5 cube.
1293
1294 >>> from psi_io.data import get_3d_data
1295 >>> from psi_io import interpolate_positions_from_hdf
1296 >>> import numpy as np
1297 >>> filepath = get_3d_data()
1298
1299 Set up positions to interpolate.
1300
1301 >>> r_vals = np.array([15, 20, 25])
1302 >>> theta_vals = np.array([np.pi/4, np.pi/2, 3*np.pi/4])
1303 >>> phi_vals = np.array([0, np.pi, 2*np.pi])
1304
1305 Interpolate at the specified positions.
1306
1307 >>> interpolate_positions_from_hdf(filepath, r_vals, theta_vals, phi_vals)
1308 [0.0008402743657585175, 0.000723875405654482, -0.00041033233811179216]
1309 """
1310 xi_ = [(np.nanmin(i), np.nanmax(i)) for i in xi]
1311 f, *scales = read_hdf_by_value(ifile, *xi_, **kwargs)
1312 interpolator = instantiate_linear_interpolator(f, *scales, bounds_error=False)
1313 return interpolator(np.stack(xi, axis=len(xi[0].shape)))
1314
1315
[docs]
1316def interpolate_point_from_1d_slice(xi, scalex, values):
1317 """
1318 Interpolate a point from a 1D slice using linear interpolation.
1319
1320 Parameters
1321 ----------
1322 xi : float
1323 The scale value at which to interpolate.
1324 scalex : np.ndarray
1325 The scale (coordinate array) for the dimension.
1326 values : np.ndarray
1327 The data array to interpolate.
1328
1329 Returns
1330 -------
1331 out : np.ndarray
1332 The interpolated data point.
1333 """
1334 if scalex[0] > scalex[-1]:
1335 scalex, values = scalex[::-1], values[::-1]
1336 xi_ = int(np.searchsorted(scalex, xi))
1337 sx_ = slice(*_check_index_ranges(len(scalex), xi_, xi_))
1338 return _np_linear_interpolation([xi], [scalex[sx_]], values[sx_])
1339
1340
[docs]
1341def interpolate_point_from_2d_slice(xi, yi, scalex, scaley, values):
1342 """
1343 Interpolate a point from a 2D slice using bilinear interpolation.
1344
1345 Parameters
1346 ----------
1347 xi : float
1348 The scale value for the first dimension.
1349 yi : float
1350 The scale value for the second dimension.
1351 scalex : np.ndarray
1352 The scale (coordinate array) for the first dimension.
1353 scaley : np.ndarray
1354 The scale (coordinate array) for the second dimension.
1355 values : np.ndarray
1356 The data array to interpolate.
1357
1358 Returns
1359 -------
1360 out : np.ndarray
1361 The interpolated data point.
1362 """
1363 values = np.transpose(values)
1364 if scalex[0] > scalex[-1]:
1365 scalex, values = scalex[::-1], values[::-1, :]
1366 if scaley[0] > scaley[-1]:
1367 scaley, values = scaley[::-1], values[:, ::-1]
1368 xi_, yi_ = int(np.searchsorted(scalex, xi)), int(np.searchsorted(scaley, yi))
1369 sx_, sy_ = slice(*_check_index_ranges(len(scalex), xi_, xi_)), slice(*_check_index_ranges(len(scaley), yi_, yi_))
1370 return _np_bilinear_interpolation([xi, yi], [scalex[sx_], scaley[sy_]], values[(sx_, sy_)])
1371
1372
1373def _rdhdf_nd(hdf_filename: str,
1374 dimensionality: int
1375 ) -> Tuple[np.ndarray, ...]:
1376 f, *scales = read_hdf_data(hdf_filename)
1377 if f.ndim != dimensionality:
1378 err = f'Expected {dimensionality}D data, got {f.ndim}D data instead.'
1379 raise ValueError(err)
1380 scales = scales or (np.empty(0) for _ in f.shape)
1381 return *scales, f
1382
1383
1384def _wrhdf_nd(hdf_filename: str,
1385 data: np.ndarray,
1386 *scales: Iterable[Union[np.ndarray, None]],
1387 dimensionality: int,
1388 sync_dtype: bool = True,
1389 **kwargs
1390 ) -> None:
1391 if data.ndim != dimensionality:
1392 err = f'Expected {dimensionality}D data, got {data.ndim}D data instead.'
1393 raise ValueError(err)
1394 write_hdf_data(hdf_filename, data, *scales, sync_dtype=sync_dtype, **kwargs)
1395
1396
1397def _get_scales_nd_h5(ifile: Union[ Path, str], /,
1398 dimensionality: int,
1399 dataset_id: Optional[str] = None,
1400 ):
1401 with h5.File(ifile, 'r') as hdf:
1402 data = hdf[dataset_id or PSI_DATA_ID['h5']]
1403 ndim = data.ndim
1404 if ndim != dimensionality:
1405 err = f'Expected {dimensionality}D data, got {ndim}D data instead.'
1406 raise ValueError(err)
1407 scales = []
1408 for dim in data.dims:
1409 if dim:
1410 scales.append(dim[0][:])
1411 else:
1412 raise ValueError(f'Dimension has no scale associated with it.')
1413 return tuple(scales)
1414
1415
1416def _get_scales_nd_h4(ifile: Union[ Path, str], /,
1417 dimensionality: int,
1418 dataset_id: Optional[str] = None,
1419 ):
1420 hdf = h4.SD(str(ifile))
1421 data = hdf.select(dataset_id or PSI_DATA_ID['h4'])
1422 ndim = data.info()[1]
1423 if ndim != dimensionality:
1424 err = f'Expected {dimensionality}D data, got {ndim}D data instead.'
1425 raise ValueError(err)
1426 scales = []
1427 for k_, v_ in reversed(data.dimensions(full=1).items()):
1428 if v_[3]:
1429 scales.append(hdf.select(k_)[:])
1430 else:
1431 raise ValueError('Dimension has no scale associated with it.')
1432 return tuple(scales)
1433
1434
1435def _read_h5_meta(ifile: Union[Path, str], /,
1436 dataset_id: Optional[str] = None
1437 ):
1438 """HDF5 (.h5) version of :func:`read_hdf_meta`."""
1439 with h5.File(ifile, 'r') as hdf:
1440 # Raises KeyError if ``dataset_id`` not found
1441 # If ``dataset_id`` is None, get all non-scale :class:`h5.Dataset`s
1442 if dataset_id:
1443 datasets = (dataset_id, hdf[dataset_id]),
1444 else:
1445 datasets = ((k, v) for k, v in hdf.items() if not v.is_scale)
1446
1447 # One should avoid multiple calls to ``dimproxy[0]`` – *e.g.* ``dimproxy[0].dtype`` and
1448 # ``dimproxy[0].shape`` – because the __getitem__ method creates and returns a new
1449 # :class:`~h5.DimensionProxy` object each time it is called. [Does this matter? Probably not.]
1450 return [HdfDataMeta(name=k,
1451 type=v.dtype,
1452 shape=v.shape,
1453 scales=[HdfScaleMeta(name=dimproxy.label,
1454 type=dim.dtype,
1455 shape=dim.shape,
1456 imin=dim[0],
1457 imax=dim[-1])
1458 for dimproxy in v.dims if dimproxy and (dim := dimproxy[0])])
1459 for k, v in datasets]
1460
1461
1462def _read_h4_meta(ifile: Union[Path, str], /,
1463 dataset_id: Optional[str] = None
1464 ):
1465 """HDF4 (.hdf) version of :func:`read_hdf_meta`."""
1466 hdf = h4.SD(str(ifile))
1467 # Raises HDF4Error if ``dataset_id`` not found
1468 # If ``dataset_id`` is None, get all non-scale :class:`pyhdf.SD.SDS`s
1469 if dataset_id:
1470 datasets = (dataset_id, hdf.select(dataset_id)),
1471 else:
1472 datasets = ((k, hdf.select(k)) for k in hdf.datasets().keys() if not hdf.select(k).iscoordvar())
1473
1474 # The inner list comprehension differs in approach from the HDF5 version because calling
1475 # ``dimensions(full=1)`` on an :class:`~pyhdf.SD.SDS` returns a dictionary of dimension
1476 # dataset identifiers (keys) and tuples containing dimension metadata (values). Even if no
1477 # coordinate-variable datasets are defined, this dictionary is still returned; the only
1478 # indication that the datasets returned do not exist is that the "type" field (within the
1479 # tuple of dimension metadata) is set to 0.
1480
1481 # Also, one cannot avoid multiple calls to ``hdf.select(k_)`` within the inner list comprehension
1482 # because :class:`~pyhdf.SD.SDS` objects do not define a ``__bool__`` method, and the fallback
1483 # behavior of Python is to assess if the __len__ method returns a non-zero value (which, in
1484 # this case, always returns 0).
1485 return [HdfDataMeta(name=k,
1486 type=SDC_TYPE_CONVERSIONS[v.info()[3]],
1487 shape=_cast_shape_tuple(v.info()[2]),
1488 scales=[HdfScaleMeta(name=k_,
1489 type=SDC_TYPE_CONVERSIONS[v_[3]],
1490 shape=_cast_shape_tuple(v_[0]),
1491 imin=hdf.select(k_)[0],
1492 imax=hdf.select(k_)[-1])
1493 for k_, v_ in v.dimensions(full=1).items() if v_[3]])
1494 for k, v in datasets]
1495
1496
1497def _read_h5_rtp(ifile: Union[ Path, str], /):
1498 """HDF5 (.h5) version of :func:`read_rtp_meta`."""
1499 with h5.File(ifile, 'r') as hdf:
1500 return {k: (hdf[v].size, hdf[v][0], hdf[v][-1])
1501 for k, v in zip('rtp', PSI_SCALE_ID['h5'])}
1502
1503
1504def _read_h4_rtp(ifile: Union[ Path, str], /):
1505 """HDF4 (.hdf) version of :func:`read_rtp_meta`."""
1506 hdf = h4.SD(str(ifile))
1507 return {k: (hdf.select(v).info()[2], hdf.select(v)[0], hdf.select(v)[-1])
1508 for k, v in zip('ptr', PSI_SCALE_ID['h4'])}
1509
1510
1511def _read_h5_data(ifile: Union[Path, str], /,
1512 dataset_id: Optional[str] = None,
1513 return_scales: bool = True,
1514 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1515 with h5.File(ifile, 'r') as hdf:
1516 data = hdf[dataset_id or PSI_DATA_ID['h5']]
1517 dataset = data[:]
1518 if return_scales:
1519 scales = [dim[0][:] for dim in data.dims if dim]
1520 return dataset, *scales
1521 return dataset
1522
1523
1524def _read_h4_data(ifile: Union[Path, str], /,
1525 dataset_id: Optional[str] = None,
1526 return_scales: bool = True,
1527 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1528 hdf = h4.SD(str(ifile))
1529 data = hdf.select(dataset_id or PSI_DATA_ID['h4'])
1530 if return_scales:
1531 out = (data[:],
1532 *[hdf.select(k_)[:] for k_, v_ in reversed(data.dimensions(full=1).items()) if v_[3]])
1533 else:
1534 out = data[:]
1535 return out
1536
1537
1538def _read_h5_by_index(ifile: Union[Path, str], /,
1539 *xi: Union[int, Tuple[Union[int, None], Union[int, None]], None],
1540 dataset_id: Optional[str] = None,
1541 return_scales: bool = True,
1542 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1543 """ HDF5(.h5) version of :func:`read_hdf_by_index`."""
1544 with h5.File(ifile, 'r') as hdf:
1545 data = hdf[dataset_id or PSI_DATA_ID['h5']]
1546 if len(xi) != data.ndim:
1547 raise ValueError(f"len(xi) must equal the number of scales for {dataset_id}")
1548 slices = [_parse_index_inputs(slice_input) for slice_input in xi]
1549 dataset = data[tuple(reversed(slices))]
1550 if return_scales:
1551 scales = [dim[0][si] for si, dim in zip(slices, data.dims) if dim]
1552 return dataset, *scales
1553 return dataset
1554
1555def _read_h4_by_index(ifile: Union[Path, str], /,
1556 *xi: Union[int, Tuple[Union[int, None], Union[int, None]], None],
1557 dataset_id: Optional[str] = None,
1558 return_scales: bool = True,
1559 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1560 """HDF4(.hdf) version of :func:`read_hdf_by_index`."""
1561 hdf = h4.SD(str(ifile))
1562 data = hdf.select(dataset_id or PSI_DATA_ID['h4'])
1563 ndim = data.info()[1]
1564 if len(xi) != ndim:
1565 raise ValueError(f"len(xi) must equal the number of scales for {dataset_id}")
1566 slices = [_parse_index_inputs(slice_input) for slice_input in xi]
1567 dataset = data[tuple(reversed(slices))]
1568 if return_scales:
1569 scales = [hdf.select(k_)[si] for si, (k_, v_) in zip(slices, reversed(data.dimensions(full=1).items())) if v_[3]]
1570 return dataset, *scales
1571 return dataset
1572
1573
1574def _read_h5_by_value(ifile: Union[Path, str], /,
1575 *xi: Union[float, Tuple[float, float], None],
1576 dataset_id: Optional[str] = None,
1577 return_scales: bool = True,
1578 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1579 """HDF5 (.h5) version of :func:`read_hdf_by_value`."""
1580 with h5.File(ifile, 'r') as hdf:
1581 data = hdf[dataset_id or PSI_DATA_ID['h5']]
1582 if len(xi) != data.ndim:
1583 raise ValueError(f"len(xi) must equal the number of scales for {dataset_id}")
1584 slices = []
1585 for dimproxy, value in zip(data.dims, xi):
1586 if dimproxy:
1587 slices.append(_parse_value_inputs(dimproxy[0], value))
1588 elif value is None:
1589 slices.append(slice(None))
1590 else:
1591 raise ValueError("Cannot slice by value on dimension without scales")
1592 dataset = data[tuple(reversed(slices))]
1593 if return_scales:
1594 scales = [dim[0][si] for si, dim in zip(slices, data.dims) if dim]
1595 return dataset, *scales
1596 return dataset
1597
1598
1599def _read_h4_by_value(ifile: Union[Path, str], /,
1600 *xi: Union[float, Tuple[float, float], None],
1601 dataset_id: Optional[str] = None,
1602 return_scales: bool = True,
1603 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1604 """HDF4 (.hdf) version of :func:`read_hdf_by_value`."""
1605 hdf = h4.SD(str(ifile))
1606 data = hdf.select(dataset_id or PSI_DATA_ID['h4'])
1607 ndim = data.info()[1]
1608 if len(xi) != ndim:
1609 raise ValueError(f"len(xi) must equal the number of scales for {dataset_id}")
1610 slices = []
1611 for (k_, v_), value in zip(reversed(data.dimensions(full=1).items()), xi):
1612 if v_[3] != 0:
1613 slices.append(_parse_value_inputs(hdf.select(k_), value))
1614 elif value is None:
1615 slices.append(slice(None))
1616 else:
1617 raise ValueError("Cannot slice by value on dimension without scales")
1618 dataset = data[tuple(reversed(slices))]
1619 if return_scales:
1620 scales = [hdf.select(k_)[si] for si, (k_, v_) in zip(slices, reversed(data.dimensions(full=1).items())) if v_[3]]
1621 return dataset, *scales
1622 return dataset
1623
1624
1625def _read_h5_by_ivalue(ifile: Union[Path, str], /,
1626 *xi: Union[float, Tuple[float, float], None],
1627 dataset_id: Optional[str] = None,
1628 return_scales: bool = True,
1629 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1630 with h5.File(ifile, 'r') as hdf:
1631 data = hdf[dataset_id or PSI_DATA_ID['h5']]
1632 if len(xi) != data.ndim:
1633 raise ValueError(f"len(xi) must equal the number of scales for {dataset_id}")
1634 slices = [_parse_ivalue_inputs(*args) for args in zip(reversed(data.shape), xi)]
1635 dataset = data[tuple(reversed(slices))]
1636 if return_scales:
1637 scales = [np.arange(si.start or 0, si.stop or size) for si, size in zip(slices, reversed(data.shape))]
1638 return dataset, *scales
1639 return dataset
1640
1641
1642def _read_h4_by_ivalue(ifile: Union[Path, str], /,
1643 *xi: Union[float, Tuple[float, float], None],
1644 dataset_id: Optional[str] = None,
1645 return_scales: bool = True,
1646 ) -> Union[np.ndarray, Tuple[np.ndarray]]:
1647 hdf = h4.SD(str(ifile))
1648 data = hdf.select(dataset_id or PSI_DATA_ID['h4'])
1649 ndim, shape = data.info()[1], _cast_shape_tuple(data.info()[2])
1650 if len(xi) != ndim:
1651 raise ValueError(f"len(xi) must equal the number of scales for {dataset_id}")
1652 slices = [_parse_ivalue_inputs(*args) for args in zip(reversed(shape), xi)]
1653 dataset = data[tuple(reversed(slices))]
1654 if return_scales:
1655 scales = [np.arange(si.start or 0, si.stop or size) for si, size in zip(slices, reversed(shape))]
1656 return dataset, *scales
1657 return dataset
1658
1659
1660def _write_h4_data(ifile: Union[Path, str], /,
1661 data: np.ndarray,
1662 *scales: Iterable[np.ndarray],
1663 dataset_id: Optional[str] = None,
1664 sync_dtype: bool = False,
1665 ) -> Path:
1666 dataid = dataset_id or PSI_DATA_ID['h4']
1667 h4file = h4.SD(str(ifile), h4.SDC.WRITE | h4.SDC.CREATE | h4.SDC.TRUNC)
1668 sds_id = h4file.create(dataid, NPTYPES_TO_SDCTYPES[data.dtype.name], data.shape)
1669
1670 if scales:
1671 for i, scale in enumerate(reversed(scales)):
1672 if scale is not None:
1673 if sync_dtype:
1674 scale = scale.astype(data.dtype)
1675 sds_id.dim(i).setscale(NPTYPES_TO_SDCTYPES[scale.dtype.name], scale.tolist())
1676
1677 sds_id.set(data)
1678 sds_id.endaccess()
1679 h4file.end()
1680
1681 return ifile
1682
1683
1684def _write_h5_data(ifile: Union[Path, str], /,
1685 data: np.ndarray,
1686 *scales: Iterable[np.ndarray],
1687 dataset_id: Optional[str] = None,
1688 sync_dtype: bool = False,
1689 ) -> Path:
1690 dataid = dataset_id or PSI_DATA_ID['h5']
1691 with h5.File(ifile, "w") as h5file:
1692 h5file.create_dataset(dataid, data=data, dtype=data.dtype, shape=data.shape)
1693
1694 if scales:
1695 for i, scale in enumerate(scales):
1696 if scale is not None:
1697 if sync_dtype:
1698 scale = scale.astype(data.dtype)
1699 h5file.create_dataset(f"dim{i+1}", data=scale, dtype=scale.dtype, shape=scale.shape)
1700 h5file[dataid].dims[i].attach_scale(h5file[f"dim{i+1}"])
1701 h5file[dataid].dims[i].label = f"dim{i+1}"
1702
1703 return ifile
1704
1705
1706def _np_linear_interpolation(xi: Iterable, scales: Iterable, values: np.ndarray):
1707 """
1708 Perform linear interpolation over one dimension.
1709
1710 Parameters
1711 ----------
1712 xi : list
1713 List of values or None for each dimension.
1714 scales : list
1715 List of scales (coordinate arrays) for each dimension.
1716 values : np.ndarray
1717 The data array to interpolate.
1718
1719 Returns
1720 -------
1721 np.ndarray
1722 The interpolated data.
1723
1724 """
1725 index0 = next((i for i, v in enumerate(xi) if v is not None), None)
1726 t = (xi[index0] - scales[index0][0])/(scales[index0][1] - scales[index0][0])
1727 f0 = [slice(None, None)]*values.ndim
1728 f1 = [slice(None, None)]*values.ndim
1729 f0[index0] = 0
1730 f1[index0] = 1
1731
1732 return (1 - t)*values[tuple(f0)] + t*values[tuple(f1)]
1733
1734
1735def _np_bilinear_interpolation(xi, scales, values):
1736 """
1737 Perform bilinear interpolation over two dimensions.
1738
1739 Parameters
1740 ----------
1741 xi : Iterable
1742 List of values or None for each dimension.
1743 scales : Iterable
1744 List of scales (coordinate arrays) for each dimension.
1745 values : np.ndarray
1746 The data array to interpolate.
1747
1748 Returns
1749 -------
1750 np.ndarray
1751 The interpolated data.
1752
1753 """
1754 index0, index1 = [i for i, v in enumerate(xi) if v is not None]
1755 t, u = [(xi[i] - scales[i][0])/(scales[i][1] - scales[i][0]) for i in (index0, index1)]
1756
1757 f00 = [slice(None, None)]*values.ndim
1758 f10 = [slice(None, None)]*values.ndim
1759 f01 = [slice(None, None)]*values.ndim
1760 f11 = [slice(None, None)]*values.ndim
1761 f00[index0], f00[index1] = 0, 0
1762 f10[index0], f10[index1] = 1, 0
1763 f01[index0], f01[index1] = 0, 1
1764 f11[index0], f11[index1] = 1, 1
1765
1766 return (
1767 (1 - t)*(1 - u)*values[tuple(f00)] +
1768 t*(1 - u)*values[tuple(f10)] +
1769 (1 - t)*u*values[tuple(f01)] +
1770 t*u*values[tuple(f11)]
1771 )
1772
1773
1774def _np_trilinear_interpolation(xi, scales, values):
1775 """
1776 Perform trilinear interpolation over three dimensions.
1777
1778 Parameters
1779 ----------
1780 xi : list
1781 List of values or None for each dimension.
1782 scales : list
1783 List of scales (coordinate arrays) for each dimension.
1784 values : np.ndarray
1785 The data array to interpolate.
1786
1787 Returns
1788 -------
1789 np.ndarray
1790 The interpolated data.
1791
1792 """
1793 index0, index1, index2 = [i for i, v in enumerate(xi) if v is not None]
1794 t, u, v = [(xi[i] - scales[i][0])/(scales[i][1] - scales[i][0]) for i in (index0, index1, index2)]
1795
1796 f000 = [slice(None, None)]*values.ndim
1797 f100 = [slice(None, None)]*values.ndim
1798 f010 = [slice(None, None)]*values.ndim
1799 f110 = [slice(None, None)]*values.ndim
1800 f001 = [slice(None, None)]*values.ndim
1801 f101 = [slice(None, None)]*values.ndim
1802 f011 = [slice(None, None)]*values.ndim
1803 f111 = [slice(None, None)]*values.ndim
1804
1805 f000[index0], f000[index1], f000[index2] = 0, 0, 0
1806 f100[index0], f100[index1], f100[index2] = 1, 0, 0
1807 f010[index0], f010[index1], f010[index2] = 0, 1, 0
1808 f110[index0], f110[index1], f110[index2] = 1, 1, 0
1809 f001[index0], f001[index1], f001[index2] = 0, 0, 1
1810 f101[index0], f101[index1], f101[index2] = 1, 0, 1
1811 f011[index0], f011[index1], f011[index2] = 0, 1, 1
1812 f111[index0], f111[index1], f111[index2] = 1, 1, 1
1813
1814 c00 = values[tuple(f000)]*(1 - t) + values[tuple(f100)]*t
1815 c10 = values[tuple(f010)]*(1 - t) + values[tuple(f110)]*t
1816 c01 = values[tuple(f001)]*(1 - t) + values[tuple(f101)]*t
1817 c11 = values[tuple(f011)]*(1 - t) + values[tuple(f111)]*t
1818
1819 c0 = c00*(1 - u) + c10*u
1820 c1 = c01*(1 - u) + c11*u
1821
1822 return c0*(1 - v) + c1*v
1823
1824
1825def _check_index_ranges(arr_size: int,
1826 i0: Union[int, np.integer],
1827 i1: Union[int, np.integer]
1828 ) -> Tuple[int, int]:
1829 """
1830 Adjust index ranges to ensure they cover at least two indices.
1831
1832 Parameters
1833 ----------
1834 arr_size : int
1835 The size of the array along the dimension.
1836 i0 : int
1837 The starting index.
1838 i1 : int
1839 The ending index.
1840
1841 Returns
1842 -------
1843 Tuple[int, int]
1844 Adjusted starting and ending indices.
1845
1846 Notes
1847 -----
1848 This function ensures that the range between `i0` and `i1` includes at least
1849 two indices for interpolation purposes.
1850
1851 """
1852 i0, i1 = int(i0), int(i1)
1853 if i0 == 0:
1854 return (i0, i1 + 2) if i1 == 0 else (i0, i1 + 1)
1855 elif i0 == arr_size:
1856 return i0 - 2, i1
1857 else:
1858 return i0 - 1, i1 + 1
1859
1860
1861def _cast_shape_tuple(input: Union[int, Iterable[int]]
1862 ) -> tuple[int, ...]:
1863 """
1864 Cast an input to a tuple of integers.
1865
1866 Parameters
1867 ----------
1868 input : int | Iterable[int]
1869 The input to cast.
1870
1871 Returns
1872 -------
1873 tuple[int]
1874 The input cast as a tuple of integers.
1875
1876 Raises
1877 ------
1878 TypeError
1879 If the input is neither an integer nor an iterable of integers.
1880 """
1881 if isinstance(input, int):
1882 return (input,)
1883 elif isinstance(input, Iterable):
1884 return tuple(int(i) for i in input)
1885 else:
1886 raise TypeError("Input must be an integer or an iterable of integers.")
1887
1888
1889def _parse_index_inputs(input: Union[int, slice, Iterable[Union[int, None]], None]
1890 ) -> slice:
1891 """
1892 Parse various slice input formats into a standard slice object.
1893
1894 Parameters
1895 ----------
1896 input : int | slice | Iterable[Union[int, None]] | None
1897 The input to parse.
1898 arr_size : int
1899 The size of the array along the dimension.
1900
1901 Returns
1902 -------
1903 slice
1904 The parsed slice object.
1905
1906 Raises
1907 ------
1908 TypeError
1909 If the input type is unsupported.
1910 ValueError
1911 If the input iterable does not have exactly two elements.
1912 """
1913 if isinstance(input, int):
1914 return slice(input, input + 1)
1915 elif isinstance(input, Iterable):
1916 return slice(*input)
1917 elif input is None:
1918 return slice(None)
1919 else:
1920 raise TypeError("Unsupported input type for slicing.")
1921
1922
1923def _parse_value_inputs(dimproxy,
1924 value,
1925 scale_exists: bool = True
1926 ) -> slice:
1927 if value is None:
1928 return slice(None)
1929 if not scale_exists:
1930 raise ValueError("Cannot parse value inputs when scale does not exist.")
1931 dim = dimproxy[:]
1932 if not isinstance(value, Iterable):
1933 insert_index = np.searchsorted(dim, value)
1934 return slice(*_check_index_ranges(dim.size, insert_index, insert_index))
1935 else:
1936 temp_range = list(value)
1937 if temp_range[0] is None:
1938 temp_range[0] = -np.inf
1939 if temp_range[-1] is None:
1940 temp_range[-1] = np.inf
1941 insert_indices = np.searchsorted(dim, temp_range)
1942 return slice(*_check_index_ranges(dim.size, *insert_indices))
1943
1944
1945def _parse_ivalue_inputs(dimsize,
1946 input: Union[Union[int, float], slice, Iterable[Union[Union[int, float], None]], None]
1947 ) -> slice:
1948 """
1949 Parse various slice input formats into a standard slice object.
1950
1951 Parameters
1952 ----------
1953 input : int | slice | Iterable[Union[int, None]] | None
1954 The input to parse.
1955 arr_size : int
1956 The size of the array along the dimension.
1957
1958 Returns
1959 -------
1960 slice
1961 The parsed slice object.
1962
1963 Raises
1964 ------
1965 TypeError
1966 If the input type is unsupported.
1967 ValueError
1968 If the input iterable does not have exactly two elements.
1969 """
1970 if input is None:
1971 return slice(None)
1972 elif isinstance(input, (int, float)):
1973 i0, i1 = math.floor(input), math.ceil(input)
1974 elif isinstance(input, Iterable):
1975 i0, i1 = math.floor(input[0]), math.ceil(input[1])
1976 else:
1977 raise TypeError("Unsupported input type for slicing.")
1978
1979 if i0 > i1:
1980 i0, i1 = i1, i0
1981 i0, i1 = max(0, i0), min(dimsize - 1, i1)
1982 if i0 > i1:
1983 i0, i1 = i1, i0
1984 i0, i1 = max(0, i0), min(dimsize - 1, i1)
1985 if (i1 - i0) < 2:
1986 return slice(i0, i1 + 2 - (i1-i0))
1987 else:
1988 return slice(i0, i1)
