DensityMatrixSimulation.py

import numpy as np
def ket(v):
    return np.array(v, dtype=complex).reshape(-1, 1)
def density_matrix(state):
    return state @ state.conj().T  # |psi><psi|
def dagger(U):
    return U.conj().T
def purity(rho):
    return float(np.trace(rho @ rho).real)  # Tr[rho^2]
def probability(rho, phi):
    P_phi = density_matrix(phi)             # projector |phi><phi|
    return float(np.trace(rho @ P_phi).real)  # Tr[rho P_phi]
def properties(rho):
    # Hermitian: rho† = rho
    hermitian = np.allclose(rho, rho.conj().T)
    # PSD: eigenvalues >= 0
    eigs = np.linalg.eigvalsh(rho)  # for Hermitian matrices
    psd = np.all(eigs >= -1e-12)
    # normalized: Tr[rho] = 1
    normalized = np.allclose(np.trace(rho), 1.0)
    return hermitian, psd, normalized, eigs
def partial_trace_B(rho_AB):
    # rho_A = Tr_B[rho_AB] for two qubits (A,B) dims 2x2
    # block trace: sum diagonal 2x2 blocks
    return rho_AB[0:2, 0:2] + rho_AB[2:4, 2:4]
# representation of states
state_0 = ket([1, 0])
state_1 = ket([0, 1])
state_plus = (state_0 + state_1) / np.sqrt(2)
alpha, beta = 1/np.sqrt(3), np.sqrt(2/3)
state_psi = alpha * state_0 + beta * state_1
rho_psi = density_matrix(state_psi)
rho_0 = density_matrix(state_0)
rho_plus = density_matrix(state_plus)
print("rho_psi =\n", rho_psi)
print("|0><0| =\n", rho_0)
print("|+><+| =\n", rho_plus)
# requirements and properties
for name, rho in [("rho_psi", rho_psi), ("rho_0", rho_0), ("rho_plus", rho_plus)]:
    herm, psd, norm, eigs = properties(rho)
    print(f"\n{name}: Hermitian={herm}, PSD={psd}, Tr=1={norm}, eigenvalues={eigs}")
# types of states (pure vs mixed)
# mixed: rho = sum_i p_i |psi_i><psi_i|
rho_mixed = 0.5 * density_matrix(state_0) + 0.5 * density_matrix(state_1)
print("\nrho_mixed =\n", rho_mixed)
# pure projector property: rho^2 = rho , mixed generally not
print("\nPure check (rho_psi^2 == rho_psi):", np.allclose(rho_psi @ rho_psi, rho_psi))
print("Mixed check (rho_mixed^2 == rho_mixed):", np.allclose(rho_mixed @ rho_mixed, rho_mixed))
# purity measure: Tr[rho^2]
print("\nPurity:")
print("Tr[rho_psi^2] =", purity(rho_psi))
print("Tr[rho_mixed^2] =", purity(rho_mixed))
# probabilities
# p_phi = Tr[rho P_phi]
state_minus = (state_0 - state_1) / np.sqrt(2)
print("\nProbabilities using p_phi = Tr[rho P_phi]:")
print("For rho_psi: p(0) =", probability(rho_psi, state_0), " p(1) =", probability(rho_psi, state_1))
print("For rho_mixed: p(0) =", probability(rho_mixed, state_0), " p(1) =", probability(rho_mixed, state_1))
print("For rho_mixed: p(+) =", probability(rho_mixed, state_plus), " p(-) =", probability(rho_mixed, state_minus))
# gates (unitary evolution)
# rho' = U rho U†
H = (1/np.sqrt(2)) * np.array([[1, 1],
                               [1, -1]], dtype=complex)
rho_prime = H @ rho_mixed @ dagger(H)
print("\nH =\n", H)
print("rho' = H rho_mixed H† =\n", rho_prime)
# entangled states + reduced states
# Bell state example used for reduced state result
state_00 = np.kron(state_0, state_0)
state_11 = np.kron(state_1, state_1)
state_bell = (state_00 + state_11) / np.sqrt(2)
rho_AB = density_matrix(state_bell)
rho_A = partial_trace_B(rho_AB)
print("\nrho_AB (Bell) =\n", rho_AB)
print("rho_A = Tr_B[rho_AB] =\n", rho_A)
print("Tr[(rho_A)^2] =", purity(rho_A))
# decoherence (system + environment interaction then reduction)
# system in |+> environment in |0> interact by a unitary (CNOT) then ignore environment
CNOT = np.array([[1,0,0,0],
                 [0,1,0,0],
                 [0,0,0,1],
                 [0,0,1,0]], dtype=complex)
state_SE = np.kron(state_plus, state_0)
rho_SE = density_matrix(CNOT @ state_SE)
rho_system = partial_trace_B(rho_SE)
print("\nrho_SE (after interaction) =\n", rho_SE)
print("Reduced system state (after ignoring environment) =\n", rho_system)