Create plotting functions for viscosity function and running viscosity

transport_analysis/tests/test_viscosity.py

@@ -147,6 +147,100 @@ def test_dimtype_error(self, ag, dimtype):
         with pytest.raises(ValueError, match=errmsg):
             VH(ag, dim_type=dimtype)
 
+    def test_plot_viscosity_function(self, visc_helfand):
+        # Expected data to be plotted
+        x_exp = visc_helfand.times
+        y_exp = visc_helfand.results.timeseries
+
+        # Actual data returned from plot
+        (line,) = visc_helfand.plot_viscosity_function()
+        x_act, y_act = line.get_xydata().T
+
+        assert_allclose(x_act, x_exp)
+        assert_allclose(y_act, y_exp)
+
+    def test_plot_viscosity_function_labels(self, visc_helfand):
+        # Expected labels
+        x_exp = "Time (ps)"
+        y_exp = "Viscosity Function"  # TODO: Specify units
+
+        # Actual labels returned from plot
+        (line,) = visc_helfand.plot_viscosity_function()
+        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_viscosity_function_start_stop_step(
+        self, visc_helfand, start=1, stop=9, step=2
+    ):
+        # Expected data to be plotted
+        x_exp = visc_helfand.times[start:stop:step]
+        y_exp = visc_helfand.results.timeseries[start:stop:step]
+
+        # Actual data returned from plot
+        (line,) = visc_helfand.plot_viscosity_function(
+            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)
+
+    def test_plot_viscosity_function_exception(self, step_vtraj_full):
+        vis_h = VH(step_vtraj_full.atoms)
+        errmsg = "Analysis must be run"
+        with pytest.raises(RuntimeError, match=errmsg):
+            vis_h.plot_viscosity_function()
+
+    def test_plot_running_viscosity(self, visc_helfand):
+        # Expected data to be plotted
+        x_exp = visc_helfand.times[1:]
+        y_exp = visc_helfand.results.timeseries[1:] / x_exp
+
+        # Actual data returned from plot
+        (line,) = visc_helfand.plot_running_viscosity()
+        x_act, y_act = line.get_xydata().T
+
+        assert_allclose(x_act, x_exp)
+        assert_allclose(y_act, y_exp)
+
+    def test_plot_running_viscosity_labels(self, visc_helfand):
+        # Expected labels
+        x_exp = "Time (ps)"
+        y_exp = "Running Viscosity"  # TODO: Specify units
+
+        # Actual labels returned from plot
+        (line,) = visc_helfand.plot_running_viscosity()
+        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_running_viscosity_start_stop_step(
+        self, visc_helfand, start=1, stop=9, step=2
+    ):
+        # Expected data to be plotted
+        x_exp = visc_helfand.times[start:stop:step]
+        y_exp = visc_helfand.results.timeseries[start:stop:step] / x_exp
+
+        # Actual data returned from plot
+        (line,) = visc_helfand.plot_running_viscosity(
+            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)
+
+    def test_plot_running_viscosity_exception(self, step_vtraj_full):
+        vis_h = VH(step_vtraj_full.atoms)
+        errmsg = "Analysis must be run"
+        with pytest.raises(RuntimeError, match=errmsg):
+            vis_h.plot_running_viscosity()
+
 
 @pytest.mark.parametrize(
     "tdim, tdim_factor",

transport_analysis/viscosity.py

@@ -13,6 +13,7 @@
 from MDAnalysis.exceptions import NoDataError
 from MDAnalysis.units import constants
 import numpy as np
+import matplotlib.pyplot as plt
 
 if TYPE_CHECKING:
     from MDAnalysis.core.universe import AtomGroup
@@ -57,6 +58,10 @@ class ViscosityHelfand(AnalysisBase):
         Number of frames analysed in the trajectory.
     n_particles : int
         Number of particles viscosity was calculated over.
+    running_viscosity : :class:`numpy.ndarray`
+        The running viscosity of the analysis calculated from dividing
+        ``results.timeseries`` by the corresponding times. Obtained after
+        calling :meth:`ViscosityHelfand.plot_running_viscosity`
     """
 
     def __init__(
@@ -87,6 +92,7 @@ def __init__(
         # local
         self.atomgroup = atomgroup
         self.n_particles = len(self.atomgroup)
+        self._run_called = False
 
     def _prepare(self):
         """
@@ -211,3 +217,86 @@ def _conclude(self):
         )
         # average over # particles and update results array
         self.results.timeseries = self.results.visc_by_particle.mean(axis=1)
+        self._run_called = True
+
+    def plot_viscosity_function(self, start=0, stop=0, step=1):
+        """
+        Returns a viscosity function plot via ``Matplotlib``. Usage
+        of this plot is recommended to help determine where to take the
+        slope of the viscosity function to obtain the viscosity.
+        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 viscosity function plotting information.
+        """
+        if not self._run_called:
+            raise RuntimeError("Analysis must be run prior to plotting")
+
+        stop = self.n_frames if stop == 0 else stop
+
+        fig, ax_visc = plt.subplots()
+        ax_visc.set_xlabel("Time (ps)")
+        ax_visc.set_ylabel("Viscosity Function")  # TODO: Specify units
+        return ax_visc.plot(
+            self.times[start:stop:step],
+            self.results.timeseries[start:stop:step],
+        )
+
+    def plot_running_viscosity(self, start=1, stop=0, step=1):
+        """
+        Returns a running viscosity plot via ``Matplotlib``. `start` is
+        set to `1` by default to avoid division by 0.
+        Usage of this plot will give an idea of the viscosity over the course
+        of the simulation but it is recommended for users to exercise their
+        best judgement and take the slope of the viscosity function to obtain
+        viscosity rather than use the running viscosity as a final result.
+        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 running viscosity plotting information.
+        """
+        if not self._run_called:
+            raise RuntimeError("Analysis must be run prior to plotting")
+
+        stop = self.n_frames if stop == 0 else stop
+
+        self.running_viscosity = (
+            self.results.timeseries[start:stop:step]
+            / self.times[start:stop:step]
+        )
+
+        fig, ax_running_visc = plt.subplots()
+        ax_running_visc.set_xlabel("Time (ps)")
+        ax_running_visc.set_ylabel("Running Viscosity")  # TODO: Specify units
+        return ax_running_visc.plot(
+            self.times[start:stop:step],
+            self.running_viscosity,
+        )