From 6aa1797a4c9a65f662f9284181ab0bed994f4022 Mon Sep 17 00:00:00 2001
From: Sudarshan Vijay <sudarshan.vijay@vasp.at>
Date: Sun, 24 Mar 2024 20:44:47 +0100
Subject: [PATCH] add is_converged for the electronic steps

---
 src/py4vasp/_raw/data.py            | 2 ++
 src/py4vasp/_raw/definition.py      | 1 +
 src/py4vasp/calculation/_OSZICAR.py | 9 +++++++++
 tests/calculation/test_oszicar.py   | 7 +++++++
 tests/conftest.py                   | 3 ++-
 5 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/src/py4vasp/_raw/data.py b/src/py4vasp/_raw/data.py
index f7f81ff1..d5e291fa 100644
--- a/src/py4vasp/_raw/data.py
+++ b/src/py4vasp/_raw/data.py
@@ -324,6 +324,8 @@ class OSZICAR:
     "All columns of the OSZICAR file stored for all ionic steps."
     label: VaspData
     "Label of all the data from the OSZICAR file."
+    EDIFF: VaspData
+    "EDIFF decides the energy converge of the SCF step."
 
 
 @dataclasses.dataclass
diff --git a/src/py4vasp/_raw/definition.py b/src/py4vasp/_raw/definition.py
index 59b16361..98aebf6b 100644
--- a/src/py4vasp/_raw/definition.py
+++ b/src/py4vasp/_raw/definition.py
@@ -383,6 +383,7 @@ def selections(quantity):
     required=raw.Version(6, 5),
     label="intermediate/ion_dynamics/oszicar_label",
     convergence_data="intermediate/ion_dynamics/oszicar",
+    EDIFF="/intermediate/ion_dynamics/EDIFF",
 )
 #
 group = "intermediate/pair_correlation"
diff --git a/src/py4vasp/calculation/_OSZICAR.py b/src/py4vasp/calculation/_OSZICAR.py
index 87f3d544..b78f2d8c 100644
--- a/src/py4vasp/calculation/_OSZICAR.py
+++ b/src/py4vasp/calculation/_OSZICAR.py
@@ -115,3 +115,12 @@ def to_graph(self, selection="E"):
             xlabel="Iteration number",
             ylabel=ylabel,
         )
+
+    @_base.data_access
+    def is_converged(self):
+        difference_energy = self._read(b"dE")
+        if self._more_than_one_ionic_step(difference_energy):
+            last_step_energy = [dE[-1] for dE in difference_energy]
+        else:
+            last_step_energy = difference_energy[-1]
+        return last_step_energy < self._raw_data.EDIFF
diff --git a/tests/calculation/test_oszicar.py b/tests/calculation/test_oszicar.py
index 18d4cac9..02a35fd8 100644
--- a/tests/calculation/test_oszicar.py
+++ b/tests/calculation/test_oszicar.py
@@ -22,6 +22,7 @@ def OSZICAR(raw_data):
     oszicar.ref.ncg = convergence_data[:, 4]
     oszicar.ref.rms = convergence_data[:, 5]
     oszicar.ref.rmsc = convergence_data[:, 6]
+    oszicar.ref.EDIFF = raw_oszicar.EDIFF
     string_rep = "N\t\tE\t\tdE\t\tdeps\t\tncg\trms\t\trms(c)\n"
     format_rep = "{0:g}\t{1:0.12E}\t{2:0.6E}\t{3:0.6E}\t{4:g}\t{5:0.3E}\t{6:0.3E}\n"
     for idx in range(len(convergence_data)):
@@ -54,3 +55,9 @@ def test_plot(OSZICAR, Assert):
 def test_print(OSZICAR, format_):
     actual, _ = format_(OSZICAR)
     assert actual["text/plain"] == OSZICAR.ref.string_rep
+
+
+def test_is_converged(OSZICAR, Assert):
+    actual = OSZICAR.is_converged()
+    expected = OSZICAR.ref.dE[-1] < OSZICAR.ref.EDIFF
+    assert actual == expected
diff --git a/tests/conftest.py b/tests/conftest.py
index fdc1b490..dd8bd451 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -663,7 +663,8 @@ def _example_OSZICAR():
     )
     convergence_data = raw.VaspData(convergence_data)
     label = raw.VaspData([b"N", b"E", b"dE", b"deps", b"ncg", b"rms", b"rms(c)"])
-    return raw.OSZICAR(convergence_data=convergence_data, label=label)
+    ediff = 0.5
+    return raw.OSZICAR(convergence_data=convergence_data, label=label, EDIFF=ediff)
 
 
 def _Sr2TiO4_CONTCAR():