psindex: Parameter Stability Index

# the multivariate normal distribution, unrestricted (h1)
# - everything is evalued under the MLEs: Mu = ybar, Sigma = S 

# 1) loglikelihood h1 (from raw data, or sample statistics)
# 4) hessian h1 around MLEs
# 5) information h1 (restricted Sigma/mu)
#    5a: (unit)    expected information h1 (A1 = Gamma.NT^{-1})
#    5b: (unit)    observed information h1 (A1 = Gamma.NT^{-1})
#    5c: (unit) first.order information h1 (B1 = A1 %*% Gamma %*% A1)
# 6) inverted information h1 mu + vech(Sigma)
#    6a: (unit) inverted expected information (A1.inv = Gamma.NT)
#    6b: (unit) inverted observed information (A1.inv = Gamma.NT)
#    6c: (unit) inverted first-order information (B1.inv)
# 7) ACOV h1 mu + vech(Sigma)
#    7a: 1/N * Gamma.NT
#    7b: 1/N * Gamma.NT
#    7c: 1/N * (Gamma.NT * Gamma^{-1} * Gamma.NT)
#    7d: 1/N * Gamma (sandwich)


# YR 25 Mar 2016: first version
# YR 19 Jan 2017: added 6) + 7)

# 1. likelihood h1

# 1a: input is raw data
lav_mvnorm_h1_loglik_data <- function(Y             = NULL,
                                      x.idx         = NULL,
                                      casewise      = FALSE,
                                      wt            = NULL,
                                      Sinv.method   = "eigen") {
    P <- NCOL(Y); N <- NROW(Y)

    # sample statistics
    if(is.null(wt)) {
        sample.mean <- colMeans(Y)
        sample.cov <- 1/N*crossprod(Y) - tcrossprod(sample.mean)
    } else {
        out <- stats::cov.wt(Y, wt = wt, method = "ML")
        if(casewise) {
            loglik <- lav_mvnorm_loglik_data(Y, Mu = out$center, 
                          Sigma = out$cov, casewise = TRUE, 
                          Sinv.method = Sinv.method)
            return(loglik * wt)
        } else {
            sample.mean <- out$center
            sample.cov  <- out$cov
        }
    }

    if(casewise) {
        LOG.2PI <- log(2 * pi)

        # invert sample.cov
        if(Sinv.method == "chol") {
            cS <- chol(sample.cov); icS <- backsolve(cS, diag(P))
            Yc <- t( t(Y) - sample.mean )
            DIST <- rowSums((Yc %*% icS)^2)
            logdet <- -2 * sum(log(diag(icS)))
        } else {
            sample.cov.inv <- lav_matrix_symmetric_inverse(S = sample.cov, 
                                  logdet = TRUE, Sinv.method = Sinv.method)
            logdet <- attr(sample.cov.inv, "logdet")
            # mahalanobis distance
            Yc <- t( t(Y) - sample.mean )
            DIST <- rowSums(Yc %*% sample.cov.inv * Yc)
        }

        loglik <- -(P * LOG.2PI + logdet + DIST)/2

    } else {
        # invert sample.cov
        sample.cov.inv <- lav_matrix_symmetric_inverse(S = sample.cov,
                              logdet = TRUE, Sinv.method = Sinv.method)
        logdet <- attr(sample.cov.inv, "logdet")

        loglik <-
            lav_mvnorm_h1_loglik_samplestats(sample.cov.logdet = logdet,
                                             sample.nvar       = P,
                                             sample.nobs       = N)
    }

    # fixed.x?
    if(!is.null(x.idx) && length(x.idx) > 0L) {
        loglik.x <- lav_mvnorm_h1_loglik_data(Y = Y[, x.idx, drop = FALSE],
                                              wt = wt, x.idx = NULL, 
                                              casewise = casewise,
                                              Sinv.method = Sinv.method)
        # subtract logl.X
        loglik <- loglik - loglik.x
    }

    loglik
}



# 1b: input are sample statistics only (logdet, N and P)
lav_mvnorm_h1_loglik_samplestats <- function(sample.cov.logdet = NULL,
                                             sample.nvar       = NULL,
                                             sample.nobs       = NULL,
                                             # or
                                             sample.cov        = NULL,
                                             x.idx             = NULL,
                                             x.cov             = NULL,
                                             Sinv.method       = "eigen") {

    if(is.null(sample.nvar)) {
        P <- NCOL(sample.cov)
    } else {
        P <- sample.nvar # number of variables
    }

    N <- sample.nobs
    stopifnot(!is.null(P), !is.null(N))

    LOG.2PI <- log(2 * pi)

    # all we need is the logdet
    if(is.null(sample.cov.logdet)) {
        sample.cov.inv <- lav_matrix_symmetric_inverse(S = sample.cov,
                              logdet = TRUE, Sinv.method = Sinv.method)
        logdet <- attr(sample.cov.inv, "logdet")
    } else {
        logdet <- sample.cov.logdet
    }

    loglik <- -N/2 * (P * LOG.2PI + logdet + P)

    # fixed.x?
    if(!is.null(x.idx) && length(x.idx) > 0L) {
        if(is.null(sample.cov)) {
            if(is.null(x.cov)) {
                stop("psindex ERROR: when x.idx is not NULL, we need sample.cov or x.cov")
            } else {
                sample.cov.x <- x.cov
            }
        } else {
            sample.cov.x  <- sample.cov[x.idx, x.idx, drop = FALSE]
        }

        loglik.x <-
            lav_mvnorm_h1_loglik_samplestats(sample.cov  = sample.cov.x,
                                             sample.nobs = sample.nobs,
                                             x.idx = NULL,
                                             Sinv.method = Sinv.method)
        # subtract logl.X
        loglik <- loglik - loglik.x
    }

    loglik
}


# 4. hessian of logl (around MLEs of Mu and Sigma)

# 4a: hessian logl Mu and vech(Sigma) from raw data
lav_mvnorm_h1_logl_hessian_data <- function(Y              = NULL,
                                            wt             = NULL,
                                            x.idx          = NULL,
                                            Sinv.method    = "eigen",
                                            sample.cov.inv = NULL,
                                            meanstructure  = TRUE) {
    N <- NROW(Y)

    # observed information
    observed <- lav_mvnorm_h1_information_observed_data(Y = Y, wt = wt,
                    x.idx = x.idx, Sinv.method = Sinv.method,
                    sample.cov.inv = sample.cov.inv,
                    meanstructure = meanstructure)
 
    -N*observed
}

# 4b: hessian Mu and vech(Sigma) from samplestats
lav_mvnorm_h1_logl_hessian_samplestats <-
    function(sample.mean    = NULL, # unused!
             sample.cov     = NULL,
             sample.nobs    = NULL,
             x.idx          = NULL,
             Sinv.method    = "eigen",
             sample.cov.inv = NULL,
             meanstructure  = TRUE) {

    N <- sample.nobs

    # observed information
    observed <- lav_mvnorm_h1_information_observed_samplestats(sample.mean = 
        sample.mean, sample.cov = sample.cov, x.idx = x.idx,
        Sinv.method = Sinv.method, sample.cov.inv = sample.cov.inv, 
        meanstructure = meanstructure)
    
    -N*observed
}



# 5) Information h1 (note: expected == observed if data is complete!)

# 5a: unit expected information h1
lav_mvnorm_h1_information_expected <- function(Y              = NULL,
                                               wt             = NULL,
                                               sample.cov     = NULL,
                                               x.idx          = NULL,
                                               Sinv.method    = "eigen",
                                               sample.cov.inv = NULL,
                                               meanstructure  = TRUE) {
    if(is.null(sample.cov.inv)) {

        if(is.null(sample.cov)) {
            if(is.null(wt)) {
                sample.mean <- colMeans(Y); N <- NROW(Y)
                sample.cov <- 1/N*crossprod(Y) - tcrossprod(sample.mean)
            } else {
                out <- stats::cov.wt(Y, wt = wt, method = "ML")
                sample.cov <- out$cov
            }
        }

        # invert sample.cov
        sample.cov.inv <- lav_matrix_symmetric_inverse(S = sample.cov, 
                              logdet = FALSE, Sinv.method = Sinv.method)
    }

    I11 <- sample.cov.inv
    I22 <- 0.5 * lav_matrix_duplication_pre_post(sample.cov.inv %x% 
                                                 sample.cov.inv)

    if(meanstructure) {
        out <- lav_matrix_bdiag(I11, I22)
    } else {
        out <- I22
    }

    # fixed.x?
    if(!is.null(x.idx) && length(x.idx) > 0L) {
        not.x <- eliminate.pstar.idx(nvar = NCOL(sample.cov.inv),
                                     el.idx = x.idx,
                                     meanstructure = meanstructure)
        out[!not.x, ] <- 0
        out[, !not.x] <- 0
    }

    out
}

# 5b: unit observed information h1
lav_mvnorm_h1_information_observed_data <- function(Y              = NULL,
                                                    wt             = NULL,
                                                    x.idx          = NULL,
                                                    Sinv.method    = "eigen",
                                                    sample.cov.inv = NULL,
                                                    meanstructure  = TRUE) {

    lav_mvnorm_h1_information_expected(Y = Y, Sinv.method = Sinv.method,
                                       wt = wt, x.idx = x.idx,
                                       sample.cov.inv = sample.cov.inv,
                                       meanstructure = meanstructure)
}

# 5b-bis: observed information h1 from sample statistics
lav_mvnorm_h1_information_observed_samplestats <-
    function(sample.mean    = NULL, # unused!
             sample.cov     = NULL, 
             x.idx          = NULL,
             Sinv.method    = "eigen",
             sample.cov.inv = NULL,
             meanstructure  = TRUE) {

    if(is.null(sample.cov.inv)) {
        # invert sample.cov
        sample.cov.inv <- lav_matrix_symmetric_inverse(S = sample.cov, 
                              logdet = FALSE, Sinv.method = Sinv.method)
    }

    I11 <- sample.cov.inv
    I22 <- 0.5 * lav_matrix_duplication_pre_post(sample.cov.inv %x% 
                                                 sample.cov.inv)

    if(meanstructure) {
        out <- lav_matrix_bdiag(I11, I22)
    } else {
        out <- I22
    }

    # fixed.x?
    if(!is.null(x.idx) && length(x.idx) > 0L) {
        not.x <- eliminate.pstar.idx(nvar = NCOL(sample.cov.inv),
                                     el.idx = x.idx,
                                     meanstructure = meanstructure)
        out[!not.x, ] <- 0
        out[, !not.x] <- 0
    }

    out
}

# 5c: unit first-order information h1
#     note: first order information h1 == A1 %*% Gamma %*% A1 
#           (where A1 = obs/exp information h1)
lav_mvnorm_h1_information_firstorder <- function(Y              = NULL,
                                                 wt             = NULL,
                                                 sample.cov     = NULL,
                                                 x.idx          = NULL,
                                                 Sinv.method    = "eigen",
                                                 sample.cov.inv = NULL,
                                                 Gamma          = NULL,
                                                 meanstructure  = TRUE) {

    if(!is.null(wt)) {
        out <- stats::cov.wt(Y, wt = wt, method = "ML")
        res <- lav_mvnorm_information_firstorder(Y = Y, wt = wt, 
                   Mu = out$center, Sigma = out$cov, x.idx = x.idx,
                   meanstructure = meanstructure)
        return( res )
    }

    # Gamma
    if(is.null(Gamma)) {
        if(!is.null(x.idx) && length(x.idx) > 0L) {
            Gamma <- lav_samplestats_Gamma(Y, x.idx = x.idx, fixed.x = TRUE,
                                           meanstructure = meanstructure)
        } else {
            Gamma <- lav_samplestats_Gamma(Y, meanstructure = meanstructure)
        }
    }

    # sample.cov.in
    if(is.null(sample.cov.inv)) {
        # invert sample.cov
        if(is.null(sample.cov)) {
            N <- NROW(Y)
            sample.mean <- colMeans(Y)
            sample.cov <- 1/N*crossprod(Y) - tcrossprod(sample.mean)
        }
        sample.cov.inv <- lav_matrix_symmetric_inverse(S = sample.cov,
                              logdet = FALSE, Sinv.method = Sinv.method)
    }

    # A1
    A1 <- lav_mvnorm_h1_information_expected(Y = Y, Sinv.method = Sinv.method,
                                             sample.cov.inv = sample.cov.inv,
                                             x.idx = x.idx,
                                             meanstructure = meanstructure)

    A1 %*% Gamma %*% A1
}

# 6) inverted information h1 mu + vech(Sigma)

#    6a: (unit) inverted expected information (A1.inv = Gamma.NT)
#    6b: (unit) inverted observed information (A1.inv = Gamma.NT)

lav_mvnorm_h1_inverted_information_expected <- 
lav_mvnorm_h1_inverted_information_observed <- function(Y              = NULL,
                                                        sample.cov     = NULL,
                                                        x.idx          = NULL) {
    # sample.cov
    if(is.null(sample.cov)) {
        sample.mean <- colMeans(Y); N <- NROW(Y)
        sample.cov <- 1/N*crossprod(Y) - tcrossprod(sample.mean)
    }

    if(!is.null(x.idx) && length(x.idx) > 0L) {
        Gamma.NT <- lav_samplestats_Gamma_NT(Y = Y, x.idx = x.idx,
                                             COV = sample.cov,
                                             meanstructure = TRUE,
                                             fixed.x = TRUE)
    } else {
        I11 <- sample.cov
        I22 <- 2 * lav_matrix_duplication_ginv_pre_post(sample.cov %x% sample.cov)

        Gamma.NT <- lav_matrix_bdiag(I11, I22)
    }

    Gamma.NT
}

#    6c: (unit) inverted first-order information (B1.inv)
#        J1.inv = Gamma.NT %*% solve(Gamma) %*%  Gamma.NT
#
lav_mvnorm_h1_inverted_information_firstorder <- function(Y          = NULL,
                                                      sample.cov     = NULL,
                                                      x.idx          = NULL,
                                                      Sinv.method    = "eigen",
                                                      sample.cov.inv = NULL,
                                                      Gamma          = NULL) {
    # Gamma
    if(is.null(Gamma)) {
        if(!is.null(x.idx) && length(x.idx) > 0L) {
            Gamma <- lav_samplestats_Gamma(Y, x.idx = x.idx, fixed.x = TRUE,
                                           meanstructure = TRUE)
        } else {
            Gamma <- lav_samplestats_Gamma(Y, meanstructure = TRUE)
        }
    }

    # Gamma.NT
    Gamma.NT <- 
        lav_mvnorm_h1_inverted_information_expected(Y          = Y, 
                                                    sample.cov = sample.cov,
                                                    x.idx      = x.idx)
    if(!is.null(x.idx) && length(x.idx) > 0L) {
        # FIXME: surely there is better way
        out <- Gamma.NT %*% MASS::ginv(Gamma) %*% Gamma.NT
    } else {
        out <- Gamma.NT %*% solve(Gamma, Gamma.NT)
    }

    out
}


# 7) ACOV h1 mu + vech(Sigma)

#    7a: 1/N * Gamma.NT
#    7b: 1/N * Gamma.NT
lav_mvnorm_h1_acov_expected <-
lav_mvnorm_h1_acov_observed <- function(Y              = NULL,
                                        sample.cov     = NULL,
                                        x.idx          = NULL) {
    N <- NROW(Y)

    Gamma.NT <- 
        lav_mvnorm_h1_inverted_information_expected(Y          = Y, 
                                                    sample.cov = sample.cov,
                                                    x.idx      = x.idx)

    (1/N) * Gamma.NT
}

#    7c: 1/N * (Gamma.NT * Gamma^{-1} * Gamma.NT)
lav_mvnorm_h1_acov_firstorder <- function(Y              = NULL,
                                          sample.cov     = NULL,
                                          Sinv.method    = "eigen",
                                          x.idx          = NULL,
                                          sample.cov.inv = NULL,
                                          Gamma          = NULL) {
    N <- NROW(Y)

    J1.inv <- lav_mvnorm_h1_inverted_information_firstorder(Y = Y,
                  sample.cov = sample.cov, 
                  x.idx = x.idx, Sinv.method = Sinv.method, 
                  sample.cov.inv = sample.cov.inv, Gamma = Gamma)

    (1/N) * J1.inv
}

#    7d: 1/N * Gamma (sandwich)
lav_mvnorm_h1_acov_sandwich <- function(Y              = NULL,
                                        sample.cov     = NULL,
                                        x.idx          = NULL,
                                        Gamma          = NULL) {
    N <- NROW(Y)

    # Gamma
    if(is.null(Gamma)) {
        if(!is.null(x.idx) && length(x.idx) > 0L) {
            Gamma <- lav_samplestats_Gamma(Y, x.idx = x.idx, fixed.x = TRUE,
                                           meanstructure = TRUE)
        } else {
            Gamma <- lav_samplestats_Gamma(Y, meanstructure = TRUE)
        }
    }

    (1/N) * Gamma
}