from qiskit import QuantumCircuit
import qiskit.quantum_info as qi

qc_AB = QuantumCircuit(1)
qc_AB.h(0)
qc_AB.s(0)
qc_AB.z(0)
qc_AB.draw()

psi_AB = qi.Statevector.from_instruction(qc_AB)
psi_AB.draw('latex', prefix='|\\psi_{AB}\\rangle = ')

rho_AB = qi.DensityMatrix.from_instruction(qc_AB)
rho_AB.draw('latex', prefix='\\rho_{AB} = ')

from qiskit.visualization import plot_state_city
plot_state_city(rho_AB.data, title='Density Matrix')

qc_AB = QuantumCircuit(2)
qc_AB.h(0)
qc_AB.h(1)
qc_AB.draw()

psi_AB = qi.Statevector.from_instruction(qc_AB)
psi_AB.draw('latex', prefix='|\\psi_{AB}\\rangle = ')

rho_AB = qi.DensityMatrix.from_instruction(qc_AB)
rho_AB.draw('latex', prefix='\\rho_{AB} = ')

from qiskit.visualization import plot_state_city
plot_state_city(rho_AB.data, title='Density Matrix')

qc_AB = QuantumCircuit(2)
qc_AB.h(0)
qc_AB.h(1)
qc_AB.draw()

psi_AB = qi.Statevector.from_instruction(qc_AB)
psi_AB.draw('latex', prefix='|\\psi_{AB}\\rangle = ')

rho_AB = qi.DensityMatrix.from_instruction(qc_AB)
rho_AB.draw('latex', prefix='\\rho_{AB} = ')

rho_B = qi.partial_trace(rho_AB,[0])
rho_A = qi.partial_trace(rho_AB,[1])

display(rho_B.draw('latex', prefix=" \\rho_{B} = "),
        rho_A.draw('latex', prefix=" \\rho_{A} = "))

