Writing Simple Datasets

Write 1D, 2D, and 3D arrays to PSI-style HDF5 or HDF4 files.

This example demonstrates the basics of writing PSI-style HDF files using write_hdf_data() for generic n-dimensional datasets, and wrhdf_1d(), wrhdf_2d(), wrhdf_3d() for dimensionality-specific writing.

import tempfile
from pathlib import Path
import numpy as np
from psi_io import write_hdf_data, wrhdf_1d, wrhdf_2d, wrhdf_3d, read_hdf_data

Construct a simple 3D dataset representing a hypothetical scalar field on a rectilinear (r, θ, φ) grid. Following PSI’s Fortran ordering convention, the data array is C-ordered as (nφ, nθ, nr) while the scales are supplied in physical dimension order (r, θ, φ):

nr, nt, np_ = 10, 20, 30
r = np.linspace(1.0, 5.0, nr, dtype=np.float32)
t = np.linspace(0.0, np.pi, nt, dtype=np.float32)
p = np.linspace(0.0, 2*np.pi, np_, dtype=np.float32)
f3d = np.random.default_rng(0).random((np_, nt, nr)).astype(np.float32)

Write the dataset to an HDF5 file with write_hdf_data(). Scales are passed as positional arguments in physical dimension order (r, θ, φ):

with tempfile.TemporaryDirectory() as tmpdir:
    out_h5 = Path(tmpdir) / "output.h5"
    write_hdf_data(out_h5, f3d, r, t, p)

    f_back, r_back, t_back, p_back = read_hdf_data(out_h5)
    print(f"Data shape (read back) : {f_back.shape}")
    print(f"r scale : shape={r_back.shape}, range=[{r_back[0]:.2f}, {r_back[-1]:.2f}]")
    print(f"t scale : shape={t_back.shape}, range=[{t_back[0]:.4f}, {t_back[-1]:.4f}]")
    print(f"p scale : shape={p_back.shape}, range=[{p_back[0]:.4f}, {p_back[-1]:.4f}]")
    print(f"Round-trip exact match : {np.array_equal(f3d, f_back)}")

Data shape (read back) : (30, 20, 10)
r scale : shape=(10,), range=[1.00, 5.00]
t scale : shape=(20,), range=[0.0000, 3.1416]
p scale : shape=(30,), range=[0.0000, 6.2832]
Round-trip exact match : True

Writing to an HDF4 file uses the identical call – only the file extension changes. Dispatch to the correct backend is handled automatically based on the extension:

with tempfile.TemporaryDirectory() as tmpdir:
    out_hdf = Path(tmpdir) / "output.hdf"
    write_hdf_data(out_hdf, f3d, r, t, p)

    f_back_h4 = read_hdf_data(out_hdf, return_scales=False)
    print(f"HDF4 data shape (read back) : {f_back_h4.shape}")
    print(f"Round-trip exact match      : {np.array_equal(f3d, f_back_h4)}")

HDF4 data shape (read back) : (30, 20, 10)
Round-trip exact match      : True

For 1D and 2D datasets the same conventions apply – the scale(s) precede the data in the positional argument list:

x = np.linspace(0.0, 2*np.pi, 64, dtype=np.float32)
f1d = np.sin(x)

y = np.linspace(0.0, np.pi, 32, dtype=np.float32)
f2d = np.outer(np.sin(x), np.cos(y)).astype(np.float32)

with tempfile.TemporaryDirectory() as tmpdir:
    write_hdf_data(Path(tmpdir) / "out1d.h5", f1d, x)
    write_hdf_data(Path(tmpdir) / "out2d.h5", f2d, x, y)

    r1d = read_hdf_data(Path(tmpdir) / "out1d.h5")
    r2d = read_hdf_data(Path(tmpdir) / "out2d.h5")
    print(f"1D : data={r1d[0].shape}, scale={r1d[1].shape}")
    print(f"2D : data={r2d[0].shape}, x={r2d[1].shape}, y={r2d[2].shape}")

1D : data=(64,), scale=(64,)
2D : data=(64, 32), x=(64,), y=(32,)

The legacy dimension-specific writers (wrhdf_1d(), wrhdf_2d(), wrhdf_3d()) provide a backward-compatible interface for existing PSI codes.

Attention

The argument order of these writers differs from write_hdf_data(): scales are interleaved between the filename and the data array (filename, x, [y, [z,]] f). They also default to sync_dtype=True, which silently casts scale arrays to match the data dtype, and always write to PSI’s standard dataset identifiers ('Data-Set-2' for HDF4, 'Data' for HDF5).

with tempfile.TemporaryDirectory() as tmpdir:
    wrhdf_1d(Path(tmpdir) / "legacy1d.h5", x, f1d)
    wrhdf_2d(Path(tmpdir) / "legacy2d.h5", x, y, f2d)
    wrhdf_3d(Path(tmpdir) / "legacy3d.h5", r, t, p, f3d)

    rl3d = read_hdf_data(Path(tmpdir) / "legacy3d.h5")
    print(f"Legacy 3D : data={rl3d[0].shape}, r={rl3d[1].shape}, t={rl3d[2].shape}, p={rl3d[3].shape}")
    print(f"Round-trip exact match : {np.array_equal(f3d, rl3d[0])}")

Legacy 3D : data=(30, 20, 10), r=(10,), t=(20,), p=(30,)
Round-trip exact match : True