def get_group_zellner(groups, X, isgmom=False):
"""Note that V=(XtX)^-1 and Vinv=XtX."""
n, p = X.shape
Vinv = jnp.zeros((p, p))
V = jnp.zeros((p, p))
for group, p_j in zip(*jnp.unique(groups, return_counts=True)):
mask = jnp.arange(p)[groups == group]
X_j = X[:, mask]
p_term = cond(isgmom, p_j, lambda x: x, p_j, lambda x: x + 2)
aux = jnp.dot(X_j.T, X_j) * n / p_term
Vinv = Vinv.at[jnp.ix_(mask, mask)].set(aux)
V = V.at[jnp.ix_(mask, mask)].set(jnp.linalg.inv(aux))
return V, Vinv
def reweighted_metrics(weights, ell, ngals, noise, cl_in, gals_per_arcmin2, fsky):
"""The sparse cinv calculation is currently the bottleneck.
"""
cl_out = reweight_cl(weights, ngals, cl_in)
nl_out = reweight_noise_cl(weights, ngals, noise, gals_per_arcmin2)
cov_out = gaussian_cl_covariance(ell, cl_out[-1], nl_out, fsky)
cinv = sparse.inv(cov_out)
mu = cl_out[-1].reshape(-1, 1)
dmu = cl_out[:-1].reshape(7, -1)
F = sparse.dot(dmu, cinv, dmu.T)
Finv = jnp.linalg.inv(F)
return {
'SNR_3x2': jnp.sqrt(sparse.dot(mu.T, cinv, mu)[0, 0]),
'FOM_3x2': 1 / (6.17 * jnp.pi * jnp.sqrt(jnp.linalg.det(Finv[jnp.ix_([0,4], [0,4])]))),
'FOM_DETF_3x2': 1 / (6.17 * jnp.pi * jnp.sqrt(jnp.linalg.det(Finv[jnp.ix_([5,6], [5,6])])))
}
def cartesian_product(*arrays):
"""
IN: any number of np arrays of same length
OUT: cartesian product of the arrays
"""
la = len(arrays)
dtype = np.result_type(*arrays)
arr = np.empty([len(a) for a in arrays] + [la], dtype=dtype)
for i, a in enumerate(np.ix_(*arrays)):
# arr[...,i] = a
arr = index_update(arr, index[..., i], a)
return arr.reshape(-1, la)
def cartesian_product(*arrays) -> jnp.ndarray:
"""calculate the n-dimensional cartesian product, i.e. create all
possible combinations of all elements in a given collection of arrays.
Args:
*arrays: the arrays to calculate the cartesian product for
Returns:
the cartesian product.
"""
ixarrays = jnp.ix_(*arrays)
barrays = jnp.broadcast_arrays(*ixarrays)
sarrays = jnp.stack(barrays, -1)
product = sarrays.reshape(-1, sarrays.shape[-1])
return product
def ix_(*args): args = [_remove_jaxarray(a) for a in args] return jnp.ix_(*args)