~jrjohansson/qutip/master

« back to all changes in this revision

Viewing changes to qutip/ptrace.py

• browse files at revision 1
• compare with another revision
• download diff
• download tarball
• view history from revision 1

Show diffs side-by-side

added

removed

qutip/ptrace.py

 

1

#This file is part of QuTIP.

 

2

#

 

3

#    QuTIP is free software: you can redistribute it and/or modify

 

4

#    it under the terms of the GNU General Public License as published by

 

5

#    the Free Software Foundation, either version 3 of the License, or

 

6

#   (at your option) any later version.

 

7

#

 

8

#    QuTIP is distributed in the hope that it will be useful,

 

9

#    but WITHOUT ANY WARRANTY; without even the implied warranty of

 

10

#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the

 

11

#    GNU General Public License for more details.

 

12

#

 

13

#    You should have received a copy of the GNU General Public License

 

14

#    along with QuTIP.  If not, see <http://www.gnu.org/licenses/>.

 

15

#

 

16

# Copyright (C) 2011, Paul D. Nation & Robert J. Johansson

 

17

#

 

18

###########################################################################

 

19

from scipy import *

 

20

import scipy.sparse as sp

 

21

from scipy.linalg import *

 

22

from qobj import *

 

23

from list2ind import *

 

24

from selct import *

 

25

from qobj import dag

 

26

 

 

27

import numpy

 

28

 

 

29

def ptrace(rho,sel):

 

30

    if isinstance(sel,int):

 

31

        sel=array([sel])

 

32

    sel=asarray(sel)

 

33

    drho=rho.dims[0]

 

34

    N=prod(drho)

 

35

    M=prod(asarray(drho).take(sel))

 

36

 

 

37

    #

 

38

    # wave function:

 

39

    # 

 

40

    if prod(rho.dims[1]) == 1:

 

41

        #print "ptrace for ket"

 

42

        # XXX: temporary solution

 

43

        rho = rho * dag(rho)

 

44

 

 

45

    #

 

46

    # density matrix (operator)

 

47

    #

 

48

    #print "ptrace for operator"

 

49

 

 

50

    perm = sp.lil_matrix(zeros((M*M,N*N)))

 

51

    rest=setdiff1d(arange(len(drho)),asarray(sel)) #all elements in range(len(drho)) not in sel set

 

52

    ilistsel=selct(sel,drho)

 

53

    indsel=list2ind(ilistsel,drho)

 

54

    ilistrest=selct(rest,drho)

 

55

    indrest=list2ind(ilistrest,drho)

 

56

    irest=(indrest-1)*N+indrest

 

57

    m=0

 

58

    for k in indsel:

 

59

        temp=(k-1)*N

 

60

        for l in indsel:

 

61

            m=m+1

 

62

            col=irest+temp+l-1

 

63

            for i in arange(len(col)):

 

64

                perm[m-1,int(col[i][0])-1]=1

 

65

    perm.tocsr()

 

66

    rws=prod(shape(rho.data))

 

67

    

 

68

    #print

 

69

    #print "perm =", perm.shape

 

70

    #print "rho  =", rho.data.shape

 

71

    #print "rws  =", rws

 

72

 

 

73

    rho1=qobj()

 

74

 

 

75

    if sp.issparse(rho.data):

 

76

        rhdata=dot(perm,rho.data.tolil().reshape((rws,1)))

 

77

        rhdata=rhdata.tolil().reshape((M,M))

 

78

        rho1.data=rhdata.tocsr()

 

79

    else:

 

80

        rhdata = perm * rho.data.reshape((rws,1))

 

81

        rhdata=rhdata.reshape((M,M))

 

82

        rho1.data=rhdata

 

83

 

 

84

    dims_kept0=asarray(rho.dims[0]).take(sel)

 

85

    dims_kept1=asarray(rho.dims[0]).take(sel)

 

86

    rho1.dims=[dims_kept0.tolist(),dims_kept1.tolist()]

 

87

    rho1.shape=[prod(dims_kept0),prod(dims_kept1)]

 

88

    rho1.size=[1,1]

 

89

    return rho1

 

90

    

 

91

    

 

92

    

• Older »

Loggerhead is a web-based interface for Breezy
Version: 2.0.1