def getitem_helper(self, sel): if type(sel) != tuple: sel = (sel, ) assert len(sel) < 3, "Too many indices in slice" detslice = sel[0] if len(sel) > 0 else slice(None) sampslice = sel[1] if len(sel) > 1 else slice(None) assert isinstance(sampslice, slice), "Sample part of slice must be slice object" res = cpy.deepcopy(self) # These will be passed to fortran, so make them contiguous res.boresight = np.ascontiguousarray( enlib.slice.slice_downgrade(res.boresight, sampslice, axis=0)) res.offsets = np.ascontiguousarray(res.offsets[detslice]) res.comps = np.ascontiguousarray(res.comps[detslice]) res.dets = res.dets[detslice] res.hwp = np.ascontiguousarray( enlib.slice.slice_downgrade(res.hwp, sampslice, axis=0)) res.hwp_phase = np.ascontiguousarray( enlib.slice.slice_downgrade(res.hwp_phase, sampslice, axis=0)) try: # The whole scan stuff is horrible and should be redesigned res.dark_tod = np.ascontiguousarray( enlib.slice.slice_downgrade(res.dark_tod, sampslice, axis=1)) res.dark_cut = res.dark_cut[sel] except AttributeError as e: pass try: res.buddy_comps = res.buddy_comps[:, detslice] res.buddy_offs = res.buddy_offs[:, detslice] except AttributeError as e: pass res.cut = res.cut[sel] res.cut_noiseest = res.cut_noiseest[sel] res.noise = res.noise[sel] return res, detslice, sampslice
def getitem_helper(self, sel): """Expands sel to a detector and sample slice. The detector slice is straightforward. The sample slice may be less so. In fourier space, its effect is a rescaling and truncation, such that find2 = find1 * n2/n1, with find2_max = find1_max * n2/n1 / step, and n2 = stop-start.""" if type(sel) != tuple: sel = (sel, ) assert len(sel) < 3, "Too many indices in slice" detslice = sel[0] if len(sel) > 0 else slice(None) sampslice = sel[1] if len(sel) > 1 else slice(None) assert isinstance(sampslice, slice), "Sample part of slice must be slice object" res = copy.deepcopy(self) return res, detslice, sampslice
def slice_helper(self, newlen, oldlen): """Returns a new Bufmap with buffer slices scaled by newlen/oldlen. This is used for defining slice operations.""" buf_info = tuple(np.array(self.buf_info) * newlen / oldlen) buf_shape = self.buf_shape * newlen / oldlen data_info = [(ind, dslice, slice(bslice.start * newlen / oldlen, bslice.stop * newlen / oldlen)) for ind, dslice, bslice in self.data_info] return Bufmap(data_info, buf_info, buf_shape)
def get_samples(self): """Return the actual detector samples. Slow! Data is read from disk, so store the result if you need to reuse it.""" with h5py.File(self.fname, "r") as hfile: tod = hfile["tod"].value[self.subdets] method = config.get("downsample_method") for s in self.sampslices: tod = resample.resample(tod, 1.0/np.abs(s.step or 1), method=method) s = slice(s.start, s.stop, np.sign(s.step) if s.step else None) tod = tod[:,s] res = np.ascontiguousarray(tod) return res
def broadcast_into(dmap, val, axis=None): """Transpose of a sum a dmap along the specified axis, or the flattened version if axis is None.""" # Full sum if axis is None: dmap[:] = val else: # Non-pixel sums if axis < 0: axis = dmap.ndim + axis if axis < dmap.ndim - 2: for itile, otile in zip(dmap.tiles, val.tiles): itile[:] = np.asarray(otile)[(slice(None), ) * axis + (None, ) + (slice(None), ) * (dmap.ndim - axis - 1)] else: # Pixel sums: Sum each tile along the specified direction. Then sum tiles # that are on the same row/column. Then stack along the remaining row/column res = np.zeros(dmap.shape[:axis] + dmap.shape[axis + 1:], dmap.dtype) paxis = axis - (dmap.ndim - 2) for tile, ind in zip(dmap.tiles, dmap.loc_inds): pbox = dmap.geometry.tile_boxes[ind] if paxis == 0: tile[:] = res[Ellipsis, None, pbox[0, 1]:pbox[1, 1]] else: tile[:] = res[Ellipsis, pbox[0, 0]:pbox[1, 0], None]
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)