1from __future__ import print_function 2import numpy as np 3import matplotlib.pyplot as plt 4import matplotlib as mpl 5import matplotlib.cm as cm 6import pylab as P 7import matplotlib 8"""The main functions you'll use here are read_dm and read_dm_offdiag. 9These take as arguments a filename, which should be the stdout of gosling as applied 10to your log file. 11""" 12################################################################ 13 14def plot_color(A,filename,ticklabels=None, annotate=True): 15 nregion1=A.shape[0] 16 nregion2=A.shape[1] 17 P.clf() 18 19 20 cdict = {'red': ((0.0, 1.0, 1.0), 21 (0.5, 1.0, 1.0), 22 (1.0, 0.4, 0.4)), 23 'green': ((0.0, 0.4, 0.4), 24 (0.5, 1.0, 1.0), 25 (1.0, 0.4, 0.4)), 26 'blue': ((0.0, 0.4, 0.4), 27 (0.5, 1.0, 1.0), 28 (1.0, 1.0, 1.0))} 29 my_cmap = matplotlib.colors.LinearSegmentedColormap('my_colormap',cdict,256) 30 vmax=np.max(A*100) 31 vmin=np.min(A*100) 32 vmax=max(-vmin,vmax) 33 vmin=min(-vmax,vmin) 34 #print vmax,vmin 35# P.pcolor(A*100,cmap=my_cmap,edgecolors='k',vmin=vmin,vmax=vmax,array=True) 36 P.imshow(A*100,cmap=my_cmap,vmin=vmin,vmax=vmax,interpolation='nearest',origin='upper') 37 38 thresh=0.05*(vmax-vmin) 39 40 if annotate: 41 for r1 in range(0,nregion1): 42 for r2 in range(0,nregion2): 43 n=A[r1,r2]*100 44 if abs(n) > thresh: 45 P.annotate("%.1f"%(n), (r2,r1),ha="center",va="center",fontsize=2) 46 a=P.gca() 47 a.get_xaxis().tick_top() 48 for i in ['top','right','left','bottom']: 49 a.spines[i].set_visible(False) 50 a.tick_params(labelsize=10,width=0) 51 P.savefig(filename) 52 53################################################################ 54 55def convert(s): 56 return float(s) 57 a=eval(s) 58 m=float(a[0]) 59 for i in a: 60 m=max(m,i) 61 return m 62 63################################################################ 64 65 66def read_dm(inpf): 67 """Read in the 1-RDM and/or 1-RDM and diagonals of the 2-RDM, if only 68 the diagonals were calculated""" 69 nmo=0 70 while True: 71 line=inpf.readline() 72 if "tbdm" in line: 73 spl=line.split() 74 nmo=int(spl[2]) 75 break 76 if line=="": 77 return None 78 data={} 79 #one-body up, one-body up err, two-body-uu, two-body-uu err 80 for nm in ['ou','oue','od','ode','tuu','tuue','tud','tude','tdu','tdue','tdd','tdde']: 81 data[nm]=np.zeros((nmo,nmo)) 82 83 while True: 84 line=inpf.readline() 85 if line=="" or "Trace" in line: 86 break; 87 if line.find("tbdm: states") != -1: 88 data['states']=np.array(list(map(int,line.split()[3:-1]))) 89 #print data['states'] 90 if line.find("One-body density") != -1: 91 line=inpf.readline() 92 for i in range(0,nmo): 93 for j in range(0,nmo): 94 a=inpf.readline().split() 95 data['ou'][i,j]=convert(a[2]) 96 data['oue'][i,j]=convert(a[3]) 97 data['od'][i,j]=convert(a[4]) 98 data['ode'][i,j]=convert(a[5]) 99 if line.find("two-body density") != -1: 100 line=inpf.readline() 101 #print "two-body",line 102 for i in range(0,nmo): 103 for j in range(0,nmo): 104 a=inpf.readline().split() 105 #print i,j,"a ",a 106 data['tuu'][i,j]=convert(a[4]) 107 data['tuue'][i,j]=convert(a[5]) 108 data['tud'][i,j]=convert(a[6]) 109 data['tude'][i,j]=convert(a[7]) 110 data['tdu'][i,j]=convert(a[8]) 111 data['tdue'][i,j]=convert(a[9]) 112 data['tdd'][i,j]=convert(a[10]) 113 data['tdde'][i,j]=convert(a[11]) 114 break 115 return data 116 117################################################################ 118 119 120def read_dm_offdiag(f): 121 """Read in the 1-RDM and 2-RDM, if the off-diagonal elements were calculated 122 for the 2-RDM""" 123 124 nmo=0 125 while True: 126 line=f.readline() 127 if line.find("tbdm")!=-1: 128 spl=line.split() 129 nmo=int(spl[2]) 130 break 131 data={} 132 #one-body up, one-body up err, two-body-uu, two-body-uu err 133 for nm in ['ou','oue','od','ode']: 134 data[nm]=np.zeros((nmo,nmo)) 135 for nm in ['tuu','tuue','tud','tude','tdu','tdue','tdd','tdde']: 136 data[nm]=np.zeros((nmo*nmo,nmo*nmo)) 137 while True: 138 line=f.readline() 139 if line=="" or "Trace" in line: 140 break; 141 if line.find("tbdm: states") != -1: 142 data['states']=np.array(list(map(int,line.split()[3:-1]))) 143 if line.find("One-body density") != -1: 144 line=f.readline() 145 for i in range(0,nmo): 146 for j in range(0,nmo): 147 a=f.readline().split() 148 data['ou'][i,j]=convert(a[2]) 149 data['oue'][i,j]=convert(a[3]) 150 data['od'][i,j]=convert(a[4]) 151 data['ode'][i,j]=convert(a[5]) 152 if line.find("two-body density") != -1: 153 line=f.readline() 154 for i in range(0,nmo): 155 for j in range(0,nmo): 156 for k in range(0,nmo): 157 for l in range(0,nmo): 158 a=f.readline().split() 159 #print i,j,"a ",a 160 for n,index in zip(['tuu','tuue','tud','tude','tdu','tdue','tdd','tdde'], 161 range(4,12)): 162 data[n][i*nmo+j,k*nmo+l]=convert(a[index]) 163 break 164 return data 165################################################################ 166 167def read_ekt(f): 168 data={} 169 while True: 170 line=f.readline() 171 if 'EKT' in line: 172 data['nmo']=int(line.split()[2]) 173 break 174 175 line=f.readline() 176 data['states']=np.array(list(map(int,line.split()[3:-1]))) 177 print(data['states'],data['nmo']) 178 nmo=data['nmo'] 179 for el in ['ou','oue','od','ode','vu','vue','vd','vde','cu','cue','cd','cde']: 180 data[el]=np.zeros((data['nmo'],data['nmo'])) 181 while True: 182 line=f.readline() 183 if 'One-body density matrix' in line: 184 f.readline() 185 for i in range(0,nmo): 186 for j in range(0,nmo): 187 a=f.readline().split() 188 data['ou'][i,j]=convert(a[2]) 189 data['oue'][i,j]=convert(a[3]) 190 data['od'][i,j]=convert(a[4]) 191 data['ode'][i,j]=convert(a[5]) 192 193 if 'Valence band matrix' in line: 194 f.readline() 195 for i in range(0,nmo): 196 for j in range(0,nmo): 197 a=f.readline().split() 198 data['vu'][i,j]=convert(a[2]) 199 data['vue'][i,j]=convert(a[3]) 200 data['vd'][i,j]=convert(a[4]) 201 data['vde'][i,j]=convert(a[5]) 202 if 'Conduction band matrix' in line: 203 f.readline() 204 for i in range(0,nmo): 205 for j in range(0,nmo): 206 a=f.readline().split() 207 data['cu'][i,j]=convert(a[2]) 208 data['cue'][i,j]=convert(a[3]) 209 data['cd'][i,j]=convert(a[4]) 210 data['cde'][i,j]=convert(a[5]) 211 if line=="" or "Trace" in line: 212 break 213 return data 214 215 216################################################################ 217 218def clean_zeros(A,Ae,nsigma=4): 219 A[np.abs(A-Ae) < nsigma*Ae]=0.0 220 221################################################################ 222 223def calc_entropy(a): 224 e=0.0 225 tiny=1e-12 226 n=int(np.sum(a)+0.5) 227 for x in a: 228 if abs(x) > tiny: 229 e-=x*np.log(x) 230 return e/float(n) 231 232################################################################ 233 234 235def gen_excitation(exc,occ): 236 ex=[] 237 for o in occ: 238 found=False 239 for e in exc: 240 if o==e[0]: 241 ex.append(e[1]) 242 found=True 243 break 244 if not found: 245 ex.append(o) 246 return ex 247 248################################################################ 249 250 251def _blob(x,y,area,colour): 252 """ 253 Draws a square-shaped blob with the given area (< 1) at 254 the given coordinates. 255 """ 256 hs = np.sqrt(area) / 2 257 xcorners = np.array([x - hs, x + hs, x + hs, x - hs]) 258 ycorners = np.array([y - hs, y - hs, y + hs, y + hs]) 259 P.fill(xcorners, ycorners, colour, edgecolor=colour) 260################################################################ 261 262def hinton(W,filename="hinton.pdf", maxWeight=None): 263 """ 264 Draws a Hinton diagram for visualizing a weight matrix. 265 Temporarily disables matplotlib interactive mode if it is on, 266 otherwise this takes forever. 267 """ 268 reenable = False 269 if P.isinteractive(): 270 P.ioff() 271 P.clf() 272 ax=P.subplot(111) 273 height, width = W.shape 274 if not maxWeight: 275 maxWeight = 2**np.ceil(np.log(np.max(np.abs(W)))/np.log(2)) 276 277 P.fill(np.array([0,width,width,0]),np.array([0,0,height,height]),'gray') 278 P.axis('off') 279 #P.axis('equal') 280 ax.set_yticklabels(['25','20','15','10','5','0']) 281 for x in xrange(width): 282 for y in xrange(height): 283 _x = x+1 284 _y = y+1 285 w = W[y,x] 286 if w > 0: 287 _blob(_x - 0.5, height - _y + 0.5, min(1,w/maxWeight),'white') 288 elif w < 0: 289 _blob(_x - 0.5, height - _y + 0.5, min(1,-w/maxWeight),'black') 290 if reenable: 291 P.ion() 292 P.title(filename) 293 P.savefig(filename) 294 # P.show() 295################################################################ 296 297 298