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()
