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 __new__(cls, global_shape, axis=0, comm=None, *args, **kwargs): # if mpiutil.world is None: # raise RuntimeError('There is no mpi4py installation. Aborting.') if comm is None: comm = mpiutil.world # Determine local section of distributed axis local_num, local_start, local_end = mpiutil.split_local( global_shape[axis], comm=comm) # Figure out the local shape and offset lshape = list(global_shape) lshape[axis] = local_num loffset = [0] * len(global_shape) loffset[axis] = local_start # Create array arr = np.ndarray.__new__(cls, lshape, *args, **kwargs) # Set attributes of class arr._global_shape = global_shape arr._axis = axis arr._local_shape = tuple(lshape) arr._local_offset = tuple(loffset) arr._comm = comm return arr
def test_H5FileSelect_distributed(container_on_disk, fsel, isel): """Load H5 into parallel container while down-selecting axes.""" sel = {"dset1": (fsel, isel, slice(None)), "dset2": (fsel, slice(None))} # Tests are designed to run for 1, 2 or 4 processes assert 4 % comm.size == 0 m = MemGroup.from_hdf5(container_on_disk, selections=sel, distributed=True, comm=comm) d1 = dset1[(fsel, isel, slice(None))] d2 = dset2[(fsel, slice(None))] n, s, e = mpiutil.split_local(d1.shape[0], comm=comm) dslice = slice(s, e) # For debugging... # Need to dereference datasets as this is collective # md1 = m["dset1"][:] # md2 = m["dset2"][:] # for ri in range(comm.size): # if ri == comm.rank: # print(comm.rank) # print(md1.shape, d1.shape, d1[dslice].shape) # print(md1[0, :2, :2] if md1.size else "Empty") # print(d1[dslice][0, :2, :2] if d1[dslice].size else "Empty") # print() # comm.Barrier() assert np.all(m["dset1"][:] == d1[dslice]) assert np.all(m["dset2"][:] == d2[dslice])
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()
def __new__(cls, global_shape, axis=0, comm=None, *args, **kwargs): # if mpiutil.world is None: # raise RuntimeError('There is no mpi4py installation. Aborting.') if comm is None: comm = mpiutil.world # Determine local section of distributed axis local_num, local_start, local_end = mpiutil.split_local(global_shape[axis], comm=comm) # Figure out the local shape and offset lshape = list(global_shape) lshape[axis] = local_num loffset = [0] * len(global_shape) loffset[axis] = local_start # Create array arr = np.ndarray.__new__(cls, lshape, *args, **kwargs) # Set attributes of class arr._global_shape = global_shape arr._axis = axis arr._local_shape = tuple(lshape) arr._local_offset = tuple(loffset) arr._comm = comm return arr
def make_clzz_array(self): p_bands, s_bands, e_bands = mpiutil.split_all(self.nbands) p, s, e = mpiutil.split_local(self.nbands) self.clarray = np.zeros( ( self.nbands, self.telescope.lmax + 1, self.telescope.nfreq, self.telescope.nfreq, ), dtype=np.float64, ) for bi in range(s, e): self.clarray[bi] = self.make_clzz(self.band_pk[bi]) bandsize = ((self.telescope.lmax + 1) * self.telescope.nfreq * self.telescope.nfreq) sizes = p_bands * bandsize displ = s_bands * bandsize MPI.COMM_WORLD.Allgatherv(MPI.IN_PLACE, [self.clarray, sizes, displ, MPI.DOUBLE])
def test_io(self): import h5py # Cleanup directories fname = 'testdset.hdf5' if mpiutil.rank0 and os.path.exists(fname): os.remove(fname) mpiutil.barrier() gshape = (19, 17) ds = mpiarray.MPIArray(gshape, dtype=np.int64) ga = np.arange(np.prod(gshape)).reshape(gshape) l0, s0, e0 = mpiutil.split_local(gshape[0]) ds[:] = ga[s0:e0] ds.redistribute(axis=1).to_hdf5(fname, 'testds', create=True) if mpiutil.rank0: with h5py.File(fname, 'r') as f: h5ds = f['testds'][:] assert (h5ds == ga).all() ds2 = mpiarray.MPIArray.from_hdf5(fname, 'testds') assert (ds2 == ds).all()
def wrap(cls, array, axis, comm=None): """Turn a set of numpy arrays into a distributed MPIArray object. This is needed for functions such as `np.fft.fft` which always return an `np.ndarray`. Parameters ---------- array : np.ndarray Array to wrap. axis : integer Axis over which the array is distributed. The lengths are checked to try and ensure this is correct. comm : MPI.Comm, optional The communicator over which the array is distributed. If `None` (default), use `MPI.COMM_WORLD`. Returns ------- dist_array : MPIArray An MPIArray view of the input. """ # from mpi4py import MPI if comm is None: comm = mpiutil.world # Get axis length, both locally, and globally axlen = array.shape[axis] totallen = mpiutil.allreduce(axlen, comm=comm) # Figure out what the distributed layout should be local_num, local_start, local_end = mpiutil.split_local(totallen, comm=comm) # Check the local layout is consistent with what we expect, and send # result to all ranks layout_issue = mpiutil.allreduce(axlen != local_num, op=mpiutil.MAX, comm=comm) if layout_issue: raise Exception("Cannot wrap, distributed axis local length is incorrect.") # Set shape and offset lshape = array.shape global_shape = list(lshape) global_shape[axis] = totallen loffset = [0] * len(lshape) loffset[axis] = local_start # Setup attributes of class dist_arr = array.view(cls) dist_arr._global_shape = tuple(global_shape) dist_arr._axis = axis dist_arr._local_shape = tuple(lshape) dist_arr._local_offset = tuple(loffset) dist_arr._comm = comm return dist_arr
def test_redistribution(self): gshape = (1, 11, 2, 14, 3, 4) nelem = np.prod(gshape) garr = np.arange(nelem).reshape(gshape) l0, s0, e0 = mpiutil.split_local(11) l1, s1, e1 = mpiutil.split_local(14) l2, s2, e2 = mpiutil.split_local(4) arr = mpiarray.MPIArray(gshape, axis=1, dtype=np.int64) arr[:] = garr[:, s0:e0] arr2 = arr.redistribute(axis=3) assert (arr2 == garr[:, :, :, s1:e1]).view(np.ndarray).all() arr3 = arr.redistribute(axis=5) assert (arr3 == garr[:, :, :, :, :, s2:e2]).view(np.ndarray).all()
def test_construction(self): arr = mpiarray.MPIArray((10, 11), axis=1) l, s, e = mpiutil.split_local(11) # Check that global shape is set correctly assert arr.global_shape == (10, 11) assert arr.shape == (10, l) assert arr.local_offset == (0, s) assert arr.local_shape == (10, l)
def make_clzz_array(self): p_bands, s_bands, e_bands = mpiutil.split_all(self.nbands) p, s, e = mpiutil.split_local(self.nbands) self.clarray = np.zeros((self.nbands, self.telescope.lmax + 1, self.telescope.nfreq, self.telescope.nfreq), dtype=np.float64) for bi in range(s, e): self.clarray[bi] = self.make_clzz(self.band_pk[bi]) bandsize = (self.telescope.lmax + 1) * self.telescope.nfreq * self.telescope.nfreq sizes = p_bands * bandsize displ = s_bands * bandsize MPI.COMM_WORLD.Allgatherv(MPI.IN_PLACE, [self.clarray, sizes, displ, MPI.DOUBLE])
def test_LoadBasicCont_selection(ss_container, mpi_tmp_path): fname = str(mpi_tmp_path / "ss.h5") ss_container.save(fname) freq_range = [1, 4] ra_range = [3, 10] task = io.LoadBasicCont() task.files = [fname] task.selections = { "freq_range": freq_range, "ra_range": ra_range, } task.setup() ss_load = task.next() ss_vis = ss_load.vis[:] ss_weight = ss_load.weight[:] # Check the datasets (only a maximum of three ranks will have anything) nf = freq_range[1] - freq_range[0] if ss_load.comm.rank < nf: # Check the freq selection n, s, e = mpiutil.split_local(nf, comm=ss_load.comm) vis_real = np.arange(*freq_range)[s:e][:, np.newaxis, np.newaxis] assert (ss_vis.real == vis_real).all() # Check the ra selection vis_imag = np.arange(*ra_range)[np.newaxis, np.newaxis, :] assert (ss_vis.imag == vis_imag).all() # Check that nothing funky happened on the stack axis weight = np.arange(ss_container.vis.shape[1])[np.newaxis, :, np.newaxis] assert (ss_weight == weight).all() # Check the attributes... assert ss_load.attrs["test_attr1"] == "hello" assert ss_load.vis.attrs["test_attr2"] == "hello2" assert ss_load.weight.attrs["test_attr3"] == "hello3" # As we only put one item into the queue, this should end the iterations with pytest.raises(pipeline.PipelineStopIteration): task.next()
def transfer_matrices(self, bl_indices, f_indices): """Calculate the spherical harmonic transfer matrices for baseline and frequency combinations. Parameters ---------- bl_indices : array_like Indices of baselines to calculate. f_indices : array_like Indices of frequencies to calculate. Must be broadcastable against `bl_indices`. Returns ------- transfer : np.ndarray, dtype=np.complex128 An array containing the transfer functions. The shape is somewhat complicated, the first indices correspond to the broadcast size of `bl_indices` and `f_indices`, then there may be some polarisation indices, then finally the (l,m) indices, range (lside, 2*lside-1). """ ## Check indices are all in range if out_of_range(bl_indices, 0, self.npairs): raise Exception("Baseline indices aren't valid") if out_of_range(f_indices, 0, self.nfreq): raise Exception("Frequency indices aren't valid") # Generate the array for the Transfer functions nfreq = f_indices.size nbl = bl_indices.size nprod = nfreq * nbl indices = list(itertools.product(f_indices, bl_indices)) lind, sind, eind = mpiutil.split_local(nprod) # local section of the array tshape = (lind, self.num_pol_sky, self.theta_size, self.phi_size) tarray = np.zeros(tshape, dtype=np.complex128) for ind, (f_ind, bl_ind) in enumerate(indices[sind:eind]): tarray[ind] = self._transfer_single(bl_ind, f_ind, self.lmax) return tarray
def test_wrap(self): ds = mpiarray.MPIArray((10, 17)) df = np.fft.rfft(ds, axis=1) assert type(df) == np.ndarray da = mpiarray.MPIArray.wrap(df, axis=0) assert type(da) == mpiarray.MPIArray assert da.global_shape == (10, 9) l0, s0, e0 = mpiutil.split_local(10) assert da.local_shape == (l0, 9) if mpiutil.rank0: df = df[:-1] with self.assertRaises(Exception): mpiarray.MPIArray.wrap(df, axis=0)
def test_io(self): import h5py # Cleanup directories fname = 'testdset.hdf5' if mpiutil.rank0 and os.path.exists(fname): os.remove(fname) mpiutil.barrier() gshape = (19, 17) ds = mpiarray.MPIArray(gshape, dtype=np.int64) ga = np.arange(np.prod(gshape)).reshape(gshape) l0, s0, e0 = mpiutil.split_local(gshape[0]) ds[:] = ga[s0:e0] ds.redistribute(axis=1).to_hdf5(fname, 'testds', create=True) if mpiutil.rank0: with h5py.File(fname, 'r') as f: h5ds = f['testds'][:] assert (h5ds == ga).all() ds2 = mpiarray.MPIArray.from_hdf5(fname, 'testds') assert (ds2 == ds).all() mpiutil.barrier() if mpiutil.rank0 and os.path.exists(fname): os.remove(fname)
def test_wrap(self): ds = mpiarray.MPIArray((10, 17)) df = np.fft.rfft(ds, axis=1) assert type(df) == np.ndarray da = mpiarray.MPIArray.wrap(df, axis=0) assert type(da) == mpiarray.MPIArray assert da.global_shape == (10, 9) l0, s0, e0 = mpiutil.split_local(10) assert da.local_shape == (l0, 9) if mpiutil.rank0: df = df[:-1] if mpiutil.size > 1: with self.assertRaises(Exception): mpiarray.MPIArray.wrap(df, axis=0)
def test_reshape(self): gshape = (1, 11, 2, 14) l0, s0, e0 = mpiutil.split_local(11) arr = mpiarray.MPIArray(gshape, axis=1, dtype=np.int64) arr2 = arr.reshape((None, 28)) # Check type assert isinstance(arr2, mpiarray.MPIArray) # Check global shape assert arr2.global_shape == (11, 28) # Check local shape assert arr2.local_shape == (l0, 28) # Check local offset assert arr2.local_offset == (s0, 0) # Check axis assert arr2.axis == 0
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 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 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
def test_io(self): import h5py # Cleanup directories fname = 'testdset.hdf5' if mpiutil.rank0 and os.path.exists(fname): os.remove(fname) mpiutil.barrier() gshape = (19, 17) ds = mpiarray.MPIArray(gshape, dtype=np.int64) ga = np.arange(np.prod(gshape)).reshape(gshape) l0, s0, e0 = mpiutil.split_local(gshape[0]) ds[:] = ga[s0:e0] ds.redistribute(axis=1).to_hdf5(fname, 'testds', create=True) if mpiutil.rank0: with h5py.File(fname, 'r') as f: h5ds = f['testds'][:] assert (h5ds == ga).all() ds2 = mpiarray.MPIArray.from_hdf5(fname, 'testds') assert (ds2 == ds).all() mpiutil.barrier() # Check that reading over another distributed axis works ds3 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=1) assert ds3.shape[0] == gshape[0] assert ds3.shape[1] == mpiutil.split_local(gshape[1])[0] ds3 = ds3.redistribute(axis=0) assert (ds3 == ds).all() mpiutil.barrier() # Check a read with an arbitrary slice in there. This only checks the shape is correct. ds4 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=1, sel=(np.s_[3:10:2], np.s_[1:16:3])) assert ds4.shape[0] == 4 assert ds4.shape[1] == mpiutil.split_local(5)[0] mpiutil.barrier() # Check the read with a slice along the axis being read ds5 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=1, sel=(np.s_[:], np.s_[3:15:2])) assert ds5.shape[0] == gshape[0] assert ds5.shape[1] == mpiutil.split_local(6)[0] ds5 = ds5.redistribute(axis=0) assert (ds5 == ds[:, 3:15:2]).all() mpiutil.barrier() # Check the read with a slice along the axis being read ds6 = mpiarray.MPIArray.from_hdf5(fname, 'testds', axis=0, sel=(np.s_[:], np.s_[3:15:2])) ds6 = ds6.redistribute(axis=0) assert (ds6 == ds[:, 3:15:2]).all() mpiutil.barrier() if mpiutil.rank0 and os.path.exists(fname): os.remove(fname)
def process(self, rt): assert isinstance(rt, RawTimestream), '%s only works for RawTimestream object currently' % self.__class__.__name__ if not 'ns_on' in rt.iterkeys(): raise RuntimeError('No noise source info, can not do noise source calibration') local_bl_size, local_bl_start, local_bl_end = mpiutil.split_local(len(rt['blorder'])) rt.redistribute('baseline') num_mean = self.params['num_mean'] phs_only = self.params['phs_only'] save_gain = self.params['save_gain'] tag_output_iter = self.params['tag_output_iter'] gain_file = self.params['gain_file'] bl_incl = self.params['bl_incl'] bl_excl = self.params['bl_excl'] freq_incl = self.params['freq_incl'] freq_excl = self.params['freq_excl'] ns_stable = self.params['ns_stable'] last_transit2stable = self.params['last_transit2stable'] #=================================== normal_gain_file = self.params['normal_gain_file'] rt.gain_file = self.params['gain_file'] absolute_gain_filename = self.params['absolute_gain_filename'] use_center_data = self.params['use_center_data'] if use_center_data and rt['ns_on'].attrs['on_time'] <= 2: warnings.warn('The period %d <= 2, cannot get rid of the beginning and ending points. Use the whole average automatically!') use_center_data = False #=================================== # apply abs gain to data if ps_first if rt.ps_first: if absolute_gain_filename is None: raise NoPsGainFile('Absent parameter absolute_gain_file. In ps_first mode, absolute_gain_filename is required to process!') if not os.path.exists(output_path(absolute_gain_filename)): raise NoPsGainFile('No absolute gain file %s, do the ps calibration first!'%output_path(absolute_gain_filename)) with h5py.File(output_path(absolute_gain_filename, 'r')) as agfilein: if not ns_stable: # try: # if there is transit in this batch of data, build up transit amp and phs # rt.normal_index = np.where(np.abs(rt['jul_date'] - agfilein.attrs['transit_time']) < 2.e-6)[0][0] # to avoid numeric error. 1.e-6 julian date is about 0.1s if rt['sec1970'][0] <= agfilein.attrs['transit_time'] and rt['sec1970'][-1] >= agfilein.attrs['transit_time']: rt.normal_index = np.int64(np.around((agfilein.attrs['transit_time'] - rt['sec1970'][0])/rt.attrs['inttime'])) if mpiutil.rank0: print('Detected transit, time index %d, build up transit normalization gain file!'%rt.normal_index) build_normal_file = True rt.normal_phs = np.zeros((rt['vis'].shape[1], local_bl_size)) rt.normal_phs.fill(np.nan) if not phs_only: rt.normal_amp = np.zeros((rt['vis'].shape[1], local_bl_size)) rt.normal_amp.fill(np.nan) # except IndexError: # no transit in this batch of data, load transit amp and phs from file else: rt.normal_index = None # no transit flag build_normal_file = False if not os.path.exists(output_path(normal_gain_file)): time_info = (aipy.phs.juldate2ephem(agfilein.attrs['transit_jul'] + 8./24.), aipy.phs.juldate2ephem(rt['jul_date'][0] + 8./24.), aipy.phs.juldate2ephem(rt['jul_date'][-1] + 8./24.)) raise TransitGainNotRecorded('The transit %s is not in time range %s to %s and no transit normalization gain was recorded!'%time_info) else: if mpiutil.rank0: print('No transit, use existing transit normalization gain file!') with h5py.File(output_path(normal_gain_file), 'r') as filein: rt.normal_phs = filein['phs'][:, local_bl_start:local_bl_end] if not phs_only: rt.normal_amp = filein['amp'][:, local_bl_start:local_bl_end] else: rt.normal_index = None # no transit flag rt.normal_phs = np.zeros((rt['vis'].shape[1], local_bl_size)) rt.normal_phs.fill(np.nan) if not phs_only: rt.normal_amp = np.zeros((rt['vis'].shape[1], local_bl_size)) rt.normal_amp.fill(np.nan) try: stable_time = rt['pointingtime'][-1,-1] except KeyError: stable_time = rt['transitsource'][-1, 0] if stable_time > 0: # the time when the antenna will be pointing at the target region was recorded, use it stable_time += 300 # plus 5 min to ensure stable else: stable_time = rt['transitsource'][-2, 0] + last_transit2stable stable_time_jul = datetime.utcfromtimestamp(stable_time) stable_time_jul = ephem.julian_date(stable_time_jul) # convert to julian date, convenient to display if not os.path.exists(output_path(normal_gain_file)): if mpiutil.rank0: print('Normalization gain file has not been built yet. Try to build it!') print('Recorded transit Time: %s'%aipy.phs.juldate2ephem(agfilein.attrs['transit_jul'] + 8./24.)) print('Last transit Time: %s'%aipy.phs.juldate2ephem(ephem.julian_date(datetime.utcfromtimestamp(rt['transitsource'][-2,0])) + 8./24.)) print('First time point of this data: %s'%aipy.phs.juldate2ephem(rt['jul_date'][0] + 8./24.)) build_normal_file = True # if rt['jul_date'][0] < stable_time: if rt.attrs['sec1970'][0] < stable_time: raise BeforeStableTime('The beginning time point is %s, but only after %s, will the noise source be stable. Abort the noise calibration!'%(aipy.phs.juldate2ephem(rt['jul_date'][0] + 8./24.), aipy.phs.juldate2ephem(stable_time_jul + 8./24.))) else: if mpiutil.rank0: print('Use existing normalization gain file!') build_normal_file = False with h5py.File(output_path(normal_gain_file), 'r') as filein: rt.normal_phs = filein['phs'][:, local_bl_start:local_bl_end] if not phs_only: rt.normal_amp = filein['amp'][:, local_bl_start:local_bl_end] # apply absolute gain absgain = agfilein['gain'][:] polfeed, _ = bl2pol_feed_inds(rt.local_bl, agfilein['gain'].attrs['feed'][:], agfilein['gain'].attrs['pol'][:]) for ii, (ipol, ifeed, jpol, jfeed) in enumerate(polfeed): # rt.local_vis[:,:,ii] /= (absgain[None,:,ipol,ifeed-1] * absgain[None,:,jpol,jfeed-1].conj()) rt.local_vis[:,:,ii] /= (absgain[None,:,ipol,ifeed] * absgain[None,:,jpol,jfeed].conj()) #=================================== nt = rt.local_vis.shape[0] if num_mean <= 0: raise RuntimeError('Invalid num_mean = %s' % num_mean) ns_on = rt['ns_on'][:] ns_on = np.where(ns_on, 1, 0) diff_ns = np.diff(ns_on) inds = np.where(diff_ns==1)[0] # NOTE: these are inds just 1 before the first ON if not rt.FRB_cal: # for FRB there might be just one noise point, avoid waste if inds[0]-1 < 0: # no off data in the beginning to use inds = inds[1:] if inds[-1]+2 > nt-1: # no on data in the end to use inds = inds[:-1] if save_gain: num_inds = len(inds) shp = (num_inds,)+rt.local_vis.shape[1:] dtype = rt.local_vis.real.dtype # create dataset to record ns_cal_time_inds rt.create_time_ordered_dataset('ns_cal_time_inds', inds) # create dataset to record ns_cal_phase ns_cal_phase = np.empty(shp, dtype=dtype) ns_cal_phase[:] = np.nan ns_cal_phase = mpiarray.MPIArray.wrap(ns_cal_phase, axis=2, comm=rt.comm) rt.create_freq_and_bl_ordered_dataset('ns_cal_phase', ns_cal_phase, axis_order=(None, 1, 2)) rt['ns_cal_phase'].attrs['unit'] = 'radians' if not phs_only: # create dataset to record ns_cal_amp ns_cal_amp = np.empty(shp, dtype=dtype) ns_cal_amp[:] = np.nan ns_cal_amp = mpiarray.MPIArray.wrap(ns_cal_amp, axis=2, comm=rt.comm) rt.create_freq_and_bl_ordered_dataset('ns_cal_amp', ns_cal_amp, axis_order=(None, 1, 2)) if bl_incl == 'all': bls_plt = [ tuple(bl) for bl in rt.bl ] else: bls_plt = [ bl for bl in bl_incl if not bl in bl_excl ] if freq_incl == 'all': freq_plt = range(rt.freq.shape[0]) else: freq_plt = [ fi for fi in freq_incl if not fi in freq_excl ] show_progress = self.params['show_progress'] progress_step = self.params['progress_step'] if rt.ps_first: rt.freq_and_bl_data_operate(self.cal, full_data=True, show_progress=show_progress, progress_step=progress_step, keep_dist_axis=False, num_mean=num_mean, inds=inds, bls_plt=bls_plt, freq_plt=freq_plt, build_normal_file = build_normal_file) else: rt.freq_and_bl_data_operate(self.cal, full_data=True, show_progress=show_progress, progress_step=progress_step, keep_dist_axis=False, num_mean=num_mean, inds=inds, bls_plt=bls_plt, freq_plt=freq_plt) if save_gain: interp_mask_ratio = mpiutil.allreduce(np.sum(rt.interp_mask_count))/1./mpiutil.allreduce(np.size(rt.interp_mask_count)) * 100. if interp_mask_ratio > 50. and rt.normal_index is not None: warnings.warn('%.1f%% of the data was masked due to shortage of noise points for interpolation(need at least 4 to perform cubic spline)! The pointsource calibration may not be done due to too many masked points!'%interp_mask_ratio, NotEnoughPointToInterpolateWarning) if interp_mask_ratio > 80.: rt.interp_all_masked = True # gather bl_order to rank0 bl_order = mpiutil.gather_array(rt['blorder'].local_data, axis=0, root=0, comm=rt.comm) # gather ns_cal_phase / ns_cal_amp to rank 0 ns_cal_phase = mpiutil.gather_array(rt['ns_cal_phase'].local_data, axis=2, root=0, comm=rt.comm) phs_unit = rt['ns_cal_phase'].attrs['unit'] rt.delete_a_dataset('ns_cal_phase', reserve_hint=False) if not phs_only: ns_cal_amp = mpiutil.gather_array(rt['ns_cal_amp'].local_data, axis=2, root=0, comm=rt.comm) rt.delete_a_dataset('ns_cal_amp', reserve_hint=False) if tag_output_iter: gain_file = output_path(gain_file, iteration=self.iteration) else: gain_file = output_path(gain_file) if rt.ps_first: phase_finite_count = mpiutil.allreduce(np.isfinite(rt.normal_phs).sum()) if not phs_only: amp_finite_count = mpiutil.allreduce(np.isfinite(rt.normal_amp).sum()) if mpiutil.rank0: if rt.ps_first and build_normal_file: if phase_finite_count == 0: raise AllMasked('All values are masked when calculating phase!') if not phs_only: if amp_finite_count == 0: raise AllMasked('All values are masked when calculating amplitude!') with h5py.File(output_path(normal_gain_file), 'w') as tapfilein: tapfilein.create_dataset('amp',(rt['vis'].shape[1], rt['vis'].shape[2])) tapfilein.create_dataset('phs',(rt['vis'].shape[1], rt['vis'].shape[2])) with h5py.File(gain_file, 'w') as f: # save time f.create_dataset('time', data=rt['jul_date'][:]) f['time'].attrs['unit'] = 'Julian date' # save freq f.create_dataset('freq', data=rt['freq'][:]) f['freq'].attrs['unit'] = rt['freq'].attrs['unit'] # save bl f.create_dataset('bl_order', data=bl_order) # save ns_cal_time_inds f.create_dataset('ns_cal_time_inds', data=rt['ns_cal_time_inds'][:]) # save ns_cal_phase f.create_dataset('ns_cal_phase', data=ns_cal_phase) f['ns_cal_phase'].attrs['unit'] = phs_unit f['ns_cal_phase'].attrs['dim'] = '(time, freq, bl)' if not phs_only: # save ns_cal_amp f.create_dataset('ns_cal_amp', data=ns_cal_amp) # save channo f.create_dataset('channo', data=rt['channo'][:]) f['channo'].attrs['dim'] = rt['channo'].attrs['dimname'] if rt.exclude_bad: f['channo'].attrs['badchn'] = rt['channo'].attrs['badchn'] if not (absolute_gain_filename is None): if not os.path.exists(output_path(absolute_gain_filename)): raise NoPsGainFile('No absolute gain file %s, do the ps calibration first!'%output_path(absolute_gain_filename)) with h5py.File(output_path(absolute_gain_filename,'r')) as abs_gain: new_gain = uni_gain(abs_gain, f, phs_only = phs_only) f.create_dataset('uni_gain', data = new_gain) f['uni_gain'].attrs['dim'] = '(time, freq, bl)' if rt.ps_first and build_normal_file: if mpiutil.rank0: print('Start write normalization gain into %s'%output_path(normal_gain_file)) for i in range(10): try: for ii in range(mpiutil.size): if ii == mpiutil.rank: with h5py.File(output_path(normal_gain_file), 'r+') as fileout: fileout['phs'][:,local_bl_start:local_bl_end] = rt.normal_phs[:,:] if not phs_only: fileout['amp'][:,local_bl_start:local_bl_end] = rt.normal_amp[:,:] mpiutil.barrier() break except IOError: time.sleep(0.5) continue rt.delete_a_dataset('ns_cal_time_inds', reserve_hint=False) return super(NsCal, self).process(rt)
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()
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 _generate_phase(self, time): ntime = len(time) freq = self.freq nfreq = len(freq) # Generate the correlation function cf_delay = self._corr_func(self.corr_length_delay, self.sigma_delay) # Check if we are simulating relative delays or common mode delays if self.sim_type == "relative": n_realisations = self.ninput_local # Generate delay fluctuations self.delay_error = gain.generate_fluctuations( time, cf_delay, n_realisations, self._prev_time, self._prev_delay) gain_phase = (2.0 * np.pi * freq[:, np.newaxis, np.newaxis] * 1e6 * self.delay_error[np.newaxis, :, :] / np.sqrt(self.ndays)) if self.sim_type == "common_mode_cyl": n_realisations = 1 ninput = self.ninput_global # Generates as many random delay errors as there are cylinders if self.comm.rank == 0: if self.common_mode_type == "sinusoidal": P1 = self.sinusoidal_period[0] P2 = self.sinusoidal_period[1] omega1 = 2 * np.pi / P1 omega2 = 2 * np.pi / P2 delay_error = (self.sigma_delay * (np.sin(omega1 * time) - np.sin(omega2 * time))[np.newaxis, :]) if self.common_mode_type == "random": delay_error = gain.generate_fluctuations( time, cf_delay, n_realisations, self._prev_time, self._prev_delay, ) else: delay_error = None # Broadcast to other ranks self.delay_error = self.comm.bcast(delay_error, root=0) # Split frequencies to processes. lfreq, sfreq, efreq = mpiutil.split_local(nfreq) # Create an array to hold all inputs, which are common-mode within # a cylinder gain_phase = np.zeros((lfreq, ninput, ntime), dtype=complex) # Since we have 2 cylinders populate half of them with a delay) # TODO: generalize this for 3 or even 4 cylinders in the future. gain_phase[:, ninput // self.ncyl:, :] = ( 2.0 * np.pi * freq[sfreq:efreq, np.newaxis, np.newaxis] * 1e6 * self.delay_error[np.newaxis, :, :] / np.sqrt(self.ndays)) gain_phase = mpiarray.MPIArray.wrap(gain_phase, axis=0, comm=self.comm) # Redistribute over input to match rest of the code gain_phase = gain_phase.redistribute(axis=1) gain_phase = gain_phase.view(np.ndarray) if self.sim_type == "common_mode_iceboard": nchannel = self.nchannel ninput = self.ninput_global # Number of channels on a board nboards = ninput // nchannel # Generates as many random delay errors as there are iceboards if self.comm.rank == 0: delay_error = gain.generate_fluctuations( time, cf_delay, nboards, self._prev_time, self._prev_delay) else: delay_error = None # Broadcast to other ranks self.delay_error = self.comm.bcast(delay_error, root=0) # Calculate the corresponding phase by multiplying with frequencies phase = (2.0 * np.pi * freq[:, np.newaxis, np.newaxis] * 1e6 * self.delay_error[np.newaxis, :] / np.sqrt(self.ndays)) # Create an array to hold all inputs, which are common-mode within # one iceboard gain_phase = mpiarray.MPIArray((nfreq, ninput, ntime), axis=1, dtype=np.complex128, comm=self.comm) gain_phase[:] = 0.0 # Loop over inputs and and group common-mode phases on every board for il, ig in gain_phase.enumerate(axis=1): # Get the board number bi bi = int(ig / nchannel) gain_phase[:, il] = phase[:, bi] gain_phase = gain_phase.view(np.ndarray) self._prev_delay = self.delay_error self._prev_time = time return gain_phase
def wrap(cls, array, axis, comm=None): """Turn a set of numpy arrays into a distributed MPIArray object. This is needed for functions such as `np.fft.fft` which always return an `np.ndarray`. Parameters ---------- array : np.ndarray Array to wrap. axis : integer Axis over which the array is distributed. The lengths are checked to try and ensure this is correct. comm : MPI.Comm, optional The communicator over which the array is distributed. If `None` (default), use `MPI.COMM_WORLD`. Returns ------- dist_array : MPIArray An MPIArray view of the input. """ # from mpi4py import MPI if comm is None: comm = mpiutil.world # Get axis length, both locally, and globally axlen = array.shape[axis] totallen = mpiutil.allreduce(axlen, comm=comm) # Figure out what the distributed layout should be local_num, local_start, local_end = mpiutil.split_local(totallen, comm=comm) # Check the local layout is consistent with what we expect, and send # result to all ranks layout_issue = mpiutil.allreduce(axlen != local_num, op=mpiutil.MAX, comm=comm) if layout_issue: raise Exception( "Cannot wrap, distributed axis local length is incorrect.") # Set shape and offset lshape = array.shape global_shape = list(lshape) global_shape[axis] = totallen loffset = [0] * len(lshape) loffset[axis] = local_start # Setup attributes of class dist_arr = array.view(cls) dist_arr._global_shape = tuple(global_shape) dist_arr._axis = axis dist_arr._local_shape = tuple(lshape) dist_arr._local_offset = tuple(loffset) dist_arr._comm = comm return dist_arr
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.rank == 0: import warnings warnings.warn( "Cannot redistribute array of compound datatypes." " Sorry!!") return self # Get a view of the array arr = self.view(np.ndarray) if self.comm.size == 1: # only one process if arr.shape[self.axis] == self.global_shape[self.axis]: # We are working on a single node and being asked to do # a trivial transpose. trans_arr = arr.copy() else: raise ValueError( "Global shape %s is incompatible with local arrays shape %s" % (self.global_shape, self.shape)) else: pc, sc, ec = mpiutil.split_local(arr.shape[axis], comm=self.comm) par, sar, ear = mpiutil.split_all(self.global_shape[self.axis], comm=self.comm) pac, sac, eac = mpiutil.split_all(arr.shape[axis], comm=self.comm) new_shape = np.asarray(self.global_shape) new_shape[axis] = pc requests_send = [] requests_recv = [] trans_arr = np.empty(new_shape, dtype=arr.dtype) mpitype = mpiutil.typemap(arr.dtype) buffers = list() # Cut out the right blocks of the local array to send around blocks = np.array_split(arr, np.insert(eac, 0, sac[0]), axis)[1:] # Iterate over all processes row wise for ir in range(self.comm.size): # Iterate over all processes column wise for ic in range(self.comm.size): # Construct a unique tag tag = ir * self.comm.size + ic # Send and receive the messages as non-blocking passes if self.comm.rank == ir: # Send the message request = self.comm.Isend( [blocks[ic].flatten(), mpitype], dest=ic, tag=tag) requests_send.append([ir, ic, request]) if self.comm.rank == ic: buffer_shape = np.asarray(new_shape) buffer_shape[axis] = eac[ic] - sac[ic] buffer_shape[self.axis] = ear[ir] - sar[ir] buffers.append( np.ndarray(buffer_shape, dtype=arr.dtype)) request = self.comm.Irecv([buffers[ir], mpitype], source=ir, tag=tag) requests_recv.append([ir, ic, request]) # Wait for all processes to have started their messages self.comm.Barrier() # For each node iterate over all sends and wait until completion for ir, ic, request in requests_send: stat = mpiutil.MPI.Status() request.Wait(status=stat) if stat.error != mpiutil.MPI.SUCCESS: print( "**** ERROR in MPI SEND (r: %i c: %i rank: %i) *****". format(ir, ic, self.comm.rank)) self.comm.Barrier() # For each frequency iterate over all receives and wait until # completion for ir, ic, request in requests_recv: stat = mpiutil.MPI.Status() request.Wait(status=stat) if stat.error != mpiutil.MPI.SUCCESS: print( "**** ERROR in MPI RECV (r: %i c: %i rank: %i) *****". format(ir, ir, self.comm.rank)) # Put together the blocks we received np.concatenate(buffers, self.axis, trans_arr) # 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 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.rank == 0: import warnings warnings.warn('Cannot redistribute array of compound datatypes.' ' Sorry!!') return self # Get a view of the array arr = self.view(np.ndarray) if self.comm.size == 1: # only one process if arr.shape[self.axis] == self.global_shape[self.axis]: # We are working on a single node and being asked to do # a trivial transpose. trans_arr = arr.copy() else: raise ValueError( 'Global shape %s is incompatible with local arrays shape %s' % (self.global_shape, self.shape)) else: pc, sc, ec = mpiutil.split_local(arr.shape[axis], comm=self.comm) par, sar, ear = mpiutil.split_all(self.global_shape[self.axis], comm=self.comm) pac, sac, eac = mpiutil.split_all(arr.shape[axis], comm=self.comm) new_shape = np.asarray(self.global_shape) new_shape[axis] = pc requests_send = [] requests_recv = [] trans_arr = np.empty(new_shape, dtype=arr.dtype) mpitype = mpiutil.typemap(arr.dtype) buffers = list() # Cut out the right blocks of the local array to send around blocks = np.array_split(arr, np.insert(eac, 0, sac[0]), axis)[1:] # Iterate over all processes row wise for ir in range(self.comm.size): # Iterate over all processes column wise for ic in range(self.comm.size): # Construct a unique tag tag = ir * self.comm.size + ic # Send and receive the messages as non-blocking passes if self.comm.rank == ir: # Send the message request = self.comm.Isend([blocks[ic].flatten(), mpitype], dest=ic, tag=tag) requests_send.append([ir, ic, request]) if self.comm.rank == ic: buffer_shape = np.asarray(new_shape) buffer_shape[axis] = eac[ic] - sac[ic] buffer_shape[self.axis] = ear[ir] - sar[ir] buffers.append(np.ndarray(buffer_shape, dtype=arr.dtype)) request = self.comm.Irecv([buffers[ir], mpitype], source=ir, tag=tag) requests_recv.append([ir, ic, request]) # Wait for all processes to have started their messages self.comm.Barrier() # For each node iterate over all sends and wait until completion for ir, ic, request in requests_send: stat = mpiutil.MPI.Status() request.Wait(status=stat) if stat.error != mpiutil.MPI.SUCCESS: print("**** ERROR in MPI SEND (r: %i c: %i rank: %i) *****" .format(ir, ic, self.comm.rank)) self.comm.Barrier() # For each frequency iterate over all receives and wait until # completion for ir, ic, request in requests_recv: stat = mpiutil.MPI.Status() request.Wait(status=stat) if stat.error != mpiutil.MPI.SUCCESS: print("**** ERROR in MPI RECV (r: %i c: %i rank: %i) *****" .format(ir, ir, self.comm.rank)) # Put together the blocks we received np.concatenate(buffers, self.axis, trans_arr) # 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
import matplotlib matplotlib.use('Agg') if mpiutil.size != 20: raise RuntimeError('Need 20 processes') np.random.seed(0) rands = np.random.rand(1000, 20) N = 1000 m = 20 # number of augmentation by rotation n = 256 dataset = np.zeros((N / 20, m, n, n, 1)) # for i in range(N): ln, s, e = mpiutil.split_local(N) for li, i in enumerate(range(s, e)): if mpiutil.rank0: print '%d of %d ...' % (li, ln) # read in reconstructed map with h5py.File('./output_sim_750MHz/map/ts/ts_%04d/map_full.hdf5' % i, 'r') as f: rec_map = f['map'][0, 0] for j in range(m): if j == 0: rot_map = rec_map.copy() else: # angle = 360 * np.random.rand() angle = 360 * rands[i, j] rot_map = rotate.np_rotate(rec_map.copy(), angle)
print 'rank %d has %s after gather_list' % (rank, lst) # parallel_map separator(sec, 'parallel_map') glist = range(6) result = mpiutil.parallel_map(lambda x: x * x, glist, root=0) if rank == 0: print 'result = %s' % result # split_all separator(sec, 'split_all') print 'rank %d has: %s' % (rank, mpiutil.split_all(6)) # split_local separator(sec, 'split_local') print 'rank %d has: %s' % (rank, mpiutil.split_local(6)) # gather_array separator(sec, 'gather_array') if rank == 0: local_ary = np.array([[0, 1], [6, 7]]) elif rank == 1: local_ary = np.array([[2], [8]]) elif rank == 2: local_ary = np.array([[3], [9]]) if rank == 3: local_ary = np.array([[4, 5], [10, 11]]) global_ary = mpiutil.gather_array(local_ary, axis=1, root=0) if rank == 0: print 'global_ary = %s' % global_ary
def test_io(self): import h5py # Cleanup directories fname = "testdset.hdf5" if mpiutil.rank0 and os.path.exists(fname): os.remove(fname) mpiutil.barrier() gshape = (19, 17) ds = mpiarray.MPIArray(gshape, dtype=np.int64) ga = np.arange(np.prod(gshape)).reshape(gshape) l0, s0, e0 = mpiutil.split_local(gshape[0]) ds[:] = ga[s0:e0] ds.redistribute(axis=1).to_hdf5(fname, "testds", create=True) if mpiutil.rank0: with h5py.File(fname, "r") as f: h5ds = f["testds"][:] assert (h5ds == ga).all() ds2 = mpiarray.MPIArray.from_hdf5(fname, "testds") assert (ds2 == ds).all() mpiutil.barrier() # Check that reading over another distributed axis works ds3 = mpiarray.MPIArray.from_hdf5(fname, "testds", axis=1) assert ds3.shape[0] == gshape[0] assert ds3.shape[1] == mpiutil.split_local(gshape[1])[0] ds3 = ds3.redistribute(axis=0) assert (ds3 == ds).all() mpiutil.barrier() # Check a read with an arbitrary slice in there. This only checks the shape is correct. ds4 = mpiarray.MPIArray.from_hdf5(fname, "testds", axis=1, sel=(np.s_[3:10:2], np.s_[1:16:3])) assert ds4.shape[0] == 4 assert ds4.shape[1] == mpiutil.split_local(5)[0] mpiutil.barrier() # Check the read with a slice along the axis being read ds5 = mpiarray.MPIArray.from_hdf5(fname, "testds", axis=1, sel=(np.s_[:], np.s_[3:15:2])) assert ds5.shape[0] == gshape[0] assert ds5.shape[1] == mpiutil.split_local(6)[0] ds5 = ds5.redistribute(axis=0) assert (ds5 == ds[:, 3:15:2]).all() mpiutil.barrier() # Check the read with a slice along the axis being read ds6 = mpiarray.MPIArray.from_hdf5(fname, "testds", axis=0, sel=(np.s_[:], np.s_[3:15:2])) ds6 = ds6.redistribute(axis=0) assert (ds6 == ds[:, 3:15:2]).all() mpiutil.barrier() if mpiutil.rank0 and os.path.exists(fname): os.remove(fname)