Eigenmath/contract.cpp

115 lines
1.7 KiB
C++

// Contract across tensor indices
#include "stdafx.h"
#include "defs.h"
void
eval_contract(void)
{
push(cadr(p1));
eval();
if (cddr(p1) == symbol(NIL)) {
push_integer(1);
push_integer(2);
} else {
push(caddr(p1));
eval();
push(cadddr(p1));
eval();
}
contract();
}
void
contract(void)
{
save();
yycontract();
restore();
}
void
yycontract(void)
{
int h, i, j, k, l, m, n, ndim, nelem;
int ai[MAXDIM], an[MAXDIM];
U **a, **b;
p3 = pop();
p2 = pop();
p1 = pop();
if (!istensor(p1)) {
if (!iszero(p1))
stop("contract: tensor expected, 1st arg is not a tensor");
push(zero);
return;
}
push(p2);
l = pop_integer();
push(p3);
m = pop_integer();
ndim = p1->u.tensor->ndim;
if (l < 1 || l > ndim || m < 1 || m > ndim || l == m
|| p1->u.tensor->dim[l - 1] != p1->u.tensor->dim[m - 1])
stop("contract: index out of range");
l--;
m--;
n = p1->u.tensor->dim[l];
// nelem is the number of elements in "b"
nelem = 1;
for (i = 0; i < ndim; i++)
if (i != l && i != m)
nelem *= p1->u.tensor->dim[i];
p2 = alloc_tensor(nelem);
p2->u.tensor->ndim = ndim - 2;
j = 0;
for (i = 0; i < ndim; i++)
if (i != l && i != m)
p2->u.tensor->dim[j++] = p1->u.tensor->dim[i];
a = p1->u.tensor->elem;
b = p2->u.tensor->elem;
for (i = 0; i < ndim; i++) {
ai[i] = 0;
an[i] = p1->u.tensor->dim[i];
}
for (i = 0; i < nelem; i++) {
push(zero);
for (j = 0; j < n; j++) {
ai[l] = j;
ai[m] = j;
h = 0;
for (k = 0; k < ndim; k++)
h = (h * an[k]) + ai[k];
push(a[h]);
add();
}
b[i] = pop();
for (j = ndim - 1; j >= 0; j--) {
if (j == l || j == m)
continue;
if (++ai[j] < an[j])
break;
ai[j] = 0;
}
}
if (nelem == 1)
push(b[0]);
else
push(p2);
}