def trans(self, imap, omap): """Transpose of apply. omap = U.T(imap). Omap argument specifies the shape of the result, but its values may be destroyed during the operation.""" imap = np.ascontiguousarray(utils.moveaxis(imap,0,-1)) omap = np.ascontiguousarray(utils.moveaxis(omap,0,-1)) interpol.map_coordinates(omap, self.skew_pix, odata=imap, trans=True, order=self.order) return np.ascontiguousarray(utils.moveaxis(omap,-1,0))
def box2pix(shape, wcs, box):
"""Convert one or several bounding boxes of shape [2,2] or [n,2,2]
into pixel counding boxes in standard python half-open format.
The box must be [{from,to},...].
"""
nphi = int(np.round(np.abs(360. / wcs.wcs.cdelt[0])))
box = np.asarray(box)
fbox = box.reshape(-1, 2, 2)
if wcs.wcs.cdelt[0] < 0: fbox[..., 1] = fbox[..., ::-1, 1]
# Must rollaxis because sky2pix expects [{dec,ra},...]
ibox = enmap.sky2pix(shape, wcs, utils.moveaxis(fbox, 2, 0), corner=True)
# FIXME: This is one of many places in the code that will break if a bounding
# box goes more than half the way around the sky. Properly fixing
# this will require a big overhaul, especially if we want to handle
# bounding boxes that are bigger than the full sky.
ibox[1] = utils.unwrap_range(ibox[1].T, nphi).T
#ibox[1] = utils.unwind(ibox[1], nphi)
# We now have [{y,x},:,{from,to}
# Truncate to integer, and add one to endpoint to make halfopen interval
ibox = np.floor(np.sort(ibox, 2)).astype(int)
# Add 1 to endpoint to make halfopen interval
ibox[:, :, 1] += 1
#ibox = np.array([np.floor(ibox[0]),np.ceil(ibox[1])]).astype(int)
# Shuffle to [:,{from,to},{y,x}]
ibox = utils.moveaxis(ibox, 0, 2)
return ibox.reshape(box.shape)
def apply(self, map): """Apply unskew operation to map, returning an array where the scanning motion goes along the vertical axis.""" work = np.ascontiguousarray(utils.moveaxis(map,0,-1)) omap = interpol.map_coordinates(work, self.skew_pix, order=self.order) omap = np.ascontiguousarray(utils.moveaxis(omap,-1,0)) return omap
def slice_downgrade(d, s, axis=-1): """Slice array d along the specified axis using the Slice s, but interpret the step part of the slice as downgrading rather than skipping.""" a = moveaxis(d, axis, 0) step = s.step or 1 a = a[s.start:s.stop:-1 if step < 0 else 1] step = abs(step) # Handle the whole blocks first a2 = a[:len(a)/step*step] a2 = np.mean(a2.reshape((len(a2)/step,step)+a2.shape[1:]),1) # Then append the incomplete block if len(a2)*step != len(a): rest = a[len(a2)*step:] a2 = np.concatenate([a2,[np.mean(rest,0)]],0) return moveaxis(a2, 0, axis)
def draw_map_field_mpl(map, args, crange=None, printer=noprint):
"""Render a map field using matplotlib. Less tested and
maintained than draw_map_field, and supports fewer features.
Returns an object one can call savefig on to draw."""
map, color = prepare_map_field(map, args, crange, printer=printer)
# Set up matplotlib. We do it locally here to
# avoid having it as a dependency in general
with printer.time("matplotplib", 3):
import matplotlib
matplotlib.use("Agg")
from matplotlib import pyplot, ticker
matplotlib.rcParams.update({'font.size': 10})
dpi, pad = args.mpl_dpi, args.mpl_pad
winch, hinch = map.shape[2]/dpi, map.shape[1]/dpi
fig = pyplot.figure(figsize=(winch+pad,hinch+pad))
box = map.box()*180/np.pi
pyplot.imshow(utils.moveaxis(color,0,2), extent=[box[0,1],box[1,1],box[1,0],box[0,0]])
# Make conformal in center of image
pyplot.axes().set_aspect(1/np.cos(np.mean(map.box()[:,0])))
if args.grid % 2:
ax = pyplot.axes()
ticks = np.full(2,1.0); ticks[:] = [float(w) for w in args.ticks.split(",")]
ax.xaxis.set_major_locator(ticker.MultipleLocator(ticks[1]))
ax.yaxis.set_major_locator(ticker.MultipleLocator(ticks[0]))
if args.subticks:
ax.xaxis.set_minor_locator(ticker.MultipleLocator(args.sub))
ax.yaxis.set_minor_locator(ticker.MultipleLocator(args.sub))
pyplot.minorticks_on()
pyplot.grid(True, which="major", linewidth=2)
pyplot.grid(True, which="minor", linewidth=1)
else:
pyplot.grid(True)
pyplot.tight_layout(pad=0.0,h_pad=0.0,w_pad=0.0)
return pyplot
def draw_colorbar(crange, width, args):
col = tuple([
int(args.font_color[i:i + 2], 16)
for i in range(0, len(args.font_color), 2)
])
font = cgrid.get_font(args.font_size)
fmt = "%g"
labels, boxes = [], []
for val in crange:
labels.append(fmt % val)
boxes.append(font.getsize(labels[-1]))
boxes = np.array(boxes, int)
lw, lh = np.max(boxes, 0)
img = PIL.Image.new("RGBA", (width, lh))
draw = PIL.ImageDraw.Draw(img)
# Draw the labels on the image
draw.text((lw - boxes[0, 0], 0), labels[0], col, font=font)
draw.text((width - lw, 0), labels[1], col, font=font)
# Draw the color bar itself
bar_data = np.zeros((lh, width - 2 * lw))
bar_data[:] = np.linspace(0, 1, bar_data.shape[-1])
bar_col = map_to_color(bar_data[None], [0, 1], args)
bar_img = PIL.Image.fromarray(utils.moveaxis(bar_col, 0,
2)).convert('RGBA')
# Overlay it on the output image
img.paste(bar_img, (lw, 0))
bounds = np.array([[0, 0], [width, lh]])
return img, bounds
def overlaps_any(box, refboxes): rdec, rra = utils.moveaxis(refboxes - box[0,:], 2,0) wdec, wra = box[1] - box[0] rra -= np.floor(rra[:,0,None]/(2*np.pi)+0.5)*(2*np.pi) for i in range(-1,2): nra = rra + i*(2*np.pi) if np.any((nra[:,1]>0)&(nra[:,0]<wra)&(rdec[:,1]>0)&(rdec[:,0]<wdec)): return True return False
def overlaps_any(box, refboxes): rdec, rra = utils.moveaxis(refboxes - box[0,:], 2,0) wdec, wra = np.abs(box[1] - box[0]) rra -= np.floor(rra[:,0,None]/(2*np.pi)+0.5)*(2*np.pi) for i in range(-1,2): nra = rra + i*(2*np.pi) if np.any((nra[:,1]>0)&(nra[:,0]<wra)&(rdec[:,1]>0)&(rdec[:,0]<wdec)): return True return False
def cellify(map, res): """Given a map [...,ny,nx] and a cell resolution [ry,rx], return map reshaped into a cell grid [...,ncelly,ncellx,ry,rx]. The map will be truncated if necessary""" res = np.array(res,int) cshape = map.shape[-2:]/res omap = map[...,:cshape[0]*res[0],:cshape[1]*res[1]] omap = omap.reshape(omap.shape[:-2]+(cshape[0],res[0],cshape[1],res[1])) omap = utils.moveaxis(omap, -3, -2) return omap
def draw_map_field(map, args, crange=None, return_layers=False, return_info=False, printer=noprint, cache=None):
"""Draw a single map field, resulting in a single image. Adds a coordinate grid
and lables as specified by args. If return_layers is True, an array will be
returned instead of an image, wich each entry being a component of the image,
such as the base image, the coordinate grid, the labels, etc. If return_bounds
is True, then the """
map, color = prepare_map_field(map, args, crange, printer=printer)
tag = (tuple(map.shape), map.wcs.to_header_string(), repr(args))
layers = []
names = []
# Image layer
with printer.time("to image", 3):
img = PIL.Image.fromarray(utils.moveaxis(color,0,2)).convert('RGBA')
layers.append((img,[[0,0],img.size]))
names.append("img")
# Contours
if args.contours:
with printer.time("draw contours", 3):
contour_levels = calc_contours(crange, args)
cimg = draw_contours(map, contour_levels, args)
layers.append((cimg, [[0,0],cimg.size]))
names.append("cont")
# Annotations
if args.annotate:
with printer.time("draw annotations", 3):
def get_aimg():
annots = parse_annotations(args.annotate)
return draw_annotations(map, annots, args)
aimg = get_cache(cache, ("annotate",tag), get_aimg)
layers.append((aimg, [[0,0],aimg.size]))
names.append("annot")
# Coordinate grid
if args.grid % 2:
with printer.time("draw grid", 3):
ginfo = get_cache(cache, ("ginfo",tag), lambda: calc_gridinfo(map.shape, map.wcs, args))
grid = get_cache(cache, ("grid", tag), lambda: draw_grid(ginfo, args))
layers.append(grid)
names.append("grid")
if not args.nolabels:
with printer.time("draw labels", 3):
labels = get_cache(cache, ("labels",tag), lambda: draw_grid_labels(ginfo, args))
layers.append(labels)
names.append("tics")
# Possibly other stuff too, like point source circles
# or contours
with printer.time("stack layers", 3):
layers, bounds = standardize_images(layers)
if not return_layers: layers = merge_images(layers)
class Info:
def __init__(self, bounds, names):
self.bounds = bounds
self.names = names
info = Info(bounds, names)
if return_info: return layers, info
else: return layers
def eval_srcs_vectorized(posmap, poss, amps, beam, cres, nhit, cell_srcs):
# Find the coordinates of each pixel in each cell
cposmap = cellify(posmap, cres) # [2,cy,cx,ry,rx]
csrcpos = poss[cell_srcs] # [cy,cx,nmax,2]
csrcpos = utils.moveaxis(csrcpos, -1, 0) # [2,cy,cx,nmax]
csrcamp = amps[cell_srcs] # [cy,cx,nmax,ncomp2]
csrcamp = utils.moveaxis(csrcamp, -1, 0) # [ncomp,cy,cx,nmax]
# Ok, we can now compute the total
# Get the beam radial position for each pixel for each source. [cy,cx,nmax,ry,rx]
# Express this in beam resolution units
diff = cposmap[:, :, :, None, :, :] - csrcpos[:, :, :, :, None, None]
bpix = (diff[0]**2 +
(diff[1] * np.cos(cposmap[0, :, :, None, :, :]))**2)**0.5
bpix = (bpix - beam[0, 0]) / (beam[0, 1] - beam[0, 0])
# Evaluate the beam at these locations
bval = utils.interpol(beam[1], bpix[None], mode="constant", order=1)
model = csrcamp[:, :, :, :, None, None] * bval
# Mask invalid sources
mask = np.arange(model.shape[-3]) < nhit[:, :, None]
model *= mask[:, :, :, None, None]
model = np.sum(model, -3) # [ncomp,cy,cx,ry,rx]
model = uncellify(model)
return model
id = ids[ind]
entry = file_db[id]
try:
stats.append(todinfo.build_tod_stats(entry))
except (errors.DataMissing, AttributeError) as e:
print "Skipping %s (%s)" % (id, e.message)
continue
print "%3d %4d/%d %5.1f%% %s" % (comm.rank, ind + 1, len(ids),
(ind + 1) / float(len(ids)) * 100, id)
stats = todinfo.merge_tod_stats(stats)
if comm.rank == 0: print "Reducing"
comm.Barrier()
for key in stats:
stats[key] = utils.allgatherv(stats[key], comm)
# Sort by id and move id index last
inds = np.argsort(stats["id"])
for key in stats:
stats[key] = utils.moveaxis(stats[key][inds], 0, -1)
stat_db = todinfo.Todinfo(stats)
# Merge with original tags. Rightmost overrides for overridable fields.
# For tags, we get the union. This means that stat_db can't override incorrect
# tags in scan_db, just add to them.
stat_db = scan_db + stat_db
if comm.rank == 0:
print "Writing"
stat_db.write(args.ofile)
print "Done"
def uncellify(cmap): omap = utils.moveaxis(cmap, -2, -3) omap = omap.reshape(omap.shape[:-4]+(omap.shape[-4]*omap.shape[-3],omap.shape[-2]*omap.shape[-1])) return omap