def project_sky(self, sky, mlist=None, threshold=None, harmonic=False):
# Set default list of m-modes (i.e. all of them), and partition
if mlist is None:
mlist = list(range(self.telescope.mmax + 1))
mpart = mpiutil.partition_list_mpi(mlist)
# Total number of sky modes.
nmodes = self.beamtransfer.nfreq * self.beamtransfer.ntel
# If sky is alm fine, if not perform spherical harmonic transform.
alm = sky if harmonic else hputil.sphtrans_sky(sky, lmax=self.telescope.lmax)
## Routine to project sky onto eigenmodes
def _proj(mi):
p1 = self.project_sky_vector_forward(mi, alm[:, :, mi], threshold)
p2 = np.zeros(nmodes, dtype=np.complex128)
p2[-p1.size :] = p1
return p2
# Map over list of m's and project sky onto eigenbasis
proj_sec = [(mi, _proj(mi)) for mi in mpart]
# Gather projections onto the rank=0 node.
proj_all = mpiutil.world.gather(proj_sec, root=0)
proj_arr = None
if mpiutil.rank0:
# Create array to put projections into
proj_arr = np.zeros(
(2 * self.telescope.mmax + 1, nmodes), dtype=np.complex128
)
# Iterate over all gathered projections and insert into the array
for proc_rank in proj_all:
for pm in proc_rank:
proj_arr[pm[0]] = pm[1]
# Return the projections (rank=0) or None elsewhere.
return proj_arr
def project_sky(self, sky, mlist = None, threshold=None, harmonic=False):
# Set default list of m-modes (i.e. all of them), and partition
if mlist is None:
mlist = range(self.telescope.mmax + 1)
mpart = mpiutil.partition_list_mpi(mlist)
# Total number of sky modes.
nmodes = self.beamtransfer.nfreq * self.beamtransfer.ntel
# If sky is alm fine, if not perform spherical harmonic transform.
alm = sky if harmonic else hputil.sphtrans_sky(sky, lmax=self.telescope.lmax)
## Routine to project sky onto eigenmodes
def _proj(mi):
p1 = self.project_sky_vector_forward(mi, alm[:, :, mi], threshold)
p2 = np.zeros(nmodes, dtype=np.complex128)
p2[-p1.size:] = p1
return p2
# Map over list of m's and project sky onto eigenbasis
proj_sec = [(mi, _proj(mi)) for mi in mpart]
# Gather projections onto the rank=0 node.
proj_all = mpiutil.world.gather(proj_sec, root=0)
proj_arr = None
if mpiutil.rank0:
# Create array to put projections into
proj_arr = np.zeros((2*self.telescope.mmax + 1, nmodes), dtype=np.complex128)
# Iterate over all gathered projections and insert into the array
for proc_rank in proj_all:
for pm in proc_rank:
proj_arr[pm[0]] = pm[1]
# Return the projections (rank=0) or None elsewhere.
return proj_arr
def generate_map(args):
import numpy as np
import h5py
import healpy
from cora.util import hputil
with h5py.File(args.in_map, 'r') as f:
in_map = f['map'][...]
nside = healpy.pixelfunc.get_nside(in_map[0])
in_alm = hputil.sphtrans_sky(in_map)
# lmax = in_alm.shape[-2] - 1
mmax = in_alm.shape[-1] - 1
print 'mmax = %d' % mmax
# l_cut = min(args.l_cut, lmax)
# m_cut = min(args.m_cut, mmax)
m_index = min(args.m_index, mmax)
temp_alm = np.zeros_like(in_alm, dtype=in_alm.dtype)
temp_alm[:, :, :, m_index] = in_alm[:, :, :, m_index]
out_map = hputil.sphtrans_inv_sky(temp_alm, nside)
out_file = args.out_file or ('m_slice_%d_' % m_index + args.in_map)
with h5py.File(out_file, 'w') as f:
f.create_dataset('map', data=out_map)
def plot_alm(args):
import os
import numpy as np
import h5py
from cora.util import hputil
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt
with h5py.File(args.in_map, 'r') as f:
in_map = f['map'][...]
alm = hputil.sphtrans_sky(in_map, lmax=args.maxl)
# plot alm
plt.figure(figsize=(args.figlength, args.figwidth))
plt.subplot(121)
plt.pcolor(alm[args.ifreq, args.pol].T.real, vmin=args.min, vmax=args.max)
lmin = args.lmin or 0
lmax = args.lmax or alm.shape[-2]
mmin = args.mmin or 0
mmax = args.mmax or alm.shape[-1]
plt.xlim(lmin, lmax)
plt.ylim(mmin, mmax)
plt.xlabel(r'$l$')
plt.ylabel(r'$m$')
plt.title(r'$\Re\left(a_{lm}\right)$')
plt.colorbar()
plt.subplot(122)
plt.pcolor(alm[args.ifreq, args.pol].T.imag, vmin=args.min, vmax=args.max)
plt.xlim(lmin, lmax)
plt.ylim(mmin, mmax)
plt.xlabel(r'$l$')
plt.ylabel(r'$m$')
plt.title(r'$\Im\left(a_{lm}\right)$')
plt.colorbar()
out_file = args.out_file or ('alm_of_' + os.path.basename(args.in_map).replace('.hdf5', '_%d_%s.%s' % (args.ifreq, ('{%d}' % args.pol).format('T', 'Q', 'U', 'V'), args.figfmt)))
plt.savefig(out_file)
def generate_map(args):
import os
import numpy as np
import h5py
import healpy
from cora.util import hputil
with h5py.File(args.in_map, 'r') as f:
in_map = f['map'][...]
nside = healpy.pixelfunc.get_nside(in_map[0])
in_alm = hputil.sphtrans_sky(in_map, lmax=args.maxl)
lmax = in_alm.shape[-2] - 1
mmax = in_alm.shape[-1] - 1
print 'lmax = %d\nmmax = %d' % (lmax, mmax)
lmin_cut = max(args.lmin, 0)
lmax_cut = min(args.lmax, lmax) if args.lmax is not None else lmax
mmin_cut = max(args.mmin, 0)
mmax_cut = min(args.mmax, mmax) if args.mmax is not None else mmax
temp_alm = np.zeros_like(in_alm, dtype=in_alm.dtype)
temp_alm[:, :, lmin_cut:(lmax_cut+1), mmin_cut:(mmax_cut+1)] = in_alm[:, :, lmin_cut:(lmax_cut+1), mmin_cut:(mmax_cut+1)]
out_map = hputil.sphtrans_inv_sky(temp_alm, nside)
out_file = args.out_file or ('lm_cut_%d_%d_%d_%d_' % (lmin_cut, lmax_cut, mmin_cut, mmax_cut) + os.path.basename(args.in_map))
with h5py.File(out_file, 'w') as f:
f.create_dataset('map', data=out_map)
with h5py.File(cm_name, 'r') as f:
cm_map = f['map'][:]
else:
map_dir = '../sky_map/'
ps_name = map_dir + 'sim_pointsource_%d_700_800_256.hdf5' % nside
ga_name = map_dir + 'sim_galaxy_%d_700_800_256.hdf5' % nside
cm_name = map_dir + 'sim_21cm_%d_700_800_256.hdf5' % nside
with h5py.File(ps_name, 'r') as f:
ps_map = f['map'][:, 0, :]
with h5py.File(ga_name, 'r') as f:
ga_map = f['map'][:, 0, :]
with h5py.File(cm_name, 'r') as f:
cm_map = f['map'][:, 0, :]
ps_alm = hputil.sphtrans_sky(ps_map)
ga_alm = hputil.sphtrans_sky(ga_map)
cm_alm = hputil.sphtrans_sky(cm_map)
ps_alm_name = 'alm_pointsource_%d_700_800_256.hdf5' % nside
ga_alm_name = 'alm_galaxy_%d_700_800_256.hdf5' % nside
cm_alm_name = 'alm_21cm_%d_700_800_256.hdf5' % nside
# save alms
with h5py.File(out_dir+ps_alm_name, 'w') as f:
ps = f.create_dataset('alm', data=ps_alm)
ps.attrs['axes'] = '(freq, l, m)'
with h5py.File(out_dir+ga_alm_name, 'w') as f:
ga = f.create_dataset('alm', data=ga_alm)
ga.attrs['axes'] = '(freq, l, m)'
with h5py.File(out_dir+cm_alm_name, 'w') as f:
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
----------
m : ProductManager object
Products of telescope to simulate.
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.
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
npol = tel.num_pol_sky
projmaps = (len(maps) > 0)
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
local_freq = range(sfreq, efreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# 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))
col_vis = np.zeros((tel.npairs, lfreq, ntime), dtype=np.complex128)
## If we want to add maps use the m-mode formalism to project a skymap
## into visibility space.
if projmaps:
# Load file to find out the map shapes.
with h5py.File(maps[0], 'r') as f:
mapshape = f['map'].shape
if lfreq > 0:
# Allocate array to store the local frequencies
row_map = np.zeros((lfreq, ) + mapshape[1:], dtype=np.float64)
# Read in and sum up the local frequencies of the supplied maps.
for mapfile in maps:
with h5py.File(mapfile, 'r') as f:
row_map += f['map'][sfreq:efreq]
# Calculate the alm's for the local sections
row_alm = hputil.sphtrans_sky(row_map, lmax=lmax).reshape(
(lfreq, npol * (lmax + 1), lmax + 1))
else:
row_alm = np.zeros((lfreq, npol * (lmax + 1), lmax + 1),
dtype=np.complex128)
# Perform the transposition to distribute different m's across processes. Neat
# tip, putting a shorter value for the number of columns, trims the array at
# the same time
col_alm = mpiutil.transpose_blocks(row_alm, (nfreq, npol *
(lmax + 1), mmax + 1))
# Transpose and reshape to shift m index first.
col_alm = np.transpose(col_alm,
(2, 0, 1)).reshape(lm, nfreq, npol, lmax + 1)
# Create storage for visibility data
vis_data = np.zeros((lm, nfreq, bt.ntel), dtype=np.complex128)
# Iterate over m's local to this process and generate the corresponding
# visibilities
for mp, mi in enumerate(range(sm, em)):
vis_data[mp] = bt.project_vector_sky_to_telescope(mi, col_alm[mp])
# Rearrange axes such that frequency is last (as we want to divide
# frequencies across processors)
row_vis = vis_data.transpose(
(0, 2, 1)) #.reshape((lm * bt.ntel, nfreq))
# Parallel transpose to get all m's back onto the same processor
col_vis_tmp = mpiutil.transpose_blocks(row_vis,
((mmax + 1), bt.ntel, nfreq))
col_vis_tmp = col_vis_tmp.reshape(mmax + 1, 2, tel.npairs, lfreq)
# Transpose the local section to make the m's the last axis and unwrap the
# positive and negative m at the same time.
col_vis[..., 0] = col_vis_tmp[0, 0]
for mi in range(1, mmax + 1):
col_vis[..., mi] = col_vis_tmp[mi, 0]
col_vis[..., -mi] = col_vis_tmp[
mi, 1].conj() # Conjugate only (not (-1)**m - see paper)
del col_vis_tmp
## If we're simulating noise, create a realisation and add it to col_vis
if ndays > 0:
# Fetch the noise powerspectrum
noise_ps = tel.noisepower(np.arange(tel.npairs)[:, np.newaxis],
np.array(local_freq)[np.newaxis, :],
ndays=ndays).reshape(tel.npairs,
lfreq)[:, :, 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_vis = (np.array([1.0, 1.0J]) *
np.random.standard_normal(col_vis.shape +
(2, ))).sum(axis=-1)
noise_vis *= (noise_ps / 2.0)**0.5
# Reset RNG
if seed is not None:
np.random.seed()
# Add into main noise sims
col_vis += noise_vis
del noise_vis
# Fourier transform m-modes back to get timestream.
vis_stream = np.fft.ifft(col_vis, axis=-1) * ntime
vis_stream = vis_stream.reshape(tel.npairs, lfreq, ntime)
# The time samples the visibility is calculated at
tphi = np.linspace(0, 2 * np.pi, ntime, endpoint=False)
# Create timestream object
tstream = Timestream(outdir, m)
## Iterate over the local frequencies and write them to disk.
for lfi, fi in enumerate(local_freq):
# Make directory if required
if not os.path.exists(tstream._fdir(fi)):
os.makedirs(tstream._fdir(fi))
# Write file contents
with h5py.File(tstream._ffile(fi), 'w') as f:
# Timestream data
f.create_dataset('/timestream', data=vis_stream[:, lfi])
f.create_dataset('/phi', data=tphi)
# 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)
# Write metadata
f.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
f.attrs['ntime'] = ntime
mpiutil.barrier()
return tstream
def simulate(m, outdir, maps=[], ndays=None, resolution=0, seed=None, **kwargs):
"""Create a simulated timestream and save it to disk.
Parameters
----------
m : ProductManager object
Products of telescope to simulate.
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.
Returns
-------
timestream : Timestream
"""
## Read in telescope system
bt = m.beamtransfer
tel = bt.telescope
lmax = tel.lmax
mmax = tel.mmax
nfreq = tel.nfreq
npol = tel.num_pol_sky
projmaps = (len(maps) > 0)
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
local_freq = range(sfreq, efreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# 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))
col_vis = np.zeros((tel.npairs, lfreq, ntime), dtype=np.complex128)
## If we want to add maps use the m-mode formalism to project a skymap
## into visibility space.
if projmaps:
# Load file to find out the map shapes.
with h5py.File(maps[0], 'r') as f:
mapshape = f['map'].shape
if lfreq > 0:
# Allocate array to store the local frequencies
row_map = np.zeros((lfreq,) + mapshape[1:], dtype=np.float64)
# Read in and sum up the local frequencies of the supplied maps.
for mapfile in maps:
with h5py.File(mapfile, 'r') as f:
row_map += f['map'][sfreq:efreq]
# Calculate the alm's for the local sections
row_alm = hputil.sphtrans_sky(row_map, lmax=lmax).reshape((lfreq, npol * (lmax+1), lmax+1))
else:
row_alm = np.zeros((lfreq, npol * (lmax+1), lmax+1), dtype=np.complex128)
# Perform the transposition to distribute different m's across processes. Neat
# tip, putting a shorter value for the number of columns, trims the array at
# the same time
col_alm = mpiutil.transpose_blocks(row_alm, (nfreq, npol * (lmax+1), mmax+1))
# Transpose and reshape to shift m index first.
col_alm = np.transpose(col_alm, (2, 0, 1)).reshape(lm, nfreq, npol, lmax+1)
# Create storage for visibility data
vis_data = np.zeros((lm, nfreq, bt.ntel), dtype=np.complex128)
# Iterate over m's local to this process and generate the corresponding
# visibilities
for mp, mi in enumerate(range(sm, em)):
vis_data[mp] = bt.project_vector_sky_to_telescope(mi, col_alm[mp])
# Rearrange axes such that frequency is last (as we want to divide
# frequencies across processors)
row_vis = vis_data.transpose((0, 2, 1))#.reshape((lm * bt.ntel, nfreq))
# Parallel transpose to get all m's back onto the same processor
col_vis_tmp = mpiutil.transpose_blocks(row_vis, ((mmax+1), bt.ntel, nfreq))
col_vis_tmp = col_vis_tmp.reshape(mmax + 1, 2, tel.npairs, lfreq)
# Transpose the local section to make the m's the last axis and unwrap the
# positive and negative m at the same time.
col_vis[..., 0] = col_vis_tmp[0, 0]
for mi in range(1, mmax+1):
col_vis[..., mi] = col_vis_tmp[mi, 0]
col_vis[..., -mi] = col_vis_tmp[mi, 1].conj() # Conjugate only (not (-1)**m - see paper)
del col_vis_tmp
## If we're simulating noise, create a realisation and add it to col_vis
if ndays > 0:
# Fetch the noise powerspectrum
noise_ps = tel.noisepower(np.arange(tel.npairs)[:, np.newaxis], np.array(local_freq)[np.newaxis, :], ndays=ndays).reshape(tel.npairs, lfreq)[:, :, 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_vis = (np.array([1.0, 1.0J]) * np.random.standard_normal(col_vis.shape + (2,))).sum(axis=-1)
noise_vis *= (noise_ps / 2.0)**0.5
# Reset RNG
if seed is not None:
np.random.seed()
# Add into main noise sims
col_vis += noise_vis
del noise_vis
# Fourier transform m-modes back to get timestream.
vis_stream = np.fft.ifft(col_vis, axis=-1) * ntime
vis_stream = vis_stream.reshape(tel.npairs, lfreq, ntime)
# The time samples the visibility is calculated at
tphi = np.linspace(0, 2*np.pi, ntime, endpoint=False)
# Create timestream object
tstream = Timestream(outdir, m)
## Iterate over the local frequencies and write them to disk.
for lfi, fi in enumerate(local_freq):
# Make directory if required
if not os.path.exists(tstream._fdir(fi)):
os.makedirs(tstream._fdir(fi))
# Write file contents
with h5py.File(tstream._ffile(fi), 'w') as f:
# Timestream data
f.create_dataset('/timestream', data=vis_stream[:, lfi])
f.create_dataset('/phi', data=tphi)
# 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)
# Write metadata
f.attrs['beamtransfer_path'] = os.path.abspath(bt.directory)
f.attrs['ntime'] = ntime
tstream.save()
mpiutil.barrier()
return tstream
## Useful output
print("==================================")
print("Projecting file:\n %s\ninto:\n %s" %
(args.mapfile, args.outfile))
print("Using beamtransfer: %s" % args.teldir)
print("Truncating to modes with S/N > %f" % cut)
print("==================================")
# Calculate alm's and broadcast
print("Read in skymap.")
f = h5py.File(args.mapfile)
skymap = f["map"][:]
f.close()
nside = healpy.get_nside(skymap[0])
alm = hputil.sphtrans_sky(skymap, lmax=cyl.lmax)
# else:
# almr = None
mpiutil.world.Bcast([alm, MPI.COMPLEX16], root=0)
cb = cyl.baselines + np.array([[cyl.u_width, 0.0]])
if cyl.positive_m_only:
taumax = (cb**2).sum(axis=-1)**0.5 / 3e8
else:
taumax = (np.concatenate((cb, cb))**2).sum(axis=-1)**0.5 / 3e8
tau = np.fft.fftfreq(cyl.nfreq,
(cyl.frequencies[1] - cyl.frequencies[0]) * 1e6)
# blmask = (np.abs(tau)[:, np.newaxis] > cut * (taumax[np.newaxis, :] + 10.0 / 3e8))
def mapmake_full(self, nside, mapname, nbin=None, dirty=False, method='svd', normalize=True, threshold=1.0e3, eps=0.01, correct_order=0, prior_map_file=None):
nfreq = self.telescope.nfreq
if nbin is None:
nbin = nfreq
else:
if (nbin < 1 or nbin > nfreq): # invalid nbin
nbin = nfreq
else:
nbin = int(nbin)
if prior_map_file is not None:
# read in the prior sky map
with h5py.File(prior_map_file, 'r') as f:
prior_map = f['map'][:] # shape (nbin, npol, npix)
# alm of the prior map
alm0 = hputil.sphtrans_sky(prior_map, lmax=self.telescope.lmax).reshape(nbin, self.telescope.num_pol_sky, self.telescope.lmax+1, self.telescope.lmax+1) # shape (nbin, npol, lmax+1, lmax+1)
else:
alm0 = None
def _make_alm(mi):
print "Making %i" % mi
mmode = self.mmode(mi)
if dirty:
sphmode = self.beamtransfer.project_vector_backward_dirty(mi, mmode, nbin, normalize, threshold)
else:
if method == 'svd':
sphmode = self.beamtransfer.project_vector_telescope_to_sky(mi, mmode, nbin)
elif method == 'tk':
# sphmode = self.beamtransfer.project_vector_telescope_to_sky_tk(mi, mmode, nbin, eps=eps)
mmode0 = alm0[:, :, :, mi] if alm0 is not None else None
sphmode = self.beamtransfer.project_vector_telescope_to_sky_tk(mi, mmode, nbin, eps=eps, correct_order=correct_order, mmode0=mmode0)
else:
raise ValueError('Unknown map-making method %s' % method)
return sphmode
alm_list = mpiutil.parallel_map(_make_alm, range(self.telescope.mmax + 1), root=0, method='rand')
if mpiutil.rank0:
# get center freq of each bin
n, s, e = mpiutil.split_m(nfreq, nbin)
cfreqs = np.array([ self.beamtransfer.telescope.frequencies[(s[i]+e[i])/2] for i in range(nbin) ])
alm = np.zeros((nbin, self.telescope.num_pol_sky, self.telescope.lmax + 1,
self.telescope.lmax + 1), dtype=np.complex128)
mlist = range(1 if self.no_m_zero else 0, self.telescope.mmax + 1)
for mi in mlist:
alm[..., mi] = alm_list[mi]
alm[:, :, 100:, 1] = 0
skymap = hputil.sphtrans_inv_sky(alm, nside)
with h5py.File(self.output_directory + '/' + mapname, 'w') as f:
f.create_dataset('/map', data=skymap)
f.attrs['frequency'] = cfreqs
f.attrs['polarization'] = np.array(['I', 'Q', 'U', 'V'])[:self.beamtransfer.telescope.num_pol_sky]
mpiutil.barrier()
if mpiutil.rank0:
## Useful output
print "=================================="
print "Projecting file:\n %s\ninto:\n %s" % (args.mapfile, args.outfile)
print "Using beamtransfer: %s" % args.teldir
print "Truncating to modes with S/N > %f" % cut
print "=================================="
# Calculate alm's and broadcast
print "Read in skymap."
f = h5py.File(args.mapfile, 'r')
skymap = f['map'][:]
f.close()
nside = healpy.get_nside(skymap[0])
alm = hputil.sphtrans_sky(skymap, lmax=cyl.lmax)
#else:
# almr = None
#mpiutil.world.Bcast([alm, MPI.COMPLEX16], root=0)
def projm(mi):
## Worker function for mapping over list and projecting onto signal modes.
print "Projecting %i" % mi
mvals, mvecs = klt.modes_m(mi, threshold=cut)
if mvals is None:
def generate(self):
for mentry in self.maps:
mfile = mentry['file']
stem = mentry['stem']
if mpiutil.rank0 and not os.path.exists(os.path.dirname(stem)):
os.makedirs(os.path.dirname(stem))
mpiutil.barrier()
if self.copy_orig:
shutil.copy(mfile, stem + 'orig.hdf5')
print "============\nProjecting file %s\n============\n" % mfile
## Load map and perform spherical harmonic transform
if mpiutil.rank0:
# Calculate alm's and broadcast
print "Read in skymap."
f = h5py.File(mfile, 'r')
skymap = f['map'][:]
f.close()
nside = healpy.get_nside(skymap[0])
alm = hputil.sphtrans_sky(skymap, lmax=self.telescope.lmax)
else:
alm = None
## Function to write out a map from the collected array of alms
def _write_map_from_almarray(almp, filename, attrs=None):
if mpiutil.rank0:
almp = np.squeeze(np.transpose(almp, axes=(2, 1, 3, 0)))
almf = np.zeros(
(almp.shape[0], almp.shape[1], almp.shape[1]),
dtype=np.complex128)
almf[:, :, :almp.shape[2]] = almp
pmap = hputil.sphtrans_inv_sky(almf, self.nside)
f = h5py.File(filename, 'w')
if attrs is not None:
for key, val in attrs.items():
f.attrs[repr(key)] = val
f.create_dataset('/map', data=pmap)
f.close()
mpiutil.barrier()
## Broadcast set of alms to the world
alm = mpiutil.world.bcast(alm, root=0)
mlist = range(self.kltransform.telescope.mmax + 1)
nevals = self.beamtransfer.ntel * self.beamtransfer.nfreq
## Construct beam projection of map
if self.beam_proj:
def proj_beam(mi):
print "Projecting %i" % mi
bproj = self.beamtransfer.project_vector_forward(
mi, alm[:, :, mi]).flatten()
return self.beamtransfer.project_vector_backward(mi, bproj)
shape = (self.telescope.nfreq, self.telescope.num_pol_sky,
self.telescope.lmax + 1)
almp = kltransform.collect_m_array(mlist, proj_beam, shape,
np.complex128)
_write_map_from_almarray(almp, stem + "beam.hdf5")
mpiutil.barrier()
## Construct EV projection of map
if self.evec_proj:
def proj_evec(mi):
## Worker function for mapping over list and projecting onto signal modes.
print "Projecting %i" % mi
p2 = np.zeros(nevals, dtype=np.complex128)
if self.kltransform.modes_m(mi)[0] is not None:
p1 = self.kltransform.project_sky_vector_forward(
mi, alm[:, :, mi])
p2[-p1.size:] = p1
return p2
shape = (nevals, )
evp = kltransform.collect_m_array(mlist, proj_evec, shape,
np.complex128)
if mpiutil.rank0:
f = h5py.File(stem + "ev.hdf5", 'w')
f.create_dataset("/evec_proj", data=evp)
f.close()
mpiutil.barrier()
## Iterate over noise cuts and filter out noise.
for cut in self.thresholds:
def filt_kl(mi):
## Worker function for mapping over list and projecting onto signal modes.
print "Projecting %i" % mi
mvals, mvecs = self.kltransform.modes_m(mi, threshold=cut)
if mvals is None:
return None
ev_vec = self.kltransform.project_sky_vector_forward(
mi, alm[:, :, mi], threshold=cut)
tel_vec = self.kltransform.project_tel_vector_backward(
mi, ev_vec, threshold=cut)
alm2 = self.beamtransfer.project_vector_backward(
mi, tel_vec)
return alm2
shape = (self.telescope.nfreq, self.telescope.num_pol_sky,
self.telescope.lmax + 1)
almp = kltransform.collect_m_array(mlist, filt_kl, shape,
np.complex128)
_write_map_from_almarray(almp, stem + ("kl_%g.hdf5" % cut),
{'threshold': cut})
mpiutil.barrier()
def generate(self):
for mentry in self.maps:
mfile = mentry['file']
stem = mentry['stem']
if mpiutil.rank0 and not os.path.exists(os.path.dirname(stem)):
os.makedirs(os.path.dirname(stem))
mpiutil.barrier()
if self.copy_orig:
shutil.copy(mfile, stem + 'orig.hdf5')
print "============\nProjecting file %s\n============\n" % mfile
## Load map and perform spherical harmonic transform
if mpiutil.rank0:
# Calculate alm's and broadcast
print "Read in skymap."
f = h5py.File(mfile, 'r')
skymap = f['map'][:]
f.close()
nside = healpy.get_nside(skymap[0])
alm = hputil.sphtrans_sky(skymap, lmax=self.telescope.lmax)
else:
alm = None
## Function to write out a map from the collected array of alms
def _write_map_from_almarray(almp, filename, attrs=None):
if mpiutil.rank0:
almp = np.squeeze(np.transpose(almp, axes=(2, 1, 3, 0)))
almf = np.zeros((almp.shape[0], almp.shape[1], almp.shape[1]), dtype=np.complex128)
almf[:, :, :almp.shape[2]] = almp
pmap = hputil.sphtrans_inv_sky(almf, self.nside)
f = h5py.File(filename, 'w')
if attrs is not None:
for key, val in attrs.items():
f.attrs[repr(key)] = val
f.create_dataset('/map', data=pmap)
f.close()
mpiutil.barrier()
## Broadcast set of alms to the world
alm = mpiutil.world.bcast(alm, root=0)
mlist = range(self.kltransform.telescope.mmax+1)
nevals = self.beamtransfer.ntel * self.beamtransfer.nfreq
## Construct beam projection of map
if self.beam_proj:
def proj_beam(mi):
print "Projecting %i" % mi
bproj = self.beamtransfer.project_vector_forward(mi, alm[:, :, mi]).flatten()
return self.beamtransfer.project_vector_backward(mi, bproj)
shape = (self.telescope.nfreq, self.telescope.num_pol_sky, self.telescope.lmax+1)
almp = kltransform.collect_m_array(mlist, proj_beam, shape, np.complex128)
_write_map_from_almarray(almp, stem + "beam.hdf5")
mpiutil.barrier()
## Construct EV projection of map
if self.evec_proj:
def proj_evec(mi):
## Worker function for mapping over list and projecting onto signal modes.
print "Projecting %i" % mi
p2 = np.zeros(nevals, dtype=np.complex128)
if self.kltransform.modes_m(mi)[0] is not None:
p1 = self.kltransform.project_sky_vector_forward(mi, alm[:, :, mi])
p2[-p1.size:] = p1
return p2
shape = (nevals,)
evp = kltransform.collect_m_array(mlist, proj_evec, shape, np.complex128)
if mpiutil.rank0:
f = h5py.File(stem + "ev.hdf5", 'w')
f.create_dataset("/evec_proj", data=evp)
f.close()
mpiutil.barrier()
## Iterate over noise cuts and filter out noise.
for cut in self.thresholds:
def filt_kl(mi):
## Worker function for mapping over list and projecting onto signal modes.
print "Projecting %i" % mi
mvals, mvecs = self.kltransform.modes_m(mi, threshold=cut)
if mvals is None:
return None
ev_vec = self.kltransform.project_sky_vector_forward(mi, alm[:, :, mi], threshold=cut)
tel_vec = self.kltransform.project_tel_vector_backward(mi, ev_vec, threshold=cut)
alm2 = self.beamtransfer.project_vector_backward(mi, tel_vec)
return alm2
shape = (self.telescope.nfreq, self.telescope.num_pol_sky, self.telescope.lmax+1)
almp = kltransform.collect_m_array(mlist, filt_kl, shape, np.complex128)
_write_map_from_almarray(almp, stem + ("kl_%g.hdf5" % cut), {'threshold' : cut})
mpiutil.barrier()
def process(self, map_):
"""Simulate a SiderealStream
Parameters
----------
map : :class:`containers.Map`
The sky map to process to into a sidereal stream. Frequencies in the
map, must match the Beam Transfer matrices.
Returns
-------
ss : SiderealStream
Stacked sidereal day.
feeds : list of CorrInput
Description of the feeds simulated.
"""
if self.done:
raise pipeline.PipelineStopIteration
# Read in telescope system
bt = self.beamtransfer
tel = self.telescope
lmax = tel.lmax
mmax = tel.mmax
nfreq = tel.nfreq
npol = tel.num_pol_sky
lfreq, sfreq, efreq = mpiutil.split_local(nfreq)
lm, sm, em = mpiutil.split_local(mmax + 1)
# Set the minimum resolution required for the sky.
ntime = 2 * mmax + 1
freqmap = map_.index_map["freq"][:]
row_map = map_.map[:]
if (tel.frequencies != freqmap["centre"]).any():
raise ValueError(
"Frequencies in map do not match those in Beam Transfers.")
# Calculate the alm's for the local sections
row_alm = hputil.sphtrans_sky(row_map, lmax=lmax).reshape(
(lfreq, npol * (lmax + 1), lmax + 1))
# Trim off excess m's and wrap into MPIArray
row_alm = row_alm[..., :(mmax + 1)]
row_alm = mpiarray.MPIArray.wrap(row_alm, axis=0)
# Perform the transposition to distribute different m's across processes. Neat
# tip, putting a shorter value for the number of columns, trims the array at
# the same time
col_alm = row_alm.redistribute(axis=2)
# Transpose and reshape to shift m index first.
col_alm = col_alm.transpose((2, 0, 1)).reshape(
(None, nfreq, npol, lmax + 1))
# Create storage for visibility data
vis_data = mpiarray.MPIArray((mmax + 1, nfreq, bt.ntel),
axis=0,
dtype=np.complex128)
vis_data[:] = 0.0
# Iterate over m's local to this process and generate the corresponding
# visibilities
for mp, mi in vis_data.enumerate(axis=0):
vis_data[mp] = bt.project_vector_sky_to_telescope(
mi, col_alm[mp].view(np.ndarray))
# Rearrange axes such that frequency is last (as we want to divide
# frequencies across processors)
row_vis = vis_data.transpose((0, 2, 1))
# Parallel transpose to get all m's back onto the same processor
col_vis_tmp = row_vis.redistribute(axis=2)
col_vis_tmp = col_vis_tmp.reshape((mmax + 1, 2, tel.npairs, None))
# Transpose the local section to make the m's the last axis and unwrap the
# positive and negative m at the same time.
col_vis = mpiarray.MPIArray((tel.npairs, nfreq, ntime),
axis=1,
dtype=np.complex128)
col_vis[:] = 0.0
col_vis[..., 0] = col_vis_tmp[0, 0]
for mi in range(1, mmax + 1):
col_vis[..., mi] = col_vis_tmp[mi, 0]
col_vis[..., -mi] = col_vis_tmp[
mi, 1].conj() # Conjugate only (not (-1)**m - see paper)
del col_vis_tmp
# Fourier transform m-modes back to get final timestream.
vis_stream = np.fft.ifft(col_vis, axis=-1) * ntime
vis_stream = vis_stream.reshape((tel.npairs, lfreq, ntime))
vis_stream = vis_stream.transpose((1, 0, 2)).copy()
# Try and fetch out the feed index and info from the telescope object.
try:
feed_index = tel.input_index
except AttributeError:
feed_index = tel.nfeed
# Construct a product map
prod_map = np.zeros(tel.uniquepairs.shape[0],
dtype=[("input_a", int), ("input_b", int)])
prod_map["input_a"] = tel.uniquepairs[:, 0]
prod_map["input_b"] = tel.uniquepairs[:, 1]
# Construct container and set visibility data
sstream = containers.SiderealStream(
freq=freqmap,
ra=ntime,
input=feed_index,
prod=prod_map,
distributed=True,
comm=map_.comm,
)
sstream.vis[:] = mpiarray.MPIArray.wrap(vis_stream, axis=0)
sstream.weight[:] = 1.0
self.done = True
return sstream
