main.py 11.3 KB
Newer Older
Davison's avatar
Davison committed
1
# Main program for IM-SRG(2).
2
3
4
5
6
7
8
9
10
11
12


# Author: Jacob Davison
# Date:   07/10/2019

# import packages, libraries, and modules
# libraries
from scipy.integrate import odeint, ode
import numpy as np
import time
import pickle
13
14

import matplotlib.pyplot as plt
15
#import tensorflow as tf
Jacob August Davison's avatar
Jacob August Davison committed
16
#tf.enable_v2_behavior()
17
#print("GPU available: ",tf.test.is_gpu_available())
18
19
import tracemalloc
import os, sys
Jacob August Davison's avatar
Jacob August Davison committed
20
#from memory_profiler import profile
21
22
import itertools
import random
23
import tensornetwork as tn
24
25
26
#tn.set_default_backend("tensorflow")
#sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(log_device_placement=True))
#sess.close()
27

28
29
30
31
32
33
34
# user files
# sys.path.append('C:\\Users\\davison\\Research\\exact_diagonalization\\')
from oop_imsrg.hamiltonian import *
from oop_imsrg.occupation_tensors import *
from oop_imsrg.generator import *
from oop_imsrg.flow import *
from oop_imsrg.plot_data import *
Jacob Davison's avatar
Jacob Davison committed
35
# from oop_imsrg.display_memory import *
36
import oop_imsrg.ci_pairing.cipy_pairing_plus_ph as ci_matrix
37
from oop_imsrg.tests2B import *
38

39
40
41
42
43
44
45
46
47
def get_vacuum_coeffs(E, f, G, basis, holes):

    H2B = G
    H1B = f - np.trace(G[np.ix_(basis,holes,basis,holes)], axis1=1,axis2=3) 

    Gnode = tn.Node(G[np.ix_(holes,holes,holes,holes)])
    Gnode[0] ^ Gnode[2]
    Gnode[1] ^ Gnode[3]
    result = Gnode @ Gnode
48

49
50
51
52
53
    H0B = E - np.trace(f[np.ix_(holes,holes)]) + 0.5*result.tensor

    return (H0B, H1B, H2B)


54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
# @profile
def derivative(t, y, hamiltonian, occ_tensors, generator, flow):
    """Defines the derivative to pass into ode object.

    Arguments:
    (required by scipy.integrate.ode)
    t -- points at which to solve for y
    y -- in this case, 1D array that contains E, f, G

    (additional parameters)
    hamiltonian -- Hamiltonian object
    occ_tensors -- OccupationTensors object
    generator -- Generator object
    flow -- Flow object

    Returns:

    dy -- next step in flow"""

    assert isinstance(hamiltonian, Hamiltonian), "Arg 2 must be Hamiltonian object"
    assert isinstance(occ_tensors, OccupationTensors), "Arg 3 must be OccupationTensors object"
    assert isinstance(generator, Generator), "Arg 4 must be Generator object"
    assert isinstance(flow, Flow), "Arg 5 must be a Flow object"

    E, f, G = ravel(y, hamiltonian.n_sp_states)

80
#    timefi = time.time()
81
    generator.f = f
82
#    timeff = time.time()
83

84
#    timegi = time.time()
85
    generator.G = G
86
87
88
#    timegf = time.time()

#    timefli = time.time()
89
    dE, df, dG = flow.flow(generator)
90
91
92
93
94
95
96
#    timeflf = time.time()

    # print("""
    # f time:     {:2.4f} s
    # G time:     {:2.4f} s
    # Flow time:  {:2.4f} s
    # """.format(timeff-timefi, timegf-timegi, timeflf-timefli))
97

98
    dy = unravel(dE, df, dG)
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

    return dy

# @profile
def unravel(E, f, G):
    """Transforms E, f, and G into a 1D array. Facilitates
    compatability with scipy.integrate.ode.

    Arguments:

    E, f, G -- normal-ordered pieces of Hamiltonian

    Returns:

    concatenation of tensors peeled into 1D arrays"""
    unravel_E = np.reshape(E, -1)
    unravel_f = np.reshape(f, -1)
    unravel_G = np.reshape(G, -1)

    return np.concatenate([unravel_E, unravel_f, unravel_G], axis=0)

# @profile
def ravel(y, bas_len):
    """Transforms 1D array into E, f, and G. Facilitates
    compatability with scipy.integrate.ode.

    Arugments:

    y -- 1D data array (output from unravel)
    bas_len -- length of single particle basis

    Returns:

    E, f, G -- normal-ordered pieces of Hamiltonian"""

    # bas_len = len(np.append(holes,particles))

    ravel_E = np.reshape(y[0], ())
    ravel_f = np.reshape(y[1:bas_len**2+1], (bas_len, bas_len))
    ravel_G = np.reshape(y[bas_len**2+1:bas_len**2+1+bas_len**4],
                         (bas_len, bas_len, bas_len, bas_len))
140
    
141
142
143
144

    return(ravel_E, ravel_f, ravel_G)

# @profile
145
def main(n_holes, n_particles, ref=[], d=1.0, g=0.5, pb=0.0, verbose=1, flow_data_log=0, generator='wegner'):
Davison's avatar
Davison committed
146
    """Main method uses scipy.integrate.ode to solve the IMSRG(2) flow
147
148
    equations."""

Jacob August Davison's avatar
Jacob August Davison committed
149
150
151
152
    start = time.time() # start full timer

    initi = time.time() # start instantiation timer

153
    if ref == []:
154
        ha = PairingHamiltonian2B(n_holes, n_particles, d=d, g=g, pb=pb)
155
        ref = ha.reference # this is just for printing
156
157
    else:
        ha = PairingHamiltonian2B(n_holes, n_particles, ref=ref, d=d, g=g, pb=pb)
Jacob August Davison's avatar
Jacob August Davison committed
158

159
    ot = OccupationTensors(ha.sp_basis, ha.reference)
160
161
162
163
164
165

    generator_dict = {'wegner':WegnerGenerator(ha, ot), 
                      'white':WhiteGenerator(ha),
                      'white_mp':WhiteGeneratorMP(ha)}

    wg = generator_dict[generator] #WegnerGenerator(ha, ot)
Jacob August Davison's avatar
Jacob August Davison committed
166
167
168
169
    fl = Flow_IMSRG2(ha, ot) 

    initf = time.time() # finish instantiation timer

170
171
172
173
174
    if verbose:
        print("Initialized objects in {:2.4f} seconds\n".format(initf-initi))

    if verbose:
        print("""Pairing model IM-SRG(2) flow:
175
        Generator      = {}
176
177
178
179
180
181
        d              = {:2.4f}
        g              = {:2.4f}
        pb             = {:2.4f}
        SP basis size  = {:2d}
        n_holes        = {:2d}
        n_particles    = {:2d}
182
        ref            = {d}""".format(generator, ha.d, ha.g, ha.pb, ha.n_sp_states,
183
184
185
186
                                       len(ha.holes), len(ha.particles),
                                       d=ref) )
    if verbose:
        print("Flowing...")
Jacob August Davison's avatar
Jacob August Davison committed
187
    flowi = time.time()
188
189

    # --- Solve the IM-SRG flow
190
    y0 = unravel(ha.E, ha.f, ha.G)
191
192
    coeffs = get_vacuum_coeffs(ha.E, ha.f, ha.G, ha.sp_basis, ha.holes)
    pickle.dump( coeffs, open( "./vac_coeffs_unevolved.p", "wb" ) )
193

194
    solver = ode(derivative,jac=None)
195
    solver.set_integrator('vode', method='bdf', order=5, nsteps=1000)
196
197
198
    solver.set_f_params(ha, ot, wg, fl)
    solver.set_initial_value(y0, 0.)

199
200
    sfinal = 50
    ds = 1
201
202
203
204
205
206
207
    s_vals = []
    E_vals = []

    iters = 0
    convergence = 0
    while solver.successful() and solver.t < sfinal:

208
209
        #print('solver success,', solver.successful())

210
211
        ys = solver.integrate(sfinal, step=True)
        Es, fs, Gs = ravel(ys, ha.n_sp_states)
212
213
        s_vals.append(solver.t)
        E_vals.append(Es)
214

215

216
217
        # if iters == 176:
        #     break
218
        if iters %10 == 0 and verbose: print("iter: {:>6d} \t scale param: {:0.4f} \t E = {:0.8f}".format(iters, solver.t, Es))
219
220
221
222
223
        
        if flow_data_log:# and iters %10 == 0:
            H0B, H1B, H2B = get_vacuum_coeffs(Es, fs, Gs, ha.sp_basis, ha.holes)
            fname = 'vac_coeffs_flow_c{}.p'.format(iters)
            pickle.dump((solver.t, H0B, H1B, H2B), open(fname, 'wb'))
224

225
226
227
228
#        if iters %20 == 0 and verbose:
#            coeffs = get_vacuum_coeffs(Es, fs, Gs, ha.sp_basis, ha.holes)
#            pickle.dump( coeffs, open( "mixed_state_test/pickled_coeffs/vac_coeffs_s{}.p".format(iters), "wb" ) )

229
        if len(E_vals) > 100 and abs(E_vals[-1] - E_vals[-2]) < 10**-8 and E_vals[-1] != E_vals[0]:
230
231

            if verbose: print("---- Energy converged at iter {:>06d} with energy {:1.8f}\n".format(iters,E_vals[-1]))
232
233
            convergence = 1
            break
234

235
236
237
        # if len(E_vals) > 100 and abs(E_vals[-1] - E_vals[-2]) > 1:
        #     if verbose: print("---- Energy diverged at iter {:>06d} with energy {:3.8f}\n".format(iters,E_vals[-1]))
        #     break
238

239
        if iters > 20000:
240
            if verbose: print("---- Energy diverged at iter {:>06d} with energy {:3.8f}\n".format(iters,E_vals[-1]))
241
            break
242

243
        if iters % 1000 == 0 and verbose:
244
            print('Iteration {:>06d}'.format(iters))
245

246
247
        iters += 1
        #print(solver.successful())
Jacob Davison's avatar
Jacob Davison committed
248
    flowf = time.time()
249
    end = time.time()
250
251
252
    time_str = "{:2.5f}".format(end-start)

    if verbose: print("IM-SRG(2) converged in {:2.5f} seconds".format(flowf-flowi))
253

254
255
256
257
258
    coeffs = get_vacuum_coeffs(Es, fs, Gs, ha.sp_basis, ha.holes)
    #pickle.dump( coeffs, open( "mixed_state_test/pickled_coeffs/vac_coeffs_evolved.p", "wb" ) )
    pickle.dump(coeffs, open('vac_coeffs_evolved.p', 'wb'))


259
    del ha, ot, wg, fl, solver, y0, sfinal, ds
Jacob August Davison's avatar
Jacob August Davison committed
260
    
261
    return (convergence, iters, d, g, pb, n_holes+n_particles, s_vals, E_vals, time_str)
262
 
263
264

if __name__ == '__main__':
Jacob Davison's avatar
Jacob Davison committed
265

266
267

    #test_exact('plots_exact_2b/', main)
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
    # refs = list(map("".join, itertools.permutations('11110000')))
    # refs = list(dict.fromkeys(refs)) # remove duplicates
    # refs = [list(map(int, list(ref))) for ref in refs]


    # TESTING MIXED STATE --------------------------------------------
    # refs = [[1,1,1,1,0,0,0,0],[1,1,0,0,1,1,0,0],[1,1,0,0,0,0,1,1],
    #         [0,0,1,1,1,1,0,0],[0,0,1,1,0,0,1,1],[0,0,0,0,1,1,1,1]]
    

    # gsws = [0.99, 0.985,0.975,0.95,0.9,0.85,0.8]#,0.75,0.7,0.65,0.6,0.55,0.5]
    
    # E_data_full = []
    # for gsw in gsws:
    #     E_data = []
    #     count = 0
    #     rsw = (1-gsw)
    #     for i in range(len(refs)):
    #         ref = np.zeros_like(refs[0])
    #         E = 0.0
    #         if i == 0: 
    #             ref = refs[0]
    #             data = main(4,4,ref=ref,verbose=0)
    #             E = (data[7])[-1]
    #             count += 1
    #         else:
    #             esum = np.zeros_like(refs[0])
    #             for state in refs[1:count+1]:
    #                 esum = np.add(esum,state)
    #             print(esum)
    #             ref = gsw*np.array(refs[0]) + rsw/count*(esum)
    #             data = main(4,4,ref=ref,verbose=0)
    #             E = (data[7])[-1]
    #             count += 1
    #         print("{:0.8f}, {}, {f}".format(E, sum(ref), f=ref))
    #         E_data.append(E)
    #     E_data_full.append(E_data)

    # exact = ci_matrix.exact_diagonalization(1.0, 0.5, 0.0)

    # pickle.dump( E_data_full, open( "save.p", "wb" ) )

    # ----------------------------------------------------------------

312
313
    refs = [[1,1,1,1,0,0,0,0],[1,1,0,0,1,1,0,0],[1,1,0,0,0,0,1,1],
            [0,0,1,1,1,1,0,0],[0,0,1,1,0,0,1,1]]
314

315
    ref = 0.5*np.asarray(refs[0]) + (1.0-0.5)/4.0*(np.asarray(refs[1]) + np.asarray(refs[2]) + np.asarray(refs[3]) + np.asarray(refs[4]))
316
317
    # main(4,4, g=5, ref=[1,1,1,1,0,0,0,0])

318
319
    
    main(4,4,g=2, ref=ref, flow_data_log=0, generator='white')
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353

    # H1B_true, H2B_true = pickle.load(open('comparison.p','rb'))
    # H1B, H2B = pickle.load(open('vac_coeffs_unevolved.p', 'rb'))
    # print(H1B, H1B_true)

    
    #print(np.array_equal(H2B_true, H2B))
    #main(4,4)
    

    # ref1 = 0.9*refs[0] + 0.1*refs[1]
    # data = main(4,4,ref=ref1)
    # E1 = data[7]

    # ref2 = 0.9*refs[0] + 0.05*(refs[1] + refs[2])
    # data = main(4,4,ref=ref2)
        
#    ref = 0.9*np.array([1,1,1,1,0,0,0,0])+0.1*np.array([0,1,1,1,1,0,0,0])
    
    #main(4,4,ref=ref)

    # exact = ci_matrix.exact_diagonalization(1.0, 0.5, 0.0)

    # fig = plt.figure()

    # for i in range(0,10):
    #     data = main(4,4)
    #     s_vals = data[6]
    #     E_vals = data[7]/exact
    #     plt.plot(s_vals, E_vals)
    # plt.ylim([1+10**-8, 1+10**-6])
    # plt.xlabel('scale param (s)')
    # plt.ylabel('energy')
    # plt.show()
354
355


356
    #-- TESTING TIMINGS -------------------------
357
358
359
360
361
362
363
364
365
366
367
    #holes = int(sys.argv[1])

    # print('convergence,iters,d,g,pb,num states,GS energy,total time')
    # for num_pairs in range(1,21):
    #     #print(n)55
    #     #n = int(num_sp/2)
    #     #holes = num_pairs*2
    #     particles = num_pairs*2

    #     convergence,iters,d,g,pb,sp_states,s_vals,E_vals,time_str = main(holes, particles, verbose=0)
    #     print('{},{},{},{},{},{},{},{}'.format(convergence,iters,d,g,pb,sp_states,E_vals[-1],time_str))
368
    
369
    #     del convergence, iters, d, g, pb, sp_states, s_vals, E_vals, time_str