[fix] change h5 r/w to with statement and use gc

* consistently use with statement for h5 r/w across all modules * add gc.collect() to ensure proper memory release after dmft_step and during csc_flow
TRIQS · Aug 7, 2024 · bcc7b23 · bcc7b23
1 parent 753c0a0
commit bcc7b23
Show file tree

Hide file tree

Showing 12 changed files with 292 additions and 260 deletions.
diff --git a/python/solid_dmft/csc_flow.py b/python/solid_dmft/csc_flow.py
@@ -29,6 +29,7 @@
 import subprocess
 import shlex
 import os
+import gc
 import numpy as np
 
 # triqs
@@ -89,8 +90,8 @@ def _run_w90converter(seedname, tolerance):
 
     # Checks if creating of rot_mat succeeded
     if mpi.is_master_node():
-        with HDFArchive(seedname+'.h5', 'r') as archive:
-            assert archive['dft_input']['use_rotations'], 'Creation of rot_mat failed in W90 converter'
+        with HDFArchive(seedname+'.h5', 'r') as ar:
+            assert ar['dft_input']['use_rotations'], 'Creation of rot_mat failed in W90 converter'
     mpi.barrier()
 
 def _full_qe_run(seedname, dft_params, mode):
@@ -248,7 +249,8 @@ def csc_flow_control(general_params, solver_params, dft_params, gw_params, advan
     """
 
     # Removes legacy file vasp.suppress_projs if present
-    vasp.remove_legacy_projections_suppressed()
+    if dft_params['dft_code'] == 'vasp':
+        vasp.remove_legacy_projections_suppressed()
 
     # if GAMMA file already exists, load it by doing extra DFT iterations
     if dft_params['dft_code'] == 'vasp' and os.path.exists('GAMMA'):
@@ -259,9 +261,9 @@ def csc_flow_control(general_params, solver_params, dft_params, gw_params, advan
     # Reads in iteration offset if restarting
     iteration_offset = 0
     if mpi.is_master_node() and os.path.isfile(general_params['seedname']+'.h5'):
-        with HDFArchive(general_params['seedname']+'.h5', 'r') as archive:
-            if 'DMFT_results' in archive and 'iteration_count' in archive['DMFT_results']:
-                iteration_offset = archive['DMFT_results']['iteration_count']
+        with HDFArchive(general_params['seedname']+'.h5', 'r') as ar:
+            if 'DMFT_results' in ar and 'iteration_count' in ar['DMFT_results']:
+                iteration_offset = ar['DMFT_results']['iteration_count']
     iteration_offset = mpi.bcast(iteration_offset)
 
     iter_dmft = iteration_offset+1
@@ -359,6 +361,9 @@ def csc_flow_control(general_params, solver_params, dft_params, gw_params, advan
         end_time_dft = timer()
         mpi.report('  solid_dmft: DFT cycle took {:10.3f} seconds'.format(end_time_dft-start_time_dft))
 
+        del sum_k
+        gc.collect()
+
     # Kills background VASP process for clean end
     if mpi.is_master_node() and dft_params['dft_code'] == 'vasp':
         print('  solid_dmft: Stopping VASP\n', flush=True)

diff --git a/python/solid_dmft/dmft_cycle.py b/python/solid_dmft/dmft_cycle.py
diff --git a/python/solid_dmft/dmft_tools/afm_mapping.py b/python/solid_dmft/dmft_tools/afm_mapping.py
@@ -24,6 +24,7 @@
 ################################################################################
 
 import numpy as np
+from h5 import HDFArchive
 import triqs.utility.mpi as mpi
 
 def determine(general_params, archive, n_inequiv_shells):
@@ -36,42 +37,43 @@ def determine(general_params, archive, n_inequiv_shells):
     afm_mapping = None
     if mpi.is_master_node():
         # Reads mapping from h5 archive if it exists already from a previous run
-        if 'afm_mapping' in archive['DMFT_input']:
-            afm_mapping = archive['DMFT_input']['afm_mapping']
-        elif len(general_params['magmom']) == n_inequiv_shells:
-            # find equal or opposite spin imps, where we use the magmom array to
-            # identity those with equal numbers or opposite
-            # [copy Yes/False, from where, switch up/down channel]
-            afm_mapping = [None] * n_inequiv_shells
-            abs_moms = np.abs(general_params['magmom'])
-
-            for icrsh in range(n_inequiv_shells):
-                # if the moment was seen before ...
-                previous_occurences = np.nonzero(np.isclose(abs_moms[:icrsh], abs_moms[icrsh]))[0]
-                if previous_occurences.size > 0:
-                    # find the source imp to copy from
-                    source = np.min(previous_occurences)
-                    # determine if we need to switch up and down channel
-                    switch = np.isclose(general_params['magmom'][icrsh], -general_params['magmom'][source])
-
-                    afm_mapping[icrsh] = [True, source, switch]
-                else:
-                    afm_mapping[icrsh] = [False, icrsh, False]
-
-
-            print('AFM calculation selected, mapping self energies as follows:')
-            print('imp  [copy sigma, source imp, switch up/down]')
-            print('---------------------------------------------')
-            for i, elem in enumerate(afm_mapping):
-                print('{}: {}'.format(i, elem))
-            print('')
-
-            archive['DMFT_input']['afm_mapping'] = afm_mapping
-
-        # if anything did not work set afm_order false
-        else:
-            print('WARNING: couldn\'t determine afm mapping. No mapping used.')
-            general_params['afm_order'] = False
+        with HDFArchive(archive, 'a') as ar:
+            if 'afm_mapping' in ar['DMFT_input']:
+                afm_mapping = ar['DMFT_input']['afm_mapping']
+            elif len(general_params['magmom']) == n_inequiv_shells:
+                # find equal or opposite spin imps, where we use the magmom array to
+                # identity those with equal numbers or opposite
+                # [copy Yes/False, from where, switch up/down channel]
+                afm_mapping = [None] * n_inequiv_shells
+                abs_moms = np.abs(general_params['magmom'])
+
+                for icrsh in range(n_inequiv_shells):
+                    # if the moment was seen before ...
+                    previous_occurences = np.nonzero(np.isclose(abs_moms[:icrsh], abs_moms[icrsh]))[0]
+                    if previous_occurences.size > 0:
+                        # find the source imp to copy from
+                        source = np.min(previous_occurences)
+                        # determine if we need to switch up and down channel
+                        switch = np.isclose(general_params['magmom'][icrsh], -general_params['magmom'][source])
+
+                        afm_mapping[icrsh] = [True, source, switch]
+                    else:
+                        afm_mapping[icrsh] = [False, icrsh, False]
+
+
+                print('AFM calculation selected, mapping self energies as follows:')
+                print('imp  [copy sigma, source imp, switch up/down]')
+                print('---------------------------------------------')
+                for i, elem in enumerate(afm_mapping):
+                    print('{}: {}'.format(i, elem))
+                print('')
+
+                ar['DMFT_input']['afm_mapping'] = afm_mapping
+
+            # if anything did not work set afm_order false
+            else:
+                print('WARNING: couldn\'t determine afm mapping. No mapping used.')
+                general_params['afm_order'] = False
 
     general_params['afm_order'] = mpi.bcast(general_params['afm_order'])
     if general_params['afm_order']:

diff --git a/python/solid_dmft/dmft_tools/convergence.py b/python/solid_dmft/dmft_tools/convergence.py
@@ -29,6 +29,7 @@
 import numpy as np
 
 # triqs
+from h5 import HDFArchive
 from triqs.gf import MeshImFreq, MeshImTime, MeshReFreq, BlockGf
 from solid_dmft.dmft_tools import solver
 
@@ -183,13 +184,14 @@ def prep_conv_obs(h5_archive):
         conv array for calculation
     """
 
-    # determine number of impurities
-    n_inequiv_shells = h5_archive['dft_input']['n_inequiv_shells']
+    with HDFArchive(h5_archive, 'r') as ar:
+        # determine number of impurities
+        n_inequiv_shells = ar['dft_input']['n_inequiv_shells']
 
-    # check for previous iterations
-    conv_prev = []
-    if 'convergence_obs' in h5_archive['DMFT_results']:
-        conv_prev = h5_archive['DMFT_results']['convergence_obs']
+        # check for previous iterations
+        conv_prev = []
+        if 'convergence_obs' in ar['DMFT_results']:
+            conv_prev = ar['DMFT_results']['convergence_obs']
 
     # prepare observable dicts
     if len(conv_prev) > 0:

diff --git a/python/solid_dmft/dmft_tools/greens_functions_mixer.py b/python/solid_dmft/dmft_tools/greens_functions_mixer.py
@@ -28,6 +28,7 @@
 
 # triqs
 import triqs.utility.mpi as mpi
+from h5 import HDFArchive
 from triqs.gf import MeshImFreq, MeshImTime, Gf, make_gf_from_fourier
 from triqs.gf.descriptors import Fourier
 from triqs.gf.tools import inverse
@@ -194,14 +195,15 @@ def mix_g0(solver, general_params, icrsh, archive, G0_freq_previous, it, deg_she
         mpi.report('\n############\n!!!! WARNING !!!! broyden mixing is still in early testing stage ! Use with caution.\n############\n')
         # TODO implement broyden mixing for Sigma
         if mpi.is_master_node():
-            broyler = archive['DMFT_results']['broyler']
-            # calculate the next G0 via broyden scheme
-            G0_broyden_update, broyler[icrsh] = _broyden_update(it, broyler[icrsh], general_params,
-                                                                deg_shell, solver.G0_freq,
-                                                                G0_freq_previous)
-            # store broyden update to h5 archive
-            archive['DMFT_results']['broyler'] = broyler
-            solver.G0_freq << G0_broyden_update
+            with HDFArchive(archive, 'a') as ar:
+                broyler = ar['DMFT_results']['broyler']
+                # calculate the next G0 via broyden scheme
+                G0_broyden_update, broyler[icrsh] = _broyden_update(it, broyler[icrsh], general_params,
+                                                                    deg_shell, solver.G0_freq,
+                                                                    G0_freq_previous)
+                # store broyden update to h5 archive
+                ar['DMFT_results']['broyler'] = broyler
+                solver.G0_freq << G0_broyden_update
 
         solver.G0_freq << mpi.bcast(solver.G0_freq)
 

diff --git a/python/solid_dmft/dmft_tools/initial_self_energies.py b/python/solid_dmft/dmft_tools/initial_self_energies.py
@@ -299,18 +299,19 @@ def _load_sigma_from_h5(h5_archive, iteration):
     internal_path = 'DMFT_results/'
     internal_path += 'last_iter' if iteration == -1 else 'it_{}'.format(iteration)
 
-    n_inequiv_shells = h5_archive['dft_input']['n_inequiv_shells']
+    with HDFArchive(h5_archive, 'r') as ar:
+        n_inequiv_shells = ar['dft_input']['n_inequiv_shells']
 
-    # Loads previous self-energies and DC
-    self_energies = [h5_archive[internal_path]['Sigma_freq_{}'.format(iineq)]
-                     for iineq in range(n_inequiv_shells)]
-    last_g0 = [h5_archive[internal_path]['G0_freq_{}'.format(iineq)]
-                     for iineq in range(n_inequiv_shells)]
-    dc_imp = h5_archive[internal_path]['DC_pot']
-    dc_energy = h5_archive[internal_path]['DC_energ']
+        # Loads previous self-energies and DC
+        self_energies = [ar[internal_path]['Sigma_freq_{}'.format(iineq)]
+                         for iineq in range(n_inequiv_shells)]
+        last_g0 = [ar[internal_path]['G0_freq_{}'.format(iineq)]
+                         for iineq in range(n_inequiv_shells)]
+        dc_imp = ar[internal_path]['DC_pot']
+        dc_energy = ar[internal_path]['DC_energ']
 
-    # Loads density_matrix to recalculate DC if dc_dmft
-    density_matrix = h5_archive[internal_path]['dens_mat_post']
+        # Loads density_matrix to recalculate DC if dc_dmft
+        density_matrix = ar[internal_path]['dens_mat_post']
 
     print('Loaded Sigma_imp0...imp{} '.format(n_inequiv_shells-1)
           + ('at last it ' if iteration == -1 else 'at it {} '.format(iteration)))
@@ -495,9 +496,8 @@ def determine_dc_and_initial_sigma(general_params, gw_params, advanced_params, s
         # Loads Sigma from different calculation
         elif general_params['load_sigma']:
             print('\nFrom {}:'.format(general_params['path_to_sigma']), end=' ')
-            with HDFArchive(general_params['path_to_sigma'], 'r') as sigma_archive:
-                (loaded_sigma, loaded_dc_imp, _,
-                 _, loaded_density_matrix) = _load_sigma_from_h5(sigma_archive, general_params['load_sigma_iter'])
+            (loaded_sigma, loaded_dc_imp, _,
+             _, loaded_density_matrix) = _load_sigma_from_h5(general_params['path_to_sigma'], general_params['load_sigma_iter'])
 
             # Recalculate double counting in case U, J or DC formula changed
             if general_params['dc']:

diff --git a/python/solid_dmft/dmft_tools/interaction_hamiltonian.py b/python/solid_dmft/dmft_tools/interaction_hamiltonian.py
@@ -327,8 +327,8 @@ def _construct_dynamic(sum_k, general_params, icrsh):
     mpi.report('###### Dynamic U calculation ######, load parameters from input archive.')
     U_onsite = None
     if mpi.is_master_node():
-        with HDFArchive(general_params['jobname']+'/'+general_params['seedname']+'.h5', 'r') as archive:
-            U_onsite = archive['dynamic_U']['U_scr']
+        with HDFArchive(general_params['jobname']+'/'+general_params['seedname']+'.h5', 'r') as ar:
+            U_onsite = ar['dynamic_U']['U_scr']
     U_onsite = mpi.bcast(U_onsite)
 
     n_orb = solver.get_n_orbitals(sum_k)[icrsh]['up']

diff --git a/python/solid_dmft/dmft_tools/manipulate_chemical_potential.py b/python/solid_dmft/dmft_tools/manipulate_chemical_potential.py
@@ -31,6 +31,7 @@
 import numpy as np
 
 import triqs.utility.mpi as mpi
+from h5 import HDFArchive
 from triqs.gf import BlockGf, GfImFreq, GfImTime, Fourier, MeshImFreq
 try:
     if mpi.is_master_node():
@@ -242,7 +243,8 @@ def _set_mu_to_gap_middle_with_maxent(general_params, sum_k, gf_lattice_iw, arch
     # Writes spectral function to archive
     if archive is not None:
         unpacked_results = maxent_gf_latt._unpack_maxent_results(maxent_results, mesh)
-        archive['DMFT_results/last_iter']['Alatt_w'] = unpacked_results
+        with HDFArchive(archive, 'a') as ar:
+            ar['DMFT_results/last_iter']['Alatt_w'] = unpacked_results
 
     # Checks if spectral function at Fermi energy below threshold
     spectral_func_threshold = general_params['beta']/np.pi * general_params['mu_gap_gb2_threshold']
@@ -304,7 +306,8 @@ def set_initial_mu(general_params, sum_k, iteration_offset, archive, broadening)
     # If continuing calculation and not updating mu, loads sold value
     if iteration_offset % general_params['mu_update_freq'] != 0:
         if mpi.is_master_node():
-            sum_k.chemical_potential = archive['DMFT_results/last_iter/chemical_potential_pre']
+            with HDFArchive(archive, 'r') as ar:
+                sum_k.chemical_potential = ar['DMFT_results/last_iter/chemical_potential_pre']
         sum_k.chemical_potential = mpi.bcast(sum_k.chemical_potential)
         mpi.report('Chemical potential not updated this step, '
                    + 'reusing loaded one of {:.3f} eV'.format(sum_k.chemical_potential))
@@ -314,7 +317,8 @@ def set_initial_mu(general_params, sum_k, iteration_offset, archive, broadening)
     # chemical_potential_pre from the last run
     previous_mu = None
     if mpi.is_master_node():
-        previous_mu = archive['DMFT_results/last_iter/chemical_potential_pre']
+        with HDFArchive(archive, 'r') as ar:
+            previous_mu = ar['DMFT_results/last_iter/chemical_potential_pre']
     previous_mu = mpi.bcast(previous_mu)
 
     # Runs maxent if spectral weight too low and occupation is close to desired one

diff --git a/python/solid_dmft/dmft_tools/observables.py b/python/solid_dmft/dmft_tools/observables.py
@@ -31,6 +31,7 @@
 
 # triqs
 import triqs.utility.mpi as mpi
+from h5 import HDFArchive
 from triqs.gf import Gf, MeshImTime
 from triqs.atom_diag import trace_rho_op
 from triqs.gf.descriptors import Fourier
@@ -52,13 +53,14 @@ def prep_observables(h5_archive, sum_k):
         observable array for calculation
     """
 
-    # determine number of impurities
-    n_inequiv_shells = h5_archive['dft_input']['n_inequiv_shells']
+    with HDFArchive(h5_archive, 'r') as ar:
+        # determine number of impurities
+        n_inequiv_shells = ar['dft_input']['n_inequiv_shells']
 
-    # check for previous iterations
-    obs_prev = []
-    if 'observables' in h5_archive['DMFT_results']:
-        obs_prev = h5_archive['DMFT_results']['observables']
+        # check for previous iterations
+        obs_prev = []
+        if 'observables' in ar['DMFT_results']:
+            obs_prev = ar['DMFT_results']['observables']
 
     # prepare observable dicts
     if len(obs_prev) > 0:

diff --git a/python/solid_dmft/dmft_tools/results_to_archive.py b/python/solid_dmft/dmft_tools/results_to_archive.py
@@ -23,6 +23,7 @@
 ################################################################################
 
 import os
+from h5 import HDFArchive
 import triqs.utility.mpi as mpi
 
 
@@ -142,34 +143,35 @@ def write(archive, sum_k, general_params, solver_params, solvers, map_imp_solver
                                        previous_mu, density_mat_pre, density_mat, deltaN, dens)
 
     # Saves the results to last_iter
-    archive['DMFT_results']['iteration_count'] = it
-    for key, value in write_to_h5.items():
-        archive['DMFT_results/last_iter'][key] = value
-
-    # Permanently saves to h5 archive every h5_save_freq iterations
-    if ((not is_sampling and it % general_params['h5_save_freq'] == 0)
-            or (is_sampling and it % general_params['sampling_h5_save_freq'] == 0)):
-
-        archive['DMFT_results'].create_group('it_{}'.format(it))
+    with HDFArchive(archive, 'a') as ar:
+        ar['DMFT_results']['iteration_count'] = it
         for key, value in write_to_h5.items():
-            # Full density matrix only written to last_iter - it is large
-            if 'full_dens_mat_' not in key and 'h_loc_diag_' not in key:
-                archive['DMFT_results/it_{}'.format(it)][key] = value
-
-        # Saves CSC input
-        if general_params['csc']:
-            for dft_var in ['dft_update', 'dft_input', 'dft_misc_input']:
-                if dft_var in archive:
-                    archive['DMFT_results/it_{}'.format(it)].create_group(dft_var)
-                    for key, value in archive[dft_var].items():
-                        # do only store changing elements
-                        if key not in ['symm_kpath', 'kpts_cart']:
-                            archive['DMFT_results/it_{}'.format(it)][dft_var][key] = value
-            for band_elem in ['_bands.dat', '_bands.dat.gnu', '_bands.projwfc_up', '_band.dat']:
-                if os.path.isfile('./{}{}'.format(general_params['seedname'], band_elem)):
-                    os.rename('./{}{}'.format(general_params['seedname'], band_elem),
-                              './{}{}_it{}'.format(general_params['seedname'], band_elem, it))
-            for w90_elem in ['_hr.dat', '.wout']:
-                if os.path.isfile('./{}{}'.format(general_params['seedname'], w90_elem)):
-                    os.rename('./{}{}'.format(general_params['seedname'], w90_elem),
-                              './{}_it{}{}'.format(general_params['seedname'], it, w90_elem))
+            ar['DMFT_results/last_iter'][key] = value
+
+        # Permanently saves to h5 archive every h5_save_freq iterations
+        if ((not is_sampling and it % general_params['h5_save_freq'] == 0)
+                or (is_sampling and it % general_params['sampling_h5_save_freq'] == 0)):
+
+            ar['DMFT_results'].create_group('it_{}'.format(it))
+            for key, value in write_to_h5.items():
+                # Full density matrix only written to last_iter - it is large
+                if 'full_dens_mat_' not in key and 'h_loc_diag_' not in key:
+                    ar['DMFT_results/it_{}'.format(it)][key] = value
+
+            # Saves CSC input
+            if general_params['csc']:
+                for dft_var in ['dft_update', 'dft_input', 'dft_misc_input']:
+                    if dft_var in ar:
+                        ar['DMFT_results/it_{}'.format(it)].create_group(dft_var)
+                        for key, value in ar[dft_var].items():
+                            # do only store changing elements
+                            if key not in ['symm_kpath', 'kpts_cart']:
+                                ar['DMFT_results/it_{}'.format(it)][dft_var][key] = value
+                for band_elem in ['_bands.dat', '_bands.dat.gnu', '_bands.projwfc_up', '_band.dat']:
+                    if os.path.isfile('./{}{}'.format(general_params['seedname'], band_elem)):
+                        os.rename('./{}{}'.format(general_params['seedname'], band_elem),
+                                  './{}{}_it{}'.format(general_params['seedname'], band_elem, it))
+                for w90_elem in ['_hr.dat', '.wout']:
+                    if os.path.isfile('./{}{}'.format(general_params['seedname'], w90_elem)):
+                        os.rename('./{}{}'.format(general_params['seedname'], w90_elem),
+                                  './{}_it{}{}'.format(general_params['seedname'], it, w90_elem))