diff --git a/transport_analysis/tests/test_viscosity.py b/transport_analysis/tests/test_viscosity.py index dace793..c0500dd 100644 --- a/transport_analysis/tests/test_viscosity.py +++ b/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", diff --git a/transport_analysis/viscosity.py b/transport_analysis/viscosity.py index b1f104c..74f97d8 100644 --- a/transport_analysis/viscosity.py +++ b/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, + )