def output_tile(prefix, tpos, info):
shape = info.model.shape[-2:]
tname = "tile%(y)03d_%(x)03d.fits" % {"y":tpos,"x":tpos[1]}
ffpad_slice = (Ellipsis,slice(0,shape[0]-info.ffpad[0]),slice(0,shape[1]-info.ffpad[1]))
for maptypename, mapgroup in [("snmap",info.snmaps),("snresid",info.snresid)]:
for srctypename, map in mapgroups:
write_padtile("%s%s_%s/%s" % (prefix, srctypename, maptypename, tname), map[ffpad_slice])
if not args.output_full_model:
if len(info.model) > 0: model = info.model[0]
else: model = enmap.zeros(info.model.shape[-2:], info.model.wcs, info.model.dtype)
write_padtile(prefix + "model" + tname, model[ffpad_slice])
else:
write_padtile(prefix + "model" + tname, info.model[ffpad_slice])
# Output total catalogue
apod_slice = (slice(args.apod_edge,-args.apod_edge), slice(args.apod_edge,-args.apod_edge))
shape, wcs = enmap.slice_geometry(info.model.shape, info.model.wcs, ffpad_slice[-2:])
shape, wcs = enmap.slice_geometry(shape, wcs, apod_slice)
box = enmap.box(shape, wcs)
jointmap.write_catalogue(prefix + "catalogue" + tname, info.catalogue, box)
def output_tile(prefix, tpos, info):
shape = info.model.shape[-2:]
tname = "tile%(y)03d_%(x)03d.fits" % {"y":tpos,"x":tpos[1]}
ffpad_slice = (Ellipsis,slice(0,shape[0]-info.ffpad[0]),slice(0,shape[1]-info.ffpad[1]))
for maptypename, mapgroup in [("snmap",info.snmaps),("snresid",info.snresid)]:
for srctypename, map in mapgroups:
write_padtile("%s%s_%s/%s" % (prefix, srctypename, maptypename, tname), map[ffpad_slice])
if not args.output_full_model:
if len(info.model) > 0: model = info.model[0]
else: model = enmap.zeros(info.model.shape[-2:], info.model.wcs, info.model.dtype)
write_padtile(prefix + "model" + tname, model[ffpad_slice])
else:
write_padtile(prefix + "model" + tname, info.model[ffpad_slice])
# Output total catalogue
apod_slice = (slice(args.apod_edge,-args.apod_edge), slice(args.apod_edge,-args.apod_edge))
shape, wcs = enmap.slice_geometry(info.model.shape, info.model.wcs, ffpad_slice[-2:])
shape, wcs = enmap.slice_geometry(shape, wcs, apod_slice)
box = enmap.box(shape, wcs)
jointmap.write_catalogue(prefix + "catalogue" + tname, info.catalogue, box)
def make_dummy_tile(shape, wcs, box, pad=0):
pbox = calc_pbox(shape, wcs, box)
if pad:
pbox[0] -= pad
pbox[1] += pad
shape2, wcs2 = enmap.slice_geometry(
shape,
wcs, (slice(pbox[0, 0], pbox[1, 0]), slice(pbox[0, 1], pbox[1, 1])),
nowrap=True)
shape2 = shape[:-2] + tuple(pbox[1] - pbox[0])
map = enmap.zeros(shape2, wcs2, dtype)
div = enmap.zeros(shape2[-2:], wcs2, dtype)
return bunch.Bunch(map=map, div=div)
def __init__(self,
shape,
wcs=None,
bbpix=None,
tshape=None,
dtype=None,
comm=None,
bbox=None,
pre=None):
"""DGeometry(shape, wcs, bbpix, tshape) or
DGeometry(dgeometry)
Construct a DMap geometry. When constructed explicitly from shape, wcs, etc.,
it is a relatively expensive operation, as MPI communication is necessary.
Constructing one based on an existing DGeometry is fast.
bbox indicates the bounds [{from,to},{lat,lon}] of the area of the map of
interest to each mpi task, and will in general be different for each task.
bbpix is the same as bbox, but expressed in units of pixels.
It is allowed to pass a list of bounding boxes, in which case each mpi task
will have multiple work spaces.
tshape specifies the tiling scheme to use. The global geometry will be
split into tiles of tshape dimensions (except at the edges, as no tile
will extend beyond the edge of the full map), and tiles will be stored
in a distributed fashion between mpi tasks based on the degree of overlap
with their workspaces."""
try:
# We are pretty lightweight, so copy everything properly. This is
# less error prone than having changes in one object suddenly affect
# another. comm must *not* be deepcopied, as that breaks it.
for key in ["shape", "tshape", "comm", "dtype"]:
setattr(self, key, getattr(shape, key))
for key in [
"wcs", "tile_pos", "tile_boxes", "tile_geometry",
"work_geometry", "tile_ownership", "tile_glob2loc",
"tile_loc2glob", "tile_bufinfo", "work_bufinfo"
]:
setattr(self, key, copy.deepcopy(getattr(shape, key)))
except AttributeError:
# 1. Set up basic properties
assert shape is not None
if wcs is None:
_, wcs = enmap.geometry(pos=np.array([[-1, -1], [1, 1]]) * 5 *
np.pi / 180,
shape=shape[-2:])
if comm is None: comm = mpi.COMM_WORLD
if tshape is None: tshape = (720, 720)
if dtype is None: dtype = np.float64
if bbpix is None:
if bbox is None: bbpix = [[0, 0], shape[-2:]]
else: bbpix = box2pix(shape, wcs, bbox)
nphi = int(np.round(np.abs(360 / wcs.wcs.cdelt[0])))
# Reorder from/to if necessary, and expand to [:,2,2]
bbpix = sanitize_pixbox(bbpix, shape)
dtype = np.dtype(dtype)
self.shape, self.wcs, self.bbpix = tuple(
shape), wcs.deepcopy(), np.array(bbpix)
self.tshape, self.dtype, self.comm = tuple(tshape), dtype, comm
# 2. Set up workspace descriptions
work_geometry = [
enmap.slice_geometry(
shape,
wcs, (slice(b[0, 0], b[1, 0]), slice(b[0, 1], b[1, 1])),
nowrap=True) for b in bbpix
]
# 3. Define global workspace ownership
nwork = utils.allgather([len(bbpix)], comm)
wown = np.concatenate(
[np.full(n, i, dtype=int) for i, n in enumerate(nwork)])
# 3. Define tiling. Each tile has shape tshape, starting from the (0,0) corner
# of the full map. Tiles at the edge are clipped, as pixels beyond the edge
# of the full map may have undefined wcs positions.
tbox = build_tiles(shape, tshape)
bshape = tbox.shape[:2]
tbox = tbox.reshape(-1, 2, 2)
ntile = len(tbox)
tile_indices = np.array(
[np.arange(ntile) / bshape[1],
np.arange(ntile) % bshape[1]]).T
# 4. Define tile ownership.
# a) For each task compute the overlap of each tile with its workspaces, and
# concatenate across tasks to form a [nworktot,ntile] array.
wslices = select_nonempty( # slices into work
utils.allgatherv(utils.box_slice(bbpix, tbox), comm,
axis=0), # normal
utils.allgatherv(utils.box_slice(bbpix + [0, nphi], tbox),
comm,
axis=0)) # wrapped
tslices = select_nonempty( # slices into tiles
utils.allgatherv(utils.box_slice(tbox, bbpix), comm,
axis=1), # normal
utils.allgatherv(utils.box_slice(tbox, bbpix + [0, nphi]),
comm,
axis=1)) # wrapped
# b) Compute the total overlap each mpi task has with each tile, and use this
# to decide who should get which tiles
overlaps = utils.box_area(wslices)
overlaps = utils.sum_by_id(overlaps, wown, 0)
town = assign_cols_round_robin(overlaps)
# Map tile indices from local to global and back
tgmap = [[] for i in range(comm.size)]
tlmap = np.zeros(ntile, dtype=int)
for ti, id in enumerate(town):
tlmap[ti] = len(tgmap[id]) # glob 2 loc
tgmap[id].append(ti) # loc 2 glob
# 5. Define tiles
tile_geometry = [
enmap.slice_geometry(
shape, wcs,
(slice(tb[0, 0], tb[1, 0]), slice(tb[0, 1], tb[1, 1])))
for tb in tbox
]
# 6. Define mapping between work<->wbuf and tiles<->tbuf
wbufinfo = np.zeros([2, comm.size], dtype=int)
tbufinfo = np.zeros([2, comm.size], dtype=int)
winfo, tinfo = [], []
woff, toff = 0, 0
prelen = np.product(shape[:-2])
for id in xrange(comm.size):
## Buffer info to send to alltoallv
wbufinfo[1, id] = woff
tbufinfo[1, id] = toff
# Slices for transfering to and from w buffer. Loop over all of my
# workspaces and determine the slices into them and how much we need
# to send.
for tloc, tglob in enumerate(np.where(town == id)[0]):
for wloc, wglob in enumerate(
np.where(wown == comm.rank)[0]):
ws = wslices[wglob, tglob]
wlen = utils.box_area(ws) * prelen
work_slice = (Ellipsis, slice(ws[0, 0], ws[1, 0]),
slice(ws[0, 1], ws[1, 1]))
wbuf_slice = slice(woff, woff + wlen)
winfo.append((wloc, work_slice, wbuf_slice))
woff += wlen
# Slices for transferring to and from t buffer. Loop over all
# my tiles, and determine how much I have to receive from each
# workspace of each task.
for tloc, tglob in enumerate(np.where(town == comm.rank)[0]):
for wloc, wglob in enumerate(np.where(wown == id)[0]):
ts = tslices[tglob, wglob]
tlen = utils.box_area(ts) * prelen
tile_slice = (Ellipsis, slice(ts[0, 0], ts[1, 0]),
slice(ts[0, 1], ts[1, 1]))
tbuf_slice = slice(toff, toff + tlen)
tinfo.append((tloc, tile_slice, tbuf_slice))
toff += tlen
wbufinfo[0, id] = woff - wbufinfo[1, id]
tbufinfo[0, id] = toff - tbufinfo[1, id]
wbufinfo, tbufinfo = tuple(wbufinfo), tuple(tbufinfo)
# TODO: store tbox? loc_geometry vs. tile_geometry, etc.
# 7. Store
# [ntile,2]: position of each (global) tile in grid
self.tile_pos = tile_indices
# [ntile,2,2]: pixel box for each (global) tile
self.tile_boxes = tbox
# [ntile,(shape,wcs)]: geometry of each (global) tile
self.tile_geometry = tile_geometry
# [nwork,(shape,wcs)]: geometry of each local workspace
self.work_geometry = work_geometry
# [ntile]: rank of owner of each tile
self.tile_ownership = town
# [ntile]: local index of each (global) tile
self.tile_glob2loc = tlmap
# [nloc]: global index of each local tile
self.tile_loc2glob = tgmap[comm.rank]
# Communication buffers
self.tile_bufinfo = Bufmap(tinfo, tbufinfo, toff)
self.work_bufinfo = Bufmap(winfo, wbufinfo, woff)
def rand_map(shape,
wcs,
ps_lensinput,
lmax=None,
maplmax=None,
dtype=np.float64,
seed=None,
oversample=2.0,
spin=2,
output="l",
geodesic=True,
verbose=False,
delta_theta=None):
import curvedsky, sharp
ctype = np.result_type(dtype, 0j)
# Restrict to target number of components
oshape = shape[-3:]
if len(oshape) == 2: shape = (1, ) + tuple(shape)
# First draw a random lensing field, and use it to compute the undeflected positions
if verbose: print("Generating alms")
alm, ainfo = curvedsky.rand_alm(ps_lensinput,
lmax=lmax,
seed=seed,
dtype=ctype,
return_ainfo=True)
phi_alm, cmb_alm = alm[0], alm[1:1 + shape[-3]]
# Truncate alm if we want a smoother map. In taylens, it was necessary to truncate
# to a lower lmax for the map than for phi, to avoid aliasing. The appropriate lmax
# for the cmb was the one that fits the resolution. FIXME: Can't slice alm this way.
#if maplmax: cmb_alm = cmb_alm[:,:maplmax]
del alm
if delta_theta is None: bsize = shape[-2]
else:
bsize = utils.nint(abs(delta_theta / utils.degree / wcs.wcs.cdelt[1]))
# Adjust bsize so we don't get any tiny blocks at the end
nblock = shape[-2] // bsize
bsize = int(shape[-2] / (nblock + 0.5))
# Allocate output maps
if "p" in output: phi_map = enmap.empty(shape[-2:], wcs, dtype=dtype)
if "k" in output:
kappa_map = enmap.empty(shape[-2:], wcs, dtype=dtype)
l = np.arange(ainfo.lmax + 1.0)
kappa_alm = ainfo.lmul(phi_alm, l * (l + 1) / 2)
for i1 in range(0, shape[-2], bsize):
curvedsky.alm2map(kappa_alm, kappa_map[..., i1:i1 + bize, :])
del kappa_alm
if "a" in output:
grad_map = enmap.empty((2, ) + shape[-2:], wcs, dtype=dtype)
if "u" in output: cmb_raw = enmap.empty(shape, wcs, dtype=dtype)
if "l" in output: cmb_obs = enmap.empty(shape, wcs, dtype=dtype)
# Then loop over dec bands
for i1 in range(0, shape[-2], bsize):
i2 = min(i1 + bsize, shape[-2])
lshape, lwcs = enmap.slice_geometry(shape, wcs,
(slice(i1, i2), slice(None)))
if "p" in output:
if verbose: print("Computing phi map")
curvedsky.alm2map(phi_alm, phi_map[..., i1:i2, :])
if verbose: print("Computing grad map")
if "a" in output: grad = grad_map[..., i1:i2, :]
else: grad = enmap.zeros((2, ) + lshape[-2:], lwcs, dtype=dtype)
curvedsky.alm2map(phi_alm, grad, deriv=True)
if "l" not in output: continue
if verbose: print("Computing observed coordinates")
obs_pos = enmap.posmap(lshape, lwcs)
if verbose: print("Computing alpha map")
raw_pos = enmap.samewcs(
offset_by_grad(obs_pos, grad, pol=shape[-3] > 1,
geodesic=geodesic), obs_pos)
del obs_pos, grad
if "u" in output:
if verbose: print("Computing unlensed map")
curvedsky.alm2map(cmb_alm, cmb_raw[..., i1:i2, :], spin=spin)
if verbose: print("Computing lensed map")
cmb_obs[..., i1:i2, :] = curvedsky.alm2map_pos(cmb_alm,
raw_pos[:2],
oversample=oversample,
spin=spin)
if raw_pos.shape[0] > 2 and np.any(raw_pos[2]):
if verbose: print("Rotating polarization")
cmb_obs[..., i1:i2, :] = enmap.rotate_pol(cmb_obs[..., i1:i2, :],
raw_pos[2])
del raw_pos
del cmb_alm, phi_alm
# Output in same order as specified in output argument
res = []
for c in output:
if c == "l": res.append(cmb_obs.reshape(oshape))
elif c == "u": res.append(cmb_raw.reshape(oshape))
elif c == "p": res.append(phi_map)
elif c == "k": res.append(kappa_map)
elif c == "a": res.append(grad_map)
return tuple(res)
# so get them based on the first one
entry = filedb.data[ids[0]]
site = actdata.read(entry, ["site"]).site
# Determine the bounding box of our selected data
bounds = db.data["bounds"].reshape(2,-1).copy()
bounds[0] = utils.rewind(bounds[0], bounds[0,0], 360)
box = utils.widen_box(utils.bounding_box(bounds.T), 4*args.rad, relative=False)
waz, wel = box[1]-box[0]
# Use fullsky horizontally if we wrap too far
if waz <= 180:
shape, wcs = enmap.geometry(pos=box[:,::-1]*utils.degree, res=args.res*utils.degree, proj="car", ref=(0,0))
else:
shape, wcs = enmap.fullsky_geometry(res=args.res*utils.degree)
y1, y2 = np.sort(enmap.sky2pix(shape, wcs, [box[:,1]*utils.degree,[0,0]])[0].astype(int))
shape, wcs = enmap.slice_geometry(shape, wcs, (slice(y1,y2),slice(None)))
sys.stderr.write("using %s workspace with resolution %.2f deg" % (str(shape), args.res) + "\n")
# Get the hitmap for each block
hits = enmap.zeros((nblock,narray)+shape, wcs)
ndig = calc_ndig(nblock)
sys.stderr.write("estimating hitmap for block %*d/%d" % (ndig,0,nblock))
for bi in range(nblock):
for ai in range(narray):
block_db = db.select(block_inds[bi,ai])
hits[bi,ai] = fastweight.fastweight(shape, wcs, block_db, array_rad=args.rad*utils.degree, site=site, weight=args.weight)
sys.stderr.write("%s%*d/%d" % ("\b"*(1+2*ndig),ndig,bi+1,nblock))
sys.stderr.write("\n")
# Build a mask for the region of interest per array
else:
def oname(name):
return args.odir + "/" + name
pixbox = np.array([[-pad, -pad],
[shape[-2] + pad, shape[-1] + pad]])
if args.cont and os.path.exists(oname("cands.txt")): continue
if comm_intra.rank == 0:
print "group %3d processing tile %3d %3d" % (comm_inter.rank, ty,
tx)
# Get the shape of the sliced, downgraded tiles
tshape, twcs = enmap.slice_geometry(
shape,
wcs, [
slice(pixbox[0, 0], pixbox[1, 0]),
slice(pixbox[0, 1], pixbox[1, 1])
],
nowrap=True)
tshape, twcs = enmap.downgrade_geometry(tshape, twcs, args.downgrade)
#print ty, tx, np.mean(enmap.box(tshape, twcs),0)/utils.degree
# Read in our tile data
frhss, kmaps, mjds = [], [], []
for ind in range(comm_intra.rank, len(idirs), comm_intra.size):
if not overlaps(pboxes[ind], pixbox, nphi):
#print idirs[ind], "does not overlap", pixbox.reshape(-1), pboxes[ind].reshape(-1)
continue
idir = idirs[ind]
lshape, lwcs = enmap.read_map_geometry(idir + "/frhs.fits")
pixbox_loc = pixbox - enmap.pixbox_of(wcs, lshape, lwcs)[0]
kmap = enmap.read_map(idir + "/kmap.fits",
# so get them based on the first one
entry = filedb.data[ids[0]]
site = actdata.read(entry, ["site"]).site
# Determine the bounding box of our selected data
bounds = db.data["bounds"].reshape(2,-1).copy()
bounds[0] = utils.rewind(bounds[0], bounds[0,0], 360)
box = utils.widen_box(utils.bounding_box(bounds.T), 4*args.rad, relative=False)
waz, wel = box[1]-box[0]
# Use fullsky horizontally if we wrap too far
if waz <= 180:
shape, wcs = enmap.geometry(pos=box[:,::-1]*utils.degree, res=args.res*utils.degree, proj="car", ref=(0,0))
else:
shape, wcs = enmap.fullsky_geometry(res=args.res*utils.degree)
y1, y2 = np.sort(enmap.sky2pix(shape, wcs, [box[:,1]*utils.degree,[0,0]])[0].astype(int))
shape, wcs = enmap.slice_geometry(shape, wcs, (slice(y1,y2),slice(None)))
sys.stderr.write("using %s workspace with resolution %.2f deg" % (str(shape), args.res) + "\n")
# Get the hitmap for each block
hits = enmap.zeros((nblock,narray)+shape, wcs)
ndig = calc_ndig(nblock)
sys.stderr.write("estimating hitmap for block %*d/%d" % (ndig,0,nblock))
for bi in range(nblock):
for ai in range(narray):
block_db = db.select(block_inds[bi,ai])
hits[bi,ai] = fastweight.fastweight(shape, wcs, block_db, array_rad=args.rad*utils.degree, site=site, weight=args.weight)
sys.stderr.write("%s%*d/%d" % ("\b"*(1+2*ndig),ndig,bi+1,nblock))
sys.stderr.write("\n")
# Build a mask for the region of interest per array