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,
+ )