main.py 13.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
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
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]
41
    result_ij = Gnode @ Gnode
42

43

44
45
46
    H0B = E - np.trace(H1B[np.ix_(holes,holes)]) - 0.5*result_ij.tensor

    #print(result_ij.tensor, result_ji.tensor)
47
48
49
    return (H0B, H1B, H2B)


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

    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)

76
#    timefi = time.time()
77
    generator.f = f
78
#    timeff = time.time()
79

80
#    timegi = time.time()
81
    generator.G = G
82
83
84
#    timegf = time.time()

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

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

    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.

121
    Arguments:
122

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

    Returns:

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

    # 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))
137
    
138
139
140
141

    return(ravel_E, ravel_f, ravel_G)

# @profile
142
def main(n_holes, n_particles, ref=[], d=1.0, g=0.5, pb=0.0, verbose=1, flow_data_log=0, generator='wegner', output_root='.'):
Davison's avatar
Davison committed
143
    """Main method uses scipy.integrate.ode to solve the IMSRG(2) flow
144
145
146
147
148
149
150
151
152
153
154
155
156
    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)
157
158
159
    verbose       -- toggles output of flow information (default: 1)
    flow_data_log -- toggles output of flow data (pickled IM-SRG coefficients every 10 integrator steps) (default: 0)
    generator     -- specify generator to produce IM-SRG flow (default: wegner)
160
    output_root   -- specify folder for output files
161
162
163
164
165
166
167
168
169
170
171
172
173

    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)
    """
174
    
Jacob August Davison's avatar
Jacob August Davison committed
175
176
177
178
    start = time.time() # start full timer

    initi = time.time() # start instantiation timer

179
180
181
    if not os.path.exists(output_root):
        os.mkdir(output_root)

182
    if ref == []:
183
        ha = PairingHamiltonian2B(n_holes, n_particles, d=d, g=g, pb=pb)
184
        ref = ha.reference # this is just for printing
185
186
    else:
        ha = PairingHamiltonian2B(n_holes, n_particles, ref=ref, d=d, g=g, pb=pb)
Jacob August Davison's avatar
Jacob August Davison committed
187

188
    ot = OccupationTensors(ha.sp_basis, ha.reference)
189
190
191

    generator_dict = {'wegner':WegnerGenerator(ha, ot), 
                      'white':WhiteGenerator(ha),
192
                      'white_mp':WhiteGeneratorMP(ha),
193
194
                      'brillouin':BrillouinGenerator(ha),
                      'imtime':ImTimeGenerator(ha)}
195
196

    wg = generator_dict[generator] #WegnerGenerator(ha, ot)
Jacob August Davison's avatar
Jacob August Davison committed
197
198
199
200
    fl = Flow_IMSRG2(ha, ot) 

    initf = time.time() # finish instantiation timer

201
202
203
204
205
    if verbose:
        print("Initialized objects in {:2.4f} seconds\n".format(initf-initi))

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

    # --- Solve the IM-SRG flow
221
    y0 = unravel(ha.E, ha.f, ha.G)
222
223
224
    H0B, H1B, H2B = get_vacuum_coeffs(ha.E, ha.f, ha.G, ha.sp_basis, ha.holes)
    zero, eta1B_vac, eta2B_vac = get_vacuum_coeffs(0.0, wg.eta1B, wg.eta2B, ha.sp_basis, ha.holes)
    pickle.dump( (H0B, H1B, H2B, eta1B_vac, eta2B_vac), open( output_root+"/vac_coeffs_unevolved.p", "wb" ) )
225

226
    solver = ode(derivative,jac=None)
227
    solver.set_integrator('vode', method='bdf', order=5, nsteps=1000)
228
229
230
    solver.set_f_params(ha, ot, wg, fl)
    solver.set_initial_value(y0, 0.)

231
    sfinal = 50
232
#    ds = 1
233
234
235
236
237
    s_vals = []
    E_vals = []

    iters = 0
    convergence = 0
238

239
240
    if verbose:
        print("iter,      \t    s, \t         E, \t         ||eta1B||, \t         ||eta2B||")
241
        
242
243
    while solver.successful() and solver.t < sfinal:

244
245
        #print('solver success,', solver.successful())

246
247
        ys = solver.integrate(sfinal, step=True)
        Es, fs, Gs = ravel(ys, ha.n_sp_states)
248
249
        s_vals.append(solver.t)
        E_vals.append(Es)
250

251
252
253
254
255
256
257
258
        if iters %10 == 0 and verbose: 
            norm_eta1B = np.linalg.norm(np.ravel(wg.eta1B))
            norm_eta2B = np.linalg.norm(np.ravel(wg.eta2B))
            print("{:>6d}, \t {:0.4f}, \t {:0.8f}, \t {:0.8f}, \t {:0.8f}".format(iters, 
                                                                                  solver.t, 
                                                                                  Es,
                                                                                  norm_eta1B,
                                                                                  norm_eta2B))
259
        
260
        if flow_data_log and iters %10 == 0:
261
            H0B, H1B, H2B = get_vacuum_coeffs(Es, fs, Gs, ha.sp_basis, ha.holes)
262
263
264
            zero, eta1B_vac, eta2B_vac = get_vacuum_coeffs(0.0, wg.eta1B, wg.eta2B, ha.sp_basis, ha.holes)
            fname = output_root+'/vac_coeffs_flow_c{}.p'.format(iters)
            pickle.dump((solver.t, H0B, H1B, H2B, eta1B_vac, eta2B_vac), open(fname, 'wb'))
265

266
267
268
269
#        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" ) )

Davison, Jacob's avatar
Davison, Jacob committed
270
        if len(E_vals) > 100 and abs(E_vals[-1] - E_vals[-2]) < 10**-8:
271
272

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

276
277
278
        # 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
279

280
        if iters > 10000:
281
            if verbose: print("---- Energy diverged at iter {:>06d} with energy {:3.8f}\n".format(iters,E_vals[-1]))
282
            break
283

284
285
        # if iters % 1000 == 0 and verbose:
        #     print('Iteration {:>06d}'.format(iters))
286

287
288
        iters += 1
        #print(solver.successful())
Jacob Davison's avatar
Jacob Davison committed
289
    flowf = time.time()
290
    end = time.time()
291
292
293
    time_str = "{:2.5f}".format(end-start)

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

295
296
    H0B, H1B, H2B = get_vacuum_coeffs(Es, fs, Gs, ha.sp_basis, ha.holes)
    zero, eta1B_vac, eta2B_vac = get_vacuum_coeffs(0.0, wg.eta1B, wg.eta2B, ha.sp_basis, ha.holes)
297
    #pickle.dump( coeffs, open( "mixed_state_test/pickled_coeffs/vac_coeffs_evolved.p", "wb" ) )
298
    pickle.dump((H0B, H1B, H2B, eta1B_vac, eta2B_vac), open(output_root+'/vac_coeffs_evolved.p', 'wb'))
299

300
    num_sp = n_holes+n_particles
301

302
    del ha, ot, wg, fl, solver, y0, sfinal
Jacob August Davison's avatar
Jacob August Davison committed
303
    
304
    return (convergence, iters, d, g, pb, num_sp, s_vals, E_vals, time_str)
305
 
306
307

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

309
310

    #test_exact('plots_exact_2b/', main)
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
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
    # 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" ) )

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

355
356
    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]]
357

358
    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]))
359
360
    # main(4,4, g=5, ref=[1,1,1,1,0,0,0,0])

361
362
    ref = 0.8*np.array([1,1,1,1,0,0,0,0])+0.2*np.array([1,1,0,0,1,1,0,0])
    main(4,4,g=2, ref=ref, pb=0.0, generator='white')
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396

    # 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()
397
398


399
    #-- TESTING TIMINGS -------------------------
400
401
402
403
404
405
406
407
408
409
410
    #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))
411
    
412
    #     del convergence, iters, d, g, pb, sp_states, s_vals, E_vals, time_str