def generate_mmodes(self, ts_data=None): """Calculate the m-modes corresponding to the Timestream. Perform an MPI transpose for efficiency. """ completed_file = self._mdir + 'COMPLETED_M' if os.path.exists(completed_file): if mpiutil.rank0: print "******* m-files already generated ********" mpiutil.barrier() return # Make directory if required # if mpiutil.rank0 and not os.path.exists(self._mdir): # os.makedirs(self._mdir) try: os.makedirs(self._mdir) except OSError: # directory exists pass tel = self.telescope mmax = tel.mmax ntime = ts_data.shape[0] if ts_data is not None else self.ntime nbl = tel.nbase nfreq = tel.nfreq indices = list(itertools.product(np.arange(nfreq), np.arange(nbl))) lind, sind, eind = mpiutil.split_local(nfreq * nbl) # load the local section of the time stream tstream = np.zeros((ntime, lind), dtype=np.complex128) for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]): if ts_data is not None: tstream[:, ind] = ts_data[:, f_ind, bl_ind] else: with h5py.File(self._tsfile, 'r') as f: tstream[:, ind] = f['/timestream'][:, f_ind, bl_ind] # FFT to get m-mode mmodes = np.fft.fft(tstream, axis=0) / ntime # m = 0 is at left mmodes = MPIArray.wrap(mmodes, axis=1) # redistribute along different m mmodes = mmodes.redistribute(axis=0) # save m-modes to file ms = np.concatenate([np.arange(0, mmax+1), np.arange(-mmax, 0)]) for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')): with h5py.File(self._mfile(mi), 'w') as f: f.create_dataset('/mmode', data=mmodes[ind].view(np.ndarray).reshape(nfreq, nbl)) f.attrs['m'] = mi mpiutil.barrier() if mpiutil.rank0: # Make file marker that the m's have been correctly generated: open(completed_file, 'a').close()
$ mpiexec -n 4 python mpiarray_demo.py """ import os import numpy as np from caput import mpiutil from caput.mpiarray import MPIArray rank = mpiutil.rank size = mpiutil.size # construct a MPIArray with global_shape (5, 4, 3) and distribute axis 0 shape = (5, 4, 3) dist_axis = 0 darr = MPIArray(global_shape=shape, axis=dist_axis, dtype=np.float32) print 'rank %d has global_shape = %s, local_shape = %s, local_offset = %s' % (rank, darr.global_shape, darr.local_shape, darr.local_offset) # from_numpy_array nparr = np.arange(6*5*4).reshape(6, 5, 4) darr1 = MPIArray.from_numpy_array(nparr, axis=0, root=None) # to_numpy_array nparr1 = darr1.to_numpy_array(root=0) if rank == 0: print 'rank 0: nparr1 == nparr: %s' % np.allclose(nparr, nparr1) else: print 'rank %d: nparr1 = %s' % (rank, nparr1) # to_hdf5 h5_file = 'test.hdf5' darr1.to_hdf5(h5_file, 'test', create=True) # remove the file
def _generate_mfiles(self, regen=False): completed_file = self._mdir + 'COMPLETED_BEAM' if os.path.exists(completed_file) and not regen: if mpiutil.rank0: print print '=' * 80 print "******* Beam transfer m-files already generated ********" mpiutil.barrier() return if mpiutil.rank0: print print '=' * 80 print 'Create beam transfer m-files...' st = time.time() # Calculate the Beam Transfer Matrices nfreq = self.telescope.nfreq nbl = self.telescope.nbase npol = self.telescope.num_pol_sky ntheta = self.telescope.theta_size nphi = self.telescope.phi_size # create file to save beam transfer matrices dsize = (nfreq, nbl, npol, ntheta) csize = (nfreq, 1, npol, ntheta) mmax = self.telescope.mmax ms = np.concatenate([np.arange(0, mmax+1), np.arange(-mmax, 0)]) # get local section of m'th for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')): with h5py.File(self._mfile(mi), 'w') as f: f.create_dataset('beam_m', dsize, chunks=csize, compression='lzf', dtype=np.complex128) f.attrs['m'] = mi # calculate the total memory needed for the transfer matrix total_memory = nfreq * nbl * npol * ntheta * nphi * 16.0 # Bytes, 16 for complex128 limit = 1.0 # GB, memory limit for each process # make each process have maximum `limit` GB sigle_memory = limit * 2**30 # Bytes # how many chunks num_chunks = np.int(np.ceil(total_memory / (mpiutil.size * sigle_memory))) # split bls to num_chunks sections if nbl < num_chunks: warnings.warn('Could not split to %d chunks for %d baselines' % (num_chunks, nbl)) num_chunks = min(num_chunks, nbl) num, start, end = mpiutil.split_m(nbl, num_chunks) for ci in range(num_chunks): if mpiutil.rank0: print "Starting chunk %i of %i" % (ci+1, num_chunks) tarray = self.telescope.transfer_matrices(np.arange(start[ci], end[ci]), np.arange(nfreq)) tarray = MPIArray.wrap(tarray, axis=0) # redistribute along different m tarray = tarray.redistribute(axis=3) # save beam transfer matrices to file for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')): with h5py.File(self._mfile(mi), 'r+') as f: f['beam_m'][:, start[ci]:end[ci]] = tarray[..., ind].view(np.ndarray).reshape(nfreq, num[ci], npol, ntheta) mpiutil.barrier() et = time.time() if mpiutil.rank0: # Make file marker that the m's have been correctly generated: open(completed_file, 'a').close() # Print out timing print "=== Create beam transfer m-files took %f s ===" % (et - st)
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
def mapmake_full(self, nside, maptype): mapfile = self._mapsdir + 'map_%s.hdf5' % maptype Tmfile = self._Tmsdir + 'Tm_%s.hdf5' % maptype if os.path.exists(mapfile): if mpiutil.rank0: print "File %s exists. Skipping..." % mapfile mpiutil.barrier() return elif os.path.exists(Tmfile): if mpiutil.rank0: print "File %s exists. Read from it..." % Tmfile Tm = MPIArray.from_hdf5(Tmfile, 'Tm') else: def _make_Tm(mi): print "Making %i" % mi mmode = self.mmode(mi) return self.beamtransfer.project_vector_telescope_to_sky(mi, mmode) # if mpiutil.rank0 and not os.path.exists(self._Tmsdir): # # Make directory for Tms file # os.makedirs(self._Tmsdir) # Make directory for Tms file try: os.makedirs(self._Tmsdir) except OSError: # directory exists pass tel = self.telescope mmax = tel.mmax lm, sm, em = mpiutil.split_local(mmax+1) nfreq = tel.nfreq npol = tel.num_pol_sky ntheta = tel.theta_size # the local Tm array Tm = np.zeros((nfreq, npol, ntheta, lm), dtype=np.complex128) for ind, mi in enumerate(range(sm, em)): Tm[..., ind] = _make_Tm(mi) Tm = MPIArray.wrap(Tm, axis=3) Tm = Tm.redistribute(axis=0) # redistribute along freq # Save Tm Tm.to_hdf5(Tmfile, 'Tm', create=True) # if mpiutil.rank0 and not os.path.exists(self._mapsdir): # # Make directory for maps file # os.makedirs(self._mapsdir) # Make directory for maps file try: os.makedirs(self._mapsdir) except OSError: # directory exists pass tel = self.telescope npol = tel.num_pol_sky ntime = self.ntime # irfft to get map # cart_map = np.fft.irfft(Tm, axis=3, n=ntime) * ntime # NOTE the normalization constant ntime here to be consistant with the simulation fft cart_map = np.fft.hfft(Tm, axis=3, n=ntime) lfreq = cart_map.shape[0] hp_map = np.zeros((lfreq, npol, 12*nside**2), dtype=cart_map.dtype) for fi in range(lfreq): for pi in range(npol): hp_map[fi, pi] = tel.cart_projector.inv_projmap(cart_map[fi, pi], nside) mpiutil.barrier() hp_map = MPIArray.wrap(hp_map, axis=0) # save map hp_map.to_hdf5(mapfile, 'map', create=True)
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
def mapmake_full(self, nside, maptype): mapfile = self._mapsdir + 'map_%s.hdf5' % maptype Tmfile = self._Tmsdir + 'Tm_%s.hdf5' % maptype if os.path.exists(mapfile): if mpiutil.rank0: print "File %s exists. Skipping..." % mapfile mpiutil.barrier() return elif os.path.exists(Tmfile): if mpiutil.rank0: print "File %s exists. Read from it..." % Tmfile Tm = MPIArray.from_hdf5(Tmfile, 'Tm') else: def _make_Tm(mi): print "Making %i" % mi mmode = self.mmode(mi) return self.beamtransfer.project_vector_telescope_to_sky( mi, mmode) # if mpiutil.rank0 and not os.path.exists(self._Tmsdir): # # Make directory for Tms file # os.makedirs(self._Tmsdir) # Make directory for Tms file try: os.makedirs(self._Tmsdir) except OSError: # directory exists pass tel = self.telescope mmax = tel.mmax lm, sm, em = mpiutil.split_local(mmax + 1) nfreq = tel.nfreq npol = tel.num_pol_sky ntheta = tel.theta_size # the local Tm array Tm = np.zeros((nfreq, npol, ntheta, lm), dtype=np.complex128) for ind, mi in enumerate(range(sm, em)): Tm[..., ind] = _make_Tm(mi) Tm = MPIArray.wrap(Tm, axis=3) Tm = Tm.redistribute(axis=0) # redistribute along freq # Save Tm Tm.to_hdf5(Tmfile, 'Tm', create=True) # if mpiutil.rank0 and not os.path.exists(self._mapsdir): # # Make directory for maps file # os.makedirs(self._mapsdir) # Make directory for maps file try: os.makedirs(self._mapsdir) except OSError: # directory exists pass tel = self.telescope npol = tel.num_pol_sky ntime = self.ntime # irfft to get map # cart_map = np.fft.irfft(Tm, axis=3, n=ntime) * ntime # NOTE the normalization constant ntime here to be consistant with the simulation fft cart_map = np.fft.hfft(Tm, axis=3, n=ntime) lfreq = cart_map.shape[0] hp_map = np.zeros((lfreq, npol, 12 * nside**2), dtype=cart_map.dtype) for fi in range(lfreq): for pi in range(npol): hp_map[fi, pi] = tel.cart_projector.inv_projmap( cart_map[fi, pi], nside) mpiutil.barrier() hp_map = MPIArray.wrap(hp_map, axis=0) # save map hp_map.to_hdf5(mapfile, 'map', create=True)
def generate_mmodes(self, ts_data=None): """Calculate the m-modes corresponding to the Timestream. Perform an MPI transpose for efficiency. """ completed_file = self._mdir + 'COMPLETED_M' if os.path.exists(completed_file): if mpiutil.rank0: print "******* m-files already generated ********" mpiutil.barrier() return # Make directory if required # if mpiutil.rank0 and not os.path.exists(self._mdir): # os.makedirs(self._mdir) try: os.makedirs(self._mdir) except OSError: # directory exists pass tel = self.telescope mmax = tel.mmax ntime = ts_data.shape[0] if ts_data is not None else self.ntime nbl = tel.nbase nfreq = tel.nfreq indices = list(itertools.product(np.arange(nfreq), np.arange(nbl))) lind, sind, eind = mpiutil.split_local(nfreq * nbl) # load the local section of the time stream tstream = np.zeros((ntime, lind), dtype=np.complex128) for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]): if ts_data is not None: tstream[:, ind] = ts_data[:, f_ind, bl_ind] else: with h5py.File(self._tsfile, 'r') as f: tstream[:, ind] = f['/timestream'][:, f_ind, bl_ind] # FFT to get m-mode mmodes = np.fft.fft(tstream, axis=0) / ntime # m = 0 is at left mmodes = MPIArray.wrap(mmodes, axis=1) # redistribute along different m mmodes = mmodes.redistribute(axis=0) # save m-modes to file ms = np.concatenate([np.arange(0, mmax + 1), np.arange(-mmax, 0)]) for ind, mi in enumerate(mpiutil.mpilist(ms, method='con')): with h5py.File(self._mfile(mi), 'w') as f: f.create_dataset('/mmode', data=mmodes[ind].view(np.ndarray).reshape( nfreq, nbl)) f.attrs['m'] = mi mpiutil.barrier() if mpiutil.rank0: # Make file marker that the m's have been correctly generated: open(completed_file, 'a').close()