fixed crash bug. rewrite of .stream open. couldn't open large stream files.

reduction_tools/stream_stats.py

@@ -19,85 +19,78 @@ import pandas as pd
 import regex as re
 import sys
 
-def count_chunks( stream ):
 
-    # get number of chunks
-    # example - ----- Begin chunk -----
-    # count them
-    try:
-        pattern = r"-----\sBegin\schunk\s-----"
-        chunks = re.findall( pattern, stream )
-        if AttributeError:
-            return len( chunks )
-    except AttributeError:
-        return np.nan
-
-def scrub_cells( stream ):
+def scrub_cells( line ):
 
     # get uc values from stream file
     # example - Cell parameters 7.71784 7.78870 3.75250 nm, 90.19135 90.77553 90.19243 deg
-    # scrub clen and return - else nan
-    try:
-        pattern = r"Cell\sparameters\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\snm,\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\sdeg"
-        cell_lst = re.findall( pattern, stream )
-        xtals = len( cell_lst )
-        if AttributeError:
-            return cell_lst, xtals
-    except AttributeError:
-        return np.nan
+    pattern = r"Cell\sparameters\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\snm,\s(\d+\.\d+)\s(\d+\.\d+)\s(\d+\.\d+)\sdeg"
+    a = re.search( pattern, line ).group(1)
+    b = re.search( pattern, line ).group(2)
+    c = re.search( pattern, line ).group(3)
+    alpha = re.search( pattern, line ).group(4)
+    beta = re.search( pattern, line ).group(5)
+    gamma = re.search( pattern, line ).group(6)
+
+    return [ a, b, c, alpha, beta, gamma ]
 
 def scrub_res( stream ):
 
     # get diffraction limit
     # example - diffraction_resolution_limit = 4.07 nm^-1 or 2.46 A
-    # scrub res_lst or return np.nan
-    try:
-        pattern = r"diffraction_resolution_limit\s=\s\d\.\d+\snm\^-1\sor\s(\d+\.\d+)\sA"
-        res_lst = re.findall( pattern, stream )
-        if AttributeError:
-            return res_lst
-    except AttributeError:
-        return np.nan
+    pattern = r"diffraction_resolution_limit\s=\s\d\.\d+\snm\^-1\sor\s(\d+\.\d+)\sA"
+    res = re.search( pattern, stream ).group(1)
+    return res
 
 def scrub_obs( stream ):
 
     # get number of reflections
     # example - num_reflections = 308
-    # scrub reflections or return np.nan
-    try:
-        pattern = r"num_reflections\s=\s(\d+)"
-        obs_lst = re.findall( pattern, stream )
-        if AttributeError:
-            return obs_lst
-    except AttributeError:
-        return np.nan
+    pattern = r"num_reflections\s=\s(\d+)"
+    obs = re.search( pattern, stream ).group(1)
+    return obs
 
 def main( stream_pwd ):
 
-    print( "reading stream file" )
-    # open stream file
-    stream = open( stream_pwd, "r" ).read()
-    print( "done" )
+    chunks = 0
+    xtals = 0
+    cells = []
+    obs_list = []
+    res_list = []
 
     print( "scrubing data" )
-    # get total number chunks
-    chunks = count_chunks( stream )
+    # open stream file
+    with open( stream_pwd ) as stream:
+        for line in stream:
+            
+            # count chunks
+            if line.startswith( "----- Begin chunk -----" ):
+                chunks = chunks + 1
 
-    # get list of cells
-    cell_lst, xtals = scrub_cells( stream )
+            # get cell
+            if line.startswith( "Cell parameters" ):
+                cell = scrub_cells( line )
+                cells.append( cell )
+                xtals = xtals + 1
 
-    # get list of cells
-    res_lst = scrub_res( stream )
+            # get res
+            if line.startswith( "diffraction_resolution_limit" ):
+                res = scrub_res( line )
+                res_list.append( res )
 
-    # get list of cells
-    obs_lst = scrub_obs( stream )
-    print( "done" )
+            # get obs
+            if line.startswith( "num_reflections" ):
+                obs = scrub_obs( line )
+                obs_list.append( obs )
+
+            if chunks % 1000 == 0:
+                print( "scrubbed {0} chunks".format( chunks ), end='\r' )
 
     # res_df
     cols = [ "a", "b", "c", "alpha", "beta", "gamma" ]
-    df = pd.DataFrame( cell_lst, columns=cols )
-    df[ "resolution" ] = res_lst
-    df[ "obs" ] = obs_lst
+    df = pd.DataFrame( cells, columns=cols )
+    df[ "resolution" ] = res_list
+    df[ "obs" ] = obs_list
 
     # convert all to floats
     df = df.astype(float)