Add a start index option to the merge_files function

tracking/monitors/tools.py

@@ -6,9 +6,10 @@ and hdf5 files.
 @author: Alexis Gamelin
 """
 
+import numpy as np
 import h5py as hp
 
-def merge_files(files_prefix, files_number, file_name=None):
+def merge_files(files_prefix, files_number, start_idx=0, file_name=None):
     """
     Merge several hdf5 files into one.
     
@@ -25,6 +26,8 @@ def merge_files(files_prefix, files_number, file_name=None):
         Name of the files to merge.
     files_number : int
         Number of files to merge.
+    start_idx : int, optional
+        Start index of the hdf5 files.        
     file_name : str, optional
         Name of the file with the merged data. If None, files_prefix without
         number is used.
@@ -35,7 +38,7 @@ def merge_files(files_prefix, files_number, file_name=None):
     f = hp.File(file_name + ".hdf5", "a")
     
     ## Create file architecture
-    f0 = hp.File(files_prefix + "_" + str(0) + ".hdf5", "r")
+    f0 = hp.File(files_prefix + "_" + str(start_idx) + ".hdf5", "r")
     for group in list(f0):
         f.require_group(group)
         for dataset_name in list(f0[group]):
@@ -52,8 +55,8 @@ def merge_files(files_prefix, files_number, file_name=None):
     f0.close()
     
     ## Copy data
-    for i in range(files_number):
-        fi = hp.File(files_prefix + "_" + str(i) + ".hdf5", "r")
+    for i, file_num in enumerate(range(start_idx, start_idx + files_number)):
+        fi = hp.File(files_prefix + "_" + str(file_num) + ".hdf5", "r")
         for group in list(fi):
             for dataset_name in list(fi[group]):
                 shape = fi[group][dataset_name].shape
@@ -65,8 +68,8 @@ def merge_files(files_prefix, files_number, file_name=None):
                 slice_list.append(slice(length*i,length*(i+1)))
                 if (dataset_name == "freq"):
                     continue
-                if (dataset_name == "time") and (i != 0):
-                    f[group][dataset_name][tuple(slice_list)] = f[group][dataset_name][(length*i) - 1] + fi[group][dataset_name]
+                if (dataset_name == "time") and (file_num != start_idx):
+                    f[group][dataset_name][tuple(slice_list)] = np.max(f[group][dataset_name][:]) + fi[group][dataset_name]
                 else:
                     f[group][dataset_name][tuple(slice_list)] = fi[group][dataset_name]
         fi.close()