Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
# If no fragments explicitly shown, grab all
if fragments is None:
fragments = self.data.fragments.keys()
# Get the data if available
self.get_fragment_data(fragments=fragments, refresh_cache=refresh_cache)
ret = {}
for frag in fragments:
tmp = {}
for k, v in self.data.fragments[frag].items():
if v in self._torsiondrive_cache:
obj = self._torsiondrive_cache[v]
if isinstance(obj, TorsionDriveRecord):
tmp[k] = obj.get_final_molecules()
elif isinstance(obj, OptimizationRecord):
tmp[k] = obj.get_final_molecule()
else:
raise TypeError("Internal type error encoured, buy a dev a coffee.")
else:
tmp[k] = None
ret[frag] = tmp
return ret
print('\nCurrent skip={}\n-----------'.format(skip))
ret = mongo_storage.get_procedures(procedure='optimization', status=None, limit=max_limit, skip=skip)
mongo_res= ret['data']
print('mongo results returned: ', len(mongo_res), ', total: ', ret['meta']['n_found'])
# check if this patch has been already stored
if is_mapped(sql_storage, ProcedureMap, mongo_res[-1]['id']):
print('Skipping first ', skip+max_limit)
continue
# load mapped ids in memory
mongo_res = get_ids_map(sql_storage, ['initial_molecule', 'final_molecule'], MoleculeMap, mongo_res)
mongo_res = get_ids_map(sql_storage, ['stdout', 'stderr', 'error'], KVStoreMap, mongo_res)
mongo_res = get_ids_map(sql_storage, ['trajectory'], ResultMap, mongo_res)
results_py = [OptimizationRecord(**res) for res in mongo_res]
sql_insered = sql_storage.add_procedures(results_py)['data']
print('Inserted in SQL:', len(sql_insered))
# store the ids mapping in the sql DB
mongo_ids = [obj['id'] for obj in mongo_res]
store_ids_map(sql_storage, mongo_ids, sql_insered, ProcedureMap)
if with_check:
with sql_storage.session_scope() as session:
proc = session.query(ProcedureMap).filter_by(mongo_id=mongo_res[0]['id']).first().sql_id
ret = sql_storage.get_procedures(id=[proc])
print('Get from SQL:', ret['data'])
ret2 = mongo_storage.get_procedures(id=[mongo_res[0]['id']])
print('Get from Mongo:', ret2['data'])
for rec in prod_proc:
if rec['initial_molecule']:
rec['initial_molecule'] = mols_map[rec['initial_molecule']]
if rec['final_molecule']:
rec['final_molecule'] = mols_map[rec['final_molecule']]
if rec['trajectory']:
rec['trajectory'] = [results_map[i] for i in rec['trajectory']]
if rec['stdout']:
rec['stdout'] = kvstore_map[rec['stdout']]
if rec['stderr']:
rec['stderr'] = kvstore_map[rec['stderr']]
if rec['error']:
rec['error'] = kvstore_map[rec['error']]
procedures_py = [OptimizationRecord(**proc) for proc in prod_proc]
staging_ids = staging_storage.add_procedures(procedures_py)['data']
if VERBOSE:
print('Inserted in SQL:', len(staging_ids))
print('---- Done copying Optimization procedures\n\n')
return {m1: m2 for m1, m2 in zip(procedure_ids, staging_ids)}
def parse_output(self, opt_outputs):
"""Save the results of the procedure.
It must make sure to save the results in the results table
including the task_id in the TaskQueue table
"""
completed_tasks = []
updates = []
for output in opt_outputs:
rec = self.storage.get_procedures(id=output["base_result"].id)["data"][0]
rec = OptimizationRecord(**rec)
procedure = output["result"]
# Add initial and final molecules
update_dict = {}
initial_mol, final_mol = self.storage.add_molecules(
[Molecule(**procedure["initial_molecule"]), Molecule(**procedure["final_molecule"])]
)["data"]
assert initial_mol == rec.initial_molecule
update_dict["final_molecule"] = final_mol
# Parse trajectory computations and add task_id
traj_dict = {k: v for k, v in enumerate(procedure["trajectory"])}
results = parse_single_tasks(self.storage, traj_dict)
for k, v in results.items():
v["task_id"] = output["task_id"]
def get_procedure_class(record):
if isinstance(record, OptimizationRecord):
procedure_class = OptimizationProcedureORM
elif isinstance(record, TorsionDriveRecord):
procedure_class = TorsionDriveProcedureORM
elif isinstance(record, GridOptimizationRecord):
procedure_class = GridOptimizationProcedureORM
else:
raise TypeError("Procedure of type {} is not valid or supported yet.".format(type(record)))
return procedure_class
results[k] = ResultRecord(**v)
ret = self.storage.add_results(list(results.values()))
update_dict["trajectory"] = ret["data"]
update_dict["energies"] = procedure["energies"]
# Save stdout/stderr
stdout, stderr, error = self.storage.add_kvstore(
[procedure["stdout"], procedure["stderr"], procedure["error"]]
)["data"]
update_dict["stdout"] = stdout
update_dict["stderr"] = stderr
update_dict["error"] = error
update_dict["provenance"] = procedure["provenance"]
rec = OptimizationRecord(**{**rec.dict(), **update_dict})
updates.append(rec)
completed_tasks.append(output["task_id"])
self.storage.update_procedures(updates)
return completed_tasks, [], []