<commit_before>import os
import getpass
import numpy as np
from lammps import lammps
import Tools
import Save
import helpers.time as tm
import itertools as it
import tempfile as temp
import time
import yaml

import helpers.git as _git
__git_hash__ = _git.get_git_hash(__file__)

class Job(object):

    def __init__(self, config_path):
        super(Job, self).__init__()
        self.config = self.read_config(config_path)
        self.config_with_setup = os.path.join(os.path.dirname(config_path), 'config_with_setup.yml')
        if os.path.exists(self.config_with_setup):
            with open(self.config_with_setup) as f:
                setup_config = yaml.load(f)
            self.config.sim_parameter['poly_sequence'] = setup_config['sim_parameter']['poly_sequence']

        self.username = getpass.getuser()

        self.config.sim_path['sim_list'] = os.path.join(self.config.sim_path['root'], 'sim.list')
        open(self.config.sim_path['sim_list'], 'a').close()
        if self.config.sim_parameter['local'] == 1:
            self._switch_to_local()
        else:
            new_paths = self.create_folders(self.config.lmp_path['root'])
            self.config.lmp_path.update(new_paths)

    def read_config(self, path):
        return Tools.config.JobConfig(path)

    def _create_fifos(self):
        fifo = {}
        for name, data in self.config.fifo.items():
            data['path'] = os.path.join(self.config.lmp_path['fifo'], data['path'])
            fifo[name] = Tools.FiFo.create_from_dict(name, self, data)
        return fifo

    def set_active_site(self):
        as_data = self.sim.activeSiteParticles(self.protein, self.config.sim_path['active_site'])
        # the mapping ensures that vanilla python objects are used as the initial data is numpy-based.
        return  {'xyz': map(int, as_data['xyz']), 
                'chain': map(str, as_data['chain']), 
                'pdb_id': map(int, as_data['pdb_id']), 
                'iCode': map(str, as_data['iCode'])}
        

    def setup_env(self):
        self.env = Tools.Environment(self.config.sim_path['config'])
        # Protein
        self.protein_creator = Tools.ProteinCreator(
            self.env,
            self.config.sim_path['protein'],
            with_ions=True,
            surface_file=self.config.sim_path['surface_db'],
            protonation_file=self.config.sim_path['protonation_db'],
            ph=self.config.sim_parameter['pH']
            )

        self.protein = self.protein_creator.create()
        # Polymer
        if 'poly_sequence' in self.config.sim_parameter:
            self.polymer_creator = Tools.PolymerCreator(self.env, self.config.sim_parameter['poly_sequence'], mode='cycle')
        else:
            self.polymer_creator = Tools.PolymerCreator(self.env, 
                self.config.sim_parameter['monomers'], weights=self.config.sim_parameter['weights'], 
                length=self.config.sim_parameter['poly_length'])
        self.poly = self.polymer_creator.create()

        self.sim = Tools.EnvManipulator(self.env, auto_repulsion=False)
        self.sim.create_random_start_positions()
        self.setup_writer = Tools.LmpWriter(self.env)

        # Update Lammps Parameters
        self.config.sim_parameter['poly_sequence'] = [monomers.type_.name for monomers in self.poly.data['monomers']]
        self.config.sim_parameter['named_sequence'] = [particle.type_.name for particle in self.poly.data['particles']]
        self.config.sim_parameter['id_sequence'] = [particle.type_.Id for particle in self.poly.data['particles']]
        self.config.sim_parameter['active_site'] = self.set_active_site()
        self.config.lmp_parameter['active_site_ids'] = self.config.sim_parameter['active_site']['xyz']
        self.config.lmp_parameter['monomer_ids'] = self._get_monomer_ids()
        self.config.lmp_parameter['bb_id'] = self.env.atom_type['BB_bb'].Id
        self.config.lmp_parameter['ghost_id'] = self.env.atom_type['BB_ghost_bb'].Id
        # the config with added information is always saved to the root directory
        self.config.save(self.config_with_setup)
        # save particle list
        p_list_path = os.path.join(self.config.sim_path['root'], 'particle_list.npy')
        self.save_particle_list(p_list_path)
        # fifos can only be created if the monomer ids are known
        self.fifo = self._create_fifos()

    def save_particle_list(self, path):
        dtype = [('xyz', np.int), ('p_id', np.int), ('name', 'S6'), 
                 ('chain', 'S1'), ('atom', 'S6'), ('res_id', np.int), ('iCode', 'S1')]

        molecules = sorted(self.env.molecules.values(), key=lambda x: x.Id)
        particle_gen = it.chain(*[molecule.data['particles'] for molecule in molecules])
        particle_count = sum([len(mol.data['particles']) for mol in self.env.molecules.values()])
        particle_dat = np.empty(particle_count, dtype)
        for i,p in enumerate(particle_gen):
            particle_dat[i] = (p.Id, p.type_.Id, p.residue.name, p.residue.chain, 
                                p.residue.id[0], p.residue.id[1], p.residue.id[2]) 
        particle_dat.tofile(path)

    def terminate_fifos(self):
        for name, fifo in self.fifo.items():
            fifo.terminate()

    def setup_job_save(self):
        self.compactor = Save.JobSave(self.config.lmp_path)

    def save(self):
        '''Save the data of the completed simulations to HDF5 database.
        '''
        if not hasattr(self, 'compactor'):
            self.setup_job_save()
        try:
            self.compactor.save()
            self.compactor.save_versions(lmp_version=lammps().version())
        finally:
            self.compactor.db.close()

    def clean_up(self):
        try:
            self.terminate_fifos()
        except OSError:
            pass
        self.compactor.clean_up()

    def set_variables(self, vars_, lmp_instance):
        for name, val in vars_.items():
            # lists are converted to strings
            if isinstance(val, list):
                val = ' '.join(map(str, val))
            lmp_instance.command('variable %s string "%s"'% (name,val))

    def set_paths(self, paths, lmp_instance):
        for name, path in paths.items():
            lmp_instance.command('variable %s string "%s"'% (name, path))

    def set_fifos(self, fifos, lmp_instance):
        for name,fifo in fifos.items():
            lmp_instance.command(fifo.lammps_string())

    def lmps_run(self, Id, parameters, paths, fifos={}):
        self._start_fifo_capture(Id)
        lammps_sim = lammps()
        # submitting parameters
        self.set_variables(parameters, lammps_sim)
        # submitting paths
        self.set_paths(paths, lammps_sim)
        # submitting run Id
        lammps_sim.command('variable num string %05d'% Id)
        # starting script
        lammps_sim.file(paths['script'])
        # specify fifo dumps
        self.set_fifos(self.fifo, lammps_sim)
        lammps_sim.command('run ${time_steps}')
        # write snapshot of end-comformation
        lammps_sim.command('write_dump solid xyz ${end_xyz}.xyz')
        lammps_sim.close()
        # report completed simulation so restarting jobs will know
        # also, it notes the machine and folder, so scattered info can be retrieved
        self._mark_complete(Id)

    def generate_new_sim(self, index):
        # Create new Start Conditions
        self.sim.create_random_start_positions()
        # Create New Setup File
        self.setup_writer.write('%s/%05d' % (self.config.lmp_path['input'], index))

    def run(self):
        start_idx = self._get_last_uncompleted_index()
        end_idx = self.config.sim_parameter['sampling_rate']
        # check if simulations is already completed
        if start_idx == -1:
            return 
        for i in xrange(start_idx,  end_idx):
            self.generate_new_sim(i)
            # start next LAMMPS run
            self.lmps_run(i, self.config.lmp_parameter, self.config.lmp_path, fifos=self.fifo)
            # mark job as completed
        if start_idx != -1:
            self._mark_complete(-1)

    def create_local_env(self, local_dir='/data/'):
        '''Create unique local job-folder and create the 
        '''
        name_comp = self.config.sim_path['root'].split('/')[-3:]
        if name_comp[1] == 'jobs':
            del name_comp[1]
        else:
            del name_comp[0]
        folder_name = '-'.join(name_comp)
        if 'local_root' in self.config.sim_path and self.config.sim_path['local_root'] != '':<commit_msg>Moved _mark_complete() from lmps_run to run.<commit_after>import os
import getpass
import numpy as np
from lammps import lammps
import Tools
import Save
import helpers.time as tm
import itertools as it
import tempfile as temp
import time
import yaml

import helpers.git as _git
__git_hash__ = _git.get_git_hash(__file__)

class Job(object):

    def __init__(self, config_path):
        super(Job, self).__init__()
        self.config = self.read_config(config_path)
        self.config_with_setup = os.path.join(os.path.dirname(config_path), 'config_with_setup.yml')
        if os.path.exists(self.config_with_setup):
            with open(self.config_with_setup) as f:
                setup_config = yaml.load(f)
            self.config.sim_parameter['poly_sequence'] = setup_config['sim_parameter']['poly_sequence']

        self.username = getpass.getuser()

        self.config.sim_path['sim_list'] = os.path.join(self.config.sim_path['root'], 'sim.list')
        open(self.config.sim_path['sim_list'], 'a').close()
        if self.config.sim_parameter['local'] == 1:
            self._switch_to_local()
        else:
            new_paths = self.create_folders(self.config.lmp_path['root'])
            self.config.lmp_path.update(new_paths)

    def read_config(self, path):
        return Tools.config.JobConfig(path)

    def _create_fifos(self):
        fifo = {}
        for name, data in self.config.fifo.items():
            data['path'] = os.path.join(self.config.lmp_path['fifo'], data['path'])
            fifo[name] = Tools.FiFo.create_from_dict(name, self, data)
        return fifo

    def set_active_site(self):
        as_data = self.sim.activeSiteParticles(self.protein, self.config.sim_path['active_site'])
        # the mapping ensures that vanilla python objects are used as the initial data is numpy-based.
        return  {'xyz': map(int, as_data['xyz']), 
                'chain': map(str, as_data['chain']), 
                'pdb_id': map(int, as_data['pdb_id']), 
                'iCode': map(str, as_data['iCode'])}
        

    def setup_env(self):
        self.env = Tools.Environment(self.config.sim_path['config'])
        # Protein
        self.protein_creator = Tools.ProteinCreator(
            self.env,
            self.config.sim_path['protein'],
            with_ions=True,
            surface_file=self.config.sim_path['surface_db'],
            protonation_file=self.config.sim_path['protonation_db'],
            ph=self.config.sim_parameter['pH']
            )

        self.protein = self.protein_creator.create()
        # Polymer
        if 'poly_sequence' in self.config.sim_parameter:
            self.polymer_creator = Tools.PolymerCreator(self.env, self.config.sim_parameter['poly_sequence'], mode='cycle')
        else:
            self.polymer_creator = Tools.PolymerCreator(self.env, 
                self.config.sim_parameter['monomers'], weights=self.config.sim_parameter['weights'], 
                length=self.config.sim_parameter['poly_length'])
        self.poly = self.polymer_creator.create()

        self.sim = Tools.EnvManipulator(self.env, auto_repulsion=False)
        self.sim.create_random_start_positions()
        self.setup_writer = Tools.LmpWriter(self.env)

        # Update Lammps Parameters
        self.config.sim_parameter['poly_sequence'] = [monomers.type_.name for monomers in self.poly.data['monomers']]
        self.config.sim_parameter['named_sequence'] = [particle.type_.name for particle in self.poly.data['particles']]
        self.config.sim_parameter['id_sequence'] = [particle.type_.Id for particle in self.poly.data['particles']]
        self.config.sim_parameter['active_site'] = self.set_active_site()
        self.config.lmp_parameter['active_site_ids'] = self.config.sim_parameter['active_site']['xyz']
        self.config.lmp_parameter['monomer_ids'] = self._get_monomer_ids()
        self.config.lmp_parameter['bb_id'] = self.env.atom_type['BB_bb'].Id
        self.config.lmp_parameter['ghost_id'] = self.env.atom_type['BB_ghost_bb'].Id
        # the config with added information is always saved to the root directory
        self.config.save(self.config_with_setup)
        # save particle list
        p_list_path = os.path.join(self.config.sim_path['root'], 'particle_list.npy')
        self.save_particle_list(p_list_path)
        # fifos can only be created if the monomer ids are known
        self.fifo = self._create_fifos()

    def save_particle_list(self, path):
        dtype = [('xyz', np.int), ('p_id', np.int), ('name', 'S6'), 
                 ('chain', 'S1'), ('atom', 'S6'), ('res_id', np.int), ('iCode', 'S1')]

        molecules = sorted(self.env.molecules.values(), key=lambda x: x.Id)
        particle_gen = it.chain(*[molecule.data['particles'] for molecule in molecules])
        particle_count = sum([len(mol.data['particles']) for mol in self.env.molecules.values()])
        particle_dat = np.empty(particle_count, dtype)
        for i,p in enumerate(particle_gen):
            particle_dat[i] = (p.Id, p.type_.Id, p.residue.name, p.residue.chain, 
                                p.residue.id[0], p.residue.id[1], p.residue.id[2]) 
        particle_dat.tofile(path)

    def terminate_fifos(self):
        for name, fifo in self.fifo.items():
            fifo.terminate()

    def setup_job_save(self):
        self.compactor = Save.JobSave(self.config.lmp_path)

    def save(self):
        '''Save the data of the completed simulations to HDF5 database.
        '''
        if not hasattr(self, 'compactor'):
            self.setup_job_save()
        try:
            self.compactor.save()
            self.compactor.save_versions(lmp_version=lammps().version())
        finally:
            self.compactor.db.close()

    def clean_up(self):
        try:
            self.terminate_fifos()
        except OSError:
            pass
        self.compactor.clean_up()

    def set_variables(self, vars_, lmp_instance):
        for name, val in vars_.items():
            # lists are converted to strings
            if isinstance(val, list):
                val = ' '.join(map(str, val))
            lmp_instance.command('variable %s string "%s"'% (name,val))

    def set_paths(self, paths, lmp_instance):
        for name, path in paths.items():
            lmp_instance.command('variable %s string "%s"'% (name, path))

    def set_fifos(self, fifos, lmp_instance):
        for name,fifo in fifos.items():
            lmp_instance.command(fifo.lammps_string())

    def lmps_run(self, Id, parameters, paths, fifos={}):
        self._start_fifo_capture(Id)
        lammps_sim = lammps()
        # submitting parameters
        self.set_variables(parameters, lammps_sim)
        # submitting paths
        self.set_paths(paths, lammps_sim)
        # submitting run Id
        lammps_sim.command('variable num string %05d'% Id)
        # starting script
        lammps_sim.file(paths['script'])
        # specify fifo dumps
        self.set_fifos(self.fifo, lammps_sim)
        lammps_sim.command('run ${time_steps}')
        # write snapshot of end-comformation
        lammps_sim.command('write_dump solid xyz ${end_xyz}.xyz')
        lammps_sim.close()

    def generate_new_sim(self, index):
        # Create new Start Conditions
        self.sim.create_random_start_positions()
        # Create New Setup File
        self.setup_writer.write('%s/%05d' % (self.config.lmp_path['input'], index))

    def run(self):
        start_idx = self._get_last_uncompleted_index()
        end_idx = self.config.sim_parameter['sampling_rate']
        # check if simulations is already completed
        if start_idx == -1:
            return 
        for i in xrange(start_idx,  end_idx):
            self.generate_new_sim(i)
            # start next LAMMPS run
            self.lmps_run(i, self.config.lmp_parameter, self.config.lmp_path, fifos=self.fifo)
            # report completed simulation so restarting jobs will know
            # also, it notes the machine and folder, so scattered info can be retrieved
            self._mark_complete(i)
        if start_idx != -1:
            self._mark_complete(-1)

    def create_local_env(self, local_dir='/data/'):
        '''Create unique local job-folder and create the 
        '''
        name_comp = self.config.sim_path['root'].split('/')[-3:]
        if name_comp[1] == 'jobs':
            del name_comp[1]
        else:
            del name_comp[0]
        folder_name = '-'.join(name_comp)
        if 'local_root' in self.config.sim_path and self.config.sim_path['local_root'] != '':
{'fraction_non_alphanumeric': 0.12272302662307333, 'fraction_numerical': 0.002335357309668379, 'mean_word_length': 2.622884940778342, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '383868', 'n_tokens_mistral': 5474, 'n_tokens_neox': 5259, 'n_words': 1043}