def generate_mmodes(self): """Calculate the m-modes corresponding to the Timestream. Perform an MPI transpose for efficiency. """ if os.path.exists(self.output_directory + "/mmodes/COMPLETED_M"): if mpiutil.rank0: print "******* m-files already generated ********" return tel = self.telescope mmax = tel.mmax nfreq = tel.nfreq lfreq, sfreq, efreq = mpiutil.split_local(nfreq) lm, sm, em = mpiutil.split_local(mmax + 1) # Load in the local frequencies of the time stream tstream = np.zeros((lfreq, tel.npairs, self.ntime), dtype=np.complex128) for lfi, fi in enumerate(range(sfreq, efreq)): tstream[lfi] = self.timestream_f(fi) # FFT to calculate the m-modes for the timestream row_mmodes = np.fft.fft(tstream, axis=-1) / self.ntime ## Combine positive and negative m parts. row_mpairs = np.zeros((lfreq, 2, tel.npairs, mmax + 1), dtype=np.complex128) row_mpairs[:, 0, ..., 0] = row_mmodes[..., 0] for mi in range(1, mmax + 1): row_mpairs[:, 0, ..., mi] = row_mmodes[..., mi] row_mpairs[:, 1, ..., mi] = row_mmodes[..., -mi].conj() # Transpose to get the entirety of an m-mode on each process (i.e. all frequencies) col_mmodes = mpiutil.transpose_blocks(row_mpairs, (nfreq, 2, tel.npairs, mmax + 1)) # Transpose the local section to make the m's first col_mmodes = np.transpose(col_mmodes, (3, 0, 1, 2)) for lmi, mi in enumerate(range(sm, em)): # Make directory for each m-mode if not os.path.exists(self._mdir(mi)): os.makedirs(self._mdir(mi)) # Create the m-file and save the result. with h5py.File(self._mfile(mi), 'w') as f: f.create_dataset('/mmode', data=col_mmodes[lmi]) f.attrs['m'] = mi if mpiutil.rank0: # Make file marker that the m's have been correctly generated: open(self.output_directory + "/mmodes/COMPLETED_M", 'a').close() mpiutil.barrier()
def generate_mmodes(self): """Calculate the m-modes corresponding to the Timestream. Perform an MPI transpose for efficiency. """ if os.path.exists(self.output_directory + "/mmodes/COMPLETED_M"): if mpiutil.rank0: print "******* m-files already generated ********" return tel = self.telescope mmax = tel.mmax nfreq = tel.nfreq lfreq, sfreq, efreq = mpiutil.split_local(nfreq) lm, sm, em = mpiutil.split_local(mmax + 1) # Load in the local frequencies of the time stream tstream = np.zeros((lfreq, tel.npairs, self.ntime), dtype=np.complex128) for lfi, fi in enumerate(range(sfreq, efreq)): tstream[lfi] = self.timestream_f(fi) # FFT to calculate the m-modes for the timestream row_mmodes = np.fft.fft(tstream, axis=-1) / self.ntime ## Combine positive and negative m parts. row_mpairs = np.zeros((lfreq, 2, tel.npairs, mmax+1), dtype=np.complex128) row_mpairs[:, 0, ..., 0] = row_mmodes[..., 0] for mi in range(1, mmax+1): row_mpairs[:, 0, ..., mi] = row_mmodes[..., mi] row_mpairs[:, 1, ..., mi] = row_mmodes[..., -mi].conj() # Transpose to get the entirety of an m-mode on each process (i.e. all frequencies) col_mmodes = mpiutil.transpose_blocks(row_mpairs, (nfreq, 2, tel.npairs, mmax + 1)) # Transpose the local section to make the m's first col_mmodes = np.transpose(col_mmodes, (3, 0, 1, 2)) for lmi, mi in enumerate(range(sm, em)): # Make directory for each m-mode if not os.path.exists(self._mdir(mi)): os.makedirs(self._mdir(mi)) # Create the m-file and save the result. with h5py.File(self._mfile(mi), 'w') as f: f.create_dataset('/mmode', data=col_mmodes[lmi]) f.attrs['m'] = mi if mpiutil.rank0: # Make file marker that the m's have been correctly generated: open(self.output_directory + "/mmodes/COMPLETED_M", 'a').close() mpiutil.barrier()
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 redistribute(self, axis): """Change the axis that the array is distributed over. Parameters ---------- axis : integer Axis to distribute over. Returns ------- array : MPIArray A new copy of the array distributed over the specified axis. Note that the local section will have changed. """ # Check to see if this is the current distributed axis if self.axis == axis or self.comm is None: return self # Test to see if the datatype is one understood by MPI, this can # probably be fixed up at somepoint by creating a datatype of the right # number of bytes try: mpiutil.typemap(self.dtype) except KeyError: if self.comm is None or self.comm.rank == 0: import warnings warnings.warn('Cannot redistribute array of compound datatypes. Sorry!!') return self # Construct the list of the axes to swap around axlist_f = list(range(len(self.shape))) # Remove the axes we are going to swap around axlist_f.remove(self.axis) axlist_f.remove(axis) # Move the current dist axis to the front, and the new to the end axlist_f.insert(0, self.axis) axlist_f.append(axis) # Perform a local transpose on the array to get the axes in the correct order trans_arr = self.view(np.ndarray).transpose(axlist_f).copy() # Perform the global transpose tmp_gshape = (self.global_shape[self.axis],) + trans_arr.shape[1:] trans_arr = mpiutil.transpose_blocks(trans_arr, tmp_gshape, comm=self.comm) axlist_b = list(range(len(self.shape))) axlist_b.pop(0) last = axlist_b.pop(-1) if self.axis < axis: # This has to awkwardly depend on the order of the axes axlist_b.insert(self.axis, 0) axlist_b.insert(axis, last) else: axlist_b.insert(axis, last) axlist_b.insert(self.axis, 0) # Perform the local transpose to get the axes back in the correct order trans_arr = trans_arr.transpose(axlist_b) # Create a new MPIArray object out of the data dist_arr = MPIArray(self.global_shape, axis=axis, dtype=self.dtype, comm=self.comm) dist_arr[:] = trans_arr return dist_arr
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
def redistribute(self, axis): """Change the axis that the array is distributed over. Parameters ---------- axis : integer Axis to distribute over. Returns ------- array : MPIArray A new copy of the array distributed over the specified axis. Note that the local section will have changed. """ # Check to see if this is the current distributed axis if self.axis == axis or self.comm is None: return self # Test to see if the datatype is one understood by MPI, this can # probably be fixed up at somepoint by creating a datatype of the right # number of bytes try: mpiutil.typemap(self.dtype) except KeyError: if self.comm is None or self.comm.rank == 0: import warnings warnings.warn( 'Cannot redistribute array of compound datatypes. Sorry!!') return self # Construct the list of the axes to swap around axlist_f = list(range(len(self.shape))) # Remove the axes we are going to swap around axlist_f.remove(self.axis) axlist_f.remove(axis) # Move the current dist axis to the front, and the new to the end axlist_f.insert(0, self.axis) axlist_f.append(axis) # Perform a local transpose on the array to get the axes in the correct order trans_arr = self.view(np.ndarray).transpose(axlist_f).copy() # Perform the global transpose tmp_gshape = (self.global_shape[self.axis], ) + trans_arr.shape[1:] trans_arr = mpiutil.transpose_blocks(trans_arr, tmp_gshape, comm=self.comm) axlist_b = list(range(len(self.shape))) axlist_b.pop(0) last = axlist_b.pop(-1) if self.axis < axis: # This has to awkwardly depend on the order of the axes axlist_b.insert(self.axis, 0) axlist_b.insert(axis, last) else: axlist_b.insert(axis, last) axlist_b.insert(self.axis, 0) # Perform the local transpose to get the axes back in the correct order trans_arr = trans_arr.transpose(axlist_b) # Create a new MPIArray object out of the data dist_arr = MPIArray(self.global_shape, axis=axis, dtype=self.dtype, comm=self.comm) dist_arr[:] = trans_arr return dist_arr