Source code for ensembler.visualisation.animationSimulation

import numpy as np
from matplotlib import animation, pyplot as plt

from ensembler.util.ensemblerTypes import systemCls, Iterable, List, Tuple, Union, Number
from ensembler.visualisation import style, dpi_animation, animation_figsize
from ensembler.potentials.OneD import metadynamicsPotential


[docs]def animation_trajectory(simulated_system:systemCls, limits_coordinate_space: Tuple[float, float]=None, limits_potential_system_energy: Tuple[float, float]=None, resolution_of_analytic_potential:int=1000, title: str = None, out_path: str = None, out_writer: str = "pillow", dpi: int = dpi_animation, every_n_frame:int=1) -> Tuple[animation.Animation, Union[str, None]]: """ this function is generating a animation out of a simulation. Parameters ---------- simulated_system: systemCls the system, that carried out a simulation to be visualized limits_coordinate_space: Tuple[float, float], optional the plot limits of the coordinate space limits_potential_system_energy: Tuple[float, float], optional the plot limits of the potential energies title: str, optional the plot title out_path: str, optional the out path storing the animation. out_writer: str, optional the animation render. dpi: int, optional the resolution of the animation. every_n_frame: int, optional use every n-th frame Returns ------- Tuple[animation.Animation, Union[str, None]] returns the animation object and the output path if saved. """ # plotting x1data = simulated_system.trajectory.position y1data = simulated_system.trajectory.total_potential_energy shift = simulated_system.trajectory.dhdpos lastindx = x1data.shape[-1] - 1 lastindy = np.array(list(y1data)).shape[-1] - 1 x_max = max(x1data) x_min = min(x1data) if (limits_coordinate_space is None): xtot_space = np.linspace(x_min + 0.2 * x_min, x_max + 0.2 * x_max + 1, resolution_of_analytic_potential) else: xtot_space = np.linspace(min(limits_coordinate_space), max(limits_coordinate_space) + 1, resolution_of_analytic_potential) ytot_space = simulated_system.potential.ene(xtot_space) #settings step_size = every_n_frame tmax = len(y1data) - 1 t0 = 0 active_dots = 20 # build figure fig = plt.figure(dpi=60, figsize=animation_figsize) ax = fig.add_subplot(111) ## setup static parts ax.plot(xtot_space, ytot_space, label="Potential", c=style.potential_light) if(isinstance(simulated_system.potential, metadynamicsPotential)): #for metadynamics, show original potential ax.plot(xtot_space, simulated_system.potential.origPotential.ene(xtot_space), c="k", alpha=style.alpha_val, zorder=10, label="original Potential") ### Params if (limits_coordinate_space != None): ax.set_xlim(limits_coordinate_space) if (limits_potential_system_energy != None): ax.set_ylim(limits_potential_system_energy) ax.set_xlabel("$r$") ax.set_ylabel("$V$") if (title != None): fig.suptitle(title) # data structures in ani ##setup data structures xdata, ydata = [], [] scatter = ax.scatter([], [], c=[], vmin=0, vmax=1, cmap=style.animation_traj) start_p, = ax.plot([], [], "bo", c=style.traj_start, ms=10) end_p, = ax.plot([], [], "bo", c=style.traj_end, ms=10) curr_p, = ax.plot([], [], "bo", c=style.traj_current, ms=10) def init(): del xdata[:], ydata[:] start_p.set_data(x1data[0], y1data[0]) end_p.set_data([], []) curr_p.set_data([], []) def data_gen(t=t0): while t < tmax: t += step_size yield x1data[t], y1data[t] def run(data): x, V = data if (x == x1data[lastindx]): curr_p.set_data([], []) end_p.set_data(x1data[lastindx], y1data[lastindy]) else: curr_p.set_data([x], [V]) xdata.append(x) ydata.append(V) #color fading effect if (len(xdata) > active_dots): c = np.concatenate( (np.array([0.6 for x in range(len(xdata) - active_dots)]), np.linspace(0.6, 0, active_dots))) else: c = np.linspace(0.6, 0, len(xdata)) #set new data scatter.set_offsets(np.c_[xdata, ydata]) scatter.set_array(c) #if necessary adapt y axis if (min(ax.get_ylim()) > V): ax.set_ylim(V + V * 0.1, max(ax.get_ylim())) return scatter, ani = animation.FuncAnimation(fig=fig, func=run, frames=data_gen, init_func=init, blit=False, interval=20, repeat=False, cache_frame_data = True) if (out_path != None): # Set up formatting for the movie files Writer = animation.writers[out_writer] writer = Writer(metadata=dict(artist='animationsMD1D_David_Hahn_Benjamin_Schroeder')) ani.save(out_path, writer=writer, dpi=dpi) return ani, out_path
[docs]def animation_EDS_trajectory(system: systemCls, limits_coordinate_space=None, limits_potential_system_energy: Tuple[float, float]=None, title: str = None, out_path: str = None, hide_legend: bool = True, s_values:List[float] = [1.0], step_size: float = 1, out_writer: str = "pillow", dpi: int = 100, total_potential_resolution_points: int = 100) -> (animation.Animation, (str or None)): """ this function is generating a animation out of an eds-simulation. Parameters ---------- system: systemCls the system, that carried out a simulation to be visualized limits_coordinate_space: Tuple[float, float], optional the plot limits of the coordinate space limits_potential_system_energy: Tuple[float, float], optional the plot limits of the potential energies title: str, optional the plot title out_path: str, optional the out path storing the animation. hide_legend: bool, optional should the plot legend be hidden? s_values: List[float], optional s-values for the eds-potential step_size: int, optional size of the steps trough the trajectory out_writer: str, optional the animation render. dpi: int, optional the resolution of the animation. total_potential_resolution_points: int, optional the number of points for the visualization of the analytical reference potential. Returns ------- Tuple[animation.Animation, Union[str, None]] returns the animation object and the output path if saved. """ # plotting x1data = np.array(system.trajectory.position) y1data = system.trajectory.totPotEnergy x_max = max(x1data) x_min = min(x1data) active_dots = 20 if (limits_coordinate_space is None): xtot_space = np.array(np.arange(x_min + 0.2 * x_min, x_max + 0.2 * x_max + 1), ndmin=1) else: xtot_space = np.array(np.linspace(min(limits_coordinate_space), max(limits_coordinate_space) + 1, total_potential_resolution_points), ndmin=1) tmax = len(y1data) - 1 - step_size t0 = 0 xdata, ydata = [], [] # figures fig = plt.figure(figsize=animation_figsize) ax = fig.add_subplot(111) from ensembler.visualisation.plotPotentials import envPot_differentS_overlay_plot _, ax = envPot_differentS_overlay_plot(eds_potential=system.potential, s_values=s_values, title=title, positions=xtot_space, axes=ax, hide_legend=hide_legend) scatter = ax.scatter([], [], c=[], vmin=0, vmax=1, cmap='inferno') start_p, = ax.plot([], [], "bo", c="g", ms=10) end_p, = ax.plot([], [], "bo", c="r", ms=10) curr_p, = ax.plot([], [], "bo", c="k", ms=10) # Params ax.set_xlabel("$r$") ax.set_ylabel("$V$") if (isinstance(title, type(None))): fig.suptitle(title) def init(): del xdata[:], ydata[:] end_p.set_data([], []) curr_p.set_data([], []) if (limits_coordinate_space != None): ax.set_xlim(limits_coordinate_space) # return line, def data_gen(t=t0): while t < tmax: t += step_size yield x1data[t], y1data[t] def run(data): # update the data x, V = data if (type(x) == type(x1data[-1]) == list and all( [xi == x1i for xi, x1i in zip(x, x1data[-1])])): # last step of traj curr_p.set_data([], []) end_p.set_data(x1data[-1], y1data[-1]) elif (type(x) == type(x1data[-1]) == Number and x == x1data[-1]): # last step of traj curr_p.set_data([], []) end_p.set_data(x1data[-1], y1data[-1]) else: curr_p.set_data([x], [V]) xdata.append(x) if (not isinstance(x, Iterable)) else xdata.append(x[0]) ydata.append(V) if (len(xdata) > active_dots + 10): c = np.concatenate( (np.array([0.6 for x in range(len(xdata) - active_dots)]), np.linspace(0.6, 0, active_dots))) else: c = np.linspace(0.6, 0, len(xdata)) scatter.set_offsets(np.c_[xdata, ydata]) scatter.set_array(c) if (min(ax.get_ylim()) > V): ax.set_ylim(V + V * 0.1, max(ax.get_ylim())) ani = animation.FuncAnimation(fig=fig, func=run, frames=data_gen, init_func=init, blit=False, interval=20, repeat=False, save_count=len(x1data)) if (out_path != None): # Set up formatting for the movie files Writer = animation.writers[out_writer] writer = Writer(fps=15, metadata=dict(artist='animationsMD1D_David_Hahn_Benjamin_Schroeder'), bitrate=1800) ani.save(out_path, writer=writer, dpi=dpi) return ani, out_path