def build_sr(ll=None,dl=None,ul=None): ''' Build a bsr or csr matrix by lower/diagonal/upper part of a tridiagonal matrix. ll/dl/ul: The lower/diagonal/upper part of tridiagonal matrix. Leave None to make any of them zeros(no all of them). *return*: a csr_matrix instance for scalar version, else a bsr_matrix instance. ''' nzarr=None for i,il in enumerate([ll,ul,dl]): if il is not None: nzarr=il n=len(il)+1 if i!=2 else len(il) break if nzarr is None: raise ValueError('At least one of ll,dl,ul should be nonzeros!') is_scalar=ndim(nzarr)==1 p=1 if is_scalar else nzarr.shape[-1] if is_scalar: mgen=csr_matrix nullval=zeros(0) else: mgen=bsr_matrix nullval=zeros([0,p,p]) indx_d=arange(n) indx=concatenate([[] if dl is None else indx_d,[] if ll is None else indx_d[1:],[] if ul is None else indx_d[:-1]],axis=0) indy=concatenate([[] if dl is None else indx_d,[] if ll is None else indx_d[:-1],[] if ul is None else indx_d[1:]],axis=0) data=concatenate([nullval if dl is None else dl,nullval if ll is None else ll,nullval if ul is None else ul],axis=0) odl=argsort(indx) L=mgen((data[odl],indy[odl],ind2ptr(indx[odl],n)),dtype=nzarr.dtype) return L
def tobsr(self): '''Transform to bsr_matrix.''' n=self.n p=self.p m=ndarray((n,n),dtype='O') indx=concatenate([arange(n-1),arange(n),arange(1,n)]) indy=concatenate([arange(1,n),arange(n),arange(n-1)]) args=argsort(indx) data=concatenate([self.upper,self.diagonal,self.lower]) res=bsr_matrix((data[args],indy[args],ind2ptr(indx[args],n)),blocksize=(p,p)) return res