main.py 16.1 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 pandas as pd
import seaborn as sns
15
16

import matplotlib.pyplot as plt
17
18
import tracemalloc
import os, sys
Jacob August Davison's avatar
Jacob August Davison committed
19
#from memory_profiler import profile
20
21
import itertools
import random
22
import tensornetwork as tn
23
#tn.set_default_backend("tensorflow")
24

25
26
# USER MODULES
from oop_imsrg.spin_sq import *
27
28
29
30
31
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
32
# from oop_imsrg.display_memory import *
33
import oop_imsrg.ci_pairing.cipy_pairing_plus_ph as ci_matrix
34
from oop_imsrg.tests2B import *
35

36
37
38
39
sys.path.append('/mnt/home/daviso53/Research/')
from pyci.density_matrix.density_matrix import density_1b, density_2b
import pyci.imsrg_ci.pyci_p3h as pyci

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]
48
    result_ij = Gnode @ Gnode
49

50

51
52
53
    H0B = E - np.trace(H1B[np.ix_(holes,holes)]) - 0.5*result_ij.tensor

    #print(result_ij.tensor, result_ji.tensor)
54
55
56
    return (H0B, H1B, H2B)


57
# @profile
58
def derivative(t, y, inputs):
59
60
61
62
63
64
65
66
67
68
    """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
69
70
    generator   -- Generator object
    flow        -- Flow object
71
72
73
74
75

    Returns:

    dy -- next step in flow"""

76
77
    hamiltonian, occ_tensors, generator, flow, tspinsq, generator_spin, flow_spin = inputs
    #assert isinstance(hamiltonian, Hamiltonian), "Arg 2 must be Hamiltonian object"
78
79
80
81
    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"

82
    half = int(len(y)/2)
83

84
85
    # Hamiltonian flow
    E, f, G = ravel(y[0:half], hamiltonian.n_sp_states)
86
87
    generator.f = f
    generator.G = G
88
    dE, df, dG = flow.flow(generator)
89

90
91
92
93
94
    # Spin-squared flow
    E_spin, f_spin, G_spin = ravel(y[half::], tspinsq.n_sp_states)
    generator_spin.f = f_spin
    generator_spin.G = G_spin
    dE_spin, df_spin, dG_spin = flow_spin.flow(generator)
95

96
    dy = np.concatenate([unravel(dE, df, dG), unravel(dE_spin, df_spin, dG_spin)], axis=0)
97
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

    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.

123
    Arguments:
124

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

    Returns:

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

    # 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, flow_data_log=0, generator='wegner', output_root='.'):
Davison's avatar
Davison committed
145
    """Main method uses scipy.integrate.ode to solve the IMSRG(2) flow
146
147
148
149
150
151
152
153
154
155
156
157
158
    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)
159
160
161
    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)
162
    output_root   -- specify folder for output files
163
164
165
166
167
168
169
170
171
172
173
174
175

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

    initi = time.time() # start instantiation timer

181
182
183
    if not os.path.exists(output_root):
        os.mkdir(output_root)

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

192
    ot = OccupationTensors(ha.sp_basis, ha.reference)
193
194
195

    generator_dict = {'wegner':WegnerGenerator(ha, ot), 
                      'white':WhiteGenerator(ha),
196
                      'white_mp':WhiteGeneratorMP(ha),
197
198
                      'brillouin':BrillouinGenerator(ha),
                      'imtime':ImTimeGenerator(ha)}
199
200

    wg = generator_dict[generator] #WegnerGenerator(ha, ot)
Jacob August Davison's avatar
Jacob August Davison committed
201
202
    fl = Flow_IMSRG2(ha, ot) 

Davison, Jacob's avatar
Davison, Jacob committed
203
    wg_spin = WhiteGenerator(ss)
204
205
    fl_spin = Flow_IMSRG2(ss, ot)

Jacob August Davison's avatar
Jacob August Davison committed
206
207
    initf = time.time() # finish instantiation timer

208
209
210
211
212
    if verbose:
        print("Initialized objects in {:2.4f} seconds\n".format(initf-initi))

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

    # --- Solve the IM-SRG flow
228
229
    y0 = np.concatenate([unravel(ha.E, ha.f, ha.G), unravel(ss.E, ss.f, ss.G)], axis=0)

230
231
    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)
232

233
    pickle.dump( (H0B, H1B, H2B, eta1B_vac, eta2B_vac), open( output_root+"/vac_coeffs_unevolved.p", "wb" ) )
234

235
    solver = ode(derivative,jac=None)
236
    solver.set_integrator('vode', method='bdf', order=5, nsteps=1000)
237
    solver.set_f_params([ha, ot, wg, fl, ss, wg_spin, fl_spin])
238
239
    solver.set_initial_value(y0, 0.)

240
    sfinal = 50
241
#    ds = 1
242
243
    s_vals = []
    E_vals = []
244
245
246
    
    E_gs_lst = []
    s_expect_lst = []
247
248
249

    iters = 0
    convergence = 0
250

251
252
    if verbose:
        print("iter,      \t    s, \t         E, \t         ||eta1B||, \t         ||eta2B||")
253
        
254
255
    while solver.successful() and solver.t < sfinal:

256
257
        #print('solver success,', solver.successful())

258
        ys = solver.integrate(sfinal, step=True)
259
260
261
262
263

        half = int(len(ys)/2)
        Es, fs, Gs = ravel(ys[0:half], ha.n_sp_states)
        E_spins, f_spins, G_spins = ravel(ys[half::], ss.n_sp_states)

264
265
        s_vals.append(solver.t)
        E_vals.append(Es)
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
        
        # compute density matrices, get expectation values
        H0B, H1B, H2B = get_vacuum_coeffs(Es, fs, Gs, ha.sp_basis, ha.holes)
        SS0B, SS1B, SS2B = get_vacuum_coeffs(E_spins, f_spins, G_spins, ss.sp_basis, ss.holes)
        hme = pyci.matrix(n_holes, n_particles, H0B, H1B, H2B, H2B, imsrg=True)
        w,v = np.linalg.eigh(hme)
        v0 = v[:,0]

        rho1b = density_1b(n_holes,n_particles, weights=v0)
        rho2b = density_2b(n_holes,n_particles, weights=v0)

        contract_1b = np.einsum('ij,ij', rho1b, SS1B)

        rho_reshape_2b = np.reshape(rho2b, (ss.n_sp_states**2,ss.n_sp_states**2))
        h2b_reshape_2b = np.reshape(SS2B, (ss.n_sp_states**2,ss.n_sp_states**2))
        contract_2b = np.einsum('ij,ij', rho_reshape_2b, h2b_reshape_2b)

        s_expect = SS0B + contract_1b + 0.25*contract_2b


        E_gs_lst.append(w[0])
        s_expect_lst.append(s_expect)
        
289

290
291
292
293
294
295
296
297
        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))
298
        
299
        if flow_data_log and iters %10 == 0:
300
            H0B, H1B, H2B = get_vacuum_coeffs(Es, fs, Gs, ha.sp_basis, ha.holes)
301
302
303
            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'))
304

305
306
307
308
#        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" ) )

309
        if len(E_vals) > 100 and abs(E_vals[-1] - E_vals[-2]) < 10**-8 and E_vals[-1] != E_vals[0]:
310
311

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

315
316
317
        # 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
318

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

323
324
        # if iters % 1000 == 0 and verbose:
        #     print('Iteration {:>06d}'.format(iters))
325

326
327
        iters += 1
        #print(solver.successful())
Jacob Davison's avatar
Jacob Davison committed
328
    flowf = time.time()
329
    end = time.time()
330
331
332
    time_str = "{:2.5f}".format(end-start)

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

334
335
    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)
336
    #pickle.dump( coeffs, open( "mixed_state_test/pickled_coeffs/vac_coeffs_evolved.p", "wb" ) )
337
    pickle.dump((H0B, H1B, H2B, eta1B_vac, eta2B_vac), open(output_root+'/vac_coeffs_evolved.p', 'wb'))
338

339
    num_sp = n_holes+n_particles
340

341
    del ha, ot, wg, fl, solver, y0, sfinal
342
343
344
345

    expect_flow_df = pd.DataFrame({'s':s_vals, 'E_gs':E_gs_lst, 's_expect':s_expect_lst})
    pickle.dump(expect_flow_df, open('expect_flow.p', 'wb'))

346
    return (convergence, iters, d, g, pb, num_sp, s_vals, E_vals, time_str)
347
 
348
349

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

351
352

    #test_exact('plots_exact_2b/', main)
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
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
    # 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" ) )

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

397
398
399
    main(4,4, generator='white')
    data = pickle.load(open('expect_flow.p', 'rb'))

Davison, Jacob's avatar
Davison, Jacob committed
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
    fig = plt.figure(figsize=(8,4))
    sns.lineplot(x='s', y=data['E_gs']/data['E_gs'][0], data=data)
    sns.lineplot(x='s', y=data['s_expect']/data['s_expect'][0], data=data)
    plt.legend(['E(s)/E(s=0)', 'SS(s)/SS(s=0)'])
    plt.savefig('flow_conservation.png')


    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, ref=ref, generator='white')
    data = pickle.load(open('expect_flow.p', 'rb'))

    fig = plt.figure(figsize=(8,4))
    sns.lineplot(x='s', y=data['E_gs']/data['E_gs'][0], data=data)
    sns.lineplot(x='s', y=data['s_expect']/data['s_expect'][0], data=data)
    plt.legend(['E(s)/E(s=0)', 'SS(s)/SS(s=0)'])
    plt.savefig('flow_conservation_ensemble.png')

417
418
    # 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]]
419

420
421
    # 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]))
    # # main(4,4, g=5, ref=[1,1,1,1,0,0,0,0])
422

423
    
424
    # main(4,4,g=1.0, pb=0.1, flow_data_log=0, generator='white')
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458

    # 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()
459
460


461
    #-- TESTING TIMINGS -------------------------
462
463
464
465
466
467
468
469
470
471
472
    #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))
473
    
474
    #     del convergence, iters, d, g, pb, sp_states, s_vals, E_vals, time_str