From 7d3356561fbdf600db920963929757d3db0b2f6a Mon Sep 17 00:00:00 2001
From: arm61 <a.r.mccluskey@bath.ac.uk>
Date: Tue, 5 Mar 2019 10:58:50 +0000
Subject: [PATCH] add maxwell boltzmann distribution visual

---
 .../simple_examples/molecular_dynamics.ipynb  |  4 +-
 pylj/sample.py                                | 73 +++++++++++++++++++
 2 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/examples/simple_examples/molecular_dynamics.ipynb b/examples/simple_examples/molecular_dynamics.ipynb
index b39a53f..93a2127 100644
--- a/examples/simple_examples/molecular_dynamics.ipynb
+++ b/examples/simple_examples/molecular_dynamics.ipynb
@@ -26,7 +26,7 @@
     "    # Initialise the system\n",
     "    system = md.initialise(number_of_particles, temperature, box_length, 'square')\n",
     "    # This sets the sampling class\n",
-    "    sample_system = sample.Energy(system)\n",
+    "    sample_system = sample.MaxBolt(system)\n",
     "    # Start at time 0\n",
     "    system.time = 0\n",
     "    # Begin the molecular dynamics loop\n",
@@ -67,7 +67,7 @@
    },
    "outputs": [],
    "source": [
-    "md_simulation(50, 273.15, 100, 5000, 25)"
+    "md_simulation(20, 1000, 100, 50000, 100)"
    ]
   },
   {
diff --git a/pylj/sample.py b/pylj/sample.py
index fb5d36b..4d06397 100644
--- a/pylj/sample.py
+++ b/pylj/sample.py
@@ -314,6 +314,44 @@ def update(self, system):
         update_energyview(self.ax[1], system)
         self.fig.canvas.draw()
 
+class MaxBolt(object):  # pragma: no cover
+    """The MaxBolt class will plot the particle positions and a histogram
+    of the particle velocities.
+
+    Parameters
+    ----------
+    system: System
+        The whole system information.
+    size: string
+        The size of the visualisation:
+            - 'small'
+            - 'medium' (default)
+            - 'large'
+    """
+
+    def __init__(self, system, size='medium'):
+        fig, ax = environment(2, size)
+        self.velocities = np.array([])
+        setup_cellview(ax[0], system)
+        setup_maxbolthist(ax[1])
+
+        plt.tight_layout()
+        self.ax = ax
+        self.fig = fig
+
+    def update(self, system):
+        """This updates the visualisation environment. Often this can be slower
+        than the cythonised force calculation so used is wisely.
+
+        Parameters
+        ----------
+        system: System
+            The whole system information.
+        """
+        update_cellview(self.ax[0], system)
+        self.velocities = update_maxbolthist(self.ax[1], system, self.velocities)
+        self.fig.canvas.draw()
+
 
 class RDF(object):  # pragma: no cover
     """The RDF class will plot the particle positions and radial distribution
@@ -518,6 +556,19 @@ def setup_tempview(ax):  # pragma: no cover
     ax.set_ylabel("Temperature/K", fontsize=16)
     ax.set_xlabel("Time/s", fontsize=16)
 
+def setup_maxbolthist(ax):  # pragma: no cover
+    """Builds the simulation velocity histogram
+    visualisation pane.
+
+    Parameters
+    ----------
+    ax: Axes object
+        The axes position that the pane should be placed in.
+    """
+
+    ax.step([0], [0], color="#34a5daff")
+    ax.set_ylabel("PDF", fontsize=16)
+    ax.set_xlabel("Velocity/ms$^{-1}$", fontsize=16)
 
 def update_cellview(ax, system):  # pragma: no cover
     """Updates the particle positions visualisation pane.
@@ -673,6 +724,28 @@ def update_tempview(ax, system):  # pragma: no cover
         np.amax(system.temperature_sample) + np.amax(system.temperature_sample) * 0.05,
     )
 
+def update_maxbolthist(ax, system, velocities):  # pragma: no cover
+    """Updates the simulation velocity histogram visualisation pane.
+
+    Parameters
+    ----------
+    ax: Axes object
+        The axes position that the pane should be placed in.
+    system: System
+        The whole system information.
+    """
+    velocities = np.append(velocities, np.sqrt(np.square(system.particles['xvelocity']) + np.square(system.particles['yvelocity'])))
+    line = ax.lines[0]
+    hist, bin_edges = np.histogram(velocities, bins=25, density=True)
+    line.set_ydata(hist)
+    line.set_xdata(bin_edges[:-1])
+    ax.set_xlim(0, bin_edges[-2])
+    ax.set_ylim(
+        0,
+        np.amax(hist) + np.amax(hist) * 0.05,
+    )
+    return velocities
+
 
 def update_pressureview(ax, system):  # pragma: no cover
     """Updates the simulation instantaneous pressure visualisation pane.