WIP: density contour

[skip CI]

src/py4vasp/calculation/_density.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from py4vasp import _config, calculation, exception
-from py4vasp._third_party import view
+from py4vasp._third_party import graph, view
 from py4vasp._util import documentation, import_, index, select
 from py4vasp.calculation import _base, _structure
 
@@ -225,6 +225,12 @@ def to_view(self, selection=None, supercell=None, **user_options):
         ]
         return viewer
 
+    @_base.data_access
+    def to_contour(self, *, a=None):
+        index = np.round(a * len(self._raw_data.charge)).astype(np.int_)
+        contour = graph.Contour(self._raw_data.charge[index], "charge")
+        return graph.Graph(contour)
+
     def _filter_noncollinear_magnetization_from_selections(self, tree):
         if self._selection or not self.is_noncollinear():
             yield from tree.selections()

tests/calculation/test_density.py

@@ -255,6 +255,18 @@ def test_plotting_supercell(supercell, reference_density, Assert):
     check_view(reference_density, expected, Assert, supercell=supercell)
 
 
+def test_contour_of_slice(nonpolarized_density, Assert):
+    graph = nonpolarized_density.to_contour(a=0.1)
+    slice_ = nonpolarized_density.ref.output["charge"][1]
+    assert len(graph) == 1
+    Assert.allclose(graph.series.data, slice_)
+
+
+# TODO: a, b, c
+#       slice missing
+#       a < 0 or a > 1
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":
@@ -307,4 +319,5 @@ def test_print(reference_density, format_):
 
 def test_factory_methods(raw_data, check_factory_methods):
     data = raw_data.density("Fe3O4 collinear")
-    check_factory_methods(calculation.density, data)
+    parameters = {"to_contour": {"a": 0.3}}
+    check_factory_methods(calculation.density, data, parameters)