From a16b57e1ba545cc21be269c9c7dfc9385c504ada Mon Sep 17 00:00:00 2001
From: jmmshn <jmmshn@gmail.com>
Date: Thu, 9 Jan 2025 21:51:49 -0800
Subject: [PATCH] fixed names for insertion jobs

The old way was pretty messy.  Refactored to be cleaner.
---
 src/atomate2/common/flows/electrode.py |  6 +++++-
 src/atomate2/common/jobs/electrode.py  | 13 ++++++------
 tests/vasp/flows/test_electrode.py     | 28 +++++++++++++-------------
 3 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/src/atomate2/common/flows/electrode.py b/src/atomate2/common/flows/electrode.py
index cb6a17ee4a..8a4f62405d 100644
--- a/src/atomate2/common/flows/electrode.py
+++ b/src/atomate2/common/flows/electrode.py
@@ -120,6 +120,10 @@ def make(
             relax = self.bulk_relax_maker.make(structure)
         else:
             relax = self.relax_maker.make(structure)
+
+        _shown_steps = str(n_steps) if n_steps else "inf"
+        relax.append_name(f" 0/{_shown_steps}")
+
         # add ignored_species to the structure matcher
         sm = _add_ignored_species(self.structure_matcher, inserted_element)
         # Get the inserted structure
@@ -131,7 +135,7 @@ def make(
             relax_maker=self.relax_maker,
             get_charge_density=self.get_charge_density,
             n_steps=n_steps,
-            insertions_per_step=insertions_per_step,
+            n_inserted=1,
         )
         relaxed_summary = RelaxJobSummary(
             structure=relax.output.structure,
diff --git a/src/atomate2/common/jobs/electrode.py b/src/atomate2/common/jobs/electrode.py
index 6c0ad5f2cd..9f6f05a86c 100644
--- a/src/atomate2/common/jobs/electrode.py
+++ b/src/atomate2/common/jobs/electrode.py
@@ -86,12 +86,14 @@ def get_stable_inserted_results(
     """
     if structure is None:
         return []
-    if n_steps is not None and n_steps <= 0:
+    if n_inserted > n_steps if n_steps is not None else float("inf"):
         return []
     # append job name
-    add_name = f"{n_inserted}"
+    _shown_steps = str(n_steps) if n_steps else "inf"
+    add_name = f"{n_inserted}/{_shown_steps}"
 
     static_job = static_maker.make(structure=structure)
+    static_job.append_name(f" {n_inserted - 1}/{_shown_steps}")
     insertion_job = get_inserted_structures(
         static_job.output.dir_name,
         get_charge_density,
@@ -107,7 +109,6 @@ def get_stable_inserted_results(
         ref_structure=structure,
         structure_matcher=structure_matcher,
     )
-    nn_step = n_steps - 1 if n_steps is not None else None
     next_step = get_stable_inserted_results(
         structure=min_en_job.output[0],
         inserted_element=inserted_element,
@@ -116,12 +117,12 @@ def get_stable_inserted_results(
         relax_maker=relax_maker,
         get_charge_density=get_charge_density,
         insertions_per_step=insertions_per_step,
-        n_steps=nn_step,
+        n_steps=n_steps,
         n_inserted=n_inserted + 1,
     )
 
-    for job_ in [static_job, insertion_job, min_en_job, relax_jobs, next_step]:
-        job_.append_name(f" {add_name}")
+    # for job_ in [static_job, insertion_job, min_en_job, relax_jobs, next_step]:
+    #     job_.append_name(f" {add_name}")
     combine_job = get_computed_entries(next_step.output, min_en_job.output)
     replace_flow = Flow(
         jobs=[
diff --git a/tests/vasp/flows/test_electrode.py b/tests/vasp/flows/test_electrode.py
index 79ab6738b2..3b35049c2c 100644
--- a/tests/vasp/flows/test_electrode.py
+++ b/tests/vasp/flows/test_electrode.py
@@ -17,26 +17,26 @@ def test_electrode_makers(mock_vasp, clean_dir, test_dir):
 
     # mapping from job name to directory containing test files
     ref_paths = {
-        "relax": "H_Graphite/relax",
-        "relax 0 (0) 0": "H_Graphite/relax_0_(0)",
-        "relax 1 (0) 1 0": "H_Graphite/relax_1_(0)",
-        "relax 1 (1) 1 0": "H_Graphite/relax_1_(1)",
-        "relax 1 (2) 1 0": "H_Graphite/relax_1_(2)",
-        "static 0": "H_Graphite/static_0",
-        "static 1 0": "H_Graphite/static_1",
+        "relax 0/2": "H_Graphite/relax",
+        "relax 1/2 (0)": "H_Graphite/relax_0_(0)",
+        "relax 2/2 (0)": "H_Graphite/relax_1_(0)",
+        "relax 2/2 (1)": "H_Graphite/relax_1_(1)",
+        "relax 2/2 (2)": "H_Graphite/relax_1_(2)",
+        "static 0/2": "H_Graphite/static_0",
+        "static 1/2": "H_Graphite/static_1",
     }
 
     fake_run_vasp_kwargs = {
-        "relax": {
+        "relax 0/2": {
             "incar_settings": ["NSW", "ISIF"],
             "check_inputs": ["incar", "poscar"],
         },
-        "relax 0 (0) 0": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
-        "relax 1 (0) 1 0": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
-        "relax 1 (1) 1 0": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
-        "relax 1 (2) 1 0": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
-        "static 0": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
-        "static 1 0": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
+        "relax 1/2 (0)": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
+        "relax 2/2 (0)": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
+        "relax 2/2 (1)": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
+        "relax 2/2 (2)": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
+        "static 0/2": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
+        "static 1/2": {"incar_settings": ["NSW"], "check_inputs": ["incar"]},
     }
 
     # automatically use fake VASP and write POTCAR.spec during the test