From 17bed8185a73ddabcaf4be7e2379cdf9f2f614eb Mon Sep 17 00:00:00 2001
From: Hao Zhang <hzhangxyz@outlook.com>
Date: Thu, 9 Mar 2023 17:14:39 +0800
Subject: [PATCH] [tet.py] Add support user defined path of scalapack.

---
 CHANGELOG.org                                  |  3 ++-
 .../tetragono/sampling_lattice/gradient.py     |  4 +++-
 .../tetragono/sampling_lattice/observer.py     | 18 ++++++------------
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/CHANGELOG.org b/CHANGELOG.org
index 93f4e2aa0..4c2017795 100644
--- a/CHANGELOG.org
+++ b/CHANGELOG.org
@@ -8,7 +8,8 @@
 + *TAT.py*: Add =dtype= and =btype= class member for =Tensor= object, which makes it easier to communicate with numpy.
 + *scalapack.py*: Add a python wrapper for scalapack.
 *** Changed
-+ *tetragono*: Use the =PyScalapack= to speed up min-SR method.
++ *tetragono*: Use the =PyScalapack= to speed up min-SR method. User need to specify the path of scalapack dynamic link
+  libraries by parameter =scalapack_libraries= for =gm_run= when =natural_gradient_by_direct_pseudo_inverse= enabled.
 + *TAT.py*: Change the module alias name convension, =float= and =complex= without bytes specified would be considered
   as double precision now.
 *** Deprecated
diff --git a/tetragono/tetragono/sampling_lattice/gradient.py b/tetragono/tetragono/sampling_lattice/gradient.py
index e308bce64..961ab0e6b 100644
--- a/tetragono/tetragono/sampling_lattice/gradient.py
+++ b/tetragono/tetragono/sampling_lattice/gradient.py
@@ -123,6 +123,7 @@ def gradient_descent(
         metric_inverse_epsilon=0.01,
         cache_natural_delta=None,
         use_natural_gradient_by_direct_pseudo_inverse=False,
+        scalapack_libraries=None,
         natural_gradient_r_pinv=1e-12,
         natural_gradient_a_pinv=0,
         # About gauge fixing
@@ -270,7 +271,8 @@ def restrict(configuration, replacement=None):
                 if use_natural_gradient:
                     if use_natural_gradient_by_direct_pseudo_inverse:
                         grad = observer.natural_gradient_by_direct_pseudo_inverse(natural_gradient_r_pinv,
-                                                                                  natural_gradient_a_pinv)
+                                                                                  natural_gradient_a_pinv,
+                                                                                  scalapack_libraries.split(","))
                     else:
                         grad = observer.natural_gradient_by_conjugate_gradient(conjugate_gradient_method_step,
                                                                                conjugate_gradient_method_error,
diff --git a/tetragono/tetragono/sampling_lattice/observer.py b/tetragono/tetragono/sampling_lattice/observer.py
index 8c9986db6..4dec553b3 100644
--- a/tetragono/tetragono/sampling_lattice/observer.py
+++ b/tetragono/tetragono/sampling_lattice/observer.py
@@ -667,7 +667,7 @@ def _array_to_delta(self, array):
                 index += size
         return result
 
-    def natural_gradient_by_direct_pseudo_inverse(self, r_pinv, a_pinv):
+    def natural_gradient_by_direct_pseudo_inverse(self, r_pinv, a_pinv, libraries):
         """
         Get the energy natural gradient for every tensor.
 
@@ -675,6 +675,8 @@ def natural_gradient_by_direct_pseudo_inverse(self, r_pinv, a_pinv):
         ----------
         r_pinv, a_pinv : float
             Parameters control how to calculate pseudo inverse.
+        libraries : list[str]
+            The dynamic link libraries containing scalapack functions.
 
         Returns
         -------
@@ -714,22 +716,14 @@ def natural_gradient_by_direct_pseudo_inverse(self, r_pinv, a_pinv):
 
         total_n_s = int(self._count)
         result = self._array_to_delta(
-            self._pseudo_inverse_kernel(
-                Delta,
-                Energy,
-                r_pinv,
-                a_pinv,
-                total_n_s,
-                dtype,
-                btype,
-            ))
+            self._pseudo_inverse_kernel(Delta, Energy, r_pinv, a_pinv, total_n_s, dtype, btype, libraries))
 
         showln("calculate natural gradient done")
         return result * 2
 
     @staticmethod
-    def _pseudo_inverse_kernel(Delta, Energy, r_pinv, a_pinv, total_n_s, dtype, btype):
-        scalapack = Scalapack("/usr/lib/libscalapack.so")
+    def _pseudo_inverse_kernel(Delta, Energy, r_pinv, a_pinv, total_n_s, dtype, btype, libraries):
+        scalapack = Scalapack(*libraries)
 
         with scalapack(b'C', -1, 1) as context:
             n_s, n_p = Delta.shape