diff --git a/resolve/time_freq_avg.py b/resolve/time_freq_avg.py new file mode 100644 index 0000000000000000000000000000000000000000..2306628ced480c8be2d975baed302adf8b520e5b --- /dev/null +++ b/resolve/time_freq_avg.py @@ -0,0 +1,72 @@ +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)