func. not check - added wilson to plot + tanh curve fits + relabelled axis

reduction_tools/partialator.py

@@ -16,12 +16,14 @@ python partialator.py -s <path-to-stream-file>
                        -b number of resolution bins - must be > 20
                        -r high-res limt. Needs a default. Default set to 1.3
                        -a max-adu. Default = 12000
+                       -R ra reservation name if available
 
 # output
 - scaled/merged files
 - an mtz file
 - useful plots
 - useful summerized .dat files
+- log file of output
 """
 
 # modules
@@ -43,12 +45,15 @@ def submit_job( job_file, reservation ):
     # submit the job
     if reservation:
         print( "using a ra beamtime reservation = {0}".format( reservation ) )
+        logger.info( "using ra reservation to process data = {0}".format( reservation ) )
         submit_cmd = [ "sbatch", "--reservation={0}".format( reservation ), "--cpus-per-task=32", "--" , job_file ]
     else:
         submit_cmd = [ "sbatch", "--cpus-per-task=32", "--" , job_file ]
-    
+    logger.info( "using slurm command = {0}".format( submit_cmd ) )
+
     try:
         job_output = subprocess.check_output( submit_cmd )
+        logger.info( "submited job = {0}".format( job_output ) )
     except subprocess.CalledProcessError as e:
         print( "please give the correct ra reservation or remove the -R from the arguements" )
         exit()
@@ -61,17 +66,17 @@ def submit_job( job_file, reservation ):
 
 def wait_for_jobs( job_ids, total_jobs ):
 
-    with tqdm(total=total_jobs, desc="Jobs Completed", unit="job") as pbar:
+    with tqdm( total=total_jobs, desc="Jobs Completed", unit="job" ) as pbar:
         while job_ids:
             completed_jobs = set()
             for job_id in job_ids:
-                status_cmd = [ "squeue", "-h", "-j", str(job_id) ]
+                status_cmd = [ "squeue", "-h", "-j", str( job_id ) ]
-                status = subprocess.check_output(status_cmd)
+                status = subprocess.check_output( status_cmd )
                 if not status:
-                    completed_jobs.add(job_id)
+                    completed_jobs.add( job_id )
-                    pbar.update(1)
+                    pbar.update( 1 )
-            job_ids.difference_update(completed_jobs)
+            job_ids.difference_update( completed_jobs )
-            time.sleep(2)
+            time.sleep( 2 )
 
 def run_partialator( proc_dir, name, stream, pointgroup, model, iterations, cell, shells, part_h_res, adu ):
 
@@ -102,6 +107,11 @@ def run_partialator( proc_dir, name, stream, pointgroup, model, iterations, cell
     # make file executable
     subprocess.call( [ "chmod", "+x", "{0}".format( part_run_file ) ] )
 
+    # add partialator script to log
+    part_input = open( part_run_file, "r" )
+    logger.info( "partialator input file =\n{0}".format( part_input.read() ) )
+    part_input.close()
+
     # return partialator file name
     return part_run_file
 
@@ -113,7 +123,7 @@ def make_process_dir( dir ):
         if e.errno != errno.EEXIST:
             raise
 
-def summary_stats( cc_dat, ccstar_dat, mult_dat, rsplit_dat ):
+def summary_stats( cc_dat, ccstar_dat, mult_dat, rsplit_dat, wilson_dat ):
 
     # read all files into pd
     # function to sort out different column names
@@ -129,37 +139,33 @@ def summary_stats( cc_dat, ccstar_dat, mult_dat, rsplit_dat ):
                     "mult", "snr", "I", "d", "min", "max" ]
         elif var == "rsplit":
             cols = [ "d(nm)", "rsplit", "nref", "d", "min", "max"  ]
+        elif var == "wilson":
+            cols = [ "bin", "s2", "d", "lnI", "nref" ]
         
         df = pd.read_csv( dat, names=cols, skiprows=1, sep="\s+" )
 
-        print(df)
         return df
-        
 
     # make df
     cc_df = read_dat( cc_dat, "cc" )
     ccstar_df = read_dat( ccstar_dat, "ccstar" )
     mult_df = read_dat( mult_dat, "mult" )
     rsplit_df = read_dat( rsplit_dat, "rsplit" )
+    wilson_df = read_dat( wilson_dat, "wilson" )
 
     # remove unwanted cols
     cc_df = cc_df[ [ "cc" ] ]
     ccstar_df = ccstar_df[ [ "ccstar" ] ]
     rsplit_df = rsplit_df[ [ "rsplit" ] ]
+    wilson_df = wilson_df[ [ "lnI" ] ]
     
     # merge dfs
-    stats_df = pd.concat( [ mult_df, cc_df, ccstar_df, rsplit_df ], axis=1, join="inner" )
+    stats_df = pd.concat( [ mult_df, cc_df, ccstar_df, rsplit_df, wilson_df ], axis=1, join="inner" )
 
     # make 1/d, 1/d^2 column
     stats_df[ "1_d" ] = 1 / stats_df.d
     stats_df[ "1_d2" ] = 1 / stats_df.d**2
 
-    # reorder cols
-    stats_df = stats_df[ [  "1_d", "1_d2", "d", "min",
-                            "max", "nref", "poss",
-                            "comp", "obs", "mult",
-                            "snr", "I", "cc", "ccstar", "rsplit" ] ]
-
     # change nan to 0 
     stats_df = stats_df.fillna(0)
 
@@ -170,7 +176,7 @@ def get_metric( d2_series, cc_series, cut_off ):
     # Define the tanh function from scitbx
     def tanh(x, r, s0):
         z = (x - s0)/r
-        return 0.5 * (1 - np.tanh(z))
+        return 0.5 * ( 1 - np.tanh(z) )
 
     def arctanh( y, r, s0 ):
         return r * np.arctanh( 1 - 2*y ) + s0
@@ -183,13 +189,15 @@ def get_metric( d2_series, cc_series, cut_off ):
 
     # calculate cut-off point 
     cc_stat = arctanh( cut_off, r, s0 )
-
     # covert back from 1/d2 to d
     cc_stat = np.sqrt( ( 1 / cc_stat ) )
 
-    return cc_stat
+    # get curve for plotting
+    cc_tanh = tanh( d2_series, r, s0 )
 
-def summary_fig( stats_df ):
+    return cc_stat, cc_tanh
+
+def summary_fig( stats_df, cc_tanh, ccstar_tanh ):
 
     # plot results
     cc_fig, axs = plt.subplots(2, 2)
@@ -198,39 +206,47 @@ def summary_fig( stats_df ):
     # cc plot
     color = "tab:red"
     axs[0,0].set_xlabel( "1/d (1/A)" )
-    axs[0,0].set_ylabel("CC" )
+    axs[0,0].set_ylabel( "CC" )
     axs[0,0].set_ylim( 0, 1 )
-    axs[0,0].axhline(y = 0.3, color="black", linestyle = "dashed")
+    axs[0,0].axhline( y = 0.3, color="black", linestyle = "dashed" )
-    axs[0,0].plot(stats_df[ "1_d" ], stats_df.cc, color=color)
+    # plot cc
-    axs[0,0].tick_params(axis="y", labelcolor=color)
+    axs[0,0].plot( stats_df[ "1_d" ], stats_df.cc, color=color )
+    # plot fit
+    axs[0,0].plot( stats_df[ "1_d2" ], cc_tanh, color="tab:grey", linestyle = "dashed" )
+    axs[0,0].xticks( stats_df[ "1_d2" ].iloc[::5, :], stats_df[ "d" ].iloc[::5, :] )
+    axs[0,0].tick_params( axis="y", labelcolor=color )
 
     # cc* plot
     color = "tab:blue"
     axs[0,1].set_xlabel( "1/d (1/A)" )
-    axs[0,1].set_ylabel("CC*", color=color)
+    axs[0,1].set_ylabel( "CC*", color=color )
     axs[0,1].set_ylim( 0, 1 )
-    axs[0,1].axhline(y = 0.7, color="black", linestyle = "dashed")
+    axs[0,1].axhline( y = 0.7, color="black", linestyle = "dashed" )
-    axs[0,1].plot(stats_df[ "1_d" ], stats_df.ccstar, color=color)
+    axs[0,1].plot( stats_df[ "1_d" ], stats_df.ccstar, color=color )
-    axs[0,1].tick_params(axis='y', labelcolor=color)
+    # plot fit
+    axs[0,0].plot( stats_df[ "1_d2" ], ccstar_tanh, color="tab:grey", linestyle = "dashed" )
+    axs[0,0].xticks( stats_df[ "1_d2" ].iloc[::5, :], stats_df[ "d" ].iloc[::5, :] )
+    axs[0,1].tick_params( axis='y', labelcolor=color )
 
     # rsplit plot
     color = "tab:green"
     axs[1,0].set_xlabel( "1/d (1/A)" )
-    axs[1,0].set_ylabel("Rsplit", color=color)
+    axs[1,0].set_ylabel( "Rsplit", color=color )
-    axs[1,0].plot(stats_df[ "1_d" ], stats_df.rsplit, color=color)
+    axs[1,0].plot( stats_df[ "1_d" ], stats_df.rsplit, color=color )
-    axs[1,0].tick_params(axis='y', labelcolor=color)
+    axs[1,0].tick_params( axis='y', labelcolor=color )
 
 
-    # rsplit plot
+    # wilson plot
     color = "tab:purple"
-    axs[1,1].set_xlabel( "1/d (1/A)" )
+    axs[1,1].set_xlabel( "d (A)" )
-    axs[1,1].set_ylabel("Multiplicity", color=color)
+    axs[1,1].set_ylabel( "lnI", color=color )
-    axs[1,1].plot(stats_df[ "1_d" ], stats_df.mult, color=color)
+    axs[1,1].plot( stats_df[ "1_d2" ], stats_df.lnI, color=color )
-    axs[1,1].tick_params(axis='y', labelcolor=color)
+#    axs[1,1].invert_xaxis()
+    axs[1,1].tick_params( axis='y', labelcolor=color )
 
     # save figure
     plt.tight_layout()
-    plt.savefig("plots.png")
+    plt.savefig( "plots.png" )
 
 def get_mean_cell( stream ):
 
@@ -287,23 +303,26 @@ def main( cwd, name, stream, pointgroup, model, iterations, cell, shells, part_h
     ccstar_dat = "ccstar.dat"
     mult_dat = "mult.dat"
     rsplit_dat = "Rsplit.dat"
+    wilson_dat = "wilson.dat"
 
     # make summary data table
-    stats_df = summary_stats( cc_dat, ccstar_dat, mult_dat, rsplit_dat )
+    stats_df = summary_stats( cc_dat, ccstar_dat, mult_dat, rsplit_dat, wilson_dat  )
+    logger.info( "stats table from .dat file =\n{0}".format( stats_df ) )
     print_df = stats_df[ [  "1_d", "d", "min",
                             "max", "nref", "poss",
                             "comp", "obs", "mult",
-                            "snr", "I", "rsplit", "cc", "ccstar"] ]
+                            "snr", "I", "rsplit", "cc", "ccstar" ] ]
     print_df.to_csv( "summary_table.csv", sep="\t", index=False )
 
+
     # calculate cc metrics
-    cc_cut = get_metric( stats_df[ "1_d2" ], stats_df.cc, 0.3 )
+    cc_cut, cc_tanh = get_metric( stats_df[ "1_d2" ], stats_df.cc, 0.3 )
-    ccstar_cut = get_metric( stats_df[ "1_d2" ], stats_df.ccstar, 0.7 )
+    ccstar_cut, ccstar_tanh = get_metric( stats_df[ "1_d2" ], stats_df.ccstar, 0.7 )
     print( "resolution at CC0.5 at 0.3 = {0}".format( cc_cut ) )
     print( "resolution at CC* at 0.7 = {0}".format( ccstar_cut ) )
 
     # show plots
-    summary_fig( stats_df )
+    summary_fig( stats_df, cc_tanh, ccstar_tanh )
 
     # move back to top dir
     os.chdir( cwd )