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
