Python API

zkit exposes three main layers:

High-level Run object that loads an entire simulation directory.
Low-level readers in zkit.io for individual HDF5 outputs.
Visualization and analysis helpers in zkit.viz and zkit.mwigner.

Top-level shortcuts

import zkit

# Convenience aliases exported from zkit.__init__
zkit.open_eig
zkit.reconstruct_knots
zkit.partition_of_unity
zkit.infinite_well_energy
zkit.find_binary
zkit.run
zkit.Run
zkit.load
zkit.read_eigen
zkit.read_evolution
zkit.read_timeseries
zkit.read_wfs
zkit.read_tdm
zkit.tdm_of_state
zkit.tdm_magnitude_of_state
zkit.reconstruct_wfs
zkit.reconstruct_static_wfs
zkit.wfs_to_vtk
zkit.compute_wfs_norm
zkit.detect_dimension
zkit.plot_wavefunction
zkit.plot_transition_diagram
zkit.plot_tdm_matrix
zkit.plot_tdm

Load a whole run

from zkit.simulation import Run

run = Run("run_dir", "h2p.inp")
print("dimension:", run.dim)
print("n_steps:", run.evolution.n_steps)
print("spectrum:", run.spectrum.values)
print("time:", run.time[:5])
print("dipoles:", run.dipoles.shape)
print("populations:", run.populations.shape)
print("energies:", run.energies.shape)
print("currents:", run.currents.shape)
print("autocorrelation:", run.autocorrelation.shape)
print("ts times:", run.ts.times)

The Run object parses the input file when available and falls back to data-derived metadata otherwise.

Read individual files

from zkit.io.eigen import read_eigen
from zkit.io.evolution import read_evolution
from zkit.io.ts import read_timeseries
from zkit.io.wavefunction import read_wfs
from zkit.io.tdm import read_tdm, tdm_of_state, tdm_magnitude_of_state
from zkit.io.base import detect_dimension

# Eigenvalues and eigenvectors
eigen = read_eigen("static/EigenData_h2p.inp.h5")
print(eigen.values.shape)
print(eigen.vectors.shape)

# Time evolution observables
evolution = read_evolution("td/TimeEvolutionData_h2p.inp.h5")
print(evolution.dimension)

# PETSc TS snapshots
ts = read_timeseries("ts_h2p.inp.h5")
print(ts.times)

# Wavefunction snapshots
wfs = read_wfs("td/wfs_h2p.inp.h5")
print(wfs.n_snapshots)
print(wfs.data.shape)

# TDM
tdm = read_tdm("static/EigenData_h2p.inp.h5")
print(tdm.keys())
print(tdm_of_state("static/EigenData_h2p.inp.h5", 0, axes="x")["tdm_x"].shape)

Visualize

from zkit.viz import plot_wavefunction, plot_tdm

plot_wavefunction(
    "td/wfs_h2p.inp.h5",
    step=0,
    outdir="figures",
    npoints=200,
    dpi=150,
    vtk="figures/wfs.vts",
)

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

Wigner analysis

from zkit.mwigner import transform, current, covariance, autocorr, analysis

# Wigner transform
W, p = transform.get_wigner_research(psi, dx, dy)

# Marginal current
Jx = current.current_x(W, p)

# Covariance matrix
Sigma = covariance.compute_4d_covariance(psi, dx, dy)
nu = covariance.symplectic_eigenvalues(Sigma)
print(covariance.gaussian_purity(nu))

# Autocorrelation analysis
A, A2, P, gap, bound = autocorr.autocorr_analysis(psi_series, dxs)