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