Phase Space Functions

def wigner(rho: Qobj, xvec: ArrayLike, yvec: ArrayLike, g: float = np.sqrt(2), 
           sparse: bool = False, method: str = 'clenshaw') -> np.ndarray:
    """
    Calculate Wigner function on phase space grid.
    
    Parameters:
    - rho: Density matrix or state vector
    - xvec: Array of x-coordinates (position-like)
    - yvec: Array of y-coordinates (momentum-like)
    - g: Scaling factor (default √2)
    - sparse: Use sparse matrices
    - method: 'clenshaw', 'iterative', or 'laguerre'
    
    Returns:
    - ndarray: Wigner function W(x,y) on grid
    """

def wigner_transform(psi: Qobj, op: Qobj, g: float = np.sqrt(2)) -> Qobj:
    """
    Calculate Wigner transform of operator with respect to state.
    
    Parameters:
    - psi: Reference state
    - op: Operator to transform
    - g: Scaling factor
    
    Returns:
    - Qobj: Wigner-transformed operator
    """

def qfunc(rho: Qobj, xvec: ArrayLike, yvec: ArrayLike, g: float = np.sqrt(2), 
          sparse: bool = False) -> np.ndarray:
    """
    Calculate Q-function (Husimi distribution) on phase space grid.
    
    Parameters:
    - rho: Density matrix or state vector
    - xvec: Array of x-coordinates
    - yvec: Array of y-coordinates  
    - g: Scaling factor (default √2)
    - sparse: Use sparse matrices
    
    Returns:
    - ndarray: Q-function Q(α) = ⟨α|ρ|α⟩/π on grid
    """

class QFunc:
    """
    Q-function calculation class for efficient repeated calculations.
    """
    def __init__(self, rho: Qobj = None, xvec: ArrayLike = None, yvec: ArrayLike = None):
        """
        Initialize Q-function calculator.
        
        Parameters:
        - rho: Density matrix
        - xvec: X-coordinate array
        - yvec: Y-coordinate array
        """
    
    def qfunc(self, rho: Qobj = None) -> np.ndarray:
        """
        Calculate Q-function for given or stored state.
        
        Parameters:
        - rho: Density matrix (uses stored if None)
        
        Returns:
        - ndarray: Q-function values
        """

def spin_wigner(rho: Qobj, theta: ArrayLike, phi: ArrayLike) -> np.ndarray:
    """
    Calculate spin Wigner function on sphere.
    
    Parameters:
    - rho: Spin density matrix
    - theta: Polar angle array (0 to π)
    - phi: Azimuthal angle array (0 to 2π)
    
    Returns:
    - ndarray: Spin Wigner function on sphere
    """

def spin_q_function(rho: Qobj, theta: ArrayLike, phi: ArrayLike) -> np.ndarray:
    """
    Calculate spin Q-function on sphere.
    
    Parameters:
    - rho: Spin density matrix
    - theta: Polar angle array
    - phi: Azimuthal angle array
    
    Returns:
    - ndarray: Spin Q-function on sphere
    """

import qutip as qt
import numpy as np
import matplotlib.pyplot as plt

# Phase space grid
x_max = 4
N_points = 100
xvec = np.linspace(-x_max, x_max, N_points)
yvec = np.linspace(-x_max, x_max, N_points)

# Various quantum states for Wigner function analysis

# 1. Vacuum state
N_levels = 20
vacuum = qt.basis(N_levels, 0)
rho_vacuum = vacuum * vacuum.dag()

W_vacuum = qt.wigner(rho_vacuum, xvec, yvec)
print(f"Vacuum Wigner function:")
print(f"  Min value: {W_vacuum.min():.3f}")
print(f"  Max value: {W_vacuum.max():.3f}")
print(f"  Should be Gaussian with max ≈ 0.318")

# 2. Coherent state
alpha = 1.5 + 1.0j
coherent = qt.coherent(N_levels, alpha)
rho_coherent = coherent * coherent.dag()

W_coherent = qt.wigner(rho_coherent, xvec, yvec)
Q_coherent = qt.qfunc(rho_coherent, xvec, yvec)

print(f"Coherent state α = {alpha}:")
print(f"  Wigner min/max: [{W_coherent.min():.3f}, {W_coherent.max():.3f}]")
print(f"  Q-function min/max: [{Q_coherent.min():.3f}, {Q_coherent.max():.3f}]")

# Find peak positions
X, Y = np.meshgrid(xvec, yvec)
max_idx = np.unravel_index(np.argmax(W_coherent), W_coherent.shape)
x_peak, y_peak = xvec[max_idx[1]], yvec[max_idx[0]]
print(f"  Peak at (x,p) ≈ ({x_peak:.2f}, {y_peak:.2f})")
print(f"  Expected: ({np.sqrt(2)*alpha.real:.2f}, {np.sqrt(2)*alpha.imag:.2f})")

# 3. Squeezed state
squeeze_param = 0.8
squeezed = qt.squeeze(N_levels, squeeze_param) * vacuum
rho_squeezed = squeezed * squeezed.dag()

W_squeezed = qt.wigner(rho_squeezed, xvec, yvec)
print(f"Squeezed vacuum (r={squeeze_param}):")
print(f"  Wigner min/max: [{W_squeezed.min():.3f}, {W_squeezed.max():.3f}]")

# 4. Fock state (non-classical)
n_photons = 2
fock = qt.basis(N_levels, n_photons)
rho_fock = fock * fock.dag()

W_fock = qt.wigner(rho_fock, xvec, yvec)
Q_fock = qt.qfunc(rho_fock, xvec, yvec)

print(f"Fock state |{n_photons}⟩:")
print(f"  Wigner min/max: [{W_fock.min():.3f}, {W_fock.max():.3f}]")  
print(f"  Has negative values: {W_fock.min() < -1e-10}")  # Non-classical
print(f"  Q-function is always ≥ 0: {Q_fock.min() >= -1e-10}")

# 5. Cat state (superposition)
cat_amp = 1.5
cat_state = (qt.coherent(N_levels, cat_amp) + qt.coherent(N_levels, -cat_amp)).unit()
rho_cat = cat_state * cat_state.dag()

W_cat = qt.wigner(rho_cat, xvec, yvec)
print(f"Schrödinger cat state (α = ±{cat_amp}):")
print(f"  Wigner min/max: [{W_cat.min():.3f}, {W_cat.max():.3f}]")
print(f"  Exhibits interference fringes with negative values")

# 6. Thermal state
n_thermal = 1.5
rho_thermal = qt.thermal_dm(N_levels, n_thermal)

W_thermal = qt.wigner(rho_thermal, xvec, yvec)
Q_thermal = qt.qfunc(rho_thermal, xvec, yvec)

print(f"Thermal state (⟨n⟩ = {n_thermal}):")
print(f"  Wigner min/max: [{W_thermal.min():.3f}, {W_thermal.max():.3f}]")
print(f"  Q-function max: {Q_thermal.max():.3f}")

# Comparison of calculation methods
# Different methods for Wigner function calculation
W_clenshaw = qt.wigner(rho_coherent, xvec, yvec, method='clenshaw')
W_iterative = qt.wigner(rho_coherent, xvec, yvec, method='iterative')

method_diff = np.abs(W_clenshaw - W_iterative).max()
print(f"Clenshaw vs iterative method difference: {method_diff:.2e}")

# Spin Wigner functions
# Spin-1/2 system
j = 0.5
theta = np.linspace(0, np.pi, 50)
phi = np.linspace(0, 2*np.pi, 100)

# Spin-up state
spin_up = qt.spin_state(j, j)  # |j,j⟩ = |1/2,1/2⟩
rho_spin_up = spin_up * spin_up.dag()

W_spin_up = qt.spin_wigner(rho_spin_up, theta, phi)
Q_spin_up = qt.spin_q_function(rho_spin_up, theta, phi)

print(f"Spin-up Wigner function:")
print(f"  Min/max: [{W_spin_up.min():.3f}, {W_spin_up.max():.3f}]")
print(f"  Peak should be at north pole (θ=0)")

# Spin coherent state
theta_coherent, phi_coherent = np.pi/3, np.pi/4  # Pointing direction
spin_coherent = qt.spin_coherent(j, theta_coherent, phi_coherent)
rho_spin_coherent = spin_coherent * spin_coherent.dag()

W_spin_coherent = qt.spin_wigner(rho_spin_coherent, theta, phi)
print(f"Spin coherent state:")
print(f"  Direction: θ={theta_coherent:.2f}, φ={phi_coherent:.2f}")
print(f"  Wigner min/max: [{W_spin_coherent.min():.3f}, {W_spin_coherent.max():.3f}]")

# Spin-1 system (larger)
j1 = 1.0
spin_states_j1 = [qt.spin_state(j1, m) for m in [1, 0, -1]]  # |1,1⟩, |1,0⟩, |1,-1⟩

# Superposition state
superpos_j1 = (spin_states_j1[0] + spin_states_j1[2]).unit()  # (|1,1⟩ + |1,-1⟩)/√2
rho_superpos_j1 = superpos_j1 * superpos_j1.dag()

W_superpos_j1 = qt.spin_wigner(rho_superpos_j1, theta, phi)
print(f"Spin-1 superposition state:")
print(f"  Wigner min/max: [{W_superpos_j1.min():.3f}, {W_superpos_j1.max():.3f}]")

# Phase space integral properties
# Wigner function should integrate to 1
X, Y = np.meshgrid(xvec, yvec)
dx = xvec[1] - xvec[0]
dy = yvec[1] - yvec[0]

integral_W_vacuum = np.sum(W_vacuum) * dx * dy / np.pi
integral_W_coherent = np.sum(W_coherent) * dx * dy / np.pi
integral_Q_coherent = np.sum(Q_coherent) * dx * dy / np.pi

print(f"Phase space integrals (should be 1):")
print(f"  Vacuum Wigner: {integral_W_vacuum:.3f}")
print(f"  Coherent Wigner: {integral_W_coherent:.3f}")
print(f"  Coherent Q-function: {integral_Q_coherent:.3f}")

# Q-function class usage
qfunc_calc = qt.QFunc(rho_coherent, xvec, yvec)
Q_from_class = qfunc_calc.qfunc()

# Verify consistency
consistency_check = np.abs(Q_from_class - Q_coherent).max()
print(f"Q-function class consistency: {consistency_check:.2e}")

# Time evolution in phase space
# Harmonic oscillator evolution
H = qt.num(N_levels)  # Number operator
times = np.linspace(0, 2*np.pi, 5)  # One period

print(f"Time evolution of coherent state:")
for i, t in enumerate(times):
    # Evolve state
    U = (-1j * H * t).expm()
    psi_t = U * coherent
    rho_t = psi_t * psi_t.dag()
    
    # Calculate Wigner function
    W_t = qt.wigner(rho_t, xvec, yvec)
    
    # Find peak
    max_idx = np.unravel_index(np.argmax(W_t), W_t.shape)
    x_peak, y_peak = xvec[max_idx[1]], yvec[max_idx[0]]
    
    print(f"  t = {t:.2f}: peak at ({x_peak:.2f}, {y_peak:.2f})")

# The coherent state should rotate in phase space with frequency ω=1
# Expected position: α(t) = α * exp(-iωt) = α * exp(-it)
alpha_expected = alpha * np.exp(-1j * times[-1])
x_expected = np.sqrt(2) * alpha_expected.real
y_expected = np.sqrt(2) * alpha_expected.imag
print(f"  Expected final position: ({x_expected:.2f}, {y_expected:.2f})")

# Visualization example (pseudo-code, actual plotting would need matplotlib)
def plot_phase_space_comparison(states_dict, xvec, yvec):
    """
    Example function showing how to visualize phase space functions.
    """
    print("Phase space visualization:")
    for name, rho in states_dict.items():
        W = qt.wigner(rho, xvec, yvec)
        Q = qt.qfunc(rho, xvec, yvec)
        
        print(f"  {name}:")
        print(f"    Wigner: min={W.min():.3f}, max={W.max():.3f}")
        print(f"    Q-func: min={Q.min():.3f}, max={Q.max():.3f}")
        print(f"    Non-classical (W<0): {W.min() < -1e-10}")

# Example usage
states_to_plot = {
    'Vacuum': rho_vacuum,
    'Coherent': rho_coherent, 
    'Squeezed': rho_squeezed,
    'Fock n=2': rho_fock,
    'Cat state': rho_cat,
    'Thermal': rho_thermal
}

plot_phase_space_comparison(states_to_plot, xvec, yvec)

# Phase space functions return numpy ndarrays containing
# the calculated quasi-probability distributions on the specified grids