From 8640710c38b463f2bd3be37643634468f959936b Mon Sep 17 00:00:00 2001
From: Tong Jiang <tj.theochem@gmail.com>
Date: Tue, 12 Sep 2023 20:32:39 -0400
Subject: [PATCH] Enhance mps/mpo reproducibility; Print git hash for easier
 debugging; fix mp dump/load

---
 renormalizer/mps/backend.py      | 26 +++++++++++++++++++++-----
 renormalizer/mps/gs.py           |  9 ++++++---
 renormalizer/mps/mp.py           |  5 ++++-
 renormalizer/mps/symbolic_mpo.py |  7 ++++---
 4 files changed, 35 insertions(+), 12 deletions(-)

diff --git a/renormalizer/mps/backend.py b/renormalizer/mps/backend.py
index 06fa7233..739b3777 100644
--- a/renormalizer/mps/backend.py
+++ b/renormalizer/mps/backend.py
@@ -3,6 +3,7 @@
 import os
 import logging
 import random
+import subprocess
 
 import numpy as np
 
@@ -49,6 +50,13 @@ def try_import_cupy():
     logger.info(f"Using GPU: {GPU_ID}")
     return True, cp
 
+def get_git_commit_hash():
+    try:
+        commit_hash = subprocess.check_output(['git', 'rev-parse', 'HEAD']).strip().decode('utf-8')
+        return commit_hash
+    except subprocess.CalledProcessError:
+        return "Unknown"
+
 
 USE_GPU, xp = try_import_cupy()
 
@@ -56,17 +64,25 @@ def try_import_cupy():
 #USE_GPU = False
 #xp = np
 
+xpseed = 2019
+npseed = 9012
+randomseed = 1092
+
+xp.random.seed(xpseed)
+np.random.seed(npseed)
+random.seed(randomseed)
+
+
 if not USE_GPU:
     logger.info("Use NumPy as backend")
+    logger.info(f"numpy random seed is {npseed}")
     OE_BACKEND = "numpy"
 else:
     logger.info("Use CuPy as backend")
+    logger.info(f"cupy random seed is {xpseed}")
     OE_BACKEND = "cupy"
-
-
-xp.random.seed(2019)
-np.random.seed(9012)
-random.seed(1092)
+logger.info(f"random seed is {randomseed}")
+logger.info("Git Commit Hash: %s", get_git_commit_hash())
 
 
 class Backend:
diff --git a/renormalizer/mps/gs.py b/renormalizer/mps/gs.py
index 4e4828f1..ec35574c 100644
--- a/renormalizer/mps/gs.py
+++ b/renormalizer/mps/gs.py
@@ -383,10 +383,10 @@ def eigh_direct(
     nroots = mps.optimize_config.nroots
     if nroots == 1:
         e = w[0]
-        c = v[:, 0]
+        c = v[:, 0] / np.sign(np.max(v[:, 0]))
     else:
         e = w[:nroots]
-        c = [v[:, iroot] for iroot in range(min(nroots, v.shape[1]))]
+        c = [v[:, iroot] / np.sign(np.max(v[:, iroot])) for iroot in range(min(nroots, v.shape[1]))]
     return e, c
 
 
@@ -556,4 +556,7 @@ def hop(x):
     else:
         assert False
     logger.debug(f"use {algo}, HC hops: {count}")
-    return e, c
+    if nroots == 1:
+        return e, c/np.sign(np.max(c))
+    else:
+        return e, [ci/np.sign(np.max(ci)) for ci in c]
\ No newline at end of file
diff --git a/renormalizer/mps/mp.py b/renormalizer/mps/mp.py
index cd0bfd64..00d15e7e 100644
--- a/renormalizer/mps/mp.py
+++ b/renormalizer/mps/mp.py
@@ -49,7 +49,7 @@ def load(cls, model: Model, fname: str):
 
         mp.qn = []
         for i in range(nsites+1):
-            subqn = npload[f"subqn_{i}"].astype(int).tolist()    
+            subqn = npload[f"subqn_{i}"].astype(int).tolist()
             mp.qn.append(subqn)
 
         mp.qnidx = int(npload["qnidx"])
@@ -1084,6 +1084,9 @@ def dump(self, fname, other_attrs=None):
         arr = np.empty(len(qn), object)
         arr[:] = qn
         data_dict['qn'] = arr
+
+        for i in range(self.site_num+1):
+            data_dict[f"subqn_{i}"] = qn[i]
         
         try:
             np.savez(fname, **data_dict)
diff --git a/renormalizer/mps/symbolic_mpo.py b/renormalizer/mps/symbolic_mpo.py
index 0c4a5a4b..3f448427 100644
--- a/renormalizer/mps/symbolic_mpo.py
+++ b/renormalizer/mps/symbolic_mpo.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import logging
 import itertools
-from collections import namedtuple
+from collections import namedtuple, OrderedDict
 from typing import List, Set, Tuple, Dict
 
 import numpy as np
@@ -136,10 +136,11 @@ def construct_symbolic_mpo(table, factor, algo="Hopcroft-Karp"):
 
     # translate the symbolic operator table to an easy to manipulate numpy array
     table = np.array(table)
+
     # unique operators with DoF names taken into consideration
     # The inclusion of DoF names is necessary for multi-dof basis.
-    unique_op: Set[Op] = set(table.ravel())
-
+    unique_op = OrderedDict.fromkeys(table.ravel())
+    unique_op = list(unique_op.keys())
     # check the index of different operators could be represented with np.uint16
     assert len(unique_op) < max_uint16