auto_namd_python/auto_namd/job.py

import os
import sys
import glob

from auto_namd.functions import *
from auto_namd.simulation import Simulation


class Job:
    """
    A class that represents a job/system residing in a unique directory path

    Attributes:
    self.path = absolute job directory path
    self.name = name of the job, derived from parent directory
    self.prefix = self.path + self.name
    self.pdb = latest .pdb file
    self.psf = latest .psf file
    self.coor = latest .coor file
    self.stage = current simulation stage
    self.step = current simulation steps
    """

    def __init__(self, path):
        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 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


    # allow the Simulation class to instantiate via the simulate() method
    def simulate(self, ffs_path, steps, namdbin, params):
        Simulation(self, ffs_path, steps, namdbin, params)


    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'
              )


"""Job control"""
def get_next_job(jobs_path):
    jobs_path = san_path(jobs_path)
    job_dirs_l = get_jobs_from_path(jobs_path)
    jobs = create_job_instances(job_dirs_l)
    next_job = get_youngest_job(jobs)
    return next_job


def get_jobs_from_path(jobs_path):
    # Only return job paths with a pdb file
    jobs_path = os.path.join(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 create_job_instances(job_paths):
    job_instances = []
    for job_path in job_paths:
        job_instances.append(Job(job_path))
    return job_instances


def get_youngest_job(jobs):
    jobs.sort(key=lambda x: (x.stage, x.step))
    return jobs[0]