Specialize jacobi() for 3 dimensions

diff --git a/src/gromacs/linearalgebra/nrjac.cpp b/src/gromacs/linearalgebra/nrjac.cpp
index f5cbc5880055a043f8d585c3bbdee56622b703c9..e75c442974902c01a0fcedc52f687d73d26c0f74 100644 (file)
--- a/src/gromacs/linearalgebra/nrjac.cpp
+++ b/src/gromacs/linearalgebra/nrjac.cpp
@@ -4,7 +4,7 @@
  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
  * Copyright (c) 2001-2004, The GROMACS development team.
  * Copyright (c) 2012,2014,2015,2017,2018 by the GROMACS development team.
- * Copyright (c) 2019,2020, by the GROMACS development team, led by
+ * Copyright (c) 2019,2020,2021, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
@@ -43,9 +43,11 @@
 #include <cmath>
 
 #include "gromacs/utility/fatalerror.h"
+#include "gromacs/utility/gmxassert.h"
 #include "gromacs/utility/smalloc.h"
 
-static inline void do_rotate(double** a, int i, int j, int k, int l, double tau, double s)
+template<typename MatrixType>
+static inline void do_rotate(MatrixType a, int i, int j, int k, int l, double tau, double s)
 {
     double g, h;
     g       = a[i][j];
@@ -54,7 +56,8 @@ static inline void do_rotate(double** a, int i, int j, int k, int l, double tau,
     a[k][l] = h + s * (g - h * tau);
 }
 
-void jacobi(double** a, int n, double d[], double** v, int* nrot)
+template<typename MatrixType>
+static int jacobi(MatrixType a, const int n, double d[], MatrixType v)
 {
     int    j, i;
     int    iq, ip;
@@ -75,7 +78,7 @@ void jacobi(double** a, int n, double d[], double** v, int* nrot)
         b[ip] = d[ip] = a[ip][ip];
         z[ip]         = 0.0;
     }
-    *nrot = 0;
+    int nrot = 0;
     for (i = 1; i <= 50; i++)
     {
         sm = 0.0;
@@ -90,7 +93,7 @@ void jacobi(double** a, int n, double d[], double** v, int* nrot)
         {
             sfree(z);
             sfree(b);
-            return;
+            return nrot;
         }
         if (i < 4)
         {
@@ -151,7 +154,7 @@ void jacobi(double** a, int n, double d[], double** v, int* nrot)
                     {
                         do_rotate(v, j, ip, j, iq, tau, s);
                     }
-                    ++(*nrot);
+                    ++nrot;
                 }
             }
         }
@@ -163,6 +166,27 @@ void jacobi(double** a, int n, double d[], double** v, int* nrot)
         }
     }
     gmx_fatal(FARGS, "Error: Too many iterations in routine JACOBI\n");
+
+    return nrot;
+}
+
+void jacobi(double** a, const int numDimensions, double* eigenvalues, double** eigenvectors, int* numRotations)
+{
+    int numRot = jacobi(a, numDimensions, eigenvalues, eigenvectors);
+
+    if (numRotations)
+    {
+        *numRotations = numRot;
+    }
+}
+
+int jacobi(gmx::ArrayRef<gmx::DVec> a, gmx::ArrayRef<double> eigenvalues, gmx::ArrayRef<gmx::DVec> eigenvectors)
+{
+    GMX_RELEASE_ASSERT(gmx::ssize(a) == DIM, "Size should be 3");
+    GMX_RELEASE_ASSERT(gmx::ssize(eigenvalues) == DIM, "Size should be 3");
+    GMX_RELEASE_ASSERT(gmx::ssize(eigenvectors) == DIM, "Size should be 3");
+
+    return jacobi(a, DIM, eigenvalues.data(), eigenvectors);
 }
 
 int m_inv_gen(real* m, int n, real* minv)

diff --git a/src/gromacs/linearalgebra/nrjac.h b/src/gromacs/linearalgebra/nrjac.h
index 460c383fc2227957b1c0be024c5e2c08536a2999..82bbb51b12941fefc37f770dbc858961f6bcc737 100644 (file)
--- a/src/gromacs/linearalgebra/nrjac.h
+++ b/src/gromacs/linearalgebra/nrjac.h
@@ -3,7 +3,8 @@
  *
  * Copyright (c) 1991-2000, University of Groningen, The Netherlands.
  * Copyright (c) 2001-2004, The GROMACS development team.
- * Copyright (c) 2010,2014,2017,2018,2019, by the GROMACS development team, led by
+ * Copyright (c) 2010,2014,2017,2018,2019 by the GROMACS development team.
+ * Copyright (c) 2021, by the GROMACS development team, led by
  * Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
  * and including many others, as listed in the AUTHORS file in the
  * top-level source directory and at http://www.gromacs.org.
@@ -37,17 +38,29 @@
 #ifndef GMX_LINEARALGEBRA_NRJAC_H
 #define GMX_LINEARALGEBRA_NRJAC_H
 
+#include "gromacs/math/vectypes.h"
+#include "gromacs/utility/arrayref.h"
 #include "gromacs/utility/real.h"
 
-void jacobi(double** a, int n, double d[], double** v, int* nrot);
-/*
- * real   **omega = input matrix a[0..n-1][0..n-1] must be symmetric
- * int     natoms = number of rows and columns
- * real      NULL = d[0]..d[n-1] are the eigenvalues of a[][]
- * real       **v = v[0..n-1][0..n-1] the eigenvectors:
- *                                    v[i][j] is component i of vector j
- * int      *irot = number of jacobi rotations
+/* Diagonalizes a symmetric matrix
+ *
+ * \param[in,out] a           Input matrix a[0..n-1][0..n-1] must be symmetric, gets modified
+ * \param[in]  numDimensions  Number of rows and columns
+ * \param[out] eigenvalues    eigenvalues[0]..eigenvalues[n-1] are the eigenvalues of a
+ * \param[out] eigenvectors   v[0..n-1][0..n-1] the eigenvectors: v[i][j] is component i of vector j
+ * \param[out] numRotations   The number of jacobi rotations, can be nullptr
+ */
+void jacobi(double** a, int numDimensions, double* eigenvalues, double** eigenvectors, int* numRotations);
+
+/* Like jacobi above, but specialized for n=3
+ *
+ * \param[in,out] a  The symmetric matrix to diagonalize, size 3, note that the contents gets modified
+ * \param[out] eigenvalues  The eigenvalues, size 3
+ * \param[out] eigenvectors The eigenvectors, size 3
+
+ * Returns the number of jacobi rotations.
  */
+int jacobi(gmx::ArrayRef<gmx::DVec> a, gmx::ArrayRef<double> eigenvalues, gmx::ArrayRef<gmx::DVec> eigenvectors);
 
 int m_inv_gen(real* m, int n, real* minv);
 /* Produces minv, a generalized inverse of m, both stored as linear arrays.