contract.tmpl

    SUBROUTINE tensorcontract${rankA}${rankB}${rankT}${rankoutT} ( A, B, ash, ain, bsh, bin, ABorder, squeeze, error, outT)
    !
    !   Purpose:
    !       Contract tensor A and B. 
    !       A and B are tensor of ranks as specified in the function name.
    !       The function name should be understood in the following way:
    !       tensorcontract4444 means contract two rank 4 tensors ('44')
    !   into an intermediate rank 4 tensor ('4') and the final result is 
    !   a rank 4 tensor ('4').     
    !       ash and bsh are the ranks of the tensor A and B.
    !       ain and bin are the common index in A and B respectively. 
    !       ABorder gives the contracted tensor sequence. 
    !       squeeze specifis the new shape of the contracted tensor.
    !       For example, A(x,y,i,m,n,p), B(a,i,m,b,c); if we want to contract the
    !   common index i and m then we need to have ain=(3,4) and bin=(2,3); 
    !   after the contraction, we will get the intermediate tensor T(xynp,abc);
    !   in order to get T(x,a,y,b,n,p,c), we need to set ABorder=(1,5,2,6,3,4,7); 
    !   it usually works in the following way:
    !   T(xynp,abc):
    !       x y n p a b c
    !       1 2 3 4 5 6 7
    !   T(x,a,y,b,n,p,c)
    !       x a y b n p c
    !       1 5 2 6 3 4 7   ! This is just ABorder
    !   
    !       If squeeze=(3,2,2), then one will squeeze the first 3(xay) index,
    !    the middle 2(bn), and the last 2(pc) index to form a rank 3 tensor.
    !       error is the error flag. 
    !       error = 0: no error
    !       error = 1: The input arrays and ranks are not consistent.
    !       error = 2: The contracted index is not equal.
    !
    !       In order to get functions for other ranks, one will have to change
    !   the declaration of the following variables:
    !   A
    !   B
    !   Anew
    !   Bnew
    !   outT
    !   T
    !   Ssize
    !   and some others as marked in the template file 'contract.tmpl'.
    !
    !       The code can be generated by 'Cheetah' template engine for differet 
    !   input and output ranks. See the python file 'contract_gen.py' for details.
 
    USE set_precisions
    USE useful_functions
    IMPLICIT NONE
    
    !   Data dictionary: declare calling parameter types & definitions
    REAL(kind=DBL), INTENT(IN), DIMENSION(${dimA}) :: A     ! Input array A
    REAL(kind=DBL), INTENT(IN), DIMENSION(${dimB}) :: B     ! Input array B
    INTEGER, INTENT(IN) :: ash                              ! Rank of A 
    INTEGER, INTENT(IN), DIMENSION(:) :: ain                ! Contraction index 
    INTEGER, INTENT(IN) :: bsh                              ! Rank of B
    INTEGER, INTENT(IN), DIMENSION(:) :: bin                ! Contraction index
    INTEGER, INTENT(IN), DIMENSION(:), OPTIONAL :: ABorder  ! Opt input
    INTEGER, INTENT(IN), DIMENSION(:), OPTIONAL :: squeeze  ! Opt input
    INTEGER, INTENT(OUT) :: error                           ! Error flag
    REAL(kind=DBL), INTENT(OUT), DIMENSION(${dimoutT}) :: outT  ! Output
    
    !   Data dictionary: declare local variable types & definitions
    INTEGER, DIMENSION(ash) :: indexA                       ! Index array A 
    INTEGER, DIMENSION(bsh) :: indexB                       ! Index array B
    INTEGER, DIMENSION(ash - SIZE(ain)) :: Aother   ! Index not contracted in A
    INTEGER, DIMENSION(bsh - SIZE(bin)) :: Bother   ! Index not contracted in B
    INTEGER :: Nain, Nbin                           ! # of index contracted
    LOGICAL :: flagA, flagB                         ! Flags for permutation
    REAL(kind=DBL), ALLOCATABLE, DIMENSION(${dimAnew}) :: Anew  ! Copy of A 
    REAL(kind=DBL), ALLOCATABLE, DIMENSION(${dimBnew}) :: Bnew  ! Copy of B
    REAL(kind=DBL), ALLOCATABLE, DIMENSION(${dimT}) :: T    ! Intermediate one     
    INTEGER, DIMENSION(ash) :: Asize                        ! Size of Anew
    INTEGER, DIMENSION(bsh) :: Bsize                        ! Size of Bnew
    INTEGER, DIMENSION(SIZE(ain)) :: SefAsize               ! Self dimension
    INTEGER, DIMENSION(SIZE(bin)) :: SefBsize               ! Self dimension
    INTEGER, DIMENSION(ash - SIZE(ain)) :: Asize1           ! Other dimension
    INTEGER, DIMENSION(bsh - SIZE(bin)) :: Bsize1           ! Other dimension
    INTEGER :: i, k
    INTEGER :: Num
    INTEGER, DIMENSION(ash + bsh - SIZE(ain) - SIZE(bin) ) :: Tsize
    INTEGER, DIMENSION(SIZE(SHAPE(T))) :: OrderSize
    INTEGER, DIMENSION(${dimSsize}) :: Ssize      
    INTEGER :: istat


    ! Check the consistency of the tensor rank
    errorif: IF ( ash /= SIZE(SHAPE(A)) .OR. bsh /= SIZE(SHAPE(B)) ) THEN
        error = 1
        WRITE(*,*) 'The input arrays and ranks are not consistent.'
    ELSE errorif
        error = 0
    END IF errorif
    
    indexA = (/ (i, i=1,ash) /)                     ! index array A 
    indexB = (/ (i, i=1,bsh) /)                     ! index array B
    Aother = setdiff( indexA, ain)                  ! index not contracted in A
    Bother = setdiff( indexB, bin)                  ! index not contracted in B
    Nain = SIZE(ain)                                ! # of index contracted
    Nbin = SIZE(bin)                                ! # of index contracted
    flagA = ALL( indexA(ash - Nain + 1: ash) == ain)
    flagB = ALL( indexB(: Nbin) == bin)
    
    ALLOCATE ( Anew(${shapeAnew} ) , STAT=istat) ! replace
    ! permute A only when necessary
    ifAllocatedA: IF ( ALLOCATED(Anew) ) THEN
        ifA: IF ( .NOT. flagA ) THEN                ! Anew(...i) 
            !Anew = RESHAPE( A, SHAPE(A), ORDER = (/ Aother, ain /) )
            Anew = RESHAPE( A, SHAPE(A), ORDER = neworder( (/ Aother, ain /) ) )
                                                    ! equivalent to permute
        ELSE ifA
            Anew = A
        END IF ifA
    ELSE ifAllocatedA
        WRITE(*,*) 'Warning-Array Anew not allocated!'
    END IF ifAllocatedA

    ALLOCATE ( Bnew(${shapeBnew} ) , STAT=istat) ! replace
    ! permute B only when necessary
    ifAllocatedB: IF ( ALLOCATED(Bnew) ) THEN
        ifB: IF ( .NOT. flagB ) THEN                ! Bnew(i...) 
            !Bnew = RESHAPE( B, SHAPE(B), ORDER = (/ bin, Bother /) )
            Bnew = RESHAPE( B, SHAPE(B), ORDER = neworder( (/ bin, Bother /) ) )
                                                    ! equivalent to permute
        ELSE ifB
            Bnew = B
        END IF ifB
    ELSE ifAllocatedB
        WRITE(*,*) 'Warning-Array Bnew not allocated!'
    END IF ifAllocatedB

    Asize = UBOUND(Anew)                            ! size of Anew 
    Bsize = UBOUND(Bnew)                            ! size of Bnew
    SefAsize = Asize( ash - Nain + 1 : ash )        ! self dimension
    SefBsize = Bsize( 1 : Nbin )                    ! self dimension
    Asize1 = Asize( 1 : ash - Nain )                ! other dimension
    Bsize1 = Bsize( Nbin + 1 : bsh  )               ! other dimension
    
    ! contract Anew(aaa, i) and Bnew(i, bbbb) to T(aaa, bbbb)
    Tsize = (/ Asize1, Bsize1 /)                    ! size of T
    prodAB: IF ( PRODUCT(SefAsize) == PRODUCT(SefBsize) ) THEN
        ABorderif: IF ( PRESENT( ABorder) ) THEN
            OrderSize = Tsize(ABorder)
            ALLOCATE(T(${shapeT} ),&
                     STAT=istat)                    ! allocate T
            ifAllocatedT: IF ( ALLOCATED(T) ) THEN
                T = RESHAPE( MATMUL( &              ! contract and reshape T
                RESHAPE( Anew, SHAPE = (/ PRODUCT(Asize1), PRODUCT(SefAsize) /) ), &
                RESHAPE( Bnew, SHAPE = (/ PRODUCT(SefBsize), PRODUCT(Bsize1) /) )  &
                            ), SHAPE(T), ORDER = neworder( ABorder ) )
            ELSE ifAllocatedT
                WRITE(*,*) 'Warning-Array T not allocated!'
            END IF ifAllocatedT
        ELSE ABorderif
            OrderSize = Tsize
            ALLOCATE(T(${shapeT} ),&
                     STAT=istat)        ! replace
            ifAllocatedT2: IF ( ALLOCATED(T) ) THEN
                T = RESHAPE( MATMUL( &
                RESHAPE( Anew, SHAPE = (/ PRODUCT(Asize1), PRODUCT(SefAsize) /) ), &
                RESHAPE( Bnew, SHAPE = (/ PRODUCT(SefBsize), PRODUCT(Bsize1) /) )  &
                            ), SHAPE(T) )
            ELSE ifAllocatedT2
                WRITE(*,*) 'Warning-Array T not allocated!'
            END IF ifAllocatedT2
        END IF ABorderif

        DEALLOCATE(Anew, STAT=istat)
        DEALLOCATE(Bnew, STAT=istat)

    ELSE prodAB
        error = 2
        WRITE(*,*) 'The contracted index is not equal.'
    END IF prodAB

    ! if optional arguments 'ABorder' and 'squeeze' are present
    IF ( PRESENT(ABorder) .AND. PRESENT(squeeze) ) THEN 
        Num = SIZE(squeeze)
        !Tsize = (/ Asize1, Bsize1 /)
        !OrderSize = Tsize(ABorder)
        k = 1
        DO i = 1, Num
            Ssize( i ) = PRODUCT( OrderSize( k: k + squeeze(i) -1 ) )
            k = k + squeeze(i)
        END DO
        $outTReshape
    END IF
    $outTNoReshape
    DEALLOCATE(T, STAT = istat)
    
    END SUBROUTINE tensorcontract${rankA}${rankB}${rankT}${rankoutT}

!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!