def collect_m_arrays(mlist, func, shapes, dtype): data = [ (mi, func(mi)) for mi in mpiutil.partition_list_mpi(mlist) ] mpiutil.barrier() if mpiutil.rank0 and mpiutil.size == 1: p_all = [data] else: p_all = mpiutil.world.gather(data, root=0) mpiutil.barrier() # Not sure if this barrier really does anything, # but hoping to stop collect breaking marrays = None if mpiutil.rank0: marrays = [np.zeros((len(mlist),) + shape, dtype=dtype) for shape in shapes] for p_process in p_all: for mi, result in p_process: for si in range(len(shapes)): if result[si] is not None: marrays[si][mi] = result[si] mpiutil.barrier() return marrays
def project_sky(self, sky, mlist = None, threshold=None, harmonic=False): # Set default list of m-modes (i.e. all of them), and partition if mlist is None: mlist = range(self.telescope.mmax + 1) mpart = mpiutil.partition_list_mpi(mlist) # Total number of sky modes. nmodes = self.beamtransfer.nfreq * self.beamtransfer.ntel # If sky is alm fine, if not perform spherical harmonic transform. alm = sky if harmonic else hputil.sphtrans_sky(sky, lmax=self.telescope.lmax) ## Routine to project sky onto eigenmodes def _proj(mi): p1 = self.project_sky_vector_forward(mi, alm[:, :, mi], threshold) p2 = np.zeros(nmodes, dtype=np.complex128) p2[-p1.size:] = p1 return p2 # Map over list of m's and project sky onto eigenbasis proj_sec = [(mi, _proj(mi)) for mi in mpart] # Gather projections onto the rank=0 node. proj_all = mpiutil.world.gather(proj_sec, root=0) proj_arr = None if mpiutil.rank0: # Create array to put projections into proj_arr = np.zeros((2*self.telescope.mmax + 1, nmodes), dtype=np.complex128) # Iterate over all gathered projections and insert into the array for proc_rank in proj_all: for pm in proc_rank: proj_arr[pm[0]] = pm[1] # Return the projections (rank=0) or None elsewhere. return proj_arr
sky_vec = bt.project_vector_forward(mi, alm[:, :, mi]) tau_vec = np.fft.fft(blackman_harris(cyl.nfreq)[:, np.newaxis] * sky_vec, axis=0) tau_vec = np.fft.ifft(tau_vec * blmask, axis=0) alm2 = bt.project_vector_backward(mi, tau_vec) return alm2 # Project m-modes across different processes mlist = list(range(mmax + 1)) mpart = mpiutil.partition_list_mpi(mlist) mproj = [[mi, projm(mi)] for mi in mpart] if mpiutil.rank0: print("Gather results onto root process") p_all = mpiutil.world.gather(mproj, root=0) # Save out results if mpiutil.rank0: palm = np.zeros_like(alm) print("Combining results.") for p_process in p_all: for mi, proj in p_process:
mvals, mvecs = klt.modes_m(mi, threshold=cut) if mvals is None: return None ev_vec = klt.project_sky_vector_forward(mi, alm[:, :, mi], threshold=cut) tel_vec = klt.project_tel_vector_backward(mi, ev_vec, threshold=cut) alm2 = bt.project_vector_backward(mi, tel_vec) return alm2 # Project m-modes across different processes mlist = range(mmax+1) mpart = mpiutil.partition_list_mpi(mlist) mproj = [[mi, projm(mi)] for mi in mpart] if mpiutil.rank0: print "Gather results onto root process" p_all = mpiutil.world.gather(mproj, root=0) # Save out results if mpiutil.rank0: palm = np.zeros_like(alm) print "Combining results." for p_process in p_all: