main.py 10.7 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
48
49
50
51
52
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
    H0B = E - np.trace(f[np.ix_(holes,holes)]) + 0.5*result.tensor

    return (H0B, H1B, H2B)


53
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
# @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)

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

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

#    timefli = time.time()
88
    dE, df, dG = flow.flow(generator)
89
90
91
92
93
94
95
#    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))
96

97
    dy = unravel(dE, df, dG)
98
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

    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))
139
    
140
141
142
143

    return(ravel_E, ravel_f, ravel_G)

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

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

    initi = time.time() # start instantiation timer

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

158
159
    ot = OccupationTensors(ha.sp_basis, ha.reference)
    wg = WegnerGenerator(ha, ot)
Jacob August Davison's avatar
Jacob August Davison committed
160
161
162
163
    fl = Flow_IMSRG2(ha, ot) 

    initf = time.time() # finish instantiation timer

164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
    if verbose:
        print("Initialized objects in {:2.4f} seconds\n".format(initf-initi))

    if verbose:
        print("""Pairing model IM-SRG(2) flow:
        d              = {:2.4f}
        g              = {:2.4f}
        pb             = {:2.4f}
        SP basis size  = {:2d}
        n_holes        = {:2d}
        n_particles    = {:2d}
        ref            = {d}""".format(ha.d, ha.g, ha.pb, ha.n_sp_states,
                                       len(ha.holes), len(ha.particles),
                                       d=ref) )
    if verbose:
        print("Flowing...")
Jacob August Davison's avatar
Jacob August Davison committed
180
    flowi = time.time()
181
182

    # --- Solve the IM-SRG flow
183
    y0 = unravel(ha.E, ha.f, ha.G)
184
185
    coeffs = get_vacuum_coeffs(ha.E, ha.f, ha.G, ha.sp_basis, ha.holes)
    pickle.dump( coeffs, open( "./vac_coeffs_unevolved.p", "wb" ) )
186

187
188
189
190
191
    solver = ode(derivative,jac=None)
    solver.set_integrator('vode', method='bdf', order=5, nsteps=500)
    solver.set_f_params(ha, ot, wg, fl)
    solver.set_initial_value(y0, 0.)

192
    sfinal = 100
193
194
195
196
197
198
199
200
201
202
    ds = 0.1
    s_vals = []
    E_vals = []

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

        ys = solver.integrate(sfinal, step=True)
        Es, fs, Gs = ravel(ys, ha.n_sp_states)
203
204
        s_vals.append(solver.t)
        E_vals.append(Es)
205
206

        iters += 1
207
208
        # if iters == 176:
        #     break
209
        if iters %10 == 0 and verbose: print("iter: {:>6d} \t scale param: {:0.4f} \t E = {:0.8f}".format(iters, solver.t, Es))
210

211
212
213
214
#        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" ) )

215
        if len(E_vals) > 100 and abs(E_vals[-1] - E_vals[-2]) < 10**-8 and E_vals[-1] != E_vals[0]:
216
217

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

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

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

229
        if iters % 1000 == 0 and verbose:
230
            print('Iteration {:>06d}'.format(iters))
231
232

        
Jacob Davison's avatar
Jacob Davison committed
233
    flowf = time.time()
234
    end = time.time()
235
236
237
    time_str = "{:2.5f}".format(end-start)

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

239
240
241
242
243
    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'))


244
    del ha, ot, wg, fl, solver, y0, sfinal, ds
Jacob August Davison's avatar
Jacob August Davison committed
245
    
246
    return (convergence, iters, d, g, pb, n_holes+n_particles, s_vals, E_vals, time_str)
247
 
248
249

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

251
252

    #test_exact('plots_exact_2b/', main)
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
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
    # 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" ) )

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

297
298
299
300
301
302
303
    # 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]]

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

    test_exact('plots_exact_2b/', main)
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337

    # 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()
338
339


340
    #-- TESTING TIMINGS -------------------------
341
342
343
344
345
346
347
348
349
350
351
    #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))
352
    
353
    #     del convergence, iters, d, g, pb, sp_states, s_vals, E_vals, time_str