contract.cpp

   1 // Contract across tensor indices
   2
   3 #include "stdafx.h"
   4 #include "defs.h"
   5
   6 void
   7 eval_contract(void)
   8 {
   9         push(cadr(p1));
  10         eval();
  11         if (cddr(p1) == symbol(NIL)) {
  12                 push_integer(1);
  13                 push_integer(2);
  14         } else {
  15                 push(caddr(p1));
  16                 eval();
  17                 push(cadddr(p1));
  18                 eval();
  19         }
  20         contract();
  21 }
  22
  23 void
  24 contract(void)
  25 {
  26         save();
  27         yycontract();
  28         restore();
  29 }
  30
  31 void
  32 yycontract(void)
  33 {
  34         int h, i, j, k, l, m, n, ndim, nelem;
  35         int ai[MAXDIM], an[MAXDIM];
  36         U **a, **b;
  37
  38         p3 = pop();
  39         p2 = pop();
  40         p1 = pop();
  41
  42         if (!istensor(p1)) {
  43                 if (!iszero(p1))
  44                         stop("contract: tensor expected, 1st arg is not a tensor");
  45                 push(zero);
  46                 return;
  47         }
  48
  49         push(p2);
  50         l = pop_integer();
  51
  52         push(p3);
  53         m = pop_integer();
  54
  55         ndim = p1->u.tensor->ndim;
  56
  57         if (l < 1 || l > ndim || m < 1 || m > ndim || l == m
  58         || p1->u.tensor->dim[l - 1] != p1->u.tensor->dim[m - 1])
  59                 stop("contract: index out of range");
  60
  61         l--;
  62         m--;
  63
  64         n = p1->u.tensor->dim[l];
  65
  66         // nelem is the number of elements in "b"
  67
  68         nelem = 1;
  69         for (i = 0; i < ndim; i++)
  70                 if (i != l && i != m)
  71                         nelem *= p1->u.tensor->dim[i];
  72
  73         p2 = alloc_tensor(nelem);
  74
  75         p2->u.tensor->ndim = ndim - 2;
  76
  77         j = 0;
  78         for (i = 0; i < ndim; i++)
  79                 if (i != l && i != m)
  80                         p2->u.tensor->dim[j++] = p1->u.tensor->dim[i];
  81
  82         a = p1->u.tensor->elem;
  83         b = p2->u.tensor->elem;
  84
  85         for (i = 0; i < ndim; i++) {
  86                 ai[i] = 0;
  87                 an[i] = p1->u.tensor->dim[i];
  88         }
  89
  90         for (i = 0; i < nelem; i++) {
  91                 push(zero);
  92                 for (j = 0; j < n; j++) {
  93                         ai[l] = j;
  94                         ai[m] = j;
  95                         h = 0;
  96                         for (k = 0; k < ndim; k++)
  97                                 h = (h * an[k]) + ai[k];
  98                         push(a[h]);
  99                         add();
 100                 }
 101                 b[i] = pop();
 102                 for (j = ndim - 1; j >= 0; j--) {
 103                         if (j == l || j == m)
 104                                 continue;
 105                         if (++ai[j] < an[j])
 106                                 break;
 107                         ai[j] = 0;
 108                 }
 109         }
 110
 111         if (nelem == 1)
 112                 push(b[0]);
 113         else
 114                 push(p2);
 115 }