def sim_srcs(shape, wcs, srcs, beam, omap=None, dtype=None, nsigma=5, rmax=None, method="loop", mmul=1,
return_padded=False):
"""Simulate a point source map in the geometry given by shape, wcs
for the given srcs[nsrc,{dec,ra,T...}], using the beam[{r,val},npoint],
which must be equispaced. If omap is specified, the sources will be
added to it in place. All angles are in radians. The beam is only evaluated up to
the point where it reaches exp(-0.5*nsigma**2) unless rmax is specified, in which
case this gives the maximum radius. mmul gives a factor to multiply the resulting
source model by. This is mostly useful in conction with omap. method can be
"loop" or "vectorized", but "loop" is both faster and uses less memory, so there's
no point in using the latter.
The source simulation is sped up by using a source lookup grid.
"""
if omap is None: omap = enmap.zeros(shape, wcs, dtype)
ishape = omap.shape
omap = omap.preflat
ncomp = omap.shape[0]
# In keeping with the rest of the functions here, srcs is [nsrc,{dec,ra,T,Q,U}].
# The beam parameters are ignored - the beam argument is used instead
amps = srcs[:,2:2+ncomp]
poss = srcs[:,:2].copy()
# Rewind positions to let us use flat-sky approximation for distance calculations
#wcs = enmap.enlib.wcs.fix_wcs(wcs)
ref = np.mean(enmap.box(shape, wcs, corner=False)[:,1])
poss[:,1] = utils.rewind(poss[:,1], ref)
beam = expand_beam(beam, nsigma, rmax)
rmax = nsigma2rmax(beam, nsigma)
# Pad our map by rmax, so we get the contribution from sources
# just ourside our area. We will later split our map into cells of size cres. Let's
# adjust the padding so we have a whole number of cells
cres = utils.nint(rmax/omap.pixshape())
epix = cres-(omap.shape[-2:]+2*cres)%cres
padding = [cres,cres+epix]
wmap, wslice = enmap.pad(omap, padding, return_slice=True)
# Overall we will have this many grid cells
cshape = wmap.shape[-2:]/cres
# Find out which sources matter for which cells
srcpix = wmap.sky2pix(poss.T).T
pixbox= np.array([[0,0],wmap.shape[-2:]],int)
nhit, cell_srcs = build_src_cells(pixbox, srcpix, cres)
posmap = wmap.posmap()
model = eval_srcs_loop(posmap, poss, amps, beam, cres, nhit, cell_srcs)
# Update our work map, through our view
if mmul != 1: model *= mmul
wmap += model
if not return_padded:
# Copy out
omap[:] = wmap[wslice]
# Restore shape
omap = omap.reshape(ishape)
return omap
else:
return wmap.reshape(ishape[:-2]+wmap.shape[-2:]), wslice
def random_catalog(shape, wcs, N, edge_avoid_deg=0.):
box = enmap.box(shape, wcs)
dec0 = min(box[0, 0], box[1, 0]) + edge_avoid_deg * np.pi / 180.
dec1 = max(box[0, 0], box[1, 0]) - edge_avoid_deg * np.pi / 180.
ra0 = min(box[0, 1], box[1, 1]) + edge_avoid_deg * np.pi / 180.
ra1 = max(box[0, 1], box[1, 1]) - edge_avoid_deg * np.pi / 180.
ras = np.random.uniform(ra0, ra1, N) * 180. / np.pi
decs = np.random.uniform(dec0, dec1, N) * 180. / np.pi
return ras, decs
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 calc_gridinfo(shape, wcs, steps=[2, 2], nstep=[200, 200], zenith=False):
"""Return an array of line segments representing a coordinate grid
for the given shape and wcs. the steps argument describes the
number of points to use along each meridian."""
steps = np.zeros([2]) + steps
nstep = np.zeros([2], dtype=int) + nstep
gridinfo = Gridinfo()
if enlib.wcs.is_plain(wcs):
box = np.sort(enmap.box(shape, wcs), 0)
start = np.floor(box[0] / steps) * steps
nline = np.floor(box[1] / steps) - np.floor(box[0] / steps) + 1
else:
box = np.array([[-90., 0.], [90., 360.]])
start = np.array([-90., 0.])
nline = np.array([180., 360.]) / steps + 1
gridinfo.lon = []
gridinfo.lat = []
gridinfo.shape = shape[-2:]
gridinfo.wcs = wcs
# Draw lines of longitude
for phi in start[1] + np.arange(nline[1]) * steps[1]:
# Loop over theta
pixs = np.array(
wcs.wcs_world2pix(
phi, np.linspace(box[0, 0], box[1, 0], nstep[0],
endpoint=True), 0)).T
if not enlib.wcs.is_plain(wcs): phi = utils.rewind(phi, 0, 360)
gridinfo.lon.append((phi, calc_line_segs(pixs)))
# Draw lines of latitude
for theta in start[0] + np.arange(nline[0]) * steps[0]:
# Loop over phi
pixs = np.array(
wcs.wcs_world2pix(
np.linspace(box[0, 1],
box[1, 1] + 0.9,
nstep[1],
endpoint=True), theta, 0)).T
if zenith: theta = 90 - theta
gridinfo.lat.append((theta, calc_line_segs(pixs)))
return gridinfo
from enlib import enmap
parser = argparse.ArgumentParser()
parser.add_argument("omap")
parser.add_argument("-D", "--diameter", type=float, default=30)
parser.add_argument("-n", "--npix", type=int, default=800)
parser.add_argument("--proj", type=str, default="cea")
args = parser.parse_args()
r = args.diameter * np.pi / 180 / 2
shape, wcs = enmap.geometry(pos=[[-r, -r], [r, r]],
shape=(args.npix, args.npix),
proj=args.proj)
def linpos(n, b):
return b[0] + (np.arange(n) + 0.5) * (b[1] - b[0]) / n
alpha = enmap.zeros((2, ) + shape, wcs)
pos = enmap.posmap(shape, wcs)
# I'm not sure how to compute this in general, so here's a specialization
# for cylindrical projections
if args.proj == "cea" or args.proj == "car":
dec = pos[0, :, 0]
ra = pos[1, 0, :]
lindec = linpos(shape[0], enmap.box(shape, wcs)[:, 0])
scale = 1 / np.cos(dec) - 1
alpha[0] = (lindec - dec)[:, None]
alpha[1] = ra[None, :] * scale[:, None]
enmap.write_map(args.omap, alpha * 180 * 60 / np.pi)
taper_percent=18.0,
pad_percent=4.0,
weight=None)
pxover = 0.5
for k, index in enumerate(my_tasks):
with bench.show("full_sky"):
fullsky = curvedsky.rand_map(fshape,
fwcs,
ps,
lmax=lmax,
dtype=np.float32)
map_south = fullsky.submap(enmap.box(shape, wcs)) * taper
if k == 0:
if rank == 0:
with bench.show("Rot init..."):
r = fmaps.MapRotatorEquator(shape,
wcs,
wdeg,
hdeg,
width_multiplier=0.6,
height_multiplier=1.2,
downsample=True,
verbose=True,
pix_target_override_arcmin=pxover)
else:
r = fmaps.MapRotatorEquator(shape,
# This program computes a simple estimation of the amount of distortion the
# flat-sky approximation involves
import numpy as np, argparse
from enlib import enmap
parser = argparse.ArgumentParser()
parser.add_argument("omap")
parser.add_argument("-D", "--diameter", type=float, default=30)
parser.add_argument("-n", "--npix", type=int, default=800)
parser.add_argument("--proj", type=str, default="cea")
args = parser.parse_args()
r = args.diameter*np.pi/180/2
shape, wcs = enmap.geometry(pos=[[-r,-r],[r,r]], shape=(args.npix, args.npix), proj=args.proj)
def linpos(n,b): return b[0] + (np.arange(n)+0.5)*(b[1]-b[0])/n
alpha = enmap.zeros((2,)+shape, wcs)
pos = enmap.posmap(shape, wcs)
# I'm not sure how to compute this in general, so here's a specialization
# for cylindrical projections
if args.proj == "cea" or args.proj =="car":
dec = pos[0,:,0]
ra = pos[1,0,:]
lindec= linpos(shape[0],enmap.box(shape,wcs)[:,0])
scale = 1/np.cos(dec)-1
alpha[0] = (lindec-dec)[:,None]
alpha[1] = ra[None,:]*scale[:,None]
enmap.write_map(args.omap, alpha*180*60/np.pi)
import numpy as np, argparse
from enlib import enmap, utils
parser = argparse.ArgumentParser()
parser.add_argument("imap")
parser.add_argument("template")
parser.add_argument("omap")
parser.add_argument("-O", "--order", type=int, default=3)
parser.add_argument("-m", "--mode", type=str, default="constant")
#parser.add_argument("-M", "--mem", type=float, default=1e8)
parser.add_argument("-I", "--ivar", action="store_true")
args = parser.parse_args()
shape, wcs = enmap.read_map_geometry(args.template)
box = enmap.box(shape, wcs)
box = utils.widen_box(box, 1 * utils.degree, relative=False)
imap = enmap.read_map(args.imap, box=box)
if args.ivar: imap /= imap.pixsizemap()
omap = imap.project(shape, wcs, order=args.order, mode=args.mode)
if args.ivar: omap *= omap.pixsizemap()
enmap.write_map(args.omap, omap)
#blockpix = np.product(shape[:-2])*shape[-1]
#bsize = max(1,utils.nint(args.mem/(blockpix*imap.dtype.itemsize)))
#
#nblock = (shape[-2]+bsize-1)//bsize
#for b in range(nblock):
# r1, r2 = b*bsize, (b+1)*bsize
# osub = omap[...,r1:r2,:]
# omap[...,r1:r2,:] = enmap.project(imap, osub.shape, osub.wcs, order=args.order, mode=args.mode)
##o = enmap.project(m, t.shape, t.wcs, order=args.order, mode=args.mode)
# Get MPI comm
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
numcores = comm.Get_size()
Ntot = 32
num_each, each_tasks = mpi_distribute(Ntot, numcores)
mpibox = MPIStats(comm, num_each, tag_start=333)
if rank == 0: print(("At most ", max(num_each), " tasks..."))
my_tasks = each_tasks[rank]
deg = 40.
px = 2.0
shape, wcs = enmap.rect_geometry(deg * 60., px, proj="car")
bbox = enmap.box(shape, wcs)
bin_edges = np.arange(100, 4000, 40)
theory_file_root = "data/Aug6_highAcc_CDM"
theory = cmb.loadTheorySpectraFromCAMB(theory_file_root,
unlensedEqualsLensed=False,
useTotal=False,
TCMB=2.7255e6,
lpad=9000,
get_dimensionless=False)
for k, i in enumerate(my_tasks):
map_file = "/gpfs01/astro/workarea/msyriac/data/sims/sigurd/cori/v6full/fullsky_curved_lensed_car_" + str(
i).zfill(2) + ".fits"
def enmaps_from_config(Config, sim_section, analysis_section, pol=False):
"""
Algorithm for deciding sim and analysis shapes and wcs:
Check if user has specified a *data* template
* If yes, use its shape and wcs for the data
- Determine ratio simpixel/anpixel
- Upsample by that ratio to make sim template
* If no,
"""
pixel_sim = Config.getfloat(sim_section, "pixel_arcmin")
buffer_sim = Config.getfloat(sim_section, "buffer")
projection = Config.get(analysis_section, "projection")
try:
pt_file = Config.get(analysis_section, "patch_template")
imap = enmap.read_map(pt_file)
shape_dat = imap.shape
wcs_dat = imap.wcs
res = np.min(imap.extent() / imap.shape[-2:]) * 60. * 180. / np.pi
if np.isclose(pixel_sim, res, 1.e-2):
shape_sim = shape_dat
wcs_sim = wcs_dat
else:
bbox = enmap.box(shape_dat, wcs_dat)
shape_sim, wcs_sim = enmap.geometry(bbox,
res=pixel_sim * np.pi / 180. /
60.,
proj=projection)
except:
pixel_analysis = Config.getfloat(analysis_section, "pixel_arcmin")
try:
width_analysis_deg = Config.getfloat(analysis_section,
"patch_degrees_width")
except:
width_analysis_deg = Config.getfloat(analysis_section,
"patch_arcmin_width") / 60.
try:
height_analysis_deg = Config.getfloat(analysis_section,
"patch_degrees_height")
except:
height_analysis_deg = Config.getfloat(analysis_section,
"patch_arcmin_height") / 60.
ra_offset = Config.getfloat(analysis_section, "ra_offset")
dec_offset = Config.getfloat(analysis_section, "dec_offset")
shape_dat, wcs_dat = maps.rect_geometry(
width_analysis_deg * 60.,
pixel_analysis,
proj=projection,
pol=pol,
height_arcmin=height_analysis_deg * 60.,
xoffset_degree=ra_offset,
yoffset_degree=dec_offset)
if np.abs(buffer_sim - 1.) < 1.e-3:
shape_sim, wcs_sim = maps.rect_geometry(
width_analysis_deg * 60.,
pixel_sim,
proj=projection,
pol=pol,
height_arcmin=height_analysis_deg * 60.,
xoffset_degree=ra_offset,
yoffset_degree=dec_offset)
else:
raise NotImplementedError("Buffer !=1 not implemented")
return shape_sim, wcs_sim, shape_dat, wcs_dat
from __future__ import division, print_function
import numpy as np, argparse, os
from enlib import retile, mpi, utils, enmap
parser = argparse.ArgumentParser()
parser.add_argument("imap")
parser.add_argument("template", nargs="?")
parser.add_argument("odir")
parser.add_argument("-t", "--tsize", type=int, default=675)
parser.add_argument("-s", "--slice", type=str, default=None)
args = parser.parse_args()
tfile = args.template or args.imap
# Find the bounds of our tiling
shape, wcs = enmap.read_map_geometry(tfile)
box = enmap.box(shape, wcs, corner=False)
tshape = (args.tsize, args.tsize)
ntile = tuple([(s + t - 1) // t for s, t in zip(shape[-2:], tshape)])
utils.mkdir(args.odir)
opathfmt = args.odir + "/tile%(y)03d_%(x)03d.fits"
retile.retile(args.imap,
opathfmt,
otilenum=ntile,
ocorner=box[0],
otilesize=tshape,
verbose=True,
slice=args.slice)
