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

@@ -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():

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)