transpose.cpp

   1 // Transpose tensor indices
   2
   3 #include "stdafx.h"
   4 #include "defs.h"
   5
   6 void
   7 eval_transpose(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         transpose();
  21 }
  22
  23 void
  24 transpose(void)
  25 {
  26         int i, j, k, l, m, ndim, nelem, t;
  27         int ai[MAXDIM], an[MAXDIM];
  28         U **a, **b;
  29
  30         save();
  31
  32         p3 = pop();
  33         p2 = pop();
  34         p1 = pop();
  35
  36         if (!istensor(p1)) {
  37                 if (!iszero(p1))
  38                         stop("transpose: tensor expected, 1st arg is not a tensor");
  39                 push(zero);
  40                 restore();
  41                 return;
  42         }
  43
  44         ndim = p1->u.tensor->ndim;
  45         nelem = p1->u.tensor->nelem;
  46
  47         // vector?
  48
  49         if (ndim == 1) {
  50                 push(p1);
  51                 restore();
  52                 return;
  53         }
  54
  55         push(p2);
  56         l = pop_integer();
  57
  58         push(p3);
  59         m = pop_integer();
  60
  61         if (l < 1 || l > ndim || m < 1 || m > ndim)
  62                 stop("transpose: index out of range");
  63
  64         l--;
  65         m--;
  66
  67         p2 = alloc_tensor(nelem);
  68
  69         p2->u.tensor->ndim = ndim;
  70
  71         for (i = 0; i < ndim; i++)
  72                 p2->u.tensor->dim[i] = p1->u.tensor->dim[i];
  73
  74         p2->u.tensor->dim[l] = p1->u.tensor->dim[m];
  75         p2->u.tensor->dim[m] = p1->u.tensor->dim[l];
  76
  77         a = p1->u.tensor->elem;
  78         b = p2->u.tensor->elem;
  79
  80         // init tensor index
  81
  82         for (i = 0; i < ndim; i++) {
  83                 ai[i] = 0;
  84                 an[i] = p1->u.tensor->dim[i];
  85         }
  86
  87         // copy components from a to b
  88
  89         for (i = 0; i < nelem; i++) {
  90
  91                 // swap indices l and m
  92
  93                 t = ai[l]; ai[l] = ai[m]; ai[m] = t;
  94                 t = an[l]; an[l] = an[m]; an[m] = t;
  95
  96                 // convert tensor index to linear index k
  97
  98                 k = 0;
  99                 for (j = 0; j < ndim; j++)
 100                         k = (k * an[j]) + ai[j];
 101
 102                 // swap indices back
 103
 104                 t = ai[l]; ai[l] = ai[m]; ai[m] = t;
 105                 t = an[l]; an[l] = an[m]; an[m] = t;
 106
 107                 // copy one element
 108
 109                 b[k] = a[i];
 110
 111                 // increment tensor index
 112
 113                 // Suppose the tensor dimensions are 2 and 3.
 114                 // Then the tensor index ai increments as follows:
 115                 // 00 -> 01
 116                 // 01 -> 02
 117                 // 02 -> 10
 118                 // 10 -> 11
 119                 // 11 -> 12
 120                 // 12 -> 00
 121
 122                 for (j = ndim - 1; j >= 0; j--) {
 123                         if (++ai[j] < an[j])
 124                                 break;
 125                         ai[j] = 0;
 126                 }
 127         }
 128
 129         push(p2);
 130         restore();
 131 }
 132
 133 #if SELFTEST
 134
 135 static char *s[] = {
 136
 137         "transpose(0)",
 138         "0",
 139
 140         "transpose(0.0)",
 141         "0",
 142
 143         "transpose(((a,b),(c,d)))",
 144         "((a,c),(b,d))",
 145
 146         "transpose(((a,b),(c,d)),1,2)",
 147         "((a,c),(b,d))",
 148
 149         "transpose(((a,b,c),(d,e,f)),1,2)",
 150         "((a,d),(b,e),(c,f))",
 151
 152         "transpose(((a,d),(b,e),(c,f)),1,2)",
 153         "((a,b,c),(d,e,f))",
 154
 155         "transpose((a,b,c))",
 156         "(a,b,c)",
 157 };
 158
 159 void
 160 test_transpose(void)
 161 {
 162         test(__FILE__, s, sizeof s / sizeof (char *));
 163 }
 164
 165 #endif