diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py index 40377a93..3c0c0d4c 100644 --- a/src/py4vasp/calculation/_density.py +++ b/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() diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py index fef5cd34..0862f3f2 100644 --- a/tests/calculation/test_density.py +++ b/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)