greenWTE.io

IO module for Green’s function transport equation (WTE) calculations.

This module provides a high-level interface for storing and retrieving precomputed Green’s operators in an HDF5 container. It supports lazy dataset creation, dynamic resizing, and optional GPU (CuPy) compatibility for data I/O.

Functions

load_phono3py_data(filename, temperature, ...)

Load data from a phono3py-generated HDF5 file.

save_solver_result(filename, solver, **kwargs)

Save a solver run to an HDF5 file.

Classes

GreenContainer(path, nat3, nq, dtype, meta, ...)

Container for HDF5-stored Green's operators.

GreenOperatorLike(*args, **kwargs)

Used for type checking Green's operator-like objects.
class greenWTE.io.GreenContainer(path: PathLike, nat3: int | None = None, nq: int | None = None, dtype: dtype = <class 'numpy.complex128'>, meta: dict | None = None, tile_B: int = 512, read_only: bool = False)

Bases: object

Container for HDF5-stored Green’s operators.

This class manages the storage of precomputed Green’s operators indexed by frequency omega, wavevector magnitude k, and a q-point index q. Data is stored in a 5D array with shape (Nw, Nk, nq, m, m), where \(m = (3*n_\mathrm{at})^2\).

Parameters:

  • path (os.PathLike) – Path to the HDF5 file. Created if it does not exist.

  • nat3 (int, optional) – Number of atoms times 3; defines the block size of the operator. Only required if this is a new Green’s function file. Read from file if None.

  • nq (int, optional) – Number of q-points. Only required if this is a new Green’s function file. Read from file if None.

  • dtype (dtype or str, optional) – Complex or real data type for storage.

  • meta (dict, optional) – Arbitrary metadata to persist as root attributes.

  • tile_B (int, optional) – Block size for chunking the matrix dimensions.

  • read_only (bool, optional) – If True, open the file in read-only mode. Default is False. This will allow multiple readers to access the file simultaneously.

f

Open HDF5 handle.

Type:

h5py.File

ds_w, ds_k

The 1-D frequency and wavevector datasets.

Type:

h5py.Dataset

ds_mask

Boolean mask indicating which (w, k, q) blocks exist.

Type:

h5py.Dataset

ds_tens

Tensor operator blocks with shape (Nw, Nk, nq, m, m).

Type:

h5py.Dataset

m

Matrix block dimension, nat3**2.

Type:

int

__init__(path: PathLike, nat3: int | None = None, nq: int | None = None, dtype: dtype = <class 'numpy.complex128'>, meta: dict | None = None, tile_B: int = 512, read_only: bool = False) → None

Initialize the GreenContainer.

_ensure_main() → None

Create /green and /mask if missing, with current Nw,Nk (>0).

close() → None

Close the underlying HDF5 file.

find_indices(w: float, k: float, atol: float = 0.0) → tuple[int, int]

Return (iw, ik) indices for the given (w, k).

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • atol (float, optional) – Absolute tolerance for comparing existing values.

Returns:

(iw, ik) indices into the omega and k datasets.

Return type:

tuple of int

Raises:

KeyError – If the omega or k value is not found within the specified tolerance.

get(w: float, k: float, q: int, atol: float = 0.0) → ndarray

Load a single Green’s operator block G[w, k, q].

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • q (int) – q-point index.

  • atol (float, optional) – Absolute tolerance for comparing existing values.

Returns:

A GPU array with shape (m, m) and complex dtype matching the on-disk dataset.

Return type:

cupy.ndarray

Raises:

KeyError – If the requested block is missing.

get_bz_block(w: float, k: float, atol: float = 0.0) → ndarray

Load all q-blocks for a given pair (w, k).

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • atol (float, optional) – Absolute tolerance for comparing existing values.

Returns:

The (nq, m, m) array of operator blocks for the specified (w, k).

Return type:

cupy.ndarray

Raises:

KeyError – If the requested block is missing.

has(w: float, k: float, q: int, atol: float = 0.0) → bool

Check whether a Green’s operator block exists for given indices.

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • q (int) – q-point index.

  • atol (float, optional) – Absolute tolerance for comparing existing values.

Returns:

True if the block exists, False otherwise.

Return type:

bool

has_bz_block(w: float, k: float, atol: float = 0.0) → bool

Check whether a Green’s operator block exists for the full Brillouin zone.

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • atol (float, optional) – Absolute tolerance for comparing existing values.

Returns:

True if the block exists, False otherwise.

Return type:

bool

indices(w: float, k: float, atol: float = 0.0) → tuple[int, int]

Return (iw, ik) indices for the given (w, k). Append if missing.

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • atol (float, optional) – Absolute tolerance for comparing existing values.

Returns:

(iw, ik) indices into the omega and k datasets.

Return type:

tuple of int

ks(w: float | None = None) → ndarray | None

Return all stored \(k\) values.

Parameters:

w (float, optional) – If provided, return only the wavevector magnitudes for which any q-block exists at this omega (according to the mask).

Returns:

One-dimensional array of wavevector magnitudes in 1/m. If w is provided but not found, returns None.

Return type:

numpy.ndarray | None

omegas(k: float | None = None) → ndarray | None

Return all stored \(\omega\) values.

Parameters:

k (float, optional) – If provided, return only the frequencies for which any q-block exists at this k (according to the mask).

Returns:

One-dimensional array of frequencies in rad/s. If k is provided but not found, returns None.

Return type:

numpy.ndarray | None

put(w: float, k: float, q: int, data: ndarray, atol: float = 0.0, flush: bool = True) → None

Store a Green’s operator block G[w, k, q].

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • q (int) – q-point index.

  • data (cupy.ndarray) – Operator block to store; must have shape (m, m).

  • atol (float, optional) – Absolute tolerance for comparing existing values.

  • flush (bool, optional) – If True, flush the file to disk after writing.

Raises:

ValueError – If the shape of data does not match (m, m).

put_bz_block(w: float, k: float, data: ndarray, atol: float = 0.0, flush: bool = True) → None

Store the full Brillouin-zone block for a pair (w, k).

Parameters:

  • w (float) – Angular frequency.

  • k (float) – Wavevector magnitude.

  • data (cupy.ndarray or greenWTE.green.RTAGreenOperator) – Array of operator blocks to store; must have shape (nq, m, m).

  • atol (float, optional) – Absolute tolerance for comparing existing values.

  • flush (bool, optional) – If True, flush the file to disk after writing.

Raises:

ValueError – If the shape of data does not match (nq, m, m).

class greenWTE.io.GreenOperatorLike(*args, **kwargs)

Bases: Protocol

Used for type checking Green’s operator-like objects.

__init__(*args, **kwargs)
_abc_impl = <_abc._abc_data object>
_green: ndarray
_is_protocol = True
_is_runtime_protocol = True
greenWTE.io._ensure(f: File, name: str, *, shape: tuple[int, ...], maxshape: tuple[int, ...] | None, chunks: tuple[int, ...], dtype: dtype | str, **kwargs) → Dataset

Ensure that a dataset exists in an HDF5 file.

If the dataset named name exists in the file f, it is returned. Otherwise, the dataset is created with the given shape, dtype, chunking, compression, and unlimited dimensions.

Parameters:

  • f (h5py.File) – Open HDF5 file object.

  • name (str) – Dataset path.

  • shape (tuple of int) – Initial dataset shape.

  • maxshape (tuple of int or None) – Maximum shape; use None for dimensions that may grow.

  • chunks (tuple of int) – Chunk sizes per dimension.

  • dtype (dtype or str) – On-disk dtype.

  • **kwargs (dict) – Additional keyword arguments forwarded to h5py.File.create_dataset().

Returns:

The existing or newly created dataset.

Return type:

h5py.Dataset

greenWTE.io._find_index_1d(dset, value: float, atol: float = 0.0)

Find the index of value in a 1-D dataset.

Parameters:

  • dset (h5py.Dataset) – One-dimensional dataset.

  • value (float) – Scalar value to find or append.

  • atol (float, optional) – Absolute tolerance for equality checks.

Returns:

Index if found; -1 otherwise.

Return type:

int

greenWTE.io._find_or_append_1d(dset, value: float, atol: float = 0.0)

Find the index of value in a 1-D dataset, appending if absent.

Parameters:

  • dset (h5py.Dataset) – One-dimensional dataset.

  • value (float) – Scalar value to find or append.

  • atol (float, optional) – Absolute tolerance for equality checks.

Returns:

The index of the matching or newly appended value.

Return type:

int

greenWTE.io._scalar_to_float(x) → float

Convert Python/NumPy/CuPy scalar or size-1 array to float.

greenWTE.io.load_phono3py_data(filename: PathLike, temperature: float, dir_idx: int, exclude_gamma: bool = True, dtyper: dtype = <class 'numpy.float64'>, dtypec: dtype = <class 'numpy.complex128'>) → tuple[ndarray, ndarray, ndarray, ndarray, ndarray, float, ndarray]

Load data from a phono3py-generated HDF5 file.

Parameters:

  • filename (os.PathLike) – Path to the :mod:phono3py HDF5 file.

  • temperature (float) – Temperature in Kelvin at which to read linewidths and heat capacities.

  • dir_idx (int) – Cartesian direction index of the velocity operator (0=x, 1=y, 2=z).

  • exclude_gamma (bool, optional) – If True, drop the \(\Gamma\) point from the beginning of each q-point-indexed dataset.

  • dtyper (dtype or str, optional) – Real data type for loading real-valued arrays.

  • dtypec (dtype or str, optional) – Complex data type for loading complex-valued arrays.

Returns:

  • qpoint (cupy.ndarray) – q-points of, shape (nq, 3).

  • velocity_operator (cupy.ndarray) – Velocity operator in m/s, shape (nq, nat3, nat3).

  • phonon_freq (cupy.ndarray) – Phonon frequencies in rad/s, shape (nq, nat3).

  • linewidth (cupy.ndarray) – Linewidths of each mode in rad/s, shape (nq, nat3).

  • heat_capacity (cupy.ndarray) – Heat capacity of each mode in J/m^3/K, shape (nq, nat3).

  • volume (float) – Volume of the system in m^3.

  • weight (cupy.ndarray) – Normalized q-point weights, shape (nq,).

Raises:

ValueError – If the specified temperature is not found in the input file.

greenWTE.io.save_solver_result(filename: PathLike, solver: SolverBase, **kwargs) → None

Save a solver run to an HDF5 file.

This is a light-weight, analysis-friendly snapshot of a full frequency-domain solve, including the temperature response, Wigner distribution, and iteration metadata.

Parameters:

  • filename (os.PathLike) – Path to the output HDF5 file.

  • solver (greenWTE.solver.SolverBase) – A solver instance (e.g., iterative.IterativeWTESolver or green.GreenWTESolver) that exposes the recorded arrays as attributes.

  • **kwargs (dict) – Additional key-value pairs to store as datasets in the root group. Scalars, lists, and arrays are supported.

Notes

The function records (if available):

  • dT: complex temperature response over frequencies (Nw,).

  • dT_init: initial guess (Nw,).

  • n: Wigner distribution (Nw, nq, m) or similar aggregate.

  • niter: number of outer iterations per frequency.

  • iter_time: wall-clock time per frequency in seconds.

  • gmres_residual: GMRES residual history (ragged) concatenated.

  • dT_iterates: stored iterates of dT (stacked).

  • n_norms: norms of n per outer iteration.

  • source: source term used in the solve.

  • omega: angular frequencies in rad/s.

  • k: wavevector magnitudes in rad/m.

In addition, command-line arguments found in solver.command_line_args are persisted as datasets if present.