import os
from typing import Any, Optional, Union
import h5py
import numpy as np
from qutip import Qobj
[docs]
class SpectrumData:
def __init__(
self,
energy_table: np.ndarray,
system_params: dict[str, Any],
param_name: Optional[str] = None,
param_vals: Optional[np.ndarray] = None,
state_table: Optional[Union[list[Qobj], np.ndarray]] = None,
matrixelem_table: Optional[dict[str, np.ndarray]] = None,
t1_table: Optional[dict[str, np.ndarray]] = None,
tphi_table: Optional[dict[tuple[int, int, str], np.ndarray]] = None,
d2E_table: Optional[dict[str, np.ndarray]] = None,
**kwargs,
) -> None:
self.system_params = system_params
self.param_name = param_name
self.param_vals = param_vals
self.energy_table = energy_table
self.state_table = state_table if state_table is not None else []
self.matrixelem_table = matrixelem_table if matrixelem_table is not None else {}
self.t1_table = t1_table if t1_table is not None else {}
self.tphi_table = tphi_table if tphi_table is not None else {}
self.d2E_table = d2E_table if d2E_table is not None else {}
for dataname, data in kwargs.items():
setattr(self, dataname, data)
[docs]
def subtract_ground(self) -> None:
"""Subtract ground state energies from spectrum.
The energy_table has dimensions (param_vals_count, evals_count).
This method subtracts the ground state energy (column 0) from all states
at each parameter value (each row) and then removes the ground state column
which becomes identically zero.
"""
# Reshape ground state energies to allow broadcasting
ground_energies = self.energy_table[:, 0].reshape(-1, 1)
self.energy_table = self.energy_table - ground_energies
[docs]
def filewrite(self, filename: str, overwrite: bool = False) -> None:
"""Save the SpectrumData to an HDF5 file."""
if not overwrite and os.path.exists(filename):
raise FileExistsError(
f"The file '{filename}' already exists and overwrite is set to False."
)
directory = os.path.dirname(filename)
if directory:
os.makedirs(directory, exist_ok=True)
with h5py.File(filename, "w") as f:
f.create_dataset("energy_table", data=self.energy_table)
f.create_dataset("param_vals", data=self.param_vals)
f.create_dataset("param_name", data=np.bytes_(self.param_name))
f.create_dataset("system_params", data=np.bytes_(str(self.system_params)))
if isinstance(self.state_table, list):
if self.state_table and isinstance(self.state_table[0], Qobj):
state_data = np.array([state.full() for state in self.state_table])
f.create_dataset("state_table", data=state_data)
else:
f.create_dataset("state_table", data=np.array(self.state_table))
elif self.state_table is not None:
f.create_dataset("state_table", data=self.state_table)
if self.matrixelem_table:
for key, value in self.matrixelem_table.items():
f.create_dataset(f"matrixelem_table/{key}", data=value)
if self.t1_table:
for key, value in self.t1_table.items():
f.create_dataset(f"t1_table/{key}", data=value)
if self.tphi_table:
for tphi_key, value in self.tphi_table.items():
key_str: str = ",".join(map(str, tphi_key))
f.create_dataset(f"tphi_table/{key_str}", data=value)
[docs]
@staticmethod
def read(filename: str) -> "SpectrumData":
"""Read SpectrumData from an HDF5 file."""
with h5py.File(filename, "r") as f:
energy_table = f["energy_table"][:]
param_vals = f["param_vals"][:]
param_name = f["param_name"][()].decode("utf-8")
system_params = eval(f["system_params"][()].decode("utf-8"))
state_table = f["state_table"][:] if "state_table" in f else None
matrixelem_table = (
{key: f[f"matrixelem_table/{key}"][:] for key in f["matrixelem_table"]}
if "matrixelem_table" in f
else {}
)
t1_table = (
{key: f[f"t1_table/{key}"][:] for key in f["t1_table"]}
if "t1_table" in f
else {}
)
tphi_table = (
{
tuple(key.split(",")): f[f"tphi_table/{key}"][:]
for key in f["tphi_table"]
}
if "tphi_table" in f
else {}
)
return SpectrumData(
energy_table=energy_table,
system_params=system_params,
param_name=param_name,
param_vals=param_vals,
state_table=state_table,
matrixelem_table=matrixelem_table,
t1_table=t1_table,
tphi_table=tphi_table,
)