Implemented Cholesky decomposition of positive semidefinite matrices …

…with complete pivoting

TensorMath.lua

@@ -1208,6 +1208,20 @@ static void THTensor_random1__(THTensor *self, THGenerator *gen, long b)
  {name=Tensor},
  {name='charoption', values={'U', 'L'}, default='U'}} -- uplo
  )
+ interface:wrap("pstrf",
+ cname("pstrf"),
+ {{name=Tensor, returned=true},
+ {name='IntTensor', returned=true},
+ {name=Tensor},
+ {name='charoption', values={'U', 'L'}, default='U'}, -- uplo
+ {name=real, default=-1}},
+ cname("pstrf"),
+ {{name=Tensor, default=true, returned=true, invisible=true},
+ {name='IntTensor', default=true, returned=true, invisible=true},
+ {name=Tensor},
+ {name='charoption', values={'U', 'L'}, default='U'}, -- uplo
+ {name=real, default=-1}}
+ )
  interface:wrap("qr",
  cname("qr"),
  {{name=Tensor, returned=true},

doc/maths.md

@@ -1859,13 +1859,20 @@ Note: Irrespective of the original strides, the returned matrices `resb` and
 <a name="torch.potrf"></a>
 ### torch.potrf([res,] A [, 'U' or 'L'] ) ###
 
-Cholesky Decomposition of 2D tensor `A`. Matrix `A` has to be a positive-definite and either symetric or complex Hermitian.
+Cholesky Decomposition of 2D tensor `A`. Matrix `A` has to be a positive-definite and either symmetric or complex Hermitian.
 
-Optional character `uplo` = {'U', 'L'} specifies whether the upper or lower triangular decomposition should be returned. By default, `uplo` = 'U'.
+The factorization has the form
 
-`X = torch.potrf(A, 'U')` returns the upper triangular Cholesky decomposition of `X`.
+  A = U**T * U, if UPLO = 'U', or
+  A = L * L**T, if UPLO = 'L',
 
-`X = torch.potrf(A, 'L')` returns the lower triangular Cholesky decomposition of `X`.
+where `U` is an upper triangular matrix and `L` is lower triangular.
+
+The optional character `uplo` = {'U', 'L'} specifies whether the upper or lower triangular decomposition should be returned. By default, `uplo` = 'U'.
+
+`U = torch.potrf(A, 'U')` returns the upper triangular Cholesky decomposition of `A`.
+
+`L = torch.potrf(A, 'L')` returns the lower triangular Cholesky decomposition of `A`.
 
 If tensor `res` is provided, the resulting decomposition will be stored therein.
 
@@ -1896,6 +1903,67 @@ If tensor `res` is provided, the resulting decomposition will be stored therein.
 [torch.DoubleTensor of size 5x5]
 ```
 
+<a name="torch.pstrf"></a>
+### torch.pstrf([res, piv, ] A [, 'U' or 'L'] ) ###
+
+Cholesky factorization with complete pivoting of a real symmetric positive semidefinite 2D tensor `A`.
+The matrix `A` has to be a positive semi-definite and symmetric. The factorization has the form
+
+ P**T * A * P = U**T * U , if UPLO = 'U',
+ P**T * A * P = L * L**T, if UPLO = 'L',
+
+where `U` is an upper triangular matrix and `L` is lower triangular, and
+`P` is stored as the vector `piv`. More specifically, `piv` is such that the nonzero entries are `P[piv[k], k] = 1`.
+
+The optional character argument `uplo` = {'U', 'L'} specifies whether the upper or lower triangular decomposition should be returned. By default, `uplo` = 'U'.
+
+`U, piv = torch.sdtrf(A, 'U')` returns the upper triangular Cholesky decomposition of `A`
+
+`L, piv = torch.potrf(A, 'L')` returns the lower triangular Cholesky decomposition of `A`.
+
+If tensors `res` and `piv` (an `IntTensor`) are provided, the resulting decomposition will be stored therein.
+
+```lua
+> A = torch.Tensor({
+ {1.2705, 0.9971, 0.4948, 0.1389, 0.2381},
+ {0.9971, 0.9966, 0.6752, 0.0686, 0.1196},
+ {0.4948, 0.6752, 1.1434, 0.0314, 0.0582},
+ {0.1389, 0.0686, 0.0314, 0.0270, 0.0526},
+ {0.2381, 0.1196, 0.0582, 0.0526, 0.3957}})
+
+> U, piv = torch.pstrf(A)
+> U
+ 1.1272 0.4390 0.2112 0.8846 0.1232
+ 0.0000 0.9750 -0.0354 0.2942 -0.0233
+ 0.0000 0.0000 0.5915 -0.0961 0.0435
+ 0.0000 0.0000 0.0000 0.3439 -0.0854
+ 0.0000 0.0000 0.0000 0.0000 0.0456
+[torch.DoubleTensor of size 5x5]
+
+> piv
+ 1
+ 3
+ 5
+ 2
+ 4
+[torch.IntTensor of size 5]
+
+> Ap = U:t() * U
+> Ap
+ 1.2705 0.4948 0.2381 0.9971 0.1389
+ 0.4948 1.1434 0.0582 0.6752 0.0314
+ 0.2381 0.0582 0.3957 0.1196 0.0526
+ 0.9971 0.6752 0.1196 0.9966 0.0686
+ 0.1389 0.0314 0.0526 0.0686 0.0270
+[torch.DoubleTensor of size 5x5]
+
+> -- Permute rows and columns
+> Ap:indexCopy(1, piv:long(), Ap:clone())
+> Ap:indexCopy(2, piv:long(), Ap:clone())
+> (Ap - A):norm()
+1.5731560566382e-16
+```
+
 <a name="torch.potrs"></a>
 ### torch.potrs([res,] B, chol [, 'U' or 'L'] ) ###
 

lib/TH/generic/THLapack.c

@@ -31,6 +31,8 @@ TH_EXTERNC void sorgqr_(int *m, int *n, int *k, float *a, int *lda, float *tau,
 TH_EXTERNC void dorgqr_(int *m, int *n, int *k, double *a, int *lda, double *tau, double *work, int *lwork, int *info);
 TH_EXTERNC void sormqr_(char *side, char *trans, int *m, int *n, int *k, float *a, int *lda, float *tau, float *c, int *ldc, float *work, int *lwork, int *info);
 TH_EXTERNC void dormqr_(char *side, char *trans, int *m, int *n, int *k, double *a, int *lda, double *tau, double *c, int *ldc, double *work, int *lwork, int *info);
+TH_EXTERNC void spstrf_(char *uplo, int *n, float *a, int *lda, int *piv, int *rank, float *tol, float *work, int *info);
+TH_EXTERNC void dpstrf_(char *uplo, int *n, double *a, int *lda, int *piv, int *rank, double *tol, double *work, int *info);
 
 
 /* Compute the solution to a real system of linear equations A * X = B */

lib/TH/generic/THLapack.h

@@ -26,6 +26,8 @@ void THLapack_(potrf)(char uplo, int n, real *a, int lda, int *info);
 void THLapack_(potri)(char uplo, int n, real *a, int lda, int *info);
 /* Solve A*X = B with a symmetric positive definite matrix A using the Cholesky factorization */
 void THLapack_(potrs)(char uplo, int n, int nrhs, real *a, int lda, real *b, int ldb, int *info);
+/* Cholesky factorization with complete pivoting. */
+void THLapack_(pstrf)(char uplo, int n, real *a, int lda, int *piv, int *rank, real tol, real *work, int *info);
 
 /* QR decomposition */
 void THLapack_(geqrf)(int m, int n, real *a, int lda, real *tau, real *work, int lwork, int *info);

lib/TH/generic/THTensorLapack.c

@@ -431,47 +431,91 @@ void THTensor_(getri)(THTensor *ra_, THTensor *a)
  THIntTensor_free(ipiv);
 } 
 
-void THTensor_(potrf)(THTensor *ra_, THTensor *a, const char *uplo)
+void THTensor_(clearUpLoTriangle)(THTensor *a, const char *uplo)
 {
- if (a == NULL) a = ra_;
  THArgCheck(a->nDimension == 2, 1, "A should be 2 dimensional");
  THArgCheck(a->size[0] == a->size[1], 1, "A should be square");
 
- int n, lda, info;
- THTensor *ra__ = NULL;
+ int n = a->size[0];
 
- ra__ = THTensor_(cloneColumnMajor)(ra_, a);
+ /* Build full matrix */
+ real *p = THTensor_(data)(a);
+ long i, j;
 
- n = ra__->size[0];
- lda = n;
+ /* Upper Triangular Case */
+ if (uplo[0] == 'U')
+ {
+ /* Clear lower triangle (excluding diagonals) */
+ for (i=0; i<n; i++) {
+ for (j=i+1; j<n; j++) {
+ p[n*i + j] = 0;
+ }
+ }
+ }
+ /* Lower Triangular Case */
+ else if (uplo[0] == 'L')
+ {
+ /* Clear upper triangle (excluding diagonals) */
+ for (i=0; i<n; i++) {
+ for (j=0; j<i; j++) {
+ p[n*i + j] = 0;
+ }
+ }
+ }
+}
 
- /* Run Factorization */
- THLapack_(potrf)(uplo[0], n, THTensor_(data)(ra__), lda, &info);
- THLapackCheck("Lapack Error %s : A(%d,%d) is 0, A cannot be factorized", "potrf", info, info);
+void THTensor_(copyUpLoTriangle)(THTensor *a, const char *uplo)
+{
+ THArgCheck(a->nDimension == 2, 1, "A should be 2 dimensional");
+ THArgCheck(a->size[0] == a->size[1], 1, "A should be square");
+
+ int n = a->size[0];
 
  /* Build full matrix */
- real *p = THTensor_(data)(ra__);
+ real *p = THTensor_(data)(a);
  long i, j;
 
  /* Upper Triangular Case */
  if (uplo[0] == 'U')
  {
+ /* Clear lower triangle (excluding diagonals) */
  for (i=0; i<n; i++) {
  for (j=i+1; j<n; j++) {
- p[n*i + j] = 0;
+ p[n*i + j] = p[n*j+i];
  }
  }
  }
  /* Lower Triangular Case */
- else
+ else if (uplo[0] == 'L')
  {
+ /* Clear upper triangle (excluding diagonals) */
  for (i=0; i<n; i++) {
  for (j=0; j<i; j++) {
- p[n*i + j] = 0;
+ p[n*i + j] = p[n*j+i];
  }
  }
  }
+}
 
+void THTensor_(potrf)(THTensor *ra_, THTensor *a, const char *uplo)
+{
+ if (a == NULL) a = ra_;
+ THArgCheck(a->nDimension == 2, 1, "A should be 2 dimensional");
+ THArgCheck(a->size[0] == a->size[1], 1, "A should be square");
+
+ int n, lda, info;
+ THTensor *ra__ = NULL;
+
+ ra__ = THTensor_(cloneColumnMajor)(ra_, a);
+
+ n = ra__->size[0];
+ lda = n;
+
+ /* Run Factorization */
+ THLapack_(potrf)(uplo[0], n, THTensor_(data)(ra__), lda, &info);
+ THLapackCheck("Lapack Error %s : A(%d,%d) is 0, A cannot be factorized", "potrf", info, info);
+
+ THTensor_(clearUpLoTriangle)(ra__, uplo);
  THTensor_(freeCopyTo)(ra__, ra_);
 }
 
@@ -522,30 +566,52 @@ void THTensor_(potri)(THTensor *ra_, THTensor *a, const char *uplo)
  THLapack_(potri)(uplo[0], n, THTensor_(data)(ra__), lda, &info);
  THLapackCheck("Lapack Error %s : A(%d,%d) is 0, A cannot be factorized", "potri", info, info);
 
- /* Build full matrix */
- real *p = THTensor_(data)(ra__);
- long i, j;
+ THTensor_(copyUpLoTriangle)(ra__, uplo);
+ THTensor_(freeCopyTo)(ra__, ra_);
+}
 
- /* Upper Triangular Case */
- if (uplo[0] == 'U')
- {
- for (i=0; i<n; i++) {
- for (j=i+1; j<n; j++) {
- p[n*i+j] = p[n*j+i];
- }
- }
- }
- /* Lower Triangular Case */
- else
- {
- for (i=0; i<n; i++) {
- for (j=0; j<i; j++) {
- p[n*i + j] = p[n*j+i];
- }
- }
- }
+/*
+ Computes the Cholesky factorization with complete pivoting of a real symmetric
+ positive semidefinite matrix.
+
+ Args:
+ * `ra_` - result Tensor in which to store the factor U or L from the
+ Cholesky factorization.
+ * `rpiv_` - result IntTensor containing sparse permutation matrix P, encoded
+ as P[rpiv_[k], k] = 1.
+ * `a` - input Tensor; the input matrix to factorize.
+ * `uplo` - string; specifies whether the upper or lower triangular part of
+ the symmetric matrix A is stored. "U"/"L" for upper/lower
+ triangular.
+ * `tol` - double; user defined tolerance, or < 0 for automatic choice.
+ The algorithm terminates when the pivot <= tol.
+ */
+void THTensor_(pstrf)(THTensor *ra_, THIntTensor *rpiv_, THTensor *a, const char *uplo, real tol) {
+ THArgCheck(a->nDimension == 2, 1, "A should be 2 dimensional");
+ THArgCheck(a->size[0] == a->size[1], 1, "A should be square");
+
+ int n = a->size[0];
+
+ THTensor *ra__ = THTensor_(cloneColumnMajor)(ra_, a);
+ THTensor_(resize1d)(rpiv_, n);
+
+ // Allocate working tensor
+ THTensor *work = THTensor_(newWithSize1d)(2 * n);
+
+ // Run Cholesky factorization
+ int lda = n;
+ int rank, info;
+
+ THLapack_(pstrf)(uplo[0], n, THTensor_(data)(ra__), lda,
+ THIntTensor_data(rpiv_), &rank, tol,
+ THTensor_(data)(work), &info);
+
+ THLapackCheck("Lapack Error %s : matrix is rank deficient or not positive semidefinite", "pstrf", info);
+
+ THTensor_(clearUpLoTriangle)(ra__, uplo);
 
  THTensor_(freeCopyTo)(ra__, ra_);
+ THTensor_(free)(work);
 }
 
 /*

lib/TH/generic/THTensorLapack.h

@@ -17,5 +17,6 @@ TH_API void THTensor_(qr)(THTensor *rq_, THTensor *rr_, THTensor *a);
 TH_API void THTensor_(geqrf)(THTensor *ra_, THTensor *rtau_, THTensor *a);
 TH_API void THTensor_(orgqr)(THTensor *ra_, THTensor *a, THTensor *tau);
 TH_API void THTensor_(ormqr)(THTensor *ra_, THTensor *a, THTensor *tau, THTensor *c, const char *side, const char *trans);
+TH_API void THTensor_(pstrf)(THTensor *ra_, THIntTensor *rpiv_, THTensor*a, const char* uplo, real tol);
 
 #endif

test/test.lua

@@ -2215,6 +2215,51 @@ function torchtest.potri()
  mytester:assertlt(inv0:dist(inv1),1e-12,"torch.potri; uplo='L'")
 end
 
+function torchtest.pstrf()
+ local function checkPsdCholesky(a, uplo, inplace)
+ local u, piv, args, a_reconstructed
+ if inplace then
+ u = torch.Tensor(a:size())
+ piv = torch.IntTensor(a:size(1))
+ args = {u, piv, a}
+ else
+ args = {a}
+ end
+
+ if uplo then table.insert(args, uplo) end
+
+ u, piv = torch.pstrf(unpack(args))
+
+ if uplo == 'L' then
+ a_reconstructed = torch.mm(u, u:t())
+ else
+ a_reconstructed = torch.mm(u:t(), u)
+ end
+
+ piv = piv:long()
+ local a_permuted = a:index(1, piv):index(2, piv)
+ mytester:assertTensorEq(a_permuted, a_reconstructed, 1e-14,
+ 'torch.pstrf did not allow rebuilding the original matrix;' ..
+ 'uplo=' .. tostring(uplo))
+ end
+
+ local dimensions = { {5, 1}, {5, 3}, {5, 5}, {10, 10} }
+ for _, dim in pairs(dimensions) do
+ local m = torch.Tensor(unpack(dim)):uniform()
+ local a = torch.mm(m, m:t())
+ -- add a small number to the diagonal to make the matrix numerically positive semidefinite
+ for i = 1, m:size(1) do
+ a[i][i] = a[i][i] + 1e-7
+ end
+ checkPsdCholesky(a, nil, false)
+ checkPsdCholesky(a, 'U', false)
+ checkPsdCholesky(a, 'L', false)
+ checkPsdCholesky(a, nil, true)
+ checkPsdCholesky(a, 'U', true)
+ checkPsdCholesky(a, 'L', true)
+ end
+end
+
 function torchtest.testNumel()
  local b = torch.ByteTensor(3, 100, 100)
  mytester:asserteq(b:nElement(), 3*100*100, "nElement not right")