Source code for CDTK.an_plt_rho2D

#*  **************************************************************************
#*
#*  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
import CDTK.Tools.Conversion as cv

[docs] def start(): """ Wrapper for starting the program. Needed for good unit test coding. """ fig = plt.figure() # -------------------------------------------------------------------------- # 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('-m','--name2', dest='name2', type='str', default=None, help='Name of second file to be found in trj directories') parser.add_option('-t','--time', dest='time', type='float', default=0.0, help='time (closest to show, in au') 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=20, help='number of discrete steps in the x axis') parser.add_option('-Z','--ysteps', dest='ysteps', type='int', default=20, help='number of discrete steps in the y axis') parser.add_option('-N','--ncontour', dest='ncontour', type='int', default=40, help='number of contours') parser.add_option('','--symmetric', action='store_true', dest='is_symmetric', default=False, help='if set symmetrize plot') 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 xx = np.linspace(opts.xmin,opts.xmax,opts.xsteps)*xf yy = np.linspace(opts.ymin,opts.ymax,opts.ysteps)*yf ZZ = np.zeros((opts.xsteps,opts.ysteps),float) count = 0 for d in dirs: try: xy = np.loadtxt(d+'/'+opts.name,unpack=True) xy2 = np.loadtxt(d+'/'+opts.name2,unpack=True) dt = xy[0][1] - xy[0][0] it = opts.time//dt gg = mat.GaussianFunction(xx,fwhm=opts.fwhm,xbar=xy[1][it]) gg2 = mat.GaussianFunction(yy,fwhm=opts.fwhm,xbar=xy2[1][it]) ZZ = ZZ + np.outer(gg,gg2) count += 1 except: pass #ZZ = ZZ/float(count) YY,XX = np.meshgrid(yy,xx) plt.xlim(xmin=opts.xmin) plt.xlim(xmax=opts.xmax) plt.ylim(ymin=opts.ymin) plt.ylim(ymax=opts.ymax) if opts.xlabel: plt.xlabel(opts.xlabel) if opts.ylabel: plt.ylabel(opts.ylabel) plt.gca().set_aspect('equal', adjustable='box') if opts.is_symmetric: ZZ = 0.5*(ZZ + np.transpose(ZZ)) 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') plt.show()
if __name__ == "__main__": start()