<commit_before>#This file is part of QuTIP.
#
#    QuTIP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QuTIP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QuTIP.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2011-2012, <NAME> & <NAME>
#
###########################################################################
from scipy import *
import scipy.sparse as sp
from scipy.linalg import *
from qutip.Qobj import *


def ptrace(rho,sel):
    """
    Compute partial trace of composite quantum object formed by :func:`qutip.tensor`
    
    Args:
    
        rho (Qobj): Input composite quantum object.
    
        sel (int or list/array): index or indices for components to keep.
    
    Returns: 
    
        The density matrix of components from sel as a Qobj. 
    """
    if isinstance(sel,int):
        sel=array([sel])
    sel=asarray(sel)
    drho=rho.dims[0]
    N=prod(drho)
    M=prod(asarray(drho).take(sel))
    if prod(rho.dims[1]) == 1:
        rho = rho * rho.dag()
    perm = sp.lil_matrix((M*M,N*N))
    rest=setdiff1d(arange(len(drho)),asarray(sel)) #all elements in range(len(drho)) not in sel set
    ilistsel=selct(sel,drho)
    indsel=list2ind(ilistsel,drho)
    ilistrest=selct(rest,drho)
    indrest=list2ind(ilistrest,drho)
    irest=(indrest-1)*N+indrest
    m=0
    for k in indsel:
        temp=(k-1)*N
        for l in indsel:
            m=m+1
            col=irest+temp+l-1
            for i in xrange(len(col)):
                perm[m-1,int(col[i][0])-1]=1
    perm.tocsr()
    rws=prod(shape(rho.data))
    rho1=Qobj()

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

    dims_kept0=asarray(rho.dims[0]).take(sel)
    dims_kept1=asarray(rho.dims[0]).take(sel)
    rho1.dims=[dims_kept0.tolist(),dims_kept1.tolist()]
    rho1.shape=[prod(dims_kept0),prod(dims_kept1)]
    rho1.size=[1,1]
    return Qobj(rho1)



def list2ind(ilist,dims):
	"""!
	Private function returning indicies
	"""
	ilist=asarray(ilist)
	dims=asarray(dims)
	irev=fliplr(ilist)-1
	fact=append(array([1]),(cumprod(flipud(dims)[:-1])))
	fact=fact.reshape(len(fact),1)
	return sort(dot(irev,fact)+1,0)

def selct(sel,dims):
	"""
	Private function finding selected components
	"""
	sel=asarray(sel)#make sure sel is array
	dims=asarray(dims)#make sure dims is array
	rlst=dims.take(sel)
	rprod=prod(rlst)
	ilist=ones((rprod,len(dims)));
	counter=arange(rprod)
	for k in xrange(len(sel)):
		ilist[:,sel[k]]=remainder(fix(counter/prod(dims[sel[k+1:]])),dims[sel[k]])+1
	return ilist
<commit_msg>added csr_to_col() function for reshaping csr matrices without converting to lil format (it is too slow for reshape).  Now reshape is much much faster (~10x).<commit_after>#This file is part of QuTIP.
#
#    QuTIP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QuTIP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QuTIP.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2011-2012, <NAME> & <NAME>
#
###########################################################################
from scipy import *
import scipy.sparse as sp
from scipy.linalg import *
from qutip.Qobj import *


def ptrace(rho,sel):
    """
    Compute partial trace of composite quantum object formed by :func:`qutip.tensor`
    
    Args:
    
        rho (Qobj): Input composite quantum object.
    
        sel (int or list/array): index or indices for components to keep.
    
    Returns: 
    
        The density matrix of components from sel as a Qobj. 
    """
    if isinstance(sel,int):
        sel=array([sel])
    sel=asarray(sel)
    drho=rho.dims[0]
    N=prod(drho)
    M=prod(asarray(drho).take(sel))
    if prod(rho.dims[1]) == 1:
        rho = rho * rho.dag()
    perm = sp.lil_matrix((M*M,N*N))
    rest=setdiff1d(arange(len(drho)),asarray(sel)) #all elements in range(len(drho)) not in sel set
    ilistsel=selct(sel,drho)
    indsel=list2ind(ilistsel,drho)
    ilistrest=selct(rest,drho)
    indrest=list2ind(ilistrest,drho)
    irest=(indrest-1)*N+indrest
    m=0
    for k in indsel:
        temp=(k-1)*N
        for l in indsel:
            m=m+1
            col=irest+temp+l-1
            for i in xrange(len(col)):
                perm[m-1,int(col[i][0])-1]=1
    perm.tocsr()
    rws=prod(shape(rho.data))
    rho1=Qobj()
    rhdata=perm.dot(csr_to_col(rho.data))
    rhdata=rhdata.tolil().reshape((M,M))
    rho1.data=rhdata.tocsr()

    dims_kept0=asarray(rho.dims[0]).take(sel)
    dims_kept1=asarray(rho.dims[0]).take(sel)
    rho1.dims=[dims_kept0.tolist(),dims_kept1.tolist()]
    rho1.shape=[prod(dims_kept0),prod(dims_kept1)]
    rho1.size=[1,1]
    return Qobj(rho1)



def list2ind(ilist,dims):
	"""!
	Private function returning indicies
	"""
	ilist=asarray(ilist)
	dims=asarray(dims)
	irev=fliplr(ilist)-1
	fact=append(array([1]),(cumprod(flipud(dims)[:-1])))
	fact=fact.reshape(len(fact),1)
	return sort(dot(irev,fact)+1,0)

def selct(sel,dims):
	"""
	Private function finding selected components
	"""
	sel=asarray(sel)#make sure sel is array
	dims=asarray(dims)#make sure dims is array
	rlst=dims.take(sel)
	rprod=prod(rlst)
	ilist=ones((rprod,len(dims)));
	counter=arange(rprod)
	for k in xrange(len(sel)):
		ilist[:,sel[k]]=remainder(fix(counter/prod(dims[sel[k+1:]])),dims[sel[k]])+1
	return ilist



def csr_to_col(mat):
    mat.sort_indices()
    rows=array([len(range(mat.indptr[i],mat.indptr[i+1])) for i in xrange(mat.shape[1])])
    rows=[[k for j in xrange(rows[k])] for k in xrange(len(rows))] 
    rows=array([item for sublist in rows for item in sublist])
    inds=mat.shape[1]*rows+mat.indices
    ptrs=array([0,len(mat.data)])
    out=sp.csr_matrix((mat.data,inds,ptrs),shape=(1, prod(mat.shape)),dtype=complex)
    return out.transpose()

{'fraction_non_alphanumeric': 0.15714285714285714, 'fraction_numerical': 0.015773809523809523, 'mean_word_length': 3.6447823082239115, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3167340', 'n_tokens_mistral': 2508, 'n_tokens_neox': 2483, 'n_words': 578}