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