Fix transpose of density and potential

vasp-dev · Feb 23, 2024 · 2a86958 · 2a86958
1 parent aa510a1
commit 2a86958
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 5 deletions.
diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
@@ -259,7 +259,7 @@ def _grid_quantity(self, selector, selection, map_, user_options):
         component_label = selector.label(selection)
         component = map_.get(component_label, -1)
         return view.GridQuantity(
-            quantity=selector[selection][np.newaxis],
+            quantity=(selector[selection].T)[np.newaxis],
             label=self._label(component_label),
             isosurfaces=self._isosurfaces(component, **user_options),
         )

diff --git a/src/py4vasp/calculation/_potential.py b/src/py4vasp/calculation/_potential.py
@@ -121,7 +121,7 @@ def _create_potential_isosurface(
         else:
             potential = potential_data[0]
         return view.GridQuantity(
-            quantity=potential[np.newaxis],
+            quantity=potential.T[np.newaxis],
             label=f"{kind} potential" + (f"({component})" if component else ""),
             isosurfaces=[view.Isosurface(isolevel, color, opacity)],
         )

diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
@@ -140,7 +140,7 @@ def check_view(density, expected, Assert, **kwargs):
     grid_scalar = view.grid_scalars[0]
     assert grid_scalar.label == expected.label
     assert grid_scalar.quantity.ndim == 4
-    Assert.allclose(grid_scalar.quantity, expected.density.T)
+    Assert.allclose(grid_scalar.quantity, expected.density)
     assert len(grid_scalar.isosurfaces) == len(expected.isosurfaces)
     assert grid_scalar.isosurfaces == expected.isosurfaces
 

diff --git a/tests/calculation/test_potential.py b/tests/calculation/test_potential.py
@@ -97,7 +97,6 @@ def test_read(reference_potential, Assert):
             continue
         Assert.allclose(actual[key], reference_potential.ref.output[key])
 
-
 def test_plot_total_potential(reference_potential, Assert):
     view = reference_potential.plot()
     expectation = Expectation(
@@ -177,7 +176,7 @@ def check_view(potential, view, expectations, Assert, supercell=None):
     for grid_scalar, expected in zip(view.grid_scalars, expectations):
         assert grid_scalar.label == expected.label
         assert grid_scalar.quantity.ndim == 4
-        Assert.allclose(grid_scalar.quantity, expected.potential.T)
+        Assert.allclose(grid_scalar.quantity, expected.potential)
         assert len(grid_scalar.isosurfaces) == 1
         assert grid_scalar.isosurfaces[0] == expected.isosurface