class MDRecordTests(unittest.TestCase): """ Testing MD Record API """ def setUp(self): fname = os.path.join(FILE_DIR, "mdrecord.oedb") ifs = oechem.oeifstream(fname) records = [] for record in read_records(ifs): records.append(record) self.assertEqual(len(records), 1) self.record = OERecord(records[0]) self.mdrecord = MDDataRecord(records[0]) self.cwd = os.getcwd() os.chdir(FILE_DIR) @pytest.mark.travis @pytest.mark.local def test_get_primary(self): mol = self.mdrecord.get_value(Fields.primary_molecule) self.assertEqual(mol.NumAtoms(), self.mdrecord.get_primary.NumAtoms()) @pytest.mark.travis @pytest.mark.local def test_set_primary(self): mol = self.mdrecord.get_value(Fields.primary_molecule) self.assertTrue(self.mdrecord.set_primary(mol)) @pytest.mark.travis @pytest.mark.local def test_get_flask(self): mol = self.mdrecord.get_value(Fields.flask) self.assertEqual(mol.NumAtoms(), self.mdrecord.get_flask.NumAtoms()) @pytest.mark.travis @pytest.mark.local def test_set_flask(self): mol = self.mdrecord.get_value(Fields.flask) self.assertTrue(self.mdrecord.set_flask(mol)) @pytest.mark.travis @pytest.mark.local def test_get_flaskid(self): id = self.mdrecord.get_value(Fields.flaskid) self.assertEqual(id, 0) @pytest.mark.travis @pytest.mark.local def test_has_flaskid(self): self.assertTrue(self.mdrecord.has_flask_id) @pytest.mark.travis @pytest.mark.local def test_set_flaskid(self): self.mdrecord.set_flask_id(5) self.assertEqual(self.mdrecord.get_flask_id, 5) @pytest.mark.travis @pytest.mark.local def test_get_lig_id(self): id = self.mdrecord.get_value(Fields.ligid) self.assertEqual(id, 0) @pytest.mark.travis @pytest.mark.local def test_has_lig_id(self): self.assertTrue(self.mdrecord.has_lig_id) @pytest.mark.travis @pytest.mark.local def test_set_lig_id(self): self.mdrecord.set_lig_id(5) self.assertEqual(self.mdrecord.get_lig_id, 5) @pytest.mark.travis @pytest.mark.local def test_get_conf_id(self): id = self.mdrecord.get_value(Fields.confid) self.assertEqual(id, 0) @pytest.mark.travis @pytest.mark.local def test_has_conf_id(self): self.assertTrue(self.mdrecord.has_conf_id) @pytest.mark.travis @pytest.mark.local def test_set_conf_id(self): self.mdrecord.set_conf_id(5) self.assertEqual(self.mdrecord.get_conf_id, 5) @pytest.mark.travis @pytest.mark.local def test_get_title(self): title = self.mdrecord.get_value(Fields.title) self.assertEqual(title, 'pPRT_ltoluene') @pytest.mark.travis @pytest.mark.local def test_has_tile(self): self.assertTrue(self.mdrecord.has_title) @pytest.mark.travis @pytest.mark.local def test_set_title(self): self.mdrecord.set_title("Pippo") self.assertEqual(self.mdrecord.get_title, 'Pippo') @pytest.mark.travis @pytest.mark.local def test_get_last_stage(self): last_stage = self.mdrecord.get_last_stage self.assertEqual(last_stage.get_value(Fields.stage_name), 'Production') self.assertEqual(last_stage.get_value(Fields.stage_type), 'NPT') @pytest.mark.travis @pytest.mark.local def test_get_stage_by_name(self): last_stage = self.mdrecord.get_stage_by_name() self.assertEqual(last_stage.get_value(Fields.stage_name), 'Production') self.assertEqual(last_stage.get_value(Fields.stage_type), 'NPT') param_stage = self.mdrecord.get_stage_by_name( stg_name='System Parametrization') self.assertEqual(param_stage.get_value(Fields.stage_name), 'System Parametrization') self.assertEqual(param_stage.get_value(Fields.stage_type), 'SETUP') param_stage = self.mdrecord.get_stage_by_name( stg_name='System Minimization') self.assertEqual(param_stage.get_value(Fields.stage_name), 'System Minimization') self.assertEqual(param_stage.get_value(Fields.stage_type), 'MINIMIZATION') with self.assertRaises(ValueError): self.mdrecord.get_stage_by_name('Error') # @pytest.mark.local # def test_delete_stage_by_name(self): # new_record = OERecord(self.record) # new_mdrecord = MDDataRecord(new_record) # # new_mdrecord.delete_stage_by_name(stg_name='System Minimization') # self.assertFalse(new_mdrecord.has_stage_name('System Minimization')) # self.assertEqual(len(new_mdrecord.get_stages), 2) # TODO THIS IS GOING TO DELETE FILES LOCALLY KEEP DISABLED @pytest.mark.travis @pytest.mark.local def test_has_stage_name(self): self.assertTrue(self.mdrecord.has_stage_name('Production')) self.assertFalse(self.mdrecord.has_stage_name('Error')) @pytest.mark.travis @pytest.mark.local def test_get_stage_by_idx(self): with self.assertRaises(ValueError): self.mdrecord.get_stage_by_idx(5) self.assertEqual( self.mdrecord.get_stage_by_idx(0).get_value(Fields.stage_name), 'System Parametrization') @pytest.mark.travis @pytest.mark.local def test_get_stage_state(self): last_stage = self.mdrecord.get_last_stage mddata_fn = os.path.join(FILE_DIR, last_stage.get_value(Fields.mddata)) with TemporaryDirectory() as out_directory: with tarfile.open(mddata_fn) as tar: tar.extractall(path=out_directory) state_fn = os.path.join(out_directory, MDFileNames.state) with open(state_fn, 'rb') as f: md_state = pickle.load(f) self.assertEqual(md_state.get_positions(), self.mdrecord.get_stage_state().get_positions()) self.assertEqual(md_state.get_velocities(), self.mdrecord.get_stage_state().get_velocities()) self.assertEqual(md_state.get_box_vectors(), self.mdrecord.get_stage_state().get_box_vectors()) @pytest.mark.travis @pytest.mark.local def test_get_stage_topology(self): par_stage = self.mdrecord.get_stage_by_idx(0) mddata_fn = os.path.join(FILE_DIR, par_stage.get_value(Fields.mddata)) with TemporaryDirectory() as out_directory: with tarfile.open(mddata_fn) as tar: tar.extractall(path=out_directory) topology_fn = os.path.join(out_directory, MDFileNames.topology) topology_mol = oechem.OEMol() with oechem.oemolistream(topology_fn) as ifs: oechem.OEReadMolecule(ifs, topology_mol) topology = self.mdrecord.get_stage_topology( stg_name='System Parametrization') for mol_at, top_at in zip(topology_mol.GetAtoms(), topology.GetAtoms()): self.assertEqual(mol_at.GetAtomicNum(), top_at.GetAtomicNum()) @pytest.mark.travis @pytest.mark.local def test_get_stage_info(self): last_stage = self.mdrecord.get_last_stage info = last_stage.get_value(Fields.log_data) self.assertEqual(info, self.mdrecord.get_stage_info()) min_stage = self.mdrecord.get_stage_by_name( stg_name='System Minimization') info = min_stage.get_value(Fields.log_data) self.assertEqual( info, self.mdrecord.get_stage_info(stg_name='System Minimization')) @pytest.mark.travis @pytest.mark.local def test_get_stage_trajectory(self): self.assertTrue(os.path.isfile(self.mdrecord.get_stage_trajectory())) @pytest.mark.travis @pytest.mark.local def test_add_new_stage(self): new_record = OERecord(self.record) new_mdrecord = MDDataRecord(new_record) topology = self.mdrecord.get_stage_topology() md_state = self.mdrecord.get_stage_state() self.assertTrue( new_mdrecord.add_new_stage("Testing", MDStageTypes.FEC, topology, md_state, "test.tar.gz", log='TestingLogs')) self.assertEqual(len(new_mdrecord.get_value(Fields.md_stages)), 4) new_last_stage = new_mdrecord.get_stage_by_name(stg_name='Testing') self.assertEqual(new_last_stage.get_value(Fields.stage_name), 'Testing') self.assertEqual(new_last_stage.get_value(Fields.stage_type), MDStageTypes.FEC) @pytest.mark.travis @pytest.mark.local def test_get_stages(self): stages = self.mdrecord.get_stages self.assertEqual(len(stages), 3) @pytest.mark.travis @pytest.mark.local def test_get_stages_names(self): ls_names = [ 'System Parametrization', 'System Minimization', 'Production' ] stg_names = self.mdrecord.get_stages_names self.assertEqual(stg_names, ls_names) @pytest.mark.travis @pytest.mark.local def test_has_stages(self): self.assertTrue(self.mdrecord.has_stages) @pytest.mark.travis @pytest.mark.local def test_get_parmed(self): pmd = self.mdrecord.get_parmed(sync_stage_name='last') self.assertEqual(len(pmd.atoms), 30439) self.assertEqual((len(pmd.residues)), 9446) self.assertEqual((len(pmd.bonds)), 21178) self.assertEqual((len(pmd.angles)), 14069) self.assertEqual((len(pmd.dihedrals)), 8028) @pytest.mark.travis @pytest.mark.local def test_set_parmed(self): new_record = OERecord(self.record) new_mdrecord = MDDataRecord(new_record) pmd = new_mdrecord.get_parmed() new_mdrecord.delete_field(Fields.pmd_structure) self.assertFalse(new_mdrecord.has_parmed) new_mdrecord.set_parmed(pmd, sync_stage_name='last') self.assertTrue(new_mdrecord.has_parmed) @pytest.mark.travis @pytest.mark.local def test_has_parmed(self): self.assertTrue(self.mdrecord.has_parmed) @pytest.mark.travis @pytest.mark.local def test_delete_parmed(self): new_record = OERecord(self.record) new_mdrecord = MDDataRecord(new_record) new_mdrecord.delete_parmed self.assertFalse(new_mdrecord.has_parmed) @pytest.mark.travis @pytest.mark.local def test_protein_traj(self): oetraj_record = self.record.get_value(OEField('OETraj', Types.Record)) prot_mol = oetraj_record.get_value(Fields.protein_traj_confs) mdrecord = MDDataRecord(oetraj_record) mdprot = mdrecord.get_protein_traj self.assertEqual(prot_mol.NumAtoms(), mdprot.NumAtoms()) @pytest.mark.travis @pytest.mark.local def test_set_protein_traj(self): oetraj_record = self.record.get_value(OEField('OETraj', Types.Record)) prot_mol = oetraj_record.get_value(Fields.protein_traj_confs) mdrecord = MDDataRecord(oetraj_record) oetraj_record.delete_field(Fields.protein_traj_confs) with self.assertRaises(ValueError): mdrecord.get_protein_traj self.assertTrue(mdrecord.set_protein_traj(prot_mol))
def process(self, record, port): try: # The copy of the dictionary option as local variable # is necessary to avoid filename collisions due to # the parallel cube processes opt = dict(self.opt) opt['CubeTitle'] = self.title # Logger string str_logger = '-' * 32 + ' MIN CUBE PARAMETERS ' + '-' * 32 str_logger += "\n{:<25} = {:<10}".format("Cube Title", opt['CubeTitle']) for k, v in sorted(self.parameters().items()): tmp_default = copy.deepcopy(v) if v.default is None: tmp_default.default = 'None' elif isinstance(v, parameters.BooleanParameter): if v.default: tmp_default.default = 'True' else: tmp_default.default = 'False' else: tmp_description = textwrap.fill(" ".join( v.description.split()), subsequent_indent=' ' * 39, width=80) str_logger += "\n{:<25} = {:<10} {}".format( k, getattr(self.args, tmp_default.name), tmp_description) str_logger += "\n{:<25} = {:<10}".format("Simulation Type", opt['SimType']) # Create the MD record to use the MD Record API mdrecord = MDDataRecord(record) system_title = mdrecord.get_title opt['system_title'] = system_title opt['system_id'] = mdrecord.get_flask_id flask = mdrecord.get_stage_topology() mdstate = mdrecord.get_stage_state() if opt['restraint_to_reference']: opt['reference_state'] = mdrecord.get_stage_state( stg_name=MDStageNames.ForceField) opt['out_directory'] = mdrecord.cwd opt['molecule'] = flask opt['str_logger'] = str_logger opt['Logger'].info('[{}] MINIMIZING System: {}'.format( opt['CubeTitle'], system_title)) # Extract the Parmed structure and synchronize it with the last MD stage state parmed_structure = mdrecord.get_parmed(sync_stage_name='last') # Run the MD simulation new_mdstate = md_simulation(mdstate, parmed_structure, opt) # Update the flask coordinates flask.SetCoords(new_mdstate.get_oe_positions()) mdrecord.set_flask(flask) data_fn = os.path.basename( mdrecord.cwd) + '_' + opt['system_title'] + '_' + str( opt['system_id']) + '-' + opt['suffix'] + '.tar.gz' if not mdrecord.add_new_stage(self.title, MDStageTypes.MINIMIZATION, flask, new_mdstate, data_fn, append=opt['save_md_stage'], log=opt['str_logger']): raise ValueError("Problems adding the new Minimization Stage") self.success.emit(mdrecord.get_record) del mdrecord except Exception as e: print("Failed to complete", str(e), flush=True) self.opt['Logger'].info('Exception {} {}'.format( str(e), self.title)) self.log.error(traceback.format_exc()) self.failure.emit(record) return
def process(self, record, port): try: # The copy of the dictionary option as local variable # is necessary to avoid filename collisions due to # the parallel cube processes opt = dict(self.opt) opt['CubeTitle'] = self.title # Logger string str_logger = '-' * 32 + ' NPT CUBE PARAMETERS ' + '-' * 32 str_logger += "\n{:<25} = {:<10}".format("Cube Title", opt['CubeTitle']) for k, v in sorted(self.parameters().items()): tmp_default = copy.deepcopy(v) if v.default is None: tmp_default.default = 'None' elif isinstance(v, parameters.BooleanParameter): if v.default: tmp_default.default = 'True' else: tmp_default.default = 'False' else: tmp_description = textwrap.fill(" ".join( v.description.split()), subsequent_indent=' ' * 39, width=80) str_logger += "\n{:<25} = {:<10} {}".format( k, getattr(self.args, tmp_default.name), tmp_description) str_logger += "\n{:<25} = {:<10}".format("Simulation Type", opt['SimType']) # Create the MD record to use the MD Record API mdrecord = MDDataRecord(record) system_title = mdrecord.get_title opt['system_title'] = system_title opt['system_id'] = mdrecord.get_flask_id flask = mdrecord.get_stage_topology() mdstate = mdrecord.get_stage_state() # Update cube simulation parameters for field in record.get_fields(include_meta=True): field_name = field.get_name() if field_name in ['temperature', 'pressure']: rec_value = record.get_value(field) opt[field_name] = rec_value opt['Logger'].info( "{} Updating parameters for molecule: {} {} = {}". format(self.title, system_title, field_name, rec_value)) if opt['restraint_to_reference']: opt['reference_state'] = mdrecord.get_stage_state( stg_name=MDStageNames.ForceField) opt['out_directory'] = mdrecord.cwd opt['molecule'] = flask opt['str_logger'] = str_logger opt['Logger'].info('[{}] START NPT SIMULATION: {}'.format( opt['CubeTitle'], system_title)) opt['out_fn'] = os.path.basename(opt['out_directory']) + '_' + \ opt['system_title'] + '_' + \ str(opt['system_id']) + '-' + \ opt['suffix'] # Trajectory file name if any generated opt['trj_fn'] = opt['out_fn'] + '_' + 'traj.tar.gz' # Extract the Parmed structure and synchronize it with the last MD stage state parmed_structure = mdrecord.get_parmed(sync_stage_name='last') # Run the MD simulation new_mdstate = md_simulation(mdstate, parmed_structure, opt) # Update the system coordinates flask.SetCoords(new_mdstate.get_oe_positions()) mdrecord.set_flask(flask) # Trajectory if opt['trajectory_interval'] or opt['trajectory_frames']: trajectory_fn = opt['trj_fn'] if opt['md_engine'] == MDEngines.OpenMM: trajectory_engine = MDEngines.OpenMM else: trajectory_engine = MDEngines.Gromacs else: # Empty Trajectory trajectory_fn = None trajectory_engine = None data_fn = opt['out_fn'] + '.tar.gz' if not mdrecord.add_new_stage( self.title, MDStageTypes.NPT, flask, new_mdstate, data_fn, append=opt['save_md_stage'], log=opt['str_logger'], trajectory_fn=trajectory_fn, trajectory_engine=trajectory_engine, trajectory_orion_ui=opt['system_title'] + '_' + str(opt['system_id']) + '-' + opt['suffix'] + '.tar.gz'): raise ValueError("Problems adding in the new NPT Stage") self.success.emit(mdrecord.get_record) del mdrecord except Exception as e: print("Failed to complete", str(e), flush=True) self.opt['Logger'].info('Exception {} {}'.format( str(e), self.title)) self.log.error(traceback.format_exc()) self.failure.emit(record) return
def test_compex_prep_split_floe(self): workfloe = WorkFloeWrapper.get_workfloe(os.path.join( FLOES_DEV_DIR, "Complex_prep.py"), run_timeout=43200, queue_timeout=2000) ligand_file = DatasetWrapper.get_dataset( os.path.join(FILE_DIR, "MCL1_lig26.oeb")) protein_file = DatasetWrapper.get_dataset( os.path.join(FILE_DIR, "MCL1_protein_ACE_NMA_caps.pdb")) output_file = OutputDatasetWrapper(extension=".oedb") fail_output_file = OutputDatasetWrapper(extension=".oedb") workfloe.start({ "promoted": { "ligands": ligand_file.identifier, "protein": protein_file.identifier, "out": output_file.identifier, "fail": fail_output_file.identifier } }) self.assertWorkFloeComplete(workfloe) fail_ifs = oechem.oeifstream() records_fail = [] for rec_fail in read_records(fail_ifs): records_fail.append(rec_fail) fail_ifs.close() count = len(records_fail) # The fail record must be empty self.assertEqual(count, 0) ifs = oeifstream(output_file.path) records = [] for rec in read_records(ifs): records.append(rec) ifs.close() count = len(records) # Check the out record list self.assertEqual(count, 1) # Each record should have the MD record interface for record in records: mdrecord = MDDataRecord(record) self.assertTrue(record.has_value(Fields.flaskid)) self.assertTrue(record.has_value(Fields.title)) self.assertTrue(record.has_value(Fields.ligand)) self.assertTrue(record.has_value(Fields.protein)) self.assertTrue(record.has_value(Fields.primary_molecule)) self.assertTrue(record.has_value(Fields.flask)) self.assertTrue(record.has_value(Fields.md_stages)) self.assertTrue(record.has_value(Fields.pmd_structure)) self.assertEqual(mdrecord.get_flask_id, 0) self.assertEqual(mdrecord.get_title, "pMCL1_l26") self.assertEqual(record.get_value(Fields.ligand).NumAtoms(), 43) self.assertEqual(record.get_value(Fields.protein).NumAtoms(), 2432) complx = mdrecord.get_flask protein_split, ligand_split, water, excipients = oeommutils.split( complx) self.assertEqual(protein_split.NumAtoms(), 2432) self.assertEqual(ligand_split.NumAtoms(), 43) self.assertEqual(water.NumAtoms(), 20022) self.assertEqual(excipients.NumAtoms(), 17) stages = mdrecord.get_stages self.assertEqual(len(stages), 1) stage = stages[0] self.assertTrue(stage.has_value(Fields.stage_name)) self.assertTrue(stage.has_value(Fields.mddata)) self.assertEqual(stage.get_value(Fields.stage_type), "SETUP") top_mol = mdrecord.get_stage_topology() self.assertEqual(top_mol.NumAtoms(), complx.NumAtoms())