def estimate_SN(d, fparams, src_groups):
gmax = max([len(g) for g in src_groups])
nsrc, ncomp = fparams[:, 2:5].shape
SN = np.zeros([nsrc, ncomp])
for i in range(gmax):
for c in range(ncomp):
flat = fparams.copy()
flat[:, 2:5] = 0
for g in src_groups:
if len(g) <= i: continue
flat[g[i], c + 2] = fparams[g[i], c + 2]
# Do all the compatible sources in parallel
mtod = d.tod * 0
apply_model(mtod, flat, d)
ntod = mtod.copy()
ptsrc_data.nmat_basis(ntod, d)
# And then extract S/N for each of them
for g in src_groups:
if len(g) <= i: continue
si = g[i]
my_sn = 0
for off in d.offsets[si]:
for ri in d.rangesets[off[0]:off[1]]:
r = d.ranges[ri]
my_sn += np.sum(ntod[r[0]:r[1]] * mtod[r[0]:r[1]])
SN[si, c] = my_sn
return SN
def make_maps(tod, data, pos, ncomp, radius, resolution):
tod = tod.copy()
pos = np.array(pos)
# Handle angle wrapping
pos = utils.rewind(pos, data.ref)
nsrc = len(pos)
dbox = np.array([[-1, -1], [1, 1]]) * radius
shape, wcs = enmap.geometry(pos=dbox, res=resolution)
# Set up pixels
n = int(np.round(2 * radius / resolution))
boxes = np.array([[p - radius, p + radius] for p in pos])
# Set up output maps
rhs = enmap.zeros((nsrc, ncomp) + shape, wcs, dtype=dtype)
div = enmap.zeros((ncomp, nsrc, ncomp) + shape, wcs, dtype=dtype)
# Build rhs
ptsrc_data.nmat_basis(tod, data)
pmat_thumbs(-1, tod, rhs, boxes, data)
# Build div
for c in range(ncomp):
idiv = div[0].copy()
idiv[:, c] = 1
wtod = data.tod.astype(dtype, copy=True)
wtod[...] = 0
pmat_thumbs(1, wtod, idiv, boxes, data)
ptsrc_data.nmat_basis(wtod, data, white=True)
pmat_thumbs(-1, wtod, div[c], boxes, data)
div = np.rollaxis(div, 1)
mask = div[:, 0] != 0
bin = rhs.copy()
bin[mask] = rhs[mask] / div[:, 0][mask] # Fixme: only works for ncomp == 1
return bin, rhs, div
def make_maps(tod, data, pos, ncomp, radius, resolution): tod = tod.copy() pos = np.array(pos) # Handle angle wrapping pos = utils.rewind(pos, data.ref) nsrc= len(pos) dbox= np.array([[-1,-1],[1,1]])*radius shape, wcs = enmap.geometry(pos=dbox, res=resolution) # Set up pixels n = int(np.round(2*radius/resolution)) boxes = np.array([[p-radius,p+radius] for p in pos]) # Set up output maps rhs = enmap.zeros((nsrc,ncomp) +shape, wcs, dtype=dtype) div = enmap.zeros((ncomp,nsrc,ncomp)+shape, wcs, dtype=dtype) # Build rhs ptsrc_data.nmat_basis(tod, data) pmat_thumbs(-1, tod, rhs, boxes, data) # Build div for c in range(ncomp): idiv = div[0].copy(); idiv[:,c] = 1 wtod = data.tod.astype(dtype,copy=True); wtod[...] = 0 pmat_thumbs( 1, wtod, idiv, boxes, data) ptsrc_data.nmat_basis(wtod, data, white=True) pmat_thumbs(-1, wtod, div[c], boxes, data) div = np.rollaxis(div,1) mask = div[:,0] != 0 bin = rhs.copy() bin[mask] = rhs[mask]/div[:,0][mask] # Fixme: only works for ncomp == 1 return bin, rhs, div
def icov_fun(x):
p = pflat.copy()
p[:, 2:-3], = dof.unzip(x)
ptsrc_data.pmat_model(tod, p, d, dir=+1)
ptsrc_data.nmat_basis(tod, d)
ptsrc_data.pmat_model(tod, p, d, dir=-1)
return dof.zip(p[:, 2:-3])
def calc_amp_dist(tod, d, params, mask=None): if mask is None: mask = params.strong if np.sum(mask) == 0: return AmpDist(np.zeros([0,0]), np.zeros([0]), DOF(Arg(mask=mask))) # rhs = P'N"d tod = tod.astype(dtype, copy=True) pflat = params.flat.copy() ptsrc_data.nmat_basis(tod, d) apply_model(tod, pflat, d, dir=-1) rhs = pflat[:,2:5].copy() dof = DOF(Arg(mask=mask)) # Set up functional form of icov def icov_fun(x): p = pflat.copy() p[:,2:5], = dof.unzip(x) tod[:] = 0 apply_model(tod, p, d, dir=+1) ptsrc_data.nmat_basis(tod, d) apply_model(tod, p, d, dir=-1) return dof.zip(p[:,2:5]) # Build A matrix in parallel. When using more than # one component, the ndof will be twice the number of sources, so # groups must be modified dgroups = groups_to_dof(params.groups, dof) icov = np.zeros([dof.n,dof.n]) nmax = max([len(g) for g in dgroups]) for i in range(nmax): # Loop through the elements of the uncorrelated groups in parallel u = np.zeros(dof.n) u[[g[i] for g in dgroups if len(g) > i]] = 1 icov_u = icov_fun(u) # Extract result into full A for g in dgroups: if len(g) > i: icov[g[i],g] = icov_u[g] return AmpDist(icov, dof.zip(rhs), dof)
def estimate_SN(d, fparams, src_groups): gmax = max([len(g) for g in src_groups]) nsrc, ncomp = fparams[:,2:5].shape SN = np.zeros([nsrc,ncomp]) for i in range(gmax): for c in range(ncomp): flat = fparams.copy() flat[:,2:5] = 0 for g in src_groups: if len(g) <= i: continue flat[g[i],c+2] = fparams[g[i],c+2] # Do all the compatible sources in parallel mtod = d.tod*0 apply_model(mtod, flat, d) ntod = mtod.copy() ptsrc_data.nmat_basis(ntod, d) # And then extract S/N for each of them for g in src_groups: if len(g) <= i: continue si = g[i] my_sn = 0 for off in d.offsets[si]: for ri in d.rangesets[off[0]:off[1]]: r = d.ranges[ri] my_sn += np.sum(ntod[r[0]:r[1]]*mtod[r[0]:r[1]]) SN[si,c] = my_sn return SN
def icov_fun(x):
p = pflat.copy()
p[:, 2:5], = dof.unzip(x)
tod[:] = 0
apply_model(tod, p, d, dir=+1)
ptsrc_data.nmat_basis(tod, d)
apply_model(tod, p, d, dir=-1)
return dof.zip(p[:, 2:5])
def icov_fun(x): p = pflat.copy() p[:,2:5], = dof.unzip(x) tod[:] = 0 apply_model(tod, p, d, dir=+1) ptsrc_data.nmat_basis(tod, d) apply_model(tod, p, d, dir=-1) return dof.zip(p[:,2:5])
def calc_marginal_amps_strong(d, p): # The marginal probability -2 log P(pos|beam,aw) = (d-Pw aw)'N"(d-Pw aw) - as' As" as # where as = (Ps'N"Ps)"Ps'N"(d-Pw aw). First compute as and As. p_weak = p.flat; p_weak[:,2:5][p.strong] = 0 tod_rest = subtract_model(d.tod, d, p_weak) # calc_amp_dist only uses strong dof by default adist_strong = calc_amp_dist(tod_rest, d, p) x_s, Ai_s = adist_strong.x, adist_strong.Ai P_s = 0.5*np.sum(x_s*Ai_s.dot(x_s)) # The remainder is tod_rest'N"tod_rest ntod_rest = tod_rest.copy() ptsrc_data.nmat_basis(ntod_rest, d) P_w = -0.5*np.sum(tod_rest*ntod_rest) return P_s, P_w, adist_strong
def calc_marginal_amps_strong(d, p):
# The marginal probability -2 log P(pos|beam,aw) = (d-Pw aw)'N"(d-Pw aw) - as' As" as
# where as = (Ps'N"Ps)"Ps'N"(d-Pw aw). First compute as and As.
p_weak = p.flat
p_weak[:, 2:5][p.strong] = 0
tod_rest = subtract_model(d.tod, d, p_weak)
# calc_amp_dist only uses strong dof by default
adist_strong = calc_amp_dist(tod_rest, d, p)
x_s, Ai_s = adist_strong.x, adist_strong.Ai
P_s = 0.5 * np.sum(x_s * Ai_s.dot(x_s))
# The remainder is tod_rest'N"tod_rest
ntod_rest = tod_rest.copy()
ptsrc_data.nmat_basis(ntod_rest, d)
P_w = -0.5 * np.sum(tod_rest * ntod_rest)
return P_s, P_w, adist_strong
def calc_amp_dist(tod, d, params, mask=None):
if mask is None: mask = params.strong
if np.sum(mask) == 0:
return AmpDist(np.zeros([0, 0]), np.zeros([0]), DOF(Arg(mask=mask)))
# rhs = P'N"d
tod = tod.astype(dtype, copy=True)
pflat = params.flat.copy()
ptsrc_data.nmat_basis(tod, d)
apply_model(tod, pflat, d, dir=-1)
rhs = pflat[:, 2:5].copy()
dof = DOF(Arg(mask=mask))
# Set up functional form of icov
def icov_fun(x):
p = pflat.copy()
p[:, 2:5], = dof.unzip(x)
tod[:] = 0
apply_model(tod, p, d, dir=+1)
ptsrc_data.nmat_basis(tod, d)
apply_model(tod, p, d, dir=-1)
return dof.zip(p[:, 2:5])
# Build A matrix in parallel. When using more than
# one component, the ndof will be twice the number of sources, so
# groups must be modified
dgroups = groups_to_dof(params.groups, dof)
icov = np.zeros([dof.n, dof.n])
nmax = max([len(g) for g in dgroups])
for i in range(nmax):
# Loop through the elements of the uncorrelated groups in parallel
u = np.zeros(dof.n)
u[[g[i] for g in dgroups if len(g) > i]] = 1
icov_u = icov_fun(u)
# Extract result into full A
for g in dgroups:
if len(g) > i:
icov[g[i], g] = icov_u[g]
return AmpDist(icov, dof.zip(rhs), dof)
def calc_posterior(tod, d, fparams):
wtod = subtract_model(tod, d, fparams)
ntod = wtod.copy()
ptsrc_data.nmat_basis(ntod, d)
return -0.5 * np.sum(ntod * wtod)
def calc_posterior(tod, d, fparams): wtod = subtract_model(tod, d, fparams) ntod = wtod.copy() ptsrc_data.nmat_basis(ntod, d) return -0.5*np.sum(ntod*wtod)
