def prepare(self):
if self.ready: return
if self.auto_grow:
# Promote everything to full-size maps.
all_geoms = self.comm.gather((self.shape, self.wcs), root=0)
if self.comm.rank == 0:
new_shape, new_wcs = coords.get_supergeom(*all_geoms)
else:
new_shape, new_wcs = None, None
new_shape, new_wcs = self.comm.bcast((new_shape, new_wcs), root=0)
for attr in ['map_rhs', 'map_div']:
submap = getattr(self, attr)
lead_shape = submap.shape[:-2]
newmap = enmap.zeros(lead_shape + new_shape,
dtype=submap.dtype,
wcs=new_wcs)
newmap.insert(submap)
setattr(self, attr, newmap)
del submap
self.shape, self.wcs = new_shape, new_wcs
print("rank %3d ntod %2d" % (self.comm.rank, len(self.junk_rhs)))
self.comm.Barrier()
self.map_rhs = utils.allreduce(self.map_rhs, self.comm)
self.map_div = utils.allreduce(self.map_div, self.comm)
self.map_idiv = safe_invert_div(self.map_div)
self.junk_offs = utils.cumsum([r.size for r in self.junk_rhs],
endpoint=True)
self.junk_rhs = np.concatenate(self.junk_rhs)
self.junk_div = np.concatenate(self.junk_div)
# Add the RHS now that the shape is finalized.
self.dof.add(MapZipper(self.map_rhs.shape, self.map_rhs.wcs))
self.dof.add(ArrayZipper(self.junk_rhs.shape), distributed=True)
self.ready = True
def A(self, x):
imap, ijunk = self.dof.unzip(x)
# This is necessary because multizipper reduces everything to a single array, and
# hence can't maintain the separation between map_dtype and tod_dtype
ijunk = ijunk.astype(self.dtype_tod)
omap, ojunk = imap * 0, ijunk * 0
for di, data in enumerate(self.observations):
j1, j2 = self.junk_offs[di:di + 2]
tod = np.zeros([data.ndet, data.nsamp], self.dtype_tod)
wmap = imap.extract(data.shape, data.wcs) * 1
t1 = time.time()
data.pmat.from_map(wmap, dest=tod)
data.pcut.forward(tod, ijunk[j1:j2])
t2 = time.time()
data.nmat.apply(tod)
t3 = time.time()
wmap[:] = 0
data.pcut.backward(tod, ojunk[j1:j2])
data.pmat.to_map(signal=tod, dest=wmap)
t4 = time.time()
omap.insert(wmap, op=np.ndarray.__iadd__)
#L.debug("A %-70s P %8.3f N %8.3f P' %8.3f %3d %6d" % (data.id, t2-t1, t3-t2, t4-t3, data.ndet, data.nsamp))
omap = utils.allreduce(omap, self.comm)
return self.dof.zip(omap, ojunk)
# Load mask
# if i==0:
# mask = enmap.read_map(savedir + "/tilec_mask.fits")
# shape,wcs = mask.shape,mask.wcs
# modlmap = mask.modlmap()
# binner = stats.bin2D(modlmap,bin_edges)
# w2 = np.mean(mask**2.)
inputs = {}
for input_name in input_names:
# kmap= enmap.fft(get_input(input_name,args.set_id,sim_index,shape,wcs) * mask,normalize='phys')
# res = binner.bin(np.real(kmap*kmap.conj()))[1]/w2
res = np.ones((5, )) * 1e14 if input_name == 'CMB' else np.ones(
(5, )) * 1e-14 #!!!!!
print(rank, task, input_name, res[3])
try:
totinputs[input_name] = totinputs[input_name] + res
except:
totinputs[input_name] = res.copy()
for key in sorted(totinputs.keys()):
totinputs[key] = putils.allreduce(totinputs[key], comm) / nsims
if rank == 0:
# cents = binner.centers
for input_name in input_names:
ii = totinputs[input_name]
print(input_name, ii)
def get_set2():
return np.random.normal(loc=1e-14, scale=0.1e-14, size=(bins, ))
comm, rank, my_tasks = mpi.distribute(nsims)
results = {}
for task in my_tasks:
vec = get_set1()
try:
results['1'] = results['1'] + vec
except:
results['1'] = vec.copy()
vec = get_set2()
try:
results['2'] = results['2'] + vec
except:
results['2'] = vec.copy()
results['1'] = utils.allreduce(results['1'], comm) / nsims
results['2'] = utils.allreduce(results['2'], comm) / nsims
if rank == 0:
print(results['1'])
print(results['2'])