def tiles(self, from_file=False):
comm = self.comm
for i in range(comm.rank, len(self.pboxes), comm.size):
eshape, ewcs = slice_geometry_by_pixbox(self.ishape, self.iwcs,
self.pboxes[i])
if from_file:
extracter = lambda x, **kwargs: self._prepare(
enmap.read_map(x,
pixbox=enmap.pixbox_of(
enmap.read_map_geometry(x)[1], eshape,
ewcs),
**kwargs))
else:
extracter = lambda x: self._prepare(
enmap.extract_pixbox(x, enmap.pixbox_of(
x.wcs, eshape, ewcs)))
inserter = lambda inp, out: enmap.insert_at(out,
self.ipboxes[i],
self._finalize(inp),
op=np.ndarray.__iadd__)
yield i, extracter, inserter, eshape, ewcs
def test_extract(self):
# Tests that extraction is sensible
shape, wcs = enmap.geometry(pos=(0, 0), shape=(500, 500), res=0.01)
imap = enmap.enmap(np.random.random(shape), wcs)
smap = imap[200:300, 200:300]
sshape, swcs = smap.shape, smap.wcs
smap2 = enmap.extract(imap, sshape, swcs)
pixbox = enmap.pixbox_of(imap.wcs, sshape, swcs)
# Do write and read test
filename = "temporary_extract_map.fits" # NOT THREAD SAFE
enmap.write_map(filename, imap)
smap3 = enmap.read_map(filename, pixbox=pixbox)
os.remove(filename)
assert np.all(np.isclose(smap, smap2))
assert np.all(np.isclose(smap, smap3))
assert wcsutils.equal(smap.wcs, smap2.wcs)
assert wcsutils.equal(smap.wcs, smap3.wcs)
def read_helper(fname, shape=None, wcs=None):
if shape is None: return enmap.read_map(fname)
mshape, mwcs = enmap.read_map_geometry(fname)
pixbox = enmap.pixbox_of(mwcs, shape, wcs)
return enmap.read_map(fname, pixbox=pixbox)
def search_maps_tiled(ifiles,
odir,
tshape=(1000, 1000),
margin=100,
padding=150,
mode="find",
icat=None,
box=None,
pixbox=None,
sel=None,
mask=None,
templates=default_templates,
cl_cmb=None,
freq0=98.0,
nmat1="constcorr",
nmat2="constcorr",
snr1=5,
snr2=4,
comps="TQU",
dtype=np.float32,
comm=None,
cont=False,
sim_cat=None,
sim_noise=False,
verbose=False):
wdir = odir + "/work"
utils.mkdir(wdir)
if comm is None: comm = bunch.Bunch(rank=0, size=1)
tshape = np.zeros(2, int) + tshape
meta = mapdata.read_meta(ifiles[0])
# Allow us to slice the map that will be tiled
geo = enmap.Geometry(*meta.map_geometry)
if pixbox is not None or box is not None:
geo = geo.submap(pixbox=pixbox, box=box)
if sel is not None: geo = geo[sel]
shape = np.array(geo.shape[-2:])
ny, nx = (shape + tshape - 1) // tshape
def is_done(ty, tx):
return os.path.isfile("%s/cat_%03d_%03d.fits" % (wdir, ty, tx))
tyxs = [(ty, tx) for ty in range(ny) for tx in range(nx)
if (not cont or not is_done(ty, tx))]
for ind in range(comm.rank, len(tyxs), comm.size):
# Get basic area of this tile
tyx = np.array(tyxs[ind])
if verbose:
print("%2d Processing tile %2d %2d of %2d %2d" %
(comm.rank, tyx[0], tyx[1], ny, nx))
yx1 = tyx * tshape
yx2 = np.minimum((tyx + 1) * tshape, shape)
# Apply padding
wyx1 = yx1 - margin - padding
wyx2 = yx2 + margin + padding
# Transform from box-relative pixbox to global pixbox
off = enmap.pixbox_of(meta.map_geometry[1], *geo)[0]
wyx1 += off
wyx2 += off
# Process this tile
res = search_maps(ifiles,
mode=mode,
icat=icat,
pixbox=[wyx1, wyx2],
templates=templates,
mask=mask,
cl_cmb=cl_cmb,
freq0=freq0,
nmat1=nmat1,
nmat2=nmat2,
snr1=snr1,
snr2=snr2,
comps=comps,
dtype=dtype,
sim_cat=sim_cat,
sim_noise=sim_noise,
verbose=verbose)
# Write tile results to work directory. We do this to avoid using too much memory,
# and to allow us to continue
write_results(wdir, res, padding=padding, tag="%03d_%03d" % tuple(tyx))
comm.Barrier()
# When everything's done, merge things into single files
if comm.rank == 0:
merge_results(wdir,
odir,
geo,
tshape=tshape,
margin=margin,
verbose=verbose)