Visualization

zkit provides matplotlib-based visualization for wavefunctions and transition dipole data.

Wavefunction plots

from zkit.viz import plot_wavefunction

path = plot_wavefunction(
    "td/wfs_h2p.inp.h5",
    step=0,
    outdir="figures",
    npoints=240,
    dpi=150,
    axes=None,
    vtk="figures/wfs.vts",
)
print(path)

Dispatch by dimension:

1D — stacked Re/Im ψ and |ψ|^2
2D — |ψ|^2, Re ψ, and phase panels
3D — |ψ|^2 on the z = z_mid slice

Transition-dipole diagrams

from zkit.viz import plot_transition_diagram

plot_transition_diagram(
    "static/EigenData_h2p.inp.h5",
    outfile="figures/tdm_diagram.png",
    axes="x",
    min_mu=1e-3,
    color_by="strength",
    dpi=150,
)

from zkit.viz.tdm import plot_tdm_matrix

plot_tdm_matrix(
    "static/EigenData_h2p.inp.h5",
    outfile="figures/tdm_matrix.png",
    axes=["x", "z"],
    dpi=150,
)

from zkit.viz import plot_tdm

out = plot_tdm(
    "static/EigenData_h2p.inp.h5",
    outdir="figures",
    axes=["x", "z"],
    min_mu=1e-3,
    color_by="strength",
    prefix="h2p",
    dpi=150,
)

Batch plotting

from pathlib import Path
from zkit.viz import plot_wavefunction

run_dir = Path("/data/run01")
td_dir = run_dir / "td"
out_dir = run_dir / "figures"
out_dir.mkdir(exist_ok=True)

for path in sorted(td_dir.glob("wfs_*.h5")):
    plot_wavefunction(path, step=0, outdir=out_dir, npoints=200, dpi=150)