Merge pull request #23 from xhgchen/plot-vacf

Create plotting function for VACF

transport_analysis/tests/test_velocityautocorr.py

@@ -1,5 +1,5 @@
 import pytest
-from numpy.testing import assert_almost_equal
+from numpy.testing import assert_almost_equal, assert_allclose
 
 from transport_analysis.velocityautocorr import (
     VelocityAutocorr as VACF,
@@ -181,6 +181,43 @@ def test_simple_start_stop_step_all_dims(
         # polynomial must take offset start into account
         assert_almost_equal(v_simple.results.timeseries, poly, decimal=4)
 
+    def test_plot_vacf(self, vacf):
+        # Expected data to be plotted
+        x_exp = vacf.times
+        y_exp = vacf.results.timeseries
+
+        # Actual data returned from plot
+        (line,) = vacf.plot_vacf()
+        x_act, y_act = line.get_xydata().T
+
+        assert_allclose(x_act, x_exp)
+        assert_allclose(y_act, y_exp)
+
+    def test_plot_vacf_labels(self, vacf):
+        # Expected labels
+        x_exp = "Time (ps)"
+        y_exp = "Velocity Autocorrelation Function (VACF) (Å)"
+
+        # Actual labels returned from plot
+        (line,) = vacf.plot_vacf()
+        x_act = line.axes.get_xlabel()
+        y_act = line.axes.get_ylabel()
+
+        assert x_act == x_exp
+        assert y_act == y_exp
+
+    def test_plot_vacf_start_stop_step(self, vacf, start=1, stop=9, step=2):
+        # Expected data to be plotted
+        x_exp = vacf.times[start:stop:step]
+        y_exp = vacf.results.timeseries[start:stop:step]
+
+        # Actual data returned from plot
+        (line,) = vacf.plot_vacf(start=start, stop=stop, step=step)
+        x_act, y_act = line.get_xydata().T
+
+        assert_allclose(x_act, x_exp)
+        assert_allclose(y_act, y_exp)
+
 
 class TestVACFFFT(object):
     @pytest.fixture(scope="class")

transport_analysis/velocityautocorr.py

@@ -49,6 +49,7 @@
 from MDAnalysis.exceptions import NoDataError
 import numpy as np
 import tidynamics
+import matplotlib.pyplot as plt
 from transport_analysis.due import due, Doi
 
 if TYPE_CHECKING:
@@ -57,7 +58,7 @@
 due.cite(
     Doi("10.21105/joss.00877"),
     description="Autocorrelation with tidynamics",
-    path="transport_analysis.analysis.velocityautocorr",
+    path="transport_analysis.velocityautocorr",
     cite_module=True,
 )
 
@@ -93,7 +94,7 @@ class VelocityAutocorr(AnalysisBase):
     start : Optional[int]
         The first frame of the trajectory used to compute the analysis.
     stop : Optional[int]
-        The frame to stop at for the analysis.
+        The frame to stop at for the analysis, non-inclusive.
     step : Optional[int]
         Number of frames to skip between each analyzed frame.
     n_frames : int
@@ -228,3 +229,36 @@ def _conclude_simple(self):
             self.results.vacf_by_particle[lag, :] = np.mean(sum_veloc, axis=0)
         # average over # particles and update results array
         self.results.timeseries = self.results.vacf_by_particle.mean(axis=1)
+
+    def plot_vacf(self, start=0, stop=0, step=1):
+        """
+        Returns a velocity autocorrelation function (VACF) plot via
+        ``Matplotlib``. Analysis must be run prior to plotting.
+
+        Parameters
+        ----------
+        start : Optional[int]
+            The first frame of ``self.results.timeseries``
+            used for the plot.
+        stop : Optional[int]
+            The frame of ``self.results.timeseries`` to stop at
+            for the plot, non-inclusive.
+        step : Optional[int]
+            Number of frames to skip between each plotted frame.
+
+        Returns
+        -------
+        :class:`matplotlib.lines.Line2D`
+            A :class:`matplotlib.lines.Line2D` instance with
+            the desired VACF plotting information.
+        """
+
+        stop = self.n_frames if stop == 0 else stop
+
+        fig, ax_vacf = plt.subplots()
+        ax_vacf.set_xlabel("Time (ps)")
+        ax_vacf.set_ylabel("Velocity Autocorrelation Function (VACF) (Å)")
+        return ax_vacf.plot(
+            self.times[start:stop:step],
+            self.results.timeseries[start:stop:step],
+        )