def test_chgcar_db_read_write(self):
# generate a doc from the test folder
drone = VaspDrone(parse_chgcar=True, parse_aeccar=True)
print(ref_dirs_si['static'])
doc = drone.assimilate(ref_dirs_si['static']+'/outputs')
# insert the doc make sure that the
cc = doc['calcs_reversed'][0]['chgcar']
self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 8.0, 4)
cc = doc['calcs_reversed'][0]['aeccar0']
self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 23.253588293583313, 4)
cc = doc['calcs_reversed'][0]['aeccar2']
self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 8.01314480789829, 4)
mmdb = VaspCalcDb.from_db_file(os.path.join(db_dir, "db.json"))
t_id = mmdb.insert_task(doc, use_gridfs=True)
# space is freed up after uploading the document
self.assertRaises(KeyError, lambda: doc['calcs_reversed'][0]['chgcar'])
self.assertRaises(KeyError, lambda: doc['calcs_reversed'][0]['aeccar0'])
self.assertRaises(KeyError, lambda: doc['calcs_reversed'][0]['aeccar2'])
cc = mmdb.get_chgcar(task_id=t_id)
self.assertAlmostEqual(cc.data['total'].sum()/cc.ngridpts, 8.0, 4)
dcc = mmdb.get_aeccar(task_id=t_id)
self.assertAlmostEqual(dcc['aeccar0'].data['total'].sum()/cc.ngridpts, 23.253588293583313, 4)
self.assertAlmostEqual(dcc['aeccar2'].data['total'].sum()/cc.ngridpts, 8.01314480789829, 4)
# check the retrieve_task function for the same fake calculation
ret_task = mmdb.retrieve_task(t_id)
ret_chgcar = ret_task['calcs_reversed'][0]['chgcar']
ret_aeccar0 = ret_task['calcs_reversed'][0]['aeccar0']
ret_aeccar2 = ret_task['calcs_reversed'][0]['aeccar2']
ret_aeccar = ret_aeccar0 + ret_aeccar2
self.assertAlmostEqual(ret_chgcar.data['total'].sum()/ret_chgcar.ngridpts, 8.0, 4)
self.assertAlmostEqual(ret_aeccar.data['total'].sum()/ret_aeccar.ngridpts, 31.2667331015, 4)
def test_chgcar_db_read_write_maggma(self):
# generate a doc from the test folder
drone = VaspDrone(parse_chgcar=True, parse_aeccar=True)
doc = drone.assimilate(ref_dirs_si['static'] + '/outputs')
mmdb = VaspCalcDb.from_db_file(os.path.join(db_dir, "db_aws.json"))
with mock_s3():
conn = boto3.client("s3")
conn.create_bucket(Bucket="test_bucket")
t_id = mmdb.insert_task(task_doc=doc)
# basic check that data was written to the stores
with mmdb._maggma_stores['chgcar_fs'] as store:
res = store.query_one()
self.assertTrue(res['data']['@class'] == "Chgcar")
with mmdb._maggma_stores['aeccar0_fs'] as store:
res = store.query_one()
self.assertTrue(res['data']['@class'] == "Chgcar")
with mmdb._maggma_stores['aeccar2_fs'] as store:
res = store.query_one()
self.assertTrue(res['data']['@class'] == "Chgcar")
with mmdb._maggma_stores['bandstructure_fs'] as store:
res = store.query_one()
self.assertTrue(res['data']['@class'] == "BandStructure")
# print(mmdb.collection.find_one({'task_id' : t_id})["calcs_reversed"][0].keys())
# complex check that the data is the same
res = mmdb.get_band_structure(task_id=t_id)
self.assertTrue(isinstance(res, BandStructure))
res = mmdb.get_dos(task_id=t_id)
self.assertTrue(isinstance(res, CompleteDos))
res = mmdb.get_chgcar(task_id=t_id)
self.assertTrue(isinstance(res, Chgcar))
res = mmdb.get_aeccar(task_id=t_id)
self.assertTrue(isinstance(res['aeccar0'], Chgcar))
self.assertTrue(isinstance(res['aeccar2'], Chgcar))