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 estimate_SN(d, fparams, src_groups):
gmax = max([len(g) for g in src_groups])
nsrc, ncomp = fparams[:, 2:-3].shape
SN = np.zeros([nsrc, ncomp])
for i in range(gmax):
for c in range(ncomp):
flat = fparams.copy()
flat[:, 2:-3] = 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.copy()
ptsrc_data.pmat_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 ri in d.rangesets[d.offsets[si, 0]:d.offsets[si, -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 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)
ptsrc_data.pmat_model(tod, pflat, d, dir=-1)
rhs = pflat[:, 2:-3].copy()
dof = DOF(Arg(mask=mask))
# Set up functional form of icov
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])
# 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 apply_model(tod, pflat, d, dir=1):
ptsrc_data.pmat_model(tod, pflat, d, dir=dir)
def subtract_model(tod, d, fparams):
mtod = tod.astype(dtype, copy=True)
p = fparams.copy()
ptsrc_data.pmat_model(mtod, p, d)
return tod - mtod
def apply_model(tod, pflat, d, dir=1): ptsrc_data.pmat_model(tod, pflat, d, dir=dir)