adding potrs and uplo option to potrf

adding tests for torch.potrs and (modified) torch.potrf

Author: j-wilson

TensorMath.lua

@@ -1109,20 +1109,33 @@ static void THTensor_random1__(THTensor *self, THGenerator *gen, long b)
                      {{name=Tensor, default=true, returned=true, invisible=true},
                       {name=Tensor}}
                   )
-     
-      interface:wrap("potri",
-                     cname("potri"),
+      interface:wrap("potrf",
+                     cname("potrf"),
                      {{name=Tensor, returned=true},
-                      {name=Tensor}},
-                     cname("potri"),
+                      {name=Tensor},
+                      {name='charoption', values={'U', 'L'}, default='U'}}, -- uplo
+                     cname("potrf"),
                      {{name=Tensor, default=true, returned=true, invisible=true},
-                      {name=Tensor}}
+                      {name=Tensor},
+                      {name='charoption', values={'U', 'L'}, default='U'}}
                   )
-      interface:wrap("potrf",
-                     cname("potrf"),
+      interface:wrap("potrs",
+                     cname("potrs"),
+                     {{name=Tensor, returned=true},
+                      {name=Tensor},
+                      {name=Tensor},
+                      {name='charoption', values={'U', 'L'}, default='U'}}, -- uplo
+                     cname("potrs"),
+                     {{name=Tensor, default=true, returned=true, invisible=true},
+                      {name=Tensor},
+                      {name=Tensor},
+                      {name='charoption', values={'U', 'L'}, default='U'}}
+                  )
+      interface:wrap("potri",
+                     cname("potri"),
                      {{name=Tensor, returned=true},
                       {name=Tensor}},
-                     cname("potrf"),
+                     cname("potri"),
                      {{name=Tensor, default=true, returned=true, invisible=true},
                       {name=Tensor}}
                   )

doc/maths.md

@@ -1621,6 +1621,102 @@ x = torch.trtrs(b, a)
 4.1895292266754e-15
 ```
 
+<a name="torch."></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.
+
+Optional character `uplo` = {'U', 'L'} specified whether the upper or lower triangular decomposition should be returned. By default, `uplo` = 'U'.
+
+`X = torch.potrf(A, 'U')` returns the upper triangular Cholesky decomposition of X.
+
+`X = torch.potrf(A, 'L')` returns the lower triangular Cholesky decomposition of X.
+
+If tensor `res` is provided, the resulting decomposition will be stored therein.
+
+```
+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}})
+
+chol = torch.potrf(A)
+> chol
+ 1.1272  0.8846  0.4390  0.1232  0.2112
+ 0.0000  0.4626  0.6200 -0.0874 -0.1453
+ 0.0000  0.0000  0.7525  0.0419  0.0738
+ 0.0000  0.0000  0.0000  0.0491  0.2199
+ 0.0000  0.0000  0.0000  0.0000  0.5255
+[torch.DoubleTensor of size 5x5]
+
+torch.potrf(chol, A, 'L')
+> chol
+ 1.1272  0.0000  0.0000  0.0000  0.0000
+ 0.8846  0.4626  0.0000  0.0000  0.0000
+ 0.4390  0.6200  0.7525  0.0000  0.0000
+ 0.1232 -0.0874  0.0419  0.0491  0.0000
+ 0.2112 -0.1453  0.0738  0.2199  0.5255
+[torch.DoubleTensor of size 5x5]
+```
+
+<a name="torch."></a>
+### torch.potrs([res,] chol [, 'U' or 'L'] ) ###
+
+Returns the solution to linear system `AX = B` using the Cholesky decomposition `chol` of 2D tensor `A`.
+
+Square matrix `chol` should be triangular; and, righthand side matrix `B` should be of full rank.
+
+Optional character `uplo` = {'U', 'L'} specified matrix `chol` as being other upper or lower triangular; and, by default, equals 'U'.
+
+If tensor `res` is provided, the resulting decomposition will be stored therein.
+
+```
+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}})
+
+B = torch.Tensor({
+    {0.6219,  0.3439,  0.0431},
+    {0.5642,  0.1756,  0.0153},
+    {0.2334,  0.8594,  0.4103},
+    {0.7556,  0.1966,  0.9637},
+    {0.1420,  0.7185,  0.7476}})
+
+chol = torch.potrf(A)
+> chol
+ 1.1272  0.8846  0.4390  0.1232  0.2112
+ 0.0000  0.4626  0.6200 -0.0874 -0.1453
+ 0.0000  0.0000  0.7525  0.0419  0.0738
+ 0.0000  0.0000  0.0000  0.0491  0.2199
+ 0.0000  0.0000  0.0000  0.0000  0.5255
+[torch.DoubleTensor of size 5x5]
+
+solve = torch.potrs(B, chol)
+> solve
+  12.1945   61.8622   92.6882
+ -11.1782  -97.0303 -138.4874
+ -15.3442  -76.6562 -116.8218
+   6.1930   13.5238   25.2056
+  29.9678  251.7346  360.2301
+[torch.DoubleTensor of size 5x3]
+
+> A*solve
+ 0.6219  0.3439  0.0431
+ 0.5642  0.1756  0.0153
+ 0.2334  0.8594  0.4103
+ 0.7556  0.1966  0.9637
+ 0.1420  0.7185  0.7476
+[torch.DoubleTensor of size 5x3]
+
+> B:dist(A*solve)
+4.6783066076306e-14
+```
+
 <a name="torch.gels"></a>
 ### torch.gels([resb, resa,] b, a) ###
 

lib/TH/generic/THLapack.c

@@ -162,31 +162,31 @@ void THLapack_(potrf)(char uplo, int n, real *a, int lda, int *info)
 #endif
 }
 
-/* Cholesky factorization based Matrix Inverse */
-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)
 {
 #ifdef  USE_LAPACK
 #if defined(TH_REAL_IS_DOUBLE)
-  dpotri_(&uplo, &n, a, &lda, info);
+  dpotrs_(&uplo, &n, &nrhs, a, &lda, b, &ldb, info);
 #else
-  spotri_(&uplo, &n, a, &lda, info);
+  spotrs_(&uplo, &n, &nrhs, a, &lda, b, &ldb, info);
 #endif
 #else
-  THError("potri: Lapack library not found in compile time\n");
+  THError("potrs: Lapack library not found in compile time\n");
 #endif
 }
 
-/* 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 based Matrix Inverse */
+void THLapack_(potri)(char uplo, int n, real *a, int lda, int *info)
 {
 #ifdef  USE_LAPACK
 #if defined(TH_REAL_IS_DOUBLE)
-  dpotrs_(&uplo, &n, &nrhs, a, &lda, b, &ldb, info);
+  dpotri_(&uplo, &n, a, &lda, info);
 #else
-  spotrs_(&uplo, &n, &nrhs, a, &lda, b, &ldb, info);
+  spotri_(&uplo, &n, a, &lda, info);
 #endif
 #else
-  THError("potrs: Lapack library not found in compile time\n");
+  THError("potri: Lapack library not found in compile time\n");
 #endif
 }
 

lib/TH/generic/THTensorLapack.c

@@ -429,16 +429,15 @@ void THTensor_(getri)(THTensor *ra_, THTensor *a)
   THTensor_(freeCopyTo)(ra__, ra_);
   THTensor_(free)(work);
   THIntTensor_free(ipiv);
-}
+} 
 
-void THTensor_(potrf)(THTensor *ra_, THTensor *a)
+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;
-  char uplo = 'U';
   THTensor *ra__ = NULL;
 
   ra__ = THTensor_(cloneColumnMajor)(ra_, a);
@@ -447,23 +446,64 @@ void THTensor_(potrf)(THTensor *ra_, THTensor *a)
   lda = n;
 
   /* Run Factorization */
-  THLapack_(potrf)(uplo, n, THTensor_(data)(ra__), lda, &info);
+  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);
 
-  /* Build full upper-triangular matrix */
+  /* Build full matrix */
+  real *p = THTensor_(data)(ra__);
+  long i, j;
+
+  /* Upper Triangular Case */
+  if (uplo[0] == 'U')
   {
-    real *p = THTensor_(data)(ra__);
-    long i,j;
     for (i=0; i<n; i++) {
-      for (j=i+1; j<n; j++) {
-        p[i*n+j] = 0;
+     for (j=i+1; j<n; j++) {
+        p[n*i + j] = 0;
+      }
+    }
+  }
+  /* Lower Triangular Case */
+  else
+  {
+    for (i=0; i<n; i++) {
+      for (j=0; j<i; j++) {
+        p[n*i + j] = 0;
       }
     }
   }
 
   THTensor_(freeCopyTo)(ra__, ra_);
 }
 
+void THTensor_(potrs)(THTensor *rb_, THTensor *b, THTensor *a, const char *uplo)
+{
+  if (b == NULL) b = rb_;
+
+  THArgCheck(a->size[0] == a->size[1], 2, "A should be square");
+  THArgCheck(b->size[0] >= b->size[1], 2, "Matrix B is rank-deficient");
+
+  int n, nrhs, lda, ldb, info;
+  THTensor *ra__; // working version of A matrix to be passed into lapack TRTRS
+  THTensor *rb__; // working version of B matrix to be passed into lapack TRTRS
+
+  ra__ = THTensor_(cloneColumnMajor)(NULL, a);
+  rb__ = THTensor_(cloneColumnMajor)(rb_, b);
+
+  n    = (int)ra__->size[0];
+  nrhs = (int)rb__->size[1];
+  lda  = n;
+  ldb  = n;
+
+  THLapack_(potrs)(uplo[0], n, nrhs, THTensor_(data)(ra__), 
+                   lda, THTensor_(data)(rb__), ldb, &info);
+
+
+  THLapackCheck("Lapack Error in %s : A(%d,%d) is zero, singular A", "potrs", info, info);
+
+  THTensor_(free)(ra__);
+  THTensor_(freeCopyTo)(rb__, rb_);
+}
+
 void THTensor_(potri)(THTensor *ra_, THTensor *a)
 {
   if (a == NULL) a = ra_;

lib/TH/generic/THTensorLapack.h

@@ -10,8 +10,9 @@ TH_API void THTensor_(geev)(THTensor *re_, THTensor *rv_, THTensor *a_, const ch
 TH_API void THTensor_(gesvd)(THTensor *ru_, THTensor *rs_, THTensor *rv_, THTensor *a, const char *jobu);
 TH_API void THTensor_(gesvd2)(THTensor *ru_, THTensor *rs_, THTensor *rv_, THTensor *ra_, THTensor *a, const char *jobu);
 TH_API void THTensor_(getri)(THTensor *ra_, THTensor *a);
+TH_API void THTensor_(potrf)(THTensor *ra_, THTensor *a, const char *uplo);
+TH_API void THTensor_(potrs)(THTensor *rb_, THTensor *b_, THTensor *a_,  const char *uplo);
 TH_API void THTensor_(potri)(THTensor *ra_, THTensor *a);
-TH_API void THTensor_(potrf)(THTensor *ra_, THTensor *a);
 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);

test/test.lua

@@ -2035,11 +2035,48 @@ function torchtest.testBoxMullerState()
 end
 
 function torchtest.testCholesky()
-    local x = torch.rand(10,10)
-    local A = torch.mm(x, x:t())
-    local C = torch.potrf(A)
-    local B = torch.mm(C:t(), C)
-    mytester:assertTensorEq(A, B, 1e-14, 'potrf did not allow rebuilding the original matrix')
+   local x = torch.rand(10,10)
+   local A = torch.mm(x, x:t())
+
+   ---- Default Case
+   local C = torch.potrf(A)
+   local B = torch.mm(C:t(), C)
+   mytester:assertTensorEq(A, B, 1e-14, 'potrf did not allow rebuilding the original matrix')
+
+    ---- Test Upper Triangular
+    local U = torch.potrf(A, 'U')
+          B = torch.mm(U:t(), U)
+    mytester:assertTensorEq(A, B, 1e-14, 'potrf (upper) did not allow rebuilding the original matrix')
+
+    ---- Test Lower Triangular
+    local L = torch.potrf(A, 'L')
+          B = torch.mm(L, L:t())
+    mytester:assertTensorEq(A, B, 1e-14, 'potrf (lower) did not allow rebuilding the original matrix')          
+end
+
+function torchtest.potrs()
+   if not torch.potrs then return end
+   local a=torch.Tensor({{6.80, -2.11,  5.66,  5.97,  8.23},
+                         {-6.05, -3.30,  5.36, -4.44,  1.08},
+                         {-0.45,  2.58, -2.70,  0.27,  9.04},
+                         {8.32,  2.71,  4.35, -7.17,  2.14},
+                         {-9.67, -5.14, -7.26,  6.08, -6.87}}):t()
+   local b=torch.Tensor({{4.02,  6.19, -8.22, -7.57, -3.03},
+                         {-1.56,  4.00, -8.67,  1.75,  2.86},
+                         {9.81, -4.09, -4.57, -8.61,  8.99}}):t()
+
+   ---- Make sure 'a' is symmetric PSD 
+   a = torch.mm(a, a:t())
+
+   ---- Upper Triangular Test
+   local U = torch.potrf(a, 'U')
+   local x = torch.potrs(b, U, 'U')
+   mytester:assertlt(b:dist(a*x),1e-12,'torch.trtrs')
+
+   ---- Lower Triangular Test
+   local L = torch.potrf(a, 'L')
+   x = torch.potrs(b, L, 'L')
+   mytester:assertlt(b:dist(a*x),1e-12,'torch.trtrs')
 end
 
 function torchtest.testNumel()