simplify integrate function interface

imsrg/__init__.py

@@ -158,9 +158,11 @@ def _get_denom2(href, hh_classifier, ph_classifier, pp_classifier):
 
 IntegrateStats = namedtuple('IntegrateStats', 's fd E0 E2')
 
-def integrate_direct(href0, s_max, generator, reference, step_monitor=None, convergence_check=None):
+def integrate_direct(href0, s_max, generator, step_monitor=None, convergence_check=None):
     from scipy.integrate import ode
 
+    reference = href0.ref
+
     solver = ode(_flow_rhs(generator, reference))
     solver.set_integrator('vode')
     solver.set_initial_value(href0.pack(symmetry='hermitian'))
@@ -193,9 +195,11 @@ def integrate_direct(href0, s_max, generator, reference, step_monitor=None, conv
     return solver.successful(), href, stats
 
 
-def integrate_magnus(href0, s_max, generator, reference, step_monitor=None, convergence_check=None):
+def integrate_magnus(href0, s_max, generator, step_monitor=None, convergence_check=None):
     from scipy.integrate import ode
 
+    reference = href0.ref
+
     solver = ode(_magnus_rhs(generator, href0, reference))
     solver.set_integrator('vode', atol=1e-6)
     solver.set_initial_value(Operator.zero(reference).pack(symmetry='antihermitian'))
@@ -231,15 +235,17 @@ def integrate_magnus(href0, s_max, generator, reference, step_monitor=None, conv
     return solver.successful(), omega, href, stats
 
 
-def integrate_magnus_dopri(href0, s_max, generator, reference, step_monitor=None, convergence_check=None):
+def integrate_magnus_dopri(href0, s_max, generator, step_monitor=None, convergence_check=None):
     from scipy.integrate import ode
 
+    reference = href0.ref
+
     stats = IntegrateStats([], [], [], [])
 
     sdbasis = SDBasis(reference.basis, reference.nparticles)
 
     def solout(s, y):
-        nonlocal stats, sdbasis
+        nonlocal stats, sdbasis, reference
 
         omega = Operator.unpack(y, reference, symmetry='antihermitian')
         href = href0.bch(omega)

pair.py

@@ -1,11 +1,32 @@
 #!/usr/bin/env python3
 
 import numpy as np
+import matplotlib.pyplot as plt
+
 
 from imsrg.operator import Basis, Reference, Operator
 from imsrg import wegner_generator, white_generator, integrate_direct, integrate_magnus, integrate_magnus_dopri
 from imsrg.ci import SDBasis, SDMatrix
 
+
+def show_twobody_part(op):
+    # Get two-body state classifiers and build a permutation that sorts states in the order hh, ph, pp
+    hh_classifier, ph_classifier, pp_classifier = op.ref.classifiers()[2:5]
+    perm = sorted(list(range(op.basis.ntpstates)), key=lambda i: 0 if hh_classifier(*op.basis.tpb[i]) else (
+        1 if ph_classifier(*op.basis.tpb[i]) else 2))
+
+    # get first ph and pp state index
+    first_ph = next(filter(lambda i: ph_classifier(*ref.basis.tpb[perm[i]]), range(ref.basis.ntpstates)), None)
+    first_pp = next(filter(lambda i: pp_classifier(*ref.basis.tpb[perm[i]]), range(ref.basis.ntpstates)), None)
+
+    plt.matshow(np.abs(op.o2[np.ix_(perm, perm)]))
+    plt.axhline(first_ph-0.5)
+    plt.axvline(first_ph-0.5)
+    plt.axhline(first_pp-0.5)
+    plt.axvline(first_pp-0.5)
+    return plt.colorbar()
+
+
 def pairing_hamiltonian(basis, g, *, energies=None, delta_e=None):
     from itertools import product
 
@@ -40,73 +61,71 @@ def pairing_hamiltonian(basis, g, *, energies=None, delta_e=None):
 
 
 basis = Basis(4)
-fermilevel = 1
+fermilevel = 2
+g = 0.5
 sMax = 15
 
 ref = Reference.single(basis, fermilevel)
+sdbasis = SDBasis(basis, ref.nparticles)
 
-hvac = pairing_hamiltonian(basis, g=0.5, delta_e=1)
-href = hvac.normalorder(ref)
+hvac = pairing_hamiltonian(basis, g, delta_e=1)
+href0 = hvac.normalorder(ref)
 
-hh_classifier, ph_classifier, pp_classifier = ref.classifiers()[2:5]
-perm = sorted(list(range(basis.ntpstates)), key=lambda i: 0 if hh_classifier(*basis.tpb[i]) else (1 if ph_classifier(*basis.tpb[i]) else 2))
+print('Starting integration for Pairing Hamiltonian (g={}) with {} levels total and {} filled.'.format(g, basis.nlevels, fermilevel))
 
-first_ph = next(filter(lambda i: ph_classifier(*basis.tpb[perm[i]]), range(basis.ntpstates)), None)
-first_pp = next(filter(lambda i: pp_classifier(*basis.tpb[perm[i]]), range(basis.ntpstates)), None)
+def step_monitor(href, stats):
+    print('s = {:.6f}, e = {:.6f}, E(2) = {:.6f}'.format(stats.s[-1], stats.E0[-1], stats.E2[-1]))
 
-import matplotlib.pyplot as plt
+# Direct integration of the ODE for H, using vode in Adams mode.
+#success, href, stats = integrate_direct(href0, sMax, white_generator, step_monitor)
 
-#plt.matshow(np.abs(href.o2[np.ix_(perm, perm)]))
-#plt.axhline(first_ph-0.5)
-#plt.axvline(first_ph-0.5)
-#plt.axhline(first_pp-0.5)
-#plt.axvline(first_pp-0.5)
-#plt.colorbar()
-#plt.show()
+# Magnus integration using vode in Adams mode.
+success, omega, href, stats = integrate_magnus(href0, sMax, white_generator, step_monitor)
 
-sdbasis = SDBasis(basis, 2*fermilevel)
-sdmat = SDMatrix(sdbasis, hvac)
+# Magnus integration using 8th order Dormand-Prince with a maximum step size of 0.1 to check numerics.
+#success, omega, href, stats = integrate_magnus_dopri(href0, sMax, white_generator, step_monitor)
 
-print(len(sdbasis.cfs))
-print(sdmat.cfmat)
-print(sdmat.eigenvalues())
-
-#y = href.pack(symmetry='hermitian')
-#print(y)
-#htest = Operator.unpack(y, ref, symmetry='hermitian')
-#print(htest.o0)
-#print(htest.o1)
-#print(htest.o2)
-
-#plt.matshow(np.abs(htest.o2[np.ix_(perm, perm)]))
-#plt.axhline(first_ph-0.5)
-#plt.axvline(first_ph-0.5)
-#plt.axhline(first_pp-0.5)
-#plt.axvline(first_pp-0.5)
-#plt.colorbar()
-#plt.show()
+print('Integration finished!')
 
-def step_monitor(href, stats):
-    print('s = {:.6f}, e = {:.6f}, E(2) = {:.6f}'.format(stats.s[-1], stats.E0[-1], stats.E2[-1]))
+print()
 
-#success, omega, href, stats = integrate_magnus(href, sMax, white_generator, ref, step_monitor)
-success, href, stats = integrate_direct(href, sMax, white_generator, ref, step_monitor)
+print('Initial Many-Body Hamiltonian')
+sdmat = SDMatrix(sdbasis, hvac)
+print(sdmat.cfmat)
+print('Eigenvalues:', sdmat.eigenvalues())
 
+print('Final Many-Body Hamiltonian')
 sdmat = SDMatrix(sdbasis, href.normalorder(basis.vacuum))
 print(sdmat.cfmat)
+print('Eigenvalues:', sdmat.eigenvalues())
 
+cb = show_twobody_part(href0)
+plt.xlabel('$i$')
+plt.ylabel('$j$')
+cb.set_label('$|V_{ij}|$')
+plt.title('Initial two-body part')
 
-plt.matshow(np.abs(href.o2[np.ix_(perm, perm)]))
-plt.axhline(first_ph-0.5)
-plt.axvline(first_ph-0.5)
-plt.axhline(first_pp-0.5)
-plt.axvline(first_pp-0.5)
-plt.colorbar()
+cb = show_twobody_part(href)
+plt.xlabel('$i$')
+plt.ylabel('$j$')
+cb.set_label('$|V_{ij}|$')
+plt.title('Final two-body part')
 
 plt.figure()
-plt.plot(stats.s, stats.E0)
-plt.plot(stats.s, stats.E2)
+plt.plot(stats.s, stats.E0, label='$E_0$')
+plt.plot(stats.s, stats.E2, label=r'$E_0 + \mathrm{MBPT}(2)$')
+plt.title('Zero-body flow')
+plt.xlabel('$s$')
+plt.ylabel('$E$ [MeV]')
+plt.legend()
 
 plt.figure()
 plt.plot(stats.s, stats.fd)
-plt.show()
+plt.xlabel('$s$')
+plt.ylabel('$E$ [MeV]')
+plt.title('Flow dependence of many-body eigenvalues')
+
+try:
+    plt.show()
+except KeyboardInterrupt:
+    pass