def compute_linear_term(self, alm, no_icov=False): ''' Return linear term at current iteration for input data alm. Parameters ---------- alm : (npol, nelem) complex array Spherical harmonic coeffients of data. no_icov : bool, optional Do not icov filter input (i.e. input is already filtered). Returns ------- lin_term : float, None Linear term of the estimator. Notes ----- Linear term is defined as sum(a C^-1 grad T[C^-1 a]). See Eq. 57 in Smith & Zaldarriaga. ''' if self.mc_gt is None: return None alm = utils.alm_return_2d(alm, self.npol, self.lmax) if no_icov: return utils.contract_almxblm(alm, np.conj(self.mc_gt)) else: return utils.contract_almxblm(self.icov(alm), np.conj(self.mc_gt))
def test_utils_contract_almxblm(self): alm = np.asarray([1, 1, 1, 1, 2j, 2j, 2j, 3j, 3j, 4j]) blm = np.asarray([2, 2, 2, 2, 3, 3, 3, 4j, 4j, 5]) ans_exp = -40 ans = utils.contract_almxblm(alm, blm) self.assertEqual(ans, ans_exp)
def test_utils_contract_almxblm_3d(self): alm = np.ones((2, 3, 10), dtype=np.complex128) blm = np.ones((2, 3, 10), dtype=np.complex128) alm *= np.asarray([1, 1, 1, 1, 2j, 2j, 2j, 3j, 3j, 4j]) blm *= np.asarray([2, 2, 2, 2, 3, 3, 3, 4j, 4j, 5]) ans_exp = -240 ans = utils.contract_almxblm(alm, blm) self.assertEqual(ans, ans_exp)
def test_utils_contract_almxblm_cl(self): # Check if contraction matches hp.alm2cl. alm = np.ones(10, dtype=np.complex128) alm += 1j * np.ones(10, dtype=np.complex128) alm[:4] = 1 lmax = 3 ells = np.asarray([0, 1, 2, 3]) cl = hp.alm2cl(alm) ans_exp = np.sum(cl * (2 * ells + 1)) ans = utils.contract_almxblm(alm, np.conj(alm)) self.assertEqual(ans, ans_exp)
def compute_fisher(self): ''' Return Fisher information at current iteration. Returns ------- fisher : float, None Fisher information. ''' if self.mc_gt_sq is None or self.mc_gt is None: return None fisher = self.mc_gt_sq fisher -= utils.contract_almxblm( self.mc_gt, self.icov(self.beam(np.conj(self.mc_gt)))) fisher /= 3. return fisher
def step(self, alm, theta_batch=25): ''' Add iteration to <grad T (C^-1 a) C^-1 grad T(C^-1 a)^*> and <grad T (C^-1 a)> Monte Carlo estimates. Parameters ---------- alm : (nelem) or (npol, nelem) complex array Healpix-ordered unfiltered alm array. Will be overwritten! theta_batch : int, optional Process loop over theta in batches of this size. Higher values take up more memory. Raises ------ ValueError If shape input alm is not understood. ''' grad_t = self._step(alm, theta_batch=theta_batch) # Add to Monte Carlo estimates. if self.mc_gt is None: self.mc_gt = grad_t else: self.__mc_gt += grad_t mc_gt_sq = utils.contract_almxblm( grad_t, self.icov(self.beam(np.conj(grad_t)))) if self.mc_gt_sq is None: self.mc_gt_sq = mc_gt_sq else: self.__mc_gt_sq += mc_gt_sq self.mc_idx += 1
def step_batch(self, alm_loader, alm_files, comm=None, verbose=False, **kwargs): ''' Add iterations to <grad T (C^-1 a) C^-1 grad T(C^-1 a)^*> and <grad T (C^-1 a)> Monte Carlo estimates by loading and processing several alms in parallel using MPI. Arguments --------- alm_loader : callable Function that returns alms on rank given filename as first argument. alm_files : array_like List of alm files to load. comm : MPI communicator, optional verbose : bool, optional Print process. kwargs : dict, optional Optional keyword arguments passed to "_step". ''' if comm is None: comm = utils.FakeMPIComm() # Monte carlo quantities local to rank. mc_idx_loc = 0 mc_gt_sq_loc = None mc_gt_loc = None # Split alm_file loop over ranks for alm_file in alm_files[comm.Get_rank():len(alm_files):comm.Get_size( )]: if verbose: print('rank {:3}: loading {}'.format(comm.Get_rank(), alm_file)) alm = alm_loader(alm_file) if verbose: print('rank {:3}: done loading'.format(comm.Get_rank())) grad_t = self._step(alm, **kwargs) if mc_gt_loc is None: mc_gt_loc = grad_t else: mc_gt_loc += grad_t mc_gt_sq = utils.contract_almxblm( grad_t, self.icov(self.beam(np.conj(grad_t)))) if mc_gt_sq_loc is None: mc_gt_sq_loc = mc_gt_sq else: mc_gt_sq_loc += mc_gt_sq mc_idx_loc += 1 # To allow allreduce when number of ranks > alm files. shape, dtype = utils.bcast_array_meta(mc_gt_loc, comm, root=0) if mc_gt_loc is None: mc_gt_loc = np.zeros(shape, dtype=dtype) if mc_gt_sq_loc is None: mc_gt_sq_loc = 0. if mc_idx_loc is None: mc_idx_loc = 0 mc_gt = utils.allreduce_array(mc_gt_loc, comm) mc_gt_sq = utils.allreduce(mc_gt_sq_loc, comm) mc_idx = utils.allreduce(mc_idx_loc, comm) # All ranks get to update the internal mc variables themselves. if self.mc_gt is None: self.mc_gt = mc_gt else: self.__mc_gt += mc_gt if self.mc_gt_sq is None: self.mc_gt_sq = mc_gt_sq else: self.__mc_gt_sq += mc_gt_sq self.mc_idx += mc_idx