tfa: basic implementation

resolve/time_freq_avg.py

+import numpy as np
+import resolve as rve
+
+
+def freq_average(obs, n_freq_chuncks):
+    splitted_freqs = np.array_split(obs.freq, n_freq_chuncks)
+    splitted_obs = []
+    for ff in splitted_freqs:
+        splitted_obs.append(obs.restrict_by_freq(ff[0], ff[-1]))
+
+    obs_avg = []
+    for obsi in splitted_obs:
+        new_vis = np.mean(obsi.vis.val, axis=2, keepdims=True)
+        cov = 1 / obsi.weight.val
+        new_cov = np.sum(cov, axis=2, keepdims=True) / (obsi.vis.shape[2]**2)
+        new_weight = 1 / new_cov
+        new_freq = np.array([np.mean(obsi.freq)])
+        new_obs = rve.Observation(
+            obsi.antenna_positions,
+            new_vis,
+            new_weight,
+            obsi.polarization,
+            new_freq,
+            obs._auxiliary_tables)
+        obs_avg.append(new_obs)
+
+    new_freq = [obs.freq[0] for obs in obs_avg]
+    new_freq = np.array(new_freq)
+    new_vis_shape = (obs.vis.shape[0], obs.vis.shape[1], len(new_freq))
+    new_vis = np.zeros(new_vis_shape, obs.vis.dtype)
+    new_weight = np.zeros(new_vis_shape, obs.weight.dtype)
+    for ii, obs in enumerate(obs_avg):
+        new_vis[:, :, ii] = obs.vis.val[:, :, 0]
+        new_weight[:, :, ii] = obs.weight.val[:, :, 0]
+
+    obs_averaged = rve.Observation(
+        obs.antenna_positions,
+        new_vis,
+        new_weight,
+        obs.polarization,
+        new_freq,
+        obs._auxiliary_tables)
+    return obs_averaged
+
+
+def time_average(obs, len_tbin):
+    tmin, tmax = rve.tmin_tmax(obs)
+    assert tmin == 0
+    n_tbins = int(tmax // len_tbin + 2)
+    tbins_endpoints = np.arange(0, n_tbins*len_tbin, len_tbin)
+    unique_times = np.unique(obs.time)
+    t_intervals = []
+    for ii in range(n_tbins-1):
+        start = tbins_endpoints[ii]
+        stop = tbins_endpoints[ii+1]
+        s = start <= unique_times
+        b = stop > unique_times
+        vis_in_inter = np.any(np.logical_and(s, b))
+        if vis_in_inter:
+            t_intervals.append([start, stop])
+
+    return obs.time_average(t_intervals)
+
+
+# obs = laod_obs
+
+n_freq_chuncks = 4
+obs = freq_average(obs, n_freq_chuncks)
+
+tmin, tmax = rve.tmin_tmax(obs)
+obs = obs.move_time(-tmin)
+obs = time_average(obs, 30)