Add polar plot of orientation vectors (#358)

* added polar plot of orientation vectors

* Add plotly boilerplate (not implemented)

* Add test for polar plot

* Fix formatting

* Fix formatting

---------

Co-authored-by: Theo Famprikis <tfamprikis@tudelft.nl>
Co-authored-by: Stef Smeets <s.smeets@esciencecenter.nl>

Author: 3 people

src/gemdat/orientations.py

@@ -272,6 +272,11 @@ def plot_rectilinear(self, *, module, **kwargs):
         """See [gemdat.plots.rectilinear][] for more info."""
         return module.rectilinear(orientations=self, **kwargs)
 
+    @plot_backend
+    def plot_polar(self, *, module, **kwargs):
+        """See [gemdat.plots.polar][] for more info."""
+        return module.polar(orientations=self, **kwargs)
+
     @plot_backend
     def plot_bond_length_distribution(self, *, module, **kwargs):
         """See [gemdat.plots.bond_length_distribution][] for more info."""

src/gemdat/plots/__init__.py

@@ -4,6 +4,7 @@
 
 from .matplotlib import (
     jumps_3d_animation,
+    polar,
 )
 from .plotly import (
     autocorrelation,
@@ -43,6 +44,7 @@
     'msd_per_element',
     'msd_per_element',
     'plot_3d',
+    'polar',
     'radial_distribution',
     'rectilinear',
     'shape',

src/gemdat/plots/matplotlib/__init__.py

@@ -16,6 +16,7 @@
 from ._jumps_vs_distance import jumps_vs_distance
 from ._jumps_vs_time import jumps_vs_time
 from ._msd_per_element import msd_per_element
+from ._polar import polar
 from ._radial_distribution import radial_distribution
 from ._rectilinear import rectilinear
 from ._shape import shape
@@ -35,6 +36,7 @@
     'jumps_vs_distance',
     'jumps_vs_time',
     'msd_per_element',
+    'polar',
     'radial_distribution',
     'rectilinear',
     'shape',

src/gemdat/plots/matplotlib/_polar.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+if TYPE_CHECKING:
+    import matplotlib.figure
+
+    from gemdat.orientations import Orientations
+
+
+def polar(
+    *,
+    orientations: Orientations,
+    shape: tuple[int, int] = (90, 360),
+    normalize_histo: bool = True,
+) -> matplotlib.figure.Figure:
+    """Plot a polar projection of a spherical function. This function uses the
+    transformed trajectory.
+
+    Parameters
+    ----------
+    orientations : Orientations
+        The unit vector trajectories
+    shape : tuple
+        The shape of the spherical sector in which the trajectory is plotted
+    normalize_histo : bool, optional
+        If True, normalize the histogram by the area of the bins, by default True
+
+    Returns
+    -------
+    fig : matplotlib.figure.Figure
+        Output figure
+    """
+    from gemdat.orientations import calculate_spherical_areas
+
+    az, el, _ = orientations.vectors_spherical.T
+    az = az.flatten()
+    el = el.flatten()
+
+    hist, *_ = np.histogram2d(el, az, shape)
+
+    if normalize_histo:
+        areas = calculate_spherical_areas(shape)
+        hist = hist / areas
+        # Drop the bins at the poles where normalization is not possible
+        hist = hist[1:-1, :]
+
+    axis_theta, axis_phi = hist.shape
+
+    phi = np.radians(np.linspace(0, 360, axis_phi))
+    theta = np.linspace(0, 180, axis_theta)
+
+    theta, phi = np.meshgrid(theta, phi)
+
+    fig, (ax1, ax2) = plt.subplots(1, 2, subplot_kw=dict(projection='polar'))
+
+    cs1 = ax1.contourf(phi, theta, hist.T)
+    ax1.set_title('θ < 90°')
+    ax1.set_rmax(90)
+    ax1.set_yticklabels([])
+
+    ax2.contourf(phi, 180 - theta, hist.T)
+    ax2.set_title('θ > 90°')
+    ax2.set_rmax(90)
+    ax2.set_yticklabels([])
+
+    fig.colorbar(cs1, ax=[ax1, ax2], orientation='horizontal', label='Areal Probability')
+
+    plt.subplots_adjust(wspace=0.5, bottom=0.35)  # Increase horizontal spacing
+
+    return fig

src/gemdat/plots/plotly/__init__.py

@@ -16,6 +16,7 @@
 from ._jumps_vs_time import jumps_vs_time
 from ._msd_per_element import msd_per_element
 from ._plot3d import plot_3d
+from ._polar import polar
 from ._radial_distribution import radial_distribution
 from ._rectilinear import rectilinear
 from ._shape import shape
@@ -36,6 +37,7 @@
     'jumps_vs_time',
     'msd_per_element',
     'plot_3d',
+    'polar',
     'radial_distribution',
     'rectilinear',
     'shape',

src/gemdat/plots/plotly/_polar.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import plotly.graph_objects as go
+
+    from gemdat.orientations import Orientations
+
+
+def polar(
+    *,
+    orientations: Orientations,
+    shape: tuple[int, int] = (90, 360),
+    normalize_histo: bool = True,
+) -> go.Figure:
+    """Plot a polar projection of a spherical function. This function uses the
+    transformed trajectory.
+
+    Parameters
+    ----------
+    orientations : Orientations
+        The unit vector trajectories
+    shape : tuple
+        The shape of the spherical sector in which the trajectory is plotted
+    normalize_histo : bool, optional
+        If True, normalize the histogram by the area of the bins, by default True
+
+    Returns
+    -------
+    fig : plotly.graph_objects.Figure
+        Output figure
+    """
+    raise NotImplementedError('Use the matplotlib version.')

tests/integration/baseline_images/plot_mpl_test/polar.png

@@ -129,3 +129,8 @@ def test_bond_length_distribution(vasp_orientations):
 @image_comparison2(baseline_images=['autocorrelation'])
 def test_autocorrelation(vasp_orientations):
     vasp_orientations.plot_autocorrelation(backend=BACKEND)
+
+
+@image_comparison2(baseline_images=['polar'])
+def test_polar(vasp_orientations):
+    vasp_orientations.plot_polar(backend=BACKEND)