diff --git a/.github/workflows/cmake.yml b/.github/workflows/cmake.yml index e3e114773..e03937f19 100644 --- a/.github/workflows/cmake.yml +++ b/.github/workflows/cmake.yml @@ -103,5 +103,5 @@ jobs: # See https://cmake.org/cmake/help/latest/manual/ctest.1.html for more detail run: | copy ${{github.workspace}}/build/${{env.BUILD_TYPE}}/pyanalisi* . - pip install pytest pytest-regressions pandas matplotlib numpy scipy testbook k3d + pip install pytest pytest-regressions pandas matplotlib numpy scipy testbook k3d ipykernel pytest -sv . diff --git a/CMakeLists.txt b/CMakeLists.txt index b8def5961..79b0fc94e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -74,6 +74,15 @@ if ("${CMAKE_SYSTEM_NAME}" STREQUAL "Linux") endif() +function(TRY_ADD_CXX_COMPILER_LINKER_FLAG FLAG TARGET) + + set (_saved_CRL ${CMAKE_REQUIRED_LIBRARIES}) + set (CMAKE_REQUIRED_LIBRARIES ${FLAG}) + check_cxx_compiler_flag ("${FLAG}" ${TARGET}) + set (CMAKE_REQUIRED_LIBRARIES ${_saved_CRL}) + message(STATUS "checked ${FLAG}") +endfunction() + include(CheckCXXCompilerFlag) include(CheckCCompilerFlag) #set(SAFE_CMAKE_REQUIRED_LINK_OPTIONS "${CMAKE_REQUIRED_LINK_OPTIONS}") @@ -81,9 +90,9 @@ include(CheckCCompilerFlag) #CHECK_CXX_COMPILER_FLAG("-fsanitize=address" COMPILER_HAS_SANITIZER_ADDR) #set(CMAKE_REQUIRED_LINK_OPTIONS "${SAFE_CMAKE_REQUIRED_LINK_OPTIONS}") #unset(SAFE_CMAKE_REQUIRED_LINK_OPTIONS) -CHECK_CXX_COMPILER_FLAG(" -fsanitize=address " COMPILER_HAS_SANITIZER_ADDR) -CHECK_CXX_COMPILER_FLAG(" -fsanitize=leak " COMPILER_HAS_SANITIZER_LEAK) -CHECK_CXX_COMPILER_FLAG(" -fsanitize=undefined " COMPILER_HAS_SANITIZER_UNDEF) +TRY_ADD_CXX_COMPILER_LINKER_FLAG(" -fsanitize=address " COMPILER_HAS_SANITIZER_ADDR) +TRY_ADD_CXX_COMPILER_LINKER_FLAG(" -fsanitize=leak " COMPILER_HAS_SANITIZER_LEAK) +TRY_ADD_CXX_COMPILER_LINKER_FLAG(" -fsanitize=undefined " COMPILER_HAS_SANITIZER_UNDEF) #CHECK_CXX_COMPILER_FLAG(" -ffpe-trap=invalid,overflow " COMPILER_HAS_FTRAP) if (COMPILER_HAS_SANITIZER_ADDR) set(COMPILER_SANITIZE_FLAGS " ${COMPILER_SANITIZE_FLAGS} -fsanitize=address") @@ -176,6 +185,7 @@ set (ANALISI_LIB lib/src/calcoliblocchi.cpp lib/src/convertibinario.cpp lib/src/convolution.cpp + lib/src/correlatorespaziale.cpp lib/src/cronometro.C lib/src/doubleround.cpp lib/src/gofrt.cpp diff --git a/analisi/main.cpp b/analisi/main.cpp index a4c01a28b..f37f412ee 100644 --- a/analisi/main.cpp +++ b/analisi/main.cpp @@ -50,6 +50,9 @@ #include "blockaverage.h" +#include +#include + namespace std{ template @@ -168,6 +171,7 @@ int main(int argc, char ** argv) std::vector factors_input; std::vector headers,output_conversion_gro; std::vector< std::pair > cvar; + std::array< size_t, 3 > start_stop_skip; options.add_options() ("input,i",boost::program_options::value(&input)->default_value(""), "input file in binary LAMMPS format: id type xu yu zu vx vy vz") @@ -205,7 +209,7 @@ int main(int argc, char ** argv) ("subBlock,k",boost::program_options::value(&n_seg)->default_value(1),"optimization option for the green-kubo calculation. It is the number of sub-blocks when calculating averages inside the big block used to calculate the variance. The performance depends on the particular system used.") ("kk",boost::program_options::bool_switch(&bench)->default_value(false),"small benchmark to find a good value of -k") ("kk-range",boost::program_options::value >(&kk_l)->multitoken(),"range where the -k benchmark has to be performed") - ("binary-convert",boost::program_options::value(&output_conversion),"perform the conversion from txt to binary LAMMPS trajectory file. The name of the output binary file is specified here") + ("binary-convert",boost::program_options::value(&output_conversion),"The name of the output binary file is specified here. If given alone, perform the conversion from txt to binary LAMMPS trajectory file. If --cut is specified, the input is a different binary trajectory.") #ifdef XDR_FILE ("binary-convert-gromacs",boost::program_options::value>(&output_conversion_gro)->multitoken(),"\ perform the conversion from gromacs format to the LAMMPS binary. Here you have to specify the output LAMMPS binary and the input type file with format:\n id0 type0\n ...\nidN typeN\nwith atomic types in the same order of the trr gromacs file. The trr gromacs file is specified with -i.") @@ -215,6 +219,7 @@ int main(int argc, char ** argv) ("spherical-harmonics-correlation,Y",boost::program_options::value(&sph)->default_value(0),"perform the calculation of the correlation function of the atomic density expanded in spherical harmonics. Note that this is a very heavy computation, first do a small test (for example with a very high -S or a low -s value). Here you have to specify the number of different bins of the considered radial distances, specified with -F. The code will calculate a correlation function for each bin.") ("buffer-size",boost::program_options::value(&buffer_size)->default_value(30),"Buffer size for sh frame values. This is machine dependend and can heavily influence the performance of the code") ("lt",boost::program_options::value (&read_lines_thread)->default_value(200),"parameter to read faster the time series column formatted txt file. It specifies the number of lines to read in one after the other for each thread") + ("cut",boost::program_options::value(&start_stop_skip)->multitoken(),"Specify [start stop skip] as parameters. Write a new binary trajectory by taking all the timesteps between index start and stop that satisfy (idx-start)%skip==0. Index of the first timestep is 0.") #ifdef EXPERIMENTAL ("spatial-correlator,A",boost::program_options::value(&nk)->default_value({0,0.0})->multitoken(),"Numero di punti della griglia ...") ("spatial-correlator-dir",boost::program_options::value(&kdir)->default_value({1.0,0.0,0.0})->multitoken(),"Direzione di k") @@ -238,14 +243,26 @@ int main(int argc, char ** argv) boost::program_options::notify(vm); - - if (argc<=1 || ( (output_conversion!="" || output_conversion_gro.size()>0 ) && input=="") ||vm.count("help")|| (vm.count("loginput")==0 && ( debug2 || heat_coeff ) ) || skip<=0 || stop_acf<0 || final<0 || (!sub_mean && (sub_mean_start!=0) ) || sub_mean_start<0 || !(kk_l.size()==0 || kk_l.size()==2)){ + std::cerr <<"cm.count " << vm.count("cut") <0 ) && input=="") + || ( output_conversion=="" && vm.count("cut") > 0 ) + || vm.count("help") + || (vm.count("loginput")==0 && ( debug2 || heat_coeff ) ) + || skip<=0 + || stop_acf<0 + || final<0 + || (!sub_mean && (sub_mean_start!=0) ) + || sub_mean_start<0 + || !(kk_l.size()==0 || kk_l.size()==2) + ){ std::cout << options << "\n"; if (vm.count("help")) return 0; return 1; } + if (cvar_list.size()%2 != 0) { std::cout << "Error: covariance indexes list must contain an even number of elements\n"; std::cout << options << "\n"; @@ -269,7 +286,7 @@ int main(int argc, char ** argv) try { - if (output_conversion!="") { + if (output_conversion!="" && vm.count("cut") == 0 ) { ConvertiBinario conv(input,output_conversion); return 0; @@ -285,6 +302,46 @@ int main(int argc, char ** argv) } } + if (vm.count("cut") > 0 ) { + if (start_stop_skip[2]==0) { + std::cerr << "Error: you must specify a positive skip"< void index_all(); int * get_lammps_id(); int *get_lammps_type(); + /* + * write a new binary file by copying the timesteps starting at timestep index start to + * timestep index stop. This is the closed interval [start, stop] taking one every skip + * steps. Starts readint at offset and return the offset of the next timestep in the file + * constantly updates written to the number of timesteps written in the new file. + * */ + size_t dump_every(size_t offset, const size_t start, const size_t stop, const size_t skip, const char * fname, size_t & written); private: std::mapid_map; size_t * timesteps; // puntatori (offset rispetto all'inizio) all'inizio di ogni timesteps diff --git a/lib/src/trajectory.cpp b/lib/src/trajectory.cpp index 08eacf564..5a885f057 100644 --- a/lib/src/trajectory.cpp +++ b/lib/src/trajectory.cpp @@ -37,7 +37,7 @@ #include #include #include "macros.h" - +#include @@ -316,6 +316,23 @@ size_t Trajectory::leggi_pezzo(const size_t &partenza /// offset da cui partire } +size_t Trajectory::dump_every(size_t offset, const size_t start, const size_t stop, const size_t skip, const char* fname, size_t & written) { + std::ofstream outfile(fname,std::ofstream::binary); + size_t timestep_size = 0; + written=0; + for (size_t i = 0; i<=stop && i(offset, header,ch); + if (i>=start && i<=stop && (i-start)%skip==0 ) { + outfile.write(file+offset,timestep_size); + written+=1; + } + offset+=timestep_size; + } + return offset; +} + /** * Sistema i puntatori con i dati del timestep e * restituisce la dimensione in byte del pezzo letto, comprensiva dei chunk con i dati veri e propri che però non vengono restituiti diff --git a/notebooks/calc_inspector.ipynb b/notebooks/calc_inspector.ipynb index f83fb6d35..7391f5399 100644 --- a/notebooks/calc_inspector.ipynb +++ b/notebooks/calc_inspector.ipynb @@ -6,7 +6,7 @@ "metadata": {}, "outputs": [], "source": [ - "import pyanalisi as pa\n", + "import pyanalisi as pa\n", "import numpy as np" ] }, @@ -95,7 +95,7 @@ }, "outputs": [], "source": [ - "density=pa.atomic_density(traj)\n", + "density=pa.atomic_density(pa.trajectory.Trajectory(traj))\n", "pa.density_field(*density)" ] }, @@ -127,6 +127,13 @@ "source": [ "shp=pa.plot_sh(0.7,1.4,times,sh,1,0,0,log=False,pre_fit=0.4)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { @@ -145,7 +152,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.12" + "version": "3.11.0" } }, "nbformat": 4, diff --git a/notebooks/example2.ipynb b/notebooks/example2.ipynb index 859b58464..ce0e252d4 100644 --- a/notebooks/example2.ipynb +++ b/notebooks/example2.ipynb @@ -26,7 +26,7 @@ "metadata": {}, "outputs": [], "source": [ - "phase_IV=read_lammps_bin('./dump2.8fs.bin',wrap=False)" + "phase_IV=read_lammps_bin('../tests/data/lammps.bin',wrap=False,dt=1.0)" ] }, { @@ -44,7 +44,7 @@ "metadata": {}, "outputs": [], "source": [ - "plt.plot(phase_IV.get_array('positions')[:,phase_IV.t.get_lammps_type()==1,0][:,:10])" + "plt.plot(phase_IV.get_array('positions')[:,phase_IV.get_analisi_traj().get_lammps_type()==1,0][:,:10])" ] }, { @@ -53,7 +53,7 @@ "metadata": {}, "outputs": [], "source": [ - "pl=show_traj(phase_IV,wrap=False)" + "pl=phase_IV.show_traj()" ] }, { @@ -66,7 +66,7 @@ "source": [ "res=multiinspect([phase_IV], plot=True, prefix='phase_IV_nose_',\n", " inspect_kw={\n", - " 'nthreads':12,\n", + " 'nthreads':4,\n", " 'do_sh':False,\n", " 'plot_st_kw': {\n", " 'transpose':True,\n", @@ -99,7 +99,7 @@ }, "outputs": [], "source": [ - "steinhardt_movie(phase_IV,skip=20,n_segments=150,\n", + "html,_=steinhardt_movie(phase_IV,skip=20,n_segments=15,\n", " plt_steinhardt_kw={\n", " 'transpose':True,\n", " 'xmax':.40,\n", @@ -107,13 +107,21 @@ " 'single':(1,1)},\n", " compute_steinhardt_kw={'nthreads':12},\n", " neigh=[(55,3.5**2,0.0),(65,3.5**2,0.0)]\n", - " )" + " )\n", + "html" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] } ], "metadata": { "kernelspec": { - "display_name": "Python 3", + "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, @@ -127,7 +135,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.7.3" + "version": "3.11.0" } }, "nbformat": 4, diff --git a/pyanalisi/__init__.py b/pyanalisi/__init__.py index 29173438e..fc908b88f 100644 --- a/pyanalisi/__init__.py +++ b/pyanalisi/__init__.py @@ -1,2 +1,2 @@ -from pyanalisi.pyanalisi import * -from pyanalisi.common import * +from .pyanalisi import * +from .common import * diff --git a/pyanalisi/analysis.py b/pyanalisi/analysis.py new file mode 100644 index 000000000..d3a265619 --- /dev/null +++ b/pyanalisi/analysis.py @@ -0,0 +1,216 @@ +import numpy as np +import pyanalisi as pa + +from .plotters import * + +class Analysis: + + @staticmethod + def pyanalisi_wrapper(Class, traj, *args): + if isinstance(traj, pa.Trajectory): + return getattr(pa, Class)(traj, *args) + elif isinstance(traj, pa.Traj): + return getattr(pa, Class + '_lammps')(traj, *args) + else: + raise RuntimeError(f"Wrapper for trajectory class '{traj.__name__}' not implemented") + + def __init__(self, trajectory,nthreads=4): + self.traj = trajectory + self.nthreads=nthreads + + def compute_steinhardt(self, ranges=None, nthreads=None, skip=10, + neigh=None, histogram=True, ls=[6, 4], averaged=False, n_segments=1, nbins=100, tskip=1): + """ + ranges : list of (min_r, max_r) radial distances where the code will compute the steinhardt parameters + for each possible pair of types (total n_types**2 elements) + neigh : list of (n_max, r_max**2, 0.0 ) neighbor list specifications for each possible pair of types. Set to use the SANN algorithm. + """ + if nthreads is None: + nthreads = self.nthreads + if ranges is None: + ranges = [ (0.0,2.0) for i in range(len(self.traj.get_atomic_types())**2) ] + atraj = self.traj.get_analisi_traj(tskip=tskip,wrapped=True) + if neigh is None: + neig=[] + stein = None + l = sorted(ls) + + if l[0] < 0: + raise IndexError(f'minimum l cannot be negative {l}') + elif l[-1] > 10: + raise IndexError( + f'spherical harmonics with l>10 are not supported {l}. Please modify and recompile source code.') + + stein = Analysis.pyanalisi_wrapper(f'SteinhardtOrderParameterHistogram_{l[-1]}', atraj, + ranges, + 1, nbins, + l, + nthreads, skip, histogram, neigh, averaged + ) + + if n_segments == 1: + stein.reset(atraj.get_nloaded_timesteps()) + stein.calculate(0) + stein_res = np.array(stein) + return stein_res + elif n_segments > 1: + segment_size = max(1, atraj.get_nloaded_timesteps() // n_segments) + stein.reset(segment_size) + #print(segment_size) + res = [] + for i in range(0, segment_size * n_segments, segment_size): + #print(f'calculating {i}...') + stein.calculate(i) + res.append(np.array(stein, copy=True)) + return np.array(res) + else: + raise IndexError('n_segments must be >= 1') + + def steinhardt_movie(self, skip=1,averaged=True,n_segments=60,plt_steinhardt_kw=None, compute_steinhardt_kw=None, plt_steinhardt_init_add_kw=None): + if compute_steinhardt_kw is None: + compute_steinhardt_kw={} + compute_steinhardt_kw.setdefault('n_segments',n_segments) + tstein = self.compute_steinhardt(**compute_steinhardt_kw) + return SteinPlot.animation(tstein,plt_steinhardt_kw=plt_steinhardt_kw,plt_steinhardt_init_add_kw=plt_steinhardt_init_add_kw) + + @staticmethod + def max_l(start, stop, tmax=0): + if tmax <= 0: + tmax = (stop - start) // 2 + n_ave = stop - start - tmax + if start >= stop: + raise RuntimeError("start index must be less than the end index") + if n_ave <= 0: + tmax = stop - start - 1 + n_ave = 1 + return tmax, n_ave + + + def compute_msd(self,start=0,stop=-1,tmax=0, tskip_msd=10,center_of_mass_MSD=True, center_of_mass_frame=False): + if stop <0: + stop = self.traj.numsteps + tmax, n_ave = Analysis.max_l(start, stop, tmax) + # mean square displacement + with Analysis.pyanalisi_wrapper('MeanSquareDisplacement', self.traj.get_analisi_traj(wrapped=False), tskip_msd, tmax, self.nthreads, center_of_mass_MSD, center_of_mass_frame, False) as msd: + msd.reset(n_ave) + msd.calculate(start) + res = np.array(msd,copy=True) + return res + + @staticmethod + def hist2gofr(gr_N, gr_dr, gr_0, gofr): + rs_m = np.arange(gr_N) * gr_dr + gr_0 + rs_p = (np.arange(gr_N) + 1) * gr_dr + gr_0 + vols = 4 * np.pi / 3 * (rs_p ** 3 - rs_m ** 3) + return gofr / vols + + + def compute_gofr(self, startr, endr, nbin,start=0,stop=None, tmax=1, tskip=10, n_segments=1, nthreads=None,return_histogram=False): + """ + startr -> first r of the g(r) histogram + endr -> maximum r of the g(r) + nbin -> number of bins on the radial g(r) histogram + tmax -> number of time lags ( if > 1, computes van hove function g(r,t) ) + """ + #TODO: uniform usage of tskip and skip + traj = self.traj.get_analisi_traj(tskip=1,wrapped=True) + if stop is None: + stop=traj.get_nloaded_timesteps() + if nthreads is None: + nthreads = self.nthreads + tmax, n_ave = Analysis.max_l(start, stop, tmax) + gofr = Analysis.pyanalisi_wrapper('Gofrt', traj, startr, endr, nbin, tmax, nthreads, tskip, False, 1) + if n_segments == 1: + gofr.reset(n_ave) + #print('calculating g(r)...', flush=True) + gofr.calculate(start) + if return_histogram: + return np.array(gofr ,copy=True) + else: + return Analysis.hist2gofr(nbin,(endr-startr)/nbin, startr, np.array(gofr, copy=True)) + elif n_segments > 1: + res = [] + segment_size = max(1, n_ave // n_segments) + gofr.reset(segment_size) + #print(segment_size) + for i in range(0, min(segment_size * n_segments, n_ave), segment_size): + gofr.calculate(start + i) + if return_histogram: + res.append(np.array(gofr, copy=True)) + else: + res.append( Analysis.hist2gofr(nbin,(endr-startr)/nbin, startr, np.array(gofr, copy=True) ) ) + return res + else: + raise IndexError(f'n_segments must be > 0 ({n_segments})') + + #### stuff to find peaks of g(r) + + @staticmethod + def compute_peaks_gofr(gr, r0, gr_0, gr_dr, peak_fractional_height): + ''' + find the peak around r0 in the g(r) array and get some info + gr -> g(r) + gr_0 -> starting value of r in g(r) array + gr_dr -> size of the bin of the histogram + peak_fractional_height -> threshold to trigger the peak width algorithm, in fractions of peak height + ''' + def peak_width(gr, NH, NH_thre, gr_r2i, min_spread=0.1): + ''' + get the width of the first peak in the gr at a given height + ''' + idx_23 = 0 + while gr[0, NH, idx_23] < NH_thre: + idx_23 += 1 + idx_spread_low = idx_23 + while gr[0, NH, idx_23] > NH_thre * 0.9: + if idx_23 == gr.shape[2] - 1 or (idx_23 - idx_spread_low > gr_r2i(min_spread) and gr[0, NH, idx_23] < NH_thre): + break + idx_23 += 1 + idx_spread_hi = idx_23 + return idx_spread_low, idx_spread_hi + def get_conv_functs(gr_0, gr_dr): + def gr_r2i(r): + return int((r - gr_0) / gr_dr) + + def gr_i2r(i): + return gr_0 + i * gr_dr + + return np.vectorize(gr_r2i), np.vectorize(gr_i2r) + gr_r2i, gr_i2r = get_conv_functs(gr_0, gr_dr) + maxs = np.argmax(gr, axis=2) + peaks = gr_i2r(maxs) + + # find nearest peak to r0 + NH = np.argmin((peaks - r0) ** 2) + # get spread at 'peak_fractional_height' of peak height + NH_peak_idx = maxs[0, NH] + NH_peak_val = gr[0, NH, NH_peak_idx] + NH_thre = NH_peak_val * peak_fractional_height + # get width of peak + idx_spread_low, idx_spread_hi = peak_width(gr, NH, NH_thre, gr_r2i) + spread_low = gr_i2r(idx_spread_low) + spread_hi = gr_i2r(idx_spread_hi) + w23h = spread_hi - spread_low + return idx_spread_low, idx_spread_hi, spread_low, spread_hi, w23h, NH_thre, NH_peak_val, NH_peak_idx, peaks, NH + + + @staticmethod + def block_average_std(data,n_b=6): + ndp=data.shape[0]//n_b + means=[] + for i in range(n_b): + means.append(np.mean(data[i*ndp:(i+1)*ndp],axis=0)) + means=np.array(means) + std=(np.var(means,axis=0)/(n_b-1))**.5 + mean=np.mean(means,axis=0) + del means + return mean,std + + def block_average(self, key, nblocks=6): + arr=self.traj.get_array(key) + return Analysis.block_average_std(arr,n_b=nblocks) + + + + + diff --git a/pyanalisi/common.py b/pyanalisi/common.py index 157c48b83..312259933 100644 --- a/pyanalisi/common.py +++ b/pyanalisi/common.py @@ -4,25 +4,28 @@ try: import aiida - HAS_AIIDA=True -except: - HAS_AIIDA=False + + HAS_AIIDA = True +except ImportError: + HAS_AIIDA = False if HAS_AIIDA: try: - from aiida import load_profile - from aiida.orm import Code, load_node - from aiida.orm import * - from aiida.engine import submit - load_profile() + from aiida import load_profile + from aiida.orm import Code, load_node + from aiida.orm import * + from aiida.engine import submit + + load_profile() except: - HAS_AIIDA=False + HAS_AIIDA = False try: import aiida_QECpWorkChain - HAS_QECPWORKCHAIN=True + + HAS_QECPWORKCHAIN = True except: - HAS_QECPWORKCHAIN=False - + HAS_QECPWORKCHAIN = False + if HAS_QECPWORKCHAIN: from aiida_QECpWorkChain.workflow import * from aiida_QECpWorkChain import write_cp_traj @@ -32,172 +35,178 @@ except: print('WARNING: cannot import numpy') try: - import matplotlib as mpl - mpl.rcParams['agg.path.chunksize'] = 1000000 - import matplotlib.pyplot as plt - import matplotlib.colors as pltc - from mpl_toolkits.axes_grid1.inset_locator import inset_axes -except: - print('WARNING: cannot import matplotlib') + import matplotlib as mpl + mpl.rcParams['agg.path.chunksize'] = 1000000 + import matplotlib.pyplot as plt + import matplotlib.colors as pltc + from mpl_toolkits.axes_grid1.inset_locator import inset_axes +except: + print('WARNING: cannot import matplotlib') import pyanalisi.pyanalisi as pa print(pa.info()) -from matplotlib import collections as mc -from IPython.core.display import display, HTML +from matplotlib import collections as mc +from IPython.display import display +from IPython.core.display import HTML import matplotlib.animation +from .plotters import * +from .analysis import * + FIGURE_PATH = '.' + + def set_figure_path(new_path): global FIGURE_PATH FIGURE_PATH = new_path -DEFAULT_PLT_STEINHARDT_KW={'transpose':True,'xmax':.30,'ymax':.60} -DEFAULT_NEIGH=[(57,3.5**2,0.0),(45,3.5**2,0.0)] -DEFAULT_SUBPLOTS_KW={'figsize':(6,4),'dpi':300} -#aiida -def plt_key(traj,key,conv=1.0,title='',ylabel='',dt=None,subplots_kw=DEFAULT_SUBPLOTS_KW): + +DEFAULT_PLT_STEINHARDT_KW = {'transpose': True, 'xmax': .30, 'ymax': .60} +DEFAULT_NEIGH = [(57, 3.5 ** 2, 0.0), (45, 3.5 ** 2, 0.0)] +DEFAULT_SUBPLOTS_KW = {'figsize': (6, 4), 'dpi': 300} +# aiida + +# legacy import +from .trajectory import Trajectory as FakeAiidaT + +analisi_cell2box = FakeAiidaT.analisi_cell2box + + +def plt_key(traj, key, conv=1.0, title='', ylabel='', dt=None, subplots_kw=DEFAULT_SUBPLOTS_KW): if not key in traj.get_arraynames(): return - if title=='': title=key - fig,ax =plt.subplots(**subplots_kw) - ax=fig.add_axes([0,0,1,1]) - axins=inset_axes(ax,width='20%',height='20%',loc='upper right') + if title == '': title = key + fig, ax = plt.subplots(**subplots_kw) + ax = fig.add_axes([0, 0, 1, 1]) + axins = inset_axes(ax, width='20%', height='20%', loc='upper right') ax.set_title(title) ax.set_xlabel('time (ps)') ax.set_ylabel(ylabel) axins.set_title('100 points of "{}"'.format(title)) - q=traj.get_array(key) + q = traj.get_array(key) if dt is None: - t=traj.get_array('times') + t = traj.get_array('times') else: - t=np.arange(q.shape[0])*dt - ax.plot(t,q*conv) - axins.plot(t,q*conv) - x1=len(t)//2 - x2=len(t)//2+100 - if x2==x1: x2=len(t)-1 - if x2>=len(t): x2=len(t)-1 - axins.set_xlim(t[x1],t[x2]) + t = np.arange(q.shape[0]) * dt + ax.plot(t, q * conv) + axins.plot(t, q * conv) + x1 = len(t) // 2 + x2 = len(t) // 2 + 100 + if x2 == x1: x2 = len(t) - 1 + if x2 >= len(t): x2 = len(t) - 1 + axins.set_xlim(t[x1], t[x2]) return ax - def get_types_id_array(types_array): - types={} - typeid=0 - res=[] - for t in types_array: - if not t in types: - types[t]=typeid - typeid=typeid+1 - res.append(types[t]) - return np.array(res,dtype='int32') - -#aiida -def get_analisi_traj_from_aiida(traj,tskip=1): - if tskip < 1: - raise IndexError(f'cannot have a trajectory skip of {tskip}') - pos=traj.get_array('positions')[::tskip].copy(order='C') - if 'velocities' in traj.get_arraynames(): - vel=traj.get_array('velocities')[::tskip].copy(order='C') - else: - vel=np.zeros(pos.shape) - cel=traj.get_array('cells')[::tskip].transpose((0,2,1)).copy(order='C') - types=get_types_id_array(traj.get_attribute('symbols')) - params= [pos, vel, types, cel] - atraj=pa.Trajectory(*params,pa.BoxFormat.CellVectors, True,True) - atraj_unw=pa.Trajectory(*params,pa.BoxFormat.CellVectors, False,True) + return FakeAiidaT.get_types_id_array(types_array) + + +# aiida +def get_analisi_traj_from_aiida(traj, tskip=1): + t = FakeAiidaT(traj, traj_skip=tskip) + atraj = t.get_analisi_traj(wrapped=True) + atraj_unw = t.get_analisi_traj(wrapped=False) return atraj, atraj_unw -#only pos and cell -def get_analisi_traj(pos,types,cell,wrap=False): - return pa.Trajectory(pos,np.zeros(pos.shape),types,cell,pa.BoxFormat.CellVectors, wrap,True) + +# only pos and cell +def get_analisi_traj(pos, types, cell, wrap=False): + return pa.Trajectory(pos, np.zeros(pos.shape), types, cell, pa.BoxFormat.CellVectors, wrap, True) + + try: import k3d - HAS_K3D=True + + HAS_K3D = True except: - HAS_K3D=False + HAS_K3D = False try: from ipywidgets import interact, interactive, fixed import ipywidgets as widgets - HAS_IPYWIDGETS=True + + HAS_IPYWIDGETS = True except: - HAS_IPYWIDGETS=False - - -def pyanalisi_wrapper(Class,traj,*args): + HAS_IPYWIDGETS = False + + +def pyanalisi_wrapper(Class, traj, *args): if isinstance(traj, pa.Trajectory): - return getattr(pa,Class)(traj,*args) + return getattr(pa, Class)(traj, *args) elif isinstance(traj, pa.Traj): - return getattr(pa,Class+'_lammps')(traj,*args) + return getattr(pa, Class + '_lammps')(traj, *args) else: raise RuntimeError(f"Wrapper for trajectory class '{traj.__name__}' not implemented") - + + def analisi_to_cell_vectors(b): - boxes=[] + boxes = [] for i in range(b.shape[0]): - bb=b[i] - boxes.append([[bb[3]*2,bb[6],bb[7]],[0,bb[4]*2,bb[8]],[0,0,bb[5]*2]]) + bb = b[i] + boxes.append([[bb[3] * 2, bb[6], bb[7]], [0, bb[4] * 2, bb[8]], [0, 0, bb[5] * 2]]) return np.array(boxes) -def fft_density(res,coeff=None): + +def fft_density(res, coeff=None): if coeff is None: - coeff=np.ones(res.shape[0]) - resg=np.einsum('ijkl,i',np.fft.fftn(res,axes=(1,2,3)),coeff) - resg[0,0,0]=0 - resg=np.fft.fftshift(resg) - resgn=np.real(resg*np.conj(resg)) - resgn/=resgn.max() + coeff = np.ones(res.shape[0]) + resg = np.einsum('ijkl,i', np.fft.fftn(res, axes=(1, 2, 3)), coeff) + resg[0, 0, 0] = 0 + resg = np.fft.fftshift(resg) + resgn = np.real(resg * np.conj(resg)) + resgn /= resgn.max() return resgn -def max_l(start,stop,tmax=0): - if tmax<=0: - tmax=(stop-start)//2 - n_ave=stop-start-tmax - if start>=stop: - raise RuntimeError("start index must be less than the end index") - if n_ave<=0: - tmax=stop-start -1 - n_ave=1 - return tmax,n_ave +def max_l(start, stop, tmax=0): + if tmax <= 0: + tmax = (stop - start) // 2 + n_ave = stop - start - tmax + if start >= stop: + raise RuntimeError("start index must be less than the end index") + if n_ave <= 0: + tmax = stop - start - 1 + n_ave = 1 + return tmax, n_ave -def analyze_msd(traj_unw,start,stop,tmax=0,nthreads=4,tskip_msd=10,print=print): - tmax,n_ave = max_l(start,stop,tmax) - #mean square displacement - msd=pyanalisi_wrapper('MeanSquareDisplacement',traj_unw,tskip_msd,tmax,nthreads,True,False,False) +def analyze_msd(traj_unw, start, stop, tmax=0, nthreads=4, tskip_msd=10, print=print): + tmax, n_ave = max_l(start, stop, tmax) + # mean square displacement + msd = pyanalisi_wrapper('MeanSquareDisplacement', traj_unw, tskip_msd, tmax, nthreads, True, False, False) msd.reset(n_ave) - print('calculating msd...',flush=True) + print('calculating msd...', flush=True) msd.calculate(start) - return np.array(msd)#,copy=True) + return np.array(msd) # ,copy=True) + -def analyze_gofr(traj,start,stop,startr,endr,nbin,tmax=1,nthreads=4,tskip=10,print=print,n_segments=1): - tmax,n_ave = max_l(start,stop,tmax) - gofr=pyanalisi_wrapper('Gofrt',traj,startr,endr,nbin,tmax,nthreads,tskip,False,1) - if n_segments==1: +def analyze_gofr(traj, start, stop, startr, endr, nbin, tmax=1, nthreads=4, tskip=10, print=print, n_segments=1): + tmax, n_ave = max_l(start, stop, tmax) + gofr = pyanalisi_wrapper('Gofrt', traj, startr, endr, nbin, tmax, nthreads, tskip, False, 1) + if n_segments == 1: gofr.reset(n_ave) - print('calculating g(r)...',flush=True) + print('calculating g(r)...', flush=True) gofr.calculate(start) - return np.array(gofr)#,copy=True) - elif n_segments>1: + return np.array(gofr) # ,copy=True) + elif n_segments > 1: res = [] - segment_size=max(1,n_ave//n_segments) + segment_size = max(1, n_ave // n_segments) gofr.reset(segment_size) print(segment_size) - for i in range(0,min(segment_size*n_segments,n_ave),segment_size): - gofr.calculate(start+i) - res.append(np.array(gofr,copy=True)) + for i in range(0, min(segment_size * n_segments, n_ave), segment_size): + gofr.calculate(start + i) + res.append(np.array(gofr, copy=True)) return res else: raise IndexError(f'n_segments must be > 0 ({n_segments})') -def analyze_vdos_single(traj,nstep=None,print=print): + +def analyze_vdos_single(traj, nstep=None, print=print): ''' - Computes vdos + Computes vdos there are no start,stop options because spettrovibrazionale.calcola do not support the features Parameters ---------- @@ -205,19 +214,20 @@ def analyze_vdos_single(traj,nstep=None,print=print): nstep : int nstep to use for the computation of vdos if None it uses all return: - vdos: numpy array (ntypes , nstep/2+1,3) + vdos: numpy array (ntypes , nstep/2+1,3) ''' - if nstep is None or nstep>=traj.get_nloaded_timesteps(): - nstep = traj.get_nloaded_timesteps() + if nstep is None or nstep >= traj.get_nloaded_timesteps(): + nstep = traj.get_nloaded_timesteps() - vdos=pyanalisi_wrapper('VibrationSpectrum',traj,False) + vdos = pyanalisi_wrapper('VibrationSpectrum', traj, False) vdos.reset(nstep) - print('calculating vdos...',flush=True) + print('calculating vdos...', flush=True) vdos.calculate(0) - return np.array(vdos)#,copy=True) + return np.array(vdos) # ,copy=True) + -def analyze_vdos_lammps(traj,nstep=None,start=0,resetAccess=True,print=print): +def analyze_vdos_lammps(traj, nstep=None, start=0, resetAccess=True, print=print): ''' Computes vdos from a lammps trajectory Parameters @@ -230,39 +240,43 @@ def analyze_vdos_lammps(traj,nstep=None,start=0,resetAccess=True,print=print): resetAccess : bool at the end reset the access to nstep and first step return: - vdos: numpy array (ntypes , nstep/2+1,3) + vdos: numpy array (ntypes , nstep/2+1,3) ''' - if nstep is None or nstep>=traj.getNtimesteps(): - nstep = traj.getNtimesteps() - if start!=0: - start=0 - print('start setted to zero because all the timesteps are have been requested') - - #set the trajectory to what we want + if nstep is None or nstep >= traj.getNtimesteps(): + nstep = traj.getNtimesteps() + if start != 0: + start = 0 + print('start setted to zero because all the timesteps are have been requested') + + # set the trajectory to what we want traj.setAccessWindowSize(nstep) traj.setAccessStart(start) - vdos = analyze_vdos_single(traj,nstep=nstep,print=print) + vdos = analyze_vdos_single(traj, nstep=nstep, print=print) - #reset Access + # reset Access if (resetAccess): - traj.setAccessWindowSize(traj.getNtimesteps()) - traj.setAccessStart(0) - print('reset Access') - print('traj.setAccessWindowSize(traj.getNtimesteps())') - print('traj.setAccessStart(0)') - - return vdos #,copy=True) - -def vdos_norm_factor(dt,N): - return dt/N/6.0 -def vdos_domega(dt,N): - return 2.0/dt/N + traj.setAccessWindowSize(traj.getNtimesteps()) + traj.setAccessStart(0) + print('reset Access') + print('traj.setAccessWindowSize(traj.getNtimesteps())') + print('traj.setAccessStart(0)') + + return vdos # ,copy=True) + + +def vdos_norm_factor(dt, N): + return dt / N / 6.0 + + +def vdos_domega(dt, N): + return 2.0 / dt / N + def analyze_vdos_numpy(pos, vel, types, box, - matrix_format=True, # matrix format for the box array - wrap=False # don't wrap the coordinates - ,nstep=None,start=0, - print=print): + matrix_format=True, # matrix format for the box array + wrap=False # don't wrap the coordinates + , nstep=None, start=0, + print=print): ''' Computes vdos from numpy arrays Parameters @@ -270,55 +284,58 @@ def analyze_vdos_numpy(pos, vel, types, box, pos : positions matrix (N_timesteps, N_atoms, 3) vel : velocities matrix (N_timesteps, N_atoms, 3) types : types vector (N_atoms) - box : box matrix + box : box matrix matrix_format :bool matrix format for the box array - wrap : bool + wrap : bool wrap coordinate nstep : int nstep to use for the computation of vdos if None it uses all start : int initial step return: - vdos: numpy array (ntypes , nstep/2+1,3) + vdos: numpy array (ntypes , nstep/2+1,3) ''' totnstep = len(pos) - if nstep is None or nstep>=totnstep: - nstep = totnstep - if start!=0: - start=0 - print('start setted to zero because all the timesteps are have been requested') - - traj = pa.Trajectory(pos[start:start+nstep,:,:], vel[start:nstep+start,:,:], types, box[start:start+nstep,:,:], - matrix_format, # matrix format for the box array - wrap # don't wrap the coordinates - ) - #set the trajectory to what we want - vdos = analyze_vdos_single(traj,nstep=nstep,print=print) - - - return vdos #,copy=True) - -def analyze_sh(traj,start,stop,startr,endr, nbin,ntypes=2, tmax=0, nthreads=4,tskip=10,print=print, - sann=[], #list to tell how to build the neighbours list: [ (max_number_of_neigh,rcut**2,not_used) ] - buffer_size=20 - ): - tmax,n_ave = max_l(start,stop,tmax) - sh=pyanalisi_wrapper('SphericalCorrelations',traj - ,[(startr #minimum of radial distance - ,endr)]*ntypes**2 #maximum of radial distance - ,nbin # number of distance bins - ,tmax # maximum length of correlation functions in timesteps - ,nthreads # number of threads - ,tskip # time skip to average over the trajectory - ,buffer_size # buffer for sh - ,True,sann) #flag that at the moment does nothing, sann if it is not empty activates the sann algorithm. - sh.reset(n_ave) # number of timesteps to average - print('calculating spherical harmonics correlation functions... (go away and have a coffee)',flush=True) - sh.calculate(start) #calculate starting at this timestep - return np.array(sh)#,copy=True) - -def analyze(traj,traj_unw,start,stop,nthreads=4, + if nstep is None or nstep >= totnstep: + nstep = totnstep + if start != 0: + start = 0 + print('start setted to zero because all the timesteps are have been requested') + + traj = pa.Trajectory(pos[start:start + nstep, :, :], vel[start:nstep + start, :, :], types, + box[start:start + nstep, :, :], + matrix_format, # matrix format for the box array + wrap # don't wrap the coordinates + ) + # set the trajectory to what we want + vdos = analyze_vdos_single(traj, nstep=nstep, print=print) + + return vdos # ,copy=True) + + +def analyze_sh(traj, start, stop, startr, endr, nbin, ntypes=2, tmax=0, nthreads=4, tskip=10, print=print, + sann=[], # list to tell how to build the neighbours list: [ (max_number_of_neigh,rcut**2,not_used) ] + buffer_size=20 + ): + tmax, n_ave = max_l(start, stop, tmax) + sh = pyanalisi_wrapper('SphericalCorrelations', traj + , [(startr # minimum of radial distance + , endr)] * ntypes ** 2 # maximum of radial distance + , nbin # number of distance bins + , tmax # maximum length of correlation functions in timesteps + , nthreads # number of threads + , tskip # time skip to average over the trajectory + , buffer_size # buffer for sh + , True, + sann) # flag that at the moment does nothing, sann if it is not empty activates the sann algorithm. + sh.reset(n_ave) # number of timesteps to average + print('calculating spherical harmonics correlation functions... (go away and have a coffee)', flush=True) + sh.calculate(start) # calculate starting at this timestep + return np.array(sh) # ,copy=True) + + +def analyze(traj, traj_unw, start, stop, nthreads=4, tmax_msd=0, tskip_msd=10, startr_gofr=0.5, @@ -331,82 +348,85 @@ def analyze(traj,traj_unw,start,stop,nthreads=4, nbin_sh=1, tskip_sh=20, tmax_sh=0): - - msd=analyze_msd(traj_unw,start,stop, - nthreads=nthreads, - tskip_msd=tskip_msd, - tmax=tmax_msd) - - #gofr (with time) - gofr=analyze_gofr(traj,start,stop,startr_gofr,endr_gofr,nbin_gofr, - tmax=tmax_gofr, + msd = analyze_msd(traj_unw, start, stop, nthreads=nthreads, - tskip=tskip_gofr) + tskip_msd=tskip_msd, + tmax=tmax_msd) + + # gofr (with time) + gofr = analyze_gofr(traj, start, stop, startr_gofr, endr_gofr, nbin_gofr, + tmax=tmax_gofr, + nthreads=nthreads, + tskip=tskip_gofr) - - #spherical harmonics correlations - sh = analyze_sh(traj,start,stop,startr_sh,endr_sh, nbin_sh, + # spherical harmonics correlations + sh = analyze_sh(traj, start, stop, startr_sh, endr_sh, nbin_sh, tmax=tmax_sh, nthreads=nthreads, tskip=tskip_sh) - print('all done',flush=True) + print('all done', flush=True) return msd, gofr, sh + try: import scipy as sp import scipy.ndimage - import scipy.optimize + import scipy.optimize except: print('WARNING: cannot import scipy') +def _int_gaussian(a, b): + return a * (np.pi / b) ** 0.5 + + +def _int_exp(a, b): + return a / b -def _int_gaussian(a,b): - return a*(np.pi/b)**0.5 -def _int_exp(a,b): - return a/b def int_fit2_(p): - freqs =[p[1]] + [ p[(i+1)*2+1] for i in range(0,(p.shape[0]-4)//2,2)] + [p[-2]] - magns = [(1-p[-1])*_int_gaussian(*p[0:2])]+[ - (1-p[-1])*_int_exp(*p[(i+1)*2:(i+2)*2]) - for i in range(0,(p.shape[0]-4)//2,2)] + [ - (1-p[-1])*_int_exp(double_exp_aux(p),p[-2]) ] - return [freqs,magns] + freqs = [p[1]] + [p[(i + 1) * 2 + 1] for i in range(0, (p.shape[0] - 4) // 2, 2)] + [p[-2]] + magns = [(1 - p[-1]) * _int_gaussian(*p[0:2])] + [ + (1 - p[-1]) * _int_exp(*p[(i + 1) * 2:(i + 2) * 2]) + for i in range(0, (p.shape[0] - 4) // 2, 2)] + [ + (1 - p[-1]) * _int_exp(double_exp_aux(p), p[-2])] + return [freqs, magns] + def _meta_generate_initial(n): - s='' + s = '' for i in range(n): - s+=f'{1.0/2**(i+1)},10*{10**-i},' - s=f'[{s}{10**-n},0]' - #print(s) + s += f'{1.0 / 2 ** (i + 1)},10*{10 ** -i},' + s = f'[{s}{10 ** -n},0]' + # print(s) return s + def _meta_generate_n_exp(n): - names={1:'single', 2:'double', 3:'triple',4:'quadruple'} - b_min='[0,0' - b_max='[np.inf,np.inf' - s=f''' -def {names.get(n,"_"+str(n))}_exp(x, *a): - return a[{n*2+1}]+(1.0-a[{n*2+1}])*(''' - m='a[0]+' - for i in range(1,n): - b_min+=',0,0' - b_max+=',np.inf,1.0' - s+=f'a[{i*2}]*np.exp(-a[{i*2+1}]*x)+' - m+=f'a[{i*2}]+' - last_mult=f'(1-({m}0))' - s+=f'a[0]*np.exp(-(a[1]*x)**2)+{last_mult}*np.exp(-a[{n*2}]*x) )' - b_min+=',0,0]' - b_max+=',np.inf,np.inf]' - s_aux=f'def {names.get(n,"_"+str(n))}_exp_aux(a): return {last_mult}' - #print(s) - #print(s_aux) + names = {1: 'single', 2: 'double', 3: 'triple', 4: 'quadruple'} + b_min = '[0,0' + b_max = '[np.inf,np.inf' + s = f''' +def {names.get(n, "_" + str(n))}_exp(x, *a): + return a[{n * 2 + 1}]+(1.0-a[{n * 2 + 1}])*(''' + m = 'a[0]+' + for i in range(1, n): + b_min += ',0,0' + b_max += ',np.inf,1.0' + s += f'a[{i * 2}]*np.exp(-a[{i * 2 + 1}]*x)+' + m += f'a[{i * 2}]+' + last_mult = f'(1-({m}0))' + s += f'a[0]*np.exp(-(a[1]*x)**2)+{last_mult}*np.exp(-a[{n * 2}]*x) )' + b_min += ',0,0]' + b_max += ',np.inf,np.inf]' + s_aux = f'def {names.get(n, "_" + str(n))}_exp_aux(a): return {last_mult}' + # print(s) + # print(s_aux) exec(s, globals()) - exec(s_aux,globals()) - s=f''' -def fit_{names.get(n,"_"+str(n))}_exp(datax,datay,p0={_meta_generate_initial(n)},bounds=(0,np.inf)): + exec(s_aux, globals()) + s = f''' +def fit_{names.get(n, "_" + str(n))}_exp(datax,datay,p0={_meta_generate_initial(n)},bounds=(0,np.inf)): #filter nan res=p0, None idxs=~np.isnan(datay) @@ -418,314 +438,344 @@ def fit_{names.get(n,"_"+str(n))}_exp(datax,datay,p0={_meta_generate_initial(n)} p0[-1] = 0 if p0[-1] > 1.0 : p0[-1] = 1.0 - res= sp.optimize.curve_fit({names.get(n,"_"+str(n))}_exp,datax[idxs],datay[idxs],p0=p0,bounds=bounds) + res= sp.optimize.curve_fit({names.get(n, "_" + str(n))}_exp,datax[idxs],datay[idxs],p0=p0,bounds=bounds) return res except: return res ''' - #print(s) - exec(s,globals()) + # print(s) + exec(s, globals()) + -for i in range(1,4): +for i in range(1, 4): _meta_generate_n_exp(i) - -def get_exp_spaced_idxs(t,scale=0.1): - idxs=[] - skips=np.array(np.exp((t-t[0])*scale),dtype=int) - idx=0 - while idxtimes[-1]: - maxt=times[-1] + + def get_idx(times, maxt): + idx_end = times.shape[0] + if maxt < 0.0 or maxt > times[-1]: + maxt = times[-1] else: - while idx_end>0 and times[idx_end-1]>maxt: - idx_end=idx_end-1 - if idx_end==0: + while idx_end > 0 and times[idx_end - 1] > maxt: + idx_end = idx_end - 1 + if idx_end == 0: pass - return maxt,idx_end - maxt,idx_end=get_idx(t,maxt) - if fitmax<0: - fitmax=maxt - fitmax,idx_fit=get_idx(t,fitmax) - if fitmin<0: - fitmin=times[0] - fitmin,idx_fit_beg=get_idx(t,fitmin) + return maxt, idx_end + + maxt, idx_end = get_idx(t, maxt) + if fitmax < 0: + fitmax = maxt + fitmax, idx_fit = get_idx(t, fitmax) + if fitmin < 0: + fitmin = times[0] + fitmin, idx_fit_beg = get_idx(t, fitmin) if idx_fit_beg >= idx_fit: raise IndexError('The provided upper time limit is too low') - print ('fit from {} ({}) to {} ({})'.format(fitmin,idx_fit_beg,fitmax,idx_fit)) - fits=[] - for l in range(lmin,lmax): - i=10-l - print('l={}: {}'.format(l, r[0,type1,type2,ibin,i])) - kwargs={} - if pre_fit>0: - idx_fit_0_t, idx_fit_0 = get_idx(t,pre_fit) - print ('pre-fit from {} ({}) to {} ({})'.format(fitmin,idx_fit_beg,idx_fit_0_t,idx_fit_0)) - if idx_fit_0 <=idx_fit_beg or idx_fit_0 >= idx_fit: + print('fit from {} ({}) to {} ({})'.format(fitmin, idx_fit_beg, fitmax, idx_fit)) + fits = [] + for l in range(lmin, lmax): + i = 10 - l + print('l={}: {}'.format(l, r[0, type1, type2, ibin, i])) + kwargs = {} + if pre_fit > 0: + idx_fit_0_t, idx_fit_0 = get_idx(t, pre_fit) + print('pre-fit from {} ({}) to {} ({})'.format(fitmin, idx_fit_beg, idx_fit_0_t, idx_fit_0)) + if idx_fit_0 <= idx_fit_beg or idx_fit_0 >= idx_fit: raise IndexError('The provided upper time limit for the pre-fit is in the wrong range') - p0,cv=fit_double_exp(t[idx_fit_beg:idx_fit_0], r[idx_fit_beg:idx_fit_0,type1, type2, ibin, i]/r[0,type1,type2,ibin,i]) + p0, cv = fit_double_exp(t[idx_fit_beg:idx_fit_0], + r[idx_fit_beg:idx_fit_0, type1, type2, ibin, i] / r[0, type1, type2, ibin, i]) kwargs['p0'] = p0 - p,cv=fit_double_exp(t[idx_fit_beg:idx_fit],r[idx_fit_beg:idx_fit,type1,type2,ibin,i]/r[0,type1,type2,ibin,i],**kwargs) - if abs(1.0-math.exp(-p[-2]*t[idx_fit-1]) )<1e-5: #safely assume that the last exponential is approximable by a constant + p, cv = fit_double_exp(t[idx_fit_beg:idx_fit], + r[idx_fit_beg:idx_fit, type1, type2, ibin, i] / r[0, type1, type2, ibin, i], **kwargs) + if abs(1.0 - math.exp(-p[-2] * t[ + idx_fit - 1])) < 1e-5: # safely assume that the last exponential is approximable by a constant p = norm_fit(p) fits.append(p) return fits -def plot_sh(startr,endr,times,res,type1,type2,ibin,lmin=0,lmax=11,title='',res_var=None, - maxt=-1.0,fitmax=-1.0,fitmin=-1.0,pre_fit=-1.0,log=True,rescale_array=True,scale_exp_coeff=0.1, - fig_ax=None,subplots_kw=DEFAULT_SUBPLOTS_KW): - t=times[:res.shape[0]]-times[0] +def plot_sh(startr, endr, times, res, type1, type2, ibin, lmin=0, lmax=11, title='', res_var=None, + maxt=-1.0, fitmax=-1.0, fitmin=-1.0, pre_fit=-1.0, log=True, rescale_array=True, scale_exp_coeff=0.1, + fig_ax=None, subplots_kw=DEFAULT_SUBPLOTS_KW): + t = times[:res.shape[0]] - times[0] r = res r_var = res_var - def get_idx(times,maxt): - idx_end=times.shape[0] - if maxt<0.0 or maxt>times[-1]: - maxt=times[-1] + + def get_idx(times, maxt): + idx_end = times.shape[0] + if maxt < 0.0 or maxt > times[-1]: + maxt = times[-1] else: - while idx_end>0 and times[idx_end-1]>maxt: - idx_end=idx_end-1 - if idx_end==0: + while idx_end > 0 and times[idx_end - 1] > maxt: + idx_end = idx_end - 1 + if idx_end == 0: pass - return maxt,idx_end - maxt,idx_end=get_idx(t,maxt) - dr=(endr-startr)/r.shape[3] + return maxt, idx_end + + maxt, idx_end = get_idx(t, maxt) + dr = (endr - startr) / r.shape[3] if fig_ax is not None: - fig, ax = fig_ax + fig, ax = fig_ax else: - fig,ax =plt.subplots(**subplots_kw) - ax=fig.add_axes([0,0,1,1]) - axins=inset_axes(ax,width='20%',height='20%',loc='upper right') - fits = fit_sh(t,res,type1,type2,ibin,lmin,lmax,maxt,fitmax,fitmin,pre_fit, - rescale_array,scale_exp_coeff) - for l in range(lmin,lmax): - i=10-l - p=fits[l-lmin] + fig, ax = plt.subplots(**subplots_kw) + ax = fig.add_axes([0, 0, 1, 1]) + axins = inset_axes(ax, width='20%', height='20%', loc='upper right') + fits = fit_sh(t, res, type1, type2, ibin, lmin, lmax, maxt, fitmax, fitmin, pre_fit, + rescale_array, scale_exp_coeff) + for l in range(lmin, lmax): + i = 10 - l + p = fits[l - lmin] + def plot(ax): with np.printoptions(precision=3, suppress=True): - ax.plot(t,double_exp(t,*p),label='fit l={} p={}'.format(l,p)) - plt_err(ax,t,r[:,type1,type2,ibin,i]/r[0,type1,type2,ibin,i], - (r_var[:,type1,type2,ibin,i]/r[0,type1,type2,ibin,i]**2) if r_var is not None else r_var,label='l='+str(l)) + ax.plot(t, double_exp(t, *p), label='fit l={} p={}'.format(l, p)) + plt_err(ax, t, r[:, type1, type2, ibin, i] / r[0, type1, type2, ibin, i], + (r_var[:, type1, type2, ibin, i] / r[ + 0, type1, type2, ibin, i] ** 2) if r_var is not None else r_var, label='l=' + str(l)) + plot(ax) plot(axins) ax.legend() - axins.set_xlim(0,maxt/50.0) - #ax.set_xscale('log') + axins.set_xlim(0, maxt / 50.0) + # ax.set_xscale('log') if log: ax.set_yscale('log') axins.set_yscale('log') - ax.set_title('{}$r \in [{:.3f} ,{:.3f}[$'.format(title,startr+dr*ibin,startr+dr*(ibin+1))) - ax.set_xlim(0,maxt) + ax.set_title('{}$r \\in [{:.3f} ,{:.3f}[$'.format(title, startr + dr * ibin, startr + dr * (ibin + 1))) + ax.set_xlim(0, maxt) ax.set_xlabel('time (ps)') - ax.set_ylabel('$c_{{\ell}}^{{{0}\,{1}}}(t)/c_{{\ell}}^{{{0}\,{1}}}(0)$'.format(type1,type2)) - return fig,ax,axins, fits + ax.set_ylabel('$c_{{\\ell}}^{{{0}\\,{1}}}(t)/c_{{\\ell}}^{{{0}\\,{1}}}(0)$'.format(type1, type2)) + return fig, ax, axins, fits -def plot_gofr(startr,endr,res,title='',res_var=None,fig_ax=None,subplots_kw=DEFAULT_SUBPLOTS_KW): +def plot_gofr(startr, endr, res, title='', res_var=None, fig_ax=None, subplots_kw=DEFAULT_SUBPLOTS_KW): if fig_ax is not None: - fig, ax = fig_ax + fig, ax = fig_ax else: - fig,ax =plt.subplots(**subplots_kw) - rm=np.arange(res.shape[-1])*(endr-startr)/res.shape[-1]+startr - rp=(np.arange(res.shape[-1])+1)*(endr-startr)/res.shape[-1]+startr - vols=4*np.pi/3*(rp**3-rm**3) - for i in range(0,res.shape[1]//2): - plt_err(ax,(rm+rp)/2,res[0,i,:]/vols,(res_var[0,i,:]/(vols**4)) if res_var is not None else res_var) - ax.set_xlabel('$\AA$') + fig, ax = plt.subplots(**subplots_kw) + rm = np.arange(res.shape[-1]) * (endr - startr) / res.shape[-1] + startr + rp = (np.arange(res.shape[-1]) + 1) * (endr - startr) / res.shape[-1] + startr + vols = 4 * np.pi / 3 * (rp ** 3 - rm ** 3) + for i in range(0, res.shape[1] // 2): + plt_err(ax, (rm + rp) / 2, res[0, i, :] / vols, + (res_var[0, i, :] / (vols ** 4)) if res_var is not None else res_var) + ax.set_xlabel(r'$\AA$') ax.set_ylabel('number of atoms in the shell / volume of the shell') ax.set_title('{}$g(r)$'.format(title)) return fig, ax -def sph_tangent(r,phi,theta,x,y,z): - a=math.cos(theta)*math.sin(phi) - b=math.sin(theta)*math.sin(phi) - c=math.cos(phi) - return [a,b,c,r-a*x-b*y-c*z] -def sph_tangents_cover(r,n,x,y,z): - res=[] - for i in range(n+1): - phi=math.pi*(i-n/2.0)/n - n_theta=max(1,int(n*math.cos(phi))) + +def sph_tangent(r, phi, theta, x, y, z): + a = math.cos(theta) * math.sin(phi) + b = math.sin(theta) * math.sin(phi) + c = math.cos(phi) + return [a, b, c, r - a * x - b * y - c * z] + + +def sph_tangents_cover(r, n, x, y, z): + res = [] + for i in range(n + 1): + phi = math.pi * (i - n / 2.0) / n + n_theta = max(1, int(n * math.cos(phi))) for j in range(n_theta): - theta=2*math.pi*j/n_theta - res.append(sph_tangent(r,phi+math.pi/2.0,theta,x,y,z)) + theta = 2 * math.pi * j / n_theta + res.append(sph_tangent(r, phi + math.pi / 2.0, theta, x, y, z)) return res -def get_max_in_mask(a,mask): - b=a.view(np.ma.MaskedArray) - b.mask=mask + +def get_max_in_mask(a, mask): + b = a.view(np.ma.MaskedArray) + b.mask = mask return np.unravel_index(np.argmax(b, axis=None), b.shape) -def idx_to_coord(idx,shape,l=(1.0,1.0,1.0),l0=np.array((0.0,0.0,0.0))): - return (np.array(idx,dtype=float)/np.array(shape,dtype=float)-l0)*np.array(l) -def is_inside(a,idx,planes): - idx_c=idx_to_coord(idx,a.shape) +def idx_to_coord(idx, shape, l=(1.0, 1.0, 1.0), l0=np.array((0.0, 0.0, 0.0))): + return (np.array(idx, dtype=float) / np.array(shape, dtype=float) - l0) * np.array(l) + + +def is_inside(a, idx, planes): + idx_c = idx_to_coord(idx, a.shape) for plane in planes: - #print(idx_c) - if np.sum(idx_c*np.array(plane[:-1]))+plane[3]<0: + # print(idx_c) + if np.sum(idx_c * np.array(plane[:-1])) + plane[3] < 0: return 1 return 0 -def is_inside_sphere(a,idx,r2,xyz,l=(1.0,1.0,1.0)): - return np.sum((idx_to_coord(idx,a.shape,l=l)-xyz)**2) mask.shape[i]: if low > 0: - idxs_low_high[i].append((low,mask.shape[i])) - idxs_low_high[i].append((0,mask.shape[i]-high)) + idxs_low_high[i].append((low, mask.shape[i])) + idxs_low_high[i].append((0, mask.shape[i] - high)) else: - idxs_low_high[i].append((0,mask.shape[i])) + idxs_low_high[i].append((0, mask.shape[i])) else: - idxs_low_high[i].append((low,high)) - for zlo,zhi in idxs_low_high[0]: - for ylo,yhi in idxs_low_high[1]: - for xlo,xhi in idxs_low_high[2]: - mask[zlo:zhi,ylo:yhi,xlo:xhi]=1 - -def d2_minimag(a,b,l): - d2=0.0 + idxs_low_high[i].append((low, high)) + for zlo, zhi in idxs_low_high[0]: + for ylo, yhi in idxs_low_high[1]: + for xlo, xhi in idxs_low_high[2]: + mask[zlo:zhi, ylo:yhi, xlo:xhi] = 1 + + +def d2_minimag(a, b, l): + d2 = 0.0 for i in range(3): - if abs(a[i]-b[i]) > l[i]/2: - d2+=(abs(a[i]-b[i])-l[i]/2)**2 + if abs(a[i] - b[i]) > l[i] / 2: + d2 += (abs(a[i] - b[i]) - l[i] / 2) ** 2 else: - d2+=(a[i]-b[i])**2 + d2 += (a[i] - b[i]) ** 2 return d2 -def mask_around_spherical(mask,idx,r): - r2=r**2 - mask_box=np.zeros_like(mask) - mask_around(mask_box,idx,(int(r),int(r),int(r))) + + +def mask_around_spherical(mask, idx, r): + r2 = r ** 2 + mask_box = np.zeros_like(mask) + mask_around(mask_box, idx, (int(r), int(r), int(r))) for index, m in np.ndenumerate(mask_box): if m: - #check if it is inside the sphere, then mask it - if d2_minimag(index,idx,mask.shape) <= r2: - mask[index]=1 - - - -def get_local_maximum_idx(a,threshold=0.0,minimum_distance=(0.1,0.1,0.1),l=(1.0,1.0,1.0)): - max_list=[] - max_list_coord=[] - mask_size=[0,0,0] + # check if it is inside the sphere, then mask it + if d2_minimag(index, idx, mask.shape) <= r2: + mask[index] = 1 + + +def get_local_maximum_idx(a, threshold=0.0, minimum_distance=(0.1, 0.1, 0.1), l=(1.0, 1.0, 1.0)): + max_list = [] + max_list_coord = [] + mask_size = [0, 0, 0] for i in range(3): - mask_size[i]=max(int(minimum_distance[i]/l[i]*a.shape[i]/2),1) - b=a.view(np.ma.MaskedArray) - mask=np.zeros(a.shape) + mask_size[i] = max(int(minimum_distance[i] / l[i] * a.shape[i] / 2), 1) + b = a.view(np.ma.MaskedArray) + mask = np.zeros(a.shape) while True: - max_last=np.unravel_index(np.argmax(b, axis=None), b.shape) - if b[max_last]3: + plot = k3d.plot() + objs = [] + if len(res.shape) > 3: for i in range(res.shape[0]): - objs.append(k3d.volume(np.array(res[i],dtype='float16'),bounds=bounds)) - plot+=objs[-1] + objs.append(k3d.volume(np.array(res[i], dtype='float16'), bounds=bounds)) + plot += objs[-1] else: - objs.append(k3d.volume(np.array(res,dtype='float16'),bounds=bounds)) - plot+=objs[-1] - - #points = k3d.points(g0_sites,point_size=0.02) - #points1 = k3d.points(g1_sites,point_size=0.005) - #plot += points - #plot += points1 - plot += plot_simulation_box(box,**box_kw) + objs.append(k3d.volume(np.array(res, dtype='float16'), bounds=bounds)) + plot += objs[-1] + + # points = k3d.points(g0_sites,point_size=0.02) + # points1 = k3d.points(g1_sites,point_size=0.005) + # plot += points + # plot += points1 + plot += plot_simulation_box(box, **box_kw) plot.display() global _w_x global _x global _n_idx - _w_x=0.05 - _x =0.14 + _w_x = 0.05 + _x = 0.14 _n_idx = 0 global _rot_a global _rot_b global _rot_c global _rot_Q - _rot_a = 0.0 + _rot_a = 0.0 _rot_b = 0.0 _rot_c = 0.0 _rot_Q = np.eye(4) - nsn=np.array(ns) - nsn=np.einsum('ij,i->ij',nsn,((nsn**2).sum(axis=1))**-.5) - nlabel=widgets.Label(value=f'n={nsn[_n_idx]}') + nsn = np.array(ns) + nsn = np.einsum('ij,i->ij', nsn, ((nsn ** 2).sum(axis=1)) ** -.5) + nlabel = widgets.Label(value=f'n={nsn[_n_idx]}') display(nlabel) - def rotate_plane(Q,ps): - prs=[] - for p in ps: - p=np.array(p) - prs.append(Q.dot(p).tolist()) - return prs + + def rotate_plane(Q, ps): + prs = [] + for p in ps: + p = np.array(p) + prs.append(Q.dot(p).tolist()) + return prs + def get_planes(): - n=nsn[_n_idx] - c1=-(r0).dot(n) + _w_x +_x - c2=-(r0).dot(n) - _w_x +_x + n = nsn[_n_idx] + c1 = -(r0).dot(n) + _w_x + _x + c2 = -(r0).dot(n) - _w_x + _x return [ - [ n[0], n[1], n[2], c1], - [-n[0],-n[1],-n[2],-c2] - ] + [n[0], n[1], n[2], c1], + [-n[0], -n[1], -n[2], -c2] + ] + def global_clipping_planes(): - return rotate_plane(_rot_Q, get_planes()) - @interact(idx=widgets.IntSlider(value=0,min=0,max=nsn.shape[0]-1,step=1,description='select index of plane')) + return rotate_plane(_rot_Q, get_planes()) + + @interact(idx=widgets.IntSlider(value=0, min=0, max=nsn.shape[0] - 1, step=1, description='select index of plane')) def g(idx): global _n_idx - _n_idx=idx - plot.clipping_planes=global_clipping_planes() - nlabel.value=f'n={nsn[_n_idx]}' - @interact(x=widgets.FloatSlider(value=box[0]+box[3],min=box[0],max=box[0]+2*box[3],step=box[3]/100)) + _n_idx = idx + plot.clipping_planes = global_clipping_planes() + nlabel.value = f'n={nsn[_n_idx]}' + + @interact(x=widgets.FloatSlider(value=box[0] + box[3], min=box[0], max=box[0] + 2 * box[3], step=box[3] / 100)) def g(x): global _x - _x=x - plot.clipping_planes=global_clipping_planes() - @interact(w_x=widgets.FloatSlider(value=box[3]/5,min=0.0,max=2*box[3],step=box[3]/100)) + _x = x + plot.clipping_planes = global_clipping_planes() + + @interact(w_x=widgets.FloatSlider(value=box[3] / 5, min=0.0, max=2 * box[3], step=box[3] / 100)) def g(w_x): global _w_x - _w_x=w_x - plot.clipping_planes=global_clipping_planes() - @interact(a=widgets.FloatSlider(value=0,min=0,max=2*np.pi,step=0.01)) + _w_x = w_x + plot.clipping_planes = global_clipping_planes() + + @interact(a=widgets.FloatSlider(value=0, min=0, max=2 * np.pi, step=0.01)) def g(a): global _rot_a global _rot_Q - _rot_a=a - _rot_Q=rotation(_rot_a,_rot_b,_rot_c) - plot.clipping_planes=global_clipping_planes() - @interact(b=widgets.FloatSlider(value=0,min=0,max=2*np.pi,step=0.01)) + _rot_a = a + _rot_Q = rotation(_rot_a, _rot_b, _rot_c) + plot.clipping_planes = global_clipping_planes() + + @interact(b=widgets.FloatSlider(value=0, min=0, max=2 * np.pi, step=0.01)) def g(b): global _rot_b global _rot_Q - _rot_b=b - _rot_Q=rotation(_rot_a,_rot_b,_rot_c) - plot.clipping_planes=global_clipping_planes() - @interact(c=widgets.FloatSlider(value=0,min=0,max=2*np.pi,step=0.01)) + _rot_b = b + _rot_Q = rotation(_rot_a, _rot_b, _rot_c) + plot.clipping_planes = global_clipping_planes() + + @interact(c=widgets.FloatSlider(value=0, min=0, max=2 * np.pi, step=0.01)) def g(c): global _rot_c global _rot_Q - _rot_c=c - _rot_Q=rotation(_rot_a,_rot_b,_rot_c) - plot.clipping_planes=global_clipping_planes() + _rot_c = c + _rot_Q = rotation(_rot_a, _rot_b, _rot_c) + plot.clipping_planes = global_clipping_planes() return plot -def density_field(res,box,box_kw={},plot=None,sph_max_idx=0,gaussian_filter=False,gaussian_filter_radius=1.2): - bounds=[ - box[0],box[0]+2*box[3], - box[1],box[1]+2*box[4], - box[2],box[2]+2*box[5] - ] + +def density_field(res, box, box_kw={}, plot=None, sph_max_idx=0, gaussian_filter=False, gaussian_filter_radius=1.2): + bounds = [ + box[0], box[0] + 2 * box[3], + box[1], box[1] + 2 * box[4], + box[2], box[2] + 2 * box[5] + ] print(bounds) if plot is None: - plot = k3d.plot() - objs=[] - if len(res.shape) >3: + plot = k3d.plot() + objs = [] + if len(res.shape) > 3: for i in range(res.shape[0]): if gaussian_filter: - objs.append(k3d.volume( - scipy.ndimage.gaussian_filter(res[i],gaussian_filter_radius,output='float32') - ,bounds=bounds)) + objs.append(k3d.volume( + scipy.ndimage.gaussian_filter(res[i], gaussian_filter_radius, output='float32') + , bounds=bounds)) else: - objs.append(k3d.volume(np.array(res[i],dtype='float16'),bounds=bounds)) - plot+=objs[-1] + objs.append(k3d.volume(np.array(res[i], dtype='float16'), bounds=bounds)) + plot += objs[-1] else: - objs.append(k3d.volume(np.array(res,dtype='float16'),bounds=bounds)) - plot+=objs[-1] - - #points = k3d.points(g0_sites,point_size=0.02) - #points1 = k3d.points(g1_sites,point_size=0.005) - #plot += points - #plot += points1 - plot += plot_simulation_box(box,**box_kw) + objs.append(k3d.volume(np.array(res, dtype='float16'), bounds=bounds)) + plot += objs[-1] + + # points = k3d.points(g0_sites,point_size=0.02) + # points1 = k3d.points(g1_sites,point_size=0.005) + # plot += points + # plot += points1 + plot += plot_simulation_box(box, **box_kw) plot.display() global _w_x global _x global _w_y global _y global _w_z - global _z + global _z global _sphere - _w_x=0.05 - _x =0.14 - _w_y=0.05 - _y =0.14 - _w_z=0.05 - _z =0.14 - _sphere=False + _w_x = 0.05 + _x = 0.14 + _w_y = 0.05 + _y = 0.14 + _w_z = 0.05 + _z = 0.14 + _sphere = False global _rot_a global _rot_b global _rot_c global _rot_Q - _rot_a = 0.0 + _rot_a = 0.0 _rot_b = 0.0 _rot_c = 0.0 _rot_Q = np.eye(4) - def rotate_plane(Q,ps): - prs=[] - for p in ps: - p=np.array(p) - prs.append(Q.dot(p).tolist()) - return prs + + def rotate_plane(Q, ps): + prs = [] + for p in ps: + p = np.array(p) + prs.append(Q.dot(p).tolist()) + return prs + def global_clipping_planes(): if _sphere: - return sph_tangents_cover(_w_x,7,_x,_y,_z) + return sph_tangents_cover(_w_x, 7, _x, _y, _z) else: - return rotate_plane(_rot_Q,[[1,0,0,_w_x-_x],[-1,0,0,_w_x+_x], - [0,1,0,_w_y-_y],[0,-1,0,_w_y+_y], - [0,0,1,_w_z-_z],[0,0,-1,_w_z+_z] - ]) + return rotate_plane(_rot_Q, [[1, 0, 0, _w_x - _x], [-1, 0, 0, _w_x + _x], + [0, 1, 0, _w_y - _y], [0, -1, 0, _w_y + _y], + [0, 0, 1, _w_z - _z], [0, 0, -1, _w_z + _z] + ]) + box_layout = widgets.Layout(display='flex', - flex_flow='row', - align_items='stretch', - border='none', - width='95%') - button=widgets.Button(description='Center on 1',layout=box_layout) + flex_flow='row', + align_items='stretch', + border='none', + width='95%') + button = widgets.Button(description='Center on 1', layout=box_layout) + @button.on_click def button1_on_click(b): global _x global _y global _z - print('finding maximum in spherical domain with r={}, center=({},{},{})'.format(_w_x,_x,_y,_z)) - idx, (_z,_y,_x)=get_max_in_spherical_domain(res[sph_max_idx],_w_x,_z,_y,_x,l=(2*box[5],2*box[4],2*box[3])) - #_x,_y,_z=-_x,-_y,-_z - print('center=({},{},{}), idx={} shape={}'.format(_x,_y,_z,idx,res[sph_max_idx].shape)) + print('finding maximum in spherical domain with r={}, center=({},{},{})'.format(_w_x, _x, _y, _z)) + idx, (_z, _y, _x) = get_max_in_spherical_domain(res[sph_max_idx], _w_x, _z, _y, _x, + l=(2 * box[5], 2 * box[4], 2 * box[3])) + # _x,_y,_z=-_x,-_y,-_z + print('center=({},{},{}), idx={} shape={}'.format(_x, _y, _z, idx, res[sph_max_idx].shape)) print('value={}'.format(res[sph_max_idx][idx])) - plot.clipping_planes=global_clipping_planes() - wlist=[button] + plot.clipping_planes = global_clipping_planes() + + wlist = [button] + def g1(spherical): global _sphere - _sphere=spherical - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g1,spherical=widgets.Checkbox(value=False,layout=box_layout))) + _sphere = spherical + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g1, spherical=widgets.Checkbox(value=False, layout=box_layout))) + def g2(a): global _rot_a global _rot_Q - _rot_a=a - _rot_Q=rotation(_rot_a,_rot_b,_rot_c) - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g2,a=widgets.FloatSlider(value=0,min=0,max=2*np.pi,step=0.01,layout=box_layout))) + _rot_a = a + _rot_Q = rotation(_rot_a, _rot_b, _rot_c) + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g2, a=widgets.FloatSlider(value=0, min=0, max=2 * np.pi, step=0.01, layout=box_layout))) + def g3(b): global _rot_b global _rot_Q - _rot_b=b - _rot_Q=rotation(_rot_a,_rot_b,_rot_c) - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g3,b=widgets.FloatSlider(value=0,min=0,max=2*np.pi,step=0.01,layout=box_layout))) + _rot_b = b + _rot_Q = rotation(_rot_a, _rot_b, _rot_c) + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g3, b=widgets.FloatSlider(value=0, min=0, max=2 * np.pi, step=0.01, layout=box_layout))) + def g4(c): global _rot_c global _rot_Q - _rot_c=c - _rot_Q=rotation(_rot_a,_rot_b,_rot_c) - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g4,c=widgets.FloatSlider(value=0,min=0,max=2*np.pi,step=0.01,layout=box_layout))) + _rot_c = c + _rot_Q = rotation(_rot_a, _rot_b, _rot_c) + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g4, c=widgets.FloatSlider(value=0, min=0, max=2 * np.pi, step=0.01, layout=box_layout))) + def g5(x): global _x - _x=x - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g5,x=widgets.FloatSlider(value=box[0]+box[3],min=box[0],max=box[0]+2*box[3],step=box[3]/100,layout=box_layout))) + _x = x + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g5, x=widgets.FloatSlider(value=box[0] + box[3], min=box[0], max=box[0] + 2 * box[3], + step=box[3] / 100, layout=box_layout))) + def g6(w_x): global _w_x - _w_x=w_x - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g6,w_x=widgets.FloatSlider(value=box[3]/5,min=0.0,max=2*box[3],step=box[3]/100,layout=box_layout))) + _w_x = w_x + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g6, w_x=widgets.FloatSlider(value=box[3] / 5, min=0.0, max=2 * box[3], step=box[3] / 100, + layout=box_layout))) + def g7(y): global _y - _y=y - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g7,y=widgets.FloatSlider(value=box[1]+box[4],min=box[1],max=box[1]+2*box[4],step=box[4]/100,layout=box_layout))) + _y = y + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g7, y=widgets.FloatSlider(value=box[1] + box[4], min=box[1], max=box[1] + 2 * box[4], + step=box[4] / 100, layout=box_layout))) + def g8(w_y): global _w_y - _w_y=w_y - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g8,w_y=widgets.FloatSlider(value=2*box[4],min=0,max=2*box[4],step=box[4]/100,layout=box_layout))) + _w_y = w_y + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g8, w_y=widgets.FloatSlider(value=2 * box[4], min=0, max=2 * box[4], step=box[4] / 100, + layout=box_layout))) + def g9(z): global _z - _z=z - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g9,z=widgets.FloatSlider(value=box[2]+box[5],min=box[2],max=box[2]+2*box[5],step=box[5]/100,layout=box_layout))) + _z = z + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g9, z=widgets.FloatSlider(value=box[2] + box[5], min=box[2], max=box[2] + 2 * box[5], + step=box[5] / 100, layout=box_layout))) + def g10(w_z): global _w_z - _w_z=w_z - plot.clipping_planes=global_clipping_planes() - wlist.append(interactive(g10,w_z=widgets.FloatSlider(value=2*box[5],min=0.0,max=2*box[5],step=box[5]/100,layout=box_layout))) + _w_z = w_z + plot.clipping_planes = global_clipping_planes() + + wlist.append(interactive(g10, w_z=widgets.FloatSlider(value=2 * box[5], min=0.0, max=2 * box[5], step=box[5] / 100, + layout=box_layout))) box_all = widgets.GridBox(children=wlist, layout=widgets.Layout( - grid_template_columns='40% 40%', - grid_template_rows='auto' + grid_template_columns='40% 40%', + grid_template_rows='auto' + ) ) - ) - #for w in wlist: display(w) + # for w in wlist: display(w) display(box_all) return plot - -def force_ratio_histogram(wf,print=print,ax=[],create_fig=lambda : plt.subplots(**DEFAULT_SUBPLOTS_KW)): - pwcalcjobs=[] - for c in [x for x in wf.called if str(x.process_class) == ""]: + + +def force_ratio_histogram(wf, print=print, ax=[], create_fig=lambda: plt.subplots(**DEFAULT_SUBPLOTS_KW)): + pwcalcjobs = [] + for c in [x for x in wf.called if + str(x.process_class) == ""]: pwcalcjobs.append(c) print(c.pk) - res,axs,figs,n_ax=analyze_forces_ratio(pwcalcjobs,minpk=wf.pk,ax_=ax,create_fig = create_fig) - return res,axs,figs,n_ax - -def plot_force_ratio(res,fig=None,ax=None,hheight=10,subplots_kw=DEFAULT_SUBPLOTS_KW): - def update_x_low_high(em,dt,d,x,x_size=0.0): - e=d.setdefault(dt,{}).setdefault(em,{}) - x_lo=e.setdefault('x_lo',x) - if x-x_sizex_hi: e['x_hi'] = x+x_size - - - #c_k=list(mcolors.TABLEAU_COLORS.keys()) - #idx=0 - if ax==None: - fig,ax=plt.subplots(**subplots_kw) - x_low_high={} - for e in dict_keys(res,level=2): - for dt in dict_keys(res,level=1): - y=[] - x=[] - y_err=[] + res, axs, figs, n_ax = analyze_forces_ratio(pwcalcjobs, minpk=wf.pk, ax_=ax, create_fig=create_fig) + return res, axs, figs, n_ax + + +def plot_force_ratio(res, fig=None, ax=None, hheight=10, subplots_kw=DEFAULT_SUBPLOTS_KW): + def update_x_low_high(em, dt, d, x, x_size=0.0): + e = d.setdefault(dt, {}).setdefault(em, {}) + x_lo = e.setdefault('x_lo', x) + if x - x_size < x_lo: e['x_lo'] = x - x_size + x_hi = e.setdefault('x_hi', x) + if x + x_size > x_hi: e['x_hi'] = x + x_size + + # c_k=list(mcolors.TABLEAU_COLORS.keys()) + # idx=0 + if ax == None: + fig, ax = plt.subplots(**subplots_kw) + x_low_high = {} + for e in dict_keys(res, level=2): + for dt in dict_keys(res, level=1): + y = [] + x = [] + y_err = [] for em in res.keys(): try: - update_x_low_high(em,dt,x_low_high,res[em][dt][e]['fratios_mean'],res[em][dt][e]['fratios_std']) + update_x_low_high(em, dt, x_low_high, res[em][dt][e]['fratios_mean'], res[em][dt][e]['fratios_std']) x.append(em) y.append(res[em][dt][e]['fratios_mean']) y_err.append(res[em][dt][e]['fratios_std']) except KeyError: pass - if len(x)>0: - ax.errorbar(x,y,yerr=y_err,label=f'{e} dt={dt:.2f}',alpha=0.7) - axins=[] - emass_min=list(res.keys())[0] - emass_max=list(res.keys())[0] - for dt in dict_keys(res,level=1): + if len(x) > 0: + ax.errorbar(x, y, yerr=y_err, label=f'{e} dt={dt:.2f}', alpha=0.7) + axins = [] + emass_min = list(res.keys())[0] + emass_max = list(res.keys())[0] + for dt in dict_keys(res, level=1): for em in res.keys(): try: - axi=inset_axes( + axi = inset_axes( ax, bbox_transform=ax.transData, bbox_to_anchor=( - em, - x_low_high[dt][em]['x_lo'], - hheight, - x_low_high[dt][em]['x_hi']-x_low_high[dt][em]['x_lo'] - ), - width="100%", height="100%",borderpad=0 - ) - for e in dict_keys(res,level=2): - x=res[em][dt][e]['hist'][0] - y=res[em][dt][e]['hist'][1] - y=(y[1:]+y[:-1])/2.0 - axi.plot(x,y,alpha=0.5) - axi.set_ylim([x_low_high[dt][em]['x_lo'],x_low_high[dt][em]['x_hi']]) + em, + x_low_high[dt][em]['x_lo'], + hheight, + x_low_high[dt][em]['x_hi'] - x_low_high[dt][em]['x_lo'] + ), + width="100%", height="100%", borderpad=0 + ) + for e in dict_keys(res, level=2): + x = res[em][dt][e]['hist'][0] + y = res[em][dt][e]['hist'][1] + y = (y[1:] + y[:-1]) / 2.0 + axi.plot(x, y, alpha=0.5) + axi.set_ylim([x_low_high[dt][em]['x_lo'], x_low_high[dt][em]['x_hi']]) axi.set_axis_off() axi.patch.set_alpha(0.5) axins.append(axi) - if ememass_max: emass_max=em + if em < emass_min: emass_min = em + if em > emass_max: emass_max = em except KeyError: pass ax.legend(loc=3) ax.grid() - ax.set_xlim(emass_min-hheight,emass_max+hheight) + ax.set_xlim(emass_min - hheight, emass_max + hheight) ax.set_xlabel('electron fictitious mass') ax.set_ylabel('CP/PW force ratio') - return fig,ax,axins + return fig, ax, axins + def length_traj(traj): - t=traj.get_array('times') - return t[-1]-t[0] -def get_cp_with_traj(wf,min_t=0.0, ignore_not_ok=False): - l=[] - lall=[] + t = traj.get_array('times') + return t[-1] - t[0] + + +def get_cp_with_traj(wf, min_t=0.0, ignore_not_ok=False): + l = [] + lall = [] for x in wf.called: try: if str(x.process_class) == "": - lall.append(x) + lall.append(x) except AttributeError as err: print('Ignoring: ', err) + except ValueError as err: + print('Ignoring: ', err) for c in lall: if 'output_trajectory' in c.outputs and (c.is_finished_ok or ignore_not_ok): if min_t > 0.0: - t = length_traj(c.outputs.output_trajectory) - if t >= min_t: - l.append(c) + t = length_traj(c.outputs.output_trajectory) + if t >= min_t: + l.append(c) else: - l.append(c) - return sorted(l,key=lambda x: x.pk) + l.append(c) + return sorted(l, key=lambda x: x.pk) -def print_cp_with_traj(wf,print=print,min_t=0.0): - l = get_cp_with_traj(wf,min_t=min_t) + +def print_cp_with_traj(wf, print=print, min_t=0.0): + l = get_cp_with_traj(wf, min_t=min_t) for c in l: - print ("===========") - print(c.pk,'traj pk={}: {:.3}ps'.format(c.outputs.output_trajectory.pk,length_traj(c.outputs.output_trajectory)) if 'output_trajectory' in c.outputs else 'No output_trajectory') + print("===========") + print(c.pk, 'traj pk={}: {:.3}ps'.format(c.outputs.output_trajectory.pk, length_traj( + c.outputs.output_trajectory)) if 'output_trajectory' in c.outputs else 'No output_trajectory') print(c.inputs.parameters.get_dict()['IONS']) print(c.inputs.parameters.get_dict()['CELL'] if 'CELL' in c.inputs.parameters.get_dict() else '') return l -#tools for converting an analisi trajectory to a aiida trajectory, or an object that behaves in a similar way - -def analisi_cell2box(box_lammps): - '''returns cell vectors arranged in ROWS. - Analisi code always rotate the system to have a triangular cell matrix.''' - cells=np.zeros((box_lammps.shape[0],3,3)) - for i in range(box_lammps.shape[0]): - lc=box_lammps[i] - if lc.shape[0] == 9: #note that this matrix is transposed: the cell vectors are arranged in ROWS - # a_x 0 0 - # b_x b_y 0 - # c_x c_y c_z - cells[i]=np.array([ [lc[3]*2, 0, 0 ], - [lc[6], lc[4]*2, 0 ], - [lc[7], lc[8], lc[5]*2] - ]) - elif lc.shape[0] == 6: #for orthorombic cells there is no difference between row arranged and column arranged cell vectors - cells[i]=np.array([ [lc[3]*2,0,0], - [0,lc[4]*2,0], - [0,0,lc[5]*2]]) - else: - raise IndexError('wrong shape of cell array') - return cells - -class FakeAiidaT: - def __init__(self,t,symbols_={},dt=1.0,pk=None,traj_skip=1): - def _copy(d,skip): - return np.copy(d[::traj_skip],order='C') if skip>1 else d - self.data={} - key_l=['positions','cells'] - key_opt=['velocities','energy_constant_motion','ionic_temperature','pressure','electronic_kinetic_energy'] - if isinstance(t,(pa.Traj,pa.Trajectory)): - self.t=t - self.symbols=[ symbols_[x] if x in symbols_ else str(x) for x in t.get_lammps_type().tolist()] - self.data['positions']=_copy(t.get_positions_copy(), traj_skip) - self.data['velocities']=_copy(t.get_velocities_copy(), traj_skip) - self.box_lammps=_copy(t.get_box_copy(), traj_skip) - self.numsteps=self.data['positions'].shape[0] - self.data['cells']=analisi_cell2box(self.box_lammps) - elif isinstance(t,dict): - self.symbols=t['symbols'] - for key in key_l: - self.data[key]=_copy(np.array(t[key]),traj_skip) - for key in key_opt: - if key in t: - self.data[key]=_copy(np.array(t[key]),traj_skip) - elif isinstance(t,list): - if isinstance(t[0], dict): - self.symbols=t[0]['symbols'] - for key in key_l: - self.data[key] = np.concatenate( [_copy(t[i][key], traj_skip) for i in range(len(t))], axis=0) - for key in key_opt: - if key in t[0]: - self.data[key]=np.concatenate( [_copy(t[i][key], traj_skip) for i in range(len(t))], axis=0) - else: - self.symbols=t[0].symbols - for key in key_l: - self.data[key] = np.concatenate( [_copy(t[i].get_array(key), traj_skip) for i in range(len(t))], axis=0) - for key in key_opt: - if key in t[0].get_arraynames(): - self.data[key]=np.concatenate( [_copy(t[i].get_array(key), traj_skip) for i in range(len(t))], axis=0) - else: - raise RuntimeError(f'first argument cannot be {str(t)}') - self.data['steps']=np.arange(0,self.data['positions'].shape[0]) - self.data['times']=self.data['steps']*dt*traj_skip - self.dt=dt*traj_skip - self.numsites=self.data['positions'].shape[1] - self.pk=pk - self.numsteps=self.data['positions'].shape[0] - - def get_array(self,name): - if name in self.data: - return self.data[name] - else: - raise KeyError(f'cannot find key {name}') - def get_attribute(self,k): - if k=='symbols': - return self.symbols - else: - raise KeyError(f'attribute {k} not present') - def get_arraynames(self): - return self.data.keys() -def analisi2aiida_traj(t,symbols): - ft = FakeAiidaT(t,symbols) - res=aiida.orm.nodes.data.array.trajectory.TrajectoryData() - res.set_attribute('symbols',ft.symbols) - for name in ['steps','positions','cells','times','velocities']: - res.set_array(name,ft.get_array(name)) - return res +# tools for converting an analisi trajectory to a aiida trajectory, or an object that behaves in a similar way + +def analisi2aiida_traj(t, symbols): + ft = FakeAiidaT(t, symbols) + return ft.get_aiida_traj() + -def read_lammps_bin(f,pk=None,symbols={},wrap=False,nsteps=0,start=0,dt=1.0, traj_skip=1): +def read_lammps_bin(f, **kwargs): ''' Read a binary lammps file. traj_skip loads a step every traj_skip steps. Note that this is done at the numpy python level, and the trajectory object still has all steps inside it ''' + return FakeAiidaT.from_lammps_binary(f,**kwargs) - t=pa.Traj(f) - t.setWrapPbc(wrap) - t.setAccessWindowSize(t.getNtimesteps() if nsteps <= 0 else nsteps) - t.setAccessStart(start) - - if pk is None: - import time - pk=str(time.time_ns()) - - return FakeAiidaT(t,symbols_=symbols,dt=dt,pk=pk,traj_skip=traj_skip) def get_type_mask(at): return get_type_mask_from_s(at.symbols) + def get_type_mask_from_s(symbols): - typesa=set(symbols) - masks={} - types=[] - types_array=np.zeros(len(symbols),dtype=int) - itype=0 - for t in typesa: - masks[t]=np.array(symbols)==t - types.append(t) - types_array[masks[t]]=itype - itype+=1 - return types,types_array,masks - -def gen_poscar(at,every=1, start=0,desc=''): - poscars=[] - types,types_array,masks=get_type_mask(at) - pos=at.get_array("positions") - vel=at.get_array("velocities") - cel=at.get_array("cells") - nt={} - for it in types: - nt[it]=np.sum(masks[it]) - for i in range(start,at.numsteps,every): - c=cel[i] - p=f'{desc}\n1.0\n' - for idim in range(3): - p+=f'{c[idim,0]} {c[idim,1]} {c[idim,2]}\n' - for it in types: - p+=it+' ' - p+='\n' - for it in types: - p+=f'{nt[it]} ' - p+='\nCartesian\n' - for it in types: - ps=pos[i,masks[it]] - for iatom in range(nt[it]): - p+= f'{ps[iatom,0]} {ps[iatom,1]} {ps[iatom,2]} \n' - p+='Cartesian\n' - for it in types: - v=vel[i,masks[it]] - for iatom in range(nt[it]): - p+= f'{v[iatom,0]} {v[iatom,1]} {v[iatom,2]} \n' - - poscars.append(p) - return poscars + return FakeAiidaT.get_type_mask_from_s(symbols) + + +def gen_poscar(at, every=1, start=0, desc=''): + traj = FakeAiidaT(at) + return traj.get_poscar(every=every, start=start, desc=desc) + def wrap_dataset(tt): '''This rotates the cell to get a triangular cell matrix, then wraps atomic positions around the cell center in a cubic region of space. Not sure about stress transformation: to check''' - nat=tt['positions'].shape[1]//3 - nts=tt['positions'].shape[0] - wrapped=pa.Trajectory(tt['positions'].reshape(nts,nat,3), - tt['forces'].reshape(nts,nat,3), - np.zeros(nat,dtype='i'), - np.array(tt['cells'].reshape(nts,3,3).transpose((0,2,1)),order='C'), - pa.BoxFormat.CellVectors, - True, - True - ) - Q=wrapped.get_rotation_matrix() - rotated_forces=np.einsum('taj,tij -> tai',tt['forces'].reshape(nts,nat,3),Q).reshape((nts,nat*3)) - rotated_stress=np.einsum('tlm,tli,tmj -> tij',tt['stress'].reshape(nts,3,3),Q,Q).reshape((nts,9)) - rotated_positions=np.einsum('taj,tij -> tai',tt['positions'].reshape(nts,nat,3),Q).reshape((nts,nat*3)) - return {'cells':analisi_cell2box(wrapped.get_box_copy()).reshape((nts,9)), - 'cells_unrotated':tt['cells'], - 'Q':Q.reshape((nts,9)), - 'positions':wrapped.get_positions_copy().reshape((nts,nat*3)), - 'forces':wrapped.get_velocities_copy().reshape((nts,nat*3)), - 'energy':tt['energy'], - 'stress':rotated_stress - } - -def show_traj(tr,wrap=True,fast=1.0,tskip=1): + nat = tt['positions'].shape[1] // 3 + nts = tt['positions'].shape[0] + wrapped = FakeAiidaT({ + 'positions': tt['positions'].reshape(nts, nat, 3), + 'forces': tt['forces'].reshape(nts, nat, 3), + 'cells': tt['cells'].reshape(nts, 3, 3) + }) + wrapped.rotate_qr() + return {'cells': wrapped.get_array('cells').reshape((nts, 9)), + 'cells_unrotated': tt['cells'], + 'positions': wrapped.get_array('positions').reshape((nts, nat * 3)), + 'forces': wrapped.get_array('forces').reshape((nts, nat * 3)), + 'energy': tt['energy'], + 'stress': wrapped.get_array('stresses').reshape((nts,9)) + } + + +def show_traj(tr, wrap=True, fast=1.0, tskip=1): if wrap: - atraj,_=get_analisi_traj_from_aiida(tr,tskip=tskip) + atraj, _ = get_analisi_traj_from_aiida(tr, tskip=tskip) else: - atraj=None - plot=k3d.plot() - plot=show_atraj(tr,atraj,wrap=wrap,plot=plot,fast=fast) + atraj = None + plot = k3d.plot() + plot = show_atraj(tr, atraj, wrap=wrap, plot=plot, fast=fast) plot.display() return plot -def show_atraj(tr,atraj,wrap=True,plot=None,fast=1.0): - atomic_species=list(set(tr.symbols)) - masks=[] + +def show_atraj(tr, atraj, wrap=True, plot=None, fast=1.0): + atomic_species = list(set(tr.symbols)) + masks = [] for sp in atomic_species: - masks.append(np.array(tr.symbols)==sp) + masks.append(np.array(tr.symbols) == sp) if wrap: - pos_m=atraj.get_positions_copy() + pos_m = atraj.get_positions_copy() else: - pos_m=tr.get_array('positions') - t=tr.get_array('steps') - #pos=t.get_positions_copy() - for sp,mask in zip(atomic_species,masks): - plot += k3d.points(positions={str(t/30.0/fast):pos_m[t,mask,:] for t in range(pos_m.shape[0])},size=1.0,name=sp) - print (sp) + pos_m = tr.get_array('positions') + t = tr.get_array('steps') + # pos=t.get_positions_copy() + for sp, mask in zip(atomic_species, masks): + plot += k3d.points(positions={str(t / 30.0 / fast): pos_m[t, mask, :] for t in range(pos_m.shape[0])}, size=1.0, + name=sp) + print(sp) return plot -import matplotlib.colors as colors -from matplotlib import cm - - -def compute_steinhardt(aiida_traj,ranges=[(2.5,3.5),(0.8,1.2),(2.2,3.0),(1.5,1.8)],nthreads=4,skip=10,neigh=[], histogram=True,ls=[6,4],averaged=False,n_segments=1,nbins=100,tskip=1): - atraj=aiida_traj - if isinstance(atraj, FakeAiidaT) and hasattr(atraj,'t'): - if tskip != 1: - raise IndexError(f'trajectory skip not implemented after Trajectory instance is created') - atraj=atraj.t - elif not isinstance(atraj,pa.Trajectory): - atraj,atraj_unw=get_analisi_traj_from_aiida(aiida_traj,tskip=tskip) - elif tskip != 1: - raise IndexError(f'trajectory skip not implemented after Trajectory instance is created') - - stein = None - l=sorted(ls) - - if l[0]<0: - raise IndexError(f'minimum l cannot be negative {l}') - elif l[-1]>10: - raise IndexError(f'spherical harmonics with l>10 are not supported {l}. Please modify and recompile source code.') - - stein=pyanalisi_wrapper(f'SteinhardtOrderParameterHistogram_{l[-1]}',atraj, - ranges, - 1,nbins, - l, - nthreads,skip,histogram,neigh,averaged - ) - - if n_segments==1: - stein.reset(atraj.get_nloaded_timesteps()) - stein.calculate(0) - stein_res=np.array(stein) - return stein_res - elif n_segments>1: - segment_size=max(1,atraj.get_nloaded_timesteps()//n_segments) - stein.reset(segment_size) - print(segment_size) - res=[] - for i in range(0,segment_size*n_segments,segment_size): - print(f'calculating {i}...') - stein.calculate(i) - res.append(np.array(stein,copy=True)) - return np.array(res) - else: - raise IndexError('n_segments must be >= 1') - - - -from mpl_toolkits.axes_grid1 import ImageGrid -from matplotlib.patches import Circle - - -def plt_steinhardt(stein_res,vmin=0.01,figsize=(6.,6.),show=True,transpose=True,xmax=0.2,ymax=0.6,inverted_type_index=False,axs=None,fig=None,single=None,plt_points=None): - if len(stein_res.shape) != 5: - raise RuntimeError('implemented only for 2d histograms!') - cmap=cm.get_cmap('inferno').copy() - cmap.set_bad('black') - nt=stein_res.shape[1] - mask=np.zeros(stein_res.shape) - #mask[:,:,:,-1,-1]=1 - masked=np.ma.masked_array(stein_res,mask=mask) - if axs is None: - if fig is None: - fig = plt.figure(figsize=figsize) - axs = ImageGrid(fig,111,nrows_ncols=(nt,nt) if single is None else (1,1),axes_pad=0.3) - idx=0 - for itype in range(nt): - for jtype in range(nt): - if single is not None: - if (itype,jtype) != single: - continue - if inverted_type_index: - itype, jtype = jtype, itype - try: - axs[idx].imshow(stein_res[0,itype,jtype].transpose() if transpose else stein_res[0,itype,jtype], - norm=colors.LogNorm(vmin=vmin, vmax=masked[0,itype,jtype].max()), - cmap=cmap, - origin='lower',extent=[0.0,1.0,0.0,1.0], - aspect=xmax/ymax) - axs[idx].set_xlim(0,xmax) - axs[idx].set_ylim(0,ymax) - if plt_points: - for x,y,r,c_kw in plt_points: - circ = Circle((x,y),radius=r,**c_kw) - axs[idx].add_patch(circ) - except Exception as e: - print(e) - idx += 1 - if show: - fig.show() - return fig,axs - -def steinhardt_movie(traj,skip=5,neigh=DEFAULT_NEIGH,averaged=True,n_segments=10,plt_steinhardt_kw=DEFAULT_PLT_STEINHARDT_KW, - compute_steinhardt_kw={'nthreads':4},tstein=None): + + + +def compute_steinhardt(aiida_traj, ranges=[(2.5, 3.5), (0.8, 1.2), (2.2, 3.0), (1.5, 1.8)], nthreads=4, skip=10, + neigh=[], histogram=True, ls=[6, 4], averaged=False, n_segments=1, nbins=100, tskip=1): + atraj = FakeAiidaT(aiida_traj,traj_skip=tskip) + a = Analysis(atraj) + + return a.compute_steinhardt(ranges=ranges,nthreads=nthreads,skip=skip,neigh=neigh,histogram=histogram,ls=ls,averaged=averaged, + n_segments=n_segments,nbins=nbins) + +def steinhardt_movie(traj, skip=5, neigh=DEFAULT_NEIGH, averaged=True, n_segments=10, + plt_steinhardt_kw=DEFAULT_PLT_STEINHARDT_KW, + compute_steinhardt_kw={'nthreads': 4}, tstein=None): if tstein is None: - tstein=compute_steinhardt(traj,skip=skip,neigh=neigh,averaged=averaged,n_segments=n_segments,**compute_steinhardt_kw) - class SteinAni: - def __init__(self,tstein,plt_steinhardt_kw): - self.tstein=tstein - self.plt_steinhardt_kw=plt_steinhardt_kw - self.fig,self.axs=plt_steinhardt(self.tstein[0],**self.plt_steinhardt_kw,show=False) - def __call__(self,i): - self.fig,self.axs=plt_steinhardt(self.tstein[i],**self.plt_steinhardt_kw,axs=self.axs,fig=self.fig,show=False) - return self.axs - - stani=SteinAni(tstein,plt_steinhardt_kw) - ani = matplotlib.animation.FuncAnimation( - stani.fig, stani, interval=200, blit=True, save_count=n_segments) - return HTML(ani.to_jshtml()), tstein - -def gofr_movie(atraj,skip=5,n_segments=10,gofr_kw={},plt_gofr_kw={},callax=lambda x:None): - gofr_kw.setdefault('gr_kw',{})['tskip']=skip - if isinstance(atraj, FakeAiidaT) and hasattr(atraj,'t'): - atraj=atraj.t - elif not isinstance(atraj,pa.Trajectory): - atraj,atraj_unw=get_analisi_traj_from_aiida(atraj) - traj=atraj - gofr=do_compute_gr_multi(traj,n_segments=n_segments,**gofr_kw) + tstein = compute_steinhardt(traj, skip=skip, neigh=neigh, averaged=averaged, n_segments=n_segments, + **compute_steinhardt_kw) + + html = SteinPlot.animation(tstein,plt_steinhardt_kw=plt_steinhardt_kw) + return html, tstein + + +def gofr_movie(atraj, skip=5, n_segments=10, gofr_kw={}, plt_gofr_kw={}, callax=lambda x: None): + gofr_kw.setdefault('gr_kw', {})['tskip'] = skip + if isinstance(atraj, FakeAiidaT) and hasattr(atraj, 't'): + atraj = atraj.t + elif not isinstance(atraj, pa.Trajectory): + atraj, atraj_unw = get_analisi_traj_from_aiida(atraj) + traj = atraj + gofr = do_compute_gr_multi(traj, n_segments=n_segments, **gofr_kw) + class SteinAni: - def __init__(self,tstein,plt_gofr_kw,callax): - self.callax=callax - self.tstein=tstein - self.plt_gofr_kw=plt_gofr_kw - self.fig,self.axs,_,_,_,_,_,_,_,_,_,_=do_plots_gr_sh(*self.tstein[0],**self.plt_gofr_kw) - def __call__(self,i): + def __init__(self, tstein, plt_gofr_kw, callax): + self.callax = callax + self.tstein = tstein + self.plt_gofr_kw = plt_gofr_kw + self.fig, self.axs, _, _, _, _, _, _, _, _, _, _ = do_plots_gr_sh(*self.tstein[0], **self.plt_gofr_kw) + + def __call__(self, i): self.axs.clear() - self.fig,self.axs,_,_,_,_,_,_,_,_,_,_=do_plots_gr_sh(*self.tstein[i],**self.plt_gofr_kw,fig_ax=(self.fig,self.axs)) + self.fig, self.axs, _, _, _, _, _, _, _, _, _, _ = do_plots_gr_sh(*self.tstein[i], **self.plt_gofr_kw, + fig_ax=(self.fig, self.axs)) self.callax(self.axs) return [self.axs] - - stani=SteinAni(gofr,plt_gofr_kw,callax) + + stani = SteinAni(gofr, plt_gofr_kw, callax) ani = matplotlib.animation.FuncAnimation( stani.fig, stani, interval=200, blit=True, save_count=n_segments) return HTML(ani.to_jshtml()) + def elastic_c(kbT, box_matrix): "returns elastic constants in GPa" - is_diagonal=False - box_ave=box_matrix.mean(axis=0) - if np.count_nonzero(box_ave-np.diag(np.diagonal(box_ave))) == 0: - is_diagonal=True - vol_mean=np.linalg.det(box_ave) - G=np.einsum('tji,tjl->til',box_matrix,box_matrix) - box_inv=np.linalg.inv(box_matrix) - box_ave_inv=box_inv.mean(axis=0) - strain=0.5*(np.einsum('ji,tjl,lm->tim',box_ave_inv,G,box_ave_inv)-np.eye(3)) + is_diagonal = False + box_ave = box_matrix.mean(axis=0) + if np.count_nonzero(box_ave - np.diag(np.diagonal(box_ave))) == 0: + is_diagonal = True + vol_mean = np.linalg.det(box_ave) + G = np.einsum('tji,tjl->til', box_matrix, box_matrix) + box_inv = np.linalg.inv(box_matrix) + box_ave_inv = box_inv.mean(axis=0) + strain = 0.5 * (np.einsum('ji,tjl,lm->tim', box_ave_inv, G, box_ave_inv) - np.eye(3)) strain_0 = strain.mean(axis=0) - nt=strain.shape[0] - #pick=[0,1,2,4,5,8] # xx, xy, xz, yy, yz, zz - pick=[0,4,8,1,2,5] # xx, yy, zz, xy, xz, yz + nt = strain.shape[0] + # pick=[0,1,2,4,5,8] # xx, xy, xz, yy, yz, zz + pick = [0, 4, 8, 1, 2, 5] # xx, yy, zz, xy, xz, yz if is_diagonal: - pick=pick[:3] - CS_minus1=np.einsum('ti,tk->ik',(strain-strain_0).reshape((nt,9))[:,pick],(strain-strain_0).reshape((nt,9))[:,pick])*vol_mean/kbT/nt*1e9 - CS=np.linalg.inv(CS_minus1) - with np.printoptions(precision=3, suppress=False,linewidth=150): + pick = pick[:3] + CS_minus1 = np.einsum('ti,tk->ik', (strain - strain_0).reshape((nt, 9))[:, pick], + (strain - strain_0).reshape((nt, 9))[:, pick]) * vol_mean / kbT / nt * 1e9 + CS = np.linalg.inv(CS_minus1) + with np.printoptions(precision=3, suppress=False, linewidth=150): print('xx, yy, zz, xy, xz, yz : GPa') - print (CS) + print(CS) print('============') print('1/GPa') print(np.linalg.det(CS_minus1)) print(CS_minus1) - if is_diagonal: #fill with zeros the 6x6 matrix not involved in the calculation - CS2=np.zeros((6,6)) - CS_minus12=np.zeros((6,6)) - CS2[:3,:3]=CS - CS_minus12[:3,:3]=CS_minus1 - CS=CS2 - CS_minus1=CS_minus12 + if is_diagonal: # fill with zeros the 6x6 matrix not involved in the calculation + CS2 = np.zeros((6, 6)) + CS_minus12 = np.zeros((6, 6)) + CS2[:3, :3] = CS + CS_minus12[:3, :3] = CS_minus1 + CS = CS2 + CS_minus1 = CS_minus12 return CS, CS_minus1 -def hist2gofr(gr_N,gr_dr,gr_0,gofr): - rs_m=np.arange(gr_N)*gr_dr+gr_0 - rs_p=(np.arange(gr_N)+1)*gr_dr+gr_0 - vols=4*np.pi/3*(rs_p**3-rs_m**3) - return gofr/vols -def peak_width(gr,NH,NH_thre,gr_r2i,min_spread=0.1): - ''' - get the width of the first peak in the gr at a given height - ''' - idx_23=0 - while gr[0,NH,idx_23] < NH_thre: - idx_23 += 1 - idx_spread_low=idx_23 - while gr[0,NH,idx_23] > NH_thre*0.9: - if idx_23 == gr.shape[2]-1 or (idx_23 -idx_spread_low > gr_r2i(min_spread) and gr[0,NH,idx_23] < NH_thre): - break - idx_23 += 1 - idx_spread_hi=idx_23 - return idx_spread_low,idx_spread_hi +def hist2gofr(gr_N, gr_dr, gr_0, gofr): + return Analysis.hist2gofr(gr_N, gr_dr, gr_0, gofr) -def analyze_peak(gr,r0,peak_fractional_height,gr_r2i,gr_i2r): - ''' - find the peak around r0 in the g(r) array and get some info - ''' - maxs=np.argmax(gr,axis=2) - peaks=gr_i2r(maxs) - - #find nearest peak to r0 - NH=np.argmin((peaks-r0)**2) - #get spread at 'peak_fractional_height' of peak height - NH_peak_idx=maxs[0,NH] - NH_peak_val=gr[0,NH,NH_peak_idx] - NH_thre=NH_peak_val*peak_fractional_height - #get width of peak - idx_spread_low,idx_spread_hi=peak_width(gr,NH,NH_thre,gr_r2i) - spread_low=gr_i2r(idx_spread_low) - spread_hi=gr_i2r(idx_spread_hi) - w23h=spread_hi-spread_low - return idx_spread_low,idx_spread_hi,spread_low,spread_hi,w23h,NH_thre,NH_peak_val,NH_peak_idx,peaks,NH - -def get_conv_functs(gr_0,gr_dr): - def gr_r2i(r): - return int((r-gr_0)/gr_dr) - def gr_i2r(i): - return gr_0+i*gr_dr - return np.vectorize(gr_r2i), np.vectorize(gr_i2r) - - -def do_compute_gr_multi(t,gr_kw={'tskip':10},n_segments=1,gr_0=0.5,gr_end=3.8,gr_N=150): + + + +def do_compute_gr_multi(t, gr_kw={'tskip': 10}, n_segments=1, gr_0=0.5, gr_end=3.8, gr_N=150): ''' computes the g(r) pair correlation function and find the peaks of it ''' - nts=t.get_nloaded_timesteps() - param_gr=(nts,gr_0,gr_end,gr_N) - #first calculate g(r), get the N-H peak, estabilish the range of sh correlation f - gr_dr=(gr_end-gr_0)/gr_N - - #get utility functions for converting from index to r value and back - gr_r2i,gr_i2r=get_conv_functs(gr_0,gr_dr) - - gofr=analyze_gofr(t,0,nts,gr_0,gr_end,gr_N, - tmax=1,**gr_kw,n_segments=n_segments) #only g(r) - #from the histogram generate the g(r) -- divide by the volumes of the spherical shells - res=[] + nts = t.get_nloaded_timesteps() + param_gr = (nts, gr_0, gr_end, gr_N) + # first calculate g(r), get the N-H peak, estabilish the range of sh correlation f + gr_dr = (gr_end - gr_0) / gr_N + + gofr = analyze_gofr(t, 0, nts, gr_0, gr_end, gr_N, + tmax=1, **gr_kw, n_segments=n_segments) # only g(r) + # from the histogram generate the g(r) -- divide by the volumes of the spherical shells + res = [] for i in range(n_segments): - gr=hist2gofr(gr_N,gr_dr,gr_0,gofr if n_segments==1 else gofr[i]) - - - #find nearest peak to 1.0 (N-H bond) - idx_spread_low,idx_spread_hi,spread_low,spread_hi,w23h,NH_thre,NH_peak_val,NH_peak_idx,peaks,NH = analyze_peak(gr,1.0,0.5,gr_r2i,gr_i2r) - - #print(f'width at 1/2 of height: {w23h}') - #calculate sh correlations of hydrogen peak - sh_low=spread_low-w23h*0.1 - sh_hi=spread_hi+w23h*0.1 - NH_peak=(NH_thre,spread_low,spread_hi,peaks,NH,w23h,NH_peak_val) - res.append( (None,param_gr,gofr if n_segments==1 else gofr[i],(None,None),None,NH_peak)) + gr = hist2gofr(gr_N, gr_dr, gr_0, gofr if n_segments == 1 else gofr[i]) + + # find nearest peak to 1.0 (N-H bond) + idx_spread_low, idx_spread_hi, spread_low, spread_hi, w23h, NH_thre, NH_peak_val, NH_peak_idx, peaks, NH = Analysis.compute_peaks_gofr( + gr, 1.0,gr_0,gr_dr, 0.5) + + # print(f'width at 1/2 of height: {w23h}') + # calculate sh correlations of hydrogen peak + sh_low = spread_low - w23h * 0.1 + sh_hi = spread_hi + w23h * 0.1 + NH_peak = (NH_thre, spread_low, spread_hi, peaks, NH, w23h, NH_peak_val) + res.append((None, param_gr, gofr if n_segments == 1 else gofr[i], (None, None), None, NH_peak)) return res -def do_compute_gr_sh(t,times,do_sh=True,neigh=[],analyze_sh_kw={'tskip':50},gr_kw={'tskip':10}): +def do_compute_gr_sh(traj, times, do_sh=True, neigh=[], analyze_sh_kw={'tskip': 50}, gr_kw={'tskip': 10}): ''' computes the g(r) pair correlation function and the spherical harmonics correlation function computed around the N-H peak of the g(r), at around 1.0 Angstrom ''' - gr_0=0.5 - gr_end=3.8 - gr_N=150 - nts=t.get_nloaded_timesteps() - param_gr=(nts,gr_0,gr_end,gr_N) - DT_PS=times[1]-times[0] - #first calculate g(r), get the N-H peak, estabilish the range of sh correlation f - gr_dr=(gr_end-gr_0)/gr_N - - #get utility functions for converting from index to r value and back - gr_r2i,gr_i2r=get_conv_functs(gr_0,gr_dr) - - gofr=analyze_gofr(t,0,nts,gr_0,gr_end,gr_N, - tmax=1,**gr_kw) #only g(r) - #from the histogram generate the g(r) -- divide by the volumes of the spherical shells - gr=hist2gofr(gr_N,gr_dr,gr_0,gofr) - - - #find nearest peak to 1.0 (N-H bond) - idx_spread_low,idx_spread_hi,spread_low,spread_hi,w23h,NH_thre,NH_peak_val,NH_peak_idx,peaks,NH = analyze_peak(gr,1.0,0.5,gr_r2i,gr_i2r) - - #print(f'width at 1/2 of height: {w23h}') - #calculate sh correlations of hydrogen peak - sh_low=spread_low-w23h*0.1 - sh_hi=spread_hi+w23h*0.1 + gr_0 = 0.5 + gr_end = 3.8 + gr_N = 150 + t=traj.get_analisi_traj(wrapped=True) + nts = t.get_nloaded_timesteps() + param_gr = (nts, gr_0, gr_end, gr_N) + DT_PS = times[1] - times[0] + # first calculate g(r), get the N-H peak, estabilish the range of sh correlation f + gr_dr = (gr_end - gr_0) / gr_N + + gofr = analyze_gofr(t, 0, nts, gr_0, gr_end, gr_N, + tmax=1, **gr_kw) # only g(r) + # from the histogram generate the g(r) -- divide by the volumes of the spherical shells + gr = hist2gofr(gr_N, gr_dr, gr_0, gofr) + + # find nearest peak to 1.0 (N-H bond) + idx_spread_low, idx_spread_hi, spread_low, spread_hi, w23h, NH_thre, NH_peak_val, NH_peak_idx, peaks, NH = Analysis.compute_peaks_gofr( + gr, 1.0, gr_0,gr_dr, 0.5) + + # print(f'width at 1/2 of height: {w23h}') + # calculate sh correlations of hydrogen peak + sh_low = spread_low - w23h * 0.1 + sh_hi = spread_hi + w23h * 0.1 if do_sh: - sh=analyze_sh(t,0,nts,sh_low,sh_hi,1,tmax=int(0.5/DT_PS),sann=neigh,**analyze_sh_kw) + sh = analyze_sh(t, 0, nts, sh_low, sh_hi, 1, tmax=int(0.5 / DT_PS), sann=neigh, **analyze_sh_kw) else: - sh=None - NH_peak=(NH_thre,spread_low,spread_hi,peaks,NH,w23h,NH_peak_val) - param_sh = (sh_low,sh_hi) - return times,param_gr,gofr,param_sh,sh,NH_peak - -def do_plots_gr_sh(times,param_gr,gofr,param_sh,sh,NH_peak,fig_ax=None,plot_gr_kw={},plot_sh_band=False): - - nts,gr_0,gr_end,gr_N=param_gr - sh_low,sh_hi=param_sh - NH_thre,spread_low,spread_hi,peaks,NH,w23h,NH_peak_val=NH_peak - - gr_dr=(gr_end-gr_0)/gr_N - #get utility functions for converting from index to r value and back - gr_r2i,gr_i2r=get_conv_functs(gr_0,gr_dr) - - fig_gr,ax_gr=plot_gofr(gr_0,gr_end,gofr,fig_ax=fig_ax,**plot_gr_kw) + sh = None + NH_peak = (NH_thre, spread_low, spread_hi, peaks, NH, w23h, NH_peak_val) + param_sh = (sh_low, sh_hi) + return times, param_gr, gofr, param_sh, sh, NH_peak + + +def do_plots_gr_sh(times, param_gr, gofr, param_sh, sh, NH_peak, fig_ax=None, plot_gr_kw={}, plot_sh_band=False): + nts, gr_0, gr_end, gr_N = param_gr + sh_low, sh_hi = param_sh + NH_thre, spread_low, spread_hi, peaks, NH, w23h, NH_peak_val = NH_peak + + gr_dr = (gr_end - gr_0) / gr_N + fig_gr, ax_gr = plot_gofr(gr_0, gr_end, gofr, fig_ax=fig_ax, **plot_gr_kw) ax_gr.grid() if plot_sh_band: - if sh_low is not None: - ax_gr.axvline(sh_low,color='r') + if sh_low is not None: + ax_gr.axvline(sh_low, color='r') if sh_hi is not None: - ax_gr.axvline(sh_hi,color='r') - - #annotate N-N peak - - #from the histogram generate the g(r) -- divide by the volumes of the spherical shells - gr=hist2gofr(gr_N,gr_dr,gr_0,gofr) - def annotate_peak(r0,h,ax_gr): - _,_,NN_spread_low,NN_spread_hi,w2NN,NN_thre,NN_peak_val,NN_peak_idx,peaks,NN = analyze_peak(gr,r0,h,gr_r2i,gr_i2r) - ax_gr.hlines(NN_thre,NN_spread_low,NN_spread_hi) - - #ax_gr.annotate(f'{w2NN:.2f}',(peaks[0,NN]*0.95,NN_thre*1.05)) - #ax_gr.annotate(f'{peaks[0,NN]:.2f}',(peaks[0,NN]*0.95,NN_peak_val*1.02)) - ax_gr.annotate(f'{peaks[0,NN]:.2f}\n{w2NN:.2f}',(peaks[0,NN]*0.95,NN_thre*1.05)) - return w2NN,peaks[0,NN] - #find nearest peak to 2.6 (N-N bond) - wNN,rNN=annotate_peak(2.6,0.5,ax_gr) - #find nearest peak to 1.7 (H-H bond) - wHH,rHH=annotate_peak(1.7,3.0/4,ax_gr) - #find nearest peak to 1.0 (N-H bond) - wNH,rNH=annotate_peak(1.0,0.5,ax_gr) + ax_gr.axvline(sh_hi, color='r') + + # annotate N-N peak + + # from the histogram generate the g(r) -- divide by the volumes of the spherical shells + gr = hist2gofr(gr_N, gr_dr, gr_0, gofr) + + def annotate_peak(r0, h, ax_gr): + _, _, NN_spread_low, NN_spread_hi, w2NN, NN_thre, NN_peak_val, NN_peak_idx, peaks, NN = Analysis.compute_peaks_gofr(gr, r0, + gr_0,gr_dr, h) + ax_gr.hlines(NN_thre, NN_spread_low, NN_spread_hi) + + # ax_gr.annotate(f'{w2NN:.2f}',(peaks[0,NN]*0.95,NN_thre*1.05)) + # ax_gr.annotate(f'{peaks[0,NN]:.2f}',(peaks[0,NN]*0.95,NN_peak_val*1.02)) + ax_gr.annotate(f'{peaks[0, NN]:.2f}\n{w2NN:.2f}', (peaks[0, NN] * 0.95, NN_thre * 1.05)) + return w2NN, peaks[0, NN] + + # find nearest peak to 2.6 (N-N bond) + wNN, rNN = annotate_peak(2.6, 0.5, ax_gr) + # find nearest peak to 1.7 (H-H bond) + wHH, rHH = annotate_peak(1.7, 3.0 / 4, ax_gr) + # find nearest peak to 1.0 (N-H bond) + wNH, rNH = annotate_peak(1.0, 0.5, ax_gr) if sh is not None: try: - fig_sh,ax_sh,axins_sh,fit_sh=plot_sh(sh_low,sh_hi,times,sh,0,1,0,log=False,pre_fit=0.4) + fig_sh, ax_sh, axins_sh, fit_sh = plot_sh(sh_low, sh_hi, times, sh, 0, 1, 0, log=False, pre_fit=0.4) except: try: - fig_sh,ax_sh,axins_sh,fit_sh=plot_sh(sh_low,sh_hi,times,sh,0,1,0,log=False,pre_fit=0.1) + fig_sh, ax_sh, axins_sh, fit_sh = plot_sh(sh_low, sh_hi, times, sh, 0, 1, 0, log=False, pre_fit=0.1) except: - fig_sh,ax_sh,axins_sh,fit_sh=plot_sh(sh_low,sh_hi,times,sh,0,1,0,log=False) + fig_sh, ax_sh, axins_sh, fit_sh = plot_sh(sh_low, sh_hi, times, sh, 0, 1, 0, log=False) ax_sh.grid() else: - fig_sh,ax_sh,axins_sh,fit_sh=None,None,None,None - return fig_gr,ax_gr,fig_sh,ax_sh,axins_sh,fit_sh,wNN,rNN,wHH,rHH,wNH,rNH - -def do_compute_msd(t_unw,times,msd_kw={}): - nts=t_unw.get_nloaded_timesteps() - msd=analyze_msd(t_unw,0,nts,**msd_kw) - #compute slope - DT_PS=times[1]-times[0] - msd_start=max(min(int(0.3/DT_PS),msd.shape[0]-100),0) - coeffs=np.polyfit(times[msd_start:msd.shape[0]]-times[0],msd[msd_start:,0,:],1) - return msd,times,msd_start,coeffs - -def do_plots_msd(msd,times,msd_start,coeffs,do_fit=True,fig_ax=None,subplots_kw=DEFAULT_SUBPLOTS_KW): - DT_PS=times[1]-times[0] - fig,ax=plot_msd(times,msd,0,fig_ax=fig_ax,subplots_kw=subplots_kw) + fig_sh, ax_sh, axins_sh, fit_sh = None, None, None, None + return fig_gr, ax_gr, fig_sh, ax_sh, axins_sh, fit_sh, wNN, rNN, wHH, rHH, wNH, rNH + + +def do_compute_msd(traj, times, msd_kw={}): + t_unw=traj.get_analisi_traj(wrapped=False) + nts = t_unw.get_nloaded_timesteps() + msd = analyze_msd(t_unw, 0, nts, **msd_kw) + # compute slope + DT_PS = times[1] - times[0] + msd_start = max(min(int(0.3 / DT_PS), msd.shape[0] - 100), 0) + coeffs = np.polyfit(times[msd_start:msd.shape[0]] - times[0], msd[msd_start:, 0, :], 1) + return msd, times, msd_start, coeffs + + +def do_plots_msd(msd, times, msd_start, coeffs, do_fit=True, fig_ax=None, subplots_kw=DEFAULT_SUBPLOTS_KW): + DT_PS = times[1] - times[0] + fig, ax = plot_msd(times, msd, 0, fig_ax=fig_ax, subplots_kw=subplots_kw) ax.grid() if do_fit: - t0=times[msd_start]-times[0] - tf=times[msd.shape[0]-1]-times[0] - lines=[[(t0,coeffs[0,0]*t0+coeffs[1,0]),(tf,coeffs[0,0]*tf+coeffs[1,0])], - [(t0,coeffs[0,1]*t0+coeffs[1,1]),(tf,coeffs[0,1]*tf+coeffs[1,1])]] - lc=mc.LineCollection(lines,zorder=-1) + t0 = times[msd_start] - times[0] + tf = times[msd.shape[0] - 1] - times[0] + lines = [[(t0, coeffs[0, 0] * t0 + coeffs[1, 0]), (tf, coeffs[0, 0] * tf + coeffs[1, 0])], + [(t0, coeffs[0, 1] * t0 + coeffs[1, 1]), (tf, coeffs[0, 1] * tf + coeffs[1, 1])]] + lc = mc.LineCollection(lines, zorder=-1) ax.add_collection(lc) - ax.annotate(f'D={coeffs[0,0]/6:.2f}',lines[0][1]) - ax.annotate(f'D={coeffs[0,1]/6:.2f}',lines[1][1]) - return fig,ax + ax.annotate(f'D={coeffs[0, 0] / 6:.2f}', lines[0][1]) + ax.annotate(f'D={coeffs[0, 1] / 6:.2f}', lines[1][1]) + return fig, ax + -def inspect(traj, only_cell=False,plot_traj=True,plot=True, - do_sh=True,do_density=False, plot_sh_h=True, - show_traj_dyn=False,dt='auto', +def inspect(traj, only_cell=False, plot_traj=True, plot=True, + do_sh=True, do_density=False, plot_sh_h=True, + show_traj_dyn=False, dt='auto', neigh=DEFAULT_NEIGH, - plot_st_kw={'transpose':True,'xmax':.20,'ymax':.60}, - analyze_sh_kw={'tskip':50}, - compute_steinhardt_kw={'skip':10}, + plot_st_kw={'transpose': True, 'xmax': .20, 'ymax': .60}, + analyze_sh_kw={'tskip': 50}, + compute_steinhardt_kw={'skip': 10}, msd_kw={}, - gr_kw={'tskip':10}, - nthreads=4,save_data=False,tskip=1, + gr_kw={'tskip': 10}, + nthreads=4, save_data=False, tskip=1, atomic_density_kwargs={}, msd_plot_kw={}, subplots_kw=DEFAULT_SUBPLOTS_KW): - results={} - analyze_sh_kw['nthreads']=nthreads - compute_steinhardt_kw['nthreads']=nthreads - msd_kw['nthreads']=nthreads - gr_kw['nthreads']=nthreads - - - natoms=traj.numsites - #conversion factor from eV to K - k_b=8.617333262145e-5 #eV/K - eV_to_K=2/(3*k_b*natoms) - - #print('traj arraynames = {}'.format(traj.get_arraynames())) + results = {} + analyze_sh_kw['nthreads'] = nthreads + compute_steinhardt_kw['nthreads'] = nthreads + msd_kw['nthreads'] = nthreads + gr_kw['nthreads'] = nthreads + + natoms = traj.numsites + # conversion factor from eV to K + k_b = 8.617333262145e-5 # eV/K + eV_to_K = 2 / (3 * k_b * natoms) + + # print('traj arraynames = {}'.format(traj.get_arraynames())) if 'ionic_temperature' in traj.get_arraynames(): - temp=traj.get_array('ionic_temperature') - T_mean=temp.mean() - results['T']=T_mean + temp = traj.get_array('ionic_temperature') + T_mean = temp.mean() + results['T'] = T_mean else: - T_mean=float('nan') + T_mean = float('nan') if 'pressure' in traj.get_arraynames(): - press=traj.get_array('pressure') - press_mean=press.mean() - results['P']=press_mean + press = traj.get_array('pressure') + press_mean = press.mean() + results['P'] = press_mean else: press_mean = float('nan') if dt is None and 'times' in traj.get_arraynames(): - t=traj.get_array('times') - elif isinstance(dt,float): - t=np.arange(temp.shape[0])*dt - elif dt=='auto' and 'times' in traj.get_arraynames(): - t=traj.get_array('times') - dt=(t[1]-t[0]) - t=np.arange(temp.shape[0])*dt + t = traj.get_array('times') + elif isinstance(dt, float): + t = np.arange(traj.get_array('positions').shape[0]) * dt + elif dt == 'auto' and 'times' in traj.get_arraynames(): + t = traj.get_array('times') + dt = (t[1] - t[0]) + t = np.arange(t.shape[0]) * dt + def t_to_timestep(x): - return np.interp(x,t,np.arange(t.shape[0])) + return np.interp(x, t, np.arange(t.shape[0])) + def timestep_to_t(x): - return np.interp(x,np.arange(t.shape[0]),t) - - - DT_FS=(t[1]-t[0])*1e3 - cell_0 =traj.get_array('cells') - volume=np.linalg.det(cell_0) - #print('cell(t=0) = {}'.format(cell_0)) - #print('cell_volume(t=0) = {} A^3'.format(volume)) - cell_0=traj.get_array('cells').mean(axis=0) - volume=np.linalg.det(cell_0) - results['V']=volume + return np.interp(x, np.arange(t.shape[0]), t) + + DT_FS = (t[1] - t[0]) * 1e3 + cell_0 = traj.get_array('cells') + volume = np.linalg.det(cell_0) + # print('cell(t=0) = {}'.format(cell_0)) + # print('cell_volume(t=0) = {} A^3'.format(volume)) + cell_0 = traj.get_array('cells').mean(axis=0) + volume = np.linalg.det(cell_0) + results['V'] = volume if 'electronic_kinetic_energy' in traj.get_arraynames(): - results['ekinc']=traj.get_array('electronic_kinetic_energy').mean() - if 'energy_constant_motion' in traj.get_arraynames(): - results['e_constant']=traj.get_array('energy_constant_motion').mean() - results['DT_FS']=DT_FS + results['ekinc'] = traj.get_array('electronic_kinetic_energy').mean() + if 'energy_constant_motion' in traj.get_arraynames(): + results['e_constant'] = traj.get_array('energy_constant_motion').mean() + results['DT_FS'] = DT_FS - plt_fname_pre=FIGURE_PATH+'/inspect_result' + plt_fname_pre = FIGURE_PATH + '/inspect_result' if 'ionic_temperature' in traj.get_arraynames() and 'pressure' in traj.get_arraynames(): display(HTML(f'

{T_mean:.0f}K {press_mean:.0f}GPa

')) - print('DT_FS = {:.3f}fs, T = {:.2f}K, P = {:.2f}GPa'.format(DT_FS,T_mean,press_mean)) - plt_fname_pre=FIGURE_PATH+f'/{T_mean:.0f}K_{press_mean:.0f}GPa_{DT_FS:.3f}_' + print('DT_FS = {:.3f}fs, T = {:.2f}K, P = {:.2f}GPa'.format(DT_FS, T_mean, press_mean)) + plt_fname_pre = FIGURE_PATH + f'/{T_mean:.0f}K_{press_mean:.0f}GPa_{DT_FS:.3f}_' print('traj pk = {}'.format(traj.pk)) - plt_fname_suff='.pdf' + plt_fname_suff = '.pdf' if traj.pk != None: - pickle_dump_name=f'traj_{traj.pk}' + pickle_dump_name = f'traj_{traj.pk}' else: - pickle_dump_name=f'traj_{T_mean:.2f}K{press_mean:.2f}GPa' - + pickle_dump_name = f'traj_{T_mean:.2f}K{press_mean:.2f}GPa' if plot_traj and plot: print('cell_average = {}'.format(cell_0)) print('cell_average_volume = {} A^3'.format(volume)) - plt_key(traj,'electronic_kinetic_energy',eV_to_K,ylabel='K',dt=dt) - plt_key(traj,'energy_constant_motion',dt=dt) - plt_key(traj,'ionic_temperature',ylabel='K',dt=dt) - plt_key(traj,'pressure',ylabel='GPa',dt=dt) + plt_key(traj, 'electronic_kinetic_energy', eV_to_K, ylabel='K', dt=dt) + plt_key(traj, 'energy_constant_motion', dt=dt) + plt_key(traj, 'ionic_temperature', ylabel='K', dt=dt) + plt_key(traj, 'pressure', ylabel='GPa', dt=dt) plt.show() - atraj,atraj_unw=get_analisi_traj_from_aiida(traj,tskip=tskip) - + #atraj, atraj_unw = get_analisi_traj_from_aiida(traj, tskip=tskip) if show_traj_dyn: - traj_dyn=show_atraj(traj,atraj,wrap=False) - - + traj_dyn = traj.show_traj() + if not only_cell: - if plot and do_density: - res=atomic_density(atraj) - plot_=density_field(*res,**atomic_density_kwargs) - - #histogram of steinhardt parameters - if plot_sh_h: - sh_h_pickle=pickle_dump_name+'_sh_h.pickle' + res = atomic_density(traj) + plot_ = density_field(*res, **atomic_density_kwargs) + + # histogram of steinhardt parameters + if plot_sh_h: + sh_h_pickle = pickle_dump_name + '_sh_h.pickle' sh_h = pickle_or_unpickle(sh_h_pickle) if sh_h is None: - sh_h=compute_steinhardt(atraj,neigh=neigh,averaged=True,**compute_steinhardt_kw) - fig_shh,axs_shh=plt_steinhardt(sh_h,**plot_st_kw) - fig_shh.savefig(plt_fname_pre+'steinhardt'+plt_fname_suff) - + sh_h = compute_steinhardt(traj, neigh=neigh, averaged=True, **compute_steinhardt_kw) + fig_shh, axs_shh = SteinPlot.plot(sh_h, **plot_st_kw) + fig_shh.savefig(plt_fname_pre + 'steinhardt' + plt_fname_suff) - #g of r / spherical correlations - grsh_pickle=pickle_dump_name+ ('_gofr_sh.pickle' if do_sh else '_gofr.pickle') + # g of r / spherical correlations + grsh_pickle = pickle_dump_name + ('_gofr_sh.pickle' if do_sh else '_gofr.pickle') gofrsh = pickle_or_unpickle(grsh_pickle) if gofrsh is None: - gofrsh=do_compute_gr_sh(atraj,t[::tskip],do_sh=do_sh,neigh=neigh,analyze_sh_kw=analyze_sh_kw,gr_kw=gr_kw) - pickle_or_unpickle(grsh_pickle,analisi = gofrsh) - - #msd - msd_pickle=pickle_dump_name+'_msd.pickle' + gofrsh = do_compute_gr_sh(traj, t, do_sh=do_sh, neigh=neigh, analyze_sh_kw=analyze_sh_kw, + gr_kw=gr_kw) + pickle_or_unpickle(grsh_pickle, analisi=gofrsh) + + # msd + msd_pickle = pickle_dump_name + '_msd.pickle' msd = pickle_or_unpickle(msd_pickle) if msd is None: - msd=do_compute_msd(atraj_unw,t[::tskip],msd_kw=msd_kw) - pickle_or_unpickle(msd_pickle,analisi = msd) + msd = do_compute_msd(traj, t, msd_kw=msd_kw) + pickle_or_unpickle(msd_pickle, analisi=msd) if save_data: - results['data_gofrsh']=gofrsh - results['data_msd']=msd - - - wNN,rNN,wHH,rHH,wNH,rNH=None,None,None,None,None,None + results['data_gofrsh'] = gofrsh + results['data_msd'] = msd + + wNN, rNN, wHH, rHH, wNH, rNH = None, None, None, None, None, None if plot: - fig_gr,ax_gr,fig_sh,ax_sh,axins_sh,fit_sh,wNN,rNN,wHH,rHH,wNH,rNH=do_plots_gr_sh(*gofrsh) + fig_gr, ax_gr, fig_sh, ax_sh, axins_sh, fit_sh, wNN, rNN, wHH, rHH, wNH, rNH = do_plots_gr_sh(*gofrsh) fig_gr.show() - fig_gr.savefig(plt_fname_pre+'gr'+plt_fname_suff) + fig_gr.savefig(plt_fname_pre + 'gr' + plt_fname_suff) if fig_sh is not None: fig_sh.show() - fig_sh.savefig(plt_fname_pre+'sh'+plt_fname_suff) - fig_msd,ax_msd=do_plots_msd(*msd,**msd_plot_kw) + fig_sh.savefig(plt_fname_pre + 'sh' + plt_fname_suff) + fig_msd, ax_msd = do_plots_msd(*msd, **msd_plot_kw) fig_msd.show() - fig_msd.savefig(plt_fname_pre+'msd'+plt_fname_suff) - - results['msd']=msd[3][0,:] - - NH_thre,spread_low,spread_hi,peaks,NH,w23h,NH_peak_val=gofrsh[5] - results['NH_peak_width']=w23h - results['NH_peak_pos']=peaks[0,NH] - results['NN_peak_width']=wNN - results['NN_peak_pos']=rNN - results['HH_peak_width']=wHH - results['HH_peak_pos']=rHH - - cells_transition=atraj.get_box_copy() - cell_transition_ave=cells_transition.mean(axis=0) - print(f'average cell xx,yy,zz = {2*cell_transition_ave[3:6]}') - a = cell_transition_ave[3]/4+cell_transition_ave[4]/(3*3**.5) - hcpx = a*4 - hcpy = a*3*3**.5 - hcpz = a*2*6**.5 + fig_msd.savefig(plt_fname_pre + 'msd' + plt_fname_suff) + + results['msd'] = msd[3][0, :] + + NH_thre, spread_low, spread_hi, peaks, NH, w23h, NH_peak_val = gofrsh[5] + results['NH_peak_width'] = w23h + results['NH_peak_pos'] = peaks[0, NH] + results['NN_peak_width'] = wNN + results['NN_peak_pos'] = rNN + results['HH_peak_width'] = wHH + results['HH_peak_pos'] = rHH + + cells_transition = traj.get_analisi_traj().get_box_copy() + cell_transition_ave = cells_transition.mean(axis=0) + print(f'average cell xx,yy,zz = {2 * cell_transition_ave[3:6]}') + a = cell_transition_ave[3] / 4 + cell_transition_ave[4] / (3 * 3 ** .5) + hcpx = a * 4 + hcpy = a * 3 * 3 ** .5 + hcpz = a * 2 * 6 ** .5 print(f'close packed ideal structure = {hcpx} {hcpy} {hcpz}') - print(f'close packed ideal structure delta = {hcpx-2*cell_transition_ave[3]} {hcpy-2*cell_transition_ave[4]} {hcpz-2*cell_transition_ave[5]}') - if cells_transition.shape[1]>6: + print( + f'close packed ideal structure delta = {hcpx - 2 * cell_transition_ave[3]} {hcpy - 2 * cell_transition_ave[4]} {hcpz - 2 * cell_transition_ave[5]}') + if cells_transition.shape[1] > 6: print(f'average cell xy,xz,yz = {cell_transition_ave[6:]}') - - if (cells_transition!=cells_transition[0]).any(): + + if (cells_transition != cells_transition[0]).any(): if plot: - fig, ax = plt.subplots(nrows=2,sharex=True,**subplots_kw) - lines=ax[0].plot(t[::tskip],cells_transition[:,3:6]*2) - for i,c in enumerate(['x','y','z']): + fig, ax = plt.subplots(nrows=2, sharex=True, **subplots_kw) + lines = ax[0].plot(t, cells_transition[:, 3:6] * 2) + for i, c in enumerate(['x', 'y', 'z']): lines[i].set_label(c) - ax[0].plot(t[::tskip],(2*cells_transition[:,3:6]).prod(axis=1)**(1.0/3.0),label=r'V$^{1/3}$') + ax[0].plot(t, (2 * cells_transition[:, 3:6]).prod(axis=1) ** (1.0 / 3.0), label=r'V$^{1/3}$') ax[0].grid() ax[0].set_title('cell parameters: cell size (up) and cell tilt (down)') ax[1].set_xlabel('t (ps)') ax[0].set_ylabel('(A)') - ax[0].secondary_xaxis('top',functions=(t_to_timestep,timestep_to_t)) + ax[0].secondary_xaxis('top', functions=(t_to_timestep, timestep_to_t)) ax[0].legend() - if cells_transition.shape[1]>6: - lines=ax[1].plot(t[::tskip],cells_transition[:,6:]) + if cells_transition.shape[1] > 6: + lines = ax[1].plot(t, cells_transition[:, 6:]) ax[1].grid() ax[1].set_ylabel('(A)') - for i,c in enumerate(['xy','xz','yz']): + for i, c in enumerate(['xy', 'xz', 'yz']): lines[i].set_label(c) ax[1].legend() fig.show() - fig.savefig(plt_fname_pre+'cell'+plt_fname_suff) + fig.savefig(plt_fname_pre + 'cell' + plt_fname_suff) try: - uma=1.66e-27 #kg - mcell=(sum( [ 1 for i in traj.get_attribute('symbols') if i=='H']) + sum( [ 14 for i in traj.get_attribute('symbols') if i=='N']))*uma - def vs_C(traj,nsteps): - density=(mcell/np.linalg.det(traj.get_array('cells')[:nsteps:tskip]*1e-10)).mean() - CS,CSm=elastic_c(traj.get_array('ionic_temperature')[:nsteps:tskip].mean()*1.38064852e-23, - traj.get_array('cells')[:nsteps:tskip]*1e-10) - vs=((CS[0,0]+4.0/3*CS[3,3])*1e9/density)**.5 - print ('sqrt((C_{xx,xx} + C_{xy,xy})/density) [m/s] ' ,vs) - return CS,CSm,vs,density - CSs,CSms,vss,ts,densities=([],[],[],[],[]) + uma = 1.66e-27 # kg + mcell = (sum([1 for i in traj.get_attribute('symbols') if i == 'H']) + sum( + [14 for i in traj.get_attribute('symbols') if i == 'N'])) * uma + + def vs_C(traj, nsteps): + density = (mcell / np.linalg.det(traj.get_array('cells')[:nsteps:tskip] * 1e-10)).mean() + CS, CSm = elastic_c(traj.get_array('ionic_temperature')[:nsteps:tskip].mean() * 1.38064852e-23, + traj.get_array('cells')[:nsteps:tskip] * 1e-10) + vs = ((CS[0, 0] + 4.0 / 3 * CS[3, 3]) * 1e9 / density) ** .5 + print('sqrt((C_{xx,xx} + C_{xy,xy})/density)', vs) + return CS, CSm, vs, density + + CSs, CSms, vss, ts, densities = ([], [], [], [], []) for i in range(10): - stop_step=traj.numsteps//tskip*(i+1)//10 - ts.append(t[::tskip][stop_step-1]) - CSi,CSmi,vsi,density=vs_C(traj,stop_step) + stop_step = traj.numsteps // tskip * (i + 1) // 10 + ts.append(t[stop_step - 1]) + CSi, CSmi, vsi, density = vs_C(traj, stop_step) CSs.append(CSi) CSms.append(CSmi) vss.append(vsi) densities.append(density) - CSs=np.array(CSs) - CSms=np.array(CSms) - vss=np.array(vss) - ts=np.array(ts) + CSs = np.array(CSs) + CSms = np.array(CSms) + vss = np.array(vss) + ts = np.array(ts) if plot: fig, ax = plt.subplots(**subplots_kw) - ax.plot(ts,vss) + ax.plot(ts, vss) ax.set_xlabel('t (ps)') ax.set_ylabel('vs (m/s)') ax.grid() fig.show() - fig.savefig(plt_fname_pre+'vs'+plt_fname_suff) - results['Vs']=vss - results['CS_GPa']=CSs - results['density']=densities[-1] + fig.savefig(plt_fname_pre + 'vs' + plt_fname_suff) + results['Vs'] = vss + results['CS_GPa'] = CSs + results['density'] = densities[-1] except: print('error in estimating elastic constant') - return atraj,atraj_unw, results + return None, None, results + -def multiinspect(nodes,plot=False,prefix='',inspect_kw={}): - all_res=[] +def multiinspect(nodes, plot=False, prefix='', inspect_kw={}): + all_res = [] for node in nodes: - _,_,res = inspect(node,plot=plot,**inspect_kw) + _, _, res = inspect(node, plot=plot, **inspect_kw) plt.show() all_res.append(res) - return all_res, print_all(all_res,prefix=prefix) -def print_all(all_res,prefix='',subplots_kw=DEFAULT_SUBPLOTS_KW): - Ts=[] - Ps=[] - MSDs=[] - Vss=[] - DTs=[] - CS_GPas=[] - rhos=[] - grp_pos=[] - grp_width=[] + return all_res, print_all(all_res, prefix=prefix) + + +def print_all(all_res, prefix='', subplots_kw=DEFAULT_SUBPLOTS_KW): + Ts = [] + Ps = [] + MSDs = [] + Vss = [] + DTs = [] + CS_GPas = [] + rhos = [] + grp_pos = [] + grp_width = [] for res in all_res: if 'T' in res: Ts.append(res['T']) @@ -1996,47 +1843,51 @@ def print_all(all_res,prefix='',subplots_kw=DEFAULT_SUBPLOTS_KW): CS_GPas.append(res['CS_GPa'][-1]) if 'density' in res: rhos.append(res['density']) - grp_pos.append([res['NH_peak_pos'],res['NN_peak_pos'],res['HH_peak_pos']]) - grp_width.append([res['NH_peak_width'],res['NN_peak_width'],res['HH_peak_width']]) - MSDs=np.array(MSDs) - CS_GPas=np.array(CS_GPas) - rhos=np.array(rhos) - grp_pos=np.array(grp_pos) - grp_width=np.array(grp_width) - alist=[('NN peak r',grp_pos[:,1]),('NN peak width',grp_width[:,1]),('msd0',MSDs[:,0]),('msd1',MSDs[:,1])] + ([('Vs',Vss)] if len(Vss) > 0 else []) + grp_pos.append([res['NH_peak_pos'], res['NN_peak_pos'], res['HH_peak_pos']]) + grp_width.append([res['NH_peak_width'], res['NN_peak_width'], res['HH_peak_width']]) + MSDs = np.array(MSDs) + CS_GPas = np.array(CS_GPas) + rhos = np.array(rhos) + grp_pos = np.array(grp_pos) + grp_width = np.array(grp_width) + alist = [('NN peak r', grp_pos[:, 1]), ('NN peak width', grp_width[:, 1]), ('msd0', MSDs[:, 0]), + ('msd1', MSDs[:, 1])] + ([('Vs', Vss)] if len(Vss) > 0 else []) try: - for k,arr in alist: - fig, ax = plt.subplots(nrows=1,**subplots_kw) + for k, arr in alist: + fig, ax = plt.subplots(nrows=1, **subplots_kw) for i in range(len(Ts)): - ax.annotate(f'dt={DTs[i]:.3f}', xy=(Ts[i], arr[i]), - xytext=(-1, 1), textcoords="offset points", - horizontalalignment="right") - ax.scatter(Ts,arr,label=k) + ax.annotate(f'dt={DTs[i]:.3f}', xy=(Ts[i], arr[i]), + xytext=(-1, 1), textcoords="offset points", + horizontalalignment="right") + ax.scatter(Ts, arr, label=k) ax.set_xlabel('T (K)') ax.set_ylabel(k) ax.grid() - fname=FIGURE_PATH+'/summary_'+prefix+k.replace(' ', '_')+'.pdf' + fname = FIGURE_PATH + '/summary_' + prefix + k.replace(' ', '_') + '.pdf' fig.show() fig.savefig(fname) - #some sqrt(elastic constants/density) - def plt_el(Ts,CS,density): - fig, ax = plt.subplots(nrows=1,**subplots_kw) -# for arr,label in [(((CS[:,0,0]+CS[:,1,1]+CS[:,2,2])/(3.0*density))**.5,'longitudinal'),(((CS[:,3,3]+CS[:,4,4]+CS[:,5,5])/(3.0*density))**.5,'shear')]: - for arr,label in [((CS[:,0,0]/density)**.5,'longitudinal'),((CS[:,3,3]/density)**.5,'shear')]: + + # some sqrt(elastic constants/density) + def plt_el(Ts, CS, density): + fig, ax = plt.subplots(nrows=1, **subplots_kw) + # for arr,label in [(((CS[:,0,0]+CS[:,1,1]+CS[:,2,2])/(3.0*density))**.5,'longitudinal'),(((CS[:,3,3]+CS[:,4,4]+CS[:,5,5])/(3.0*density))**.5,'shear')]: + for arr, label in [((CS[:, 0, 0] / density) ** .5, 'longitudinal'), + ((CS[:, 3, 3] / density) ** .5, 'shear')]: for i in range(len(Ts)): - ax.annotate(f'dt={DTs[i]:.3f}', xy=(Ts[i], arr[i]), - xytext=(-1, 1), textcoords="offset points", - horizontalalignment="right") - ax.scatter(Ts,arr,label=label) + ax.annotate(f'dt={DTs[i]:.3f}', xy=(Ts[i], arr[i]), + xytext=(-1, 1), textcoords="offset points", + horizontalalignment="right") + ax.scatter(Ts, arr, label=label) ax.set_xlabel('T (K)') ax.set_ylabel('v (m/s)') ax.legend() ax.grid() + if len(CS_GPas) > 0: - plt_el(Ts,CS_GPas*1e9,rhos) + plt_el(Ts, CS_GPas * 1e9, rhos) except: pass - - return Ts,Ps,MSDs,Vss + + return Ts, Ps, MSDs, Vss diff --git a/pyanalisi/plotters.py b/pyanalisi/plotters.py new file mode 100644 index 000000000..676b7cd58 --- /dev/null +++ b/pyanalisi/plotters.py @@ -0,0 +1,75 @@ +import matplotlib.pyplot as plt +from mpl_toolkits.axes_grid1 import ImageGrid +from matplotlib.patches import Circle +import matplotlib.colors as colors +from matplotlib import cm +import numpy as np +import matplotlib +from IPython.core.display import HTML + +class SteinPlot: + @staticmethod + def plot(stein_res, vmin=0.01, figsize=(6., 6.), show=True, transpose=True, xmax=0.2, ymax=0.6, + inverted_type_index=False, axs=None, fig=None, single=None, plt_points=None): + if len(stein_res.shape) != 5: + raise RuntimeError('implemented only for 2d histograms!') + cmap = cm.get_cmap('inferno').copy() + cmap.set_bad('black') + nt = stein_res.shape[1] + mask = np.zeros(stein_res.shape) + # mask[:,:,:,-1,-1]=1 + masked = np.ma.masked_array(stein_res, mask=mask) + if axs is None: + if fig is None: + fig = plt.figure(figsize=figsize) + axs = ImageGrid(fig, 111, nrows_ncols=(nt, nt) if single is None else (1, 1), axes_pad=0.3) + idx = 0 + for itype in range(nt): + for jtype in range(nt): + if single is not None: + if (itype, jtype) != single: + continue + if inverted_type_index: + itype, jtype = jtype, itype + try: + axs[idx].imshow(stein_res[0, itype, jtype].transpose() if transpose else stein_res[0, itype, jtype], + norm=colors.LogNorm(vmin=vmin, vmax=masked[0, itype, jtype].max()), + cmap=cmap, + origin='lower', extent=[0.0, 1.0, 0.0, 1.0], + aspect=xmax / ymax) + axs[idx].set_xlim(0, xmax) + axs[idx].set_ylim(0, ymax) + if plt_points: + for x, y, r, c_kw in plt_points: + circ = Circle((x, y), radius=r, **c_kw) + axs[idx].add_patch(circ) + except Exception as e: + print(e) + idx += 1 + if show: + fig.show() + return fig, axs + + @staticmethod + def animation(tstein, plt_steinhardt_kw=None, plt_steinhardt_init_add_kw=None): + if plt_steinhardt_kw is None: + plt_steinhardt_kw = {} + if plt_steinhardt_init_add_kw is None: + plt_steinhardt_init_add_kw = {} + n_segments = tstein.shape[0] + class SteinAni: + def __init__(self, tstein, plt_steinhardt_kw): + self.tstein = tstein + self.plt_steinhardt_kw = plt_steinhardt_kw + self.fig, self.axs = SteinPlot.plot(self.tstein[0],**plt_steinhardt_init_add_kw, **self.plt_steinhardt_kw, show=False) + + def __call__(self, i): + self.fig, self.axs = SteinPlot.plot(self.tstein[i], **self.plt_steinhardt_kw, axs=self.axs, fig=self.fig, + show=False) + return self.axs + + stani = SteinAni(tstein, plt_steinhardt_kw) + ani = matplotlib.animation.FuncAnimation( + stani.fig, stani, interval=200, blit=True, save_count=n_segments) + return HTML(ani.to_jshtml()) + diff --git a/pyanalisi/src/calculations.py b/pyanalisi/src/calculations.py new file mode 100644 index 000000000..e69de29bb diff --git a/pyanalisi/trajectory.py b/pyanalisi/trajectory.py new file mode 100644 index 000000000..16e01088a --- /dev/null +++ b/pyanalisi/trajectory.py @@ -0,0 +1,376 @@ +import numpy as np +import pyanalisi as pa +import time +from .analysis import Analysis + +class Trajectory: + KEYS = {'positions', 'cells'} + KEYS_OPT = {'velocities', 'energy_constant_motion', 'ionic_temperature', + 'pressure', 'electronic_kinetic_energy', 'forces', 'stresses'} + TIME_K = 'times' + STEP_K = 'steps' + + DEFAULT_SAVE_REFERENCE = True + + @classmethod + def from_lammps_binary(cls, fname, dt=None, traj_skip=1, symbols=None, wrap=False, start=0,nsteps=0, pk=None): + ''' + Read a binary lammps file. traj_skip loads a step every traj_skip steps. + Note that this is done at the numpy python level, and the trajectory object still has all steps inside it + ''' + t = pa.Traj(fname) + t.setWrapPbc(wrap) + step_max=t.getNtimesteps() + if nsteps > 0: + if start+nstep > step_max: + nsteps = step_max-start + else: + nsteps = step_max + t.setAccessWindowSize(nsteps) + t.setAccessStart(start) + + return cls(t,symbols_=symbols,dt=dt, pk=pk,traj_skip=traj_skip) + + @staticmethod + def _copy(d, skip): + return np.copy(d[::skip], order='C') if skip > 1 else d + + @staticmethod + def analisi_cell2box(box_lammps): + '''returns cell vectors arranged in ROWS. + Analisi code always rotate the system to have a triangular cell matrix.''' + cells = np.zeros((box_lammps.shape[0], 3, 3)) + for i in range(box_lammps.shape[0]): + lc = box_lammps[i] + if lc.shape[0] == 9: # note that this matrix is transposed: the cell vectors are arranged in ROWS + # a_x 0 0 + # b_x b_y 0 + # c_x c_y c_z + cells[i] = np.array([[lc[3] * 2, 0, 0], + [lc[6], lc[4] * 2, 0], + [lc[7], lc[8], lc[5] * 2] + ]) + elif lc.shape[ + 0] == 6: # for orthorombic cells there is no difference between row arranged and column arranged cell vectors + cells[i] = np.array([[lc[3] * 2, 0, 0], + [0, lc[4] * 2, 0], + [0, 0, lc[5] * 2]]) + else: + raise IndexError('wrong shape of cell array') + return cells + + def _init_from_pa_traj(self, t, symbols_=None, traj_skip=1, save_reference=DEFAULT_SAVE_REFERENCE): + if symbols_ is None: + symbols_ = {} + self.data['symbols'] = [symbols_[x] if x in symbols_ else str(x) for x in t.get_lammps_type().tolist()] + self.symbols = self.data['symbols'] + self.data['positions'] = Trajectory._copy(t.get_positions_copy(), traj_skip) + self.data['velocities'] = Trajectory._copy(t.get_velocities_copy(), traj_skip) + box_lammps = Trajectory._copy(t.get_box_copy(), traj_skip) + self.data['cells'] = Trajectory.analisi_cell2box(box_lammps) + if save_reference: + self.paTraj = t + + + def __init__(self, t, symbols_=None, dt=None, pk=None, traj_skip=None, max_len=None, save_reference=DEFAULT_SAVE_REFERENCE,lazy_load=False): + self.paTrajWrapped = None + self.paTraj = None + self.to_load_keys=[] + self.lazy_load_obj=None + if symbols_ is None: + symbols_ = {} + if traj_skip is None and max_len is None: + traj_skip=1 + if traj_skip is not None and max_len is not None: + raise AttributeError("You can't specify both traj_skip and max_len keywords args") + + self.data = {} + if isinstance(t, list): ##list of whatever supported object + nlist = [ Trajectory(it,lazy_load=lazy_load) for it in t ] + if max_len is not None: + tot_len = sum([_.numsteps for _ in nlist]) + traj_skip=max(tot_len//max_len,1) + + self.symbols = nlist[0].symbols + if dt is None: + dt = nlist[0].dt + if pk is None: + pks= [f'{_.pk}' for _ in nlist if _.pk is not None] + if len(pks)>0: + pk='_'.join(pks) + del pks + for key in self.KEYS: + self.data[key] = Trajectory._copy(np.concatenate([it.get_array(key) for it in nlist], axis=0), + traj_skip) + for key in self.KEYS_OPT: + if key in t[0].get_arraynames(): + self.data[key] = Trajectory._copy(np.concatenate([it.get_array(key) for it in nlist], axis=0), + traj_skip) + del nlist + elif isinstance(t, (pa.Traj, pa.Trajectory)): ##pyanalisi trajectory + if max_len is not None: + traj_skip=max(t.get_nloaded_timesteps()//max_len,1) + self._init_from_pa_traj(t, symbols_=symbols_, traj_skip=traj_skip, save_reference=save_reference) + elif isinstance(t, dict): ##simple dict with the data + self.symbols = t['symbols'] + if max_len is not None: + try: + traj_skip = max(t[self.KEYS[0]].shape[0]//max_len,1) + except: + traj_skip = max( len(t[self.KEYS[0]])//max_len,1) + for key in self.KEYS: + self.data[key] = Trajectory._copy(np.array(t[key]), traj_skip) + for key in self.KEYS_OPT: + if key in t: + self.data[key] = Trajectory._copy(np.array(t[key]), traj_skip) + else: ##Trajectory-like interface, similar to aiida one + try: + self.symbols = t.symbols + if max_len is not None: + traj_skip=max(1, t.numsteps//max_len) + if not lazy_load: + for key in self.KEYS: + self.data[key] = Trajectory._copy(t.get_array(key), traj_skip) + for key in self.KEYS_OPT: + if key in t.get_arraynames(): + self.data[key] = Trajectory._copy(t.get_array(key), traj_skip) + else: + self.lazy_load_skip=traj_skip + self.lazy_load_obj=t + to_load_keys=t.get_arraynames() + self.to_load_keys=set([ key for key in self.KEYS if key in to_load_keys] + [ key for key in self.KEYS_OPT if key in to_load_keys]) + self.will_load_keys=list(self.to_load_keys) + if Trajectory.TIME_K in t.get_arraynames() and dt is None: + times=t.get_array(Trajectory.TIME_K) + dt=times[1]-times[0] + if hasattr(t,'pk') and pk is None: + pk = t.pk + except: + raise RuntimeError(f'first argument cannot be {str(t)}') + if pk is None: + import time + pk = str(time.time_ns()) + if dt is None: + dt=1.0 + if 'positions' in self.data.keys(): + nsteps= self.data['positions'].shape[0] + nsites= self.data['positions'].shape[1] + elif lazy_load: + try: + nsteps=t.numsteps//traj_skip + nsites = t.numsites + except: + raise RuntimeError("Lazy load of trajectory of type {type(t)} not implemented") + self.data[Trajectory.STEP_K] = np.arange(0, nsteps) + self.data[Trajectory.TIME_K] = self.data[Trajectory.STEP_K] * dt * traj_skip + self.dt = dt * traj_skip + self.numsites = nsites + self.pk = pk + self.numsteps = nsteps + + + def get_array(self, name): + if name in self.to_load_keys: + self.data[name] = Trajectory._copy(self.lazy_load_obj.get_array(name), self.lazy_load_skip) + self.to_load_keys.remove(name) + if len(self.to_load_keys)==0: + self.lazy_load_obj=None + if name in self.data: + return self.data[name] + else: + raise KeyError(f'cannot find key {name}') + + def get_attribute(self, k): + if k == 'symbols': + return self.symbols + else: + raise KeyError(f'attribute {k} not present') + + def get_arraynames(self): + if self.lazy_load_obj is None: + return self.data.keys() + else: + return self.will_load_keys + + @staticmethod + def get_types_id_array(types_array): + types = {} + typeid = 0 + res = [] + for t in types_array: + if not t in types: + types[t] = typeid + typeid = typeid + 1 + res.append(types[t]) + return np.array(res, dtype='int32') + + def get_analisi_traj(self, tskip=1, wrapped=False, save_reference=DEFAULT_SAVE_REFERENCE): + if self.paTraj and tskip==1 and not wrapped: + return self.paTraj + elif self.paTrajWrapped and tskip==1 and wrapped: + return self.paTrajWrapped + + if tskip < 1: + raise IndexError(f'cannot have a trajectory skip of {tskip}') + pos = self.get_array('positions')[::tskip].copy(order='C') + if 'velocities' in self.get_arraynames(): + vel = self.get_array('velocities')[::tskip].copy(order='C') + else: + vel = np.zeros(pos.shape) + cel = self.get_array('cells')[::tskip].transpose((0, 2, 1)).copy(order='C') + types = Trajectory.get_types_id_array(self.get_attribute('symbols')) + params = [pos, vel, types, cel] + atraj = pa.Trajectory(*params, pa.BoxFormat.CellVectors, wrapped, True) + if save_reference and tskip==1: + if wrapped: + self.paTrajWrapped = atraj + else: + self.paTraj = atraj + return atraj + + @staticmethod + def get_type_mask_from_s(symbols): + typesa = set(symbols) + masks = {} + types = [] + types_array = np.zeros(len(symbols), dtype=int) + itype = 0 + for t in typesa: + masks[t] = np.array(symbols) == t + types.append(t) + types_array[masks[t]] = itype + itype += 1 + return types, types_array, masks + + def get_poscar(self, every=1, start=0, desc=''): + """ + return array of poscar file strings. + """ + poscars = [] + types, types_array, masks = Trajectory.get_type_mask_from_s(self.symbols) + pos = self.get_array("positions") + vel = self.get_array("velocities") + cel = self.get_array("cells") + nt = {} + for it in types: + nt[it] = np.sum(masks[it]) + for i in range(start, self.numsteps, every): + c = cel[i] + p = f'{desc}\n1.0\n' + for idim in range(3): + p += f'{c[idim, 0]} {c[idim, 1]} {c[idim, 2]}\n' + for it in types: + p += it + ' ' + p += '\n' + for it in types: + p += f'{nt[it]} ' + p += '\nCartesian\n' + for it in types: + ps = pos[i, masks[it]] + for iatom in range(nt[it]): + p += f'{ps[iatom, 0]} {ps[iatom, 1]} {ps[iatom, 2]} \n' + p += 'Cartesian\n' + for it in types: + v = vel[i, masks[it]] + for iatom in range(nt[it]): + p += f'{v[iatom, 0]} {v[iatom, 1]} {v[iatom, 2]} \n' + + poscars.append(p) + return poscars + + def rotate_qr(self,wrap=False,save_reference=DEFAULT_SAVE_REFERENCE): + """ + transform the cell time series of this object to a triangular matrix. + Rotate position, velocities, forces and stresses accordingly + """ + wrapped=self.get_analisi_traj(wrapped=wrap,save_reference=save_reference) + q = wrapped.get_rotation_matrix() + new_cell = Trajectory.analisi_cell2box(wrapped.get_box_copy()) + self.data['cells']=new_cell + self.data['positions']=wrapped.get_positions_copy() + if 'forces' in self.get_arraynames(): + rotated_forces = np.einsum('taj,tij -> tai', self.get_array('forces'), q) + self.data['forces'] = rotated_forces + if 'stresses' in self.get_arraynames(): + rotated_stress = np.einsum('tlm,tli,tmj -> tij', self.get_array('stresses'), q, q) + self.data['stresses']=rotated_stress + if 'velocities' in self.get_arraynames(): + self.data['velocities'] = wrapped.get_velocities_copy() + + def get_aiida_traj(self): + import aiida + res = aiida.orm.nodes.data.array.trajectory.TrajectoryData() + res.set_attribute('symbols', ft.symbols) + for name in self.KEYS : + if name in self.get_arraynames(): + res.set_array(name, ft.get_array(name)) + for name in self.KEYS_OPT : + if name in self.get_arraynames(): + res.set_array(name, ft.get_array(name)) + for name in (self.TIME_K,self.STEP_K): + if name in self.get_arraynames(): + res.set_array(name, ft.get_array(name)) + return res + + def get_lammps_data(self, timestep=0, comment=None, symb_id=None): + def write_cell(a): + res = f'{a[0]} {a[0] + 2 * a[3]} xlo xhi\n{a[1]} {a[1] + 2 * a[4]} ylo yhi\n{a[2]} {a[2] + 2 * a[5]} zlo zhi\n' + if len(a) > 6: + res += f'{a[6]} {a[7]} {a[8]} xy xz yz\n' + return res + + wrapped = self.get_analisi_traj(wrapped=False,save_reference=self.DEFAULT_SAVE_REFERENCE) + cel = wrapped.get_box_copy() + pos = wrapped.get_positions_copy() + vel = wrapped.get_velocities_copy() + + if timestep >= pos.shape[0] or timestep < -pos.shape[0]: + raise IndexError('out of range index') + if comment is None: + comment = f'timestep {timestep}' + if symb_id is None: + all_sym = list(set(self.symbols)) + symb_id = {s: str(i + 1) for i, s in enumerate(all_sym)} + symbols = [symb_id[s] for s in self.symbols] + nt = len(symb_id) + nat = pos.shape[1] + + res = f'{comment}\n{nat} atoms\n{nt} atom types\n{write_cell(cel[timestep])}\nAtoms\n\n' + for iatom, itype, xyz in zip(range(nat), symbols, pos[timestep]): + res += f'{iatom + 1} {itype} {xyz[0]} {xyz[1]} {xyz[2]}\n' + res += '\nVelocities\n\n' + for iatom, xyz in zip(range(nat), vel[timestep]): + res += f'{iatom + 1} {xyz[0]} {xyz[1]} {xyz[2]}\n' + + return res + + def write_lammps_binary(self,fname, start=0, end=-1): + if end > self.numsteps: + raise IndexError(f'end must be <= {self.numsteps}') + if start < 0 or start > self.numsteps: + raise IndexError(f'start must be in [0, self.numsteps]') + if 0 <= end < start: + raise IndexError('start > end !') + wrapped = self.get_analisi_traj(wrapped=False,save_reference=self.DEFAULT_SAVE_REFERENCE) + wrapped.write_lammps_binary(fname, start, end) + + def get_atomic_types(self): + return list(set(self.symbols)) + + def show_traj(self, fast=1.0, plot=None): + import k3d + if plot is None: + plot = k3d.plot() + atomic_species = self.get_atomic_types() + masks = [] + for sp in atomic_species: + masks.append(np.array(self.symbols) == sp) + pos_m = self.get_array('positions') + t = self.get_array('steps') + for sp, mask in zip(atomic_species, masks): + plot += k3d.points(positions={str(t / 30.0 / fast): pos_m[t, mask, :] for t in range(pos_m.shape[0])}, size=1.0, name=sp, shader='mesh',point_size=0.6) + return plot + + def get_analysis(self,**kwargs): + return Analysis(self,**kwargs) +