def plot_power_spectrum(power_spectrum, ax=None, show=False, color_scheme='jmol', **kwargs): if ax is None: fig = plt.figure(**kwargs) ax = fig.add_subplot(1,1,1) attrs = power_spectrum.get_attrs() ax.set_xlabel(r'$\omega [GHz]$') ax.set_ylabel(r'Signal $[\AA^2 fs^{-1}]$') species_of_interest = attrs['species_of_interest'] nr_of_trajectories = attrs['nr_of_trajectories'] frequencies = [power_spectrum.get_array('frequency_{}'.format(itraj)) for itraj in range(nr_of_trajectories)] for index_of_species, atomic_species in enumerate(species_of_interest): color = get_color(atomic_species, scheme=color_scheme) for itraj in range(nr_of_trajectories): freq = frequencies[itraj] periodogram = power_spectrum.get_array('periodogram_{}_{}'.format( atomic_species, itraj)) for signal in periodogram: ax.plot(freq, signal,color=color, alpha=0.1) try: periodogram_mean = power_spectrum.get_array('periodogram_{}_mean'.format( atomic_species)) periodogram_sem = power_spectrum.get_array('periodogram_{}_sem'.format( atomic_species)) ax.plot(freq, signal,color=color, alpha=1, linewidth=1) ax.fill_between(freq, periodogram_mean-periodogram_sem, periodogram_mean+periodogram_sem, facecolor=color, alpha=.2, linewidth=1) except Exception as e: print(e) if show: plt.show()
def plot_msd_isotropic(msd, ax=None, no_legend=False, species_of_interest=None, show=False, label=None, no_label=False, alpha_fill=0.2, alpha_block=0.3, alpha_fit=0.4, color_scheme='jmol', exclude_from_label=None, color_dict={}, decimals=1, **kwargs): if ax is None: fig = plt.figure(**kwargs) ax = fig.add_subplot(1,1,1) attrs = msd.get_attrs() if attrs['decomposed']: raise NotImplementedError("Plotting decomposed trajectories is not implemented") nr_of_trajectories = attrs['nr_of_trajectories'] t_start_fit_dt = attrs['t_start_fit_dt'] stepsize = attrs.get('stepsize_t', 1) timestep_fs = attrs['timestep_fs'] ax.set_ylabel(r'$\mathrm{MSD}(t)$ $\left( \mathrm{\AA}^2 \right) $ ') ax.set_xlabel(r'$t$ $\left( \mathrm{ps}\right)$') times_msd = timestep_fs/1000.0*stepsize*np.arange( attrs.get('t_start_dt')/stepsize, attrs.get('t_end_dt')/stepsize ) times_fit = timestep_fs/1000.0*stepsize*np.arange( attrs.get('t_start_fit_dt')/stepsize, attrs.get('t_end_fit_dt')/stepsize ) if species_of_interest is None: species_of_interest = attrs['species_of_interest'] for index_of_species, atomic_species in enumerate(species_of_interest): diff = attrs[atomic_species]['diffusion_mean_cm2_s'] diff_sem = attrs[atomic_species]['diffusion_sem_cm2_s'] diff_std = attrs[atomic_species]['diffusion_std_cm2_s'] if atomic_species in color_dict: color = color_dict[atomic_species] else: color = get_color(atomic_species, scheme=color_scheme) msd_mean = msd.get_array('msd_isotropic_{}_mean'.format(atomic_species)) msd_sem = msd.get_array('msd_isotropic_{}_sem'.format(atomic_species)) p1 = ax.fill_between( times_msd, msd_mean-msd_sem, msd_mean+msd_sem, facecolor=color, alpha=alpha_fill, linewidth=1, ) if no_label or (exclude_from_label and atomic_species in exclude_from_label): label_this_species = None elif label is None: # label_this_species=r'$D_{{{}}}=\left( {} \pm {} \right) \frac{{cm^2}}{{s}}$'.format( # atomic_species, my_format(diff, decimals=decimals), my_format(diff_sem, decimals=decimals)) label_this_species=r'$D_{{{}}}={} \, \frac{{cm^2}}{{s}}$'.format( atomic_species, format_mean_err(diff, diff_sem, decimals=decimals)) else: label_this_species = '{}'.format(label) if attrs.get('do_long', False): # reduce number of lines in plot, customize for later! # Keep the legend though! ax.plot([],[], color=color, linewidth=3., label=label_this_species) else: ax.plot(times_msd,msd_mean, color=color, linewidth=3., label=label_this_species) for itraj in range(nr_of_trajectories): msd_this_traj = msd.get_array('msd_isotropic_{}_{}'.format(atomic_species, itraj)) slopes_intercepts_this_traj = msd.get_array('slopes_intercepts_isotropic_{}_{}'.format(atomic_species, itraj)) for iblock in range(len(msd_this_traj)): slope_this_block, intercept_this_block = slopes_intercepts_this_traj[iblock] ax.plot(times_msd, msd_this_traj[iblock], color=color, alpha=alpha_block,) ax.plot(times_fit, [1000.*slope_this_block*x+intercept_this_block for x in times_fit], color=color, linestyle='--', alpha=alpha_fit) if attrs.get('do_long', False): times_long = timestep_fs/1000.0*stepsize*np.arange(0, attrs['nr_of_t_long_list'][itraj]) ax.plot(times_long, msd.get_array('msd_long_{}_{}'.format(atomic_species, itraj)), color=color, linestyle='-', lw=3) if not(no_legend): leg = ax.legend(loc=2,labelspacing=0.01) leg.get_frame().set_alpha(0.) if show: plt.show()
def plot_vaf_isotropic(vaf, ax=None, no_legend=False, species_of_interest=None, show=False, color_scheme='jmol', **kwargs): from matplotlib.ticker import ScalarFormatter f = ScalarFormatter() f.set_powerlimits((-1,1)) if ax is None: fig = plt.figure(**kwargs) ax = fig.add_subplot(1,1,1) attrs = vaf.get_attrs() nr_of_trajectories = attrs['nr_of_trajectories'] t_start_fit_dt = attrs['t_start_fit_dt'] t_end_fit_dt = attrs['t_end_fit_dt'] stepsize = attrs.get('stepsize_t', 1) timestep_fs = attrs['timestep_fs'] ax.set_ylabel(r'VAF $\left[ \AA^2 fs^{-2} \right]$') ax.set_xlabel('Time $t$ [fs]') axes_D = ax.twinx() axes_D.yaxis.set_major_formatter(f) axes_D.set_ylabel(r"$\int^t_0 VAF(t') dt' \quad \left[ \frac{cm^2}{s} \right]$") maxy = 0 # to set reasonable ylimits for axes_D, I track the max diff by hand times = timestep_fs*stepsize*np.arange( attrs.get('t_start_dt')/stepsize, attrs.get('t_end_dt')/stepsize ) times_fit = timestep_fs*stepsize*np.arange( attrs.get('t_start_fit_dt')/stepsize, attrs.get('t_end_fit_dt')/stepsize ) if species_of_interest is None: species_of_interest = attrs['species_of_interest'] for index_of_species, atomic_species in enumerate(species_of_interest): diff = attrs[atomic_species]['diffusion_mean_cm2_s'] diff_sem = attrs[atomic_species]['diffusion_sem_cm2_s'] diff_std = attrs[atomic_species]['diffusion_std_cm2_s'] color = get_color(atomic_species, scheme=color_scheme) vaf_mean = vaf.get_array('vaf_isotropic_{}_mean'.format(atomic_species)) vaf_sem = vaf.get_array('vaf_isotropic_{}_sem'.format(atomic_species)) vaf_integral_mean = vaf.get_array('vaf_integral_isotropic_{}_mean'.format(atomic_species)) vaf_integral_sem = vaf.get_array('vaf_integral_isotropic_{}_sem'.format(atomic_species)) ax.fill_between( times, vaf_mean-vaf_sem, vaf_mean+vaf_sem, facecolor=color, alpha=.2, linewidth=1, ) ax.plot(times,vaf_mean, color=color, linewidth=3., label=r'VAF ({})'.format(atomic_species)) maxy = max((maxy, (vaf_integral_mean+vaf_integral_sem).max())) for itraj in range(nr_of_trajectories): vaf_this_traj = vaf.get_array('vaf_isotropic_{}_{}'.format(atomic_species, itraj)) vaf_integral_this_traj = vaf.get_array('vaf_integral_isotropic_{}_{}'.format(atomic_species, itraj)) maxy = max((maxy, vaf_integral_this_traj.max())) for iblock in range(len(vaf_this_traj)): ax.plot(times, vaf_this_traj[iblock], color=color, alpha=0.1,) axes_D.plot(times, vaf_integral_this_traj[iblock], color=color, alpha=0.1, linestyle='--',) axes_D.plot(times ,vaf_integral_mean, color=color, linewidth=3., linestyle='--', label=r'$D_{{{}}}^{{VAF}}=( {:.2e} \pm {:.2e}) \frac{{cm^2}}{{s}}$'.format(atomic_species, diff, diff_sem)) axes_D.fill_between(times, vaf_integral_mean-vaf_integral_sem, vaf_integral_mean+vaf_integral_sem, facecolor=color, alpha=.2, linewidth=1) axes_D.set_ylim(0,maxy) axes_D.axvline(t_start_fit_dt*timestep_fs*stepsize, color='grey', linewidth=2, alpha=0.2,) axes_D.axvline(t_end_fit_dt*timestep_fs*stepsize, color='grey', linewidth=2, alpha=0.2) if not(no_legend): leg = ax.legend(loc=4) leg.get_frame().set_alpha(0.) leg = axes_D.legend(loc=1) leg.get_frame().set_alpha(0.) if show: plt.show()