main.py 12.4 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
16
import tracemalloc
import os, sys
Jacob August Davison's avatar
Jacob August Davison committed
17
#from memory_profiler import profile
18
19
import itertools
import random
20
import tensornetwork as tn
21
#tn.set_default_backend("tensorflow")
22

23
24
25
26
27
28
# user files
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
29
# from oop_imsrg.display_memory import *
30
import oop_imsrg.ci_pairing.cipy_pairing_plus_ph as ci_matrix
31
from oop_imsrg.tests2B import *
32

33
34
35
36
37
38
39
40
41
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
42

43
44
45
46
47
    H0B = E - np.trace(f[np.ix_(holes,holes)]) + 0.5*result.tensor

    return (H0B, H1B, H2B)


48
49
50
51
52
53
54
55
56
57
58
59
# @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
60
61
    generator   -- Generator object
    flow        -- Flow object
62
63
64
65
66
67
68
69
70
71
72
73

    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)

74
#    timefi = time.time()
75
    generator.f = f
76
#    timeff = time.time()
77

78
#    timegi = time.time()
79
    generator.G = G
80
81
82
#    timegf = time.time()

#    timefli = time.time()
83
    dE, df, dG = flow.flow(generator)
84
85
86
87
88
89
90
#    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))
91

92
    dy = unravel(dE, df, dG)
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

    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.

119
    Arguments:
120

121
    y       -- 1D data array (output from unravel)
122
123
124
125
    bas_len -- length of single particle basis

    Returns:

126
127
    E, f, G -- normal-ordered pieces of Hamiltonian
    """
128
129
130
131
132
133
134

    # 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))
135
    
136
137
138
139

    return(ravel_E, ravel_f, ravel_G)

# @profile
140
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
141
    """Main method uses scipy.integrate.ode to solve the IMSRG(2) flow
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
    equations.

    Arguments:
    
    n_holes -- number of hole states in the SP basis (int)
    n_particles -- number of particle states in the SP basis (int)
    
    Keyword arguments:
    
    ref           -- reference state for the IM-SRG flow (1D array)
    d             -- energy spacing in Pairing model (default: 1.0)
    g             -- pairing strength in Pairing model (default: 0.5)
    pb            -- pair-breaking in Pairing-plus-ph model (default: 0.0)
    verbose       -- toggles output of flow information
    flow_data_log -- toggles output of flow data (pickled IM-SRG coefficients every 10 integrator steps)
    generator     -- specify generator to produce IM-SRG flow

    Returns:

    convergence -- 0 if diverged, 1 if converged (little bit outdated)
    iters       -- number of iterations before integrator stopped
    d           -- energy spacing in pairing model
    g           -- pairing strength in pairing model
    pb          -- pair-breaking strength in Pairing-plus-ph model
    num_sp      -- number of single particle states
    s_vals      -- 1D array of flow parameter values
    E_vals      -- 1D array of zero-body energy values
    time_str    -- time taken for flow completion (string)
    """
171

Jacob August Davison's avatar
Jacob August Davison committed
172
173
174
175
    start = time.time() # start full timer

    initi = time.time() # start instantiation timer

176
    if ref == []:
177
        ha = PairingHamiltonian2B(n_holes, n_particles, d=d, g=g, pb=pb)
178
        ref = ha.reference # this is just for printing
179
180
    else:
        ha = PairingHamiltonian2B(n_holes, n_particles, ref=ref, d=d, g=g, pb=pb)
Jacob August Davison's avatar
Jacob August Davison committed
181

182
    ot = OccupationTensors(ha.sp_basis, ha.reference)
183
184
185

    generator_dict = {'wegner':WegnerGenerator(ha, ot), 
                      'white':WhiteGenerator(ha),
186
                      'white_mp':WhiteGeneratorMP(ha),
187
188
                      'brillouin':BrillouinGenerator(ha),
                      'imtime':ImTimeGenerator(ha)}
189
190

    wg = generator_dict[generator] #WegnerGenerator(ha, ot)
Jacob August Davison's avatar
Jacob August Davison committed
191
192
193
194
    fl = Flow_IMSRG2(ha, ot) 

    initf = time.time() # finish instantiation timer

195
196
197
198
199
    if verbose:
        print("Initialized objects in {:2.4f} seconds\n".format(initf-initi))

    if verbose:
        print("""Pairing model IM-SRG(2) flow:
200
        Generator      = {}
201
202
203
204
205
206
        d              = {:2.4f}
        g              = {:2.4f}
        pb             = {:2.4f}
        SP basis size  = {:2d}
        n_holes        = {:2d}
        n_particles    = {:2d}
207
        ref            = {d}""".format(generator, ha.d, ha.g, ha.pb, ha.n_sp_states,
208
209
210
211
                                       len(ha.holes), len(ha.particles),
                                       d=ref) )
    if verbose:
        print("Flowing...")
Jacob August Davison's avatar
Jacob August Davison committed
212
    flowi = time.time()
213
214

    # --- Solve the IM-SRG flow
215
    y0 = unravel(ha.E, ha.f, ha.G)
216
217
    coeffs = get_vacuum_coeffs(ha.E, ha.f, ha.G, ha.sp_basis, ha.holes)
    pickle.dump( coeffs, open( "./vac_coeffs_unevolved.p", "wb" ) )
218

219
    solver = ode(derivative,jac=None)
220
    solver.set_integrator('vode', method='bdf', order=5, nsteps=1000)
221
222
223
    solver.set_f_params(ha, ot, wg, fl)
    solver.set_initial_value(y0, 0.)

224
225
    sfinal = 50
    ds = 1
226
227
228
229
230
231
232
    s_vals = []
    E_vals = []

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

233
234
        #print('solver success,', solver.successful())

235
236
        ys = solver.integrate(sfinal, step=True)
        Es, fs, Gs = ravel(ys, ha.n_sp_states)
237
238
        s_vals.append(solver.t)
        E_vals.append(Es)
239

240

241
242
        # if iters == 176:
        #     break
243
        if iters %10 == 0 and verbose: print("iter: {:>6d} \t scale param: {:0.4f} \t E = {:0.8f}".format(iters, solver.t, Es))
244
245
246
247
248
        
        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'))
249

250
251
252
253
#        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" ) )

254
        if len(E_vals) > 100 and abs(E_vals[-1] - E_vals[-2]) < 10**-8 and E_vals[-1] != E_vals[0]:
255
256

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

260
261
262
        # 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
263

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

268
        if iters % 1000 == 0 and verbose:
269
            print('Iteration {:>06d}'.format(iters))
270

271
272
        iters += 1
        #print(solver.successful())
Jacob Davison's avatar
Jacob Davison committed
273
    flowf = time.time()
274
    end = time.time()
275
276
277
    time_str = "{:2.5f}".format(end-start)

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

279
280
281
282
    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'))

283
    num_sp = n_holes+n_particles
284

285
    del ha, ot, wg, fl, solver, y0, sfinal, ds
Jacob August Davison's avatar
Jacob August Davison committed
286
    
287
    return (convergence, iters, d, g, pb, num_sp, s_vals, E_vals, time_str)
288
 
289
290

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

292
293

    #test_exact('plots_exact_2b/', main)
294
295
296
297
298
299
300
301
302
303
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
    # 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" ) )

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

338
339
    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]]
340

341
    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]))
342
343
    # main(4,4, g=5, ref=[1,1,1,1,0,0,0,0])

344
    
345
    main(4,4,g=0.5, flow_data_log=0, generator='imtime')
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379

    # 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()
380
381


382
    #-- TESTING TIMINGS -------------------------
383
384
385
386
387
388
389
390
391
392
393
    #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))
394
    
395
    #     del convergence, iters, d, g, pb, sp_states, s_vals, E_vals, time_str