def test_torsiondrive_procedure(storage_socket, session):
"""
Torsiondrive procedure
"""
assert session.query(TorsionDriveProcedureORM).count() == 0
water = ptl.data.get_molecule("water_dimer_minima.psimol")
ret = storage_socket.add_molecules([water])
assert ret["meta"]["success"] is True
assert ret["meta"]["n_inserted"] == 1
data1 = {
"keywords": None,
"program": "p9",
"qc_spec": {
"basis": "b1",
"program": "p1",
"method": "m1",
"driver": "energy"
},
"status": "COMPLETE",
"protocols": {},
}
torj_proc = TorsionDriveProcedureORM(**data1)
session.add(torj_proc)
session.commit()
# Add optimization_history
data1["initial_molecule"] = ret["data"][0]
opt_proc = OptimizationProcedureORM(**data1)
opt_proc2 = OptimizationProcedureORM(**data1)
session.add_all([opt_proc, opt_proc2])
session.commit()
assert opt_proc.id
opt_hist = OptimizationHistory(torsion_id=torj_proc.id,
opt_id=opt_proc.id,
key="20")
opt_hist2 = OptimizationHistory(torsion_id=torj_proc.id,
opt_id=opt_proc2.id,
key="20")
torj_proc.optimization_history_obj = [opt_hist, opt_hist2]
session.commit()
torj_proc = session.query(TorsionDriveProcedureORM).options(
joinedload("optimization_history_obj")).first()
assert torj_proc.optimization_history == {
"20": [str(opt_proc.id), str(opt_proc2.id)]
}
# clean up
session_delete_all(session, OptimizationProcedureORM)
def test_torsiondrive_procedure(storage_socket, session):
"""
Torsiondrive procedure
"""
assert session.query(TorsionDriveProcedureORM).count() == 0
# assert Keywords.objects().count() == 0
# molecules = MoleculeORM.objects(molecular_formula='H4O2')
# assert molecules.count() == 2
data1 = {
# "molecule": molecules[0],
"keywords": None,
"program": "p9",
"qc_spec": {
"basis": "b1",
"program": "p1",
"method": "m1",
"driver": "energy"
},
"status": "COMPLETE",
}
torj_proc = TorsionDriveProcedureORM(**data1)
session.add(torj_proc)
session.commit()
# Add optimization_history
opt_proc = OptimizationProcedureORM(**data1)
opt_proc2 = OptimizationProcedureORM(**data1)
session.add_all([opt_proc, opt_proc2])
session.commit()
assert opt_proc.id
opt_hist = OptimizationHistory(torsion_id=torj_proc.id,
opt_id=opt_proc.id,
key='20')
opt_hist2 = OptimizationHistory(torsion_id=torj_proc.id,
opt_id=opt_proc2.id,
key='20')
torj_proc.optimization_history_obj = [opt_hist, opt_hist2]
session.commit()
torj_proc = session.query(TorsionDriveProcedureORM).options(
joinedload('optimization_history_obj')).first()
assert torj_proc.optimization_history == {
'20': [str(opt_proc.id), str(opt_proc2.id)]
}
# clean up
session_delete_all(session, OptimizationProcedureORM)
session_delete_all(session, TorsionDriveProcedureORM)
def test_optimization_procedure(storage_socket, session, molecules_H4O2):
"""
Optimization procedure
"""
assert session.query(OptimizationProcedureORM).count() == 0
# assert Keywords.objects().count() == 0
data1 = {
"procedure": "optimization",
"initial_molecule": molecules_H4O2[0],
"keywords": None,
"program": "p7",
"qc_spec": {
"basis": "b1",
"program": "p1",
"method": "m1",
"driver": "energy"
},
"status": "COMPLETE",
"protocols": {},
}
result1 = {
"procedure": "single",
"molecule": molecules_H4O2[0],
"method": "m1",
"basis": "b1",
"keywords": None,
"program": "p1",
"driver": "energy",
"status": "COMPLETE",
"protocols": {},
}
procedure = OptimizationProcedureORM(**data1)
session.add(procedure)
session.commit()
proc = session.query(OptimizationProcedureORM).options(
joinedload("initial_molecule_obj")).first()
assert proc.initial_molecule_obj.molecular_formula == "H4O2"
assert proc.procedure == "optimization"
# add a trajectory result
result = ResultORM(**result1)
session.add(result)
session.commit()
assert result.id
# link result to the trajectory
proc.trajectory_obj = [Trajectory(opt_id=proc.id, result_id=result.id)]
session.commit()
proc = session.query(OptimizationProcedureORM).options(
joinedload("trajectory_obj")).first()
assert proc.trajectory_obj
# clean up
session_delete_all(session, ResultORM)
session_delete_all(session, OptimizationProcedureORM)
def test_services(storage_socket, session):
assert session.query(OptimizationProcedureORM).count() == 0
water = ptl.data.get_molecule("water_dimer_minima.psimol")
ret = storage_socket.add_molecules([water])
assert ret["meta"]["success"] is True
assert ret["meta"]["n_inserted"] == 1
proc_data = {
"initial_molecule": ret["data"][0],
"keywords": None,
"program": "p7",
"qc_spec": {
"basis": "b1",
"program": "p1",
"method": "m1",
"driver": "energy"
},
"status": "COMPLETE",
"protocols": {},
}
service_data = {
"tag": "tag1 tag2",
"hash_index": "123",
"status": "COMPLETE",
"optimization_program": "gaussian",
# extra fields
"torsiondrive_state": {},
"dihedral_template": "1", # Not realistic?
"optimization_template": "2", # Not realistic?
"molecule_template": "",
"logger": None,
"storage_socket": storage_socket,
"task_priority": 0,
}
procedure = OptimizationProcedureORM(**proc_data)
session.add(procedure)
session.commit()
assert procedure.id
service_pydantic = TorsionDriveService(**service_data)
doc = ServiceQueueORM(**service_pydantic.dict(
include=set(ServiceQueueORM.__dict__.keys())))
doc.extra = service_pydantic.dict(
exclude=set(ServiceQueueORM.__dict__.keys()))
doc.procedure_id = procedure.id
doc.priority = doc.priority.value # Special case where we need the value not the enum
session.add(doc)
session.commit()
session.delete(doc)
session.delete(procedure)
session.commit()
assert session.query(ServiceQueueORM).count() == 0
def test_services(storage_socket, session):
assert session.query(OptimizationProcedureORM).count() == 0
proc_data = {
"initial_molecule": None,
"keywords": None,
"program": "p7",
"qc_spec": {
"basis": "b1",
"program": "p1",
"method": "m1",
"driver": "energy"
},
"status": "COMPLETE",
}
service_data = {
"tag": "tag1 tag2",
"hash_index": "123",
"status": "COMPLETE",
"optimization_program": "gaussian",
# extra fields
"torsiondrive_state": {},
"dihedral_template": "1",
"optimization_template": "2",
"molecule_template": "",
"logger": None,
"storage_socket": storage_socket,
"task_priority": 0
}
procedure = OptimizationProcedureORM(**proc_data)
session.add(procedure)
session.commit()
assert procedure.id
service_pydantic = TorsionDriveService(**service_data)
doc = ServiceQueueORM(**service_pydantic.dict(
include=set(ServiceQueueORM.__dict__.keys())))
doc.extra = service_pydantic.dict(
exclude=set(ServiceQueueORM.__dict__.keys()))
doc.procedure_id = procedure.id
doc.priority = doc.priority.value # Special case where we need the value not the enum
session.add(doc)
session.commit()
session.delete(doc)
session.delete(procedure)
session.commit()
assert session.query(ServiceQueueORM).count() == 0