<jupyter_start><jupyter_text>Excitation-number-restricted states: Jaynes-Cummings chain Authors: <NAME> (<EMAIL>) and <NAME> (<EMAIL>)  IntroductionThe ENR functions construct a basis set for multipartite systems which contains only states that have an overall number of excitations.  This is particularly useful for systems where the model conserves excitation number, as in the JC-chain example below. However it can also help reduce memory cost even if this is not the case;  however then one must be careful in choosing a large enough number of excitations to obtain convergence.For example, consider a system consisting of 4 modes, eachwith 5 states. The total hilbert space size is $5^4 = 625$. If we areonly interested in states that contain up to 2 excitations, we only needto include states such as    (0, 0, 0, 0)    (0, 0, 0, 1)    (0, 0, 0, 2)    (0, 0, 1, 0)    (0, 0, 1, 1)    (0, 0, 2, 0)    ...The enr fucntions create operators and states for the 4 modes that actwithin this state space. For example,a1, a2, a3, a4 = enr_destroy([5, 5, 5, 5], excitations=2)creates destruction operators for each mode.From this point onwards, the annihiltion operators a1, ..., a4 can beused to setup a Hamiltonian, collapse operators and expectation-valueoperators, etc., following the usual pattern.However, many functions in QuTiP will fail on states/operators constructed with this method.  So far functions which exist are:enr_fock:  create a fock state in the restricted basisenr_identity:  create an identity operatorenr_thermal_dm:  create a thermal stateenr_destroy:  create an annilation operatorenr_ptrace: perform a partial trace<jupyter_code>%matplotlib inline
import numpy as np
from qutip import *
from qutip.ipynbtools import HTMLProgressBar
def solve(d, psi0):
    
    # list of annihilation operators for cavity modes
    a = d[::2]

    # list of atomic annihilation operators
    sm = d[1::2]

    H0 = sum([aa.dag() * aa for aa in a]) + sum([s.dag() * s for s in sm])

    # atom-cavity couplings
    Hint_ac = 0
    for n in range(N):
        Hint_ac += 0.5 * (a[n].dag() * sm[n] + sm[n].dag() * a[n])

    # coupling beetween cavities
    Hint_cc = 0
    for n in range(N-1):
        Hint_cc += 0.9 * (a[n].dag() * a[n+1] + a[n+1].dag() * a[n])

    H = H0 + Hint_ac + Hint_cc

    e_ops = [x.dag() * x for x in d]
    c_ops = [0.01 * x for x in a]
    c_ops = [0.05 * x for x in sm]

    times = np.linspace(0, 250, 1000)

    L = liouvillian_ref(H, c_ops)
    
    result = mesolve(H, psi0, times, c_ops, e_ops, options=Options(nsteps=5000,store_states=True))

    return result, H, L
N = 4                        # number of systems
M = 2                        # number of cavity states

#The dimensions of the JC spin chain
dims = [M, 2] * N
#The number of excitations in the chain
excitations = 1              # total number of excitations
initial_excitiations = 1     # initial number of excitations

psi0 = tensor([basis(d, initial_excitiations) if idx == 1 else basis(d, 0) for idx, d in enumerate(dims)])



d = [tensor([destroy(d1) if idx1 == idx2 else identity(d1)
             for idx1, d1 in enumerate(dims)])
     for idx2, d2 in enumerate(dims)]



#This is the normal QuTip solution, without using ENR states
%time result1, H1, L1 = solve(d, psi0)

#Construct ENR operators and states
d = enr_destroy(dims, excitations)
psi0 = enr_fock(dims, excitations, [initial_excitiations if m == 1 else 0 for m in range(2*N)])
[initial_excitiations if m == 1 else 0 for m in range(2*N)]
#This is the solution using the ENR states/operators.

%time result2, H2, L2 = solve(d, psi0)
fig, axes = plot_expectation_values([result1, result2], figsize=(10, 8))
for ax in axes[:, 0]:
    ax.set_ylim(-0.1, 1.1)
fig.tight_layout();

def ENR_ptrace(rho, sel, excitations):
    """
    Partial trace for ENR states.   
    Parameters
    ----------
    rho : Qobj
        Qobj of an ENR system defined through  enr_fock enr_identity enr_thermal_dm enr_destroy or output from a solver
        
    sel : int/list
            An ``int`` or ``list`` of components to keep after partial trace.  
        
    excitations : integer
        The maximum number of excitations that are to be included in the
        state space.  Should be consistent with the same assumptions used to construct rho.
    
    Returns
    -------
    oper : qobj
        Quantum object representing partial trace with selected components
        remaining.
    
    Notes
    -----
    The default Qobj.ptrace() will fail with ENR systems.  This should be used instead.
    
    """
    if isinstance(sel, int):
        sel = np.array([sel])
    else:
        sel = np.asarray(sel)

    if (sel < 0).any() or (sel >= len(rho.dims[0])).any():
        raise TypeError("Invalid selection index in ptrace.")
    
    drho=rho.dims[0]
    
    ############
    #dimensions
    ####################
    #enr definitions for the original state
    nstates, state2idx, idx2state = enr_state_dictionaries(drho, excitations)
    

    
    ################
    #definition of number of states in selection
    ######################
    #
    
    dims_short= np.asarray(drho).take(sel)
    nstates2, state2idx2, idx2state2 = enr_state_dictionaries(dims_short.tolist(), excitations)

    
    # this is a list of the dimensions of the system one has traced out
    rest = np.setdiff1d(np.arange(len(drho)), sel)
    
       
    #construct matrix to return the new Density matrix
    rhout = np.zeros((nstates2,nstates2),dtype=np.complex64)
    
    for ind,state in idx2state.items():
        for ind2,state2 in idx2state.items():
            #if the parts of the states of the systems(s) being traced out are diagonal, add this to the new DM
            if  np.all(np.asarray(state).take(rest) == np.asarray(state2).take(rest)):

                rhout[state2idx2[tuple(np.asarray(state).take(sel))],
                      state2idx2[tuple(np.asarray(state2).take(sel))]] += rho.data[state2idx[state],
                                                                                    state2idx[state2]]
                
    
    dims_kept0 = np.asarray(drho).take(sel)
    rho1_dims = [dims_kept0.tolist(), dims_kept0.tolist()]
    rho1_shape = [nstates2, nstates2]
    
    return Qobj(rhout,rho1_dims,rho1_shape)
        
result1.states[10].ptrace([0]) 
ENR_ptrace(result2.states[10], [0], excitations)
result1.states[10].ptrace([0,1,4])


ENR_ptrace(result2.states[10],[0,1,4],excitations)<jupyter_output><empty_output>