153 lines
4.5 KiB
Python
153 lines
4.5 KiB
Python
import os
|
|
import sys
|
|
import glob
|
|
|
|
from functions import natural_key, abs_path
|
|
|
|
class Job:
|
|
"""
|
|
A class that represents a job/system residing in a unique directory path
|
|
|
|
|
|
Calling prepare_sim() assigns the additional instance variables:
|
|
steps = number of steps to run the simulation
|
|
ffs = a list of forcefield files for the simulation
|
|
next_step, final_step_num = the type of step and the final number of steps
|
|
conf, out file = the configuration file and the simulation output file
|
|
"""
|
|
|
|
def __init__(self, path):
|
|
self.set_empty_vars()
|
|
self.path = path
|
|
self.name = self.get_job_name()
|
|
self.prefix = os.path.join(self.path, self.name)
|
|
self.pdb, self.psf = self.get_pdb_and_psf()
|
|
self.stage, self.step, self.coor = self.get_stage()
|
|
|
|
|
|
def set_empty_vars(self):
|
|
self.pdb, self.psf = '', ''
|
|
self.steps = ''
|
|
self.ffs_path, self.ffs = '', ''
|
|
self.next_stage, self.next_step = '', ''
|
|
self.conf, self.out = '', ''
|
|
|
|
|
|
def get_job_name(self):
|
|
'''Get job name from path'''
|
|
name = os.path.basename(self.path)
|
|
return name
|
|
|
|
|
|
def get_pdb_and_psf(self):
|
|
pdb = glob.glob(f'{self.prefix}*_solv_ion.pdb')
|
|
psf = glob.glob(f'{self.prefix}*_solv_ion.psf')
|
|
if len(pdb) == 1 and len(psf) == 1:
|
|
return pdb[0], psf[0]
|
|
pdb = glob.glob(f'{self.prefix}.pdb')
|
|
psf = glob.glob(f'{self.prefix}.psf')
|
|
if len(pdb) == 1 and len(psf) == 1:
|
|
return pdb[0], psf[0]
|
|
else:
|
|
print('No PDB or PSF files found, exiting...')
|
|
sys.exit(1)
|
|
|
|
|
|
def get_stage(self):
|
|
coors_path = self.prefix + '*.coor'
|
|
coors = glob.glob(coors_path)
|
|
coors.sort(key=natural_key)
|
|
if len(coors) >= 1:
|
|
coor = coors[-1]
|
|
stage = coor.split('_')[-2].split('.')[0]
|
|
step = int(coor.split('_')[-1].split('.')[0])
|
|
else:
|
|
coor = ''
|
|
stage = ''
|
|
step = 0
|
|
return stage, int(step), coor
|
|
|
|
|
|
def prepare_sim(self, ffs_path, steps):
|
|
self.steps = steps
|
|
self.ffs_path = abs_path(ffs_path)
|
|
self.ffs = self.get_ffs()
|
|
self.next_stage, self.next_step = self.get_next_stage()
|
|
self.conf, self.out = self.conf_out_file()
|
|
|
|
|
|
def get_next_stage(self):
|
|
if self.stage == '4-sim':
|
|
next_stage = '4-sim'
|
|
next_step = self.step + self.steps
|
|
elif self.stage == '3-heat':
|
|
next_stage = '4-sim'
|
|
next_step = steps
|
|
elif self.stage == '2-min':
|
|
next_stage = '3-heat'
|
|
next_step = 0
|
|
elif self.stage == '1-min':
|
|
next_stage = '2-min'
|
|
next_step = 0
|
|
elif self.stage == '':
|
|
next_stage = '1-min'
|
|
next_step = 0
|
|
return next_stage, int(next_step)
|
|
|
|
|
|
def get_ffs(self):
|
|
ffs = glob.glob(os.path.join(self.ffs_path, '*'))
|
|
return ffs
|
|
|
|
|
|
def conf_out_file(self):
|
|
if self.next_stage == '4-sim':
|
|
conf = f'{self.prefix}_{self.next_stage}' \
|
|
f'_{str(self.next_step)}.conf'
|
|
out = f'{self.prefix}_{self.next_stage}' \
|
|
f'_{str(self.next_step)}.out'
|
|
else:
|
|
conf = f'{self.prefix}_{self.next_stage}.conf'
|
|
out = f'{self.prefix}_{self.next_stage}.out'
|
|
|
|
return conf, out
|
|
|
|
|
|
def info(self):
|
|
print(f'Job Path: {self.path}\n'
|
|
f'Name: {self.name}\n'
|
|
f'Working PDB: {self.pdb}\n'
|
|
f'Working PSF: {self.psf}\n'
|
|
f'Previous Step: {self.stage}\n'
|
|
f'Previous Step Number: {str(self.step)}\n'
|
|
f'Next Stage: {self.next_stage}\n'
|
|
f'Next Step Number: {str(self.next_step)}\n'
|
|
f'FF Path: {self.ffs_path}\n'
|
|
f'FFs: {self.ffs}\n'
|
|
)
|
|
|
|
|
|
|
|
def get_jobs_from_path(jobs_path):
|
|
jobs_path += '/**/*.pdb'
|
|
job_dirs = glob.glob(jobs_path, recursive=True)
|
|
job_dirs_l = []
|
|
for job_dir in job_dirs:
|
|
job_dirs_l.append(os.path.dirname(job_dir))
|
|
if len(job_dirs_l) < 1:
|
|
print('No valid jobs found in jobs path!')
|
|
sys.exit(1)
|
|
return job_dirs_l
|
|
|
|
|
|
def get_next_job(jobs):
|
|
jobs.sort(key=lambda x: (x.stage, x.step))
|
|
return jobs[0]
|
|
|
|
|
|
def create_job_instances(job_paths):
|
|
job_instances = []
|
|
for job_path in job_paths:
|
|
job_instances.append(Job(job_path))
|
|
return job_instances
|