fixed names for insertion jobs

The old way was pretty messy. Refactored to be cleaner.
esoteric-ephemera · Jan 10, 2025 · a16b57e · a16b57e
1 parent 5196e93
commit a16b57e
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 21 deletions.
diff --git 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
@@ -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
@@ -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