hdf5shenanigans.py

#%%
import h5py
from shutil import copyfile
import numpy as np

#%%
N_WORM_GROUPS = 1
FILE_NAME = "g4iw_avg.hdf5"

#%%
file = h5py.File(FILE_NAME, "r+")

#%%
str = "worm-001/ineq-001/g4iw-worm/00001/value"
a = file[str]

#%%
b = a[...]
print(b[0,0,0])
b = b*2
print(b[0,0,0])

#%%
# Go through all g4iw-worm subgroups
g4iw_group_name = "worm-001/ineq-001/g4iw-worm/"
super_indexes = list(file[g4iw_group_name].keys())
g4iw_value_names = [g4iw_group_name + i + "/value" for i in super_indexes]
g4iw_values = np.array([file[g4iw_value_name] for g4iw_value_name
                        in g4iw_value_names])
print(g4iw_values)


#%%
def get_g4iw_worm_sample_generator(g4iw_hdf5_file):
    """Return a generator yielding greens functions from the file."""
    file = g4iw_hdf5_file
    def g4iw_worm_sample_generator():
        g4iw_worm_groups = [s + "/ineq-001/g4iw-worm/" for s in file.keys()
                            if ("worm-" in s) and (s != "worm-last")]
        super_indexes = list(file[g4iw_worm_groups[0]].keys())
        for worm_group in g4iw_worm_groups:
            value_names = [worm_group + i + "/value" for i in super_indexes]
            yield np.array([file[value_name] for value_name in value_names])
    return g4iw_worm_sample_generator


#%%
def get_n_worm_samples(g4iw_hdf5_file):
    """Return the number of worm samples in the given HDF5 file."""
    file = g4iw_hdf5_file
    return sum(1 for s in file.keys() if ("worm-" in s) and (s != "worm-last"))


#%%
# Test generator
file = h5py.File("g4iw_avg.hdf5")
g4iw_samples = get_g4iw_worm_sample_generator(file)
g_sum = 0
for g in g4iw_samples():
    g_sum += g

print(g_sum)


#%%
file.close()