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 + ' 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 process(self, record, port):
try:
opt = self.opt
opt['CubeTitle'] = self.title
if not record.has_value(Fields.md_components):
raise ValueError("MD Components Field is missing")
md_components = record.get_value(Fields.md_components)
flask, map_comp = md_components.create_flask
opt['Logger'].info(md_components.get_info)
if not record.has_value(Fields.title):
self.log.warn("Missing record Title field")
flask_title = flask.GetTitle()[0:12]
else:
flask_title = record.get_value(Fields.title)
# Parametrize the whole flask
flask_pmd_structure = md_components.parametrize_components(protein_ff=opt['protein_forcefield'],
ligand_ff=opt['ligand_forcefield'])
# Set Parmed structure box_vectors
is_periodic = True
if md_components.get_box_vectors is not None:
flask_pmd_structure.box_vectors = md_components.get_box_vectors
else:
is_periodic = False
self.log.warn("Flask {} has been parametrize without periodic box vectors ".format(flask_title))
if flask.NumAtoms() != flask_pmd_structure.topology.getNumAtoms():
raise ValueError("The flask {} and the generated Parmed structure "
"have mismatch atom numbers: {} vs {}".
format(flask_title,
flask.NumAtoms(),
flask_pmd_structure.topology.getNumAtoms()))
# Check Formal vs Partial charges
flask_formal_charge = 0
for at in flask.GetAtoms():
flask_formal_charge += at.GetFormalCharge()
flask_partial_charge = 0.0
for at in flask_pmd_structure.atoms:
flask_partial_charge += at.charge
if abs(flask_formal_charge - flask_partial_charge) > 0.01:
raise ValueError("Flask Formal charge and flask Partial charge mismatch: {} vs {}".format(
flask_formal_charge, flask_partial_charge))
# Copying the charges between the parmed structure and the oemol
for parm_at, oe_at in zip(flask_pmd_structure.atoms, flask.GetAtoms()):
if parm_at.atomic_number != oe_at.GetAtomicNum():
raise ValueError(
"Atomic number mismatch between the Parmed and the OpenEye topologies: {} - {}".
format(parm_at.atomic_number, oe_at.GetAtomicNum()))
oe_at.SetPartialCharge(parm_at.charge)
# Set the component charges
for comp_name, comp in md_components.get_components.items():
for at_comp in comp.GetAtoms():
pred = oechem.OEHasAtomIdx(map_comp[comp_name][at_comp.GetIdx()])
at_flask = flask.GetAtom(pred)
if at_flask.GetAtomicNum() != at_comp.GetAtomicNum():
"Atomic number mismatch between the component {} atom {} and the flask atom {}".\
format(comp_name, at_comp, at_flask)
at_comp.SetPartialCharge(at_flask.GetPartialCharge())
md_components.set_component_by_name(comp_name, comp)
# Update the components after setting the charges
record.set_value(Fields.md_components, md_components)
# Check if it is possible to create the OpenMM System
if is_periodic:
omm_flask = flask_pmd_structure.createSystem(nonbondedMethod=app.CutoffPeriodic,
nonbondedCutoff=10.0 * unit.angstroms,
constraints=None,
removeCMMotion=False,
rigidWater=False)
else:
omm_flask = flask_pmd_structure.createSystem(nonbondedMethod=app.NoCutoff,
constraints=None,
removeCMMotion=False,
rigidWater=False)
mdrecord = MDDataRecord(record)
sys_id = mdrecord.get_flask_id
mdrecord.set_flask(flask)
mdrecord.set_parmed(flask_pmd_structure, shard_name="Parmed_" + flask_title + '_' + str(sys_id))
data_fn = os.path.basename(mdrecord.cwd) + '_' + flask_title+'_' + str(sys_id) + '-' + opt['suffix']+'.tar.gz'
if not mdrecord.add_new_stage(MDStageNames.ForceField,
MDStageTypes.SETUP,
flask,
MDState(flask_pmd_structure),
data_fn):
raise ValueError("Problems adding the new Parametrization 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 + ' 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