ZH: read/write TimeDiffs from hdf5

simulations/airshower_beacon_simulation/aa_generate_beacon.py

@@ -152,6 +152,78 @@ def append_antenna_hdf5(fname, antenna, columns = [], name='traces', prepend_tim
         for i, col in enumerate(columns, 1*prepend_time):
             dset[i] = col
 
+
+def read_baseline_time_diffs_hdf5(fname):
+    """
+    Read Baseline Time Diff information from HDF5 storage.
+    """
+    with h5py.File(fname, 'r') as fp:
+        group_name = 'baseline_time_diffs'
+        base_dset_name = 'baselines'
+        dset_name = 'time_diffs'
+
+        group = fp[group_name]
+
+        names = group[base_dset_name][:]
+
+        dset = group[dset_name]
+        f_beacon = dset[:,0]
+        true_phase_diffs = dset[:,1]
+        k_periods = dset[:,2]
+
+        return names, f_beacon, true_phase_diffs, k_periods
+
+
+def write_baseline_time_diffs_hdf5(fname, baselines, true_phase_diffs, k_periods, f_beacon, overwrite=True):
+    """
+    Write a combination of baselines, phase_diff, k_period and f_beacon to file.
+
+    Note that f_beacon is allowed to broadcast, but the others are not.
+    """
+
+    if not hasattr(baselines[0], '__len__'):
+        # this is a single baseline
+        N_baselines = 1
+        baselines = [baselines]
+        true_phase_diffs = [true_phase_diffs]
+        k_periods = [k_periods]
+        f_beacon = [f_beacon]
+
+    else:
+        N_baselines = len(baselines)
+
+        # Expand the f_beacon list
+        if not hasattr(f_beacon, '__len__'):
+            f_beacon = [f_beacon]*N_baselines
+
+    assert len(baselines) == len(true_phase_diffs) == len(k_periods) == len(f_beacon)
+
+    with h5py.File(fname, 'a') as fp:
+        group_name = 'baseline_time_diffs'
+        base_dset_name = 'baselines'
+        dset_name = 'time_diffs'
+
+        group = fp.require_group(group_name)
+
+        if base_dset_name in group:
+            if not overwrite:
+                raise NotImplementedError
+            del group[base_dset_name]
+
+        if dset_name in group:
+            if not overwrite:
+                raise NotImplementedError
+            del group[dset_name]
+
+        # save baselines list
+        basenames = np.array([ [b[0].name, b[1].name] for b in baselines ], dtype='S')
+
+        base_dset = group.create_dataset(base_dset_name, data=basenames)
+
+        data = np.vstack( (f_beacon, true_phase_diffs, k_periods) ).T
+        dset = group.create_dataset(dset_name, data=data)
+
+
 if __name__ == "__main__":
     from os import path