API tools faq
paste
wbrigg

basis_set_plotter.py

wbrigg
Jan 28th, 2015
481
0
Never
Add comment
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 7.99 KB | None | 0 0
raw download clone embed print report
  1. import matplotlib.pyplot as plt
  2. import argparse, os, sys
  3. import numpy as np
  4. from scipy.misc import factorial2
  5.  
  6. parser = argparse.ArgumentParser(description="Plot Basis Sets")
  7. parser.add_argument("-a", "--atom", type=str, dest="atom", help="atom set to plot", default="H")
  8. parser.add_argument("-b", "--bases", type=str, nargs="+", dest="bases", help="basis sets to plot, maximum of four", default=["3-21G"])
  9. parser.add_argument("-l", "--list", action='store_true', dest="list", help="list available basis sets", default=False)
  10. parser.add_argument("--lmax", type=int, dest="l_max", help="Only plot up to a maximum l", default=-1)
  11. parser.add_argument("--lmin", type=int, dest="l_min", help="Only plot above a minimum l", default=-1)
  12. parser.add_argument("-gs", action='store_true', dest="gs", help="Plot a comparison between a GTO and an STO")
  13.  
  14. parser.add_argument("-r", "--rrange", type=float, nargs=2, dest="r", help="limits for r", default=[-4, 4])
  15.  
  16. parser.add_argument("-pub", action="store_true", dest="pub")
  17.  
  18. parser.add_argument("--inset", action='store_true', dest="inset", help="Add an axes zoom inset")
  19. parser.add_argument("--inset-extent", type=float, nargs=4, dest="inset_extent", help="Axes limits for the zoomed plot", default=[None, None, None, None])
  20. parser.add_argument("--inset-loc", type=int, dest="inset_loc", help="Axes inset location", default=1)
  21. parser.add_argument("--inset-zoom", type=int, dest="inset_zoom", help="Axes inset zoom", default=1)
  22. parser.add_argument("--inset-aspect", type=float, dest="inset_aspect", help="Alter the inset's aspect ratio. A larger number will widen the inset axes", default=1)
  23.  
  24. args = parser.parse_args()
  25.  
  26. atom_names = {"H":"Hydrogen", "He":"Helium", "Li":"Lithium", "Be":"Beryllium", "B":"Boron", "C":"Carbon", "N":"Nitrogen", "O":"Oxygen", "F":"Fluorine", "Ne":"Neon", "Na":"Sodium", "Mg":"Magnesium", "Al":"Aluminum", "Si":"Silicon", "P":"Phosphorous", "S":"Sulfur", "Cl":"Chlorine", "Ar":"Argon", "K":"Potassium", "Ca":"Calcium", "Sc":"Scandium", "Ti":"Titanium", "V":"Vanadium", "Cr":"Chromium", "Mn":"Manganese", "Fe":"Iron", "Co":"Cobalt", "Ni":"Nickel", "Cu":"Copper", "Zn":"Zinc", "Ga":"Gallium", "Ge":"Germanium", "As":"Arsenic", "Se":"Selenium", "Br":"Bromine", "Kr":"Krypton", "Rb":"Rubidium", "Sr":"Strontium", "Y":"Yttrium", "Zr":"Zirconium", "Nb":"Niobium", "Mo":"Molybdenum", "Tc":"Technetium", "Ru":"Ruthenium", "Rh":"Rhodium", "Pd":"Palladium", "Ag":"Silver", "Cd":"Cadmium", "In":"Indium", "Sn":"Tin", "Sb":"Antimony", "Te":"Tellurium", "I":"Iodine", "Xe":"Xenon", "Cs":"Cesium", "Ba":"Barium", "La":"Lanthanum", "Ce":"Cerium", "Pr":"Praseodymium", "Nd":"Neodymium", "Pm":"Promethium", "Sm":"Samarium", "Eu":"Europium", "Gd":"Gadolinium", "Tb":"Terbium", "Dy":"Dysprosium", "Ho":"Holmium", "Er":"Erbium", "Tm":"Thulium", "Yb":"Ytterbium", "Lu":"Lutetium", "Hf":"Hafnium", "Ta":"Tantalum", "W":"Tungsten", "Re":"Rhenium", "Os":"Osmium", "Ir":"Iridium", "Pt":"Platinum", "Au":"Gold", "Hg":"Mercury", "Tl":"Thallium", "Pb":"Lead", "Bi":"Bismuth", "Po":"Polonium", "At":"Astatine", "Rn":"Radon", "Fr":"Francium", "Ra":"Radium", "Ac":"Actinium", "Th":"Thorium", "Pa":"Protactinium", "U":"Uranium", "Np":"Neptunium", "Pu":"Plutonium", "Am":"Americium", "Cm":"Curium", "Bk":"Berkelium", "Cf":"Californium", "Es":"Einsteinium", "Fm":"Fermium", "Md":"Mendelevium", "No":"Nobelium", "Lr":"Lawrencium", "Rf":"Rutherfordium", "Db":"Dubnium", "Sg":"Seaborgium", "Bh":"Bohrium", "Hs":"Hassium", "Mt":"Meitnerium", "Ds":"Darmstadtium", "Rg":"Roentgenium", "Cp":"Copernicium", "Uut":"Ununtrium", "Fl":"Flerovium", "Uup":"Ununpentium", "Lv":"Liermorium", "Uus":"Ununseptium", "Uuo":"Ununoctium"}
  27.  
  28. inv_atom_names = {v:k for k,v in atom_names.items()}
  29.  
  30.  
  31.  
  32. rs = np.linspace(args.r[0],args.r[1],10000)
  33.  
  34.  
  35. if args.gs:
  36.   fig = plt.figure()
  37.   ax1 = fig.add_subplot(1,2,1)
  38.   ax2 = fig.add_subplot(1,2,2)
  39.  
  40.   ax1.plot(rs, np.exp(-(rs**2)), color="blue")
  41.   ax2.plot(rs, np.exp(-abs(rs)), color="orange")
  42.  
  43.   ax2.set_yticklabels([])
  44. #  ax1.set_xticklabels(ax1.get_xticks().tolist()[:-1] + [""])
  45. #  ax2.set_xticklabels([""] + ax2.get_xticks().tolist()[1:])
  46.  
  47.  
  48.  
  49.   plt.subplots_adjust(wspace=0.1)
  50.  
  51.   if args.pub:
  52.     plt.savefig("GTO_vs_STO.pdf")  
  53.   else:
  54.     plt.show()
  55.   sys.exit(0)
  56.  
  57.  
  58.  
  59.  
  60.  
  61.  
  62. if args.list:
  63.   for _, _, files in os.walk("./basis_sets/"):
  64.     for basis_set_file_name in files:
  65.       sys.stdout.write(basis_set_file_name[:-4] + "\n")
  66.      
  67.   if len(args.bases) == 1:
  68.     sys.stdout.write("\nCurrent selected basis: " + args.bases[0] + "\n")
  69.   else:
  70.     sys.stdout.write("\nCurrent selected bases: " + " ".join(args.bases) + "\n")
  71.    
  72.   sys.stdout.write("Available atoms:\n")  
  73.  
  74.  
  75.      
  76. basis_sets = {}
  77. for basis in args.bases:
  78.   basis_sets[basis] = {}
  79.   with open("./basis_sets/" + basis + ".dat") as f:
  80.     l_dict = dict([(letter, i) for i,letter in enumerate("SPDFGHIJKL")])
  81.     L = 0
  82.     atom = None
  83.     for line in f:
  84.       if line.startswith("$"):
  85.         if "-TYPE FUNCTIONS" in line:
  86.           line = f.next()
  87.           primitives, functions, _ = [int(e) for e in line.split()]
  88.           basis_sets[basis][atom].append([])
  89.           for i in range(primitives):
  90.             line = f.next()
  91.             basis_sets[basis][atom][-1].append([float(e) for e in line.replace("D","E").split()])
  92.            
  93.            
  94.           tmp = zip(*basis_sets[basis][atom][-1])
  95.           basis_sets[basis][atom][-1] = {"primitives":tmp[0], "contractions":tmp[1:]}
  96.          
  97.         else:
  98.           atom = line.split()[1]
  99.           if atom not in basis_sets[basis]:
  100.             basis_sets[basis][atom] = []
  101.            
  102.            
  103.            
  104. if args.list:
  105.  
  106.   atom_sets = [[key for key in basis_set] for basis_set in basis_sets.values()]
  107.  
  108.   atom_set = set(atom_sets[0])
  109.   for atoms in atom_sets[1:]:
  110.     atom_set = atom_set & set(atoms)
  111.    
  112.   atoms = sorted(list(atom_set))
  113.  
  114.  
  115.   atoms = [inv_atom_names[atom.title()] for atom in atoms]
  116.  
  117.   sys.stdout.write(" ".join(atoms) + "\n")  
  118.  
  119.   sys.exit(0)
  120.  
  121.  
  122. def N(alpha, l):
  123.   return np.sqrt(np.power(2*alpha/np.pi, 1.5) * np.power(4*alpha, l) / factorial2(2*l-1))
  124.  
  125. def R(l, alpha, r):
  126.   return N(alpha,l) * r**l * np.exp(-alpha * r**2)
  127.  
  128. def GTO(l, alphas, cs, r):
  129.   if type(r) is list:
  130.     result = np.zeros_like(r)
  131.   else:
  132.     result = 0
  133.    
  134.   for alpha, c in zip(alphas, cs):
  135.     result += c * R(l, alpha, r)
  136.  
  137.   return result# / len(alphas)
  138.  
  139.  
  140.  
  141.  
  142. fig = plt.figure()
  143. ax = fig.add_subplot(1,1,1)
  144.  
  145.  
  146. if args.inset:
  147.   from mpl_toolkits.axes_grid1.inset_locator import inset_axes, zoomed_inset_axes, mark_inset
  148.  
  149.   extent = args.inset
  150.   axins = zoomed_inset_axes(ax, args.inset_zoom, loc=args.inset_loc)
  151.  
  152.  
  153. colours = "red blue green cyan gold purple".split()
  154. #colours = "red orange yellow green blue purple violet".split()
  155.  
  156. line_styles = "- :".split()
  157. if len(args.bases) == 3:
  158.   line_styles = "- -- :".split()
  159. elif len(args.bases) == 4:
  160.   line_styles = "- -- -. :".split()
  161. elif len(args.bases) > 4:
  162.   raise Exception("Maximum comparison is of 4 basis sets")
  163.  
  164. for line_no, basis in enumerate(args.bases):
  165.   for l, basis_set_segment in enumerate(basis_sets[basis][atom_names[args.atom].upper()]):
  166.     if l < args.l_min: continue
  167.     elif l > args.l_max and args.l_max is not -1: break
  168.    
  169.     alphas = basis_set_segment["primitives"]
  170.     for contraction in basis_set_segment["contractions"]:
  171.       ax.plot(rs, GTO(l, alphas, contraction, rs), color=colours[l], linestyle=line_styles[line_no])
  172.       if args.inset: axins.plot(rs, GTO(l, alphas, contraction, rs), color=colours[l], linestyle=line_styles[line_no])
  173.    
  174.    
  175.  
  176. if args.inset:
  177.   if args.inset_extent[0] is None: pass
  178.   else:
  179.     axins.axis(args.inset_extent)
  180.  
  181.   axins.set_xticklabels([])
  182.  
  183.   mark_inset(ax, axins, loc1=2, loc2=3, fc="none", ec="0.5")
  184.   pos = axins.get_position()
  185.   pos = [pos.x0, pos.y0,  pos.width * 2, pos.height]
  186.   axins.set_position(pos)
  187.   print pos
  188.   print axins.get_position()
  189.  
  190. if args.pub:
  191.   plt.savefig("GTO_" + args.atom + "_" + "_".join(args.bases) + ".pdf")  
  192. else:
  193.   plt.show()
Advertisement
Add Comment
Please, Sign In to add comment
Public Pastes
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy.  OK, I Understand
Not a member of Pastebin yet?
Sign Up, it unlocks many cool features!