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