From eb4e021e837eec3b9213a5f1a213810e65054f92 Mon Sep 17 00:00:00 2001 From: Martin Schlipf Date: Mon, 4 Dec 2023 15:42:10 +0100 Subject: [PATCH] Fix plotting of density (#119) * Plotting function should return Viewer instead of ViewerWrapper * NGL already transposed density, so do not transpose for plot --- src/py4vasp/_data/density.py | 13 +++++-------- tests/data/test_density.py | 12 +++++++----- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/src/py4vasp/_data/density.py b/src/py4vasp/_data/density.py index 7c8eca1c..b13eb7e8 100644 --- a/src/py4vasp/_data/density.py +++ b/src/py4vasp/_data/density.py @@ -86,9 +86,9 @@ def plot(self, selection="charge", **user_options): _raise_error_if_no_data(self._raw_data.charge) viewer = self._structure.plot() if selection == "charge": - return self._plot_charge(_ViewerWrapper(viewer), **user_options) + self._plot_charge(_ViewerWrapper(viewer), **user_options) elif selection == "magnetization": - return self._plot_magnetism(_ViewerWrapper(viewer), **user_options) + self._plot_magnetism(_ViewerWrapper(viewer), **user_options) else: msg = f"'{selection}' is an unknown option, please use 'charge' or 'magnetization' instead." raise exception.IncorrectUsage(msg) @@ -110,8 +110,7 @@ def noncollinear(self): return len(self._raw_data.charge) == 4 def _plot_charge(self, viewer, **user_options): - viewer.show_isosurface(self._raw_data.charge[0].T, **user_options) - return viewer + viewer.show_isosurface(self._raw_data.charge[0], **user_options) def _plot_magnetism(self, viewer, **user_options): if self.nonpolarized(): @@ -123,15 +122,13 @@ def _plot_magnetism(self, viewer, **user_options): def _plot_collinear_magnetism(self, viewer, **user_options): _raise_error_if_color_is_specified(**user_options) - magnetization = self._raw_data.charge[1].T + magnetization = self._raw_data.charge[1] viewer.show_isosurface(magnetization, color="blue", **user_options) viewer.show_isosurface(-magnetization, color="red", **user_options) - return viewer def _plot_noncollinear_magnetism(self, viewer, **user_options): - magnetization = np.linalg.norm(self._raw_data.charge[1:], axis=0).T + magnetization = np.linalg.norm(self._raw_data.charge[1:], axis=0) viewer.show_isosurface(magnetization, **user_options) - return viewer def _raise_is_nonpolarized_error(): diff --git a/tests/data/test_density.py b/tests/data/test_density.py index fcdf9c5b..5f9dd184 100644 --- a/tests/data/test_density.py +++ b/tests/data/test_density.py @@ -80,12 +80,13 @@ def test_charge_plot(reference_density, Assert, not_core): cm_cell = patch.object(obj, "show_cell") cm_surface = patch.object(obj, "show_isosurface") with cm_init as init, cm_cell as cell, cm_surface as surface: - reference_density.plot() + result = reference_density.plot() + assert isinstance(result, viewer3d.Viewer3d) init.assert_called_once() cell.assert_called_once() surface.assert_called_once() args, kwargs = surface.call_args - Assert.allclose(args[0], reference_density.ref.output["charge"]) + Assert.allclose(args[0], reference_density.ref.output["charge"].T) assert kwargs == {"isolevel": 0.2, "color": "yellow", "opacity": 0.6} @@ -107,9 +108,10 @@ def check_plotting_magnetization_density(polarized_density, Assert): cm_cell = patch.object(obj, "show_cell") cm_surface = patch.object(obj, "show_isosurface") with cm_init as init, cm_cell as cell, cm_surface as surface: - polarized_density.plot(selection="magnetization", isolevel=0.1, smooth=1) + result = polarized_density.plot("magnetization", isolevel=0.1, smooth=1) + assert isinstance(result, viewer3d.Viewer3d) calls = surface.call_args_list - reference_magnetization = polarized_density.ref.output["magnetization"] + reference_magnetization = polarized_density.ref.output["magnetization"].T if polarized_density.collinear(): check_collinear_plot(reference_magnetization, calls, Assert) elif polarized_density.noncollinear(): @@ -127,7 +129,7 @@ def check_collinear_plot(magnetization, calls, Assert): def check_noncollinear_plot(magnetization, calls, Assert): - magnetization = np.linalg.norm(magnetization, axis=0) + magnetization = np.linalg.norm(magnetization, axis=-1) assert len(calls) == 1 args, kwargs = calls[0] Assert.allclose(args[0], magnetization)