mxarray/mdeig.c at master · dhr/mxarray · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
#include <mex.h>
#include <string.h>
#include <lapack.h>

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {
  double *data, *dataTemp, *dataEnd;
  mxArray *eigvals, *eigvecs;
  double *evalsr, *evalsi, *evecsr, *evecsi, *evecsTemp, *work;
  ptrdiff_t m, n, ndims;
  const int *dims;
  int *eigvalDims;
  int dataInc, evalInc, evecInc;
  ptrdiff_t lwork, info;
  double *dummy;
  int count = 0;

  if (nrhs != 1) mexErrMsgTxt("Exactly one argument is expected.");

  ndims = mxGetNumberOfDimensions(prhs[0]);
  dims = mxGetDimensions(prhs[0]);
  m = dims[0];
  n = dims[1];

  if (m != n) mexErrMsgTxt("Matrices must be square.");

  data = mxGetPr(prhs[0]);
  dataInc = n*n;
  dataEnd = data + mxGetNumberOfElements(prhs[0]);
  dataTemp = mxMalloc(dataInc*sizeof(double));

  eigvecs = mxCreateNumericArray(ndims, dims, mxDOUBLE_CLASS, mxCOMPLEX);

  eigvalDims = mxMalloc(ndims*sizeof(int));
  memcpy(eigvalDims, dims, ndims*sizeof(int));
  eigvalDims[1] = 1;
  eigvals = mxCreateNumericArray(ndims, eigvalDims, mxDOUBLE_CLASS, mxCOMPLEX);

  evecInc = n*n;
  evalInc = n;

  evecsr = mxGetPr(eigvecs);
  evecsi = mxGetPi(eigvecs);
  evalsr = mxGetPr(eigvals);
  evalsi = mxGetPi(eigvals);
  evecsTemp = mxMalloc(evecInc*sizeof(double));

  lwork = 6*n;
  work = mxMalloc(lwork*sizeof(double));

  dummy = mxMalloc(evecInc*sizeof(double));

  for (; data < dataEnd; data += dataInc,
                         evalsr += evalInc,
                         evalsi += evalInc,
                         evecsr += evecInc,
                         evecsi += evecInc) {
    int i;

    memcpy(dataTemp, data, dataInc*sizeof(double));
    dgeev("N", "V", &n, dataTemp, &n, evalsr, evalsi,
          dummy, &n, evecsTemp, &n, work, &lwork, &info);

    memcpy(evecsr, evecsTemp, evecInc*sizeof(double));
    for (i = 0; i < n; ++i) {
      if (evalsi[i] != 0) {
        int j;
        for (j = 0; j < n; j++) {
          evecsr[(i + 1)*n + j] = evecsTemp[i*n + j];
          evecsi[i*n + j] = evecsTemp[(i + 1)*n + j];
          evecsi[(i + 1)*n + j] = -evecsi[i*n + j];
        }

        ++i;
      }
    }
  }

  mxFree(work);
  mxFree(dummy);
  mxFree(evecsTemp);
  mxFree(dataTemp);

  plhs[0] = eigvecs;

  if (nlhs > 1)
    plhs[1] = eigvals;
  else
    mxDestroyArray(eigvals);
}