zkit.mwigner.transform

zkit.mwigner.transform — research-grade marginal Wigner transform (ħ = 1) and the associated phase-space metrics, ported from wigner_quantum.py and integrated into zkit. No plotting here; figures live in zkit.mwigner.analysis.

The x-marginal Wigner function is

W(x, p_x) = (dx*dy/π) ∫∫ ψ*(x+ξ, y) ψ(x-ξ, y) e^{+2 i p_x ξ} dy dξ

evaluated on a uniform momentum grid p_x ∈ [-p_max, p_max] with spacing dp. The wavefunction is zero-padded by M_max on each side before the sliding-window transform, so the returned W covers the whole x-domain and no tail is cropped.

zkit.mwigner.transform.get_device()

Best available torch device (cuda > mps > cpu).

zkit.mwigner.transform.marginal_wigner(psi_np, dxs, axis=0, device=None, chunk_size=64, M_max=None, window_type='none', warn_on_alias=True)

Marginal Wigner transform of an N-dimensional wavefunction.

Computes the Wigner function marginalised over every axis EXCEPT axis (the “volume integration” over all transverse coordinates):

W(x_a, p_a) = (prod_i dx_i / pi) int … int psi*(x_a+xi, x_T) psi(x_a-xi, x_T) x e^{+2 i p_a xi} d xi prod_{T!=a} dx_T

For a 2D wavefunction with axis=0 this reduces exactly to the x-marginal Wigner used throughout zkit.mwigner; for a 3D wavefunction it integrates out both transverse coordinates, giving the true 3D marginal Wigner (no separate 3D routine needed).

Parameters:

  • psi_np (np.ndarray) – Spatial wavefunction, shape (n_0, …, n_{d-1}), any d >= 2.

  • dxs (sequence of float) – Grid spacings, one per spatial axis (length == psi_np.ndim).

  • axis (int) – Active axis to keep (marginalise over all others).

  • (chunk_size (...) – routine; see get_wigner_research for semantics).

  • M_max – routine; see get_wigner_research for semantics).

  • window_type – routine; see get_wigner_research for semantics).

  • 2D (warn_on_alias as in the) – routine; see get_wigner_research for semantics).

Returns:

  • W ((n_axis, window_size) np.ndarray (normalised by prod dx / pi))

  • p_v ((window_size,) np.ndarray momentum grid (increasing, dp>0))

zkit.mwigner.transform.get_wigner_research(psi_2d_np, dx, dy, device=None, chunk_size=64, M_max=None, window_type='none', warn_on_alias=True)

Research-grade 2D marginal Wigner transform (ħ = 1).

Thin wrapper over marginal_wigner() for the common 2D case (active axis 0). Computes

W(x, p_x) = (dx*dy/pi) int int psi*(x+xi, y) psi(x-xi, y) e^{+2 i p_x xi} dy dxi

over the FULL input domain (the wavefunction is zero-padded by M_max on each side along x so no tail is cropped). See marginal_wigner() for the full N-dimensional generalisation (including 3D volume integration).

Parameters:

  • psi_2d_np (ndarray)

  • dx (float)

  • dy (float)

  • device (torch.device)

  • chunk_size (int)

  • M_max (int)

  • window_type (str)

  • warn_on_alias (bool)

Return type:

tuple

zkit.mwigner.transform.load_psi(file_path, snap_idx)

Read snapshot snap_idx from an h5 wavefunction file as complex.

Handles both the 2D layout (nx, ny, 2) and the general N-D PETSc layout (..., 2) (the trailing axis is [real, imag]).

Parameters:

  • file_path (str)

  • snap_idx (int)

Return type:

ndarray

zkit.mwigner.transform.measure_wigner_negativity(W, dx, dp, normalize=True)

Wigner negativity (volume of the negative regions).

Parameters:

  • W (ndarray)

  • dx (float)

  • dp (float)

  • normalize (bool)

Return type:

float

zkit.mwigner.transform.verify_purities_and_entropy(psi, W, dxs, dp, axis=0, M_max=None)

Global norm, reduced purities, and Von Neumann entanglement entropy.

Works for an N-dimensional wavefunction. The reduced (axis-marginal) density matrix is formed over the FULL active-axis domain (independent of M_max) by tracing out every transverse coordinate:

rho[a, b] = (prod_{T} dT) * sum_T psi[a, T] psi*[b, T]

Parameters:

  • psi (np.ndarray) – Spatial wavefunction, shape (n_0, …, n_{d-1}).

  • W (np.ndarray) – Marginal Wigner from marginal_wigner() (shape (n_axis, n_p)).

  • dxs (float or sequence of float) – Grid spacings; a single float or one per spatial axis.

  • dp (float) – Momentum-grid spacing of W.

  • axis (int) – Active axis (the one NOT traced out).