Improve to_converged (#154)

* Implement __str__ method for OSZICAR
* Add is_converged for the electronic steps
* Pass is_elmin_converged instead of inferring convergence
* Modify tests for OSZICAR

src/py4vasp/_raw/data.py

@@ -322,6 +322,10 @@ class OSZICAR:
 
     convergence_data: VaspData
     "All columns of the OSZICAR file stored for all ionic steps."
+    label: VaspData
+    "Label of all the data from the OSZICAR file."
+    is_elmin_converged: VaspData
+    "Is the electronic minimization step converged?"
 
 
 @dataclasses.dataclass

src/py4vasp/_raw/definition.py

@@ -381,7 +381,9 @@ def selections(quantity):
 schema.add(
     raw.OSZICAR,
     required=raw.Version(6, 5),
+    label="intermediate/ion_dynamics/oszicar_label",
     convergence_data="intermediate/ion_dynamics/oszicar",
+    is_elmin_converged="/intermediate/ion_dynamics/electronic_step_converged",
 )
 #
 group = "intermediate/pair_correlation"

src/py4vasp/calculation/_OSZICAR.py

@@ -5,18 +5,7 @@
 
 from py4vasp import exception, raw
 from py4vasp._third_party import graph
-from py4vasp._util import convert
-from py4vasp.calculation import _base, _slice, _structure
-
-INDEXING_OSZICAR = {
-    "iteration_number": 0,
-    "free_energy": 1,
-    "free_energy_change": 2,
-    "bandstructure_energy_change": 3,
-    "number_hamiltonian_evaluations": 4,
-    "norm_residual": 5,
-    "difference_charge_density": 6,
-}
+from py4vasp.calculation import _base, _slice
 
 
 class OSZICAR(_slice.Mixin, _base.Refinery, graph.Mixin):
@@ -26,6 +15,35 @@ class OSZICAR(_slice.Mixin, _base.Refinery, graph.Mixin):
     Please check the vasp-wiki (https://www.vasp.at/wiki/index.php/OSZICAR) for more
     details about the exact outputs generated for each combination of INCAR tags."""
 
+    def _more_than_one_ionic_step(self, data):
+        return any(isinstance(_data, list) for _data in data) == True
+
+    @_base.data_access
+    def __str__(self):
+        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"
+        label_rep = "{}\t\t{}\t\t{}\t\t{}\t\t{}\t{}\t\t{}\n"
+        string = ""
+        labels = [label.decode("utf-8") for label in getattr(self._raw_data, "label")]
+        data = self.to_dict()
+        electronic_iterations = data["N"]
+        if not self._more_than_one_ionic_step(electronic_iterations):
+            electronic_iterations = [electronic_iterations]
+        ionic_steps = len(electronic_iterations)
+        for ionic_step in range(ionic_steps):
+            string += label_rep.format(*labels)
+            electronic_steps = len(electronic_iterations[ionic_step])
+            for electronic_step in range(electronic_steps):
+                _data = []
+                for label in self._raw_data.label:
+                    _values_electronic = data[label.decode("utf-8")]
+                    if not self._more_than_one_ionic_step(_values_electronic):
+                        _values_electronic = [_values_electronic]
+                    _value = _values_electronic[ionic_step][electronic_step]
+                    _data.append(_value)
+                _data = [float(_value) for _value in _data]
+                string += format_rep.format(*_data)
+        return string
+
     @_base.data_access
     def to_dict(self, selection=None):
         """Extract convergence data from the HDF5 file and make it available in a dict
@@ -45,20 +63,23 @@ def to_dict(self, selection=None):
         """
         return_data = {}
         if selection is None:
-            keys_to_include = INDEXING_OSZICAR
+            keys_to_include = self._from_bytes_to_utf(self._raw_data.label)
         else:
-            if keys_to_include not in INDEXING_OSZICAR:
+            labels_as_str = self._from_bytes_to_utf(self._raw_data.label)
+            if selection not in labels_as_str:
                 message = """\
 Please choose a selection including at least one of the following keywords:
-iteration_number, free_energy, free_energy_change, bandstructure_energy_change,
-number_hamiltonian_evaluations, norm_residual, difference_charge_density. Else do not
-select anything and all OSZICAR outputs will be provided."""
+N, E, dE, deps, ncg, rms, rms(c)"""
                 raise exception.RefinementError(message)
-            keys_to_include = selection
-        for key in INDEXING_OSZICAR:
+            keys_to_include = [selection]
+        for key in keys_to_include:
             return_data[key] = self._read(key)
         return return_data
 
+    def _from_bytes_to_utf(self, quantity: list):
+        return [_quantity.decode("utf-8") for _quantity in quantity]
+
+    @_base.data_access
     def _read(self, key):
         # data represents all of the electronic steps for all ionic steps
         data = getattr(self._raw_data, "convergence_data")
@@ -69,23 +90,21 @@ def _read(self, key):
             data = [raw.VaspData(_data) for _data in data]
         else:
             data = [raw.VaspData(data)]
-        data_index = INDEXING_OSZICAR[key]
+        labels = [label.decode("utf-8") for label in self._raw_data.label]
+        data_index = labels.index(key)
         return_data = [list(_data[:, data_index]) for _data in data]
         is_none = [_data.is_none() for _data in data]
         if len(return_data) == 1:
             return_data = return_data[0]
         return return_data if not np.all(is_none) else {}
 
-    def to_graph(self, selection="free_energy"):
+    def to_graph(self, selection="E"):
         """Graph the change in parameter with iteration number.
 
         Parameters
         ----------
         selection: str
-            Choose from either iteration_number, free_energy, free_energy_change,
-            bandstructure_energy_change, number_hamiltonian_evaluations, norm_residual,
-            difference_charge_density to get specific columns of the OSZICAR file. In
-            case no selection is provided, the free energy is plotted.
+            Choose strings consistent with the OSZICAR format
 
         Returns
         -------
@@ -94,10 +113,18 @@ def to_graph(self, selection="free_energy"):
             the x-axis.
         """
         data = self.to_dict()
-        series = graph.Series(data["iteration_number"], data[selection], selection)
+        series = graph.Series(data["N"], data[selection], selection)
         ylabel = " ".join(select.capitalize() for select in selection.split("_"))
         return graph.Graph(
             series=[series],
             xlabel="Iteration number",
             ylabel=ylabel,
         )
+
+    @_base.data_access
+    def is_converged(self):
+        is_elmin_converged = self._raw_data.is_elmin_converged[self._steps]
+        converged = is_elmin_converged == 0
+        if isinstance(converged, bool):
+            converged = np.array([converged])
+        return converged.flatten()

tests/calculation/test_oszicar.py

@@ -3,9 +3,10 @@
 
 import types
 
+import numpy as np
 import pytest
 
-from py4vasp import calculation
+from py4vasp import calculation, exception
 
 
 @pytest.fixture
@@ -14,39 +15,75 @@ def OSZICAR(raw_data):
     oszicar = calculation.OSZICAR.from_data(raw_oszicar)
     oszicar.ref = types.SimpleNamespace()
     convergence_data = raw_oszicar.convergence_data
-    oszicar.ref.iteration_number = convergence_data[:, 0]
-    oszicar.ref.free_energy = convergence_data[:, 1]
-    oszicar.ref.free_energy_change = convergence_data[:, 2]
-    oszicar.ref.bandstructure_energy_change = convergence_data[:, 3]
-    oszicar.ref.number_hamiltonian_evaluations = convergence_data[:, 4]
-    oszicar.ref.norm_residual = convergence_data[:, 5]
-    oszicar.ref.difference_charge_density = convergence_data[:, 6]
+    oszicar.ref.N = np.int64(convergence_data[:, 0])
+    oszicar.ref.E = convergence_data[:, 1]
+    oszicar.ref.dE = convergence_data[:, 2]
+    oszicar.ref.deps = convergence_data[:, 3]
+    oszicar.ref.ncg = convergence_data[:, 4]
+    oszicar.ref.rms = convergence_data[:, 5]
+    oszicar.ref.rmsc = convergence_data[:, 6]
+    oszicar.ref.is_elmin_converged = [raw_oszicar.is_elmin_converged == [0.0]]
+    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)):
+        string_rep += format_rep.format(*convergence_data[idx])
+    oszicar.ref.string_rep = str(string_rep)
     return oszicar
 
 
 def test_read(OSZICAR, Assert):
     actual = OSZICAR.read()
     expected = OSZICAR.ref
-    Assert.allclose(actual["iteration_number"], expected.iteration_number)
-    Assert.allclose(actual["free_energy"], expected.free_energy)
-    Assert.allclose(actual["free_energy_change"], expected.free_energy_change)
-    Assert.allclose(
-        actual["bandstructure_energy_change"], expected.bandstructure_energy_change
-    )
-    Assert.allclose(
-        actual["number_hamiltonian_evaluations"],
-        expected.number_hamiltonian_evaluations,
-    )
-    Assert.allclose(actual["norm_residual"], expected.norm_residual)
-    Assert.allclose(
-        actual["difference_charge_density"], expected.difference_charge_density
-    )
+    Assert.allclose(actual["N"], expected.N)
+    Assert.allclose(actual["E"], expected.E)
+    Assert.allclose(actual["dE"], expected.dE)
+    Assert.allclose(actual["deps"], expected.deps)
+    Assert.allclose(actual["ncg"], expected.ncg)
+    Assert.allclose(actual["rms"], expected.rms)
+    Assert.allclose(actual["rms(c)"], expected.rmsc)
+
+
+@pytest.mark.parametrize(
+    "quantity_name", ["N", "E", "dE", "deps", "ncg", "rms", "rms(c)"]
+)
+def test_read_selection(quantity_name, OSZICAR, Assert):
+    actual = OSZICAR.read(quantity_name)
+    expected = getattr(OSZICAR.ref, quantity_name.replace("(", "").replace(")", ""))
+    Assert.allclose(actual[quantity_name], expected)
+
+
+def test_read_incorrect_selection(OSZICAR):
+    with pytest.raises(exception.RefinementError):
+        OSZICAR.read("forces")
+
+
+def test_slice(OSZICAR, Assert):
+    actual = OSZICAR[0:1].read()
+    expected = OSZICAR.ref
+    Assert.allclose(actual["N"], expected.N)
+    Assert.allclose(actual["E"], expected.E)
+    Assert.allclose(actual["dE"], expected.dE)
+    Assert.allclose(actual["deps"], expected.deps)
+    Assert.allclose(actual["ncg"], expected.ncg)
+    Assert.allclose(actual["rms"], expected.rms)
+    Assert.allclose(actual["rms(c)"], expected.rmsc)
 
 
 def test_plot(OSZICAR, Assert):
     graph = OSZICAR.plot()
     assert graph.xlabel == "Iteration number"
-    assert graph.ylabel == "Free Energy"
+    assert graph.ylabel == "E"
     assert len(graph.series) == 1
-    Assert.allclose(graph.series[0].x, OSZICAR.ref.iteration_number)
-    Assert.allclose(graph.series[0].y, OSZICAR.ref.free_energy)
+    Assert.allclose(graph.series[0].x, OSZICAR.ref.N)
+    Assert.allclose(graph.series[0].y, OSZICAR.ref.E)
+
+
+def test_print(OSZICAR, format_):
+    actual, _ = format_(OSZICAR)
+    assert actual["text/plain"] == OSZICAR.ref.string_rep
+
+
+def test_is_converged(OSZICAR):
+    actual = OSZICAR.is_converged()
+    expected = OSZICAR.ref.is_elmin_converged
+    assert actual == expected

tests/conftest.py

@@ -661,11 +661,21 @@ def _Sr2TiO4_cell():
 
 
 def _example_OSZICAR():
-    random_convergence_data = np.random.rand(9, 6)
+    random_convergence_data = np.random.rand(9, 3)
     iteration_number = np.arange(1, 10)[:, np.newaxis]
-    convergence_data = np.hstack([iteration_number, random_convergence_data])
+    ncg = np.random.randint(4, 10, (9, 1))
+    random_rms = np.random.rand(9, 2)
+    convergence_data = np.hstack(
+        [iteration_number, random_convergence_data, ncg, random_rms]
+    )
     convergence_data = raw.VaspData(convergence_data)
-    return raw.OSZICAR(convergence_data=convergence_data)
+    label = raw.VaspData([b"N", b"E", b"dE", b"deps", b"ncg", b"rms", b"rms(c)"])
+    is_elmin_converged = [0]
+    return raw.OSZICAR(
+        convergence_data=convergence_data,
+        label=label,
+        is_elmin_converged=is_elmin_converged,
+    )
 
 
 def _Sr2TiO4_CONTCAR():