def test_create_dataset(self): global_data = np.arange(size*5*10, dtype=np.float32) local_data = global_data.reshape(size, -1, 10)[rank] d_array = mpiarray.MPIArray.wrap(local_data, axis=0) d_array_T = d_array.redistribute(axis=1) # Check that we must specify in advance if the dataset is distributed g = memh5.MemGroup() if comm is not None: self.assertRaises(RuntimeError, g.create_dataset, 'data', data=d_array) g = memh5.MemGroup(distributed=True) # Create an array from data g.create_dataset('data', data=d_array, distributed=True) # Create an array from data with a different distribution g.create_dataset('data_T', data=d_array, distributed=True, distributed_axis=1) # Create an empty array with a specified shape g.create_dataset('data2', shape=(size*5, 10), dtype=np.float64, distributed=True, distributed_axis=1) self.assertTrue(np.allclose(d_array, g['data'][:])) self.assertTrue(np.allclose(d_array_T, g['data_T'][:])) if comm is not None: self.assertEqual(d_array_T.local_shape, g['data2'].local_shape) # Test global indexing self.assertTrue((g['data'][rank*5] == local_data[0]).all())
def test_nested(self): root = memh5.MemGroup() l1 = root.create_group("level1") l2 = l1.require_group("level2") self.assertTrue(root["level1"] == l1) self.assertTrue(root["level1/level2"] == l2) self.assertEqual(root["level1/level2"].name, "/level1/level2")
def test_nested(self): root = memh5.MemGroup() l1 = root.create_group('level1') l2 = l1.require_group('level2') self.assertTrue(root['level1'] == l1) self.assertTrue(root['level1/level2'] == l2) self.assertEqual(root['level1/level2'].name, '/level1/level2')
def test_failure(self): """Test that we get a TypeError if we try to serialize something else""" m = memh5.MemGroup() m.attrs["non_serializable"] = {"datetime": self} with self.assertRaises(TypeError): m.to_hdf5(self.fname)
def test_to_from_hdf5(self): udata = np.array(["Test", "this", "works"]) sdata = udata.astype("S") self.assertEqual(udata.dtype.kind, "U") self.assertEqual(sdata.dtype.kind, "S") m = memh5.MemGroup() udset = m.create_dataset("udata", data=udata) sdset = m.create_dataset("sdata", data=sdata) self.assertEqual(udset.dtype.kind, "U") self.assertEqual(sdset.dtype.kind, "S") m.to_hdf5(self.fname) with h5py.File(self.fname, "r") as fh: self.assertEqual(fh["udata"].dtype.kind, "S") self.assertEqual(fh["sdata"].dtype.kind, "S") self.assertIn("__memh5_unicode", fh["udata"].attrs) self.assertNotIn("__memh5_unicode", fh["sdata"].attrs) self.assertTrue(fh["udata"].attrs["__memh5_unicode"]) m2 = memh5.MemGroup.from_hdf5(self.fname) self.assertEqual(m2["udata"].dtype.kind, "U") self.assertEqual(m2["sdata"].dtype.kind, "S") self.assertTrue((m["udata"].data == m2["udata"].data).all()) self.assertTrue((m["sdata"].data == m2["sdata"].data).all())
def test_recursive_create_dataset(self): g = memh5.MemGroup() data = np.arange(10) g.create_dataset('a/ra', data=data) self.assertTrue(memh5.is_group(g['a'])) self.assertTrue(np.all(g['a/ra'][:] == data)) g['a'].create_dataset('/ra', data=data) print g.keys() self.assertTrue(np.all(g['ra'][:] == data))
def test_failure(self): # Test that we fail when trying to write a non ASCII character udata = np.array(["\u03B2"]) m = memh5.MemGroup() m.create_dataset("udata", data=udata) with self.assertRaises(TypeError): m.to_hdf5(self.fname)
def test_recursive_create(self): g = memh5.MemGroup() self.assertRaises(ValueError, g.create_group, "") g2 = g.create_group("level2/") self.assertRaises(ValueError, g2.create_group, "/") g2.create_group("/level22") self.assertEqual(set(g.keys()), {"level22", "level2"}) g.create_group("/a/b/c/d/") gd = g["/a/b/c/d/"] self.assertEqual(gd.name, "/a/b/c/d")
def test_recursive_create(self): g = memh5.MemGroup() self.assertRaises(ValueError, g.create_group, '') g2 = g.create_group('level2/') self.assertRaises(ValueError, g2.create_group, '/') g2.create_group('/level22') self.assertEqual(set(g.keys()), {'level22', 'level2'}) g.create_group('/a/b/c/d/') gd = g['/a/b/c/d/'] self.assertEqual(gd.name, '/a/b/c/d')
def test_io(self): # Create distributed memh5 object g = memh5.MemGroup(distributed=True) g.attrs['rank'] = rank # Create an empty array with a specified shape pdset = g.create_dataset('parallel_data', shape=(size, 10), dtype=np.float64, distributed=True, distributed_axis=0) pdset[:] = rank pdset.attrs['rank'] = rank # Create an empty array with a specified shape sdset = g.create_dataset('serial_data', shape=(size*5, 10), dtype=np.float64) sdset[:] = rank sdset.attrs['rank'] = rank # Create nested groups g.create_group('hello/world') g.to_hdf5(self.fname) # Test that the HDF5 file has the correct structure with h5py.File(self.fname, 'r') as f: # Test that the file attributes are correct self.assertTrue(f['parallel_data'].attrs['rank'] == 0) # Test that the parallel dataset has been written correctly self.assertTrue((f['parallel_data'][:, 0] == np.arange(size)).all()) self.assertTrue(f['parallel_data'].attrs['rank'] == 0) # Test that the common dataset has been written correctly (i.e. by rank=0) self.assertTrue((f['serial_data'][:] == 0).all()) self.assertTrue(f['serial_data'].attrs['rank'] == 0) # Check group structure is correct self.assertIn('hello', f) self.assertIn('world', f['hello']) # Test that the read in group has the same structure as the original g2 = memh5.MemGroup.from_hdf5(self.fname, distributed=True) # Check that the parallel data is still the same self.assertTrue((g2['parallel_data'][:] == g['parallel_data'][:]).all()) # Check that the serial data is all zeros (should not be the same as before) self.assertTrue((g2['serial_data'][:] == np.zeros_like(sdset[:])).all()) # Check group structure is correct self.assertIn('hello', g2) self.assertIn('world', g2['hello']) # Check the attributes self.assertTrue(g2['parallel_data'].attrs['rank'] == 0) self.assertTrue(g2['serial_data'].attrs['rank'] == 0)
def test_recursive_create_dataset(self): g = memh5.MemGroup() data = np.arange(10) g.create_dataset("a/ra", data=data) self.assertTrue(memh5.is_group(g["a"])) self.assertTrue(np.all(g["a/ra"][:] == data)) g["a"].create_dataset("/ra", data=data) self.assertTrue(np.all(g["ra"][:] == data)) self.assertIsInstance(g["a/ra"].parent, memh5.MemGroup) # Check that d keeps g in scope. d = g["a/ra"] del g gc.collect() self.assertTrue(np.all(d.file["ra"][:] == data))
def test_to_from_hdf5(self): json_prefix = "!!_memh5_json:" data = {"foo": {"bar": [1, 2, 3], "fu": "1"}} time = datetime.datetime.now() m = memh5.MemGroup() m.attrs["data"] = data m.attrs["datetime"] = {"datetime": time} m.to_hdf5(self.fname) with h5py.File(self.fname, "r") as f: assert f.attrs["data"] == json_prefix + json.dumps(data) assert f.attrs["datetime"] == json_prefix + json.dumps( {"datetime": time.isoformat()}) m2 = memh5.MemGroup.from_hdf5(self.fname) assert m2.attrs["data"] == data assert m2.attrs["datetime"] == {"datetime": time.isoformat()}
def test_to_from_hdf5(self): udata = np.array(["Test", "this", "works"]) sdata = udata.astype("S") self.assertEqual(udata.dtype.kind, "U") self.assertEqual(sdata.dtype.kind, "S") m = memh5.MemGroup() udset = m.create_dataset("udata", data=udata) sdset = m.create_dataset("sdata", data=sdata) self.assertEqual(udset.dtype.kind, "U") self.assertEqual(sdset.dtype.kind, "S") # Test a write without conversion. This should throw an exception with self.assertRaises(TypeError): m.to_hdf5(self.fname) # Write with conversion m.to_hdf5(self.fname, convert_attribute_strings=True, convert_dataset_strings=True) with h5py.File(self.fname, "r") as fh: self.assertEqual(fh["udata"].dtype.kind, "S") self.assertEqual(fh["sdata"].dtype.kind, "S") # Test a load without conversion, types should be bytestrings m2 = memh5.MemGroup.from_hdf5(self.fname) self.assertEqual(m2["udata"].dtype.kind, "S") self.assertEqual(m2["sdata"].dtype.kind, "S") # Check the dtype here, for some reason Python 2 thinks the arrays are equal # and Python 3 does not even though both agree that the datatypes are different self.assertTrue(m["udata"].dtype != m2["udata"].dtype) self.assertTrue((m["sdata"].data == m2["sdata"].data).all()) # Test a load *with* conversion, types should be unicode m3 = memh5.MemGroup.from_hdf5(self.fname, convert_attribute_strings=True, convert_dataset_strings=True) self.assertEqual(m3["udata"].dtype.kind, "U") self.assertEqual(m3["sdata"].dtype.kind, "U") self.assertTrue((m["udata"].data == m3["udata"].data).all()) self.assertTrue((m["udata"].data == m3["sdata"].data).all())
def process(self, ts): mask_daytime = self.params['mask_daytime'] mask_time_range = self.params['mask_time_range'] tsys = self.params['tsys'] accuracy_boost = self.params['accuracy_boost'] l_boost = self.params['l_boost'] bl_range = self.params['bl_range'] auto_correlations = self.params['auto_correlations'] time_avg = self.params['time_avg'] pol = self.params['pol'] interp = self.params['interp'] beam_dir = output_path(self.params['beam_dir']) use_existed_beam = self.params['use_existed_beam'] gen_inv = self.params['gen_invbeam'] noise_weight = self.params['noise_weight'] ts_dir = output_path(self.params['ts_dir']) ts_name = self.params['ts_name'] no_m_zero = self.params['no_m_zero'] simulate = self.params['simulate'] input_maps = self.params['input_maps'] prior_map = self.params['prior_map'] add_noise = self.params['add_noise'] dirty_map = self.params['dirty_map'] nbin = self.params['nbin'] method = self.params['method'] normalize = self.params['normalize'] threshold = self.params['threshold'] eps = self.params['epsilon'] correct_order = self.params['correct_order'] if use_existed_beam: # load the saved telescope from disk tel = None else: assert isinstance(ts, Timestream), '%s only works for Timestream object' % self.__class__.__name__ ts.redistribute('baseline') lat = ts.attrs['sitelat'] # lon = ts.attrs['sitelon'] lon = 0.0 # lon = np.degrees(ts['ra_dec'][0, 0]) # the first ra local_origin = False freqs = ts.freq[:] # MHz nfreq = freqs.shape[0] band_width = ts.attrs['freqstep'] # MHz try: ndays = ts.attrs['ndays'] except KeyError: ndays = 1 feeds = ts['feedno'][:] bl_order = mpiutil.gather_array(ts.local_bl, axis=0, root=None, comm=ts.comm) bls = [ tuple(bl) for bl in bl_order ] az, alt = ts['az_alt'][0] az = np.degrees(az) alt = np.degrees(alt) pointing = [az, alt, 0.0] feedpos = ts['feedpos'][:] if ts.is_dish: from tlpipe.map.drift.telescope import tl_dish dish_width = ts.attrs['dishdiam'] tel = tl_dish.TlUnpolarisedDishArray(lat, lon, freqs, band_width, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, local_origin, dish_width, feedpos, pointing) elif ts.is_cylinder: from tlpipe.map.drift.telescope import tl_cylinder # factor = 1.2 # suppose an illumination efficiency, keep same with that in timestream_common factor = 0.79 # for xx # factor = 0.88 # for yy cyl_width = factor * ts.attrs['cywid'] tel = tl_cylinder.TlUnpolarisedCylinder(lat, lon, freqs, band_width, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, local_origin, cyl_width, feedpos) else: raise RuntimeError('Unknown array type %s' % ts.attrs['telescope']) if not simulate: # select the corresponding vis and vis_mask if pol == 'xx': local_vis = ts.local_vis[:, :, 0, :] local_vis_mask = ts.local_vis_mask[:, :, 0, :] elif pol == 'yy': local_vis = ts.local_vis[:, :, 1, :] local_vis_mask = ts.local_vis_mask[:, :, 1, :] elif pol == 'I': xx_vis = ts.local_vis[:, :, 0, :] xx_vis_mask = ts.local_vis_mask[:, :, 0, :] yy_vis = ts.local_vis[:, :, 1, :] yy_vis_mask = ts.local_vis_mask[:, :, 1, :] local_vis = np.zeros_like(xx_vis) for ti in xrange(local_vis.shape[0]): for fi in xrange(local_vis.shape[1]): for bi in xrange(local_vis.shape[2]): if xx_vis_mask[ti, fi, bi] != yy_vis_mask[ti, fi, bi]: if xx_vis_mask[ti, fi, bi]: local_vis[ti, fi, bi] = yy_vis[ti, fi, bi] else: local_vis[ti, fi, bi] = xx_vis[ti, fi, bi] else: local_vis[ti, fi, bi] = 0.5 * (xx_vis[ti, fi, bi] + yy_vis[ti, fi, bi]) local_vis_mask = xx_vis_mask | yy_vis_mask else: raise ValueError('Invalid pol: %s' % pol) if interp != 'none': for fi in xrange(local_vis.shape[1]): for bi in xrange(local_vis.shape[2]): # interpolate for local_vis true_inds = np.where(local_vis_mask[:, fi, bi])[0] # masked inds if len(true_inds) > 0: false_inds = np.where(~local_vis_mask[:, fi, bi])[0] # un-masked inds if len(false_inds) > 0.1 * local_vis.shape[0]: # nearest interpolate for local_vis if interp in ('linear', 'nearest'): itp_real = interp1d(false_inds, local_vis[false_inds, fi, bi].real, kind=interp, fill_value='extrapolate', assume_sorted=True) itp_imag = interp1d(false_inds, local_vis[false_inds, fi, bi].imag, kind=interp, fill_value='extrapolate', assume_sorted=True) elif interp == 'rbf': itp_real = Rbf(false_inds, local_vis[false_inds, fi, bi].real, smooth=10) itp_imag = Rbf(false_inds, local_vis[false_inds, fi, bi].imag, smooth=10) else: raise ValueError('Unknown interpolation method: %s' % interp) local_vis[true_inds, fi, bi] = itp_real(true_inds) + 1.0J * itp_imag(true_inds) # the interpolated vis else: local_vis[:, fi, bi] = 0 # TODO: may need to take special care # average data nt = ts['sec1970'].shape[0] phi_size = 2*tel.mmax + 1 # phi = np.zeros((phi_size,), dtype=ts['ra_dec'].dtype) phi = np.linspace(0, 2*np.pi, phi_size, endpoint=False) vis = np.zeros((phi_size,)+local_vis.shape[1:], dtype=local_vis.dtype) if time_avg == 'avg': nt_m = float(nt) / phi_size # roll data to have phi=0 near the first roll_len = np.int(np.around(0.5*nt_m)) local_vis[:] = np.roll(local_vis[:], roll_len, axis=0) if interp == 'none': local_vis_mask[:] = np.roll(local_vis_mask[:], roll_len, axis=0) # ts['ra_dec'][:] = np.roll(ts['ra_dec'][:], roll_len, axis=0) repeat_inds = np.repeat(np.arange(nt), phi_size) num, start, end = mpiutil.split_m(nt*phi_size, phi_size) # average over time for idx in xrange(phi_size): inds, weight = unique(repeat_inds[start[idx]:end[idx]], return_counts=True) if interp == 'none': vis[idx] = average(np.ma.array(local_vis[inds], mask=local_vis_mask[inds]), axis=0, weights=weight) # time mean else: vis[idx] = average(local_vis[inds], axis=0, weights=weight) # time mean # phi[idx] = np.average(ts['ra_dec'][:, 0][inds], axis=0, weights=weight) elif time_avg == 'fft': if interp == 'none': raise ValueError('Can not do fft average without first interpolation') Vm = np.fft.fftshift(np.fft.fft(local_vis, axis=0), axes=0) vis[:] = np.fft.ifft(np.fft.ifftshift(Vm[nt/2-tel.mmax:nt/2+tel.mmax+1], axes=0), axis=0) / (1.0 * nt / phi_size) # for fi in xrange(vis.shape[1]): # for bi in xrange(vis.shape[2]): # # plot local_vis and vis # import matplotlib # matplotlib.use('Agg') # import matplotlib.pyplot as plt # phi0 = np.linspace(0, 2*np.pi, nt, endpoint=False) # phi1 = np.linspace(0, 2*np.pi, phi_size, endpoint=False) # plt.figure() # plt.subplot(211) # plt.plot(phi0, local_vis[:, fi, bi].real, label='v0.real') # plt.plot(phi1, vis[:, fi, bi].real, label='v1.real') # plt.legend() # plt.subplot(212) # plt.plot(phi0, local_vis[:, fi, bi].imag, label='v0.imag') # plt.plot(phi1, vis[:, fi, bi].imag, label='v1.imag') # plt.legend() # plt.savefig('vis_fft/vis_%d_%d.png' % (fi, bi)) # plt.close() else: raise ValueError('Unknown time_avg: %s' % time_avg) del local_vis del local_vis_mask # mask daytime data if mask_daytime: day_or_night = np.where(ts['local_hour'][:]>=mask_time_range[0] & ts['local_hour'][:]<=mask_time_range[1], True, False) day_inds = np.where(np.repeat(day_or_night, phi_size).reshape(nt, phi_size).astype(np.int).sum(axis=1).astype(bool))[0] vis[day_inds] = 0 del ts # no longer need ts # redistribute vis to time axis vis = mpiarray.MPIArray.wrap(vis, axis=2).redistribute(0).local_array allpairs = tel.allpairs redundancy = tel.redundancy nrd = len(redundancy) # reorder bls according to allpairs vis_tmp = np.zeros_like(vis) for ind, (a1, a2) in enumerate(allpairs): try: b_ind = bls.index((feeds[a1], feeds[a2])) vis_tmp[:, :, ind] = vis[:, :, b_ind] except ValueError: b_ind = bls.index((feeds[a2], feeds[a1])) vis_tmp[:, :, ind] = vis[:, :, b_ind].conj() del vis # average over redundancy vis_stream = np.zeros(vis_tmp.shape[:-1]+(nrd,), dtype=vis_tmp.dtype) red_bin = np.cumsum(np.insert(redundancy, 0, 0)) # redundancy bin # average over redundancy for ind in xrange(nrd): vis_stream[:, :, ind] = np.sum(vis_tmp[:, :, red_bin[ind]:red_bin[ind+1]], axis=2) / redundancy[ind] del vis_tmp # beamtransfer bt = beamtransfer.BeamTransfer(beam_dir, tel, noise_weight, True) if not use_existed_beam: bt.generate() if tel is None: tel = bt.telescope if simulate: ndays = 733 tstream = timestream.simulate(bt, ts_dir, ts_name, input_maps, ndays, add_noise=add_noise) else: # timestream and map-making tstream = timestream.Timestream(ts_dir, ts_name, bt, no_m_zero) parent_path = os.path.dirname(tstream._fdir(0)) if os.path.exists(parent_path + '/COMPLETED'): if mpiutil.rank0: print 'Use existed timestream_f files in %s' % parent_path else: for fi in mpiutil.mpirange(nfreq): # Make directory if required if not os.path.exists(tstream._fdir(fi)): os.makedirs(tstream._fdir(fi)) # create memh5 object and write data to temporary file vis_h5 = memh5.MemGroup(distributed=True) vis_h5.create_dataset('/timestream', data=mpiarray.MPIArray.wrap(vis_stream, axis=0)) tmp_file = parent_path +'/vis_stream_temp.hdf5' vis_h5.to_hdf5(tmp_file, hints=False) del vis_h5 # re-organize data as need for tstream # make load even among nodes for fi in mpiutil.mpirange(nfreq, method='rand'): # read the needed data from the temporary file with h5py.File(tmp_file, 'r') as f: vis_fi = f['/timestream'][:, fi, :] # Write file contents with h5py.File(tstream._ffile(fi), 'w') as f: # Timestream data # allocate space for vis_stream shp = (nrd, phi_size) f.create_dataset('/timestream', data=vis_fi.T) f.create_dataset('/phi', data=phi) # Telescope layout data f.create_dataset('/feedmap', data=tel.feedmap) f.create_dataset('/feedconj', data=tel.feedconj) f.create_dataset('/feedmask', data=tel.feedmask) f.create_dataset('/uniquepairs', data=tel.uniquepairs) f.create_dataset('/baselines', data=tel.baselines) # Telescope frequencies f.create_dataset('/frequencies', data=freqs) # Write metadata f.attrs['beamtransfer_path'] = os.path.abspath(bt.directory) f.attrs['ntime'] = phi_size mpiutil.barrier() # remove temp file if mpiutil.rank0: os.remove(tmp_file) # mark all frequencies tstream files are saved correctly open(parent_path + '/COMPLETED', 'a').close() tstream.generate_mmodes() nside = hputil.nside_for_lmax(tel.lmax, accuracy_boost=tel.accuracy_boost) if dirty_map: tstream.mapmake_full(nside, 'map_full_dirty.hdf5', nbin, dirty=True, method=method, normalize=normalize, threshold=threshold) else: tstream.mapmake_full(nside, 'map_full.hdf5', nbin, dirty=False, method=method, normalize=normalize, threshold=threshold, eps=eps, correct_order=correct_order, prior_map_file=prior_map) # ts.add_history(self.history) return tstream
def test_io(self): # Create distributed memh5 object g = memh5.MemGroup(distributed=True) g.attrs["rank"] = rank # Create an empty array with a specified shape pdset = g.create_dataset( "parallel_data", shape=(size, 10), dtype=np.float64, distributed=True, distributed_axis=0, ) pdset[:] = rank pdset.attrs["const"] = 17 # Create an empty array with a specified shape sdset = g.create_dataset("serial_data", shape=(size * 5, 10), dtype=np.float64) sdset[:] = rank sdset.attrs["const"] = 18 # Create nested groups g.create_group("hello/world") # Test round tripping unicode data g.create_dataset("unicode_data", data=np.array(["hello"])) g.to_hdf5(self.fname, convert_attribute_strings=True, convert_dataset_strings=True) # Test that the HDF5 file has the correct structure with h5py.File(self.fname, "r") as f: # Test that the file attributes are correct self.assertTrue(f["parallel_data"].attrs["const"] == 17) # Test that the parallel dataset has been written correctly self.assertTrue((f["parallel_data"][:, 0] == np.arange(size)).all()) self.assertTrue(f["parallel_data"].attrs["const"] == 17) # Test that the common dataset has been written correctly (i.e. by rank=0) self.assertTrue((f["serial_data"][:] == 0).all()) self.assertTrue(f["serial_data"].attrs["const"] == 18) # Check group structure is correct self.assertIn("hello", f) self.assertIn("world", f["hello"]) # Test that the read in group has the same structure as the original g2 = memh5.MemGroup.from_hdf5( self.fname, distributed=True, convert_attribute_strings=True, convert_dataset_strings=True, ) # Check that the parallel data is still the same self.assertTrue( (g2["parallel_data"][:] == g["parallel_data"][:]).all()) # Check that the serial data is all zeros (should not be the same as before) self.assertTrue( (g2["serial_data"][:] == np.zeros_like(sdset[:])).all()) # Check group structure is correct self.assertIn("hello", g2) self.assertIn("world", g2["hello"]) # Check the unicode dataset self.assertEqual(g2["unicode_data"].dtype.kind, "U") self.assertEqual(g2["unicode_data"][0], "hello") # Check the attributes self.assertTrue(g2["parallel_data"].attrs["const"] == 17) self.assertTrue(g2["serial_data"].attrs["const"] == 18)
def process(self, ts): mask_daytime = self.params['mask_daytime'] mask_time_range = self.params['mask_time_range'] beam_theta_range = self.params['beam_theta_range'] tsys = self.params['tsys'] accuracy_boost = self.params['accuracy_boost'] l_boost = self.params['l_boost'] bl_range = self.params['bl_range'] auto_correlations = self.params['auto_correlations'] pol = self.params['pol'] beam_dir = output_path(self.params['beam_dir']) gen_inv = self.params['gen_invbeam'] noise_weight = self.params['noise_weight'] ts_dir = output_path(self.params['ts_dir']) ts_name = self.params['ts_name'] simulate = self.params['simulate'] input_maps = self.params['input_maps'] add_noise = self.params['add_noise'] ts.redistribute('frequency') lat = ts.attrs['sitelat'] # lon = ts.attrs['sitelon'] lon = 0.0 # lon = np.degrees(ts['ra_dec'][0, 0]) # the first ra freqs = ts.freq.data.to_numpy_array(root=None) band_width = ts.attrs['freqstep'] # MHz ndays = ts.attrs['ndays'] feeds = ts['feedno'][:] az, alt = ts['az_alt'][0] az = np.degrees(az) alt = np.degrees(alt) pointing = [az, alt, 0.0] feedpos = ts['feedpos'][:] if ts.is_dish: dish_width = ts.attrs['dishdiam'] tel = tl_dish.TlUnpolarisedDishArray(lat, lon, freqs, beam_theta_range, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, dish_width, feedpos, pointing) elif ts.is_cylinder: # factor = 1.2 # suppose an illumination efficiency, keep same with that in timestream_common factor = 0.79 # for xx # factor = 0.88 # for yy cyl_width = factor * ts.attrs['cywid'] tel = tl_cylinder.TlUnpolarisedCylinder( lat, lon, freqs, beam_theta_range, tsys, ndays, accuracy_boost, l_boost, bl_range, auto_correlations, cyl_width, feedpos) else: raise RuntimeError('Unknown array type %s' % ts.attrs['telescope']) # import matplotlib # matplotlib.use('Agg') # import matplotlib.pyplot as plt # plt.figure() # plt.plot(ts['ra_dec'][:]) # # plt.plot(ts['az_alt'][:]) # plt.savefig('ra_dec1.png') if not simulate: # mask daytime data if mask_daytime: day_inds = np.where( np.logical_and( ts['local_hour'][:] >= mask_time_range[0], ts['local_hour'][:] <= mask_time_range[1]))[0] ts.local_vis_mask[ day_inds] = True # do not change vis directly # average data nt = ts['sec1970'].shape[0] phi_size = tel.phi_size nt_m = float(nt) / phi_size # roll data to have phi=0 near the first roll_len = np.int(np.around(0.5 * nt_m)) ts.local_vis[:] = np.roll(ts.local_vis[:], roll_len, axis=0) ts.local_vis_mask[:] = np.roll(ts.local_vis_mask[:], roll_len, axis=0) ts['ra_dec'][:] = np.roll(ts['ra_dec'][:], roll_len, axis=0) repeat_inds = np.repeat(np.arange(nt), phi_size) num, start, end = mpiutil.split_m(nt * phi_size, phi_size) # phi = np.zeros((phi_size,), dtype=ts['ra_dec'].dtype) phi = np.linspace(0, 2 * np.pi, phi_size, endpoint=False) vis = np.zeros((phi_size, ) + ts.local_vis.shape[1:], dtype=ts.vis.dtype) # average over time for idx in xrange(phi_size): inds, weight = unique(repeat_inds[start[idx]:end[idx]], return_counts=True) vis[idx] = average(np.ma.array(ts.local_vis[inds], mask=ts.local_vis_mask[inds]), axis=0, weights=weight) # time mean # phi[idx] = np.average(ts['ra_dec'][:, 0][inds], axis=0, weights=weight) if pol == 'xx': vis = vis[:, :, 0, :] elif pol == 'yy': vis = vis[:, :, 1, :] elif pol == 'I': vis = 0.5 * (vis[:, :, 0, :] + vis[:, :, 1, :]) elif pol == 'all': vis = np.sum(vis, axis=2) # sum over all pol else: raise ValueError('Invalid pol: %s' % pol) allpairs = tel.allpairs redundancy = tel.redundancy # reorder bls according to allpairs vis_tmp = np.zeros_like(vis) bls = [tuple(bl) for bl in ts['blorder'][:]] for ind, (a1, a2) in enumerate(allpairs): try: b_ind = bls.index((feeds[a1], feeds[a2])) vis_tmp[:, :, ind] = vis[:, :, b_ind] except ValueError: b_ind = bls.index((feeds[a2], feeds[a1])) vis_tmp[:, :, ind] = vis[:, :, b_ind].conj() # average over redundancy vis_stream = np.zeros(vis.shape[:-1] + (len(redundancy), ), dtype=vis_tmp.dtype) red_bin = np.cumsum(np.insert(redundancy, 0, 0)) # redundancy bin # average over redundancy for ind in xrange(len(redundancy)): vis_stream[:, :, ind] = np.sum( vis_tmp[:, :, red_bin[ind]:red_bin[ind + 1]], axis=2) / redundancy[ind] del vis del vis_tmp vis_stream = mpiarray.MPIArray.wrap(vis_stream, axis=1) vis_h5 = memh5.MemGroup(distributed=True) vis_h5.create_dataset('/timestream', data=vis_stream) vis_h5.create_dataset('/phi', data=phi) # Telescope layout data vis_h5.create_dataset('/feedmap', data=tel.feedmap) vis_h5.create_dataset('/feedconj', data=tel.feedconj) vis_h5.create_dataset('/feedmask', data=tel.feedmask) vis_h5.create_dataset('/uniquepairs', data=tel.uniquepairs) vis_h5.create_dataset('/baselines', data=tel.baselines) # Telescope frequencies vis_h5.create_dataset('/frequencies', data=freqs) # Write metadata # vis_h5.attrs['beamtransfer_path'] = os.path.abspath(bt.directory) vis_h5.attrs['ntime'] = phi_size # beamtransfer bt = beamtransfer.BeamTransfer(beam_dir, tel, gen_inv, noise_weight) bt.generate() if simulate: ndays = 733 print ndays ts = timestream.simulate(bt, ts_dir, ts_name, input_maps, ndays, add_noise=add_noise) else: # timestream and map-making ts = timestream.Timestream(ts_dir, ts_name, bt) # Make directory if required try: os.makedirs(ts._tsdir) except OSError: # directory exists pass vis_h5.to_hdf5(ts._tsfile) # ts.generate_mmodes(vis_stream.to_numpy_array(root=None)) ts.generate_mmodes() nside = hputil.nside_for_lmax(tel.lmax, accuracy_boost=tel.accuracy_boost) ts.mapmake_full(nside, 'full') # ts.add_history(self.history) return ts
def process(self, ts): calibrator = self.params['calibrator'] catalog = self.params['catalog'] span = self.params['span'] save_gain = self.params['save_gain'] gain_file = self.params['gain_file'] ts.redistribute('frequency') lfreq = ts.local_freq[:] # local freq feedno = ts['feedno'][:].tolist() pol = ts['pol'][:].tolist() bl = ts.bl[:] bls = [tuple(b) for b in bl] # # antpointing = np.radians(ts['antpointing'][-1, :, :]) # radians # transitsource = ts['transitsource'][:] # transit_time = transitsource[-1, 0] # second, sec1970 # int_time = ts.attrs['inttime'] # second # calibrator srclist, cutoff, catalogs = a.scripting.parse_srcs(calibrator, catalog) cat = a.src.get_catalog(srclist, cutoff, catalogs) assert (len(cat) == 1), 'Allow only one calibrator' s = cat.values()[0] if mpiutil.rank0: print 'Calibrating for source %s with' % calibrator, print 'strength', s._jys, 'Jy', print 'measured at', s.mfreq, 'GHz', print 'with index', s.index # get transit time of calibrator # array aa = ts.array aa.set_jultime(ts['jul_date'][0]) # the first obs time point next_transit = aa.next_transit(s) transit_time = a.phs.ephem2juldate(next_transit) # Julian date if transit_time > ts['jul_date'][-1]: local_next_transit = ephem.Date(next_transit + 8.0 * ephem.hour) raise RuntimeError( 'Data does not contain local transit time %s of source %s' % (local_next_transit, calibrator)) # the first transit index transit_inds = [np.searchsorted(ts['jul_date'][:], transit_time)] # find all other transit indices aa.set_jultime(ts['jul_date'][0] + 1.0) transit_time = a.phs.ephem2juldate(aa.next_transit(s)) # Julian date cnt = 2 while (transit_time <= ts['jul_date'][-1]): transit_inds.append( np.searchsorted(ts['jul_date'][:], transit_time)) aa.set_jultime(ts['jul_date'][0] + 1.0 * cnt) transit_time = a.phs.ephem2juldate( aa.next_transit(s)) # Julian date cnt += 1 print transit_inds ### now only use the first transit point to do the cal ### may need to improve in the future transit_ind = transit_inds[0] int_time = ts.attrs['inttime'] # second start_ind = transit_ind - np.int(span / int_time) end_ind = transit_ind + np.int(span / int_time) nt = end_ind - start_ind nfeed = len(feedno) eigval = np.empty((nt, nfeed, 2, len(lfreq)), dtype=np.float64) eigval[:] = np.nan gain = np.empty((nt, nfeed, 2, len(lfreq)), dtype=np.complex128) gain[:] = complex(np.nan, np.nan) # construct visibility matrix for a single time, pol, freq Vmat = np.zeros((nfeed, nfeed), dtype=ts.main_data.dtype) for ind, ti in enumerate(range(start_ind, end_ind)): # when noise on, just pass if 'ns_on' in ts.iterkeys() and ts['ns_on'][ti]: continue aa.set_jultime(ts['jul_date'][ti]) s.compute(aa) # get fluxes vs. freq of the calibrator Sc = s.get_jys() # get the topocentric coordinate of the calibrator at the current time s_top = s.get_crds('top', ncrd=3) aa.sim_cache(cat.get_crds( 'eq', ncrd=3)) # for compute bm_response and sim for pi in [pol.index('xx'), pol.index('yy')]: # xx, yy aa.set_active_pol(pol[pi]) for fi, freq in enumerate(lfreq): # mpi among freq for i, ai in enumerate(feedno): for j, aj in enumerate(feedno): # uij = aa.gen_uvw(i, j, src='z').squeeze() # (rj - ri)/lambda uij = aa.gen_uvw(i, j, src='z')[:, 0, :] # (rj - ri)/lambda # bmij = aa.bm_response(i, j).squeeze() # will get error for only one local freq # import pdb # pdb.set_trace() bmij = aa.bm_response(i, j).reshape(-1) try: bi = bls.index((ai, aj)) # Vmat[i, j] = ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(-2.0J * np.pi * np.dot(s_top, uij[:, fi]))) # xx, yy Vmat[i, j] = ts.local_vis[ti, fi, pi, bi] / ( Sc[fi] * bmij[fi] * np.exp(2.0J * np.pi * np.dot( s_top, uij[:, fi]))) # xx, yy except ValueError: bi = bls.index((aj, ai)) # Vmat[i, j] = np.conj(ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(-2.0J * np.pi * np.dot(s_top, uij[:, fi])))) # xx, yy Vmat[i, j] = np.conj( ts.local_vis[ti, fi, pi, bi] / (Sc[fi] * bmij[fi] * np.exp(2.0J * np.pi * np.dot( s_top, uij[:, fi])))) # xx, yy # Eigen decomposition Vmat = np.where(np.isfinite(Vmat), Vmat, 0) e, U = eigh(Vmat) eigval[ind, :, pi, fi] = e[::-1] # descending order # max eigen-val lbd = e[-1] # lambda # the gain vector for this freq gvec = np.sqrt( lbd ) * U[:, -1] # only eigen-vector corresponding to the maximum eigen-val gain[ind, :, pi, fi] = gvec # apply gain to vis # get the time mean gain tgain = np.ma.mean(np.ma.masked_invalid(gain), axis=0) # time mean tgain = mpiutil.gather_array(tgain, axis=-1, root=None) ts.redistribute('baseline') ts.pol_and_bl_data_operate(cal, tgain=tgain) # save gain if required: if save_gain: gain_file = output_path(gain_file) eigval = mpiarray.MPIArray.wrap(eigval, axis=3) gain = mpiarray.MPIArray.wrap(gain, axis=3) mem_gain = memh5.MemGroup(distributed=True) mem_gain.create_dataset('eigval', data=eigval) mem_gain.create_dataset('gain', data=gain) # add attris mem_gain.attrs['jul_data'] = ts['jul_date'][start_ind:end_ind] mem_gain.attrs['feed'] = np.array(feedno) mem_gain.attrs['pol'] = np.array(['xx', 'yy']) mem_gain.attrs['freq'] = ts.freq[:] # freq should be common # save to file mem_gain.to_hdf5(gain_file, hints=False) return super(PsCal, self).process(ts)
def test_create_dataset(self): g = memh5.MemGroup() data = np.arange(100, dtype=np.float32) g.create_dataset("data", data=data) self.assertTrue(np.allclose(data, g["data"]))
def simulate(beamtransfer, outdir, tsname, maps=[], ndays=None, resolution=0, add_noise=True, seed=None, **kwargs): """Create a simulated timestream and save it to disk. Parameters ---------- beamtransfer : fmmode.core.beamtransfer.BeamTransfer BeamTransfer object containing the analysis products. outdir : directoryname Directory that we will save the timestream into. maps : list List of map filenames. The sum of these form the simulated sky. ndays : int, optional Number of days of observation. Setting `ndays = None` (default) uses the default stored in the telescope object; `ndays = 0`, assumes the observation time is infinite so that the noise is zero. resolution : scalar, optional Approximate time resolution in seconds. Setting `resolution = 0` (default) calculates the value from the mmax. add_noise : bool, optional Weather to add random noise to the simulated visibilities. Default True. Returns ------- timestream : Timestream """ # Create timestream object tstream = Timestream(outdir, tsname, beamtransfer) completed_file = tstream._tsdir + '/COMPLETED_TIMESTREAM' if os.path.exists(completed_file): if mpiutil.rank0: print "******* timestream-files already generated ********" mpiutil.barrier() return tstream # Make directory if required try: os.makedirs(tstream._tsdir) except OSError: # directory exists pass if mpiutil.rank0: # if not os.path.exists(tstream._tsdir): # os.makedirs(tstream._tsdir) tstream.save() ## Read in telescope system bt = beamtransfer tel = bt.telescope lmax = tel.lmax mmax = tel.mmax nfreq = tel.nfreq nbl = tel.nbase npol = tel.num_pol_sky # If ndays is not set use the default value. if ndays is None: ndays = tel.ndays # Calculate the number of timesamples from the resolution if resolution == 0: # Set the minimum resolution required for the sky. ntime = 2 * mmax + 1 else: # Set the cl ntime = int(np.round(24 * 3600.0 / resolution)) indices = list(itertools.product(np.arange(nfreq), np.arange(npol))) lind, sind, eind = mpiutil.split_local(nfreq * npol) # local section of the Tm array theta_size = tel.theta_size phi_size = tel.phi_size Tm = np.zeros((lind, theta_size, phi_size), dtype=np.complex128) for ind, (f_ind, p_ind) in enumerate(indices[sind:eind]): hp_map = None for idx, mapfile in enumerate(maps): with h5py.File(mapfile, 'r') as f: if idx == 0: hp_map = f['map'][f_ind, p_ind, :] else: hp_map += f['map'][f_ind, p_ind, :] if hp_map is not None: cart_map = hpproj.cartesian_proj(hp_map, tel.cart_projector) # Calculate the Tm's for the local sections Tm[ind] = np.fft.ifft(cart_map, axis=1) # / phi_size # m = 0 is at left Tm = MPIArray.wrap(Tm, axis=0) # redistribute along different m Tm = Tm.redistribute(axis=2) Tm = Tm.reshape((nfreq, npol, theta_size, None)) Tm = Tm.reshape((nfreq, npol * theta_size, None)) ms = np.concatenate([np.arange(0, mmax + 1), np.arange(-mmax, 0)]) lm, sm, em = mpiutil.split_local(phi_size) # local section of mmode # mmode = np.zeros((lm, nbl, nfreq), dtype=np.complex128) mmode = np.zeros((lm, nfreq, nbl), dtype=np.complex128) for ind, mi in enumerate(ms[sm:em]): mmode[ind] = bt.project_vector_sky_to_telescope( mi, Tm[:, :, ind].view(np.ndarray)) mmode = MPIArray.wrap(mmode, axis=0) mmode = mmode.redistribute(axis=2) # distribute along bl # add noise if required if add_noise: lbl, sbl, ebl = mpiutil.split_local(nbl) # Fetch the noise powerspectrum noise_ps = tel.noisepower(np.arange(sbl, ebl)[:, np.newaxis], np.arange(nfreq)[np.newaxis, :], ndays=ndays).reshape( lbl, nfreq).T[np.newaxis, :, :] # Seed random number generator to give consistent noise if seed is not None: # Must include rank such that we don't have massive power deficit from correlated noise np.random.seed(seed + mpiutil.rank) # Create and weight complex noise coefficients noise_mode = (np.array([1.0, 1.0J]) * np.random.standard_normal(mmode.shape + (2, ))).sum(axis=-1) noise_mode *= (noise_ps / 2.0)**0.5 mmode += noise_mode del noise_mode # Reset RNG if seed is not None: np.random.seed() # The time samples the visibility is calculated at tphi = np.linspace(0, 2 * np.pi, ntime, endpoint=False) # inverse FFT to get timestream vis_stream = np.fft.ifft(mmode, axis=0) * ntime vis_stream = MPIArray.wrap(vis_stream, axis=2) # save vis_stream to file vis_h5 = memh5.MemGroup(distributed=True) vis_h5.create_dataset('/timestream', data=vis_stream) vis_h5.create_dataset('/phi', data=tphi) # Telescope layout data vis_h5.create_dataset('/feedmap', data=tel.feedmap) vis_h5.create_dataset('/feedconj', data=tel.feedconj) vis_h5.create_dataset('/feedmask', data=tel.feedmask) vis_h5.create_dataset('/uniquepairs', data=tel.uniquepairs) vis_h5.create_dataset('/baselines', data=tel.baselines) # Telescope frequencies vis_h5.create_dataset('/frequencies', data=tel.frequencies) # Write metadata vis_h5.attrs['beamtransfer_path'] = os.path.abspath(bt.directory) vis_h5.attrs['ntime'] = ntime # save to file vis_h5.to_hdf5(tstream._tsfile) if mpiutil.rank0: # Make file marker that the m's have been correctly generated: open(completed_file, 'a').close() mpiutil.barrier() return tstream