Change how electrode flow retrieves charge density data

pyproject.toml

@@ -103,7 +103,7 @@ strict = [
     "jobflow==0.1.19",
     "lobsterpy==0.4.9",
     "mdanalysis==2.7.0",
-    "monty==2024.10.21",
+    "monty==2025.1.9",
     "mp-api==0.43.0",
     "numpy",
     "openmm-mdanalysis-reporter==0.1.0",
@@ -112,7 +112,7 @@ strict = [
     "pydantic-settings==2.7.0",
     "pydantic==2.9.2",
     "pymatgen-analysis-defects==2024.10.22",
-    "pymatgen==2024.11.13",
+    "pymatgen==2025.1.9",
     "python-ulid==3.0.0",
     "seekpath==2.1.0",
     "tblite==0.3.0; python_version < '3.12'",
@@ -177,6 +177,7 @@ exclude_lines = [
     '^\s*@overload( |$)',
     '^\s*assert False(,|$)',
     'if typing.TYPE_CHECKING:',
+    'if TYPE_CHECKING:',
 ]
 
 [tool.ruff]

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
@@ -132,6 +136,7 @@ def make(
             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,

src/atomate2/common/jobs/electrode.py

@@ -21,7 +21,7 @@
     from pymatgen.analysis.structure_matcher import StructureMatcher
     from pymatgen.core import Structure
     from pymatgen.entries.computed_entries import ComputedEntry
-    from pymatgen.io.vasp.outputs import VolumetricData
+    from pymatgen.io.common import VolumetricData
 
 
 logger = logging.getLogger(__name__)
@@ -84,17 +84,21 @@ def get_stable_inserted_results(
         The number of ions inserted so far, used to help assign a unique name to the
         different jobs.
     """
-    if structure is None:
-        return []
-    if n_steps is not None and n_steps <= 0:
+    if (
+        (structure is None)
+        or (n_steps is not None and n_steps <= 0)
+        or (n_inserted > n_steps)
+    ):
         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)
-    chg_job = get_charge_density_job(static_job.output.dir_name, get_charge_density)
+    static_job.append_name(f" {n_inserted - 1}/{_shown_steps}")
     insertion_job = get_inserted_structures(
-        chg_job.output,
+        static_job.output.dir_name,
+        get_charge_density,
         inserted_species=inserted_element,
         insertions_per_step=insertions_per_step,
     )
@@ -107,7 +111,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,17 +119,14 @@ 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, chg_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=[
             static_job,
-            chg_job,
             insertion_job,
             relax_jobs,
             min_en_job,
@@ -204,7 +204,8 @@ def get_insertion_electrode_doc(
 
 @job
 def get_inserted_structures(
-    chg: VolumetricData,
+    prev_dir: Path | str,
+    get_charge_density: Callable[[str | Path], VolumetricData],
     inserted_species: ElementLike,
     insertions_per_step: int = 4,
     charge_insertion_generator: ChargeInterstitialGenerator | None = None,
@@ -213,7 +214,8 @@ def get_inserted_structures(
 
     Parameters
     ----------
-    chg: The charge density.
+    prev_dir: The previous directory where the static calculation was performed.
+    get_charge_density: A function to get the charge density from a run directory.
     inserted_species: The species to insert.
     insertions_per_step: The maximum number of ion insertion sites to attempt.
     charge_insertion_generator: The charge insertion generator to use,
@@ -226,6 +228,7 @@ def get_inserted_structures(
     """
     if charge_insertion_generator is None:
         charge_insertion_generator = ChargeInterstitialGenerator()
+    chg = get_charge_density(prev_dir)
     gen = charge_insertion_generator.generate(chg, insert_species=[inserted_species])
     inserted_structures = [defect.defect_structure for defect in gen]
     return inserted_structures[:insertions_per_step]
@@ -297,22 +300,3 @@ def get_min_energy_summary(
         return None
 
     return min(topotactic_summaries, key=lambda x: x.entry.energy_per_atom)
-
-
-@job
-def get_charge_density_job(
-    prev_dir: Path | str,
-    get_charge_density: Callable,
-) -> VolumetricData:
-    """Get the charge density from a task document.
-
-    Parameters
-    ----------
-    prev_dir: The previous directory where the static calculation was performed.
-    get_charge_density: A function to get the charge density from a task document.
-
-    Returns
-    -------
-        The charge density.
-    """
-    return get_charge_density(prev_dir)

tests/abinit/conftest.py

@@ -126,13 +126,20 @@ def check_run_abi(ref_path: str | Path):
 
     user = AbinitInputFile.from_file("run.abi")
     assert user.ndtset == 1, f"'run.abi' has multiple datasets (ndtset={user.ndtset})."
-    with zopen(ref_path / "inputs" / "run.abi.gz") as file:
+    with zopen(ref_path / "inputs" / "run.abi.gz", "rt", encoding="utf-8") as file:
         ref_str = file.read()
-    ref = AbinitInputFile.from_string(ref_str.decode("utf-8"))
+    ref = AbinitInputFile.from_string(ref_str)
     # Ignore the pseudos as the directory depends on the pseudo root directory
-    diffs = user.get_differences(ref, ignore_vars=["pseudos"])
+    # The `get_differences` method doesn't take into account floating point
+    # precision - this check might fail, e.g., (real CI example):
+    # ```
+    # The variable 'tsmear' is different in the two files:
+    # - this file:  '0.0036749322175665 Ha'
+    # - other file: '0.0036749322175655 Ha'`
+    # ```
+    diffs = user.get_differences(ref, ignore_vars=["pseudos", "tsmear"])
     # TODO: should we still add some check on the pseudos here ?
-    assert diffs == [], "'run.abi' is different from reference."
+    assert diffs == [], f"'run.abi' is different from reference:\n{diffs}"
 
 
 def check_abinit_input_json(ref_path: str | Path):

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