jiri/LA_library
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
3d8386e30a
LA_library/tensor.cc

729 lines
18 KiB
C++
Raw Blame History

/*
LA: linear algebra C++ interface library
Copyright (C) 2024 Jiri Pittner <jiri.pittner@jh-inst.cas.cz> or <jiri@pittnerovi.com>
This program 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.
This program 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 this program. If not, see <http://www.gnu.org/licenses/>.
*/
#include <iostream>
#include "tensor.h"
#include "laerror.h"
#include "qsort.h"
#include "miscfunc.h"
#include <complex>
namespace LA {
template<typename T>
int Tensor<T>:: calcrank()
{
int r=0;
for(int i=0; i<shape.size(); ++i)
{
const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&shape))[i];
if(sh->number==0) laerror("empty index group");
r+=sh->number;
}
myrank=r;
return r;
}
template<typename T>
LA_largeindex Tensor<T>::calcsize()
{
groupsizes.resize(shape.size());
cumsizes.resize(shape.size());
LA_largeindex s=1;
for(int i=0; i<shape.size(); ++i)
{
const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&shape))[i];
if(sh->number==0) laerror("empty index group");
if(sh->range==0) return 0;
cumsizes[i]=s;
switch(sh->symmetry)
{
case 0:
s *= groupsizes[i] = longpow(sh->range,sh->number);
break;
case 1:
s *= groupsizes[i] = simplicial(sh->number,sh->range);
break;
case -1:
s *= groupsizes[i] = simplicial(sh->number,sh->range-sh->number+1);
break;
default:
laerror("illegal index group symmetry");
break;
}
}
return s;
}
LA_largeindex subindex(int *sign, const INDEXGROUP &g, const NRVec<LA_index> &I) //index of one subgroup
{
#ifdef DEBUG
if(I.size()<=0) laerror("empty index group in subindex");
if(g.number!=I.size()) laerror("mismatch in the number of indices in a group");
for(int i=0; i<I.size(); ++i) if(I[i]<g.offset || I[i] >= g.offset+g.range) laerror("index out of range in tensor subindex");
#endif
switch(I.size()) //a few special cases for efficiency
{
case 0:
*sign=0;
return 0;
break;
case 1:
*sign=1;
return I[0]-g.offset;
break;
case 2:
{
*sign=1;
if(g.symmetry==0) return (I[1]-g.offset)*g.range+I[0]-g.offset;
LA_index i0,i1;
if(I[0]>I[1]) {i1=I[0]; i0=I[1]; if(g.symmetry<0) *sign = -1;} else {i1=I[1]; i0=I[0];}
i0 -= g.offset;
i1 -= g.offset;
if(g.symmetry<0)
{
if(i0==i1) {*sign=0; return -1;}
return i1*(i1-1)/2+i0;
}
else
{
return i1*(i1+1)/2+i0;
}
}
break;
default: //general case
{
*sign=1;
if(g.symmetry==0) //rectangular case
{
LA_largeindex r=0;
for(int i=I.size()-1; i>=0; --i)
{
r*= g.range;
r+= I[i]-g.offset;
}
return r;
}
}
//compressed storage case
NRVec<LA_index> II(I);
II.copyonwrite();
if(g.offset!=0) II -= g.offset;
int parity=netsort(II.size(),&II[0]);
if(g.symmetry<0 && (parity&1)) *sign= -1;
if(g.symmetry<0) //antisymmetric
{
for(int i=0; i<I.size()-1; ++i)
if(II[i]==II[i+1])
{*sign=0; return -1;} //identical indices of antisymmetric tensor
LA_largeindex r=0;
for(int i=0; i<II.size(); ++i) r += simplicial(i+1,II[i]-i);
return r;
}
else //symmetric
{
LA_largeindex r=0;
for(int i=0; i<II.size(); ++i) r += simplicial(i+1,II[i]);
return r;
}
break;
}
laerror("this error should not happen");
return -1;
}
//inverse map of group superindex to canonically ordered index list
NRVec<LA_index> inverse_subindex(const INDEXGROUP &g, LA_largeindex s)
{
NRVec<LA_index> I(g.number);
if(g.number==1) {I[0]=s+g.offset; return I;}
switch(g.symmetry)
{
case 0:
for(int i=0; i<g.number; ++i)
{
I[i] = s%g.range;
s /= g.range;
}
break;
case 1:
for(int i=g.number; i>0; --i)
{
I[i-1] = inverse_simplicial(i,s);
s -= simplicial(i,I[i-1]);
}
break;
case -1:
for(int i=g.number-1; i>=0; --i)
{
I[i] = i + inverse_simplicial(i+1,s);
s -= simplicial(i+1,I[i]-i);
}
break;
default: laerror("illegal index symmetry");
}
if(g.offset!=0) I += g.offset;
return I;
}
template<typename T>
SUPERINDEX Tensor<T>::inverse_index(LA_largeindex s) const
{
SUPERINDEX I(shape.size());
for(int g=shape.size()-1; g>=0; --g)
{
LA_largeindex groupindex;
if(g>0)
{
groupindex = s/cumsizes[g];
s %= cumsizes[g];
}
else groupindex=s;
I[g] = inverse_subindex(shape[g],groupindex);
}
return I;
}
template<typename T>
LA_largeindex Tensor<T>::index(int *sign, const SUPERINDEX &I) const
{
//check index structure and ranges
#ifdef DEBUG
if(I.size()!=shape.size()) laerror("mismatch in the number of tensor index groups");
for(int i=0; i<I.size(); ++i)
{
if(shape[i].number!=I[i].size()) {std::cerr<<"error in index group no. "<<i<<std::endl; laerror("mismatch in the size of tensor index group");}
for(int j=0; j<shape[i].number; ++j)
{
if(I[i][j] <shape[i].offset || I[i][j] >= shape[i].offset+shape[i].range)
{
std::cerr<<"error in index group no. "<<i<<" index no. "<<j<<std::endl;
laerror("tensor index out of range");
}
}
}
#endif
LA_largeindex r=0;
*sign=1;
for(int g=0; g<shape.size(); ++g) //loop over index groups
{
int gsign;
LA_largeindex groupindex = subindex(&gsign,shape[g],I[g]);
//std::cout <<"INDEX TEST group "<<g<<" cumsizes "<< cumsizes[g]<<" groupindex "<<groupindex<<std::endl;
*sign *= gsign;
if(groupindex == -1) return -1;
r += groupindex * cumsizes[g];
}
return r;
}
template<typename T>
LA_largeindex Tensor<T>::index(int *sign, const FLATINDEX &I) const
{
#ifdef DEBUG
if(rank()!=I.size()) laerror("tensor rank mismatch in index");
#endif
LA_largeindex r=0;
*sign=1;
int gstart=0;
for(int g=0; g<shape.size(); ++g) //loop over index groups
{
int gsign;
int gend= gstart+shape[g].number-1;
NRVec<LA_index> subI = I.subvector(gstart,gend);
gstart=gend+1;
LA_largeindex groupindex = subindex(&gsign,shape[g],subI);
//std::cout <<"FLATINDEX TEST group "<<g<<" cumsizes "<< cumsizes[g]<<" groupindex "<<groupindex<<std::endl;
*sign *= gsign;
if(groupindex == -1) return -1;
r += groupindex * cumsizes[g];
}
return r;
}
template<typename T>
LA_largeindex Tensor<T>::vindex(int *sign, LA_index i1, va_list args) const
{
NRVec<LA_index> I(rank());
I[0]=i1;
for(int i=1; i<rank(); ++i)
{
I[i] = va_arg(args,LA_index);
}
va_end(args);
return index(sign,I);
}
//binary I/O
template<typename T>
void Tensor<T>::put(int fd) const
{
shape.put(fd,true);
groupsizes.put(fd,true);
cumsizes.put(fd,true);
data.put(fd,true);
}
template<typename T>
void Tensor<T>::get(int fd)
{
shape.get(fd,true);
myrank=calcrank(); //is not stored but recomputed
groupsizes.put(fd,true);
cumsizes.get(fd,true);
data.get(fd,true);
}
template<typename T>
Tensor<T>::Tensor(const NRVec<T> &x)
: data(x)
{
myrank=1;
shape.resize(1);
shape[0].number=1;
shape[0].symmetry=0;
#ifndef LA_TENSOR_ZERO_OFFSET
shape[0].offset=0;
#endif
shape[0].range=x.size();
calcsize();
}
template<typename T>
Tensor<T>::Tensor(const NRMat<T> &x)
: data(&x(0,0),x.nrows()*x.ncols())
{
myrank=2;
if(x.nrows()==x.ncols())
{
shape.resize(1);
shape[0].number=2;
shape[0].symmetry=0;
#ifndef LA_TENSOR_ZERO_OFFSET
shape[0].offset=0;
#endif
shape[0].range=x.nrows();
}
else
{
shape.resize(2);
shape[0].number=1; shape[1].number=1;
shape[0].symmetry=0; shape[1].symmetry=0;
#ifndef LA_TENSOR_ZERO_OFFSET
shape[0].offset=0; shape[1].offset=0;
#endif
shape[0].range=x.ncols();
shape[1].range=x.nrows();
}
calcsize();
}
template<typename T>
Tensor<T>::Tensor(const NRSMat<T> &x)
: data(NRVec<T>(x))
{
myrank=2;
shape.resize(1);
shape[0].number=2;
shape[0].symmetry=1;
#ifndef LA_TENSOR_ZERO_OFFSET
shape[0].offset=0;
#endif
shape[0].range=x.nrows();
calcsize();
}
template<typename T>
void loopingroups(Tensor<T> &t, int ngroup, int igroup, T **p, SUPERINDEX &I, void (*callback)(const SUPERINDEX &, T *))
{
LA_index istart,iend;
const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&t.shape))[ngroup];
switch(sh->symmetry)
{
case 0:
istart= sh->offset;
iend= sh->offset+sh->range-1;
break;
case 1:
istart= sh->offset;
if(igroup==sh->number-1) iend= sh->offset+sh->range-1;
else iend = I[ngroup][igroup+1];
break;
case -1:
istart= sh->offset + igroup;
if(igroup==sh->number-1) iend= sh->offset+sh->range-1;
else iend = I[ngroup][igroup+1]-1;
break;
}
for(LA_index i = istart; i<=iend; ++i)
{
I[ngroup][igroup]=i;
if(ngroup==0 && igroup==0)
{
int sign;
//std::cout <<"TEST "<<t.index(&sign,I)<<" ";
(*callback)(I,(*p)++);
}
else
{
int newigroup= igroup-1;
int newngroup=ngroup;
if(newigroup<0)
{
--newngroup;
const indexgroup *sh2 = &(* const_cast<const NRVec<indexgroup> *>(&t.shape))[newngroup];
newigroup=sh2->number-1;
}
loopingroups(t,newngroup,newigroup,p,I,callback);
}
}
}
template<typename T>
void Tensor<T>::loopover(void (*callback)(const SUPERINDEX &, T *))
{
SUPERINDEX I(shape.size());
for(int i=0; i<I.size(); ++i)
{
const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&shape))[i];
I[i].resize(sh->number);
I[i] = sh->offset;
}
T *pp=&data[0];
int ss=shape.size()-1;
const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&shape))[ss];
loopingroups(*this,ss,sh->number-1,&pp,I,callback);
}
static std::ostream *sout;
template<typename T>
static void outputcallback(const SUPERINDEX &I, T *v)
{
//print indices flat
for(int i=0; i<I.size(); ++i)
for(int j=0; j<I[i].size(); ++j) *sout << I[i][j]<<" ";
*sout<<" "<< *v<<std::endl;
}
std::ostream & operator<<(std::ostream &s, const INDEXGROUP &x)
{
s<<x.number <<" "<<x.symmetry<<" ";
#ifndef LA_TENSOR_ZERO_OFFSET
s<<x.offset<<" ";
#endif
s<< x.range<<std::endl;
return s;
}
std::istream & operator>>(std::istream &s, INDEXGROUP &x)
{
s>>x.number>>x.symmetry;
#ifndef LA_TENSOR_ZERO_OFFSET
s>>x.offset;
#endif
s>>x.range;
return s;
}
template<typename T>
std::ostream & operator<<(std::ostream &s, const Tensor<T> &x)
{
s<<x.shape;
sout= &s;
const_cast<Tensor<T> *>(&x)->loopover(&outputcallback<T>);
return s;
}
template <typename T>
std::istream & operator>>(std::istream &s, Tensor<T> &x)
{
s>>x.shape;
x.data.resize(x.calcsize()); x.calcrank();
FLATINDEX I(x.rank());
for(LA_largeindex i=0; i<x.data.size(); ++i)
{
for(int j=0; j<I.size(); ++j) s>>I[j];
T val; s>>val;
x.lhs(I) = val;
}
return s;
}
template<typename T>
void loopovergroups(Tensor<T> &t, int ngroup, T **p, GROUPINDEX &I, void (*callback)(const GROUPINDEX &, T *))
{
for(LA_largeindex i = 0; i<t.groupsizes[ngroup]; ++i)
{
I[ngroup]=i;
if(ngroup==0) (*callback)(I,(*p)++);
else loopovergroups(t,ngroup-1,p,I,callback);
}
}
template<typename T>
void Tensor<T>::grouploopover(void (*callback)(const GROUPINDEX &, T *))
{
GROUPINDEX I(shape.size());
T *pp=&data[0];
loopovergroups(*this,shape.size()-1,&pp,I,callback);
}
const NRPerm<int> *help_p;
template<typename T>
Tensor<T> *help_t;
template<typename T>
const Tensor<T> *help_tt;
template<typename T>
static void permutecallback(const GROUPINDEX &I, T *v)
{
LA_largeindex target=0;
for(int i=0; i< help_t<T>->shape.size(); ++i)
{
target += I[(*help_p)[i+1]-1] * help_t<T>->cumsizes[i];
}
help_t<T>->data[target] = *v;
}
template<typename T>
Tensor<T> Tensor<T>::permute_index_groups(const NRPerm<int> &p) const
{
NRVec<indexgroup> newshape=shape.permuted(p);
Tensor<T> r(newshape);
//prepare statics for the callback
help_p = &p;
help_t<T> = &r;
//now rearrange the data
const_cast<Tensor<T> *>(this)->grouploopover(permutecallback<T>);
return r;
}
FLATINDEX superindex2flat(const SUPERINDEX &I)
{
int rank=0;
for(int i=0; i<I.size(); ++i) rank += I[i].size();
FLATINDEX J(rank);
int ii=0;
for(int i=0; i<I.size(); ++i)
{
for(int j=0; j<I[i].size(); ++j) J[ii++] = I[i][j];
}
return J;
}
template<typename T>
static void unwind_callback(const SUPERINDEX &I, T *v)
{
FLATINDEX J = superindex2flat(I);
FLATINDEX JP = J.permuted(*help_p,true);
//std::cout <<"TEST unwind_callback: from "<<JP<<" TO "<<J<<std::endl;
*v = (*help_tt<T>)(JP); //rhs operator() generates the redundant elements for the unwinded lhs tensor
}
template<typename T>
Tensor<T> Tensor<T>::unwind_index(int group, int index) const
{
if(group<0||group>=shape.size()) laerror("wrong group number in unwind_index");
if(index<0||index>=shape[group].number) laerror("wrong index number in unwind_index");
if(shape[group].number==1) //single index in the group
{
if(group==0) return *this; //is already the least significant group
NRPerm<int> p(shape.size());
p[1]= 1+group;
int ii=1;
if(ii==1+group) ii++; //skip this
for(int i=2; i<=shape.size(); ++i)
{
p[i]=ii++;
if(ii==1+group) ii++; //skip this
}
if(!p.is_valid()) laerror("internal error in unwind_index");
return permute_index_groups(p);
}
//general case - recalculate the shape and allocate the new tensor
NRVec<indexgroup> newshape(shape.size()+1);
newshape[0].number=1;
newshape[0].symmetry=0;
newshape[0].range=shape[group].range;
#ifndef LA_TENSOR_ZERO_OFFSET
newshape[0].offset = shape[group].offset;
#endif
int flatindex=0; //(group,index) in flat form
for(int i=0; i<shape.size(); ++i)
{
newshape[i+1] = shape[i];
if(i==group)
{
--newshape[i+1].number;
flatindex += index;
}
else flatindex += shape[i].number;
}
Tensor<T> r(newshape);
if(r.rank()!=rank()) laerror("internal error 2 in unwind_index");
//compute the corresponding permutation of FLATINDEX for use in the callback
NRPerm<int> indexperm(rank());
indexperm[1]=flatindex+1;
int ii=1;
if(ii==flatindex+1) ii++;
for(int i=2; i<=rank(); ++i)
{
indexperm[i] = ii++;
if(ii==flatindex+1) ii++; //skip this
}
if(!indexperm.is_valid())
{
std::cout << "indexperm = "<<indexperm<<std::endl;
laerror("internal error 3 in unwind_index");
}
//loop recursively and do the unwinding
help_tt<T> = this;
help_p = &indexperm;
r.loopover(unwind_callback);
return r;
}
template<typename T>
static void auxmatmult(int nn, int mm, int kk, T *r, T *a, T *b, T alpha=1, T beta=0) //R(nn,mm) = A * B^T
{
for(int i=0; i<nn; ++i) for(int j=0; j<mm; ++j)
{
if(beta==0) r[i*mm+j]=0; else r[i*mm+j] *= beta;
for(int k=0; k<kk; ++k) r[i*mm+j] += alpha * a[i*kk+k] * b[j*kk+k];
}
}
template<>
void auxmatmult<double>(int nn, int mm, int kk, double *r, double *a, double *b, double alpha, double beta)
{
cblas_dgemm(CblasRowMajor, CblasNoTrans, CblasTrans, nn, mm, kk, alpha, a, kk, b, kk, beta, r, mm);
}
template<>
void auxmatmult<std::complex<double> >(int nn, int mm, int kk, std::complex<double> *r, std::complex<double> *a, std::complex<double> *b, std::complex<double> alpha, std::complex<double> beta)
{
cblas_zgemm(CblasRowMajor, CblasNoTrans, CblasTrans, nn, mm, kk, &alpha, a, kk, b, kk, &beta, r, mm);
}
//Conntraction could be implemented without the temporary storage for unwinding, but then we would need
//double recursion over indices of both tensors. Hopefully using the matrix multiplication here
//makes it also more efficient, even for (anti)symmetric indices
//The index unwinding is unfortunately a big burden, and in principle could be eliminated in case of non-symmetric indices
//
template<typename T>
void Tensor<T>::addcontraction(const Tensor &rhs1, int group, int index, const Tensor &rhs, int rhsgroup, int rhsindex, T alpha, T beta, bool doresize)
{
if(group<0||group>=rhs1.shape.size()) laerror("wrong group number in contraction");
if(rhsgroup<0||rhsgroup>=rhs.shape.size()) laerror("wrong rhsgroup number in contraction");
if(index<0||index>=rhs1.shape[group].number) laerror("wrong index number in conntraction");
if(rhsindex<0||rhsindex>=rhs.shape[rhsgroup].number) laerror("wrong index number in conntraction");
if(rhs1.shape[group].offset != rhs.shape[rhsgroup].offset) laerror("incompatible index offset in contraction");
if(rhs1.shape[group].range != rhs.shape[rhsgroup].range) laerror("incompatible index range in contraction");
Tensor<T> u = rhs1.unwind_index(group,index);
Tensor<T> rhsu = rhs.unwind_index(rhsgroup,rhsindex);
NRVec<indexgroup> newshape(u.shape.size()+rhsu.shape.size()-2);
int ii=0;
for(int i=1; i<rhsu.shape.size(); ++i) newshape[ii++] = rhsu.shape[i];
for(int i=1; i<u.shape.size(); ++i) newshape[ii++] = u.shape[i]; //this tensor will have more significant indices than the rhs one
if(doresize)
{
if(beta!= (T)0) laerror("resize in addcontraction requires beta=0");
resize(newshape);
}
else
{
if(shape!=newshape) laerror("tensor shape mismatch in addcontraction");
}
int nn,mm,kk;
kk=u.groupsizes[0];
if(kk!=rhsu.groupsizes[0]) laerror("internal error in contraction");
nn=1; for(int i=1; i<u.shape.size(); ++i) nn*= u.groupsizes[i];
mm=1; for(int i=1; i<rhsu.shape.size(); ++i) mm*= rhsu.groupsizes[i];
auxmatmult<T>(nn,mm,kk,&data[0],&u.data[0], &rhsu.data[0],alpha,beta);
}
template class Tensor<double>;
template class Tensor<std::complex<double> >;
template std::ostream & operator<<(std::ostream &s, const Tensor<double> &x);
template std::ostream & operator<<(std::ostream &s, const Tensor<std::complex<double> > &x);
template std::istream & operator>>(std::istream &s, Tensor<double> &x);
template std::istream & operator>>(std::istream &s, Tensor<std::complex<double> > &x);
}//namespace
Reference in New Issue View Git Blame Copy Permalink