528 lines
18 KiB
Python
528 lines
18 KiB
Python
#!/usr/bin/python
|
|
|
|
# author J.Beale
|
|
|
|
"""
|
|
# aim
|
|
to process a batch of data very fast by splitting it into a number of chunks and submitting
|
|
these jobs separately to the cluster - but now all with the ability to change crystfel parameters
|
|
from the command line
|
|
|
|
# usage
|
|
python crystfel_split.py -l <path-to-list-file>
|
|
-k <chunk-size> -default 1000
|
|
-g <path-to-geom-file>
|
|
-c <path-to-cell-file>
|
|
-n <job-name> -default split
|
|
-t crystfel threshold -default 10
|
|
-s crystfel min-snr -default 5
|
|
-i crystfel int-radius -default 3,5,9
|
|
-m crystfel multi or no-multi (True/False) -default False (no-multi)
|
|
-r crystfel retry or no-retry (True/False) -default False (no-retry)
|
|
-x crystfel min-pix-count -default 2
|
|
|
|
# output
|
|
a series of stream files from crystfel in the current working directory
|
|
"""
|
|
|
|
# modules
|
|
import pandas as pd
|
|
import numpy as np
|
|
import subprocess
|
|
import os, errno
|
|
import time
|
|
import argparse
|
|
from tqdm import tqdm
|
|
import regex as re
|
|
from loguru import logger
|
|
|
|
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:
|
|
logger.debug( "count_chunks error" )
|
|
return np.nan
|
|
|
|
def scrub_cells( stream ):
|
|
|
|
# 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:
|
|
logger.debug( "scrub_cells error" )
|
|
return np.nan
|
|
|
|
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:
|
|
logger.debug( "scrub_res error" )
|
|
return np.nan
|
|
|
|
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:
|
|
logger.debug( "scrub_obs error" )
|
|
return np.nan
|
|
|
|
def calculate_stats( stream_pwd ):
|
|
|
|
# open stream file
|
|
stream = open( stream_pwd, "r" ).read()
|
|
|
|
# get total number chunks
|
|
chunks = count_chunks( stream )
|
|
|
|
# get list of cells
|
|
cell_lst, xtals = scrub_cells( stream )
|
|
|
|
# get list of cells
|
|
res_lst = scrub_res( stream )
|
|
|
|
# get list of cells
|
|
obs_lst = scrub_obs( stream )
|
|
|
|
# res_df
|
|
cols = [ "a", "b", "c", "alpha", "beta", "gamma" ]
|
|
df = pd.DataFrame( cell_lst, columns=cols )
|
|
df[ "resolution" ] = res_lst
|
|
df[ "obs" ] = obs_lst
|
|
|
|
# convert all to floats
|
|
df = df.astype(float)
|
|
|
|
return df, xtals, chunks
|
|
|
|
def h5_split( lst, chunk_size ):
|
|
|
|
# read h5.lst - note - removes // from image column
|
|
# scrub file name
|
|
lst_name = os.path.basename( lst )
|
|
|
|
cols = [ "h5", "image" ]
|
|
df = pd.read_csv( lst, sep="\s//", engine="python", names=cols )
|
|
|
|
# re-add // to image columm and drop other columns
|
|
df[ "h5_path" ] = df.h5 + " //" + df.image.astype( str )
|
|
df = df[ [ "h5_path" ] ]
|
|
|
|
# split df into a lst
|
|
list_df = [df[i:i + chunk_size] for i in range( 0, len(df), chunk_size)]
|
|
|
|
return list_df
|
|
|
|
def write_crystfel_run( proc_dir, name, chunk, chunk_lst_file,
|
|
geom_file, cell_file, indexer, peakfinder,
|
|
integrator, tolerance, threshold, min_snr,
|
|
int_rad, multi, retry, min_pix, bg_rad, min_res ):
|
|
|
|
# stream file name
|
|
stream_file = "{0}_{1}.stream".format( name, chunk )
|
|
|
|
# crystfel file name
|
|
cryst_run_file = "{0}/{1}_{2}.sh".format( proc_dir, name, chunk )
|
|
# write file
|
|
run_sh = open( cryst_run_file, "w" )
|
|
run_sh.write( "#!/bin/sh\n\n" )
|
|
run_sh.write( "module purge\n" )
|
|
run_sh.write( "module use MX unstable\n" )
|
|
run_sh.write( "module load crystfel/0.10.2-rhel8\n" )
|
|
run_sh.write( "indexamajig -i {0} \\\n".format( chunk_lst_file ) )
|
|
run_sh.write( " --output={0} \\\n".format( stream_file ) )
|
|
run_sh.write( " --geometry={0} \\\n".format( geom_file ) )
|
|
run_sh.write( " --pdb={0} \\\n".format( cell_file ) )
|
|
run_sh.write( " --indexing={0} \\\n".format( indexer ) )
|
|
run_sh.write( " --peaks={0} \\\n".format( peakfinder ) )
|
|
run_sh.write( " --integration={0} \\\n".format( integrator ) )
|
|
run_sh.write( " --tolerance={0},{1},{2},{3},{4},{5} \\\n".format( tolerance[0], tolerance[1], tolerance[2],
|
|
tolerance[3], tolerance[4], tolerance[5] ) )
|
|
run_sh.write( " --threshold={0} \\\n".format( threshold ) )
|
|
run_sh.write( " --min-snr={0} \\\n".format( min_snr ) )
|
|
run_sh.write( " --int-radius={0},{1},{2} \\\n".format( int_rad[0], int_rad[1], int_rad[2] ) )
|
|
run_sh.write( " -j 32 \\\n" )
|
|
run_sh.write( " --{0} \\\n".format( multi ) )
|
|
run_sh.write( " --check-peaks \\\n" )
|
|
run_sh.write( " --{0} \\\n".format( retry ) )
|
|
run_sh.write( " --max-res=3000 \\\n" )
|
|
run_sh.write( " --min-pix-count={0} \\\n".format( min_pix ) )
|
|
run_sh.write( " --local-bg-radius={0} \\\n".format( bg_rad ) )
|
|
run_sh.write( " --min-res={0}".format( min_res ) )
|
|
run_sh.close()
|
|
|
|
# make file executable
|
|
subprocess.call( [ "chmod", "+x", "{0}".format( cryst_run_file ) ] )
|
|
|
|
# return crystfel file name
|
|
return cryst_run_file, stream_file
|
|
|
|
def make_process_dir( proc_dir ):
|
|
# make process directory
|
|
try:
|
|
os.makedirs( proc_dir )
|
|
except OSError as e:
|
|
if e.errno != errno.EEXIST:
|
|
logger.debug( "making directory error" )
|
|
raise
|
|
|
|
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", "-p", "hour", "--reservation={0}".format( reservation ), "--cpus-per-task=32", "--" , job_file ]
|
|
else:
|
|
submit_cmd = [ "sbatch", "-p", "hour", "--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 -v from the arguements" )
|
|
exit()
|
|
|
|
# scrub job id from - example Submitted batch job 742403
|
|
pattern = r"Submitted batch job (\d+)"
|
|
job_id = re.search( pattern, job_output.decode().strip() ).group(1)
|
|
|
|
return int( job_id )
|
|
|
|
def wait_for_jobs( job_ids, total_jobs ):
|
|
|
|
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 = subprocess.check_output(status_cmd)
|
|
if not status:
|
|
completed_jobs.add(job_id)
|
|
pbar.update(1)
|
|
job_ids.difference_update(completed_jobs)
|
|
time.sleep(2)
|
|
|
|
def run_splits( list_df, cwd, name, geom_file, cell_file,
|
|
indexer, peakfinder, integrator, tolerance, threshold,
|
|
min_snr, int_rad, multi, retry, min_pix, bg_rad,
|
|
min_res, reservation ):
|
|
|
|
# set chunk counter
|
|
chunk = 0
|
|
|
|
# submitted job set
|
|
submitted_job_ids = set()
|
|
|
|
# stream file list
|
|
stream_lst = []
|
|
|
|
for chunk_lst in list_df:
|
|
|
|
logger.info( "chunk {0} = {1} images".format( chunk, len( chunk_lst ) ) )
|
|
# define process directory
|
|
proc_dir = "{0}/{1}/{1}_{2}".format( cwd, name, chunk )
|
|
|
|
# make process directory
|
|
make_process_dir(proc_dir)
|
|
|
|
# move to process directory
|
|
os.chdir( proc_dir )
|
|
|
|
# write list to file
|
|
chunk_lst_file = "{0}/{1}_{2}.lst".format( proc_dir, name, chunk )
|
|
chunk_lst.to_csv( chunk_lst_file, index=False, header=False )
|
|
|
|
# write crystfel file and append path to list
|
|
cryst_run_file, stream_file = write_crystfel_run( proc_dir, name, chunk, chunk_lst_file,
|
|
geom_file, cell_file, indexer, peakfinder,
|
|
integrator, tolerance, threshold, min_snr,
|
|
int_rad, multi, retry, min_pix, bg_rad, min_res )
|
|
stream_lst.append( "{0}/{1}".format( proc_dir, stream_file ) )
|
|
|
|
# submit jobs
|
|
job_id = submit_job( cryst_run_file, reservation )
|
|
submitted_job_ids.add( job_id )
|
|
|
|
# increase chunk counter
|
|
chunk = chunk +1
|
|
|
|
# move back to top dir
|
|
os.chdir( cwd )
|
|
|
|
return submitted_job_ids, chunk, stream_lst
|
|
|
|
def main( cwd, name, lst, chunk_size, geom_file, cell_file,
|
|
indexer, peakfinder, integrator, tolerance, threshold,
|
|
min_snr, int_rad, multi, retry, min_pix, bg_rad,
|
|
min_res, reservation ):
|
|
|
|
print( "reading SwissFEL lst file" )
|
|
print( "creating {0} image chunks of lst".format( chunk_size ) )
|
|
list_df = h5_split( lst, chunk_size )
|
|
print( "done" )
|
|
|
|
# run crystfel runs on individual splits
|
|
print( "submitting jobs to cluster" )
|
|
submitted_job_ids, chunk, stream_lst = run_splits( list_df, cwd, name, geom_file, cell_file,
|
|
indexer, peakfinder, integrator, tolerance, threshold,
|
|
min_snr, int_rad, multi, retry, min_pix, bg_rad,
|
|
min_res, reservation )
|
|
|
|
# monitor progress of jobs
|
|
time.sleep( 30 )
|
|
wait_for_jobs( submitted_job_ids, chunk )
|
|
print( "done" )
|
|
|
|
# make composite .stream file
|
|
output_file = "{0}.stream".format( name )
|
|
|
|
print( "concatenating .streams from separate runs." )
|
|
try:
|
|
# Open the output file in 'append' mode
|
|
with open(output_file, "a") as output:
|
|
for file_name in stream_lst:
|
|
try:
|
|
with open(file_name, "r") as input_file:
|
|
# Read the contents of the input file and append to the output file
|
|
output.write(input_file.read())
|
|
except FileNotFoundError:
|
|
logger.debug(f"File {file_name} not found. Skipping.")
|
|
except IOError as e:
|
|
logger.debug(f"An error occurred while appending files: {e}")
|
|
|
|
print( "done" )
|
|
|
|
df, xtals, chunks = calculate_stats( output_file )
|
|
|
|
# stats
|
|
index_rate = round( xtals/chunks*100, 2 )
|
|
mean_res, std_res = round( df.resolution.mean(), 2 ), round( df.resolution.std(), 2 )
|
|
median_res = df.resolution.median()
|
|
mean_obs, std_obs = round( df.obs.mean(), 2 ), round( df.obs.std(), 2)
|
|
mean_a, std_a = round( df.a.mean()*10, 2 ), round( df.a.std()*10, 2 )
|
|
mean_b, std_b = round( df.b.mean()*10, 2 ), round( df.b.std()*10, 2 )
|
|
mean_c, std_c = round( df.c.mean()*10, 2 ), round( df.c.std()*10, 2 )
|
|
mean_alpha, std_alpha = round( df.alpha.mean(), 2 ), round( df.alpha.std(), 2 )
|
|
mean_beta, std_beta = round(df.beta.mean(), 2 ), round( df.beta.std(), 2 )
|
|
mean_gamma, std_gamma = round( df.gamma.mean(), 2 ), round( df.gamma.std(), 2 )
|
|
|
|
logger.info( "image = {0}".format( chunks ) )
|
|
logger.info( "crystals = {0}".format( xtals ) )
|
|
logger.info( "indexing rate = {0} %".format( index_rate ) )
|
|
logger.info( "mean resolution = {0} +/- {1} A".format( mean_res, std_res ) )
|
|
logger.info( "median resolution = {0} A".format( median_res ) )
|
|
logger.info( "mean observations = {0} +/- {1}".format( mean_obs, std_obs ) )
|
|
logger.info( "mean a = {0} +/- {1} A".format( mean_a, std_a ) )
|
|
logger.info( "mean b = {0} +/- {1} A".format( mean_b, std_b ) )
|
|
logger.info( "mean c = {0} +/- {1} A".format( mean_c, std_c ) )
|
|
logger.info( "mean alpha = {0} +/- {1} A".format( mean_alpha, std_alpha ) )
|
|
logger.info( "mean beta = {0} +/- {1} A".format( mean_beta, std_beta ) )
|
|
logger.info( "mean gamma = {0} +/- {1} A".format( mean_gamma, std_gamma ) )
|
|
|
|
print( "printing stats" )
|
|
print( "image = {0}".format( chunks ) )
|
|
print( "crystals = {0}".format( xtals ) )
|
|
print( "indexing rate = {0} %".format( index_rate ) )
|
|
print( "mean resolution = {0} +/- {1} A".format( mean_res, std_res ) )
|
|
print( "median resolution = {0} A".format( median_res ) )
|
|
print( "mean observations = {0} +/- {1}".format( mean_obs, std_obs ) )
|
|
print( "mean a = {0} +/- {1} A".format( mean_a, std_a ) )
|
|
print( "mean b = {0} +/- {1} A".format( mean_b, std_b ) )
|
|
print( "mean c = {0} +/- {1} A".format( mean_c, std_c ) )
|
|
print( "mean alpha = {0} +/- {1} deg".format( mean_alpha, std_alpha ) )
|
|
print( "mean beta = {0} +/- {1} deg".format( mean_beta, std_beta ) )
|
|
print( "mean gamma = {0} +/- {1} deg".format( mean_gamma, std_gamma ) )
|
|
|
|
def list_of_ints(arg):
|
|
return list(map(int, arg.split(',')))
|
|
|
|
def list_of_floats(arg):
|
|
return list(map(float, arg.split(',')))
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"-n",
|
|
"--job_name",
|
|
help="the name of the job to be done. Default = split",
|
|
type=str,
|
|
default="split"
|
|
)
|
|
parser.add_argument(
|
|
"-l",
|
|
"--lst_file",
|
|
help="file from SwissFEL output to be processed quickly. Requried.",
|
|
type=os.path.abspath,
|
|
required=True
|
|
)
|
|
parser.add_argument(
|
|
"-k",
|
|
"--chunk_size",
|
|
help="how big should each image split be? Default = 500. Fewer will be faster.",
|
|
type=int,
|
|
default=500
|
|
)
|
|
parser.add_argument(
|
|
"-g",
|
|
"--geom_file",
|
|
help="path to geom file to be used in the refinement. Requried.",
|
|
type=os.path.abspath,
|
|
required=True
|
|
)
|
|
parser.add_argument(
|
|
"-c",
|
|
"--cell_file",
|
|
help="path to cell file of the crystals used in the refinement. Requried.",
|
|
type=os.path.abspath,
|
|
required=True
|
|
)
|
|
parser.add_argument(
|
|
"-x",
|
|
"--indexer",
|
|
help="indexer to use. Default = xgandalf-latt-cell",
|
|
type=str,
|
|
default="xgandalf-latt-cell"
|
|
)
|
|
parser.add_argument(
|
|
"-f",
|
|
"--peakfinder",
|
|
help="peakfinder to use. Default = peakfinder8",
|
|
type=str,
|
|
default="peakfinder8"
|
|
)
|
|
parser.add_argument(
|
|
"-a",
|
|
"--integrator",
|
|
help="integrator to use. Default = rings-nocen-nograd",
|
|
type=str,
|
|
default="rings-nocen-nograd"
|
|
)
|
|
parser.add_argument(
|
|
"-y",
|
|
"--tolerance",
|
|
help="tolerance to use. Default = 10.0,10.0,10.0,2.0,3.0,2.0",
|
|
type=list_of_floats,
|
|
default=[10.0,10.0,10.0,2.0,3.0,2.0]
|
|
)
|
|
parser.add_argument(
|
|
"-t",
|
|
"--threshold",
|
|
help="peaks must be above this to be found during spot-finding. Default = 20",
|
|
type=int,
|
|
default=20
|
|
)
|
|
parser.add_argument(
|
|
"-s",
|
|
"--min_snr",
|
|
help="peaks must to above this to be counted. Default = 5.",
|
|
type=int,
|
|
default=5
|
|
)
|
|
parser.add_argument(
|
|
"-i",
|
|
"--int_radius",
|
|
help="integration ring radii. Default = 2,3,5 = 2 for spot and then 3 and 5 to calculate background.",
|
|
type=list_of_ints,
|
|
default=[2,3,5]
|
|
)
|
|
parser.add_argument(
|
|
"-m",
|
|
"--multi",
|
|
help="do you wnat to look for multiple lattices. Default = True",
|
|
type=bool,
|
|
default=False
|
|
)
|
|
parser.add_argument(
|
|
"-r",
|
|
"--retry",
|
|
help="do you want to retry failed indexing patterns. Default = False",
|
|
type=bool,
|
|
default=False
|
|
)
|
|
parser.add_argument(
|
|
"-p",
|
|
"--min_pix",
|
|
help="minimum number of pixels a spot should contain in peak finding.Default = 2",
|
|
type=int,
|
|
default=2
|
|
)
|
|
parser.add_argument(
|
|
"-b",
|
|
"--bg_rad",
|
|
help="radius (in pixels) used for the estimation of the local background. Default = 4",
|
|
type=int,
|
|
default=4
|
|
)
|
|
parser.add_argument(
|
|
"-q",
|
|
"--min_res",
|
|
help="min-res for spot-finding in pixels. Default = 85.",
|
|
type=int,
|
|
default=85
|
|
)
|
|
parser.add_argument(
|
|
"-v",
|
|
"--reservation",
|
|
help="reservation name for ra cluster. Usually along the lines of P11111_2024-12-10",
|
|
type=str,
|
|
default=None
|
|
)
|
|
parser.add_argument(
|
|
"-d",
|
|
"--debug",
|
|
help="output debug to terminal.",
|
|
type=bool,
|
|
default=False
|
|
)
|
|
args = parser.parse_args()
|
|
# run geom converter
|
|
cwd = os.getcwd()
|
|
# set loguru
|
|
if not args.debug:
|
|
logger.remove()
|
|
logfile = "{0}.log".format( args.job_name )
|
|
logger.add( logfile, format="{message}", level="INFO")
|
|
# log geometry file
|
|
geom = open( args.geom_file, "r" ).read()
|
|
logger.info( geom )
|
|
if args.multi == True:
|
|
multi = "multi"
|
|
else:
|
|
multi = "no-multi"
|
|
if args.retry == True:
|
|
retry = "retry"
|
|
else:
|
|
retry = "no-retry"
|
|
main( cwd, args.job_name, args.lst_file, args.chunk_size,
|
|
args.geom_file, args.cell_file, args.indexer, args.peakfinder,
|
|
args.integrator, args.tolerance, args.threshold,
|
|
args.min_snr, args.int_radius, multi, retry, args.min_pix, args.bg_rad,
|
|
args.min_res, args.reservation ) |