R/fitted.heckmange.R

In heckmanGE: Estimation and Inference for Heckman Selection Models with Cluster-Robust Variance

#' fitted.heckmanGE
#' Extract Fitted Values of the Generalized Heckman Model
#'
#' This function extracts the fitted values from a `heckmanGE` object. You can specify which part of the model you want to retrieve the fitted values for: selection, outcome, dispersion, or correlation. By default, it returns the fitted values for the outcome part of the model.
#'
#' @param object An object of class `heckmanGE`. This object should contain the fitted values for different parts of the model.
#' @param part A character vector specifying which part of the model to return the fitted values for. Options are "selection", "outcome" (default), "dispersion", or "correlation". If multiple parts are provided, only the "outcome" part will be returned.
#' @param \dots Additional arguments passed to or from other methods. These are not used in this method but must be included to match the generic method signature.
#'
#' @return A vector of fitted values corresponding to the specified part of the `heckmanGE` model. The type of the returned values depends on the part specified.
#'
#' @details
#' - If `part` is "selection", the function returns the fitted values from the selection part of the model.
#' - If `part` is "outcome", the function returns the fitted values from the outcome part of the model.
#' - If `part` is "dispersion", the function returns the fitted values from the dispersion part of the model.
#' - If `part` is "correlation", the function returns the fitted values from the correlation part of the model.
#'
#' If `part` is not one of the specified options, an error will be raised. If multiple parts are provided, the function defaults to returning the fitted values for the outcome part of the model.
#'
#' @importFrom stats fitted
#' @export
fitted.heckmanGE = function(object,
                                part = c("selection", "outcome", "dispersion", "correlation"), ...){

        if(!all(part %in% c("selection", "outcome", "dispersion", "correlation"))) {
                stop("part must be 'selection', 'outcome', 'dispersion', or 'correlation'")
        }

        if(length(part) > 1) {
                part = 'outcome'
        }

        if(part == "selection"){
                fitted = object$fitted.values$fit.selection
        }

        if(part == "outcome"){
                fitted = object$fitted.values$fit.outcome
        }

        if(part == "dispersion"){
                fitted = object$fitted.values$fit.dispersion
        }

        if(part == "correlation"){
                fitted = object$fitted.values$fit.correlation
        }

        return(fitted)
}