Tensor Operations and Composite Systems

def tensor(*args) -> Qobj:
    """
    Calculate tensor product of quantum objects.
    
    Parameters:
    - *args: Variable number of Qobj instances to tensor together
    
    Returns:
    - Qobj: Tensor product of input objects
    """

def super_tensor(*args) -> Qobj:
    """
    Calculate tensor product of superoperators.
    
    Parameters:
    - *args: Variable number of superoperator Qobj instances
    
    Returns:
    - Qobj: Tensor product of superoperators
    """

def ptrace(obj: Qobj, sel) -> Qobj:
    """
    Partial trace over specified subsystems.
    
    Parameters:
    - obj: Composite quantum object (state or operator)
    - sel: Integer, list of integers, or boolean mask specifying subsystems to keep
    
    Returns:
    - Qobj: Reduced quantum object after tracing out unselected subsystems
    """

def expect(oper: Qobj, state: Qobj) -> float:
    """
    Calculate expectation value ⟨ψ|Ô|ψ⟩ or Tr(ρÔ).
    
    Parameters:
    - oper: Observable operator
    - state: Quantum state (ket or density matrix)
    
    Returns:
    - float: Expectation value
    """

def variance(oper: Qobj, state: Qobj) -> float:
    """
    Calculate variance ⟨Ô²⟩ - ⟨Ô⟩².
    
    Parameters:
    - oper: Observable operator
    - state: Quantum state
    
    Returns:
    - float: Variance of observable
    """

def composite(*args) -> Qobj:
    """
    Create composite quantum object from components.
    
    Parameters:
    - *args: Quantum objects representing subsystems
    
    Returns:
    - Qobj: Composite quantum object
    """

def expand_operator(oper: Qobj, N: int, targets: list, dims: list = None) -> Qobj:
    """
    Expand operator to act on larger Hilbert space.
    
    Parameters:
    - oper: Operator to expand
    - N: Total number of subsystems
    - targets: List of target subsystem indices
    - dims: Dimensions of each subsystem
    
    Returns:
    - Qobj: Expanded operator
    """

def tensor_swap(obj: Qobj, mask) -> Qobj:
    """
    Swap subsystems in tensor product structure.
    
    Parameters:
    - obj: Quantum object with tensor structure
    - mask: Permutation mask for swapping subsystems
    
    Returns:
    - Qobj: Object with swapped subsystem order
    """

def tensor_contract(obj: Qobj, *pairs) -> Qobj:
    """
    Contract tensor indices in quantum object.
    
    Parameters:
    - obj: Quantum object to contract
    - *pairs: Pairs of indices to contract
    
    Returns:
    - Qobj: Contracted quantum object  
    """

import qutip as qt
import numpy as np

# Basic tensor products
psi0 = qt.basis(2, 0)  # |0⟩
psi1 = qt.basis(2, 1)  # |1⟩
psi_00 = qt.tensor(psi0, psi0)  # |00⟩
psi_01 = qt.tensor(psi0, psi1)  # |01⟩
psi_10 = qt.tensor(psi1, psi0)  # |10⟩
psi_11 = qt.tensor(psi1, psi1)  # |11⟩

# Bell state construction
bell = (psi_00 + psi_11).unit()  # (|00⟩ + |11⟩)/√2

# Tensor product of operators
sx = qt.sigmax()
sy = qt.sigmay()
sz = qt.sigmaz()
I = qt.qeye(2)

# Two-qubit operators
sx_I = qt.tensor(sx, I)    # σₓ ⊗ I (X on first qubit)
I_sy = qt.tensor(I, sy)    # I ⊗ σᵧ (Y on second qubit)
sx_sz = qt.tensor(sx, sz)  # σₓ ⊗ σᵤ (X-Z interaction)

# Three-qubit system
psi_000 = qt.tensor(psi0, psi0, psi0)  # |000⟩
H_3qubit = qt.tensor(sz, I, I) + qt.tensor(I, sz, I) + qt.tensor(I, I, sz)

# Partial trace examples
rho_total = bell * bell.dag()  # Bell state density matrix
rho_A = qt.ptrace(rho_total, 0)  # Trace out second qubit, keep first
rho_B = qt.ptrace(rho_total, 1)  # Trace out first qubit, keep second

print(f"Original state purity: {(rho_total**2).tr():.3f}")
print(f"Reduced state A purity: {(rho_A**2).tr():.3f}")
print(f"Reduced state B purity: {(rho_B**2).tr():.3f}")

# Expectation values
exp_x_total = qt.expect(sx_I, bell)  # ⟨σₓ ⊗ I⟩
exp_x_reduced = qt.expect(sx, rho_A)  # ⟨σₓ⟩ on reduced state
print(f"Expectation values: total={exp_x_total:.3f}, reduced={exp_x_reduced:.3f}")

# Variance calculation
var_x = qt.variance(sx, psi0)  # Variance of σₓ in |0⟩ state
print(f"Variance of σₓ in |0⟩: {var_x}")

# Multi-qubit expectation values
correlator = qt.expect(sx_sz, bell)  # ⟨σₓ ⊗ σᵤ⟩
print(f"X-Z correlation: {correlator}")

# Multiple subsystem partial trace
# Three-qubit system, trace out middle qubit
psi_3q = qt.ghz_state(3)  # GHZ state |000⟩ + |111⟩
rho_3q = psi_3q * psi_3q.dag()
rho_13 = qt.ptrace(rho_3q, [0, 2])  # Keep qubits 0 and 2, trace out 1

# Expand operator to larger space
single_qubit_op = sx
three_qubit_op = qt.expand_operator(single_qubit_op, 3, [1])  # Apply to qubit 1 in 3-qubit system
print(f"Expanded operator shape: {three_qubit_op.shape}")

# Harmonic oscillator tensor products
N = 5
a1 = qt.tensor(qt.destroy(N), qt.qeye(N))    # Annihilation on mode 1
a2 = qt.tensor(qt.qeye(N), qt.destroy(N))    # Annihilation on mode 2
n1 = a1.dag() * a1  # Number operator for mode 1
n2 = a2.dag() * a2  # Number operator for mode 2

# Two-mode squeezed state
r = 0.5  # Squeezing parameter
S_two = (r * (a1.dag() * a2.dag() - a1 * a2)).expm()  # Two-mode squeezing
vacuum = qt.tensor(qt.basis(N, 0), qt.basis(N, 0))
squeezed_vacuum = S_two * vacuum

# Expectation values in two-mode system
n1_exp = qt.expect(n1, squeezed_vacuum)
n2_exp = qt.expect(n2, squeezed_vacuum)
correlation = qt.expect(a1.dag() * a2, squeezed_vacuum)
print(f"Mode populations: n1={n1_exp:.3f}, n2={n2_exp:.3f}")
print(f"Mode correlation: {correlation:.3f}")

# Composite system with different subsystem dimensions
# Qubit (dim=2) + harmonic oscillator (dim=N)
qubit_state = qt.basis(2, 0)
ho_state = qt.coherent(N, 1.0)
composite_state = qt.tensor(qubit_state, ho_state)

# Operators on composite system
qubit_op = qt.tensor(sx, qt.qeye(N))  # Pauli-X on qubit
ho_op = qt.tensor(qt.qeye(2), qt.num(N))  # Number op on oscillator
interaction = qt.tensor(sz, qt.create(N) + qt.destroy(N))  # Jaynes-Cummings type

# Measurement on subsystems
rho_composite = composite_state * composite_state.dag()
rho_qubit = qt.ptrace(rho_composite, 0)  # Qubit reduced state
rho_ho = qt.ptrace(rho_composite, 1)     # Oscillator reduced state

print(f"Qubit state: {rho_qubit}")
print(f"Oscillator coherence: {qt.expect(qt.destroy(N), rho_ho):.3f}")

# All tensor operations return Qobj instances with appropriate 
# tensor structure and dimensions automatically determined