def setUp(self):
# This is a CHGCAR_sum file with reduced grid size
chgcar_path = os.path.join(test_dir, "CHGCAR.FePO4")
chg_FePO4 = Chgcar.from_file(chgcar_path)
self.chgcar_path = chgcar_path
self.chg_FePO4 = chg_FePO4
self.cia_FePO4 = ChargeInsertionAnalyzer(chg_FePO4)
def run_task(self, fw_spec):
attempt_insertions = fw_spec.get("attempt_insertions", 4)
volumetric_data_type = fw_spec.get("volumetric_data_type")
base_task_id = fw_spec.get("base_task_id")
logger.info(f"Identifying sites for task : {base_task_id}")
cia_kwargs = fw_spec.get("ChargeInsertionAnalyzer_kwargs", dict())
# get the database connection
db_file = env_chk(DB_FILE, fw_spec)
logger.info(f"DB_FILE: {db_file}")
mmdb = VaspCalcDb.from_db_file(db_file, admin=True)
if volumetric_data_type == "CHGCAR":
chgcar = mmdb.get_chgcar(task_id=base_task_id)
elif volumetric_data_type == "AECCAR":
chgcar = mmdb.get_aeccar(task_id=base_task_id)
chgcar = chgcar["aeccar0"] + chgcar["aeccar2"]
cia = ChargeInsertionAnalyzer(chgcar, **cia_kwargs)
cia.get_labels()
insert_sites = []
seent = set()
cia._extrema_df.sort_values(by=["avg_charge_den"], inplace=True)
for itr, li_site in cia._extrema_df.iterrows():
if len(insert_sites) >= attempt_insertions:
break
li_site = cia._extrema_df.iloc[itr]
lab = li_site["site_label"]
if lab not in seent:
insert_sites.append([li_site["a"], li_site["b"], li_site["c"]])
seent.add(lab)
logger.info(
f"Found {len(insert_sites)} insertion sites for task : {base_task_id}"
)
# Since we are only analyzing a single charge density
# The we will only perform one set of atoms insertions at a given time
# Thus, we just have to update the global fw_spec
# (no need to pass this information from fw to fw)
return FWAction(
update_spec={
"insert_sites": insert_sites,
"base_task_id": base_task_id,
"base_structure": chgcar.structure.as_dict(),
})
class TestChargeInsertionAnalyzer(unittest.TestCase):
def setUp(self):
# This is a CHGCAR_sum file with reduced grid size
chgcar_path = os.path.join(test_dir, "CHGCAR.FePO4")
chg_FePO4 = Chgcar.from_file(chgcar_path)
self.chgcar_path = chgcar_path
self.chg_FePO4 = chg_FePO4
self.cia_FePO4 = ChargeInsertionAnalyzer(chg_FePO4)
def test_get_label(self):
self.cia_FePO4.get_labels()
self.assertEqual(len(self.cia_FePO4._extrema_df), 12)
counts = Counter(self.cia_FePO4._extrema_df.site_label)
self.assertEqual(counts, {0: 4, 1: 4, 2: 4})
def test_filter(self):
filtered_FePO4 = ChargeInsertionAnalyzer(self.chg_FePO4, max_avg_charge=0.05)
filtered_FePO4.get_labels()
counts = Counter(filtered_FePO4._extrema_df.site_label)
self.assertEqual(counts, {0: 4, 1: 4})
def test_filter(self):
filtered_FePO4 = ChargeInsertionAnalyzer(self.chg_FePO4, max_avg_charge=0.05)
filtered_FePO4.get_labels()
counts = Counter(filtered_FePO4._extrema_df.site_label)
self.assertEqual(counts, {0: 4, 1: 4})