diff --git a/reduction_tools/stream_stats.py b/reduction_tools/stream_stats.py index 1c062f1..594c573 100644 --- a/reduction_tools/stream_stats.py +++ b/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)