From 0dd36f774983331b16d7797d70a80c27b31739e1 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 8 Apr 2024 09:53:45 +0200
Subject: [PATCH 01/46] Import Contour class from STM branch

---
 src/py4vasp/_third_party/graph/__init__.py |   2 +
 src/py4vasp/_third_party/graph/contour.py  |  93 ++++++++++++++++++
 src/py4vasp/_third_party/graph/graph.py    |  41 ++++++--
 src/py4vasp/_third_party/graph/series.py   |   8 +-
 src/py4vasp/_third_party/graph/trace.py    |  14 +++
 tests/third_party/graph/test_graph.py      | 106 ++++++++++++++++++++-
 6 files changed, 252 insertions(+), 12 deletions(-)
 create mode 100644 src/py4vasp/_third_party/graph/contour.py
 create mode 100644 src/py4vasp/_third_party/graph/trace.py

diff --git a/src/py4vasp/_third_party/graph/__init__.py b/src/py4vasp/_third_party/graph/__init__.py
index f1972421..eb2a084a 100644
--- a/src/py4vasp/_third_party/graph/__init__.py
+++ b/src/py4vasp/_third_party/graph/__init__.py
@@ -3,6 +3,7 @@
 from py4vasp._config import VASP_COLORS
 from py4vasp._util import import_
 
+from .contour import Contour
 from .graph import Graph
 from .mixin import Mixin
 from .plot import plot
@@ -15,4 +16,5 @@
     axis_format = {"showexponent": "all", "exponentformat": "power"}
     layout = {"colorway": VASP_COLORS, "xaxis": axis_format, "yaxis": axis_format}
     pio.templates["vasp"] = go.layout.Template(layout=layout)
+    pio.templates["ggplot2"].layout.shapedefaults = {}
     pio.templates.default = "ggplot2+vasp"
diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
new file mode 100644
index 00000000..01f26af8
--- /dev/null
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -0,0 +1,93 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+import dataclasses
+
+import numpy as np
+
+from py4vasp import _config
+from py4vasp._third_party.graph import trace
+from py4vasp._util import import_
+
+go = import_.optional("plotly.graph_objects")
+interpolate = import_.optional("scipy.interpolate")
+
+
+@dataclasses.dataclass
+class Contour(trace.Trace):
+    """Represents data on a 2d slice through the unit cell.
+
+    This class creates a visualization of the data within the unit cell based on its
+    configuration. Currently it supports the creation of heatmaps where each datapoint
+    corresponds to one point on the grid.
+    """
+
+    interpolation_factor = 2
+    """If the lattice does not align with the cartesian axes, the data is interpolated
+    to by approximately this factor along each line."""
+
+    data: np.array
+    "2d grid data in the plane spanned by the lattice vectors."
+    lattice: np.array
+    """2 vectors spanning the plane in which the data is represented. Each vector should
+    have two components, so remove any element normal to the plane."""
+    label: str
+    "Assign a label to the visualization that may be used to identify one among multiple plots."
+    supercell: np.array = (1, 1)
+    "Multiple of each lattice vector to be drawn."
+    show_cell: bool = True
+    "Show the unit cell in the resulting visualization."
+
+    def to_plotly(self):
+        lattice_supercell = np.diag(self.supercell) @ self.lattice
+        # swap a and b axes because that is the way plotly expects the data
+        data = np.tile(self.data, self.supercell).T
+        if self._interpolation_required():
+            x, y, z = self._interpolate_data(lattice_supercell, data)
+        else:
+            x, y, z = self._use_data_without_interpolation(lattice_supercell, data)
+        heatmap = go.Heatmap(x=x, y=y, z=z, name=self.label, colorscale="turbid_r")
+        yield heatmap, self._options()
+
+    def _interpolation_required(self):
+        return not np.allclose((self.lattice[1, 0], self.lattice[0, 1]), 0)
+
+    def _interpolate_data(self, lattice, data):
+        area_cell = abs(np.cross(lattice[0], lattice[1]))
+        points_per_area = data.size / area_cell
+        points_per_line = np.sqrt(points_per_area) * self.interpolation_factor
+        lengths = np.sum(np.abs(lattice), axis=0)
+        shape = np.ceil(points_per_line * lengths).astype(int)
+        line_mesh_a = self._make_mesh(lattice, data.shape[1], 0)
+        line_mesh_b = self._make_mesh(lattice, data.shape[0], 1)
+        x_in, y_in = (line_mesh_a[:, np.newaxis] + line_mesh_b[np.newaxis, :]).T
+        x_in = x_in.flatten()
+        y_in = y_in.flatten()
+        z_in = data.flatten()
+        x_out, y_out = np.meshgrid(
+            np.linspace(x_in.min(), x_in.max(), shape[0]),
+            np.linspace(y_in.min(), y_in.max(), shape[1]),
+        )
+        z_out = interpolate.griddata((x_in, y_in), z_in, (x_out, y_out), method="cubic")
+        return x_out[0], y_out[:, 0], z_out
+
+    def _use_data_without_interpolation(self, lattice, data):
+        x = self._make_mesh(lattice, data.shape[1], 0)
+        y = self._make_mesh(lattice, data.shape[0], 1)
+        return x, y, data
+
+    def _make_mesh(self, lattice, num_point, index):
+        vector = index if self._interpolation_required() else (index, index)
+        return (
+            np.linspace(0, lattice[vector], num_point, endpoint=False)
+            + 0.5 * lattice[vector] / num_point
+        )
+
+    def _options(self):
+        if not self.show_cell:
+            return {}
+        pos_to_str = lambda pos: f"{pos[0]} {pos[1]}"
+        corners = (self.lattice[0], self.lattice[0] + self.lattice[1], self.lattice[1])
+        to_corners = (f"L {pos_to_str(corner)}" for corner in corners)
+        path = f"M 0 0 {' '.join(to_corners)} Z"
+        unit_cell = {"type": "path", "line": {"color": _config.VASP_GRAY}, "path": path}
+        return {"shapes": [unit_cell]}
diff --git a/src/py4vasp/_third_party/graph/graph.py b/src/py4vasp/_third_party/graph/graph.py
index 6824658e..84ec2f3e 100644
--- a/src/py4vasp/_third_party/graph/graph.py
+++ b/src/py4vasp/_third_party/graph/graph.py
@@ -9,7 +9,9 @@
 
 from py4vasp import exception
 from py4vasp._config import VASP_COLORS
+from py4vasp._third_party.graph.contour import Contour
 from py4vasp._third_party.graph.series import Series
+from py4vasp._third_party.graph.trace import Trace
 from py4vasp._util import import_
 
 go = import_.optional("plotly.graph_objects")
@@ -25,7 +27,7 @@ class Graph(Sequence):
     parameters set in this class.
     """
 
-    series: Series or Sequence[Series]
+    series: Trace or Sequence[Trace]
     "One or more series shown in the graph."
     xlabel: str = None
     "Label for the x axis."
@@ -71,11 +73,12 @@ def to_plotly(self):
         "Convert the graph to a plotly figure."
         figure = self._make_plotly_figure()
         for trace, options in self._generate_plotly_traces():
-            if options["row"] is None:
+            if options.get("row") is None:
                 figure.add_trace(trace)
             else:
                 figure.add_trace(trace, row=options["row"], col=1)
-
+            for shape in options.get("shapes", ()):
+                figure.add_shape(**shape)
         return figure
 
     def show(self):
@@ -107,9 +110,8 @@ def _ipython_display_(self):
     def _generate_plotly_traces(self):
         colors = itertools.cycle(VASP_COLORS)
         for series in self:
-            if not series.color:
-                series = replace(series, color=next(colors))
-            yield from series._generate_traces()
+            series = _set_color_if_not_present(series, colors)
+            yield from series.to_plotly()
 
     def _make_plotly_figure(self):
         figure = self._figure_with_one_or_two_y_axes()
@@ -120,12 +122,13 @@ def _make_plotly_figure(self):
         return figure
 
     def _figure_with_one_or_two_y_axes(self):
+        has_secondary_y_axis = lambda series: isinstance(series, Series) and series.y2
         if self._subplot_on:
             max_row = max(series.subplot for series in self)
             figure = subplots.make_subplots(rows=max_row, cols=1)
             figure.update_layout(showlegend=False)
             return figure
-        elif any(series.y2 for series in self):
+        elif any(has_secondary_y_axis(series) for series in self):
             return subplots.make_subplots(specs=[[{"secondary_y": True}]])
         else:
             return go.Figure()
@@ -142,6 +145,8 @@ def _set_xaxis_options(self, figure):
             figure.layout.xaxis.tickmode = "array"
             figure.layout.xaxis.tickvals = tuple(self.xticks.keys())
             figure.layout.xaxis.ticktext = self._xtick_labels()
+        if self._all_are_contour():
+            figure.layout.xaxis.visible = False
 
     def _xtick_labels(self):
         # empty labels will be overwritten by plotly so we put a single space in them
@@ -156,6 +161,17 @@ def _set_yaxis_options(self, figure):
             figure.layout.yaxis.title.text = self.ylabel
             if self.y2label:
                 figure.layout.yaxis2.title.text = self.y2label
+        if self._all_are_contour():
+            figure.layout.yaxis.visible = False
+        if self._any_are_contour():
+            figure.layout.yaxis.scaleanchor = "x"
+            figure.layout.height = 500
+
+    def _all_are_contour(self):
+        return all(isinstance(series, Contour) for series in self)
+
+    def _any_are_contour(self):
+        return any(isinstance(series, Contour) for series in self)
 
     def to_frame(self):
         """Convert graph to a pandas dataframe.
@@ -213,7 +229,16 @@ def _name_column(self, series, suffix, idx=None):
 
     @property
     def _subplot_on(self):
-        return any(series.subplot for series in self)
+        has_subplot = lambda series: isinstance(series, Series) and series.subplot
+        return any(has_subplot(series) for series in self)
+
+
+def _set_color_if_not_present(series, color_iterator):
+    if isinstance(series, Contour):
+        return series
+    if not series.color:
+        series = replace(series, color=next(color_iterator))
+    return series
 
 
 Graph._fields = tuple(field.name for field in fields(Graph))
diff --git a/src/py4vasp/_third_party/graph/series.py b/src/py4vasp/_third_party/graph/series.py
index 802af191..16a9e890 100644
--- a/src/py4vasp/_third_party/graph/series.py
+++ b/src/py4vasp/_third_party/graph/series.py
@@ -5,13 +5,14 @@
 import numpy as np
 
 from py4vasp import exception
+from py4vasp._third_party.graph import trace
 from py4vasp._util import import_
 
 go = import_.optional("plotly.graph_objects")
 
 
 @dataclass
-class Series:
+class Series(trace.Trace):
     """Represents a single series in a graph.
 
     Typically this corresponds to a single line of x-y data with an optional name used
@@ -55,7 +56,7 @@ def __setattr__(self, key, value):
         assert not self._frozen or hasattr(self, key)
         super().__setattr__(key, value)
 
-    def _generate_traces(self):
+    def to_plotly(self):
         first_trace = True
         for item in enumerate(np.atleast_2d(np.array(self.y))):
             yield self._make_trace(*item, first_trace), {"row": self.subplot}
@@ -123,5 +124,8 @@ def _common_options(self, first_trace):
             "yaxis": "y2" if self.y2 else "y",
         }
 
+    def _generate_shapes(self):
+        return ()
+
 
 Series._fields = tuple(field.name for field in fields(Series))
diff --git a/src/py4vasp/_third_party/graph/trace.py b/src/py4vasp/_third_party/graph/trace.py
new file mode 100644
index 00000000..503d0cff
--- /dev/null
+++ b/src/py4vasp/_third_party/graph/trace.py
@@ -0,0 +1,14 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+from abc import ABC, abstractmethod
+
+
+class Trace(ABC):
+    """Defines a base class with all methods that need to be implemented for Graph to
+    work as intended"""
+
+    @abstractmethod
+    def to_plotly(self):
+        """Use yield to generate one or more plotly traces. Each returned element should
+        be a tuple (trace, dict) where the trace can be used as data for plotly and the
+        options modify the generation of the figure."""
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index 26d05829..c6a82ecb 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -6,8 +6,11 @@
 import numpy as np
 import pytest
 
-from py4vasp import exception
-from py4vasp._third_party.graph import Graph, Series
+from py4vasp import _config, exception
+from py4vasp._third_party.graph import Contour, Graph, Series
+from py4vasp._util import import_
+
+px = import_.optional("plotly.express")
 
 
 @pytest.fixture
@@ -59,6 +62,26 @@ def subplot():
     return Series(x1, y1, subplot=1), Series(x2, y2, subplot=2)
 
 
+@pytest.fixture
+def rectangle_contour():
+    return Contour(
+        data=np.linspace(0, 10, 20 * 18).reshape((20, 18)),
+        lattice=np.diag([4.0, 3.6]),
+        label="rectangle contour",
+    )
+
+
+@pytest.fixture
+def tilted_contour():
+    return Contour(
+        data=np.linspace(0, 5, 16 * 20).reshape((16, 20)),
+        lattice=np.array([[2, 3], [2, -3]]),
+        label="tilted contour",
+        supercell=(2, 1),
+        show_cell=False,
+    )
+
+
 def test_basic_graph(parabola, Assert, not_core):
     graph = Graph(parabola)
     fig = graph.to_plotly()
@@ -373,3 +396,82 @@ def test_nonexisting_attribute_raises_error(parabola):
     graph = Graph(parabola)
     with pytest.raises(AssertionError):
         graph.nonexisting = "not possible"
+
+
+def test_normal_series_does_not_set_contour_layout(parabola, not_core):
+    graph = Graph(parabola)
+    fig = graph.to_plotly()
+    assert fig.layout.xaxis.visible is None
+    assert fig.layout.yaxis.visible is None
+    assert fig.layout.yaxis.scaleanchor is None
+
+
+def test_contour(rectangle_contour, Assert, not_core):
+    graph = Graph(rectangle_contour)
+    fig = graph.to_plotly()
+    assert len(fig.data) == 1
+    # plotly expects y-x order
+    Assert.allclose(fig.data[0].z, rectangle_contour.data.T)
+    # shift because the points define the centers of the rectangles
+    Assert.allclose(fig.data[0].x, np.linspace(0, 4, 20, endpoint=False) + 0.1)
+    Assert.allclose(fig.data[0].y, np.linspace(0, 3.6, 18, endpoint=False) + 0.1)
+    assert fig.data[0].name == rectangle_contour.label
+    expected_colorscale = px.colors.get_colorscale("turbid_r")
+    assert len(fig.data[0].colorscale) == len(expected_colorscale)
+    for actual, expected in zip(fig.data[0].colorscale, expected_colorscale):
+        Assert.allclose(actual[0], expected[0])
+        assert actual[1] == expected[1]
+    # text explicitly that it is False to prevent None passing the test
+    assert fig.layout.xaxis.visible == False
+    assert fig.layout.yaxis.visible == False
+    assert len(fig.layout.shapes) == 1
+    check_unit_cell(fig.layout.shapes[0])
+    assert fig.layout.yaxis.scaleanchor == "x"
+
+
+def test_contour_supercell(rectangle_contour, not_core):
+    supercell = np.asarray((3, 5))
+    rectangle_contour.supercell = supercell
+    graph = Graph(rectangle_contour)
+    fig = graph.to_plotly()
+    assert len(fig.data) == 1
+    # plotly expects y-x order
+    assert all(fig.data[0].z.T.shape == supercell * rectangle_contour.data.shape)
+    assert len(fig.data[0].x) == 60
+    assert len(fig.data[0].y) == 90
+    assert len(fig.layout.shapes) == 1
+    check_unit_cell(fig.layout.shapes[0])
+
+
+def check_unit_cell(unit_cell):
+    assert unit_cell.type == "path"
+    assert unit_cell.path == "M 0 0 L 4.0 0.0 L 4.0 3.6 L 0.0 3.6 Z"
+    assert unit_cell.line.color == _config.VASP_GRAY
+
+
+def test_contour_interpolate(tilted_contour, not_core):
+    graph = Graph(tilted_contour)
+    fig = graph.to_plotly()
+    area_cell = 12.0
+    points_per_area = tilted_contour.data.size / area_cell
+    points_per_line = np.sqrt(points_per_area) * tilted_contour.interpolation_factor
+    lengths = np.array([6, 9])  # this accounts for the 2 x 1 supercell
+    expected_shape = np.ceil(points_per_line * lengths).astype(int)
+    expected_average = np.average(tilted_contour.data)
+    assert len(fig.data) == 1
+    # plotly expects y-x order
+    assert all(fig.data[0].z.T.shape == expected_shape)
+    assert fig.data[0].x.size == expected_shape[0]
+    assert fig.data[0].y.size == expected_shape[1]
+    finite = np.isfinite(fig.data[0].z)
+    assert np.isclose(np.average(fig.data[0].z[finite]), expected_average)
+    assert len(fig.layout.shapes) == 0
+
+
+def test_mix_contour_and_series(two_lines, rectangle_contour, not_core):
+    graph = Graph([rectangle_contour, two_lines])
+    fig = graph.to_plotly()
+    assert len(fig.data) == 3
+    assert fig.layout.xaxis.visible is None
+    assert fig.layout.yaxis.visible is None
+    assert fig.layout.yaxis.scaleanchor == "x"

From f78d76eb4bb53cef10e508ffe8229a1a498dea67 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 8 Apr 2024 21:33:36 +0200
Subject: [PATCH 02/46] Add test for quiver plot

---
 src/py4vasp/_third_party/graph/contour.py | 28 ++++++++--
 tests/conftest.py                         |  4 +-
 tests/third_party/graph/test_graph.py     | 64 +++++++++++++++++++++++
 3 files changed, 91 insertions(+), 5 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index 01f26af8..c776a7b5 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -1,6 +1,7 @@
 # Copyright © VASP Software GmbH,
 # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
 import dataclasses
+import itertools
 
 import numpy as np
 
@@ -8,6 +9,7 @@
 from py4vasp._third_party.graph import trace
 from py4vasp._util import import_
 
+ff = import_.optional("plotly.figure_factory")
 go = import_.optional("plotly.graph_objects")
 interpolate = import_.optional("scipy.interpolate")
 
@@ -41,12 +43,21 @@ def to_plotly(self):
         lattice_supercell = np.diag(self.supercell) @ self.lattice
         # swap a and b axes because that is the way plotly expects the data
         data = np.tile(self.data, self.supercell).T
+        if self._is_heatmap():
+            yield self._make_heatmap(lattice_supercell, data)
+        else:
+            yield self._make_quiver(lattice_supercell, data)
+
+    def _is_heatmap(self):
+        return self.data.ndim == 2
+
+    def _make_heatmap(self, lattice, data):
         if self._interpolation_required():
-            x, y, z = self._interpolate_data(lattice_supercell, data)
+            x, y, z = self._interpolate_data(lattice, data)
         else:
-            x, y, z = self._use_data_without_interpolation(lattice_supercell, data)
+            x, y, z = self._use_data_without_interpolation(lattice, data)
         heatmap = go.Heatmap(x=x, y=y, z=z, name=self.label, colorscale="turbid_r")
-        yield heatmap, self._options()
+        return heatmap, self._options()
 
     def _interpolation_required(self):
         return not np.allclose((self.lattice[1, 0], self.lattice[0, 1]), 0)
@@ -91,3 +102,14 @@ def _options(self):
         path = f"M 0 0 {' '.join(to_corners)} Z"
         unit_cell = {"type": "path", "line": {"color": _config.VASP_GRAY}, "path": path}
         return {"shapes": [unit_cell]}
+
+    def _make_quiver(self, lattice, data):
+        u = data[:, :, 0].flatten()
+        v = data[:, :, 1].flatten()
+        meshes = [
+            np.linspace(np.zeros(2), vector, num_points, endpoint=False)
+            for vector, num_points in zip(lattice, data.shape[1::-1])
+        ]
+        x, y = np.array([sum(points) for points in itertools.product(*meshes)]).T
+        fig = ff.create_quiver(x, y, u, v, scale=1)
+        return fig.data[0], {}
diff --git a/tests/conftest.py b/tests/conftest.py
index 67e10984..3f63cae2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -47,7 +47,7 @@ def not_core(is_core):
 
 class _Assert:
     @staticmethod
-    def allclose(actual, desired):
+    def allclose(actual, desired, tolerance=1):
         if _is_none(actual):
             assert _is_none(desired)
         elif getattr(actual, "dtype", None) == np.bool_:
@@ -58,7 +58,7 @@ def allclose(actual, desired):
             actual, mask_actual = _finite_subset(actual)
             desired, mask_desired = _finite_subset(desired)
             assert np.all(mask_actual == mask_desired)
-            assert_array_almost_equal_nulp(actual, desired, 30)
+            assert_array_almost_equal_nulp(actual, desired, tolerance * 30)
 
     @staticmethod
     def same_structure(actual, desired):
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index c6a82ecb..decc534d 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -82,6 +82,16 @@ def tilted_contour():
     )
 
 
+@pytest.fixture
+def quiver_plot():
+    return Contour(
+        data=np.linspace(-3, 3, 2 * 12 * 10).reshape((2, 12, 10)),
+        lattice=np.array([[3, 2], [-3, 2]]),
+        label="quiver plot",
+        supercell=(3, 2),
+    )
+
+
 def test_basic_graph(parabola, Assert, not_core):
     graph = Graph(parabola)
     fig = graph.to_plotly()
@@ -475,3 +485,57 @@ def test_mix_contour_and_series(two_lines, rectangle_contour, not_core):
     assert fig.layout.xaxis.visible is None
     assert fig.layout.yaxis.visible is None
     assert fig.layout.yaxis.scaleanchor == "x"
+
+
+def test_quiver_plot(quiver_plot, Assert, not_core):
+    graph = Graph(quiver_plot)
+    fig = graph.to_plotly()
+    data_size = np.prod(quiver_plot.supercell) * quiver_plot.data.size // 2
+    expected_positions = np.array(
+        [
+            a * quiver_plot.lattice[0] / quiver_plot.data.shape[1]
+            + b * quiver_plot.lattice[1] / quiver_plot.data.shape[2]
+            for a in np.arange(quiver_plot.supercell[0] * quiver_plot.data.shape[1])
+            for b in np.arange(quiver_plot.supercell[1] * quiver_plot.data.shape[2])
+        ]
+    )
+    work = quiver_plot.data
+    work = np.block([[work, work], [work, work], [work, work]]).T
+    expected_tips = expected_positions + work.reshape(expected_positions.shape)
+    expected_barb_length = 0.3 * np.linalg.norm(work, axis=-1).flatten()
+    #
+    assert len(fig.data) == 1
+    data = fig.data[0]
+    assert data.mode == "lines"
+    # The data contains first a line between each grid point and the tip of the arrow;
+    # each of the lines are separated by a None. The second part is the tip of the arrow
+    # consisting of two lines; again tips are separated by None.
+    assert len(data.x) == len(data.y) == 7 * data_size
+    # first element should be positions
+    slice_ = slice(0, 3 * data_size, 3)
+    actual_positions = np.array((data.x[slice_], data.y[slice_])).T
+    Assert.allclose(actual_positions, expected_positions, tolerance=10)
+    # second element of both parts should be the tip of the arrows
+    slice_ = slice(1, 3 * data_size, 3)
+    actual_tips = np.array((data.x[slice_], data.y[slice_])).T
+    Assert.allclose(actual_tips, expected_tips, tolerance=10)
+    slice_ = slice(3 * data_size + 1, None, 4)
+    actual_tips = np.array((data.x[slice_], data.y[slice_])).T
+    Assert.allclose(actual_tips, expected_tips, tolerance=10)
+    # third element of first part and fourth element of second part should be none
+    slice_ = slice(2, 3 * data_size, 3)
+    assert all(element is None for element in data.x[slice_])
+    assert all(element is None for element in data.y[slice_])
+    slice_ = slice(3 * data_size + 3, None, 4)
+    assert all(element is None for element in data.x[slice_])
+    assert all(element is None for element in data.y[slice_])
+    # the first and third element of the second part contain the barb of the arrow,
+    # which should be close to tip
+    slice_ = slice(3 * data_size, None, 4)
+    first_barb = np.array((data.x[slice_], data.y[slice_])).T
+    actual_barb_length = np.linalg.norm(first_barb - actual_tips, axis=-1)
+    Assert.allclose(actual_barb_length, expected_barb_length)
+    slice_ = slice(3 * data_size + 2, None, 4)
+    second_barb = np.array((data.x[slice_], data.y[slice_])).T
+    actual_barb_length = np.linalg.norm(second_barb - actual_tips, axis=-1)
+    Assert.allclose(actual_barb_length, expected_barb_length)

From af38bcc3254f814bccaea0e1fe23642b9cd3d319 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 8 Apr 2024 21:41:12 +0200
Subject: [PATCH 03/46] Refactor test

---
 tests/third_party/graph/test_graph.py | 21 ++++++++++-----------
 1 file changed, 10 insertions(+), 11 deletions(-)

diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index decc534d..44fc368c 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -83,7 +83,7 @@ def tilted_contour():
 
 
 @pytest.fixture
-def quiver_plot():
+def complex_quiver():
     return Contour(
         data=np.linspace(-3, 3, 2 * 12 * 10).reshape((2, 12, 10)),
         lattice=np.array([[3, 2], [-3, 2]]),
@@ -487,19 +487,18 @@ def test_mix_contour_and_series(two_lines, rectangle_contour, not_core):
     assert fig.layout.yaxis.scaleanchor == "x"
 
 
-def test_quiver_plot(quiver_plot, Assert, not_core):
-    graph = Graph(quiver_plot)
+def test_complex_quiver(complex_quiver, Assert, not_core):
+    graph = Graph(complex_quiver)
     fig = graph.to_plotly()
-    data_size = np.prod(quiver_plot.supercell) * quiver_plot.data.size // 2
+    data_size = np.prod(complex_quiver.supercell) * complex_quiver.data.size // 2
+    step_a = complex_quiver.lattice[0] / complex_quiver.data.shape[1]
+    mesh_a = np.arange(complex_quiver.supercell[0] * complex_quiver.data.shape[1])
+    step_b = complex_quiver.lattice[1] / complex_quiver.data.shape[2]
+    mesh_b = np.arange(complex_quiver.supercell[1] * complex_quiver.data.shape[2])
     expected_positions = np.array(
-        [
-            a * quiver_plot.lattice[0] / quiver_plot.data.shape[1]
-            + b * quiver_plot.lattice[1] / quiver_plot.data.shape[2]
-            for a in np.arange(quiver_plot.supercell[0] * quiver_plot.data.shape[1])
-            for b in np.arange(quiver_plot.supercell[1] * quiver_plot.data.shape[2])
-        ]
+        [a * step_a + b * step_b for a in mesh_a for b in mesh_b]
     )
-    work = quiver_plot.data
+    work = complex_quiver.data
     work = np.block([[work, work], [work, work], [work, work]]).T
     expected_tips = expected_positions + work.reshape(expected_positions.shape)
     expected_barb_length = 0.3 * np.linalg.norm(work, axis=-1).flatten()

From 43a36c5966df6a6949bba72a66b5d1ce94fae114 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 8 Apr 2024 22:30:22 +0200
Subject: [PATCH 04/46] Fix transpose of axis

---
 src/py4vasp/_third_party/graph/contour.py |  3 +-
 tests/third_party/graph/test_graph.py     | 72 +++++++++++++++++------
 2 files changed, 55 insertions(+), 20 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index c776a7b5..d66b84f3 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -108,7 +108,8 @@ def _make_quiver(self, lattice, data):
         v = data[:, :, 1].flatten()
         meshes = [
             np.linspace(np.zeros(2), vector, num_points, endpoint=False)
-            for vector, num_points in zip(lattice, data.shape[1::-1])
+            for vector, num_points in zip(reversed(lattice), data.shape)
+            # remember that b and a axis are swapped
         ]
         x, y = np.array([sum(points) for points in itertools.product(*meshes)]).T
         fig = ff.create_quiver(x, y, u, v, scale=1)
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index 44fc368c..929cb95c 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -82,6 +82,15 @@ def tilted_contour():
     )
 
 
+@pytest.fixture
+def simple_quiver():
+    return Contour(
+        data=np.array([[(y, x) for x in range(3)] for y in range(5)]).T,
+        lattice=np.diag((3, 5)),
+        label="quiver plot",
+    )
+
+
 @pytest.fixture
 def complex_quiver():
     return Contour(
@@ -487,6 +496,18 @@ def test_mix_contour_and_series(two_lines, rectangle_contour, not_core):
     assert fig.layout.yaxis.scaleanchor == "x"
 
 
+def test_simple_quiver(simple_quiver, Assert, not_core):
+    graph = Graph(simple_quiver)
+    fig = graph.to_plotly()
+    data_size = simple_quiver.data.size // 2
+    assert len(fig.data) == 1
+    actual = split_data(fig.data[0], data_size, Assert)
+    arrows = actual.tips - actual.positions
+    for (x, y), (u, v) in zip(actual.positions, arrows):
+        Assert.allclose(x, v)
+        Assert.allclose(y, u)
+
+
 def test_complex_quiver(complex_quiver, Assert, not_core):
     graph = Graph(complex_quiver)
     fig = graph.to_plotly()
@@ -496,7 +517,7 @@ def test_complex_quiver(complex_quiver, Assert, not_core):
     step_b = complex_quiver.lattice[1] / complex_quiver.data.shape[2]
     mesh_b = np.arange(complex_quiver.supercell[1] * complex_quiver.data.shape[2])
     expected_positions = np.array(
-        [a * step_a + b * step_b for a in mesh_a for b in mesh_b]
+        [a * step_a + b * step_b for b in mesh_b for a in mesh_a]
     )
     work = complex_quiver.data
     work = np.block([[work, work], [work, work], [work, work]]).T
@@ -504,37 +525,50 @@ def test_complex_quiver(complex_quiver, Assert, not_core):
     expected_barb_length = 0.3 * np.linalg.norm(work, axis=-1).flatten()
     #
     assert len(fig.data) == 1
-    data = fig.data[0]
+    actual = split_data(fig.data[0], data_size, Assert)
+    Assert.allclose(actual.positions, expected_positions, tolerance=10)
+    Assert.allclose(actual.tips, expected_tips, tolerance=10)
+    Assert.allclose(actual.barb_length, expected_barb_length)
+
+
+@dataclasses.dataclass
+class ContourData:
+    positions: np.ndarray = None
+    first_tips: np.ndarray = None
+    second_tips: np.ndarray = None
+    first_barb_length: np.ndarray = None
+    second_barb_length: np.ndarray = None
+
+
+def split_data(data, data_size, Assert):
+    actual = ContourData()
     assert data.mode == "lines"
     # The data contains first a line between each grid point and the tip of the arrow;
     # each of the lines are separated by a None. The second part is the tip of the arrow
     # consisting of two lines; again tips are separated by None.
     assert len(data.x) == len(data.y) == 7 * data_size
-    # first element should be positions
+    # first element contains positions
     slice_ = slice(0, 3 * data_size, 3)
-    actual_positions = np.array((data.x[slice_], data.y[slice_])).T
-    Assert.allclose(actual_positions, expected_positions, tolerance=10)
-    # second element of both parts should be the tip of the arrows
+    actual.positions = np.array((data.x[slice_], data.y[slice_])).T
+    # second element of both parts contain the tip of the arrows
     slice_ = slice(1, 3 * data_size, 3)
-    actual_tips = np.array((data.x[slice_], data.y[slice_])).T
-    Assert.allclose(actual_tips, expected_tips, tolerance=10)
+    actual.tips = np.array((data.x[slice_], data.y[slice_])).T
     slice_ = slice(3 * data_size + 1, None, 4)
-    actual_tips = np.array((data.x[slice_], data.y[slice_])).T
-    Assert.allclose(actual_tips, expected_tips, tolerance=10)
-    # third element of first part and fourth element of second part should be none
+    other_tips = np.array((data.x[slice_], data.y[slice_])).T
+    Assert.allclose(other_tips, actual.tips)
+    # third element of first part and fourth element of second part should be None (=separator)
     slice_ = slice(2, 3 * data_size, 3)
     assert all(element is None for element in data.x[slice_])
     assert all(element is None for element in data.y[slice_])
     slice_ = slice(3 * data_size + 3, None, 4)
     assert all(element is None for element in data.x[slice_])
     assert all(element is None for element in data.y[slice_])
-    # the first and third element of the second part contain the barb of the arrow,
-    # which should be close to tip
+    # the first and third element of the second part contain the barb of the arrow
     slice_ = slice(3 * data_size, None, 4)
-    first_barb = np.array((data.x[slice_], data.y[slice_])).T
-    actual_barb_length = np.linalg.norm(first_barb - actual_tips, axis=-1)
-    Assert.allclose(actual_barb_length, expected_barb_length)
+    barb = np.array((data.x[slice_], data.y[slice_])).T
+    actual.barb_length = np.linalg.norm(barb - actual.tips, axis=-1)
     slice_ = slice(3 * data_size + 2, None, 4)
-    second_barb = np.array((data.x[slice_], data.y[slice_])).T
-    actual_barb_length = np.linalg.norm(second_barb - actual_tips, axis=-1)
-    Assert.allclose(actual_barb_length, expected_barb_length)
+    other_barb = np.array((data.x[slice_], data.y[slice_])).T
+    other_barb_length = np.linalg.norm(other_barb - actual.tips, axis=-1)
+    Assert.allclose(other_barb_length, actual.barb_length, tolerance=10)
+    return actual

From 92b2f6e5aa654edd38eae5ffca5cbd8037728c84 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 8 Apr 2024 23:21:10 +0200
Subject: [PATCH 05/46] Add unit cell to quiver plot

---
 src/py4vasp/_third_party/graph/contour.py |  2 +-
 tests/third_party/graph/test_graph.py     | 11 +++++++----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index d66b84f3..1670f2a7 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -113,4 +113,4 @@ def _make_quiver(self, lattice, data):
         ]
         x, y = np.array([sum(points) for points in itertools.product(*meshes)]).T
         fig = ff.create_quiver(x, y, u, v, scale=1)
-        return fig.data[0], {}
+        return fig.data[0], self._options()
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index 929cb95c..5c0ba892 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -444,7 +444,7 @@ def test_contour(rectangle_contour, Assert, not_core):
     assert fig.layout.xaxis.visible == False
     assert fig.layout.yaxis.visible == False
     assert len(fig.layout.shapes) == 1
-    check_unit_cell(fig.layout.shapes[0])
+    check_unit_cell(fig.layout.shapes[0], x="4.0", y="3.6", zero="0.0")
     assert fig.layout.yaxis.scaleanchor == "x"
 
 
@@ -459,12 +459,12 @@ def test_contour_supercell(rectangle_contour, not_core):
     assert len(fig.data[0].x) == 60
     assert len(fig.data[0].y) == 90
     assert len(fig.layout.shapes) == 1
-    check_unit_cell(fig.layout.shapes[0])
+    check_unit_cell(fig.layout.shapes[0], x="4.0", y="3.6", zero="0.0")
 
 
-def check_unit_cell(unit_cell):
+def check_unit_cell(unit_cell, x, y, zero):
     assert unit_cell.type == "path"
-    assert unit_cell.path == "M 0 0 L 4.0 0.0 L 4.0 3.6 L 0.0 3.6 Z"
+    assert unit_cell.path == f"M 0 0 L {x} {zero} L {x} {y} L {zero} {y} Z"
     assert unit_cell.line.color == _config.VASP_GRAY
 
 
@@ -506,6 +506,9 @@ def test_simple_quiver(simple_quiver, Assert, not_core):
     for (x, y), (u, v) in zip(actual.positions, arrows):
         Assert.allclose(x, v)
         Assert.allclose(y, u)
+    assert len(fig.layout.shapes) == 1
+    check_unit_cell(fig.layout.shapes[0], x="3", y="5", zero="0")
+    assert fig.layout.yaxis.scaleanchor == "x"
 
 
 def test_complex_quiver(complex_quiver, Assert, not_core):

From 3dea206abeedeed6940bacc573ecbab97b80d8d8 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 9 Apr 2024 09:17:44 +0200
Subject: [PATCH 06/46] Update documentation

---
 src/py4vasp/_third_party/graph/contour.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index 1670f2a7..7532af5d 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -19,8 +19,9 @@ class Contour(trace.Trace):
     """Represents data on a 2d slice through the unit cell.
 
     This class creates a visualization of the data within the unit cell based on its
-    configuration. Currently it supports the creation of heatmaps where each datapoint
-    corresponds to one point on the grid.
+    configuration. Currently it supports the creation of heatmaps and quiver plots.
+    For heatmaps each data point corresponds to one point on the grid. For quiver plots
+    each data point should be a 2d vector within the plane.
     """
 
     interpolation_factor = 2
@@ -28,7 +29,10 @@ class Contour(trace.Trace):
     to by approximately this factor along each line."""
 
     data: np.array
-    "2d grid data in the plane spanned by the lattice vectors."
+    """2d or 3d grid data in the plane spanned by the lattice vectors. If the data is
+    the dimensions should be the ones of the grid, if the data is 3d the first dimension
+    should be a 2 for a vector in the plane of the grid and the other two dimensions
+    should be the grid."""
     lattice: np.array
     """2 vectors spanning the plane in which the data is represented. Each vector should
     have two components, so remove any element normal to the plane."""

From d17e16ce8d4ea561ad81da046fbec5a1749b9303 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 9 Apr 2024 10:38:16 +0200
Subject: [PATCH 07/46] Begin implement utility to construct 2d plane from
 slice

---
 src/py4vasp/_util/slicing.py | 10 ++++++++++
 tests/util/test_slicing.py   | 14 ++++++++++++++
 2 files changed, 24 insertions(+)
 create mode 100644 src/py4vasp/_util/slicing.py
 create mode 100644 tests/util/test_slicing.py

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
new file mode 100644
index 00000000..2121c446
--- /dev/null
+++ b/src/py4vasp/_util/slicing.py
@@ -0,0 +1,10 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+import numpy as np
+
+INDICES = {"a": 0, "b": 1, "c": 2}
+
+
+def plane(cell, cut):
+    index = INDICES[cut]
+    return np.delete(np.delete(cell, index, axis=0), index, axis=1)
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
new file mode 100644
index 00000000..7a47a98a
--- /dev/null
+++ b/tests/util/test_slicing.py
@@ -0,0 +1,14 @@
+# Copyright © VASP Software GmbH,
+# Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+import numpy as np
+import pytest
+
+from py4vasp._util import slicing
+
+
+@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)])
+def test_orthorhombic(cut, index, Assert):
+    cell = np.diag((3, 4, 5))
+    expected_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1)
+    actual_plane = slicing.plane(cell, cut)
+    Assert.allclose(actual_plane, expected_plane)

From 64b466d627a9d41e4a6800aa5877bd77f253a1f6 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 9 Apr 2024 10:54:37 +0200
Subject: [PATCH 08/46] Make slicing work for unusual order of vectors

---
 src/py4vasp/_util/slicing.py | 7 +++++--
 tests/util/test_slicing.py   | 8 ++++++++
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index 2121c446..fcaed97e 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -6,5 +6,8 @@
 
 
 def plane(cell, cut):
-    index = INDICES[cut]
-    return np.delete(np.delete(cell, index, axis=0), index, axis=1)
+    remaining_vectors = np.delete(cell, INDICES[cut], axis=0)
+    for i in range(3):
+        if np.allclose(remaining_vectors[:, i], 0):
+            return np.delete(remaining_vectors, i, axis=1)
+    raise NotImplementedError()
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 7a47a98a..88d3c8fc 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -12,3 +12,11 @@ def test_orthorhombic(cut, index, Assert):
     expected_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1)
     actual_plane = slicing.plane(cell, cut)
     Assert.allclose(actual_plane, expected_plane)
+
+
+@pytest.mark.parametrize("cut, indices", [("a", (0, 2)), ("b", (1, 0)), ("c", (2, 1))])
+def test_unusual_orthorhombic(cut, indices, Assert):
+    cell = np.roll(np.diag((3, 4, 5)), shift=1, axis=0)
+    expected_plane = np.delete(np.delete(cell, indices[0], axis=0), indices[1], axis=1)
+    actual_plane = slicing.plane(cell, cut)
+    Assert.allclose(actual_plane, expected_plane)

From 18f200324134925732ae920d5648700cf9dfee21 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 9 Apr 2024 12:04:36 +0200
Subject: [PATCH 09/46] Make transformation work for nearly orthorhombic cell

---
 src/py4vasp/_util/slicing.py | 12 ++++++++----
 tests/util/test_slicing.py   | 12 ++++++++++++
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index fcaed97e..1fc9ba90 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -6,8 +6,12 @@
 
 
 def plane(cell, cut):
-    remaining_vectors = np.delete(cell, INDICES[cut], axis=0)
+    old_vectors = np.delete(cell, INDICES[cut], axis=0)
     for i in range(3):
-        if np.allclose(remaining_vectors[:, i], 0):
-            return np.delete(remaining_vectors, i, axis=1)
-    raise NotImplementedError()
+        if np.allclose(old_vectors[:, i], 0):
+            return np.delete(old_vectors, i, axis=1)
+    U, S, _ = np.linalg.svd(old_vectors)
+    new_vectors = U @ np.diag(S)
+    if np.linalg.det(new_vectors) < 0:
+        new_vectors = new_vectors[:, ::-1]
+    return new_vectors
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 88d3c8fc..029dc7cd 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -20,3 +20,15 @@ def test_unusual_orthorhombic(cut, indices, Assert):
     expected_plane = np.delete(np.delete(cell, indices[0], axis=0), indices[1], axis=1)
     actual_plane = slicing.plane(cell, cut)
     Assert.allclose(actual_plane, expected_plane)
+
+
+@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)])
+def test_nearly_orthorhombic(cut, index, Assert):
+    cell = 0.01 * np.ones((3, 3)) + np.diag((3, 4, 5))
+    expected_plane = np.delete(cell, index, axis=0)
+    # the expected plane is in 3d and the actual plane in 2d should have same angles
+    # and lengths
+    approximate_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1)
+    actual_plane = slicing.plane(cell, cut)
+    Assert.allclose(actual_plane @ actual_plane.T, expected_plane @ expected_plane.T)
+    assert np.max(np.abs(approximate_plane - actual_plane)) < 0.1

From 4da0acd8fcc86f6bc863a8cea376a60757343f91 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Wed, 10 Apr 2024 13:47:21 +0200
Subject: [PATCH 10/46] wip: slicing

---
 src/py4vasp/_util/slicing.py | 35 +++++++++++++++++++++++++++++------
 tests/util/test_slicing.py   |  2 ++
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index 1fc9ba90..fac5f3c3 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -6,12 +6,35 @@
 
 
 def plane(cell, cut):
+    # old_det = np.linalg.det(cell)
     old_vectors = np.delete(cell, INDICES[cut], axis=0)
-    for i in range(3):
-        if np.allclose(old_vectors[:, i], 0):
-            return np.delete(old_vectors, i, axis=1)
-    U, S, _ = np.linalg.svd(old_vectors)
+    # expected = old_vectors @ old_vectors.T
+    # print(old_vectors[0])
+    # print(old_vectors[1])
+    # print(expected)
+    # print(np.cross(old_vectors[0], old_vectors[1]))
+    # for i in range(3):
+    #     if np.allclose(old_vectors[:, i], 0):
+    #         return np.delete(old_vectors, i, axis=1)
+    U, S, Vh = np.linalg.svd(old_vectors, full_matrices=False)
+    print(old_vectors)
+    print(Vh)
+    print("new", old_vectors @ Vh.T)
     new_vectors = U @ np.diag(S)
-    if np.linalg.det(new_vectors) < 0:
-        new_vectors = new_vectors[:, ::-1]
+    print(new_vectors)
+    print(U)
+    # if np.abs(U[0,0]) < np.abs(U[1,0]):
+    #     new_vectors = new_vectors[:, ::-1]
+    print(new_vectors)
+    # actual = new_vectors @ new_vectors.T
+    # print(new_vectors[0])
+    # print(new_vectors[1])
+    # print(np.linalg.det(new_vectors))
+    # print(actual)
+    # print(np.allclose(actual, expected))
+    # if not np.allclose(actual, expected):
+    # # print(new_vectors)
+    # # print(np.cross(*new_vectors))
+    # # if np.linalg.det(new_vectors) < 0:
+    # print(new_vectors)
     return new_vectors
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 029dc7cd..eac8580e 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -19,6 +19,8 @@ def test_unusual_orthorhombic(cut, indices, Assert):
     cell = np.roll(np.diag((3, 4, 5)), shift=1, axis=0)
     expected_plane = np.delete(np.delete(cell, indices[0], axis=0), indices[1], axis=1)
     actual_plane = slicing.plane(cell, cut)
+    print(expected_plane)
+    print(actual_plane)
     Assert.allclose(actual_plane, expected_plane)
 
 

From ca51c5ee50cffe8648b287f0e37f1c4b7ae897ea Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Wed, 10 Apr 2024 21:17:46 +0200
Subject: [PATCH 11/46] Generalize slicing algorithm so that everything uses
 the same code

---
 src/py4vasp/_util/slicing.py | 82 ++++++++++++++++++++++--------------
 tests/util/test_slicing.py   |  2 -
 2 files changed, 50 insertions(+), 34 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index fac5f3c3..21c3b16f 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -6,35 +6,53 @@
 
 
 def plane(cell, cut):
-    # old_det = np.linalg.det(cell)
-    old_vectors = np.delete(cell, INDICES[cut], axis=0)
-    # expected = old_vectors @ old_vectors.T
-    # print(old_vectors[0])
-    # print(old_vectors[1])
-    # print(expected)
-    # print(np.cross(old_vectors[0], old_vectors[1]))
-    # for i in range(3):
-    #     if np.allclose(old_vectors[:, i], 0):
-    #         return np.delete(old_vectors, i, axis=1)
-    U, S, Vh = np.linalg.svd(old_vectors, full_matrices=False)
-    print(old_vectors)
-    print(Vh)
-    print("new", old_vectors @ Vh.T)
-    new_vectors = U @ np.diag(S)
-    print(new_vectors)
-    print(U)
-    # if np.abs(U[0,0]) < np.abs(U[1,0]):
-    #     new_vectors = new_vectors[:, ::-1]
-    print(new_vectors)
-    # actual = new_vectors @ new_vectors.T
-    # print(new_vectors[0])
-    # print(new_vectors[1])
-    # print(np.linalg.det(new_vectors))
-    # print(actual)
-    # print(np.allclose(actual, expected))
-    # if not np.allclose(actual, expected):
-    # # print(new_vectors)
-    # # print(np.cross(*new_vectors))
-    # # if np.linalg.det(new_vectors) < 0:
-    # print(new_vectors)
-    return new_vectors
+    """Takes a 2d slice of a 3d cell and projects it onto 2d coordinates.
+
+    For simplicity in the documentation, we will assume that the cut is in the plane of
+    the a and b lattice vector. We use Greek letter α and β to refer to the vectors in
+    the 2d coordinate system. This routine computes the transformation Q such that
+
+    .. math::
+
+        \alpha &= Q a \\
+        \beta  &= Q b
+
+    Here, the matrix Q has a 2 × 3 shape so that the vectors α and β have only 2
+    dimensions. Furthermore, we want that the matrix Q fulfills certain properties,
+    namely that the lengths of the vectors and there angle is not changed. If the
+    vectors lie within e.g. the x-y plane, α and β should be the same as a and b just
+    with the zero for the z component removed.
+
+    Parameters
+    ----------
+    cell : np.ndarray
+        A 3 × 3 array defining the three lattice vectors of the unit cell.
+    cut : str
+        Either "a", "b", or "c" to define which lattice vector is removed to get the slice.
+
+    Returns
+    -------
+    np.ndarray
+        A 2 × 2 array defining the two lattice vectors spanning the plane.
+    """
+    vectors = np.delete(cell, INDICES[cut], axis=0)
+    *_, Q = np.linalg.svd(vectors, full_matrices=False)
+    Q = _make_all_largest_components_positive(Q)
+    Q = _ensure_same_order_as_input_array(Q)
+    return vectors @ Q.T
+
+def _make_all_largest_components_positive(vectors):
+    return np.array([_make_largest_component_positive(vector) for vector in vectors])
+
+def _make_largest_component_positive(vector):
+    if np.max(vector) > -np.min(vector):
+        return vector
+    else:
+        return -vector
+
+def _ensure_same_order_as_input_array(Q):
+    i, j = np.argmax(Q, axis=1)
+    if i <= j:
+        return Q
+    else:
+        return Q[::-1]
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index eac8580e..029dc7cd 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -19,8 +19,6 @@ def test_unusual_orthorhombic(cut, indices, Assert):
     cell = np.roll(np.diag((3, 4, 5)), shift=1, axis=0)
     expected_plane = np.delete(np.delete(cell, indices[0], axis=0), indices[1], axis=1)
     actual_plane = slicing.plane(cell, cut)
-    print(expected_plane)
-    print(actual_plane)
     Assert.allclose(actual_plane, expected_plane)
 
 

From 9e5516c73e63d408dcb13c4ca71de54ce4d53dd6 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Fri, 12 Apr 2024 15:57:37 +0200
Subject: [PATCH 12/46] Refactor routine to rotate to largest component

---
 src/py4vasp/_util/slicing.py | 37 +++++++++++++++++-------------------
 1 file changed, 17 insertions(+), 20 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index 21c3b16f..8748dcca 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -36,23 +36,20 @@ def plane(cell, cut):
         A 2 × 2 array defining the two lattice vectors spanning the plane.
     """
     vectors = np.delete(cell, INDICES[cut], axis=0)
-    *_, Q = np.linalg.svd(vectors, full_matrices=False)
-    Q = _make_all_largest_components_positive(Q)
-    Q = _ensure_same_order_as_input_array(Q)
-    return vectors @ Q.T
-
-def _make_all_largest_components_positive(vectors):
-    return np.array([_make_largest_component_positive(vector) for vector in vectors])
-
-def _make_largest_component_positive(vector):
-    if np.max(vector) > -np.min(vector):
-        return vector
-    else:
-        return -vector
-
-def _ensure_same_order_as_input_array(Q):
-    i, j = np.argmax(Q, axis=1)
-    if i <= j:
-        return Q
-    else:
-        return Q[::-1]
+    axis = np.cross(*vectors).astype(np.float_)
+    axis /= np.linalg.norm(axis)
+    closest_cartesian_axis = np.argmax(np.abs(axis))
+    cartesian_axis = np.zeros(3)
+    cartesian_axis[closest_cartesian_axis] = np.sign(axis[closest_cartesian_axis])
+    rotation_matrix = _calculate_rotation_matrix((axis, cartesian_axis))
+    new_vectors = vectors @ rotation_matrix.T
+    return np.delete(new_vectors, closest_cartesian_axis, axis=1)
+
+
+def _calculate_rotation_matrix(vectors):
+    cos_angle = np.dot(*vectors)
+    v = np.cross(*vectors)
+    if np.linalg.norm(v) < 1e-10:
+        return np.eye(3)
+    V = np.cross(np.eye(3), v)
+    return np.eye(3) + V + V @ V / (1 + cos_angle)

From 350b4bd6b64c8326655f5441ff5017bf896a163b Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Fri, 12 Apr 2024 16:12:18 +0200
Subject: [PATCH 13/46] Raise error when axis is not obvious

---
 src/py4vasp/_util/slicing.py | 30 +++++++++++++++++++++++++-----
 tests/util/test_slicing.py   |  6 ++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index 8748dcca..b313e198 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -2,10 +2,12 @@
 # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
 import numpy as np
 
+from py4vasp import exception
+
 INDICES = {"a": 0, "b": 1, "c": 2}
 
 
-def plane(cell, cut):
+def plane(cell, cut, direction="auto"):
     """Takes a 2d slice of a 3d cell and projects it onto 2d coordinates.
 
     For simplicity in the documentation, we will assume that the cut is in the plane of
@@ -29,6 +31,10 @@ def plane(cell, cut):
         A 3 × 3 array defining the three lattice vectors of the unit cell.
     cut : str
         Either "a", "b", or "c" to define which lattice vector is removed to get the slice.
+    direction : str
+        Set the Cartesian direction "x", "y", or "c" parallel to which the normal of
+        the plane is rotated. Alteratively, set it to "auto" to rotate to the closest
+        Cartesian axis.
 
     Returns
     -------
@@ -38,12 +44,26 @@ def plane(cell, cut):
     vectors = np.delete(cell, INDICES[cut], axis=0)
     axis = np.cross(*vectors).astype(np.float_)
     axis /= np.linalg.norm(axis)
-    closest_cartesian_axis = np.argmax(np.abs(axis))
-    cartesian_axis = np.zeros(3)
-    cartesian_axis[closest_cartesian_axis] = np.sign(axis[closest_cartesian_axis])
+    index_axis, cartesian_axis = _get_cartesian_axis(axis, direction)
     rotation_matrix = _calculate_rotation_matrix((axis, cartesian_axis))
     new_vectors = vectors @ rotation_matrix.T
-    return np.delete(new_vectors, closest_cartesian_axis, axis=1)
+    return np.delete(new_vectors, index_axis, axis=1)
+
+
+def _get_cartesian_axis(axis, direction):
+    if direction == "auto":
+        closest_cartesian_axis = np.argmax(np.abs(axis))
+        if np.abs(axis[closest_cartesian_axis]) < 0.8:
+            message = """\
+You did not specify the Cartesian direction to which the normal of the cut plane will
+be rotated to. py4vasp tries to determine the axis automatically but in this case no
+axis is close to the normal of the plane. Please pass the additional argument *direction*
+("x", "y", or "z") to specify to which axis py4vasp should use as normal vector for the
+plane."""
+            raise exception.IncorrectUsage(message)
+        cartesian_axis = np.zeros(3)
+        cartesian_axis[closest_cartesian_axis] = np.sign(axis[closest_cartesian_axis])
+        return closest_cartesian_axis, cartesian_axis
 
 
 def _calculate_rotation_matrix(vectors):
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 029dc7cd..3c8f9db3 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -3,6 +3,7 @@
 import numpy as np
 import pytest
 
+from py4vasp import exception
 from py4vasp._util import slicing
 
 
@@ -32,3 +33,8 @@ def test_nearly_orthorhombic(cut, index, Assert):
     actual_plane = slicing.plane(cell, cut)
     Assert.allclose(actual_plane @ actual_plane.T, expected_plane @ expected_plane.T)
     assert np.max(np.abs(approximate_plane - actual_plane)) < 0.1
+
+
+def test_raise_error_if_direction_is_not_obvious():
+    with pytest.raises(exception.IncorrectUsage):
+        slicing.plane(np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]), "a")

From b0a55939cf5b2f2c7e48221395307ae7968de3d5 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 15 Apr 2024 09:07:33 +0200
Subject: [PATCH 14/46] WIP: slicing

---
 src/py4vasp/_util/slicing.py | 4 +++-
 tests/util/test_slicing.py   | 3 +++
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index b313e198..5b170edb 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -5,7 +5,8 @@
 from py4vasp import exception
 
 INDICES = {"a": 0, "b": 1, "c": 2}
-
+AXIS = {"x": (1, 0, 0), "+x": (1, 0, 0), "-x": (-1, 0, 0), "y": (0, 1, 0), "+y": (0, 1, 0),
+        "-y": (0, -1, 0), "z": (0, 0, 1), "+z": (0, 0, 1), "-z": (0, 0, -1)}
 
 def plane(cell, cut, direction="auto"):
     """Takes a 2d slice of a 3d cell and projects it onto 2d coordinates.
@@ -44,6 +45,7 @@ def plane(cell, cut, direction="auto"):
     vectors = np.delete(cell, INDICES[cut], axis=0)
     axis = np.cross(*vectors).astype(np.float_)
     axis /= np.linalg.norm(axis)
+    print(axis)
     index_axis, cartesian_axis = _get_cartesian_axis(axis, direction)
     rotation_matrix = _calculate_rotation_matrix((axis, cartesian_axis))
     new_vectors = vectors @ rotation_matrix.T
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 3c8f9db3..b876402f 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -38,3 +38,6 @@ def test_nearly_orthorhombic(cut, index, Assert):
 def test_raise_error_if_direction_is_not_obvious():
     with pytest.raises(exception.IncorrectUsage):
         slicing.plane(np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]), "a")
+
+@pytest.mark.parametrize("direction", ["x", ""])
+def test_
\ No newline at end of file

From 43f3e70d2935eb81fe6ae1e5ec23cc01b9d0d692 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 15 Apr 2024 09:48:00 +0200
Subject: [PATCH 15/46] Allow manual selection of normal direction

---
 src/py4vasp/_util/slicing.py | 46 ++++++++++++++++++++----------------
 tests/util/test_slicing.py   | 18 ++++++++++++--
 2 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index 5b170edb..da88c928 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -5,10 +5,10 @@
 from py4vasp import exception
 
 INDICES = {"a": 0, "b": 1, "c": 2}
-AXIS = {"x": (1, 0, 0), "+x": (1, 0, 0), "-x": (-1, 0, 0), "y": (0, 1, 0), "+y": (0, 1, 0),
-        "-y": (0, -1, 0), "z": (0, 0, 1), "+z": (0, 0, 1), "-z": (0, 0, -1)}
+AXIS = ("x", "y", "z")
 
-def plane(cell, cut, direction="auto"):
+
+def plane(cell, cut, normal="auto"):
     """Takes a 2d slice of a 3d cell and projects it onto 2d coordinates.
 
     For simplicity in the documentation, we will assume that the cut is in the plane of
@@ -32,7 +32,7 @@ def plane(cell, cut, direction="auto"):
         A 3 × 3 array defining the three lattice vectors of the unit cell.
     cut : str
         Either "a", "b", or "c" to define which lattice vector is removed to get the slice.
-    direction : str
+    normal : str
         Set the Cartesian direction "x", "y", or "c" parallel to which the normal of
         the plane is rotated. Alteratively, set it to "auto" to rotate to the closest
         Cartesian axis.
@@ -45,27 +45,21 @@ def plane(cell, cut, direction="auto"):
     vectors = np.delete(cell, INDICES[cut], axis=0)
     axis = np.cross(*vectors).astype(np.float_)
     axis /= np.linalg.norm(axis)
-    print(axis)
-    index_axis, cartesian_axis = _get_cartesian_axis(axis, direction)
+    index_axis, cartesian_axis = _get_cartesian_axis(axis, normal)
     rotation_matrix = _calculate_rotation_matrix((axis, cartesian_axis))
     new_vectors = vectors @ rotation_matrix.T
     return np.delete(new_vectors, index_axis, axis=1)
 
 
-def _get_cartesian_axis(axis, direction):
-    if direction == "auto":
-        closest_cartesian_axis = np.argmax(np.abs(axis))
-        if np.abs(axis[closest_cartesian_axis]) < 0.8:
-            message = """\
-You did not specify the Cartesian direction to which the normal of the cut plane will
-be rotated to. py4vasp tries to determine the axis automatically but in this case no
-axis is close to the normal of the plane. Please pass the additional argument *direction*
-("x", "y", or "z") to specify to which axis py4vasp should use as normal vector for the
-plane."""
-            raise exception.IncorrectUsage(message)
-        cartesian_axis = np.zeros(3)
-        cartesian_axis[closest_cartesian_axis] = np.sign(axis[closest_cartesian_axis])
-        return closest_cartesian_axis, cartesian_axis
+def _get_cartesian_axis(axis, normal):
+    if normal in AXIS:
+        index = AXIS.index(normal)
+    elif normal == "auto":
+        index = np.argmax(np.abs(axis))
+        _raise_error_if_direction_is_not_obvious(np.abs(axis[index]))
+    cartesian_axis = np.zeros(3)
+    cartesian_axis[index] = np.sign(axis[index])
+    return index, cartesian_axis
 
 
 def _calculate_rotation_matrix(vectors):
@@ -75,3 +69,15 @@ def _calculate_rotation_matrix(vectors):
         return np.eye(3)
     V = np.cross(np.eye(3), v)
     return np.eye(3) + V + V @ V / (1 + cos_angle)
+
+
+def _raise_error_if_direction_is_not_obvious(largest_element):
+    if largest_element > 0.95:
+        return
+    message = """\
+You did not specify the Cartesian direction to which the normal of the cut plane will
+be rotated to. py4vasp tries to determine the axis automatically but in this case no
+axis is close to the normal of the plane. Please pass the additional argument *normal*
+("x", "y", or "z") to specify to which axis py4vasp should use as normal vector for the
+plane."""
+    raise exception.IncorrectUsage(message)
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index b876402f..ecaae2a9 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -39,5 +39,19 @@ def test_raise_error_if_direction_is_not_obvious():
     with pytest.raises(exception.IncorrectUsage):
         slicing.plane(np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]), "a")
 
-@pytest.mark.parametrize("direction", ["x", ""])
-def test_
\ No newline at end of file
+
+XX = np.sqrt(1 - np.sqrt(0.75))
+
+
+@pytest.mark.parametrize(
+    "normal, expected_plane",
+    [
+        ("x", [[XX, 1 + XX], [1 + XX, XX]]),
+        ("y", [[1, 1], [1 + XX, -XX]]),
+        ("z", [[XX + 1, -XX], [1, 1]]),
+    ],
+)
+def test_rotate_to_user_defined_axis(normal, expected_plane, Assert):
+    cell = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
+    actual_plane = slicing.plane(cell, "a", normal)
+    Assert.allclose(actual_plane, expected_plane)

From fdf9a1383a48644561c0c63319692b444aea4d82 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 15 Apr 2024 14:09:55 +0200
Subject: [PATCH 16/46] Add test for providing normal with orthorhombic cell

---
 src/py4vasp/_util/slicing.py |  3 ++-
 tests/util/test_slicing.py   | 11 +++++++++++
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index da88c928..922f3e81 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -58,7 +58,8 @@ def _get_cartesian_axis(axis, normal):
         index = np.argmax(np.abs(axis))
         _raise_error_if_direction_is_not_obvious(np.abs(axis[index]))
     cartesian_axis = np.zeros(3)
-    cartesian_axis[index] = np.sign(axis[index])
+    # do not use sign function because it is 0 if axis[index] == 0
+    cartesian_axis[index] = 1 if axis[index] >= 0 else -1
     return index, cartesian_axis
 
 
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index ecaae2a9..c3080acf 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -55,3 +55,14 @@ def test_rotate_to_user_defined_axis(normal, expected_plane, Assert):
     cell = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
     actual_plane = slicing.plane(cell, "a", normal)
     Assert.allclose(actual_plane, expected_plane)
+
+
+@pytest.mark.parametrize(
+    "normal, rotation",
+    (("x", [[1, 0], [0, 1]]), ("y", [[-1, 0], [0, 1]]), ("z", [[0, 1], [-1, 0]])),
+)
+def test_normal_with_orthonormal_cell(normal, rotation, Assert):
+    cell = np.diag((3, 4, 5))
+    expected_plane = np.dot(np.delete(np.delete(cell, 0, axis=0), 0, axis=1), rotation)
+    actual_plane = slicing.plane(cell, "a", normal)
+    Assert.allclose(actual_plane, expected_plane)

From db5e923468a5c268f10b961d8a2b248bd817fe96 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 15 Apr 2024 22:01:06 +0200
Subject: [PATCH 17/46] Allow disabling rotation altogether

---
 src/py4vasp/_util/slicing.py | 66 ++++++++++++++++++++++++------------
 tests/util/test_slicing.py   | 19 +++++++++++
 2 files changed, 63 insertions(+), 22 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index 922f3e81..ac3c9b2b 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -33,7 +33,7 @@ def plane(cell, cut, normal="auto"):
     cut : str
         Either "a", "b", or "c" to define which lattice vector is removed to get the slice.
     normal : str
-        Set the Cartesian direction "x", "y", or "c" parallel to which the normal of
+        Set the Cartesian direction "x", "y", or "z" parallel to which the normal of
         the plane is rotated. Alteratively, set it to "auto" to rotate to the closest
         Cartesian axis.
 
@@ -43,33 +43,39 @@ def plane(cell, cut, normal="auto"):
         A 2 × 2 array defining the two lattice vectors spanning the plane.
     """
     vectors = np.delete(cell, INDICES[cut], axis=0)
-    axis = np.cross(*vectors).astype(np.float_)
-    axis /= np.linalg.norm(axis)
-    index_axis, cartesian_axis = _get_cartesian_axis(axis, normal)
-    rotation_matrix = _calculate_rotation_matrix((axis, cartesian_axis))
+    if normal is not None:
+        return _rotate_normal_to_cartesian_axis(vectors, normal)
+    else:
+        return _rotate_first_vector_to_x_axis(vectors)
+
+
+def _rotate_first_vector_to_x_axis(vectors):
+    u, v = np.linalg.norm(vectors, axis=1)
+    x = np.dot(*vectors) / (u * v)
+    return np.array([[u, 0], [v * x, v * np.sqrt(1 - x**2)]])
+
+
+def _rotate_normal_to_cartesian_axis(vectors, normal):
+    old_normal = _get_old_normal(vectors)
+    index_axis = _get_index_axis(old_normal, normal)
+    new_normal = _get_new_normal_from_cartesian_axis(old_normal, index_axis)
+    rotation_matrix = _get_rotation_matrix((old_normal, new_normal))
     new_vectors = vectors @ rotation_matrix.T
     return np.delete(new_vectors, index_axis, axis=1)
 
 
-def _get_cartesian_axis(axis, normal):
-    if normal in AXIS:
-        index = AXIS.index(normal)
-    elif normal == "auto":
-        index = np.argmax(np.abs(axis))
-        _raise_error_if_direction_is_not_obvious(np.abs(axis[index]))
-    cartesian_axis = np.zeros(3)
-    # do not use sign function because it is 0 if axis[index] == 0
-    cartesian_axis[index] = 1 if axis[index] >= 0 else -1
-    return index, cartesian_axis
+def _get_old_normal(vectors):
+    old_normal = np.cross(*vectors).astype(np.float_)
+    return old_normal / np.linalg.norm(old_normal)
 
 
-def _calculate_rotation_matrix(vectors):
-    cos_angle = np.dot(*vectors)
-    v = np.cross(*vectors)
-    if np.linalg.norm(v) < 1e-10:
-        return np.eye(3)
-    V = np.cross(np.eye(3), v)
-    return np.eye(3) + V + V @ V / (1 + cos_angle)
+def _get_index_axis(old_normal, normal):
+    if normal in AXIS:
+        return AXIS.index(normal)
+    elif normal == "auto":
+        index = np.argmax(np.abs(old_normal))
+        _raise_error_if_direction_is_not_obvious(np.abs(old_normal[index]))
+        return index
 
 
 def _raise_error_if_direction_is_not_obvious(largest_element):
@@ -82,3 +88,19 @@ def _raise_error_if_direction_is_not_obvious(largest_element):
 ("x", "y", or "z") to specify to which axis py4vasp should use as normal vector for the
 plane."""
     raise exception.IncorrectUsage(message)
+
+
+def _get_new_normal_from_cartesian_axis(old_normal, index):
+    new_normal = np.zeros(3)
+    # do not use sign function because it is 0 if old_normal[index] == 0
+    new_normal[index] = 1 if old_normal[index] >= 0 else -1
+    return new_normal
+
+
+def _get_rotation_matrix(vectors):
+    cos_angle = np.dot(*vectors)
+    v = np.cross(*vectors)
+    if np.linalg.norm(v) < 1e-10:
+        return np.eye(3)
+    V = np.cross(np.eye(3), v)
+    return np.eye(3) + V + V @ V / (1 + cos_angle)
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index c3080acf..6050b3d8 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -66,3 +66,22 @@ def test_normal_with_orthonormal_cell(normal, rotation, Assert):
     expected_plane = np.dot(np.delete(np.delete(cell, 0, axis=0), 0, axis=1), rotation)
     actual_plane = slicing.plane(cell, "a", normal)
     Assert.allclose(actual_plane, expected_plane)
+
+
+@pytest.mark.parametrize("cut, diagonal", [("a", [4, 5]), ("b", [3, 5]), ("c", [3, 4])])
+def test_no_rotation_orthorhombic_cell(cut, diagonal, Assert):
+    cell = np.diag((3, 4, 5))
+    expected_plane = np.diag(diagonal)
+    actual_plane = slicing.plane(cell, cut, normal=None)
+    Assert.allclose(actual_plane, expected_plane)
+
+
+@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)])
+def test_no_rotation_nontrivial_cell(cut, index, Assert):
+    cell = np.array([[0, 1, 1.1], [0.9, 0, 1.1], [0.9, 1, 0]])
+    vectors = np.delete(cell, index, axis=0)
+    expected_products = vectors @ vectors.T
+    actual_plane = slicing.plane(cell, cut, normal=None)
+    Assert.allclose(actual_plane[0, 1], 0)
+    actual_products = actual_plane @ actual_plane.T
+    Assert.allclose(actual_products, expected_products)

From 9b2d1ef31e27c431914e6f68a0078dbeb8904073 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 15 Apr 2024 22:18:18 +0200
Subject: [PATCH 18/46] Add sanity checks and documentation

---
 src/py4vasp/_util/slicing.py | 58 ++++++++++++++++++++++++++----------
 tests/util/test_slicing.py   | 17 +++++++----
 2 files changed, 54 insertions(+), 21 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index ac3c9b2b..fb9bd759 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -22,9 +22,14 @@ def plane(cell, cut, normal="auto"):
 
     Here, the matrix Q has a 2 × 3 shape so that the vectors α and β have only 2
     dimensions. Furthermore, we want that the matrix Q fulfills certain properties,
-    namely that the lengths of the vectors and there angle is not changed. If the
+    namely that the lengths of the vectors and their angle is not changed. If the
     vectors lie within e.g. the x-y plane, α and β should be the same as a and b just
-    with the zero for the z component removed.
+    with the zero for the z component removed. However, to avoid irregular behavior
+    where small changes to the unit cell cause large changes to the selected vectors
+    this mode only works if the normal vector is close (~15°) to a Cartesian axis.
+    Otherwise the user needs to provide the axis with the *normal* argument.
+    Alternatively, one can align the first remaining lattice vector with the x axis by
+    passing None for *normal*.
 
     Parameters
     ----------
@@ -35,13 +40,15 @@ def plane(cell, cut, normal="auto"):
     normal : str
         Set the Cartesian direction "x", "y", or "z" parallel to which the normal of
         the plane is rotated. Alteratively, set it to "auto" to rotate to the closest
-        Cartesian axis.
+        Cartesian axis. If you set it to None, the normal will not be considered and
+        the first remaining lattice vector will be aligned with the x axis instead.
 
     Returns
     -------
     np.ndarray
         A 2 × 2 array defining the two lattice vectors spanning the plane.
     """
+    _raise_error_if_cut_unknown(cut)
     vectors = np.delete(cell, INDICES[cut], axis=0)
     if normal is not None:
         return _rotate_normal_to_cartesian_axis(vectors, normal)
@@ -76,19 +83,8 @@ def _get_index_axis(old_normal, normal):
         index = np.argmax(np.abs(old_normal))
         _raise_error_if_direction_is_not_obvious(np.abs(old_normal[index]))
         return index
-
-
-def _raise_error_if_direction_is_not_obvious(largest_element):
-    if largest_element > 0.95:
-        return
-    message = """\
-You did not specify the Cartesian direction to which the normal of the cut plane will
-be rotated to. py4vasp tries to determine the axis automatically but in this case no
-axis is close to the normal of the plane. Please pass the additional argument *normal*
-("x", "y", or "z") to specify to which axis py4vasp should use as normal vector for the
-plane."""
-    raise exception.IncorrectUsage(message)
-
+    else:
+        _raise_unknown_normal_error(normal)
 
 def _get_new_normal_from_cartesian_axis(old_normal, index):
     new_normal = np.zeros(3)
@@ -104,3 +100,33 @@ def _get_rotation_matrix(vectors):
         return np.eye(3)
     V = np.cross(np.eye(3), v)
     return np.eye(3) + V + V @ V / (1 + cos_angle)
+
+
+def _raise_error_if_cut_unknown(cut):
+    if cut in INDICES:
+        return
+    message = """\
+The selected choice {cut} is invalid. Please select one of the three lattice vectors
+with passing *cut* ("a", "b", or "c").""".format(
+        cut=cut
+    )
+    raise exception.IncorrectUsage(message)
+
+
+def _raise_error_if_direction_is_not_obvious(largest_element):
+    if largest_element > 0.95:
+        return
+    message = """\
+You did not specify the Cartesian direction to which the normal of the cut plane will
+be rotated to. py4vasp tries to determine the axis automatically but in this case no
+axis is close to the normal of the plane. Please pass the additional argument *normal*
+("x", "y", or "z") to specify to which axis py4vasp should use as normal vector for the
+plane."""
+    raise exception.IncorrectUsage(message)
+
+def _raise_unknown_normal_error(normal):
+    message = """\
+The selected normal {normal} is invalid. Please select one of the Cartesian axis ("x",
+"y", or "z") or "auto" to specify to which axis the normal is rotated. "auto" will use
+a closeby Cartesian axis if possible.""".format(normal=normal)
+    raise exception.IncorrectUsage(message)
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 6050b3d8..1acff760 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -35,11 +35,6 @@ def test_nearly_orthorhombic(cut, index, Assert):
     assert np.max(np.abs(approximate_plane - actual_plane)) < 0.1
 
 
-def test_raise_error_if_direction_is_not_obvious():
-    with pytest.raises(exception.IncorrectUsage):
-        slicing.plane(np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]), "a")
-
-
 XX = np.sqrt(1 - np.sqrt(0.75))
 
 
@@ -85,3 +80,15 @@ def test_no_rotation_nontrivial_cell(cut, index, Assert):
     Assert.allclose(actual_plane[0, 1], 0)
     actual_products = actual_plane @ actual_plane.T
     Assert.allclose(actual_products, expected_products)
+
+
+def test_raise_error_if_normal_is_not_obvious():
+    with pytest.raises(exception.IncorrectUsage):
+        slicing.plane(np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]]), "a")
+
+
+def test_raise_error_for_unknown_choices():
+    with pytest.raises(exception.IncorrectUsage):
+        slicing.plane(np.eye(3), "unknown")
+    with pytest.raises(exception.IncorrectUsage):
+        slicing.plane(np.eye(3), "a", normal="unknown")

From 77d5873a71a5e99219ea429b54861456d5632dbe Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Mon, 15 Apr 2024 22:23:37 +0200
Subject: [PATCH 19/46] Update format

---
 src/py4vasp/_util/slicing.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index fb9bd759..f7f009c4 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -86,6 +86,7 @@ def _get_index_axis(old_normal, normal):
     else:
         _raise_unknown_normal_error(normal)
 
+
 def _get_new_normal_from_cartesian_axis(old_normal, index):
     new_normal = np.zeros(3)
     # do not use sign function because it is 0 if old_normal[index] == 0
@@ -107,10 +108,8 @@ def _raise_error_if_cut_unknown(cut):
         return
     message = """\
 The selected choice {cut} is invalid. Please select one of the three lattice vectors
-with passing *cut* ("a", "b", or "c").""".format(
-        cut=cut
-    )
-    raise exception.IncorrectUsage(message)
+with passing *cut* ("a", "b", or "c")."""
+    raise exception.IncorrectUsage(message.format(cut=cut))
 
 
 def _raise_error_if_direction_is_not_obvious(largest_element):
@@ -124,9 +123,10 @@ def _raise_error_if_direction_is_not_obvious(largest_element):
 plane."""
     raise exception.IncorrectUsage(message)
 
+
 def _raise_unknown_normal_error(normal):
     message = """\
 The selected normal {normal} is invalid. Please select one of the Cartesian axis ("x",
 "y", or "z") or "auto" to specify to which axis the normal is rotated. "auto" will use
-a closeby Cartesian axis if possible.""".format(normal=normal)
-    raise exception.IncorrectUsage(message)
+a closeby Cartesian axis if possible."""
+    raise exception.IncorrectUsage(message.format(normal=normal))

From ac53e52ad6b7b85c173baf3b9a1fedd3b2972874 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 16 Apr 2024 11:21:50 +0200
Subject: [PATCH 20/46] WIP: density contour

[skip CI]
---
 src/py4vasp/calculation/_density.py |  8 +++++++-
 tests/calculation/test_density.py   | 15 ++++++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

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)

From a021abd350b77a5d8af0a3dec4709e49d137b2b6 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 16 Apr 2024 12:04:07 +0200
Subject: [PATCH 21/46] Implement helper routine for slicing

---
 src/py4vasp/_util/slicing.py | 34 ++++++++++++++++++++++++++++++++++
 tests/util/test_slicing.py   | 17 +++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index f7f009c4..cfe0db30 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -8,6 +8,36 @@
 AXIS = ("x", "y", "z")
 
 
+def grid_data(data, cut, fraction):
+    """Takes a 2d slice of a 3d grid data.
+
+    Often, we want to generate a 2d slice of a 3d data on a grid. One example would be
+    to visualize a slice through a plane as a contour plot. This routine facilitates
+    this task implementing taking the cut of the 3d data.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        3d data on a grid from which a 2d plane is extracted.
+    cut : str
+        Defines along which dimension of the grid to cut (either "a", "b", or "c").
+    fraction : float
+        Determines which plane of the grid data is used. Periodic boundaries are assumed.
+
+    Returns
+    -------
+    np.ndarray
+        A 2d array where the dimension selected by cut has been removed by selecting
+        the plane according to the specificied fraction.
+    """
+    _raise_error_if_cut_unknown(cut)
+    index = INDICES[cut]
+    length = data.shape[index]
+    slice_ = [slice(None), slice(None), slice(None)]
+    slice_[index] = np.round(length * fraction).astype(np.int_) % length
+    return data[tuple(slice_)]
+
+
 def plane(cell, cut, normal="auto"):
     """Takes a 2d slice of a 3d cell and projects it onto 2d coordinates.
 
@@ -103,6 +133,10 @@ def _get_rotation_matrix(vectors):
     return np.eye(3) + V + V @ V / (1 + cos_angle)
 
 
+def _get_slice(shape, cut, fraction):
+    return slice_
+
+
 def _raise_error_if_cut_unknown(cut):
     if cut in INDICES:
         return
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 1acff760..4870803e 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -92,3 +92,20 @@ def test_raise_error_for_unknown_choices():
         slicing.plane(np.eye(3), "unknown")
     with pytest.raises(exception.IncorrectUsage):
         slicing.plane(np.eye(3), "a", normal="unknown")
+
+
+@pytest.mark.parametrize("cut", ("a", "b", "c"))
+@pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2))
+def test_slice_data(cut, fraction, Assert):
+    grid_data = np.random.random((10, 12, 14)) + 0.1
+    if cut == "a":
+        index = np.round(fraction * 10).astype(np.int_) % 10
+        expected_data = grid_data[index, :, :]
+    elif cut == "b":
+        index = np.round(fraction * 12).astype(np.int_) % 12
+        expected_data = grid_data[:, index, :]
+    else:
+        index = np.round(fraction * 14).astype(np.int_) % 14
+        expected_data = grid_data[:, :, index]
+    actual_data = slicing.grid_data(grid_data, cut, fraction)
+    Assert.allclose(actual_data, expected_data)

From a2f5811921fe4cdc714cc26a58ed60c93b444fa3 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 16 Apr 2024 14:15:01 +0200
Subject: [PATCH 22/46] Expose lattice vectors and positions to users

---
 src/py4vasp/calculation/_structure.py | 45 +++++++++++++++++++--------
 tests/calculation/test_structure.py   | 12 +++++++
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/src/py4vasp/calculation/_structure.py b/src/py4vasp/calculation/_structure.py
index ecffaa28..81510c3d 100644
--- a/src/py4vasp/calculation/_structure.py
+++ b/src/py4vasp/calculation/_structure.py
@@ -144,8 +144,8 @@ def to_dict(self):
         {examples}
         """
         return {
-            "lattice_vectors": self._lattice_vectors(),
-            "positions": self._positions(),
+            "lattice_vectors": self.lattice_vectors(),
+            "positions": self.positions(),
             "elements": self._topology().elements(),
             "names": self._topology().names(),
         }
@@ -171,8 +171,8 @@ def to_view(self, supercell=None):
         make_3d = lambda array: array if array.ndim == 3 else array[np.newaxis]
         return view.View(
             elements=np.atleast_2d(self._topology().elements()),
-            lattice_vectors=make_3d(self._lattice_vectors()),
-            positions=make_3d(self._positions()),
+            lattice_vectors=make_3d(self.lattice_vectors()),
+            positions=make_3d(self.positions()),
             supercell=self._parse_supercell(supercell),
         )
 
@@ -261,6 +261,32 @@ def to_POSCAR(self):
             message = "Converting multiple structures to a POSCAR is currently not implemented."
             raise exception.NotImplemented(message)
 
+    @_base.data_access
+    def lattice_vectors(self):
+        """Return the lattice vectors spanning the unit cell
+
+        Returns
+        -------
+        np.ndarray
+            Lattice vectors of the unit cell in Å.
+        """
+        lattice_vectors = _LatticeVectors(self._raw_data.cell.lattice_vectors)
+        return self._scale() * lattice_vectors[self._get_steps()]
+
+    @_base.data_access
+    def positions(self):
+        """Return the direct coordinates of all ions in the unit cell.
+
+        Direct or fractional coordinates measure the position of the ions in terms of
+        the lattice vectors. Hence they are dimensionless quantities.
+
+        Returns
+        -------
+        np.ndarray
+            Positions of all ions in terms of the lattice vectors.
+        """
+        return self._raw_data.positions[self._get_steps()]
+
     @_base.data_access
     @documentation.format(examples=_slice.examples("structure", "cartesian_positions"))
     def cartesian_positions(self):
@@ -273,7 +299,7 @@ def cartesian_positions(self):
 
         {examples}
         """
-        return self._positions() @ self._lattice_vectors()
+        return self.positions() @ self.lattice_vectors()
 
     @_base.data_access
     @documentation.format(examples=_slice.examples("structure", "volume"))
@@ -287,7 +313,7 @@ def volume(self):
 
         {examples}
         """
-        return np.abs(np.linalg.det(self._lattice_vectors()))
+        return np.abs(np.linalg.det(self.lattice_vectors()))
 
     @_base.data_access
     def number_atoms(self):
@@ -331,10 +357,6 @@ def _parse_supercell(self, supercell):
     def _topology(self):
         return calculation.topology.from_data(self._raw_data.topology)
 
-    def _lattice_vectors(self):
-        lattice_vectors = _LatticeVectors(self._raw_data.cell.lattice_vectors)
-        return self._scale() * lattice_vectors[self._get_steps()]
-
     def _scale(self):
         if isinstance(self._raw_data.cell.scale, np.float_):
             return self._raw_data.cell.scale
@@ -343,9 +365,6 @@ def _scale(self):
         else:
             return 1.0
 
-    def _positions(self):
-        return self._raw_data.positions[self._get_steps()]
-
     def _get_steps(self):
         return self._steps if self._is_trajectory else ()
 
diff --git a/tests/calculation/test_structure.py b/tests/calculation/test_structure.py
index 08d99f1e..2ee5bbe7 100644
--- a/tests/calculation/test_structure.py
+++ b/tests/calculation/test_structure.py
@@ -264,6 +264,18 @@ def test_supercell_wrong_size(Sr2TiO4, not_a_supercell, not_core):
         Sr2TiO4.to_ase(not_a_supercell)
 
 
+@pytest.mark.parametrize("steps", [-1, 0, slice(1, 3), slice(None)])
+def test_lattice_vectors(Sr2TiO4, steps, Assert):
+    structure = Sr2TiO4 if steps == -1 else Sr2TiO4[steps]
+    Assert.allclose(structure.lattice_vectors(), Sr2TiO4.ref.lattice_vectors[steps])
+
+
+@pytest.mark.parametrize("steps", [-1, 0, slice(1, 3), slice(None)])
+def test_positions(Sr2TiO4, steps, Assert):
+    structure = Sr2TiO4 if steps == -1 else Sr2TiO4[steps]
+    Assert.allclose(structure.positions(), Sr2TiO4.ref.positions[steps])
+
+
 def test_cartesian_positions(Sr2TiO4, Fe3O4, Ca3AsBr3, Assert, not_core):
     check_cartesian_positions(Sr2TiO4, Assert)
     check_cartesian_positions(Fe3O4, Assert)

From f08146e0380f96621c3d732bb3496b4e914aa45f Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 16 Apr 2024 14:20:57 +0200
Subject: [PATCH 23/46] Implement basic Contour plot

---
 src/py4vasp/calculation/_density.py | 7 ++++---
 tests/calculation/test_density.py   | 4 ++++
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index 3c0c0d4c..4cb77d1d 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -4,7 +4,7 @@
 
 from py4vasp import _config, calculation, exception
 from py4vasp._third_party import graph, view
-from py4vasp._util import documentation, import_, index, select
+from py4vasp._util import documentation, import_, index, select, slicing
 from py4vasp.calculation import _base, _structure
 
 pretty = import_.optional("IPython.lib.pretty")
@@ -227,8 +227,9 @@ def to_view(self, selection=None, supercell=None, **user_options):
 
     @_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")
+        data = slicing.grid_data(self._raw_data.charge[0].T, "a", a)
+        lattice = slicing.plane(self._structure.lattice_vectors(), "a", normal=None)
+        contour = graph.Contour(data, lattice, "charge")
         return graph.Graph(contour)
 
     def _filter_noncollinear_magnetization_from_selections(self, tree):
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 0862f3f2..6a84e71e 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -260,6 +260,10 @@ def test_contour_of_slice(nonpolarized_density, Assert):
     slice_ = nonpolarized_density.ref.output["charge"][1]
     assert len(graph) == 1
     Assert.allclose(graph.series.data, slice_)
+    lattice_vectors = nonpolarized_density.ref.structure.lattice_vectors()
+    expected_products = lattice_vectors[1:] @ lattice_vectors[1:].T
+    actual_products = graph.series.lattice @ graph.series.lattice.T
+    Assert.allclose(actual_products, expected_products)
 
 
 # TODO: a, b, c

From 24c5f8da1a43c3135f1afe95aaea943277c599e2 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 16 Apr 2024 14:47:43 +0200
Subject: [PATCH 24/46] Implement selecting different cuts

---
 src/py4vasp/calculation/_density.py | 40 ++++++++++++++++++++++++-----
 tests/calculation/test_density.py   | 27 +++++++++++++------
 2 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index 4cb77d1d..a124561e 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -225,13 +225,6 @@ def to_view(self, selection=None, supercell=None, **user_options):
         ]
         return viewer
 
-    @_base.data_access
-    def to_contour(self, *, a=None):
-        data = slicing.grid_data(self._raw_data.charge[0].T, "a", a)
-        lattice = slicing.plane(self._structure.lattice_vectors(), "a", normal=None)
-        contour = graph.Contour(data, lattice, "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()
@@ -296,6 +289,22 @@ def _use_symmetric_isosurface(self, component):
             _raise_is_collinear_error()
         return component > 0
 
+    @_base.data_access
+    def to_contour(self, *, a=None, b=None, c=None):
+        cut, fraction = self._get_cut(a, b, c)
+        data = slicing.grid_data(self._raw_data.charge[0].T, cut, fraction)
+        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
+        contour = graph.Contour(data, lattice, "charge")
+        return graph.Graph(contour)
+
+    def _get_cut(self, a, b, c):
+        _raise_error_cut_selection_incorrect(a, b, c)
+        if a is not None:
+            return "a", a
+        if b is not None:
+            return "b", b
+        return "c", c
+
     @_base.data_access
     def is_nonpolarized(self):
         "Returns whether the density is not spin polarized."
@@ -355,3 +364,20 @@ def _raise_error_if_no_data(data):
             'common issue is when you create `Calculation.from_file("vaspout.h5")` '
             "because this will overwrite the default file behavior."
         )
+
+
+def _raise_error_cut_selection_incorrect(*selections):
+    # only a single element may be selected
+    selected_elements = sum(selection is not None for selection in selections)
+    if selected_elements == 0:
+        raise exception.IncorrectUsage(
+            "You have not selected a lattice vector along which the slice should be "
+            "constructed. Please set exactly one of the keyword arguments (a, b, c) "
+            "to a real number that specifies at which fraction of the lattice vector "
+            "the plane is."
+        )
+    if selected_elements > 1:
+        raise exception.IncorrectUsage(
+            "You have selected more than a single element. Please use only one of "
+            "(a, b, c) and not multiple choices."
+        )
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 6a84e71e..8b92a7a9 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -255,20 +255,31 @@ 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]
+@pytest.mark.parametrize(
+    "kwargs, index, position",
+    (({"a": 0.1}, 0, 1), ({"b": 0.7}, 1, 8), ({"c": 1.3}, 2, 4)),
+)
+def test_contour_of_charge(nonpolarized_density, kwargs, index, position, Assert):
+    graph = nonpolarized_density.to_contour(**kwargs)
+    slice_ = [slice(None), slice(None), slice(None)]
+    slice_[index] = position
+    data = nonpolarized_density.ref.output["charge"][tuple(slice_)]
     assert len(graph) == 1
-    Assert.allclose(graph.series.data, slice_)
+    Assert.allclose(graph.series.data, data)
     lattice_vectors = nonpolarized_density.ref.structure.lattice_vectors()
-    expected_products = lattice_vectors[1:] @ lattice_vectors[1:].T
+    lattice_vectors = np.delete(lattice_vectors, index, axis=0)
+    expected_products = lattice_vectors @ lattice_vectors.T
     actual_products = graph.series.lattice @ graph.series.lattice.T
     Assert.allclose(actual_products, expected_products)
 
 
-# TODO: a, b, c
-#       slice missing
-#       a < 0 or a > 1
+def test_incorrect_slice_raises_error(nonpolarized_density):
+    with pytest.raises(exception.IncorrectUsage):
+        nonpolarized_density.to_contour()
+    with pytest.raises(exception.IncorrectUsage):
+        nonpolarized_density.to_contour(a=1, b=2)
+    with pytest.raises(exception.IncorrectUsage):
+        nonpolarized_density.to_contour(3)
 
 
 def test_to_numpy(reference_density, Assert):

From d674603abeb99ec027f1bac5cdaef394b2cba4cf Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 16 Apr 2024 15:33:48 +0200
Subject: [PATCH 25/46] Implement selecting specific contour plot

---
 src/py4vasp/calculation/_density.py | 14 ++++++--
 tests/calculation/test_density.py   | 53 +++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 3 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index a124561e..bfa0ef9a 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -290,11 +290,19 @@ def _use_symmetric_isosurface(self, component):
         return component > 0
 
     @_base.data_access
-    def to_contour(self, *, a=None, b=None, c=None):
+    def to_contour(self, selection=None, *, a=None, b=None, c=None):
         cut, fraction = self._get_cut(a, b, c)
-        data = slicing.grid_data(self._raw_data.charge[0].T, cut, fraction)
+        map_ = self._create_map()
+        selector = index.Selector({0: map_}, self._raw_data.charge)
+        tree = select.Tree.from_selection(selection)
+        for selection in tree.selections():
+            density = selector[selection].T
+            break
+        print(density.shape)
+        data = slicing.grid_data(density, cut, fraction)
         lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
-        contour = graph.Contour(data, lattice, "charge")
+        label = self._label(selector.label(selection)) or "charge"
+        contour = graph.Contour(data, lattice, label)
         return graph.Graph(contour)
 
     def _get_cut(self, a, b, c):
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 8b92a7a9..9fbfc07c 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -271,6 +271,7 @@ def test_contour_of_charge(nonpolarized_density, kwargs, index, position, Assert
     expected_products = lattice_vectors @ lattice_vectors.T
     actual_products = graph.series.lattice @ graph.series.lattice.T
     Assert.allclose(actual_products, expected_products)
+    assert graph.series.label == "charge"
 
 
 def test_incorrect_slice_raises_error(nonpolarized_density):
@@ -282,6 +283,58 @@ def test_incorrect_slice_raises_error(nonpolarized_density):
         nonpolarized_density.to_contour(3)
 
 
+@pytest.mark.parametrize(
+    "selection", ["3", "sigma_z", "z", "sigma_3", "magnetization", "mag", "m"]
+)
+def test_collinear_to_contour(selection, collinear_density, Assert):
+    source = collinear_density.ref.source
+    if source == "charge":
+        expected_label = selection
+        expected_data = collinear_density.ref.output["magnetization"][:, :, 7]
+    else:
+        expected_label = f"{source}({selection})"
+        expected_data = collinear_density.ref.output[source][1, :, :, 7]
+        if selection in ("magnetization", "mag", "m"):
+            # magnetization not allowed for tau
+            return
+    expected_lattice = collinear_density.ref.structure.lattice_vectors()[:2, :2]
+    graph = collinear_density.to_contour(selection, c=-0.5)
+    assert len(graph) == 1
+    Assert.allclose(graph.series.data, expected_data)
+    Assert.allclose(graph.series.lattice, expected_lattice)
+    assert graph.series.label == expected_label
+
+
+# @pytest.mark.parametrize(
+#     "selections",
+#     [
+#         ("1", "2", "3"),
+#         ("sigma_x", "sigma_y", "sigma_z"),
+#         ("x", "y", "z"),
+#         ("sigma_1", "sigma_2", "sigma_3"),
+#         (
+#             "m(1)",
+#             "mag(2)",
+#             "magnetization(3)",
+#         ),  # the magnetization label should be ignored
+#     ],
+# )
+# def test_contour_with_selection(noncollinear_density, selections, Assert):
+#     source = noncollinear_density.ref.source
+#     if source == "charge":
+#         if "(" in selections[0]:  # magnetization filtered from selections
+#             expected_labels = ("1", "2", "3")
+#         else:
+#             expected_labels = selections
+#
+#         expected_density = noncollinear_density.ref.output["magnetization"]
+#     else:
+#         expected_labels = (f"{source}({selection})" for selection in selections)
+#         expected_density = noncollinear_density.ref.output[source][1:]
+#         if "(" in selections[0]:  # magnetization not allowed for tau
+#             return
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":

From 4531074c5ee2bb7e5561ef06114fa0663d685aec Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Wed, 17 Apr 2024 10:19:56 +0200
Subject: [PATCH 26/46] Implement noncollinear contour plots

Allow for multiple selections
---
 src/py4vasp/calculation/_density.py | 18 ++++---
 tests/calculation/test_density.py   | 80 ++++++++++++++++-------------
 2 files changed, 55 insertions(+), 43 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index bfa0ef9a..b25a7afe 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -292,18 +292,22 @@ def _use_symmetric_isosurface(self, component):
     @_base.data_access
     def to_contour(self, selection=None, *, a=None, b=None, c=None):
         cut, fraction = self._get_cut(a, b, c)
+        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
         map_ = self._create_map()
         selector = index.Selector({0: map_}, self._raw_data.charge)
         tree = select.Tree.from_selection(selection)
-        for selection in tree.selections():
-            density = selector[selection].T
-            break
-        print(density.shape)
+        selections = self._filter_noncollinear_magnetization_from_selections(tree)
+        contours = [
+            self._contour(selector, selection, cut, fraction, lattice)
+            for selection in selections
+        ]
+        return graph.Graph(contours)
+
+    def _contour(self, selector, selection, cut, fraction, lattice):
+        density = selector[selection].T
         data = slicing.grid_data(density, cut, fraction)
-        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
         label = self._label(selector.label(selection)) or "charge"
-        contour = graph.Contour(data, lattice, label)
-        return graph.Graph(contour)
+        return graph.Contour(data, lattice, label)
 
     def _get_cut(self, a, b, c):
         _raise_error_cut_selection_incorrect(a, b, c)
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 9fbfc07c..49087562 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -265,13 +265,14 @@ def test_contour_of_charge(nonpolarized_density, kwargs, index, position, Assert
     slice_[index] = position
     data = nonpolarized_density.ref.output["charge"][tuple(slice_)]
     assert len(graph) == 1
-    Assert.allclose(graph.series.data, data)
+    series = graph.series[0]
+    Assert.allclose(series.data, data)
     lattice_vectors = nonpolarized_density.ref.structure.lattice_vectors()
     lattice_vectors = np.delete(lattice_vectors, index, axis=0)
     expected_products = lattice_vectors @ lattice_vectors.T
-    actual_products = graph.series.lattice @ graph.series.lattice.T
+    actual_products = series.lattice @ series.lattice.T
     Assert.allclose(actual_products, expected_products)
-    assert graph.series.label == "charge"
+    assert series.label == "charge"
 
 
 def test_incorrect_slice_raises_error(nonpolarized_density):
@@ -300,39 +301,46 @@ def test_collinear_to_contour(selection, collinear_density, Assert):
     expected_lattice = collinear_density.ref.structure.lattice_vectors()[:2, :2]
     graph = collinear_density.to_contour(selection, c=-0.5)
     assert len(graph) == 1
-    Assert.allclose(graph.series.data, expected_data)
-    Assert.allclose(graph.series.lattice, expected_lattice)
-    assert graph.series.label == expected_label
-
-
-# @pytest.mark.parametrize(
-#     "selections",
-#     [
-#         ("1", "2", "3"),
-#         ("sigma_x", "sigma_y", "sigma_z"),
-#         ("x", "y", "z"),
-#         ("sigma_1", "sigma_2", "sigma_3"),
-#         (
-#             "m(1)",
-#             "mag(2)",
-#             "magnetization(3)",
-#         ),  # the magnetization label should be ignored
-#     ],
-# )
-# def test_contour_with_selection(noncollinear_density, selections, Assert):
-#     source = noncollinear_density.ref.source
-#     if source == "charge":
-#         if "(" in selections[0]:  # magnetization filtered from selections
-#             expected_labels = ("1", "2", "3")
-#         else:
-#             expected_labels = selections
-#
-#         expected_density = noncollinear_density.ref.output["magnetization"]
-#     else:
-#         expected_labels = (f"{source}({selection})" for selection in selections)
-#         expected_density = noncollinear_density.ref.output[source][1:]
-#         if "(" in selections[0]:  # magnetization not allowed for tau
-#             return
+    series = graph.series[0]
+    Assert.allclose(series.data, expected_data)
+    Assert.allclose(series.lattice, expected_lattice)
+    assert series.label == expected_label
+
+
+@pytest.mark.parametrize(
+    "selections",
+    [
+        ("1", "2", "3"),
+        ("sigma_x", "sigma_y", "sigma_z"),
+        ("x", "y", "z"),
+        ("sigma_1", "sigma_2", "sigma_3"),
+        (
+            "m(1)",
+            "mag(2)",
+            "magnetization(3)",
+        ),  # the magnetization label should be ignored
+    ],
+)
+def test_noncollinear_to_contour(noncollinear_density, selections, Assert):
+    source = noncollinear_density.ref.source
+    if source == "charge":
+        if "(" in selections[0]:  # magnetization filtered from selections
+            expected_labels = ("1", "2", "3")
+        else:
+            expected_labels = selections
+        expected_data = noncollinear_density.ref.output["magnetization"][:, :, 5, :]
+    else:
+        expected_labels = (f"{source}({selection})" for selection in selections)
+        expected_data = noncollinear_density.ref.output[source][1:, :, 5, :]
+        if "(" in selections[0]:  # magnetization not allowed for tau
+            return
+    graph = noncollinear_density.to_contour(" ".join(selections), b=0.4)
+    expected_lattice = noncollinear_density.ref.structure.lattice_vectors()[::2, ::2]
+    assert len(graph) == len(expected_data)
+    for density, label, series in zip(expected_data, expected_labels, graph.series):
+        Assert.allclose(series.data, density)
+        Assert.allclose(series.lattice, expected_lattice)
+        assert series.label == label
 
 
 def test_to_numpy(reference_density, Assert):

From 94914a7c35126bdbf66e8ab45400484187140e41 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Wed, 17 Apr 2024 11:15:27 +0200
Subject: [PATCH 27/46] Implement supercell for contour plots

---
 src/py4vasp/calculation/_density.py | 11 +++++++----
 tests/calculation/test_density.py   |  7 +++++++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index b25a7afe..535bcce6 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -290,7 +290,7 @@ def _use_symmetric_isosurface(self, component):
         return component > 0
 
     @_base.data_access
-    def to_contour(self, selection=None, *, a=None, b=None, c=None):
+    def to_contour(self, selection=None, *, a=None, b=None, c=None, supercell=None):
         cut, fraction = self._get_cut(a, b, c)
         lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
         map_ = self._create_map()
@@ -298,16 +298,19 @@ def to_contour(self, selection=None, *, a=None, b=None, c=None):
         tree = select.Tree.from_selection(selection)
         selections = self._filter_noncollinear_magnetization_from_selections(tree)
         contours = [
-            self._contour(selector, selection, cut, fraction, lattice)
+            self._contour(selector, selection, cut, fraction, lattice, supercell)
             for selection in selections
         ]
         return graph.Graph(contours)
 
-    def _contour(self, selector, selection, cut, fraction, lattice):
+    def _contour(self, selector, selection, cut, fraction, lattice, supercell):
         density = selector[selection].T
         data = slicing.grid_data(density, cut, fraction)
         label = self._label(selector.label(selection)) or "charge"
-        return graph.Contour(data, lattice, label)
+        contour = graph.Contour(data, lattice, label)
+        if supercell is not None:
+            contour.supercell = np.ones(2, dtype=np.int_) * supercell
+        return contour
 
     def _get_cut(self, a, b, c):
         _raise_error_cut_selection_incorrect(a, b, c)
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 49087562..95b593fb 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -343,6 +343,13 @@ def test_noncollinear_to_contour(noncollinear_density, selections, Assert):
         assert series.label == label
 
 
+def test_to_contour_supercell(nonpolarized_density, Assert):
+    graph = nonpolarized_density.to_contour(a=0, supercell=2)
+    Assert.allclose(graph.series[0].supercell, (2, 2))
+    graph = nonpolarized_density.to_contour(a=0, supercell=(2, 1))
+    Assert.allclose(graph.series[0].supercell, (2, 1))
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":

From 8a3a728ca05122cfb0e1c982a51b1ec34f0b30db Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Wed, 17 Apr 2024 11:40:17 +0200
Subject: [PATCH 28/46] Allow passing normal vector

---
 src/py4vasp/calculation/_density.py |  6 ++++--
 tests/calculation/test_density.py   | 24 ++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index 535bcce6..249a7c8c 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -290,9 +290,11 @@ def _use_symmetric_isosurface(self, component):
         return component > 0
 
     @_base.data_access
-    def to_contour(self, selection=None, *, a=None, b=None, c=None, supercell=None):
+    def to_contour(
+        self, selection=None, *, a=None, b=None, c=None, supercell=None, normal=None
+    ):
         cut, fraction = self._get_cut(a, b, c)
-        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
+        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal)
         map_ = self._create_map()
         selector = index.Selector({0: map_}, self._raw_data.charge)
         tree = select.Tree.from_selection(selection)
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 95b593fb..9af1ff00 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -350,6 +350,30 @@ def test_to_contour_supercell(nonpolarized_density, Assert):
     Assert.allclose(graph.series[0].supercell, (2, 1))
 
 
+def test_raise_error_if_normal_not_reasonable(nonpolarized_density):
+    with pytest.raises(exception.IncorrectUsage):
+        # cell is not close to cartesian vectors
+        nonpolarized_density.to_contour(a=0, normal="auto")
+    with pytest.raises(exception.IncorrectUsage):
+        nonpolarized_density.to_contour(a=0, normal="unknown choice")
+
+
+@pytest.mark.parametrize(
+    "normal, rotation",
+    [
+        ("auto", np.eye(2)),
+        ("x", np.array([[0, -1], [1, 0]])),
+        ("y", np.diag((1, -1))),
+        ("z", np.eye(2)),
+    ],
+)
+def test_to_contour_normal_vector(collinear_density, normal, rotation, Assert):
+    graph = collinear_density.to_contour(c=0.5, normal=normal)
+    lattice_vectors = collinear_density.ref.structure.lattice_vectors()
+    expected_lattice = lattice_vectors[:2, :2] @ rotation
+    Assert.allclose(graph.series[0].lattice, expected_lattice)
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":

From d83b0b738ca217af7650abf46522743a0b6d0690 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Wed, 17 Apr 2024 12:10:08 +0200
Subject: [PATCH 29/46] Use contour instead of heatmap for density

---
 src/py4vasp/_third_party/graph/contour.py | 24 +++++++++++++++++++----
 src/py4vasp/calculation/_density.py       |  2 +-
 tests/calculation/test_density.py         |  1 +
 tests/third_party/graph/test_graph.py     | 14 +++++++------
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index 7532af5d..653d45c1 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -38,6 +38,8 @@ class Contour(trace.Trace):
     have two components, so remove any element normal to the plane."""
     label: str
     "Assign a label to the visualization that may be used to identify one among multiple plots."
+    isolevels: bool = False
+    "Defines whether isolevels should be added or a heatmap is used."
     supercell: np.array = (1, 1)
     "Multiple of each lattice vector to be drawn."
     show_cell: bool = True
@@ -47,21 +49,35 @@ def to_plotly(self):
         lattice_supercell = np.diag(self.supercell) @ self.lattice
         # swap a and b axes because that is the way plotly expects the data
         data = np.tile(self.data, self.supercell).T
-        if self._is_heatmap():
+        if self._is_contour():
+            yield self._make_contour(lattice_supercell, data)
+        elif self._is_heatmap():
             yield self._make_heatmap(lattice_supercell, data)
         else:
             yield self._make_quiver(lattice_supercell, data)
 
+    def _is_contour(self):
+        return self.data.ndim == 2 and self.isolevels
+
     def _is_heatmap(self):
-        return self.data.ndim == 2
+        return self.data.ndim == 2 and not self.isolevels
+
+    def _make_contour(self, lattice, data):
+        x, y, z = self._interpolate_data_if_necessary(lattice, data)
+        contour = go.Contour(x=x, y=y, z=z, name=self.label, autocontour=True)
+        return contour, self._options()
 
     def _make_heatmap(self, lattice, data):
+        x, y, z = self._interpolate_data_if_necessary(lattice, data)
+        heatmap = go.Heatmap(x=x, y=y, z=z, name=self.label, colorscale="turbid_r")
+        return heatmap, self._options()
+
+    def _interpolate_data_if_necessary(self, lattice, data):
         if self._interpolation_required():
             x, y, z = self._interpolate_data(lattice, data)
         else:
             x, y, z = self._use_data_without_interpolation(lattice, data)
-        heatmap = go.Heatmap(x=x, y=y, z=z, name=self.label, colorscale="turbid_r")
-        return heatmap, self._options()
+        return x, y, z
 
     def _interpolation_required(self):
         return not np.allclose((self.lattice[1, 0], self.lattice[0, 1]), 0)
diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index 249a7c8c..57f5bd47 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -309,7 +309,7 @@ def _contour(self, selector, selection, cut, fraction, lattice, supercell):
         density = selector[selection].T
         data = slicing.grid_data(density, cut, fraction)
         label = self._label(selector.label(selection)) or "charge"
-        contour = graph.Contour(data, lattice, label)
+        contour = graph.Contour(data, lattice, label, isolevels=True)
         if supercell is not None:
             contour.supercell = np.ones(2, dtype=np.int_) * supercell
         return contour
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 9af1ff00..2019c099 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -305,6 +305,7 @@ def test_collinear_to_contour(selection, collinear_density, Assert):
     Assert.allclose(series.data, expected_data)
     Assert.allclose(series.lattice, expected_lattice)
     assert series.label == expected_label
+    assert series.isolevels
 
 
 @pytest.mark.parametrize(
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index 5c0ba892..b2bd9753 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -68,6 +68,7 @@ def rectangle_contour():
         data=np.linspace(0, 10, 20 * 18).reshape((20, 18)),
         lattice=np.diag([4.0, 3.6]),
         label="rectangle contour",
+        isolevels=True,
     )
 
 
@@ -435,11 +436,7 @@ def test_contour(rectangle_contour, Assert, not_core):
     Assert.allclose(fig.data[0].x, np.linspace(0, 4, 20, endpoint=False) + 0.1)
     Assert.allclose(fig.data[0].y, np.linspace(0, 3.6, 18, endpoint=False) + 0.1)
     assert fig.data[0].name == rectangle_contour.label
-    expected_colorscale = px.colors.get_colorscale("turbid_r")
-    assert len(fig.data[0].colorscale) == len(expected_colorscale)
-    for actual, expected in zip(fig.data[0].colorscale, expected_colorscale):
-        Assert.allclose(actual[0], expected[0])
-        assert actual[1] == expected[1]
+    assert fig.data[0].autocontour
     # text explicitly that it is False to prevent None passing the test
     assert fig.layout.xaxis.visible == False
     assert fig.layout.yaxis.visible == False
@@ -468,7 +465,7 @@ def check_unit_cell(unit_cell, x, y, zero):
     assert unit_cell.line.color == _config.VASP_GRAY
 
 
-def test_contour_interpolate(tilted_contour, not_core):
+def test_contour_interpolate(tilted_contour, Assert, not_core):
     graph = Graph(tilted_contour)
     fig = graph.to_plotly()
     area_cell = 12.0
@@ -485,6 +482,11 @@ def test_contour_interpolate(tilted_contour, not_core):
     finite = np.isfinite(fig.data[0].z)
     assert np.isclose(np.average(fig.data[0].z[finite]), expected_average)
     assert len(fig.layout.shapes) == 0
+    expected_colorscale = px.colors.get_colorscale("turbid_r")
+    assert len(fig.data[0].colorscale) == len(expected_colorscale)
+    for actual, expected in zip(fig.data[0].colorscale, expected_colorscale):
+        Assert.allclose(actual[0], expected[0])
+        assert actual[1] == expected[1]
 
 
 def test_mix_contour_and_series(two_lines, rectangle_contour, not_core):

From 5b14f250d3705c4494efb61d485d4cc3a68862df Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Thu, 18 Apr 2024 10:03:42 +0200
Subject: [PATCH 30/46] Apply VASP color scheme

---
 src/py4vasp/_third_party/graph/__init__.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/__init__.py b/src/py4vasp/_third_party/graph/__init__.py
index eb2a084a..5db9f5ad 100644
--- a/src/py4vasp/_third_party/graph/__init__.py
+++ b/src/py4vasp/_third_party/graph/__init__.py
@@ -1,6 +1,8 @@
 # Copyright © VASP Software GmbH,
 # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
-from py4vasp._config import VASP_COLORS
+import copy
+
+from py4vasp._config import VASP_BLUE, VASP_COLORS, VASP_GRAY, VASP_RED
 from py4vasp._util import import_
 
 from .contour import Contour
@@ -14,7 +16,10 @@
 
 if import_.is_imported(go) and import_.is_imported(pio):
     axis_format = {"showexponent": "all", "exponentformat": "power"}
+    contour = copy.copy(pio.templates["ggplot2"].data.contour[0])
+    contour.colorscale = [[0, VASP_BLUE], [0.5, VASP_GRAY], [1, VASP_RED]]
+    data = {"contour": (contour,)}
     layout = {"colorway": VASP_COLORS, "xaxis": axis_format, "yaxis": axis_format}
-    pio.templates["vasp"] = go.layout.Template(layout=layout)
+    pio.templates["vasp"] = go.layout.Template(data=data, layout=layout)
     pio.templates["ggplot2"].layout.shapedefaults = {}
     pio.templates.default = "ggplot2+vasp"

From c4435200ea58cf8d6068612bcbee60f0684e502d Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Thu, 18 Apr 2024 10:03:58 +0200
Subject: [PATCH 31/46] Add test for subtraction in contour plot

---
 tests/calculation/test_density.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 2019c099..d7d96364 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -375,6 +375,17 @@ def test_to_contour_normal_vector(collinear_density, normal, rotation, Assert):
     Assert.allclose(graph.series[0].lattice, expected_lattice)
 
 
+def test_to_contour_operation(collinear_density, Assert):
+    graph = collinear_density.to_contour("0 - 3", c=0)
+    source = collinear_density.ref.source
+    output = collinear_density.ref.output
+    if source == "charge":
+        expected_data = output["charge"][:, :, 0] - output["magnetization"][:, :, 0]
+    else:
+        expected_data = output[source][0, :, :, 0] - output[source][1, :, :, 0]
+    Assert.allclose(graph.series[0].data, expected_data)
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":

From 46d8805499c9b06eee3812d2f0943c6b044b53df Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Thu, 18 Apr 2024 15:17:10 +0200
Subject: [PATCH 32/46] Add annotations to identify lattice vectors

---
 src/py4vasp/_third_party/graph/contour.py | 93 ++++++++++++++++-------
 src/py4vasp/_third_party/graph/graph.py   |  2 +
 src/py4vasp/_util/slicing.py              | 17 ++++-
 tests/calculation/test_density.py         |  8 +-
 tests/third_party/graph/test_graph.py     | 36 ++++++---
 tests/util/test_slicing.py                | 17 +++--
 6 files changed, 125 insertions(+), 48 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index 653d45c1..17301340 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -1,5 +1,7 @@
 # Copyright © VASP Software GmbH,
 # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+from __future__ import annotations
+
 import dataclasses
 import itertools
 
@@ -24,16 +26,18 @@ class Contour(trace.Trace):
     each data point should be a 2d vector within the plane.
     """
 
-    interpolation_factor = 2
+    _interpolation_factor = 2
     """If the lattice does not align with the cartesian axes, the data is interpolated
     to by approximately this factor along each line."""
+    _shift_label_pixels = 10
+    "Shift the labels by this many pixels to avoid overlap."
 
     data: np.array
     """2d or 3d grid data in the plane spanned by the lattice vectors. If the data is
     the dimensions should be the ones of the grid, if the data is 3d the first dimension
     should be a 2 for a vector in the plane of the grid and the other two dimensions
     should be the grid."""
-    lattice: np.array
+    lattice: Lattice
     """2 vectors spanning the plane in which the data is represented. Each vector should
     have two components, so remove any element normal to the plane."""
     label: str
@@ -46,15 +50,15 @@ class Contour(trace.Trace):
     "Show the unit cell in the resulting visualization."
 
     def to_plotly(self):
-        lattice_supercell = np.diag(self.supercell) @ self.lattice
+        lattice_supercell = np.diag(self.supercell) @ self.lattice.vectors
         # swap a and b axes because that is the way plotly expects the data
         data = np.tile(self.data, self.supercell).T
         if self._is_contour():
-            yield self._make_contour(lattice_supercell, data)
+            yield self._make_contour(lattice_supercell, data), self._options()
         elif self._is_heatmap():
-            yield self._make_heatmap(lattice_supercell, data)
+            yield self._make_heatmap(lattice_supercell, data), self._options()
         else:
-            yield self._make_quiver(lattice_supercell, data)
+            yield self._make_quiver(lattice_supercell, data), self._options()
 
     def _is_contour(self):
         return self.data.ndim == 2 and self.isolevels
@@ -64,13 +68,23 @@ def _is_heatmap(self):
 
     def _make_contour(self, lattice, data):
         x, y, z = self._interpolate_data_if_necessary(lattice, data)
-        contour = go.Contour(x=x, y=y, z=z, name=self.label, autocontour=True)
-        return contour, self._options()
+        return go.Contour(x=x, y=y, z=z, name=self.label, autocontour=True)
 
     def _make_heatmap(self, lattice, data):
         x, y, z = self._interpolate_data_if_necessary(lattice, data)
-        heatmap = go.Heatmap(x=x, y=y, z=z, name=self.label, colorscale="turbid_r")
-        return heatmap, self._options()
+        return go.Heatmap(x=x, y=y, z=z, name=self.label, colorscale="turbid_r")
+
+    def _make_quiver(self, lattice, data):
+        u = data[:, :, 0].flatten()
+        v = data[:, :, 1].flatten()
+        meshes = [
+            np.linspace(np.zeros(2), vector, num_points, endpoint=False)
+            for vector, num_points in zip(reversed(lattice), data.shape)
+            # remember that b and a axis are swapped
+        ]
+        x, y = np.array([sum(points) for points in itertools.product(*meshes)]).T
+        fig = ff.create_quiver(x, y, u, v, scale=1)
+        return fig.data[0]
 
     def _interpolate_data_if_necessary(self, lattice, data):
         if self._interpolation_required():
@@ -80,12 +94,14 @@ def _interpolate_data_if_necessary(self, lattice, data):
         return x, y, z
 
     def _interpolation_required(self):
-        return not np.allclose((self.lattice[1, 0], self.lattice[0, 1]), 0)
+        y_position_first_vector = self.lattice.vectors[0, 1]
+        x_position_second_vector = self.lattice.vectors[1, 0]
+        return not np.allclose((y_position_first_vector, x_position_second_vector), 0)
 
     def _interpolate_data(self, lattice, data):
         area_cell = abs(np.cross(lattice[0], lattice[1]))
         points_per_area = data.size / area_cell
-        points_per_line = np.sqrt(points_per_area) * self.interpolation_factor
+        points_per_line = np.sqrt(points_per_area) * self._interpolation_factor
         lengths = np.sum(np.abs(lattice), axis=0)
         shape = np.ceil(points_per_line * lengths).astype(int)
         line_mesh_a = self._make_mesh(lattice, data.shape[1], 0)
@@ -114,23 +130,48 @@ def _make_mesh(self, lattice, num_point, index):
         )
 
     def _options(self):
+        return {
+            "shapes": self._create_unit_cell(),
+            "annotations": self._label_unit_cell_vectors(),
+        }
+
+    def _create_unit_cell(self):
         if not self.show_cell:
-            return {}
+            return ()
         pos_to_str = lambda pos: f"{pos[0]} {pos[1]}"
-        corners = (self.lattice[0], self.lattice[0] + self.lattice[1], self.lattice[1])
+        vectors = self.lattice.vectors
+        corners = (vectors[0], vectors[0] + vectors[1], vectors[1])
         to_corners = (f"L {pos_to_str(corner)}" for corner in corners)
         path = f"M 0 0 {' '.join(to_corners)} Z"
         unit_cell = {"type": "path", "line": {"color": _config.VASP_GRAY}, "path": path}
-        return {"shapes": [unit_cell]}
-
-    def _make_quiver(self, lattice, data):
-        u = data[:, :, 0].flatten()
-        v = data[:, :, 1].flatten()
-        meshes = [
-            np.linspace(np.zeros(2), vector, num_points, endpoint=False)
-            for vector, num_points in zip(reversed(lattice), data.shape)
-            # remember that b and a axis are swapped
+        return (unit_cell,)
+
+    def _label_unit_cell_vectors(self):
+        if self.lattice.labels is None:
+            return []
+        vectors = self.lattice.vectors
+        return [
+            {
+                "text": label,
+                "showarrow": False,
+                "x": vectors[i, 0],
+                "y": vectors[i, 1],
+                **self._shift_label(vectors[i], vectors[1 - i]),
+            }
+            for i, label in enumerate(self.lattice.labels)
         ]
-        x, y = np.array([sum(points) for points in itertools.product(*meshes)]).T
-        fig = ff.create_quiver(x, y, u, v, scale=1)
-        return fig.data[0], self._options()
+
+    def _shift_label(self, current_vector, other_vector):
+        invert = np.cross(current_vector, other_vector) < 0
+        norm = np.linalg.norm(current_vector)
+        shifts = self._shift_label_pixels * current_vector[::-1] / norm
+        if invert:
+            return {
+                "xshift": -shifts[0],
+                "yshift": shifts[1],
+            }
+        else:
+            return {
+                "xshift": shifts[0],
+                "yshift": -shifts[1],
+            }
diff --git a/src/py4vasp/_third_party/graph/graph.py b/src/py4vasp/_third_party/graph/graph.py
index 84ec2f3e..8d44152c 100644
--- a/src/py4vasp/_third_party/graph/graph.py
+++ b/src/py4vasp/_third_party/graph/graph.py
@@ -79,6 +79,8 @@ def to_plotly(self):
                 figure.add_trace(trace, row=options["row"], col=1)
             for shape in options.get("shapes", ()):
                 figure.add_shape(**shape)
+            for annotation in options.get("annotations", ()):
+                figure.add_annotation(**annotation)
         return figure
 
     def show(self):
diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index cfe0db30..4f66cc12 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -1,5 +1,7 @@
 # Copyright © VASP Software GmbH,
 # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
+import dataclasses
+
 import numpy as np
 
 from py4vasp import exception
@@ -8,6 +10,16 @@
 AXIS = ("x", "y", "z")
 
 
+@dataclasses.dataclass
+class Lattice:
+    "Stores the 2d lattice vectors and their labels"
+    vectors: np.array
+    "Vectors spanning the plane."
+    labels: tuple = None
+    """Labels that could be added to the lattice vectors e.g. in a plot. If set to None
+    no labels are added."""
+
+
 def grid_data(data, cut, fraction):
     """Takes a 2d slice of a 3d grid data.
 
@@ -80,10 +92,11 @@ def plane(cell, cut, normal="auto"):
     """
     _raise_error_if_cut_unknown(cut)
     vectors = np.delete(cell, INDICES[cut], axis=0)
+    labels = tuple("abc".replace(cut, ""))
     if normal is not None:
-        return _rotate_normal_to_cartesian_axis(vectors, normal)
+        return Lattice(_rotate_normal_to_cartesian_axis(vectors, normal), labels)
     else:
-        return _rotate_first_vector_to_x_axis(vectors)
+        return Lattice(_rotate_first_vector_to_x_axis(vectors), labels)
 
 
 def _rotate_first_vector_to_x_axis(vectors):
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index d7d96364..fe65d82d 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -270,7 +270,7 @@ def test_contour_of_charge(nonpolarized_density, kwargs, index, position, Assert
     lattice_vectors = nonpolarized_density.ref.structure.lattice_vectors()
     lattice_vectors = np.delete(lattice_vectors, index, axis=0)
     expected_products = lattice_vectors @ lattice_vectors.T
-    actual_products = series.lattice @ series.lattice.T
+    actual_products = series.lattice.vectors @ series.lattice.vectors.T
     Assert.allclose(actual_products, expected_products)
     assert series.label == "charge"
 
@@ -303,7 +303,7 @@ def test_collinear_to_contour(selection, collinear_density, Assert):
     assert len(graph) == 1
     series = graph.series[0]
     Assert.allclose(series.data, expected_data)
-    Assert.allclose(series.lattice, expected_lattice)
+    Assert.allclose(series.lattice.vectors, expected_lattice)
     assert series.label == expected_label
     assert series.isolevels
 
@@ -340,7 +340,7 @@ def test_noncollinear_to_contour(noncollinear_density, selections, Assert):
     assert len(graph) == len(expected_data)
     for density, label, series in zip(expected_data, expected_labels, graph.series):
         Assert.allclose(series.data, density)
-        Assert.allclose(series.lattice, expected_lattice)
+        Assert.allclose(series.lattice.vectors, expected_lattice)
         assert series.label == label
 
 
@@ -372,7 +372,7 @@ def test_to_contour_normal_vector(collinear_density, normal, rotation, Assert):
     graph = collinear_density.to_contour(c=0.5, normal=normal)
     lattice_vectors = collinear_density.ref.structure.lattice_vectors()
     expected_lattice = lattice_vectors[:2, :2] @ rotation
-    Assert.allclose(graph.series[0].lattice, expected_lattice)
+    Assert.allclose(graph.series[0].lattice.vectors, expected_lattice)
 
 
 def test_to_contour_operation(collinear_density, Assert):
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index b2bd9753..0d12dae3 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -8,7 +8,7 @@
 
 from py4vasp import _config, exception
 from py4vasp._third_party.graph import Contour, Graph, Series
-from py4vasp._util import import_
+from py4vasp._util import import_, slicing
 
 px = import_.optional("plotly.express")
 
@@ -66,7 +66,7 @@ def subplot():
 def rectangle_contour():
     return Contour(
         data=np.linspace(0, 10, 20 * 18).reshape((20, 18)),
-        lattice=np.diag([4.0, 3.6]),
+        lattice=slicing.Lattice(np.diag([4.0, 3.6]), ("a", "b")),
         label="rectangle contour",
         isolevels=True,
     )
@@ -76,7 +76,7 @@ def rectangle_contour():
 def tilted_contour():
     return Contour(
         data=np.linspace(0, 5, 16 * 20).reshape((16, 20)),
-        lattice=np.array([[2, 3], [2, -3]]),
+        lattice=slicing.Lattice(np.array([[2, 3], [2, -3]]), ("a", "c")),
         label="tilted contour",
         supercell=(2, 1),
         show_cell=False,
@@ -87,7 +87,7 @@ def tilted_contour():
 def simple_quiver():
     return Contour(
         data=np.array([[(y, x) for x in range(3)] for y in range(5)]).T,
-        lattice=np.diag((3, 5)),
+        lattice=slicing.Lattice(np.diag((3, 5)), ("b", "c")),
         label="quiver plot",
     )
 
@@ -96,7 +96,7 @@ def simple_quiver():
 def complex_quiver():
     return Contour(
         data=np.linspace(-3, 3, 2 * 12 * 10).reshape((2, 12, 10)),
-        lattice=np.array([[3, 2], [-3, 2]]),
+        lattice=slicing.Lattice([[3, 2], [-3, 2]], None),
         label="quiver plot",
         supercell=(3, 2),
     )
@@ -442,10 +442,11 @@ def test_contour(rectangle_contour, Assert, not_core):
     assert fig.layout.yaxis.visible == False
     assert len(fig.layout.shapes) == 1
     check_unit_cell(fig.layout.shapes[0], x="4.0", y="3.6", zero="0.0")
+    check_annotations(rectangle_contour.lattice, fig.layout.annotations, Assert)
     assert fig.layout.yaxis.scaleanchor == "x"
 
 
-def test_contour_supercell(rectangle_contour, not_core):
+def test_contour_supercell(rectangle_contour, Assert, not_core):
     supercell = np.asarray((3, 5))
     rectangle_contour.supercell = supercell
     graph = Graph(rectangle_contour)
@@ -457,6 +458,7 @@ def test_contour_supercell(rectangle_contour, not_core):
     assert len(fig.data[0].y) == 90
     assert len(fig.layout.shapes) == 1
     check_unit_cell(fig.layout.shapes[0], x="4.0", y="3.6", zero="0.0")
+    check_annotations(rectangle_contour.lattice, fig.layout.annotations, Assert)
 
 
 def check_unit_cell(unit_cell, x, y, zero):
@@ -465,12 +467,24 @@ def check_unit_cell(unit_cell, x, y, zero):
     assert unit_cell.line.color == _config.VASP_GRAY
 
 
+def check_annotations(lattice, annotations, Assert):
+    assert len(lattice.labels) == len(annotations)
+    sign = np.sign(np.cross(*lattice.vectors))
+    for vector, label, annotation in zip(lattice.vectors, lattice.labels, annotations):
+        assert annotation.showarrow == False
+        assert annotation.text == label
+        Assert.allclose((annotation.x, annotation.y), vector)
+        expected_shift = sign * 10 * vector[::-1] / np.linalg.norm(vector)
+        Assert.allclose((annotation.xshift, -annotation.yshift), expected_shift)
+        sign *= -1
+
+
 def test_contour_interpolate(tilted_contour, Assert, not_core):
     graph = Graph(tilted_contour)
     fig = graph.to_plotly()
     area_cell = 12.0
     points_per_area = tilted_contour.data.size / area_cell
-    points_per_line = np.sqrt(points_per_area) * tilted_contour.interpolation_factor
+    points_per_line = np.sqrt(points_per_area) * tilted_contour._interpolation_factor
     lengths = np.array([6, 9])  # this accounts for the 2 x 1 supercell
     expected_shape = np.ceil(points_per_line * lengths).astype(int)
     expected_average = np.average(tilted_contour.data)
@@ -487,6 +501,7 @@ def test_contour_interpolate(tilted_contour, Assert, not_core):
     for actual, expected in zip(fig.data[0].colorscale, expected_colorscale):
         Assert.allclose(actual[0], expected[0])
         assert actual[1] == expected[1]
+    check_annotations(tilted_contour.lattice, fig.layout.annotations, Assert)
 
 
 def test_mix_contour_and_series(two_lines, rectangle_contour, not_core):
@@ -510,6 +525,7 @@ def test_simple_quiver(simple_quiver, Assert, not_core):
         Assert.allclose(y, u)
     assert len(fig.layout.shapes) == 1
     check_unit_cell(fig.layout.shapes[0], x="3", y="5", zero="0")
+    check_annotations(simple_quiver.lattice, fig.layout.annotations, Assert)
     assert fig.layout.yaxis.scaleanchor == "x"
 
 
@@ -517,9 +533,10 @@ def test_complex_quiver(complex_quiver, Assert, not_core):
     graph = Graph(complex_quiver)
     fig = graph.to_plotly()
     data_size = np.prod(complex_quiver.supercell) * complex_quiver.data.size // 2
-    step_a = complex_quiver.lattice[0] / complex_quiver.data.shape[1]
+    vectors = np.array(complex_quiver.lattice.vectors)
+    step_a = vectors[0] / complex_quiver.data.shape[1]
     mesh_a = np.arange(complex_quiver.supercell[0] * complex_quiver.data.shape[1])
-    step_b = complex_quiver.lattice[1] / complex_quiver.data.shape[2]
+    step_b = vectors[1] / complex_quiver.data.shape[2]
     mesh_b = np.arange(complex_quiver.supercell[1] * complex_quiver.data.shape[2])
     expected_positions = np.array(
         [a * step_a + b * step_b for b in mesh_b for a in mesh_a]
@@ -534,6 +551,7 @@ def test_complex_quiver(complex_quiver, Assert, not_core):
     Assert.allclose(actual.positions, expected_positions, tolerance=10)
     Assert.allclose(actual.tips, expected_tips, tolerance=10)
     Assert.allclose(actual.barb_length, expected_barb_length)
+    assert len(fig.layout.annotations) == 0
 
 
 @dataclasses.dataclass
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 4870803e..f3f93408 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -12,7 +12,8 @@ def test_orthorhombic(cut, index, Assert):
     cell = np.diag((3, 4, 5))
     expected_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1)
     actual_plane = slicing.plane(cell, cut)
-    Assert.allclose(actual_plane, expected_plane)
+    Assert.allclose(actual_plane.vectors, expected_plane)
+    assert all(actual_plane.labels == np.delete(("a", "b", "c"), index))
 
 
 @pytest.mark.parametrize("cut, indices", [("a", (0, 2)), ("b", (1, 0)), ("c", (2, 1))])
@@ -20,7 +21,8 @@ def test_unusual_orthorhombic(cut, indices, Assert):
     cell = np.roll(np.diag((3, 4, 5)), shift=1, axis=0)
     expected_plane = np.delete(np.delete(cell, indices[0], axis=0), indices[1], axis=1)
     actual_plane = slicing.plane(cell, cut)
-    Assert.allclose(actual_plane, expected_plane)
+    Assert.allclose(actual_plane.vectors, expected_plane)
+    assert all(actual_plane.labels == np.delete(("a", "b", "c"), indices[0]))
 
 
 @pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)])
@@ -30,7 +32,7 @@ def test_nearly_orthorhombic(cut, index, Assert):
     # the expected plane is in 3d and the actual plane in 2d should have same angles
     # and lengths
     approximate_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1)
-    actual_plane = slicing.plane(cell, cut)
+    actual_plane = slicing.plane(cell, cut).vectors
     Assert.allclose(actual_plane @ actual_plane.T, expected_plane @ expected_plane.T)
     assert np.max(np.abs(approximate_plane - actual_plane)) < 0.1
 
@@ -49,7 +51,7 @@ def test_nearly_orthorhombic(cut, index, Assert):
 def test_rotate_to_user_defined_axis(normal, expected_plane, Assert):
     cell = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
     actual_plane = slicing.plane(cell, "a", normal)
-    Assert.allclose(actual_plane, expected_plane)
+    Assert.allclose(actual_plane.vectors, expected_plane)
 
 
 @pytest.mark.parametrize(
@@ -60,7 +62,8 @@ def test_normal_with_orthonormal_cell(normal, rotation, Assert):
     cell = np.diag((3, 4, 5))
     expected_plane = np.dot(np.delete(np.delete(cell, 0, axis=0), 0, axis=1), rotation)
     actual_plane = slicing.plane(cell, "a", normal)
-    Assert.allclose(actual_plane, expected_plane)
+    Assert.allclose(actual_plane.vectors, expected_plane)
+    assert actual_plane.labels == ("b", "c")
 
 
 @pytest.mark.parametrize("cut, diagonal", [("a", [4, 5]), ("b", [3, 5]), ("c", [3, 4])])
@@ -68,7 +71,7 @@ def test_no_rotation_orthorhombic_cell(cut, diagonal, Assert):
     cell = np.diag((3, 4, 5))
     expected_plane = np.diag(diagonal)
     actual_plane = slicing.plane(cell, cut, normal=None)
-    Assert.allclose(actual_plane, expected_plane)
+    Assert.allclose(actual_plane.vectors, expected_plane)
 
 
 @pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)])
@@ -76,7 +79,7 @@ def test_no_rotation_nontrivial_cell(cut, index, Assert):
     cell = np.array([[0, 1, 1.1], [0.9, 0, 1.1], [0.9, 1, 0]])
     vectors = np.delete(cell, index, axis=0)
     expected_products = vectors @ vectors.T
-    actual_plane = slicing.plane(cell, cut, normal=None)
+    actual_plane = slicing.plane(cell, cut, normal=None).vectors
     Assert.allclose(actual_plane[0, 1], 0)
     actual_products = actual_plane @ actual_plane.T
     Assert.allclose(actual_products, expected_products)

From 2ff3f025fc35810107f40d23562a0ff7e5d0b127 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Thu, 18 Apr 2024 15:19:23 +0200
Subject: [PATCH 33/46] Show label at middle instead of at end

---
 src/py4vasp/_third_party/graph/contour.py | 4 ++--
 tests/third_party/graph/test_graph.py     | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index 17301340..4766a78d 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -154,8 +154,8 @@ def _label_unit_cell_vectors(self):
             {
                 "text": label,
                 "showarrow": False,
-                "x": vectors[i, 0],
-                "y": vectors[i, 1],
+                "x": 0.5 * vectors[i, 0],
+                "y": 0.5 * vectors[i, 1],
                 **self._shift_label(vectors[i], vectors[1 - i]),
             }
             for i, label in enumerate(self.lattice.labels)
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index 0d12dae3..1113d5e7 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -473,7 +473,7 @@ def check_annotations(lattice, annotations, Assert):
     for vector, label, annotation in zip(lattice.vectors, lattice.labels, annotations):
         assert annotation.showarrow == False
         assert annotation.text == label
-        Assert.allclose((annotation.x, annotation.y), vector)
+        Assert.allclose((annotation.x, annotation.y), 0.5 * vector)
         expected_shift = sign * 10 * vector[::-1] / np.linalg.norm(vector)
         Assert.allclose((annotation.xshift, -annotation.yshift), expected_shift)
         sign *= -1

From 3850fa224bcb3fd84addc763a5be1a434690e6fc Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Thu, 18 Apr 2024 15:47:24 +0200
Subject: [PATCH 34/46] Change default color for unit cell to dark

---
 src/py4vasp/_config.py                    | 4 ++--
 src/py4vasp/_third_party/graph/contour.py | 2 +-
 tests/third_party/graph/test_graph.py     | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/py4vasp/_config.py b/src/py4vasp/_config.py
index 3d55cae7..2c3404e2 100644
--- a/src/py4vasp/_config.py
+++ b/src/py4vasp/_config.py
@@ -1,4 +1,4 @@
 # Copyright © VASP Software GmbH,
 # Licensed under the Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
-VASP_COLORS = ("#4C265F", "#2FB5AB", "#2C68FC", "#A82C35", "#808080")
-VASP_PURPLE, VASP_CYAN, VASP_BLUE, VASP_RED, VASP_GRAY = VASP_COLORS
+VASP_COLORS = ("#4C265F", "#2FB5AB", "#2C68FC", "#A82C35", "#808080", "#212529")
+VASP_PURPLE, VASP_CYAN, VASP_BLUE, VASP_RED, VASP_GRAY, VASP_DARK = VASP_COLORS
diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index 4766a78d..082facaa 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -143,7 +143,7 @@ def _create_unit_cell(self):
         corners = (vectors[0], vectors[0] + vectors[1], vectors[1])
         to_corners = (f"L {pos_to_str(corner)}" for corner in corners)
         path = f"M 0 0 {' '.join(to_corners)} Z"
-        unit_cell = {"type": "path", "line": {"color": _config.VASP_GRAY}, "path": path}
+        unit_cell = {"type": "path", "line": {"color": _config.VASP_DARK}, "path": path}
         return (unit_cell,)
 
     def _label_unit_cell_vectors(self):
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index 1113d5e7..1765bfda 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -464,7 +464,7 @@ def test_contour_supercell(rectangle_contour, Assert, not_core):
 def check_unit_cell(unit_cell, x, y, zero):
     assert unit_cell.type == "path"
     assert unit_cell.path == f"M 0 0 L {x} {zero} L {x} {y} L {zero} {y} Z"
-    assert unit_cell.line.color == _config.VASP_GRAY
+    assert unit_cell.line.color == _config.VASP_DARK
 
 
 def check_annotations(lattice, annotations, Assert):

From ccc6203375749c60fbdeab597e14dddada17697f Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Thu, 18 Apr 2024 16:00:11 +0200
Subject: [PATCH 35/46] Fix broken test

---
 tests/util/test_convert.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/util/test_convert.py b/tests/util/test_convert.py
index eade333b..390acee5 100644
--- a/tests/util/test_convert.py
+++ b/tests/util/test_convert.py
@@ -42,6 +42,7 @@ def test_hex_to_rgb(Assert):
         [44, 104, 252],
         [168, 44, 53],
         [128, 128, 128],
+        [33, 37, 41],
     ]
     expected = np.array(colors) / 255
     actual = np.array([to_rgb(color) for color in VASP_COLORS])

From 7e1576eb4d6af52d1aa95e1003ffac93137ac3d1 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Thu, 18 Apr 2024 16:32:37 +0200
Subject: [PATCH 36/46] Begin implementation of grid vectors for quiver plots

---
 src/py4vasp/_util/slicing.py        |  2 +-
 src/py4vasp/calculation/_density.py | 12 +++++++++++-
 tests/calculation/test_density.py   | 19 +++++++++++++++++++
 tests/util/test_slicing.py          | 29 +++++++++++++++++++++++------
 4 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index 4f66cc12..f548cf29 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -20,7 +20,7 @@ class Lattice:
     no labels are added."""
 
 
-def grid_data(data, cut, fraction):
+def grid_scalar(data, cut, fraction):
     """Takes a 2d slice of a 3d grid data.
 
     Often, we want to generate a 2d slice of a 3d data on a grid. One example would be
diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index 57f5bd47..a4aaa71d 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -307,13 +307,23 @@ def to_contour(
 
     def _contour(self, selector, selection, cut, fraction, lattice, supercell):
         density = selector[selection].T
-        data = slicing.grid_data(density, cut, fraction)
+        data = slicing.grid_scalar(density, cut, fraction)
         label = self._label(selector.label(selection)) or "charge"
         contour = graph.Contour(data, lattice, label, isolevels=True)
         if supercell is not None:
             contour.supercell = np.ones(2, dtype=np.int_) * supercell
         return contour
 
+    @_base.data_access
+    def to_quiver(self, *, a=None):
+        cut, fraction = self._get_cut(a, b=None, c=None)
+        data = slicing.grid_scalar(self._raw_data.charge[1].T, cut, fraction)
+        data = np.array((np.zeros_like(data), data))
+        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
+        label = self._selection or "magnetization"
+        quiver_plots = [graph.Contour(data, lattice, label)]
+        return graph.Graph(quiver_plots)
+
     def _get_cut(self, a, b, c):
         _raise_error_cut_selection_incorrect(a, b, c)
         if a is not None:
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index fe65d82d..86857551 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -386,6 +386,25 @@ def test_to_contour_operation(collinear_density, Assert):
     Assert.allclose(graph.series[0].data, expected_data)
 
 
+def test_collinear_to_quiver(collinear_density, Assert):
+    source = collinear_density.ref.source
+    if source == "charge":
+        expected_label = "magnetization"
+        expected_data = collinear_density.ref.output["magnetization"][8]
+    else:
+        expected_label = source
+        expected_data = collinear_density.ref.output[source][1, 8]
+    expected_data = np.array((np.zeros_like(expected_data), expected_data))
+    lattice_vectors = collinear_density.ref.structure.lattice_vectors()[1:]
+    expected_lattice = np.diag(np.linalg.norm(lattice_vectors, axis=1))
+    graph = collinear_density.to_quiver(a=-0.2)
+    assert len(graph) == 1
+    series = graph.series[0]
+    Assert.allclose(series.data, expected_data)
+    Assert.allclose(series.lattice.vectors, expected_lattice)
+    assert series.label == expected_label
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index f3f93408..0207abdb 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -99,16 +99,33 @@ def test_raise_error_for_unknown_choices():
 
 @pytest.mark.parametrize("cut", ("a", "b", "c"))
 @pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2))
-def test_slice_data(cut, fraction, Assert):
-    grid_data = np.random.random((10, 12, 14)) + 0.1
+def test_slice_grid_scalar(cut, fraction, Assert):
+    grid_scalar = np.random.random((10, 12, 14)) + 0.1
     if cut == "a":
         index = np.round(fraction * 10).astype(np.int_) % 10
-        expected_data = grid_data[index, :, :]
+        expected_data = grid_scalar[index, :, :]
     elif cut == "b":
         index = np.round(fraction * 12).astype(np.int_) % 12
-        expected_data = grid_data[:, index, :]
+        expected_data = grid_scalar[:, index, :]
     else:
         index = np.round(fraction * 14).astype(np.int_) % 14
-        expected_data = grid_data[:, :, index]
-    actual_data = slicing.grid_data(grid_data, cut, fraction)
+        expected_data = grid_scalar[:, :, index]
+    actual_data = slicing.grid_scalar(grid_scalar, cut, fraction)
     Assert.allclose(actual_data, expected_data)
+
+
+# @pytest.mark.parametrize("cut", ("a", "b", "c"))
+# @pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2))
+# def test_slice_grid_vector(cut, fraction, Assert):
+#     grid_vector = np.random.random((3, 10, 12, 14)) + 0.1
+#     if cut == "a":
+#         index = np.round(fraction * 10).astype(np.int_) % 10
+#         expected_data = grid_vector[:, index, :, :]
+#     elif cut == "b":
+#         index = np.round(fraction * 12).astype(np.int_) % 12
+#         expected_data = grid_vector[:, :, index, :]
+#     else:
+#         index = np.round(fraction * 14).astype(np.int_) % 14
+#         expected_data = grid_vector[:, :, :, index]
+#     actual_data = slicing.grid_vector(grid_vector, cut, fraction)
+#     Assert.allclose(actual_data, expected_data)

From 70da109e20e9da5edc62834ec9214b5d373e7b39 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Fri, 19 Apr 2024 10:43:18 +0200
Subject: [PATCH 37/46] Rename lattice to plane and add cut info

---
 src/py4vasp/_third_party/graph/contour.py |  5 ++--
 src/py4vasp/_util/slicing.py              | 31 +++++++++++------------
 tests/third_party/graph/test_graph.py     | 13 +++++-----
 tests/util/test_slicing.py                |  9 ++++---
 4 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/src/py4vasp/_third_party/graph/contour.py b/src/py4vasp/_third_party/graph/contour.py
index 082facaa..3f34a931 100644
--- a/src/py4vasp/_third_party/graph/contour.py
+++ b/src/py4vasp/_third_party/graph/contour.py
@@ -147,9 +147,10 @@ def _create_unit_cell(self):
         return (unit_cell,)
 
     def _label_unit_cell_vectors(self):
-        if self.lattice.labels is None:
+        if self.lattice.cut is None:
             return []
         vectors = self.lattice.vectors
+        labels = tuple("abc".replace(self.lattice.cut, ""))
         return [
             {
                 "text": label,
@@ -158,7 +159,7 @@ def _label_unit_cell_vectors(self):
                 "y": 0.5 * vectors[i, 1],
                 **self._shift_label(vectors[i], vectors[1 - i]),
             }
-            for i, label in enumerate(self.lattice.labels)
+            for i, label in enumerate(labels)
         ]
 
     def _shift_label(self, current_vector, other_vector):
diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index f548cf29..f61b39e0 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -11,16 +11,16 @@
 
 
 @dataclasses.dataclass
-class Lattice:
-    "Stores the 2d lattice vectors and their labels"
+class Plane:
+    """Defines a plane in 2d produced from cutting the 3d unit cell by removing one
+    lattice vector."""
     vectors: np.array
-    "Vectors spanning the plane."
-    labels: tuple = None
-    """Labels that could be added to the lattice vectors e.g. in a plot. If set to None
-    no labels are added."""
+    "2d vectors spanning the plane."
+    cut: str = None
+    "Lattice vector cut to get the plane, if not set, no labels will be added"
 
 
-def grid_scalar(data, cut, fraction):
+def grid_scalar(data, plane, fraction):
     """Takes a 2d slice of a 3d grid data.
 
     Often, we want to generate a 2d slice of a 3d data on a grid. One example would be
@@ -31,8 +31,8 @@ def grid_scalar(data, cut, fraction):
     ----------
     data : np.ndarray
         3d data on a grid from which a 2d plane is extracted.
-    cut : str
-        Defines along which dimension of the grid to cut (either "a", "b", or "c").
+    plane : Plane
+        Defines the 2d plane to which the data is reduced.
     fraction : float
         Determines which plane of the grid data is used. Periodic boundaries are assumed.
 
@@ -42,8 +42,8 @@ def grid_scalar(data, cut, fraction):
         A 2d array where the dimension selected by cut has been removed by selecting
         the plane according to the specificied fraction.
     """
-    _raise_error_if_cut_unknown(cut)
-    index = INDICES[cut]
+    _raise_error_if_cut_unknown(plane.cut)
+    index = INDICES[plane.cut]
     length = data.shape[index]
     slice_ = [slice(None), slice(None), slice(None)]
     slice_[index] = np.round(length * fraction).astype(np.int_) % length
@@ -87,16 +87,15 @@ def plane(cell, cut, normal="auto"):
 
     Returns
     -------
-    np.ndarray
-        A 2 × 2 array defining the two lattice vectors spanning the plane.
+    Plane
+        A 2d representation of the plane with some information to transform data to it.
     """
     _raise_error_if_cut_unknown(cut)
     vectors = np.delete(cell, INDICES[cut], axis=0)
-    labels = tuple("abc".replace(cut, ""))
     if normal is not None:
-        return Lattice(_rotate_normal_to_cartesian_axis(vectors, normal), labels)
+        return Plane(_rotate_normal_to_cartesian_axis(vectors, normal), cut)
     else:
-        return Lattice(_rotate_first_vector_to_x_axis(vectors), labels)
+        return Plane(_rotate_first_vector_to_x_axis(vectors), cut)
 
 
 def _rotate_first_vector_to_x_axis(vectors):
diff --git a/tests/third_party/graph/test_graph.py b/tests/third_party/graph/test_graph.py
index 1765bfda..1fd8b09f 100644
--- a/tests/third_party/graph/test_graph.py
+++ b/tests/third_party/graph/test_graph.py
@@ -66,7 +66,7 @@ def subplot():
 def rectangle_contour():
     return Contour(
         data=np.linspace(0, 10, 20 * 18).reshape((20, 18)),
-        lattice=slicing.Lattice(np.diag([4.0, 3.6]), ("a", "b")),
+        lattice=slicing.Plane(np.diag([4.0, 3.6]), cut="c"),
         label="rectangle contour",
         isolevels=True,
     )
@@ -76,7 +76,7 @@ def rectangle_contour():
 def tilted_contour():
     return Contour(
         data=np.linspace(0, 5, 16 * 20).reshape((16, 20)),
-        lattice=slicing.Lattice(np.array([[2, 3], [2, -3]]), ("a", "c")),
+        lattice=slicing.Plane(np.array([[2, 3], [2, -3]]), cut="b"),
         label="tilted contour",
         supercell=(2, 1),
         show_cell=False,
@@ -87,7 +87,7 @@ def tilted_contour():
 def simple_quiver():
     return Contour(
         data=np.array([[(y, x) for x in range(3)] for y in range(5)]).T,
-        lattice=slicing.Lattice(np.diag((3, 5)), ("b", "c")),
+        lattice=slicing.Plane(np.diag((3, 5)), cut="a"),
         label="quiver plot",
     )
 
@@ -96,7 +96,7 @@ def simple_quiver():
 def complex_quiver():
     return Contour(
         data=np.linspace(-3, 3, 2 * 12 * 10).reshape((2, 12, 10)),
-        lattice=slicing.Lattice([[3, 2], [-3, 2]], None),
+        lattice=slicing.Plane([[3, 2], [-3, 2]]),  # cut not set
         label="quiver plot",
         supercell=(3, 2),
     )
@@ -468,9 +468,10 @@ def check_unit_cell(unit_cell, x, y, zero):
 
 
 def check_annotations(lattice, annotations, Assert):
-    assert len(lattice.labels) == len(annotations)
+    assert len(lattice.vectors) == len(annotations)
     sign = np.sign(np.cross(*lattice.vectors))
-    for vector, label, annotation in zip(lattice.vectors, lattice.labels, annotations):
+    labels = "abc".replace(lattice.cut, "")
+    for vector, label, annotation in zip(lattice.vectors, labels, annotations):
         assert annotation.showarrow == False
         assert annotation.text == label
         Assert.allclose((annotation.x, annotation.y), 0.5 * vector)
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 0207abdb..f8953ee9 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -13,7 +13,7 @@ def test_orthorhombic(cut, index, Assert):
     expected_plane = np.delete(np.delete(cell, index, axis=0), index, axis=1)
     actual_plane = slicing.plane(cell, cut)
     Assert.allclose(actual_plane.vectors, expected_plane)
-    assert all(actual_plane.labels == np.delete(("a", "b", "c"), index))
+    assert actual_plane.cut == cut
 
 
 @pytest.mark.parametrize("cut, indices", [("a", (0, 2)), ("b", (1, 0)), ("c", (2, 1))])
@@ -22,7 +22,7 @@ def test_unusual_orthorhombic(cut, indices, Assert):
     expected_plane = np.delete(np.delete(cell, indices[0], axis=0), indices[1], axis=1)
     actual_plane = slicing.plane(cell, cut)
     Assert.allclose(actual_plane.vectors, expected_plane)
-    assert all(actual_plane.labels == np.delete(("a", "b", "c"), indices[0]))
+    assert actual_plane.cut == cut
 
 
 @pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)])
@@ -63,7 +63,7 @@ def test_normal_with_orthonormal_cell(normal, rotation, Assert):
     expected_plane = np.dot(np.delete(np.delete(cell, 0, axis=0), 0, axis=1), rotation)
     actual_plane = slicing.plane(cell, "a", normal)
     Assert.allclose(actual_plane.vectors, expected_plane)
-    assert actual_plane.labels == ("b", "c")
+    assert actual_plane.cut == "a"
 
 
 @pytest.mark.parametrize("cut, diagonal", [("a", [4, 5]), ("b", [3, 5]), ("c", [3, 4])])
@@ -110,7 +110,8 @@ def test_slice_grid_scalar(cut, fraction, Assert):
     else:
         index = np.round(fraction * 14).astype(np.int_) % 14
         expected_data = grid_scalar[:, :, index]
-    actual_data = slicing.grid_scalar(grid_scalar, cut, fraction)
+    plane = slicing.Plane(vectors=None, cut=cut)  # vectors should not be necessary
+    actual_data = slicing.grid_scalar(grid_scalar, plane, fraction)
     Assert.allclose(actual_data, expected_data)
 
 

From 8f8b136c61949c2a2b4a0805cd2f6cf0f946f579 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Fri, 19 Apr 2024 11:26:47 +0200
Subject: [PATCH 38/46] Implement slicing of vectorial data for orthorhombic
 cell

---
 src/py4vasp/_util/slicing.py | 34 ++++++++++++++++++++++++++++++++++
 tests/util/test_slicing.py   | 31 ++++++++++++++++---------------
 2 files changed, 50 insertions(+), 15 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index f61b39e0..b856e630 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -14,6 +14,7 @@
 class Plane:
     """Defines a plane in 2d produced from cutting the 3d unit cell by removing one
     lattice vector."""
+
     vectors: np.array
     "2d vectors spanning the plane."
     cut: str = None
@@ -50,6 +51,39 @@ def grid_scalar(data, plane, fraction):
     return data[tuple(slice_)]
 
 
+def grid_vector(data, plane, fraction):
+    """Takes a 2d slice of grid data where every datapoint is a 3d vector.
+
+    Often, we want to generate a 2d slice of data on a grid. One example would be to
+    visualize a slice through a plane as a quiver plot. This routine facilitates this
+    task implementing taking the cut of the 3d data.
+
+    Parameters
+    ----------
+    data : np.ndarray
+        Data on a grid where every point is a vector. The dimensions should be
+        (3d vector, grid_a, grid_b, grid_c) and will be reduced to a 2d vector on a
+        2d grid.
+    plane : Plane
+        Defines the 2d plane to which the data is reduced.
+    fraction : float
+        Determines which plane of the grid data is used. Periodic boundaries are assumed.
+
+    Returns
+    -------
+    np.ndarray
+        A 3d array where the dimension selected by cut has been removed by selecting
+        the plane according to the specificied fraction. The vector is projected onto
+        the plane for visualization.
+    """
+    _raise_error_if_cut_unknown(plane.cut)
+    index = INDICES[plane.cut]
+    length = data.shape[index + 1]  # add 1 to account for the vector dimension
+    slice_ = [np.delete(np.arange(3), index), slice(None), slice(None), slice(None)]
+    slice_[index + 1] = np.round(length * fraction).astype(np.int_) % length
+    return data[tuple(slice_)]
+
+
 def plane(cell, cut, normal="auto"):
     """Takes a 2d slice of a 3d cell and projects it onto 2d coordinates.
 
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index f8953ee9..51617d96 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -115,18 +115,19 @@ def test_slice_grid_scalar(cut, fraction, Assert):
     Assert.allclose(actual_data, expected_data)
 
 
-# @pytest.mark.parametrize("cut", ("a", "b", "c"))
-# @pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2))
-# def test_slice_grid_vector(cut, fraction, Assert):
-#     grid_vector = np.random.random((3, 10, 12, 14)) + 0.1
-#     if cut == "a":
-#         index = np.round(fraction * 10).astype(np.int_) % 10
-#         expected_data = grid_vector[:, index, :, :]
-#     elif cut == "b":
-#         index = np.round(fraction * 12).astype(np.int_) % 12
-#         expected_data = grid_vector[:, :, index, :]
-#     else:
-#         index = np.round(fraction * 14).astype(np.int_) % 14
-#         expected_data = grid_vector[:, :, :, index]
-#     actual_data = slicing.grid_vector(grid_vector, cut, fraction)
-#     Assert.allclose(actual_data, expected_data)
+@pytest.mark.parametrize("cut", ("a", "b", "c"))
+@pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2))
+def test_slice_grid_vector(cut, fraction, Assert):
+    grid_vector = np.random.random((3, 10, 12, 14)) + 0.1
+    if cut == "a":
+        index = np.round(fraction * 10).astype(np.int_) % 10
+        expected_data = grid_vector[1:, index, :, :]
+    elif cut == "b":
+        index = np.round(fraction * 12).astype(np.int_) % 12
+        expected_data = grid_vector[::2, :, index, :]
+    else:
+        index = np.round(fraction * 14).astype(np.int_) % 14
+        expected_data = grid_vector[:2, :, :, index]
+    plane = slicing.Plane(vectors=[[2, 0], [0, 3]], cut=cut)
+    actual_data = slicing.grid_vector(grid_vector, plane, fraction)
+    Assert.allclose(actual_data, expected_data)

From fbf32ddff13e0b30c36824f12b8b6ce10beded19 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Fri, 19 Apr 2024 11:44:06 +0200
Subject: [PATCH 39/46] Implement basic to_quiver for noncollinear calculations

---
 src/py4vasp/calculation/_density.py | 29 +++++++++++++++++------------
 tests/calculation/test_density.py   | 25 ++++++++++++++++++++++++-
 2 files changed, 41 insertions(+), 13 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index a4aaa71d..24adb79d 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -294,35 +294,40 @@ def to_contour(
         self, selection=None, *, a=None, b=None, c=None, supercell=None, normal=None
     ):
         cut, fraction = self._get_cut(a, b, c)
-        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal)
+        plane = slicing.plane(self._structure.lattice_vectors(), cut, normal)
         map_ = self._create_map()
         selector = index.Selector({0: map_}, self._raw_data.charge)
         tree = select.Tree.from_selection(selection)
         selections = self._filter_noncollinear_magnetization_from_selections(tree)
         contours = [
-            self._contour(selector, selection, cut, fraction, lattice, supercell)
+            self._contour(selector, selection, plane, fraction, supercell)
             for selection in selections
         ]
         return graph.Graph(contours)
 
-    def _contour(self, selector, selection, cut, fraction, lattice, supercell):
+    def _contour(self, selector, selection, plane, fraction, supercell):
         density = selector[selection].T
-        data = slicing.grid_scalar(density, cut, fraction)
+        data = slicing.grid_scalar(density, plane, fraction)
         label = self._label(selector.label(selection)) or "charge"
-        contour = graph.Contour(data, lattice, label, isolevels=True)
+        contour = graph.Contour(data, plane, label, isolevels=True)
         if supercell is not None:
             contour.supercell = np.ones(2, dtype=np.int_) * supercell
         return contour
 
     @_base.data_access
-    def to_quiver(self, *, a=None):
-        cut, fraction = self._get_cut(a, b=None, c=None)
-        data = slicing.grid_scalar(self._raw_data.charge[1].T, cut, fraction)
-        data = np.array((np.zeros_like(data), data))
-        lattice = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
+    def to_quiver(self, *, a=None, b=None, c=None, supercell=None):
+        cut, fraction = self._get_cut(a, b, c)
+        plane = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
+        if self.is_collinear():
+            data = slicing.grid_scalar(self._raw_data.charge[1].T, plane, fraction)
+            data = np.array((np.zeros_like(data), data))
+        else:
+            data = slicing.grid_vector(self.to_numpy()[1:], plane, fraction)
         label = self._selection or "magnetization"
-        quiver_plots = [graph.Contour(data, lattice, label)]
-        return graph.Graph(quiver_plots)
+        quiver_plot = graph.Contour(5.0 * data, plane, label)
+        if supercell is not None:
+            quiver_plot.supercell = np.ones(2, dtype=np.int_) * supercell
+        return graph.Graph([quiver_plot])
 
     def _get_cut(self, a, b, c):
         _raise_error_cut_selection_incorrect(a, b, c)
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 86857551..4b436d64 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -400,11 +400,34 @@ def test_collinear_to_quiver(collinear_density, Assert):
     graph = collinear_density.to_quiver(a=-0.2)
     assert len(graph) == 1
     series = graph.series[0]
-    Assert.allclose(series.data, expected_data)
+    Assert.allclose(series.data, 5.0 * expected_data)
     Assert.allclose(series.lattice.vectors, expected_lattice)
     assert series.label == expected_label
 
 
+def test_noncollinear_to_quiver(noncollinear_density, Assert):
+    source = noncollinear_density.ref.source
+    if source == "charge":
+        expected_label = "magnetization"
+        expected_data = noncollinear_density.ref.output["magnetization"][:2, :, :, 6]
+    else:
+        expected_label = source
+        expected_data = noncollinear_density.ref.output[source][1:3, :, :, 6]
+    expected_lattice = noncollinear_density.ref.structure.lattice_vectors()[:2, :2]
+    graph = noncollinear_density.to_quiver(c=1.4)
+    assert len(graph) == 1
+    series = graph.series[0]
+    Assert.allclose(series.data, 5.0 * expected_data)
+    Assert.allclose(series.lattice.vectors, expected_lattice)
+    assert series.label == expected_label
+
+def test_to_quiver_supercell(collinear_density, Assert):
+    graph = collinear_density.to_quiver(a=0, supercell=2)
+    Assert.allclose(graph.series[0].supercell, (2, 2))
+    graph = collinear_density.to_quiver(a=0, supercell=(2, 1))
+    Assert.allclose(graph.series[0].supercell, (2, 1))
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":

From cd09a03fb7d51f10e2d188f411ccab4c4b82c859 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Fri, 19 Apr 2024 13:26:13 +0200
Subject: [PATCH 40/46] Add cell argument to Plane

---
 src/py4vasp/_util/slicing.py | 6 ++++--
 tests/util/test_slicing.py   | 8 ++++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index b856e630..e4f78a69 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -17,6 +17,8 @@ class Plane:
 
     vectors: np.array
     "2d vectors spanning the plane."
+    cell: np.array = None
+    "Lattice vectors of the unit cell."
     cut: str = None
     "Lattice vector cut to get the plane, if not set, no labels will be added"
 
@@ -127,9 +129,9 @@ def plane(cell, cut, normal="auto"):
     _raise_error_if_cut_unknown(cut)
     vectors = np.delete(cell, INDICES[cut], axis=0)
     if normal is not None:
-        return Plane(_rotate_normal_to_cartesian_axis(vectors, normal), cut)
+        return Plane(_rotate_normal_to_cartesian_axis(vectors, normal), cell, cut)
     else:
-        return Plane(_rotate_first_vector_to_x_axis(vectors), cut)
+        return Plane(_rotate_first_vector_to_x_axis(vectors), cell, cut)
 
 
 def _rotate_first_vector_to_x_axis(vectors):
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 51617d96..6df04640 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -52,6 +52,7 @@ def test_rotate_to_user_defined_axis(normal, expected_plane, Assert):
     cell = np.array([[0, 1, 1], [1, 0, 1], [1, 1, 0]])
     actual_plane = slicing.plane(cell, "a", normal)
     Assert.allclose(actual_plane.vectors, expected_plane)
+    Assert.allclose(actual_plane.cell, cell)
 
 
 @pytest.mark.parametrize(
@@ -63,6 +64,7 @@ def test_normal_with_orthonormal_cell(normal, rotation, Assert):
     expected_plane = np.dot(np.delete(np.delete(cell, 0, axis=0), 0, axis=1), rotation)
     actual_plane = slicing.plane(cell, "a", normal)
     Assert.allclose(actual_plane.vectors, expected_plane)
+    Assert.allclose(actual_plane.cell, cell)
     assert actual_plane.cut == "a"
 
 
@@ -110,7 +112,9 @@ def test_slice_grid_scalar(cut, fraction, Assert):
     else:
         index = np.round(fraction * 14).astype(np.int_) % 14
         expected_data = grid_scalar[:, :, index]
-    plane = slicing.Plane(vectors=None, cut=cut)  # vectors should not be necessary
+    plane = slicing.Plane(
+        vectors=None, cut=cut
+    )  # vectors & cell should not be necessary
     actual_data = slicing.grid_scalar(grid_scalar, plane, fraction)
     Assert.allclose(actual_data, expected_data)
 
@@ -128,6 +132,6 @@ def test_slice_grid_vector(cut, fraction, Assert):
     else:
         index = np.round(fraction * 14).astype(np.int_) % 14
         expected_data = grid_vector[:2, :, :, index]
-    plane = slicing.Plane(vectors=[[2, 0], [0, 3]], cut=cut)
+    plane = slicing.Plane(vectors=[[2, 0], [0, 3]], cell=np.eye(3), cut=cut)
     actual_data = slicing.grid_vector(grid_vector, plane, fraction)
     Assert.allclose(actual_data, expected_data)

From 838a0c3ea229213cb0286fe5bc4bde97462c0bc7 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 23 Apr 2024 10:18:03 +0200
Subject: [PATCH 41/46] Implement projecting vectors onto generic plane

---
 src/py4vasp/_util/slicing.py      | 28 ++++++++++++++++++++++++++--
 tests/calculation/test_density.py |  1 +
 tests/util/test_slicing.py        | 31 ++++++++++++++++++++++++++++++-
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/src/py4vasp/_util/slicing.py b/src/py4vasp/_util/slicing.py
index e4f78a69..4b797a6d 100644
--- a/src/py4vasp/_util/slicing.py
+++ b/src/py4vasp/_util/slicing.py
@@ -81,9 +81,33 @@ def grid_vector(data, plane, fraction):
     _raise_error_if_cut_unknown(plane.cut)
     index = INDICES[plane.cut]
     length = data.shape[index + 1]  # add 1 to account for the vector dimension
-    slice_ = [np.delete(np.arange(3), index), slice(None), slice(None), slice(None)]
+    slice_ = [slice(None), slice(None), slice(None), slice(None)]
     slice_[index + 1] = np.round(length * fraction).astype(np.int_) % length
-    return data[tuple(slice_)]
+    return _project_vectors_to_plane(plane, data[tuple(slice_)])
+
+
+def _project_vectors_to_plane(plane, data):
+    # We want to want to project the vector r onto the plane spanned by the vectors
+    # u and v. Let the result be s = a u + b v. We can obtain the projected vector by
+    # minimizing the length |r - s| which leads to two conditions
+    # (r - a u - b v).u = 0
+    # (r - a u - b v).v = 0
+    # solving for a and b yields
+    # a = (v^2 u.r - u.v v.r) / N
+    # b = (u^2 v.r - u.v u.r) / N
+    # N = u^2 v^2 - (u.v)^2
+    u2 = np.dot(plane.vectors[0], plane.vectors[0])
+    v2 = np.dot(plane.vectors[1], plane.vectors[1])
+    uv = np.dot(plane.vectors[0], plane.vectors[1])
+    vectors = np.delete(plane.cell, INDICES[plane.cut], axis=0)
+    ur = np.tensordot(vectors[0], data, axes=(0, 0))
+    vr = np.tensordot(vectors[1], data, axes=(0, 0))
+    N = u2 * v2 - uv**2
+    a = (v2 * ur - uv * vr) / N
+    b = (u2 * vr - uv * ur) / N
+    au = np.multiply.outer(plane.vectors[0], a)
+    bv = np.multiply.outer(plane.vectors[1], b)
+    return au + bv
 
 
 def plane(cell, cut, normal="auto"):
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 4b436d64..92a81d3d 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -421,6 +421,7 @@ def test_noncollinear_to_quiver(noncollinear_density, Assert):
     Assert.allclose(series.lattice.vectors, expected_lattice)
     assert series.label == expected_label
 
+
 def test_to_quiver_supercell(collinear_density, Assert):
     graph = collinear_density.to_quiver(a=0, supercell=2)
     Assert.allclose(graph.series[0].supercell, (2, 2))
diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 6df04640..5834241a 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -123,15 +123,44 @@ def test_slice_grid_scalar(cut, fraction, Assert):
 @pytest.mark.parametrize("fraction", (-0.4, 0, 0.4, 0.8, 1.2))
 def test_slice_grid_vector(cut, fraction, Assert):
     grid_vector = np.random.random((3, 10, 12, 14)) + 0.1
+    ignored = 99
     if cut == "a":
         index = np.round(fraction * 10).astype(np.int_) % 10
+        cell = np.diag((ignored, 2, 3))
         expected_data = grid_vector[1:, index, :, :]
     elif cut == "b":
         index = np.round(fraction * 12).astype(np.int_) % 12
         expected_data = grid_vector[::2, :, index, :]
+        cell = np.diag((2, ignored, 3))
     else:
         index = np.round(fraction * 14).astype(np.int_) % 14
         expected_data = grid_vector[:2, :, :, index]
-    plane = slicing.Plane(vectors=[[2, 0], [0, 3]], cell=np.eye(3), cut=cut)
+        cell = np.diag((2, 3, ignored))
+    plane = slicing.Plane(vectors=[[2, 0], [0, 3]], cell=cell, cut=cut)
+    actual_data = slicing.grid_vector(grid_vector, plane, fraction)
+    Assert.allclose(actual_data, expected_data)
+
+
+def test_slice_grid_vector_nontrivial_cell(Assert):
+    cut = "a"
+    fraction = 0
+    cell = np.array([[0, 1, 1.1], [0.9, 0, 1.1], [0.9, 1, 0]])
+    normal = np.cross(cell[1], cell[2])
+    plane = slicing.plane(cell, cut, normal=None)
+    alpha = np.linspace(-1, 2, 10)
+    beta = np.linspace(-2, 1, 12)
+    gamma = np.linspace(-1.5, 1.5, 14)
+    grid_vector = np.array(
+        [
+            [[a * cell[1] + b * cell[2] + c * normal for a in alpha] for b in beta]
+            for c in gamma
+        ]
+    ).T
+    expected_data = np.array(
+        [
+            [alpha[0] * plane.vectors[0] + b * plane.vectors[1] for b in beta]
+            for c in gamma
+        ]
+    ).T
     actual_data = slicing.grid_vector(grid_vector, plane, fraction)
     Assert.allclose(actual_data, expected_data)

From 3911473fe4ad7f6ab2853b708f17214e74559417 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 23 Apr 2024 11:39:57 +0200
Subject: [PATCH 42/46] Add test for different cuts of quiver plot

---
 tests/util/test_slicing.py | 38 +++++++++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 11 deletions(-)

diff --git a/tests/util/test_slicing.py b/tests/util/test_slicing.py
index 5834241a..9eb1034d 100644
--- a/tests/util/test_slicing.py
+++ b/tests/util/test_slicing.py
@@ -141,26 +141,42 @@ def test_slice_grid_vector(cut, fraction, Assert):
     Assert.allclose(actual_data, expected_data)
 
 
-def test_slice_grid_vector_nontrivial_cell(Assert):
-    cut = "a"
+@pytest.mark.parametrize("cut, index", [("a", 0), ("b", 1), ("c", 2)])
+def test_slice_grid_vector_nontrivial_cell(cut, index, Assert):
     fraction = 0
     cell = np.array([[0, 1, 1.1], [0.9, 0, 1.1], [0.9, 1, 0]])
-    normal = np.cross(cell[1], cell[2])
+    vectors = np.delete(cell, index, axis=0)
+    normal = np.cross(*vectors)
     plane = slicing.plane(cell, cut, normal=None)
     alpha = np.linspace(-1, 2, 10)
     beta = np.linspace(-2, 1, 12)
     gamma = np.linspace(-1.5, 1.5, 14)
     grid_vector = np.array(
         [
-            [[a * cell[1] + b * cell[2] + c * normal for a in alpha] for b in beta]
-            for c in gamma
-        ]
-    ).T
-    expected_data = np.array(
-        [
-            [alpha[0] * plane.vectors[0] + b * plane.vectors[1] for b in beta]
+            [
+                [a * vectors[0] + b * vectors[1] + c * normal for a in alpha]
+                for b in beta
+            ]
             for c in gamma
         ]
     ).T
+    if cut == "a":
+        expected_data = np.array(
+            [
+                [alpha[0] * plane.vectors[0] + b * plane.vectors[1] for b in beta]
+                for c in gamma
+            ]
+        ).T
+    elif cut == "b":
+        expected_data = np.array(
+            [
+                [a * plane.vectors[0] + beta[0] * plane.vectors[1] for a in alpha]
+                for c in gamma
+            ]
+        ).T
+    else:
+        expected_data = np.array(
+            [[a * plane.vectors[0] + b * plane.vectors[1] for a in alpha] for b in beta]
+        ).T
     actual_data = slicing.grid_vector(grid_vector, plane, fraction)
-    Assert.allclose(actual_data, expected_data)
+    Assert.allclose(actual_data, expected_data, tolerance=3)

From f55dc4701663edccdc1f175e51ab890038654dd9 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 23 Apr 2024 11:53:44 +0200
Subject: [PATCH 43/46] Add normal vector for quiver plot

---
 src/py4vasp/calculation/_density.py |  4 ++--
 tests/calculation/test_density.py   | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index 24adb79d..d5c63f21 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -315,9 +315,9 @@ def _contour(self, selector, selection, plane, fraction, supercell):
         return contour
 
     @_base.data_access
-    def to_quiver(self, *, a=None, b=None, c=None, supercell=None):
+    def to_quiver(self, *, a=None, b=None, c=None, supercell=None, normal=None):
         cut, fraction = self._get_cut(a, b, c)
-        plane = slicing.plane(self._structure.lattice_vectors(), cut, normal=None)
+        plane = slicing.plane(self._structure.lattice_vectors(), cut, normal)
         if self.is_collinear():
             data = slicing.grid_scalar(self._raw_data.charge[1].T, plane, fraction)
             data = np.array((np.zeros_like(data), data))
diff --git a/tests/calculation/test_density.py b/tests/calculation/test_density.py
index 92a81d3d..e2586f06 100644
--- a/tests/calculation/test_density.py
+++ b/tests/calculation/test_density.py
@@ -429,6 +429,24 @@ def test_to_quiver_supercell(collinear_density, Assert):
     Assert.allclose(graph.series[0].supercell, (2, 1))
 
 
+@pytest.mark.parametrize(
+    "normal, rotation",
+    [
+        ("auto", np.eye(2)),
+        ("x", np.array([[0, -1], [1, 0]])),
+        ("y", np.diag((1, -1))),
+        ("z", np.eye(2)),
+    ],
+)
+def test_to_quiver_normal_vector(noncollinear_density, normal, rotation, Assert):
+    unrotated_graph = noncollinear_density.to_quiver(c=0.5)
+    rotated_graph = noncollinear_density.to_quiver(c=0.5, normal=normal)
+    expected_lattice = unrotated_graph.series[0].lattice.vectors @ rotation
+    Assert.allclose(rotated_graph.series[0].lattice.vectors, expected_lattice)
+    expected_data = (unrotated_graph.series[0].data.T @ rotation).T
+    Assert.allclose(rotated_graph.series[0].data, expected_data)
+
+
 def test_to_numpy(reference_density, Assert):
     source = reference_density.ref.source
     if source == "charge":

From 9099f823cbc7cdfb2669242d4388f0178939da7c Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 23 Apr 2024 12:35:13 +0200
Subject: [PATCH 44/46] Add documentation to to_contour

---
 src/py4vasp/calculation/_density.py | 67 +++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index d5c63f21..c2af82ca 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -293,6 +293,73 @@ def _use_symmetric_isosurface(self, component):
     def to_contour(
         self, selection=None, *, a=None, b=None, c=None, supercell=None, normal=None
     ):
+        """Generate a contour plot of the selected component of the density.
+
+        You need to specify a plane defined by two of the lattice vectors by selecting
+        a *cut* along the third one. You must only select a single cut and the value
+        should correspond to the fractional length along this third lattice vector.
+        py4vasp will then create a plane from the other two lattice vectors and
+        generate a contour plot within this plane.
+
+        Usually, the first remaining lattice vector is aligned with the x-axis and the
+        second one such that the angle between the vectors is preserved. You can
+        overwrite this choice by defining a normal direction. Then py4vasp will rotate
+        the normal vector of the plane to align with the specified direction. This is
+        particularly useful if the lattice vector you cut is aligned with a Cartesian
+        direction.
+
+        Parameters
+        ----------
+        selection : str
+            Select which component of the density you want to visualize. Please use the
+            `selections` method to get all available choices.
+
+        a, b, c : float
+            You must select exactly one of these to specify which of the three lattice
+            vectors you want to remove to form a plane. The assigned value represents
+            the fractional length along this lattice vector, so `a = 0.3` will remove
+            the first lattice vector and then take the grid points at 30% of the length
+            of the first vector in the b-c plane. The fractional height uses periodic
+            boundary conditions.
+
+        supercell : int or np.ndarray
+            Replicate the contour plot periodically a given number of times. If you
+            provide two different numbers, the resulting cell will be the two remaining
+            lattice vectors multiplied by the specific number.
+
+        normal : str or None
+            If not set, py4vasp will align the first remaining lattice vector with the
+            x-axis and the second one such that the angle between the lattice vectors
+            is preserved. You can set it to "x", "y", or "z"; then py4vasp will rotate
+            the plane in such a way that the normal direction aligns with the specified
+            Cartesian axis. This may look better if the normal direction is close to a
+            Cartesian axis. You may also set it to "auto" so that py4vasp chooses a
+            close Cartesian axis if it can find any.
+
+        Returns
+        -------
+        graph
+            A contour plot in the plane spanned by the 2 remaining lattice vectors.
+
+
+        Examples
+        --------
+
+        Cut a plane through the magnetization density at the origin of the third lattice
+        vector.
+
+        >>> calc.density.to_contour("3", c=0)
+
+        Replicate a plane in the middle of the second lattice vector 2 times in each
+        direction.
+
+        >>> calc.density.to_contour(b=0.5, supercell=2)
+
+        Take a slice of the kinetic energy density along the first lattice vector and
+        rotate it such that the normal of the plane aligns with the x axis.
+
+        >>> calc.density.to_contour("tau", a=0.3, normal="x")
+        """
         cut, fraction = self._get_cut(a, b, c)
         plane = slicing.plane(self._structure.lattice_vectors(), cut, normal)
         map_ = self._create_map()

From 8fae0b1e27317a6fb970c935a2042b7230c1d778 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 23 Apr 2024 13:26:47 +0200
Subject: [PATCH 45/46] Add documentation for to_quiver plot

---
 src/py4vasp/calculation/_density.py | 67 +++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index c2af82ca..e21d2156 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -383,6 +383,73 @@ def _contour(self, selector, selection, plane, fraction, supercell):
 
     @_base.data_access
     def to_quiver(self, *, a=None, b=None, c=None, supercell=None, normal=None):
+        """Generate a quiver plot of magnetization density.
+
+        You need to specify a plane defined by two of the lattice vectors by selecting
+        a *cut* along the third one. You must only select a single cut and the value
+        should correspond to the fractional length along this third lattice vector.
+        py4vasp will then create a plane from the other two lattice vectors and
+        generate a contour plot within this plane.
+
+        Usually, the first remaining lattice vector is aligned with the x-axis and the
+        second one such that the angle between the vectors is preserved. You can
+        overwrite this choice by defining a normal direction. Then py4vasp will rotate
+        the normal vector of the plane to align with the specified direction. This is
+        particularly useful if the lattice vector you cut is aligned with a Cartesian
+        direction.
+
+        For a collinear calculation, the magnetization density will be aligned with the
+        y axis of the plane. For noncollinear calculations, the magnetization density
+        is projected into the plane.
+
+        Parameters
+        ----------
+        a, b, c : float
+            You must select exactly one of these to specify which of the three lattice
+            vectors you want to remove to form a plane. The assigned value represents
+            the fractional length along this lattice vector, so `a = 0.3` will remove
+            the first lattice vector and then take the grid points at 30% of the length
+            of the first vector in the b-c plane. The fractional height uses periodic
+            boundary conditions.
+
+        supercell : int or np.ndarray
+            Replicate the contour plot periodically a given number of times. If you
+            provide two different numbers, the resulting cell will be the two remaining
+            lattice vectors multiplied by the specific number.
+
+        normal : str or None
+            If not set, py4vasp will align the first remaining lattice vector with the
+            x-axis and the second one such that the angle between the lattice vectors
+            is preserved. You can set it to "x", "y", or "z"; then py4vasp will rotate
+            the plane in such a way that the normal direction aligns with the specified
+            Cartesian axis. This may look better if the normal direction is close to a
+            Cartesian axis. You may also set it to "auto" so that py4vasp chooses a
+            close Cartesian axis if it can find any.
+
+        Returns
+        -------
+        graph
+            A quiver plot in the plane spanned by the 2 remaining lattice vectors.
+
+
+        Examples
+        --------
+
+        Cut a plane at the origin of the third lattice vector.
+
+        >>> calc.density.to_quiver(c=0)
+
+        Replicate a plane in the middle of the second lattice vector 2 times in each
+        direction.
+
+        >>> calc.density.to_quiver(b=0.5, supercell=2)
+
+        Take a slice of the spin components of the kinetic energy density along the
+        first lattice vector and rotate it such that the normal of the plane aligns with
+        the x axis.
+
+        >>> calc.density.to_quiver("tau", a=0.3, normal="x")
+        """
         cut, fraction = self._get_cut(a, b, c)
         plane = slicing.plane(self._structure.lattice_vectors(), cut, normal)
         if self.is_collinear():

From e6a218c45f74c7349f6a29c44eda4e657ef2a476 Mon Sep 17 00:00:00 2001
From: Martin Schlipf <martin.schlipf@gmail.com>
Date: Tue, 23 Apr 2024 13:39:28 +0200
Subject: [PATCH 46/46] Refactor common parts of documentation

---
 src/py4vasp/calculation/_density.py | 111 +++++++++++-----------------
 1 file changed, 45 insertions(+), 66 deletions(-)

diff --git a/src/py4vasp/calculation/_density.py b/src/py4vasp/calculation/_density.py
index e21d2156..19fb5c46 100644
--- a/src/py4vasp/calculation/_density.py
+++ b/src/py4vasp/calculation/_density.py
@@ -20,6 +20,45 @@
 }
 _MAGNETIZATION = ("magnetization", "mag", "m")
 
+_PLANE = """\
+You need to specify a plane defined by two of the lattice vectors by selecting
+a *cut* along the third one. You must only select a single cut and the value
+should correspond to the fractional length along this third lattice vector.
+py4vasp will then create a plane from the other two lattice vectors and
+generate a contour plot within this plane.
+
+Usually, the first remaining lattice vector is aligned with the x-axis and the
+second one such that the angle between the vectors is preserved. You can
+overwrite this choice by defining a normal direction. Then py4vasp will rotate
+the normal vector of the plane to align with the specified direction. This is
+particularly useful if the lattice vector you cut is aligned with a Cartesian
+direction.
+"""
+
+_COMMON_PARAMETERS = """\
+a, b, c : float
+    You must select exactly one of these to specify which of the three lattice
+    vectors you want to remove to form a plane. The assigned value represents
+    the fractional length along this lattice vector, so `a = 0.3` will remove
+    the first lattice vector and then take the grid points at 30% of the length
+    of the first vector in the b-c plane. The fractional height uses periodic
+    boundary conditions.
+
+supercell : int or np.ndarray
+    Replicate the contour plot periodically a given number of times. If you
+    provide two different numbers, the resulting cell will be the two remaining
+    lattice vectors multiplied by the specific number.
+
+normal : str or None
+    If not set, py4vasp will align the first remaining lattice vector with the
+    x-axis and the second one such that the angle between the lattice vectors
+    is preserved. You can set it to "x", "y", or "z"; then py4vasp will rotate
+    the plane in such a way that the normal direction aligns with the specified
+    Cartesian axis. This may look better if the normal direction is close to a
+    Cartesian axis. You may also set it to "auto" so that py4vasp chooses a
+    close Cartesian axis if it can find any.
+"""
+
 
 def _join_with_emphasis(data):
     emph_data = [f"*{x}*" for x in filter(lambda key: key != _INTERNAL, data)]
@@ -290,23 +329,13 @@ def _use_symmetric_isosurface(self, component):
         return component > 0
 
     @_base.data_access
+    @documentation.format(plane=_PLANE, common_parameters=_COMMON_PARAMETERS)
     def to_contour(
         self, selection=None, *, a=None, b=None, c=None, supercell=None, normal=None
     ):
         """Generate a contour plot of the selected component of the density.
 
-        You need to specify a plane defined by two of the lattice vectors by selecting
-        a *cut* along the third one. You must only select a single cut and the value
-        should correspond to the fractional length along this third lattice vector.
-        py4vasp will then create a plane from the other two lattice vectors and
-        generate a contour plot within this plane.
-
-        Usually, the first remaining lattice vector is aligned with the x-axis and the
-        second one such that the angle between the vectors is preserved. You can
-        overwrite this choice by defining a normal direction. Then py4vasp will rotate
-        the normal vector of the plane to align with the specified direction. This is
-        particularly useful if the lattice vector you cut is aligned with a Cartesian
-        direction.
+        {plane}
 
         Parameters
         ----------
@@ -314,27 +343,7 @@ def to_contour(
             Select which component of the density you want to visualize. Please use the
             `selections` method to get all available choices.
 
-        a, b, c : float
-            You must select exactly one of these to specify which of the three lattice
-            vectors you want to remove to form a plane. The assigned value represents
-            the fractional length along this lattice vector, so `a = 0.3` will remove
-            the first lattice vector and then take the grid points at 30% of the length
-            of the first vector in the b-c plane. The fractional height uses periodic
-            boundary conditions.
-
-        supercell : int or np.ndarray
-            Replicate the contour plot periodically a given number of times. If you
-            provide two different numbers, the resulting cell will be the two remaining
-            lattice vectors multiplied by the specific number.
-
-        normal : str or None
-            If not set, py4vasp will align the first remaining lattice vector with the
-            x-axis and the second one such that the angle between the lattice vectors
-            is preserved. You can set it to "x", "y", or "z"; then py4vasp will rotate
-            the plane in such a way that the normal direction aligns with the specified
-            Cartesian axis. This may look better if the normal direction is close to a
-            Cartesian axis. You may also set it to "auto" so that py4vasp chooses a
-            close Cartesian axis if it can find any.
+        {common_parameters}
 
         Returns
         -------
@@ -382,21 +391,11 @@ def _contour(self, selector, selection, plane, fraction, supercell):
         return contour
 
     @_base.data_access
+    @documentation.format(plane=_PLANE, common_parameters=_COMMON_PARAMETERS)
     def to_quiver(self, *, a=None, b=None, c=None, supercell=None, normal=None):
         """Generate a quiver plot of magnetization density.
 
-        You need to specify a plane defined by two of the lattice vectors by selecting
-        a *cut* along the third one. You must only select a single cut and the value
-        should correspond to the fractional length along this third lattice vector.
-        py4vasp will then create a plane from the other two lattice vectors and
-        generate a contour plot within this plane.
-
-        Usually, the first remaining lattice vector is aligned with the x-axis and the
-        second one such that the angle between the vectors is preserved. You can
-        overwrite this choice by defining a normal direction. Then py4vasp will rotate
-        the normal vector of the plane to align with the specified direction. This is
-        particularly useful if the lattice vector you cut is aligned with a Cartesian
-        direction.
+        {plane}
 
         For a collinear calculation, the magnetization density will be aligned with the
         y axis of the plane. For noncollinear calculations, the magnetization density
@@ -404,27 +403,7 @@ def to_quiver(self, *, a=None, b=None, c=None, supercell=None, normal=None):
 
         Parameters
         ----------
-        a, b, c : float
-            You must select exactly one of these to specify which of the three lattice
-            vectors you want to remove to form a plane. The assigned value represents
-            the fractional length along this lattice vector, so `a = 0.3` will remove
-            the first lattice vector and then take the grid points at 30% of the length
-            of the first vector in the b-c plane. The fractional height uses periodic
-            boundary conditions.
-
-        supercell : int or np.ndarray
-            Replicate the contour plot periodically a given number of times. If you
-            provide two different numbers, the resulting cell will be the two remaining
-            lattice vectors multiplied by the specific number.
-
-        normal : str or None
-            If not set, py4vasp will align the first remaining lattice vector with the
-            x-axis and the second one such that the angle between the lattice vectors
-            is preserved. You can set it to "x", "y", or "z"; then py4vasp will rotate
-            the plane in such a way that the normal direction aligns with the specified
-            Cartesian axis. This may look better if the normal direction is close to a
-            Cartesian axis. You may also set it to "auto" so that py4vasp chooses a
-            close Cartesian axis if it can find any.
+        {common_parameters}
 
         Returns
         -------