def test_kron2sum_large_outcome():
random = RandomState(2)
n = 50
A = random.randn(3, 3)
A = A @ A.T
F = random.randn(n, 2)
G = random.randn(n, 4)
B = random.randn(2, 3)
C0 = random.randn(3, 3)
C0 = C0 @ C0.T
C1 = random.randn(3, 3)
C1 = C1 @ C1.T
K = kron(C0, (G @ G.T)) + kron(C1, eye(n))
y = multivariate_normal(random, kron(A, F) @ vec(B), K)
Y = unvec(y, (n, 3))
Y = Y / Y.std(0)
lmm = Kron2Sum(Y, A, F, G, restricted=False)
lmm.fit(verbose=False)
assert_allclose(lmm.lml(), -12.163158697588926)
assert_allclose(lmm.C0[0, 1], -0.004781646218546575, rtol=1e-3, atol=1e-5)
assert_allclose(lmm.C1[0, 1], 0.03454122242999587, rtol=1e-3, atol=1e-5)
assert_allclose(lmm.beta[2], -0.02553979383437496, rtol=1e-3, atol=1e-5)
assert_allclose(lmm.beta_covariance[0, 1],
0.0051326042358990865,
rtol=1e-3,
atol=1e-5)
assert_allclose(lmm.mean()[3], 0.3442913781854699, rtol=1e-2, atol=1e-5)
assert_allclose(lmm.covariance()[0, 1],
0.0010745698663887468,
rtol=1e-3,
atol=1e-5)
def scan(self, A1, X1):
"""
LML, fixed-effect sizes, and scale of the candidate set.
Parameters
----------
A1 : (p, e) array_like
Trait-by-environments design matrix.
X1 : (n, m) array_like
Variants set matrix.
Returns
-------
lml : float
Log of the marginal likelihood for the set.
effsizes0 : (c, p) ndarray
Fixed-effect sizes for the covariates.
effsizes0_se : (c, p) ndarray
Fixed-effect size standard errors for the covariates.
effsizes1 : (m, e) ndarray
Fixed-effect sizes for the candidates.
effsizes1_se : (m, e) ndarray
Fixed-effect size standard errors for the candidates.
scale : float
Optimal scale.
"""
from numpy import empty
from numpy.linalg import multi_dot
from numpy_sugar import epsilon, is_all_finite
from scipy.linalg import cho_solve
A1 = asarray(A1, float)
X1 = asarray(X1, float)
if not is_all_finite(A1):
raise ValueError("A1 parameter has non-finite elements.")
if not is_all_finite(X1):
raise ValueError("X1 parameter has non-finite elements.")
if A1.shape[1] == 0:
beta_se = sqrt(self.null_beta_covariance.diagonal())
return {
"lml": self.null_lml(),
"effsizes0": unvec(self.null_beta, (self._ncovariates, -1)),
"effsizes0_se": unvec(beta_se, (self._ncovariates, -1)),
"effsizes1": empty((0, )),
"effsizes1_se": empty((0, )),
"scale": self.null_scale,
}
X1X1 = X1.T @ X1
XX1 = self._X.T @ X1
AWA1 = self._WA.T @ A1
A1W = A1.T @ self._W
GX1 = self._G.T @ X1
MRiM1 = kron(AWA1, XX1)
M1RiM1 = kron(A1W @ A1, X1X1)
M1Riy = vec(multi_dot([X1.T, self._Y, A1W.T]))
XRiM1 = kron(self._WL0.T @ A1, GX1)
ZiXRiM1 = cho_solve(self._Lz, XRiM1)
MRiXZiXRiM1 = self._XRiM.T @ ZiXRiM1
M1RiXZiXRiM1 = XRiM1.T @ ZiXRiM1
M1RiXZiXRiy = XRiM1.T @ self._ZiXRiy
T0 = [[self._MRiM, MRiM1], [MRiM1.T, M1RiM1]]
T1 = [[self._MRiXZiXRiM, MRiXZiXRiM1], [MRiXZiXRiM1.T, M1RiXZiXRiM1]]
T2 = [self._MRiy, M1Riy]
T3 = [self._MRiXZiXRiy, M1RiXZiXRiy]
MKiM = block(T0) - block(T1)
MKiy = block(T2) - block(T3)
beta = rsolve(MKiM, MKiy)
mKiy = beta.T @ MKiy
cp = self._ntraits * self._ncovariates
effsizes0 = unvec(beta[:cp], (self._ncovariates, self._ntraits))
effsizes1 = unvec(beta[cp:], (X1.shape[1], A1.shape[1]))
np = self._nsamples * self._ntraits
sqrtdot = self._yKiy - mKiy
scale = clip(sqrtdot / np, epsilon.tiny, inf)
lml = self._static_lml() / 2 - np * safe_log(scale) / 2 - np / 2
effsizes_se = sqrt(
clip(scale * pinv(MKiM).diagonal(), epsilon.tiny, inf))
effsizes0_se = unvec(effsizes_se[:cp],
(self._ncovariates, self._ntraits))
effsizes1_se = unvec(effsizes_se[cp:], (X1.shape[1], A1.shape[1]))
return {
"lml": lml,
"effsizes0": effsizes0,
"effsizes1": effsizes1,
"scale": scale,
"effsizes0_se": effsizes0_se,
"effsizes1_se": effsizes1_se,
}
def _terms(self):
from numpy_sugar.linalg import ddot, lu_slogdet, sum2diag
if self._cache["terms"] is not None:
return self._cache["terms"]
L0 = self._cov.C0.L
S, U = self._cov.C1.eigh()
W = ddot(U, 1 / S) @ U.T
S = 1 / sqrt(S)
Y = self._Y
A = self._mean.A
WL0 = W @ L0
YW = Y @ W
WA = W @ A
L0WA = L0.T @ WA
Z = kron(L0.T @ WL0, self._GG)
Z = sum2diag(Z, 1)
Lz = lu_factor(Z, check_finite=False)
# Юљ▓рхђRРЂ╗┬╣Юљ▓ = vec(YW)рхђЮљ▓
yRiy = (YW * self._Y).sum()
# MрхђRРЂ╗┬╣M = AрхђWA РіЌ XрхђX
MRiM = kron(A.T @ WA, self._XX)
# XрхђRРЂ╗┬╣Юљ▓ = vec(GрхђYWLРѓђ)
XRiy = vec(self._GY @ WL0)
# XрхђRРЂ╗┬╣M = (LРѓђрхђWA) РіЌ (GрхђX)
XRiM = kron(L0WA, self._GX)
# MрхђRРЂ╗┬╣Юљ▓ = vec(XрхђYWA)
MRiy = vec(self._XY @ WA)
ZiXRiM = lu_solve(Lz, XRiM)
ZiXRiy = lu_solve(Lz, XRiy)
MRiXZiXRiy = ZiXRiM.T @ XRiy
MRiXZiXRiM = XRiM.T @ ZiXRiM
yKiy = yRiy - XRiy @ ZiXRiy
MKiy = MRiy - MRiXZiXRiy
H = MRiM - MRiXZiXRiM
Lh = lu_factor(H, check_finite=False)
b = lu_solve(Lh, MKiy)
B = unvec(b, (self.ncovariates, -1))
self._mean.B = B
XRim = XRiM @ b
ZiXRim = ZiXRiM @ b
mRiy = b.T @ MRiy
mRim = b.T @ MRiM @ b
logdetK = lu_slogdet(Lz)[1]
logdetK -= 2 * log(S).sum() * self.nsamples
mKiy = mRiy - XRim.T @ ZiXRiy
mKim = mRim - XRim.T @ ZiXRim
self._cache["terms"] = {
"logdetK": logdetK,
"mKiy": mKiy,
"mKim": mKim,
"b": b,
"Z": Z,
"B": B,
"Lz": Lz,
"S": S,
"W": W,
"WA": WA,
"YW": YW,
"WL0": WL0,
"yRiy": yRiy,
"MRiM": MRiM,
"XRiy": XRiy,
"XRiM": XRiM,
"ZiXRiM": ZiXRiM,
"ZiXRiy": ZiXRiy,
"ZiXRim": ZiXRim,
"MRiy": MRiy,
"mRim": mRim,
"mRiy": mRiy,
"XRim": XRim,
"yKiy": yKiy,
"H": H,
"Lh": Lh,
"MRiXZiXRiy": MRiXZiXRiy,
"MRiXZiXRiM": MRiXZiXRiM,
}
return self._cache["terms"]
def B(self):
"""
Effect-sizes parameter, B.
"""
return unvec(self._vecB.value, (self.X.shape[1], self.A.shape[0]))