Added saving and initialization using all ADM's to HEOM (before just tier/order 0 density matrix). Also added option to print every nprint steps instead of every step (default = 1).

Author: suarez-va

src/libra_py/dynamics/heom/compute.py

@@ -144,7 +144,8 @@ def transform_adm(rho, rho_scaled, aux_memory, params, direction):
         pack_mtx(aux_memory["rho_unpacked"], rho)
 
 
-def run_dynamics(dyn_params, Ham, rho_init):
+#def run_dynamics(dyn_params, Ham, rho_init):
+def run_dynamics(dyn_params, Ham, rho_init, adm_init=None):
     """
     This functions integrates the HEOM for a given system's Hamiltonian, initial conditions, and bath parameters
 
@@ -188,6 +189,9 @@ def run_dynamics(dyn_params, Ham, rho_init):
             * **dyn_params["nsteps"]** ( int )
                 How many steps of the dynamics to perform [ default: 10 ]
 
+            * **dyn_params["nprint"]** ( int )
+                The number of steps in between data saves [ default: 1 ]
+
             * **dyn_params["verbosity"]** ( int )
                 The level of the run-time printout of any useful information [ default: -1 ]
 
@@ -317,7 +321,7 @@ def run_dynamics(dyn_params, Ham, rho_init):
                       "temperature": 300.0,
                       "el_phon_couplings": initialize_el_phonon_couplings(Ham.num_of_cols),
 
-                      "dt": 0.1 * units.fs2au, "nsteps": 10,
+                      "dt": 0.1 * units.fs2au, "nsteps": 10, "nprint": 1,
                       "verbosity": -1, "progress_frequency": 0.1,
 
                       "truncation_scheme": 1, "do_scale": 0,
@@ -334,6 +338,7 @@ def run_dynamics(dyn_params, Ham, rho_init):
     comn.check_input(params, default_params, critical_params)
 
     nsteps = params["nsteps"]
+    #nprint = params["nprint"]
     print_freq = int(params["progress_frequency"] * nsteps)
 
     # ============= System ======================
@@ -389,9 +394,16 @@ def run_dynamics(dyn_params, Ham, rho_init):
         aux_memory["drho_unpacked_scaled"].append(CMATRIX(nquant, nquant))
 
     # Initial conditions
-    x_ = Py2Cpp_int(list(range(nquant)))
-    y_ = Py2Cpp_int(list(range(nquant)))
-    push_submatrix(rho, rho_init, x_, y_)
+    if adm_init is None:
+        x_ = Py2Cpp_int(list(range(nquant)))
+        y_ = Py2Cpp_int(list(range(nquant)))
+        push_submatrix(rho, rho_init, x_, y_)
+    else:
+        # should test to make sure that len(adm_init) == nn_tot-1:
+        aux_memory["rho_unpacked"][0] = rho_init
+        for n in range(1,nn_tot):
+            aux_memory["rho_unpacked"][n] = adm_init[n-1]
+        pack_mtx(aux_memory["rho_unpacked"], rho)
 
     # unpack_mtx(aux_memory["rho_unpacked"], rho)
 
@@ -423,7 +435,8 @@ def run_dynamics(dyn_params, Ham, rho_init):
             aux_print_matrices(0, aux_memory["rho_unpacked_scaled"])
 
     # Initialize savers
-    _savers = save.init_heom_savers(params, nquant)
+    #_savers = save.init_heom_savers(params, nquant)
+    _savers = save.init_heom_savers(params, nquant, nn_tot)
 
     # ============== Propagation =============
 

src/libra_py/dynamics/heom/save.py

@@ -54,7 +54,7 @@
 
 # ===================== HEOM output ====================
 
-def init_heom_data(saver, hdf5_output_level, _nsteps, _nquant):
+def init_heom_data(saver, hdf5_output_level, _nsteps, _nquant, _nn_tot):
 
     if hdf5_output_level >= 1:
         # Time axis (integer steps)
@@ -67,8 +67,13 @@ def init_heom_data(saver, hdf5_output_level, _nsteps, _nquant):
         # System's density matrix
         saver.add_dataset("denmat", (_nsteps, _nquant, _nquant), "C")
 
+    if hdf5_output_level >= 4:
+        # Auxiliary density matrices
+        saver.add_dataset("auxdenmat", (_nsteps, _nn_tot-1, _nquant, _nquant), "C")
 
-def init_heom_savers(params, nquant):
+
+#def init_heom_savers(params, nquant):
+def init_heom_savers(params, nquant, nn_tot):
 
     # ================ Create savers ==================
     prefix = params["prefix"]
@@ -87,7 +92,8 @@ def init_heom_savers(params, nquant):
     if hdf5_output_level > 0:
         _savers["hdf5_saver"] = data_savers.hdf5_saver(F"{prefix}/data.hdf", properties_to_save)
         _savers["hdf5_saver"].set_compression_level(params["use_compression"], params["compression_level"])
-        init_heom_data(_savers["hdf5_saver"], hdf5_output_level, params["nsteps"], nquant)
+        #init_heom_data(_savers["hdf5_saver"], hdf5_output_level, params["nsteps"], nquant)
+        init_heom_data(_savers["hdf5_saver"], hdf5_output_level, math.ceil(params["nsteps"] / params["nprint"]), nquant, nn_tot)
 
     # ====== TXT ========
     if params["txt_output_level"] > 0:
@@ -98,26 +104,40 @@ def init_heom_savers(params, nquant):
 
     if mem_output_level > 0:
         _savers["mem_saver"] = data_savers.mem_saver(properties_to_save)
-        init_heom_data(_savers["mem_saver"], mem_output_level, params["nsteps"], nquant)
+        #init_heom_data(_savers["mem_saver"], mem_output_level, params["nsteps"], nquant)
+        init_heom_data(_savers["mem_saver"], mem_output_level, math.ceil(params["nsteps"] / params["nprint"]), nquant, nn_tot)
 
     return _savers
 
 
-def save_heom_hdf5(step, saver, params, denmat):
+#def save_heom_hdf5(step, saver, params, denmat):
+def save_heom_hdf5(step, saver, params, rho_unpacked):
 
     dt = params["dt"]
+    nprint = params["nprint"]
     hdf5_output_level = params["hdf5_output_level"]
 
+    if step % nprint == 0:
+        step = int(step / nprint)
+    else:
+        return None
+
     if hdf5_output_level >= 1:
         # Timestep
-        saver.save_scalar(step, "timestep", step)
+        saver.save_scalar(step, "timestep", nprint * step)
 
         # Actual time
-        saver.save_scalar(step, "time", step * dt)
+        saver.save_scalar(step, "time", nprint * step * dt)
 
     if hdf5_output_level >= 3:
         # Average adiabatic density matrices
-        saver.save_matrix(step, "denmat", denmat)
+        #saver.save_matrix(step, "denmat", denmat)
+        saver.save_matrix(step, "denmat", rho_unpacked[0])
+
+    if hdf5_output_level >= 4:
+        nn_tot = len(rho_unpacked)
+        for n in range(1,nn_tot):
+            saver.save_multi_matrix(step, n-1, "auxdenmat", rho_unpacked[n])
 
 
 def save_heom_data(_savers, step, print_freq, params, rho_unpacked):
@@ -129,12 +149,14 @@ def save_heom_data(_savers, step, print_freq, params, rho_unpacked):
 
     # Save properties
     if _savers["hdf5_saver"] is not None:
-        save_heom_hdf5(step, _savers["hdf5_saver"], params, rho_unpacked[0])
+        #save_heom_hdf5(step, _savers["hdf5_saver"], params, rho_unpacked[0])
+        save_heom_hdf5(step, _savers["hdf5_saver"], params, rho_unpacked)
 
     if _savers["txt_saver"] is not None:
         pass
 
     if _savers["mem_saver"] is not None:
         prms = dict(params)
         prms["hdf5_output_level"] = prms["mem_output_level"]
-        save_heom_hdf5(step, _savers["mem_saver"], prms, rho_unpacked[0])
+        #save_heom_hdf5(step, _savers["mem_saver"], prms, rho_unpacked[0])
+        save_heom_hdf5(step, _savers["mem_saver"], prms, rho_unpacked)