Initial commit

job.py

@@ -0,0 +1,152 @@
+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