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 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 = list(range(self.telescope.mmax + 1)) mpart = mpiutil.partition_list_mpi(mlist) # Total number of sky modes. nmodes = self.beamtransfer.nfreq * self.beamtransfer.ntel # If sky is alm fine, if not perform spherical harmonic transform. alm = sky if harmonic else hputil.sphtrans_sky(sky, lmax=self.telescope.lmax) ## Routine to project sky onto eigenmodes def _proj(mi): p1 = self.project_sky_vector_forward(mi, alm[:, :, mi], threshold) p2 = np.zeros(nmodes, dtype=np.complex128) p2[-p1.size :] = p1 return p2 # Map over list of m's and project sky onto eigenbasis proj_sec = [(mi, _proj(mi)) for mi in mpart] # Gather projections onto the rank=0 node. proj_all = mpiutil.world.gather(proj_sec, root=0) proj_arr = None if mpiutil.rank0: # Create array to put projections into proj_arr = np.zeros( (2 * self.telescope.mmax + 1, nmodes), dtype=np.complex128 ) # Iterate over all gathered projections and insert into the array for proc_rank in proj_all: for pm in proc_rank: proj_arr[pm[0]] = pm[1] # Return the projections (rank=0) or None elsewhere. return proj_arr
def project_sky(self, sky, mlist = None, threshold=None, harmonic=False): # Set default list of m-modes (i.e. all of them), and partition if mlist is None: mlist = range(self.telescope.mmax + 1) mpart = mpiutil.partition_list_mpi(mlist) # Total number of sky modes. nmodes = self.beamtransfer.nfreq * self.beamtransfer.ntel # If sky is alm fine, if not perform spherical harmonic transform. alm = sky if harmonic else hputil.sphtrans_sky(sky, lmax=self.telescope.lmax) ## Routine to project sky onto eigenmodes def _proj(mi): p1 = self.project_sky_vector_forward(mi, alm[:, :, mi], threshold) p2 = np.zeros(nmodes, dtype=np.complex128) p2[-p1.size:] = p1 return p2 # Map over list of m's and project sky onto eigenbasis proj_sec = [(mi, _proj(mi)) for mi in mpart] # Gather projections onto the rank=0 node. proj_all = mpiutil.world.gather(proj_sec, root=0) proj_arr = None if mpiutil.rank0: # Create array to put projections into proj_arr = np.zeros((2*self.telescope.mmax + 1, nmodes), dtype=np.complex128) # Iterate over all gathered projections and insert into the array for proc_rank in proj_all: for pm in proc_rank: proj_arr[pm[0]] = pm[1] # Return the projections (rank=0) or None elsewhere. return proj_arr
def generate(self, regen=False): """Calculate the total Fisher matrix and bias and save to a file. Parameters ---------- regen : boolean, optional Force regeneration if products already exist (default `False`). """ if mpiutil.rank0: st = time.time() print("======== Starting PS calculation ========") ffile = self.psdir + "/fisher.hdf5" if os.path.exists(ffile) and not regen: print("Fisher matrix file: %s exists. Skipping..." % ffile) return mpiutil.barrier() # Pre-compute all the angular power spectra for the bands self.genbands() # Calculate Fisher and bias for each m # Pair up each list item with its position. zlist = list(enumerate(range(self.telescope.mmax + 1))) # Partition list based on MPI rank llist = mpiutil.partition_list_mpi(zlist) # Operate on sublist fisher_bias_list = [self.fisher_bias_m(item) for ind, item in llist] # Unpack into separate lists of the Fisher matrix and bias fisher_loc, bias_loc = zip(*fisher_bias_list) # Sum over all local m-modes to get the over all Fisher and bias pe process fisher_loc = np.sum(np.array(fisher_loc), axis=0).real # Be careful of the .real here bias_loc = np.sum(np.array(bias_loc), axis=0).real # Be careful of the .real here self.fisher = mpiutil.allreduce(fisher_loc, op=MPI.SUM) self.bias = mpiutil.allreduce(bias_loc, op=MPI.SUM) # Write out all the PS estimation products if mpiutil.rank0: et = time.time() print("======== Ending PS calculation (time=%f) ========" % (et - st)) # Check to see ensure that Fisher matrix isn't all zeros. if not (self.fisher == 0).all(): # Generate derived quantities (covariance, errors..) cv = la.pinv(self.fisher, rcond=1e-8) err = cv.diagonal()**0.5 cr = cv / np.outer(err, err) else: cv = np.zeros_like(self.fisher) err = cv.diagonal() cr = np.zeros_like(self.fisher) f = h5py.File(self.psdir + "/fisher.hdf5", "w") f.attrs["bandtype"] = np.string_( self.bandtype) # HDF5 string issues f.create_dataset("fisher/", data=self.fisher) f.create_dataset("bias/", data=self.bias) f.create_dataset("covariance/", data=cv) f.create_dataset("errors/", data=err) f.create_dataset("correlation/", data=cr) f.create_dataset("band_power/", data=self.band_power) if self.bandtype == "polar": f.create_dataset("k_start/", data=self.k_start) f.create_dataset("k_end/", data=self.k_end) f.create_dataset("k_center/", data=self.k_center) f.create_dataset("theta_start/", data=self.theta_start) f.create_dataset("theta_end/", data=self.theta_end) f.create_dataset("theta_center/", data=self.theta_center) f.create_dataset("k_bands", data=self.k_bands) f.create_dataset("theta_bands", data=self.theta_bands) elif self.bandtype == "cartesian": f.create_dataset("kpar_start/", data=self.kpar_start) f.create_dataset("kpar_end/", data=self.kpar_end) f.create_dataset("kpar_center/", data=self.kpar_center) f.create_dataset("kperp_start/", data=self.kperp_start) f.create_dataset("kperp_end/", data=self.kperp_end) f.create_dataset("kperp_center/", data=self.kperp_center) f.create_dataset("kpar_bands", data=self.kpar_bands) f.create_dataset("kperp_bands", data=self.kperp_bands) f.close()