tensor: permute_index_groups and some necessary static casts

This commit is contained in:
Jiri Pittner 2024-04-24 17:43:11 +02:00
parent 0ff55b66bb
commit da0b3116f6
3 changed files with 121 additions and 25 deletions

30
t.cc
View File

@ -3200,7 +3200,7 @@ cout <<d;
} }
if(1) if(0)
{ {
INDEXGROUP g; INDEXGROUP g;
g.number=3; g.number=3;
@ -3243,4 +3243,32 @@ cout <<epsilon.data;
cout <<epsilon; cout <<epsilon;
} }
if(1)
{
int n=3;
NRVec<INDEXGROUP> s(4);
for(int i=0; i<4; ++i)
{
s[i].number=1;
s[i].symmetry=0;
s[i].offset=0;
s[i].range=n;
}
Tensor<double> t(s);
t.randomize(1.);
cout <<t;
NRPerm<int> p({3,1,4,2});
Tensor<double> tt=t.permute_index_groups(p);
cout <<tt;
for(int i=0; i<n; ++i)
for(int j=0; j<n; ++j)
for(int k=0; k<n; ++k)
for(int l=0; l<n; ++l)
{
if(t(i,j,k,l)!=tt(k,i,l,j)) laerror("error in permute_index_groups");
}
}
} }

101
tensor.cc
View File

@ -32,8 +32,9 @@ int Tensor<T>:: calcrank()
int r=0; int r=0;
for(int i=0; i<shape.size(); ++i) for(int i=0; i<shape.size(); ++i)
{ {
if(shape[i].number==0) laerror("empty index group"); const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&shape))[i];
r+=shape[i].number; if(sh->number==0) laerror("empty index group");
r+=sh->number;
} }
myrank=r; myrank=r;
return r; return r;
@ -49,19 +50,20 @@ cumsizes.resize(shape.size());
LA_largeindex s=1; LA_largeindex s=1;
for(int i=0; i<shape.size(); ++i) for(int i=0; i<shape.size(); ++i)
{ {
if(shape[i].number==0) laerror("empty index group"); const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&shape))[i];
if(shape[i].range==0) return 0; if(sh->number==0) laerror("empty index group");
if(sh->range==0) return 0;
cumsizes[i]=s; cumsizes[i]=s;
switch(shape[i].symmetry) switch(sh->symmetry)
{ {
case 0: case 0:
s *= groupsizes[i] = longpow(shape[i].range,shape[i].number); s *= groupsizes[i] = longpow(sh->range,sh->number);
break; break;
case 1: case 1:
s *= groupsizes[i] = simplicial(shape[i].number,shape[i].range); s *= groupsizes[i] = simplicial(sh->number,sh->range);
break; break;
case -1: case -1:
s *= groupsizes[i] = simplicial(shape[i].number,shape[i].range-shape[i].number+1); s *= groupsizes[i] = simplicial(sh->number,sh->range-sh->number+1);
break; break;
default: default:
laerror("illegal index group symmetry"); laerror("illegal index group symmetry");
@ -376,20 +378,21 @@ template<typename T>
void loopingroups(Tensor<T> &t, int ngroup, int igroup, T **p, SUPERINDEX &I, void (*callback)(const SUPERINDEX &, T *)) void loopingroups(Tensor<T> &t, int ngroup, int igroup, T **p, SUPERINDEX &I, void (*callback)(const SUPERINDEX &, T *))
{ {
LA_index istart,iend; LA_index istart,iend;
switch(t.shape[ngroup].symmetry) const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&t.shape))[ngroup];
switch(sh->symmetry)
{ {
case 0: case 0:
istart= t.shape[ngroup].offset; istart= sh->offset;
iend= t.shape[ngroup].offset+t.shape[ngroup].range-1; iend= sh->offset+sh->range-1;
break; break;
case 1: case 1:
istart= t.shape[ngroup].offset; istart= sh->offset;
if(igroup==t.shape[ngroup].number-1) iend= t.shape[ngroup].offset+t.shape[ngroup].range-1; if(igroup==sh->number-1) iend= sh->offset+sh->range-1;
else iend = I[ngroup][igroup+1]; else iend = I[ngroup][igroup+1];
break; break;
case -1: case -1:
istart= t.shape[ngroup].offset + igroup; istart= sh->offset + igroup;
if(igroup==t.shape[ngroup].number-1) iend= t.shape[ngroup].offset+t.shape[ngroup].range-1; if(igroup==sh->number-1) iend= sh->offset+sh->range-1;
else iend = I[ngroup][igroup+1]-1; else iend = I[ngroup][igroup+1]-1;
break; break;
} }
@ -410,7 +413,8 @@ for(LA_index i = istart; i<=iend; ++i)
if(newigroup<0) if(newigroup<0)
{ {
--newngroup; --newngroup;
newigroup=t.shape[newngroup].number-1; const indexgroup *sh2 = &(* const_cast<const NRVec<indexgroup> *>(&t.shape))[newngroup];
newigroup=sh2->number-1;
} }
loopingroups(t,newngroup,newigroup,p,I,callback); loopingroups(t,newngroup,newigroup,p,I,callback);
} }
@ -422,9 +426,16 @@ template<typename T>
void Tensor<T>::loopover(void (*callback)(const SUPERINDEX &, T *)) void Tensor<T>::loopover(void (*callback)(const SUPERINDEX &, T *))
{ {
SUPERINDEX I(shape.size()); SUPERINDEX I(shape.size());
for(int i=0; i<I.size(); ++i) {I[i].resize(shape[i].number); I[i] = shape[i].offset;} 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]; T *pp=&data[0];
loopingroups(*this,shape.size()-1,shape[shape.size()-1].number-1,&pp,I,callback); int ss=shape.size()-1;
const indexgroup *sh = &(* const_cast<const NRVec<indexgroup> *>(&shape))[ss];
loopingroups(*this,ss,sh->number-1,&pp,I,callback);
} }
@ -487,6 +498,60 @@ 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>
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;
}
template class Tensor<double>; template class Tensor<double>;
template class Tensor<std::complex<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<double> &x);

View File

@ -97,6 +97,7 @@ class LA_traits<indexgroup> {
typedef NRVec<LA_index> FLATINDEX; //all indices but in a single vector typedef NRVec<LA_index> FLATINDEX; //all indices but in a single vector
typedef NRVec<NRVec<LA_index> > SUPERINDEX; //all indices in the INDEXGROUP structure typedef NRVec<NRVec<LA_index> > SUPERINDEX; //all indices in the INDEXGROUP structure
typedef NRVec<LA_largeindex> GROUPINDEX; //set of indices in the symmetry groups
template<typename T> template<typename T>
@ -104,10 +105,9 @@ class Tensor {
public: public:
NRVec<indexgroup> shape; NRVec<indexgroup> shape;
NRVec<T> data; NRVec<T> data;
private:
int myrank; int myrank;
NRVec<LA_largeindex> groupsizes; //group sizes of symmetry index groups (a function of shape but precomputed for efficiency) NRVec<LA_largeindex> groupsizes; //group sizes of symmetry index groups (a function of shape but precomputed for efficiency)
NRVec<LA_largeindex> cumsizes; //cumulative sizes of symmetry index groups (a function of shape but precomputed for efficiency) NRVec<LA_largeindex> cumsizes; //cumulative sizes of symmetry index groups (a function of shape but precomputed for efficiency); always cumsizes[0]=1, index group 0 is the innermost-loop one
public: public:
LA_largeindex index(int *sign, const SUPERINDEX &I) const; //map the tensor indices to the position in data LA_largeindex index(int *sign, const SUPERINDEX &I) const; //map the tensor indices to the position in data
@ -117,7 +117,7 @@ public:
//constructors //constructors
Tensor() : myrank(0) {}; Tensor() : myrank(0) {};
Tensor(const NRVec<indexgroup> &s) : shape(s), data((int)calcsize()) {calcrank();}; //general tensor Tensor(const NRVec<indexgroup> &s) : shape(s) { data.resize(calcsize()); calcrank();}; //general tensor
Tensor(const indexgroup &g) {shape.resize(1); shape[0]=g; data.resize(calcsize()); calcrank();}; //tensor with a single index group Tensor(const indexgroup &g) {shape.resize(1); shape[0]=g; data.resize(calcsize()); calcrank();}; //tensor with a single index group
Tensor(const Tensor &rhs): myrank(rhs.myrank), shape(rhs.shape), groupsizes(rhs.groupsizes), cumsizes(rhs.cumsizes), data(rhs.data) {}; Tensor(const Tensor &rhs): myrank(rhs.myrank), shape(rhs.shape), groupsizes(rhs.groupsizes), cumsizes(rhs.cumsizes), data(rhs.data) {};
Tensor(int xrank, const NRVec<indexgroup> &xshape, const NRVec<LA_largeindex> &xgroupsizes, const NRVec<LA_largeindex> xcumsizes, const NRVec<T> &xdata) : myrank(xrank), shape(xshape), groupsizes(xgroupsizes), cumsizes(xcumsizes), data(xdata) {}; Tensor(int xrank, const NRVec<indexgroup> &xshape, const NRVec<LA_largeindex> &xgroupsizes, const NRVec<LA_largeindex> xcumsizes, const NRVec<T> &xdata) : myrank(xrank), shape(xshape), groupsizes(xgroupsizes), cumsizes(xcumsizes), data(xdata) {};
@ -173,14 +173,17 @@ public:
inline void randomize(const typename LA_traits<T>::normtype &x) {data.randomize(x);}; inline void randomize(const typename LA_traits<T>::normtype &x) {data.randomize(x);};
void loopover(void (*callback)(const SUPERINDEX &, T *)); void loopover(void (*callback)(const SUPERINDEX &, T *)); //loop over all elements
void grouploopover(void (*callback)(const GROUPINDEX &, T *)); //loop over all elements disregarding the internal structure of index groups
Tensor permute_index_groups(const NRPerm<int> &p) const; //rearrange the tensor storage permuting index groups as a whole
//@@@TODO - unwinding to full size in a specified index //@@@TODO - unwinding to full size in a specified index
//@@@contraction by a whole index group or by individual single index //@@@contraction by a whole index group or by individual single index
//@@@ general antisymmetrization operator Kucharski style
//@@@TODO - contractions - basic and efficient? first contraction in a single index; between a given group+index in group at each tensor //@@@TODO - contractions - basic and efficient? first contraction in a single index; between a given group+index in group at each tensor
//@@@outer product and product with a contraction //@@@outer product and product with a contraction
//@@@@symmetrize a group, antisymmetrize a group, expand a (anti)symmetric grtoup - obecne symmetry change krome +1 na -1 vse mozne //@@@@symmetrize a group, antisymmetrize a group, expand a (anti)symmetric grtoup - obecne symmetry change krome +1 na -1 vse mozne
//@@@@@@permute index groups
}; };