Source code for CDTK.an_plt_rho1D

#*  **************************************************************************
#*
#*  CDTK, Chemical Dynamics Toolkit
#*  A modular system for chemical dynamics applications and more
#*
#*  Copyright (C) 2011, 2012, 2013, 2014, 2015, 2016
#*  Oriol Vendrell, DESY, <oriol.vendrell@desy.de>
#*
#*  Copyright (C) 2017, 2018, 2019
#*  Ralph Welsch, DESY, <ralph.welsch@desy.de>
#*
#*  Copyright (C) 2020, 2021, 2022, 2023
#*  Ludger Inhester, DESY, ludger.inhester@cfel.de
#*
#*  This file is part of CDTK.
#*
#*  CDTK is free software: you can redistribute it and/or modify
#*  it under the terms of the GNU General Public License as published by
#*  the Free Software Foundation, either version 3 of the License, or
#*  (at your option) any later version.
#*
#*  This program is distributed in the hope that it will be useful,
#*  but WITHOUT ANY WARRANTY; without even the implied warranty of
#*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#*  GNU General Public License for more details.
#*
#*  You should have received a copy of the GNU General Public License
#*  along with this program.  If not, see <http://www.gnu.org/licenses/>.
#*
#*  **************************************************************************

import os
import os.path
import sys
from optparse import OptionParser

import numpy as np
import matplotlib.pylab as plt
from mpl_toolkits.mplot3d import Axes3D

import CDTK.Tools.Utils as uti
import CDTK.Tools.Mathematics as mat

[docs] def start(): """ Wrapper for starting the program. Needed for good unit test coding. """ fig = plt.figure() #ax = Axes3D(fig) # -------------------------------------------------------------------------- # Parse command line options # -------------------------------------------------------------------------- parser=OptionParser() parser.add_option('-n','--name', dest='name', type='str', default=None, help='Name of file to be found in trj directories') parser.add_option('-a','--xmin', dest='xmin', type='float', default=0.0, help='lower limit, x axis') parser.add_option('-x','--xmax', dest='xmax', type='float', default=10.0, help='upper limit, x axis') parser.add_option('-b','--ymin', dest='ymin', type='float', default=0.0, help='lower limit, y axis') parser.add_option('-y','--ymax', dest='ymax', type='float', default=10.0, help='upper limit, y axis') parser.add_option('-l','--xlabel', dest='xlabel', type='str', default=None, help='text label for x axis') parser.add_option('-L','--ylabel', dest='ylabel', type='str', default=None, help='text label for y axis') parser.add_option('-u','--xfactor', dest='xfactor', type='float', default=1.0, help='scaling factor for x axis') parser.add_option('-v','--yfactor', dest='yfactor', type='float', default=1.0, help='scaling factor for y axis') parser.add_option('-w','--fwhm', dest='fwhm', type='float', default=1.0, help='full with at half maximum for Gaussian smoothing') parser.add_option('-z','--xsteps', dest='xsteps', type='int', default=500, help='number of time steps being processed') parser.add_option('-Z','--ysteps', dest='ysteps', type='int', default=40, help='number of discrete steps in the y axis') parser.add_option('-N','--ncountour', dest='ncontour', type='int', default=40, help='number of contours') opts, args = parser.parse_args(sys.argv[1:]) if not sys.argv[1:]: # called without any option parser.print_help() sys.exit(0) # -------------------------------------------------------------------------- # Obtain list of trj directories # -------------------------------------------------------------------------- dirs = uti.getDirectoryList() # -------------------------------------------------------------------------- # Loop over directories and plot data file # -------------------------------------------------------------------------- xf = opts.xfactor yf = opts.yfactor yy = np.linspace(opts.ymin,opts.ymax,opts.ysteps)*yf ZZ = np.zeros((opts.xsteps,opts.ysteps),float) count = np.zeros(opts.xsteps,int) for d in dirs: try: xy = np.loadtxt(d+'/'+opts.name,unpack=True) dx = xy[0][1] - xy[0][0] x0 = xy[0][0] for i,y in enumerate(xy[1]): gg = mat.GaussianFunction(yy,fwhm=opts.fwhm,xbar=y) ZZ[i] = ZZ[i] + gg count[i] += 1 except: pass for i,c in enumerate(count): if c != 0: ZZ[i] = ZZ[i]/float(count[i]) x1 = (opts.xsteps-1)*dx xx = np.linspace(x0,x1,opts.xsteps) YY,XX = np.meshgrid(yy,xx) plt.contour(XX/xf, YY/yf, ZZ, opts.ncontour, cmap='Spectral_r') #ax.plot_surface(XX, YY, ZZ/count, rstride=1, cstride=1, cmap='Spectral_r') if opts.xmin: plt.xlim(xmin=opts.xmin) if opts.xmax: plt.xlim(xmax=opts.xmax) if opts.ymin: plt.ylim(ymin=opts.ymin) if opts.ymax: plt.ylim(ymax=opts.ymax) if opts.xlabel: plt.xlabel(opts.xlabel) if opts.ylabel: plt.ylabel(opts.ylabel) plt.show()
if __name__ == "__main__": start()