From 37cad6ee8f9c4309564c01701ac7bd0b743c4f68 Mon Sep 17 00:00:00 2001
From: "@jmmshn" <jmmshn@gmail.com>
Date: Tue, 19 Dec 2023 12:22:39 -0800
Subject: [PATCH 1/2] update

---
 atomate/utils/database.py                 |  1 -
 atomate/vasp/firetasks/electrode_tasks.py | 17 ++++-------------
 setup.py                                  |  1 +
 3 files changed, 5 insertions(+), 14 deletions(-)

diff --git a/atomate/utils/database.py b/atomate/utils/database.py
index 33db98959..1c1fcbd88 100644
--- a/atomate/utils/database.py
+++ b/atomate/utils/database.py
@@ -57,7 +57,6 @@ def __init__(
                         "compress" : Whether compression is used
                         "endpoint_url" : the url used to access the S3 store
             maggma_store_prefix: when using maggma stores, you can set the prefix string.
-
             **kwargs:
         """
         if maggma_store_kwargs is None:
diff --git a/atomate/vasp/firetasks/electrode_tasks.py b/atomate/vasp/firetasks/electrode_tasks.py
index d94103d79..601d6ff79 100644
--- a/atomate/vasp/firetasks/electrode_tasks.py
+++ b/atomate/vasp/firetasks/electrode_tasks.py
@@ -98,21 +98,12 @@ def run_task(self, fw_spec):
             chgcar = chgcar["aeccar0"] + chgcar["aeccar2"]
 
         cia = ChargeInsertionAnalyzer(chgcar, **cia_kwargs)
-        cia.get_labels()
+        avg_chg_groups = list(cia.filter_and_group())
 
         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)
-
+        for _, g in avg_chg_groups[:attempt_insertions]:
+            insert_sites.append(g[0])
+        
         logger.info(
             f"Found {len(insert_sites)} insertion sites for task : {base_task_id}"
         )
diff --git a/setup.py b/setup.py
index 1f028bae2..b47d45780 100644
--- a/setup.py
+++ b/setup.py
@@ -56,6 +56,7 @@
                 "moto>=4.1.14",
                 "pytest-cov>=4.1.0",
                 "pytest>=7.4.0",
+                "pymatgen-analysis-defects>=2022.8.18",
             ],
         },
         classifiers=[

From 583dfaf83d7d8d320426e96671c2fabe406beffa Mon Sep 17 00:00:00 2001
From: "@jmmshn" <jmmshn@gmail.com>
Date: Tue, 19 Dec 2023 12:26:40 -0800
Subject: [PATCH 2/2] updated insertion with defect code

---
 atomate/vasp/firetasks/electrode_tasks.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/atomate/vasp/firetasks/electrode_tasks.py b/atomate/vasp/firetasks/electrode_tasks.py
index 601d6ff79..032904e63 100644
--- a/atomate/vasp/firetasks/electrode_tasks.py
+++ b/atomate/vasp/firetasks/electrode_tasks.py
@@ -102,7 +102,7 @@ def run_task(self, fw_spec):
 
         insert_sites = []
         for _, g in avg_chg_groups[:attempt_insertions]:
-            insert_sites.append(g[0])
+            insert_sites.append(g[ 0])
         
         logger.info(
             f"Found {len(insert_sites)} insertion sites for task : {base_task_id}"