Source code for mdadash.backend.analyses.energies

"""
Widgets for various simulation energies
"""

import logging
from collections import deque
from typing import ClassVar

import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg

from mdadash.backend.widgets.base import WidgetBase

logger = logging.getLogger(__name__)


[docs] class EnergyWidgetBase: """ **Base class for Energy Widgets** This is the base class for all Energy widgets. Energy values are extracted from `MDAnalysis.coordinates.timestep.Timestep.data`_. The following keys are used for their respective plots: * ``temperature`` - Absolute Temperature * ``total_energy`` - Total Energy * ``potential_energy`` - Potential Energy * ``coulomb_energy`` - Coulomb Interaction Energy * ``bonds_energy`` - Bonds Energy * ``angles_energy`` - Angles Energy * ``dihedrals_energy`` - Dihedrals Energy * ``improper_dihedrals_energy`` - Improper Dihedrals Energy * ``van_der_walls_energy`` - Van Der Waals Energy .. note:: Energies are only available in the timestep data for streaming trajectories. The simulation engine must also be explicitly configured to send them .. _MDAnalysis.coordinates.timestep.Timestep.data: https://docs.mdanalysis.org/stable/ documentation_pages/coordinates/timestep.html#MDAnalysis.coordinates.timestep.Timestep.data **Inputs** Run frequency .. compound:: The frequency with which the widget is run - `every-frame` or `batch` Default: ``every-frame`` Max values Max values to show in plot Default: ``100`` Title Title for the plot Default: '' Plot refresh frequency The frequency with which the plot is refreshed (every n frames). This only applies when the run frequency is ``every-frame`` Default: ``1`` Reset on connect Reset the plot on every connect Default: ``False`` X-axis X-axis value - `time` or `step` Default: ``time`` **Outputs** Here is an example output plot of the **Absolute Temperature** widget: .. figure:: /_static/images/absolute_temperature_output.jpg :alt: Absolute Temperature output Here is an example output plot of the **Total Energy** widget: .. figure:: /_static/images/total_energy_output.jpg :alt: Total Energy output """ name = "" data_key = "" y_label = "Energy ( kJ / mol )" _doclink = ( "https://mdadash.readthedocs.io/en/latest/autosummary/" "mdadash.backend.analyses.energies.html" "#mdadash.backend.analyses.energies.EnergyWidgetBase" ) _notes = ( "Energies are only available for streaming trajectories and only if the " "simulation engine is configured to send them." ) _inputs: ClassVar = [ { "attribute": "_run_frequency", "name": "Run frequency", "description": "The frequency with which the widget is run", "type": "select", "items": [ "every-frame", "batch", ], }, { "attribute": "maxlen", "name": "Max values", "description": "Max values to show in plot", "type": "int", "validations": ["min:0"], }, { "attribute": "title", "name": "Title", "description": "Title for the plot", "type": "str", }, { "attribute": "plot_refresh_frequency", "name": "Plot refresh frequency", "description": "The frequency with which the plot is refreshed (every n frames)", "type": "int", "validations": ["min:1"], }, { "attribute": "reset_on_connect", "name": "Reset on connect", "description": "Reset the plot on every connect", "type": "bool", }, { "attribute": "x_type", "name": "X-axis", "type": "toggle", "options": [ {"name": "Time", "value": "time"}, {"name": "Step", "value": "step"}, ], }, ] def __init__(self): super().__init__() self.title = self.name self.default_maxlen = 100 self.maxlen = self.default_maxlen self.plot_refresh_count = 1 self.plot_refresh_frequency = 1 self.reset_on_connect = False self.x_type = "time" self.x_values = None self._setup_plot() self._reset_plot_values() def _setup_plot(self): """Setup matplotlib plot""" self.fig, self.ax = plt.subplots() self.canvas = FigureCanvasAgg(self.fig) (self.plot,) = self.ax.plot([], []) self._set_title() self.ax.set_ylabel(self.y_label) self.ax.grid(True) def _reset_plot_values(self): """Reset plot values""" self.steps = deque(maxlen=self.maxlen) self.times = deque(maxlen=self.maxlen) self.y_values = deque(maxlen=self.maxlen) self.plot_refresh_count = 1 self._set_x_values() def _set_title(self): """Set plot title""" self.ax.set_title(self.title) def _set_x_values(self): """Set the values for the x-axis""" if self.x_type == "step": x_label = "Step" self.x_values = self.steps else: x_label = "Time (ps)" self.x_values = self.times self.ax.set_xlabel(x_label)
[docs] def on_post_create(self): """:meth:`~mdadash.backend.widgets.base.WidgetBase.on_post_create` handler""" self._set_title() self._reset_plot_values()
[docs] def on_post_connect(self): """:meth:`~mdadash.backend.widgets.base.WidgetBase.on_post_connect` handler""" self._update_plot(self._compute_current_frame()) self.plot_refresh_count = 1 if self.reset_on_connect: # pragma: no cover self._reset_plot_values()
[docs] def on_input_change(self, attribute, _old_value, new_value): """:meth:`~mdadash.backend.widgets.base.WidgetBase.on_input_change` handler""" if attribute == "maxlen": if new_value < 0: self.maxlen = self.default_maxlen self._reset_plot_values() elif attribute == "title": self._set_title() elif attribute == "x_type": self._set_x_values() elif attribute == "plot_refresh_frequency": self.plot_refresh_count = 1
def _compute_current_frame(self): """Compute for current frame""" ts = self.u.trajectory.ts # pylint: disable=no-member return ( ts.data["step"], ts.data["time"], ts.data.get(self.data_key), ) def _compute_batch(self): """Compute for current batch""" u = self.u # pylint: disable=no-member values = [] for i in range(u.trajectory.buffer_size): _ = u.trajectory[i] values.append(self._compute_current_frame()) return values def _update_plot(self, values): """Append values and update plot""" if isinstance(values, tuple): values = [values] for value in values: (steps, times, v) = value self.steps.append(steps) self.times.append(times) self.y_values.append(v) # update plot points if self._run_frequency == "batch" or ( # pylint: disable=no-member self.plot_refresh_count == 1 or self.plot_refresh_count % self.plot_refresh_frequency == 0 ): self.plot.set_data(self.x_values, self.y_values) self.ax.relim() self.ax.autoscale_view() self.display_canvas(self.canvas) # pylint: disable=no-member self.plot_refresh_count += 1
[docs] def run_every_frame(self): """:meth:`~mdadash.backend.widgets.base.WidgetBase.run_every_frame` handler""" self._update_plot(self._compute_current_frame())
[docs] def run_batch(self): """:meth:`~mdadash.backend.widgets.base.WidgetBase.run_batch` handler""" self._update_plot(self._compute_batch())
[docs] class AbsoluteTemperature(EnergyWidgetBase, WidgetBase): """Absolute Temperature See :class:`EnergyWidgetBase` for more details. """ name = "Absolute Temperature" description = "Plot of Absolute Temperature" data_key = "temperature" y_label = "Temperature ( K )"
[docs] class TotalEnergy(EnergyWidgetBase, WidgetBase): """Total Energy See :class:`EnergyWidgetBase` for more details. """ name = "Total Energy" description = "Plot of Total Energy" data_key = "total_energy"
[docs] class PotentialEnergy(EnergyWidgetBase, WidgetBase): """Potential energy See :class:`EnergyWidgetBase` for more details. """ name = "Potential energy" description = "Plot of Potential Energy" data_key = "potential_energy"
[docs] class VanDerWaalsEnergy(EnergyWidgetBase, WidgetBase): """Van Der Waals Energy See :class:`EnergyWidgetBase` for more details. """ name = "Van Der Waals Energy" description = "Plot of Van Der Waals Energy" data_key = "van_der_walls_energy"
[docs] class CoulombInteractionEnergy(EnergyWidgetBase, WidgetBase): """Coulomb Interaction Energy See :class:`EnergyWidgetBase` for more details. """ name = "Coulomb Interaction Energy" description = "Plot of Coulomb Interaction Energy" data_key = "coulomb_energy"
[docs] class BondsEnergy(EnergyWidgetBase, WidgetBase): """Bonds Energy See :class:`EnergyWidgetBase` for more details. """ name = "Bonds Energy" description = "Plot of Bonds Energy" data_key = "bonds_energy"
[docs] class AnglesEnergy(EnergyWidgetBase, WidgetBase): """Angles Energy See :class:`EnergyWidgetBase` for more details. """ name = "Angles Energy" description = "Plot of Angles Energy" data_key = "angles_energy"
[docs] class DihedralsEnergy(EnergyWidgetBase, WidgetBase): """Dihedrals Energy See :class:`EnergyWidgetBase` for more details. """ name = "Dihedrals Energy" description = "Plot of Dihedrals Energy" data_key = "dihedrals_energy"
[docs] class ImproperDihedralsEnergy(EnergyWidgetBase, WidgetBase): """Improper Dihedrals Energy See :class:`EnergyWidgetBase` for more details. """ name = "Improper Dihedrals Energy" description = "Plot of Improper Dihedrals Energy" data_key = "improper_dihedrals_energy"