Refactor plotting code

Author: stefsmeets

src/gemdat/plots/_shared.py

@@ -11,6 +11,7 @@
 from scipy.stats import skewnorm
 
 if TYPE_CHECKING:
+    from gemdat.jumps import Jumps
     from gemdat.orientations import Orientations
     from gemdat.trajectory import Trajectory
 
@@ -127,7 +128,7 @@ def dataframe(self):
 def _get_vibrational_amplitudes_hist(
     *, trajectories: list[Trajectory], bins: int
 ) -> VibrationalAmplitudeHist:
-    """Calculate vabrational amplitudes histogram.
+    """Calculate vibrational amplitudes histogram.
 
     Helper for `vibrational_amplitudes`.
     """
@@ -150,3 +151,38 @@ def _get_vibrational_amplitudes_hist(
     std = np.std(data, axis=0)
 
     return VibrationalAmplitudeHist(amplitudes=amplitudes, counts=mean, std=std)
+
+
+def _jumps_vs_distance(jumps: Jumps, *, resolution: float, n_parts: int) -> pd.DataFrame:
+    """Calculate jumps vs distance histogram.
+
+    Helper for `jumps_vs_distance`.
+    """
+    sites = jumps.sites
+    trajectory = jumps.trajectory
+    lattice = trajectory.get_lattice()
+
+    pdist = lattice.get_all_distances(sites.frac_coords, sites.frac_coords)
+
+    bin_max = (1 + pdist.max() // resolution) * resolution
+    n_bins = int(bin_max / resolution) + 1
+    x = np.linspace(0, bin_max, n_bins)
+
+    bin_idx = np.digitize(pdist, bins=x)
+    data = []
+    for transitions_part in jumps.split(n_parts=n_parts):
+        counts = np.zeros_like(x)
+        for idx, n in zip(bin_idx.flatten(), transitions_part.matrix().flatten()):
+            counts[idx] += n
+        for idx in range(n_bins):
+            if counts[idx] > 0:
+                data.append((x[idx], counts[idx]))
+
+    df = pd.DataFrame(data=data, columns=['Displacement', 'count'])
+
+    grouped = df.groupby(['Displacement'])
+    mean = grouped.mean().reset_index().rename(columns={'count': 'mean'})
+    std = grouped.std().reset_index().rename(columns={'count': 'std'})
+    df = mean.merge(std, how='inner')
+
+    return df

src/gemdat/plots/matplotlib/_jumps_vs_distance.py

@@ -3,13 +3,13 @@
 from typing import TYPE_CHECKING
 
 import matplotlib.pyplot as plt
-import numpy as np
-import pandas as pd
+
+from .._shared import _jumps_vs_distance
 
 if TYPE_CHECKING:
     import matplotlib.figure
 
-    from gemdat import Jumps
+    from gemdat.jumps import Jumps
 
 
 def jumps_vs_distance(
@@ -34,32 +34,7 @@ def jumps_vs_distance(
     fig : matplotlib.figure.Figure
         Output figure
     """
-    sites = jumps.sites
-    trajectory = jumps.trajectory
-    lattice = trajectory.get_lattice()
-
-    pdist = lattice.get_all_distances(sites.frac_coords, sites.frac_coords)
-
-    bin_max = (1 + pdist.max() // jump_res) * jump_res
-    n_bins = int(bin_max / jump_res) + 1
-    x = np.linspace(0, bin_max, n_bins)
-
-    bin_idx = np.digitize(pdist, bins=x)
-    data = []
-    for transitions_part in jumps.split(n_parts=n_parts):
-        counts = np.zeros_like(x)
-        for idx, n in zip(bin_idx.flatten(), transitions_part.matrix().flatten()):
-            counts[idx] += n
-        for idx in range(n_bins):
-            if counts[idx] > 0:
-                data.append((x[idx], counts[idx]))
-
-    df = pd.DataFrame(data=data, columns=['Displacement', 'count'])
-
-    grouped = df.groupby(['Displacement'])
-    mean = grouped.mean().reset_index().rename(columns={'count': 'mean'})
-    std = grouped.std().reset_index().rename(columns={'count': 'std'})
-    df = mean.merge(std, how='inner')
+    df = _jumps_vs_distance(jumps=jumps, resolution=jump_res, n_parts=n_parts)
 
     fig, ax = plt.subplots()
 

src/gemdat/plots/plotly/_jumps_vs_distance.py

@@ -2,11 +2,11 @@
 
 from typing import TYPE_CHECKING
 
-import numpy as np
-import pandas as pd
 import plotly.express as px
 import plotly.graph_objects as go
 
+from .._shared import _jumps_vs_distance
+
 if TYPE_CHECKING:
     from gemdat import Jumps
 
@@ -33,32 +33,7 @@ def jumps_vs_distance(
     fig : plotly.graph_objects.Figure
         Output figure
     """
-    sites = jumps.sites
-    trajectory = jumps.trajectory
-    lattice = trajectory.get_lattice()
-
-    pdist = lattice.get_all_distances(sites.frac_coords, sites.frac_coords)
-
-    bin_max = (1 + pdist.max() // jump_res) * jump_res
-    n_bins = int(bin_max / jump_res) + 1
-    x = np.linspace(0, bin_max, n_bins)
-
-    bin_idx = np.digitize(pdist, bins=x)
-    data = []
-    for transitions_part in jumps.split(n_parts=n_parts):
-        counts = np.zeros_like(x)
-        for idx, n in zip(bin_idx.flatten(), transitions_part.matrix().flatten()):
-            counts[idx] += n
-        for idx in range(n_bins):
-            if counts[idx] > 0:
-                data.append((x[idx], counts[idx]))
-
-    df = pd.DataFrame(data=data, columns=['Displacement', 'count'])
-
-    grouped = df.groupby(['Displacement'])
-    mean = grouped.mean().reset_index().rename(columns={'count': 'mean'})
-    std = grouped.std().reset_index().rename(columns={'count': 'std'})
-    df = mean.merge(std, how='inner')
+    df = _jumps_vs_distance(jumps=jumps, resolution=jump_res, n_parts=n_parts)
 
     if n_parts == 1:
         fig = px.bar(df, x='Displacement', y='mean', barmode='stack')