# Path: glimix_core/mean/_kron.py
# class KronMean(Function):
#     """
#     Kronecker mean function, (A⊗X)vec(B).
# 
#     Let
# 
#     - n be the number of samples;
#     - p the number of traits; and
#     - c the number of covariates.
# 
#     The mathematical representation is
# 
#         𝐦 = (A⊗X)vec(B)
# 
#     where A is a p×p trait design matrix of fixed effects and X is a n×c sample design
#     matrix of fixed effects. B is a c×p matrix of fixed-effect sizes.
#     """
# 
#     def __init__(self, A, X):
#         """
#         Constructor.
# 
#         Parameters
#         ----------
#         A : array_like
#             p×p array.
#         X : array_like
#             n×c array.
#         """
#         self._A = asarray(A, float)
#         self._X = asarray(X, float)
#         vecB = zeros((X.shape[1], A.shape[0])).ravel()
#         self._vecB = Vector(vecB)
#         self._nparams = vecB.size
#         Function.__init__(self, "KronMean", vecB=self._vecB)
# 
#     @property
#     def nparams(self):
#         """
#         Number of parameters.
#         """
#         return self._nparams
# 
#     @property
#     def A(self):
#         """
#         Matrix A.
#         """
#         return self._A
# 
#     @property
#     def X(self):
#         """
#         Matrix X.
#         """
#         return self._X
# 
#     @property
#     def AX(self):
#         """
#         A ⊗ X.
#         """
#         return kron(self.A, self.X)
# 
#     def value(self):
#         """
#         Kronecker mean function.
# 
#         Returns
#         -------
#         𝐦 : ndarray
#             (A⊗X)vec(B).
#         """
#         return self.AX @ self._vecB.value
# 
#     def gradient(self):
#         """
#         Gradient of the linear mean function.
# 
#         Returns
#         -------
#         vecB : ndarray
#             Derivative of M over vec(B).
#         """
#         return {"vecB": self.AX}
# 
#     @property
#     def B(self):
#         """
#         Effect-sizes parameter, B.
#         """
#         return unvec(self._vecB.value, (self.X.shape[1], self.A.shape[0]))
# 
#     @B.setter
#     def B(self, v):
#         self._vecB.value = vec(asarray(v, float))
# 
#     def __str__(self):
#         tname = type(self).__name__
#         msg = "{}(A=..., X=...)".format(tname)
#         if self.name is not None:
#             msg += ": {}".format(self.name)
#         msg += "\n"
#         mat = format(self.B)
#         msg += "  B: " + "\n     ".join(mat.split("\n"))
#         return msg


from numpy import kron, ravel
from numpy.random import RandomState
from numpy.testing import assert_allclose
from glimix_core.mean import KronMean



def test_mean_kron():

    random = RandomState(0)
    # number of trais
    p = 2
    # number of covariates
    c = 3
    # sample size
    n = 4

    A = random.randn(p, p)
    X = random.randn(n, c)
    B = random.randn(p, c)

    mean = KronMean(A, X)
    mean.B = B
    assert_allclose(mean.value(), kron(A, X) @ ravel(B, order="F"))

# Path: glimix_core/mean/_kron.py
# class KronMean(Function):
#     """
#     Kronecker mean function, (A⊗X)vec(B).
# 
#     Let
# 
#     - n be the number of samples;
#     - p the number of traits; and
#     - c the number of covariates.
# 
#     The mathematical representation is
# 
#         𝐦 = (A⊗X)vec(B)
# 
#     where A is a p×p trait design matrix of fixed effects and X is a n×c sample design
#     matrix of fixed effects. B is a c×p matrix of fixed-effect sizes.
#     """
# 
#     def __init__(self, A, X):
#         """
#         Constructor.
# 
#         Parameters
#         ----------
#         A : array_like
#             p×p array.
#         X : array_like
#             n×c array.
#         """
#         self._A = asarray(A, float)
#         self._X = asarray(X, float)
#         vecB = zeros((X.shape[1], A.shape[0])).ravel()
#         self._vecB = Vector(vecB)
#         self._nparams = vecB.size
#         Function.__init__(self, "KronMean", vecB=self._vecB)
# 
#     @property
#     def nparams(self):
#         """
#         Number of parameters.
#         """
#         return self._nparams
# 
#     @property
#     def A(self):
#         """
#         Matrix A.
#         """
#         return self._A
# 
#     @property
#     def X(self):
#         """
#         Matrix X.
#         """
#         return self._X
# 
#     @property
#     def AX(self):
#         """
#         A ⊗ X.
#         """
#         return kron(self.A, self.X)
# 
#     def value(self):
#         """
#         Kronecker mean function.
# 
#         Returns
#         -------
#         𝐦 : ndarray
#             (A⊗X)vec(B).
#         """
#         return self.AX @ self._vecB.value
# 
#     def gradient(self):
#         """
#         Gradient of the linear mean function.
# 
#         Returns
#         -------
#         vecB : ndarray
#             Derivative of M over vec(B).
#         """
#         return {"vecB": self.AX}
# 
#     @property
#     def B(self):
#         """
#         Effect-sizes parameter, B.
#         """
#         return unvec(self._vecB.value, (self.X.shape[1], self.A.shape[0]))
# 
#     @B.setter
#     def B(self, v):
#         self._vecB.value = vec(asarray(v, float))
# 
#     def __str__(self):
#         tname = type(self).__name__
#         msg = "{}(A=..., X=...)".format(tname)
#         if self.name is not None:
#             msg += ": {}".format(self.name)
#         msg += "\n"
#         mat = format(self.B)
#         msg += "  B: " + "\n     ".join(mat.split("\n"))
#         return msg


# Path: glimix_core/mean/test/test_mean_kron.py
from numpy import kron, ravel
from numpy.random import RandomState
from numpy.testing import assert_allclose
from glimix_core.mean import KronMean




def test_mean_kron():

    random = RandomState(0)
    # number of trais
    p = 2
    # number of covariates
    c = 3
    # sample size
    n = 4

    A = random.randn(p, p)
    X = random.randn(n, c)
    B = random.randn(p, c)

    mean = KronMean(A, X)
    mean.B = B
    assert_allclose(mean.value(), kron(A, X) @ ravel(B, order="F"))

    assert_allclose(mean._check_grad(), 0, atol=1e-5)