Feat: Add OSZICAR information to HDF5 (#136)

* Create raw and schema addition for OSZICAR
* Implement hdf5 to py4vasp OSZICAR _data
* Implement a to_graph method for the OSZICAR data

Author: sudarshanv01

src/py4vasp/_raw/data.py

@@ -312,6 +312,18 @@ class Magnetism:
     "Contains the orbital magnetization for all atoms"
 
 
+@dataclasses.dataclass
+class OSZICAR:
+    """The OSZICAR data as generated by VASP.
+
+    All data generated by VASP and traditionally stored in the OSZICAR file will be
+    stored here. See https://www.vasp.at/wiki/index.php/OSZICAR for more details about
+    what quantities to expect."""
+
+    convergence_data: VaspData
+    "All columns of the OSZICAR file stored for all ionic steps."
+
+
 @dataclasses.dataclass
 class PairCorrelation:
     """The pair-correlation function calculated during a MD simulation.

src/py4vasp/_raw/definition.py

@@ -378,6 +378,12 @@ def selections(quantity):
     orbital_moments="intermediate/ion_dynamics/magnetism/orbital_moments/values",
 )
 #
+schema.add(
+    raw.OSZICAR,
+    required=raw.Version(6, 5),
+    convergence_data="intermediate/ion_dynamics/oszicar",
+)
+#
 group = "intermediate/pair_correlation"
 schema.add(
     raw.PairCorrelation,

src/py4vasp/_util/convert.py

@@ -21,7 +21,7 @@ def to_complex(array):
 
 
 def quantity_name(quantity):
-    if quantity == "CONTCAR":
+    if quantity in ["CONTCAR", "OSZICAR"]:
         return quantity
     else:
         return _to_snakecase(quantity)

src/py4vasp/calculation/_OSZICAR.py

@@ -0,0 +1,103 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+import numpy as np
+
+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,
+}
+
+
+class OSZICAR(_slice.Mixin, _base.Refinery, graph.Mixin):
+    """Access the convergence data for each electronic step.
+
+    The OSZICAR file written out by VASP stores information related to convergence.
+    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."""
+
+    @_base.data_access
+    def to_dict(self, selection=None):
+        """Extract convergence data from the HDF5 file and make it available in a dict
+
+        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, supply all columns.
+
+        Returns
+        -------
+        dict
+            Contains a dict from the HDF5 related to OSZICAR convergence data
+        """
+        return_data = {}
+        if selection is None:
+            keys_to_include = INDEXING_OSZICAR
+        else:
+            if keys_to_include not in INDEXING_OSZICAR:
+                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."""
+                raise exception.RefinementError(message)
+            keys_to_include = selection
+        for key in INDEXING_OSZICAR:
+            return_data[key] = self._read(key)
+        return return_data
+
+    def _read(self, key):
+        # data represents all of the electronic steps for all ionic steps
+        data = getattr(self._raw_data, "convergence_data")
+        iteration_number = data[:, 0]
+        split_index = np.where(iteration_number == 1)[0]
+        data = np.vsplit(data, split_index)[1:][self._steps]
+        if isinstance(self._steps, slice):
+            data = [raw.VaspData(_data) for _data in data]
+        else:
+            data = [raw.VaspData(data)]
+        data_index = INDEXING_OSZICAR[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"):
+        """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.
+
+        Returns
+        -------
+        Graph
+            The Graph with the quantity plotted on y-axis and the iteration number of
+            the x-axis.
+        """
+        data = self.to_dict()
+        series = graph.Series(data["iteration_number"], data[selection], selection)
+        ylabel = " ".join(select.capitalize() for select in selection.split("_"))
+        return graph.Graph(
+            series=[series],
+            xlabel="Iteration number",
+            ylabel=ylabel,
+        )

src/py4vasp/calculation/__init__.py

@@ -60,6 +60,7 @@ class provides a more flexible interface with which you can determine the source
     "internal_strain",
     "kpoint",
     "magnetism",
+    "OSZICAR",
     "pair_correlation",
     "phonon_band",
     "phonon_dos",

tests/calculation/test_oszicar.py

@@ -0,0 +1,52 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+
+import types
+
+import pytest
+
+from py4vasp import calculation
+
+
+@pytest.fixture
+def OSZICAR(raw_data):
+    raw_oszicar = raw_data.OSZICAR()
+    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]
+    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
+    )
+
+
+def test_plot(OSZICAR, Assert):
+    graph = OSZICAR.plot()
+    assert graph.xlabel == "Iteration number"
+    assert graph.ylabel == "Free Energy"
+    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)

tests/conftest.py

@@ -123,6 +123,10 @@ def CONTCAR(selection):
         else:
             raise exception.NotImplemented()
 
+    @staticmethod
+    def OSZICAR(selection=None):
+        return _example_OSZICAR()
+
     @staticmethod
     def density(selection):
         parts = selection.split()
@@ -649,6 +653,14 @@ def _Sr2TiO4_cell():
     )
 
 
+def _example_OSZICAR():
+    random_convergence_data = np.random.rand(9, 6)
+    iteration_number = np.arange(1, 10)[:, np.newaxis]
+    convergence_data = np.hstack([iteration_number, random_convergence_data])
+    convergence_data = raw.VaspData(convergence_data)
+    return raw.OSZICAR(convergence_data=convergence_data)
+
+
 def _Sr2TiO4_CONTCAR():
     structure = _Sr2TiO4_structure()
     structure.cell.lattice_vectors = structure.cell.lattice_vectors[-1]