import qutip
import numpy as np
import scipy

for mod in [qutip, np, scipy]:
    print(mod.__name__, mod.__version__)
!jupyter --version

N = 5
jx = qutip.jmat(N,'x')

H0 = jx * jx

H0_expi = (-(1j) * (H0)).expm()

_, estates_H0 = H0.eigenstates()
_, estates_expi = H0_expi.eigenstates()


identity = [[np.abs(state.overlap(U0state)) for state in estates_H0] for U0state in estates_expi]

identity = qutip.Qobj(identity)
display(identity.tidyup())