From 145299d52d048e9b765b5282cf8b864973e63e5a Mon Sep 17 00:00:00 2001
From: Nathan W Johnson
 <115745236+nathanjohnsongithub@users.noreply.github.com>
Date: Thu, 29 Jun 2023 11:22:30 -0500
Subject: [PATCH] Program no longer assumes where integrals start in fcidump
 (#63)

* only master loads wave function and integrals

* added # because it was missing

* Program no longer assumes where integrals start in fcidump

* Program no longer assumes where integrals start in fcidump: updated
---
 qe/io.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/qe/io.py b/qe/io.py
index ff90ce6..c5ee2b2 100644
--- a/qe/io.py
+++ b/qe/io.py
@@ -9,6 +9,7 @@
 from collections import defaultdict
 from qe.integral_indexing_utils import compound_idx4
 import math
+from itertools import takewhile
 
 #   _____      _ _   _       _ _          _   _
 #  |_   _|    (_) | (_)     | (_)        | | (_)
@@ -18,6 +19,12 @@
 #   \___/_| |_|_|\__|_|\__,_|_|_/___\__,_|\__|_|\___/|_| |_|
 
 
+# Takes in the raw lines and manipulates them if needed based on using zip or not.
+def manipulate_line(raw_line, used_zip):
+    line = raw_line.decode("utf-8") if used_zip else raw_line
+    return line.strip()
+
+
 # ~
 # Integrals of the Hamiltonian over molecular orbitals
 # ~
@@ -41,6 +48,9 @@ def load_integrals(
         for i in glob.glob(fcidump_path):
             print(i)
 
+    # Add used_zip boolean so we know if we need to decode the lines.
+    used_zip = True
+
     # Use an iterator to avoid storing everything in memory twice.
     if fcidump_path.split(".")[-1] == "gz":
         import gzip
@@ -52,13 +62,15 @@ def load_integrals(
         f = bz2.open(fcidump_path)
     else:
         f = open(fcidump_path)
+        used_zip = False
 
     # Only non-zero integrals are stored in the fci_dump.
     # Hence we use a defaultdict to handle the sparsity
     n_orb = int(next(f).split()[2])
 
-    for _ in range(3):
-        next(f)
+    # Using takewhile we 'take lines from the file' while '/' has not been found
+    # Needed so we can start reading data on the integrals.
+    lines = list(takewhile(lambda line: "/" not in manipulate_line(line, used_zip), f))
 
     d_one_e_integral = defaultdict(int)
     d_two_e_integral = defaultdict(int)