R/adjustment.R

In caugi: Causal Graph Interface

# ──────────────────────────────────────────────────────────────────────────────
# ─────────────────────────────── d-separation ─────────────────────────────────
# ──────────────────────────────────────────────────────────────────────────────

#' @title Are X and Y d-separated given Z?
#'
#' @description Checks whether every node in `X` is d-separated from every node
#' in `Y` given `Z` in a DAG.
#'
#' @param cg A `caugi` object.
#' @param X,Y,Z Character vectors of node names, or `NULL`. Use `*_index` to
#'   pass 1-based indices. If `Z` is `NULL` or missing, no nodes are conditioned
#'   on.
#' @param X_index,Y_index,Z_index Optional numeric 1-based indices (exclusive
#'   with `X`,`Y`,`Z` respectively).
#'
#' @returns `TRUE` if d-separated, `FALSE` otherwise.
#'
#' @examples
#' cg <- caugi(
#'   C %-->% X,
#'   X %-->% F,
#'   X %-->% D,
#'   A %-->% X,
#'   A %-->% K,
#'   K %-->% Y,
#'   D %-->% Y,
#'   D %-->% G,
#'   Y %-->% H,
#'   class = "DAG"
#' )
#'
#' d_separated(cg, "X", "Y", Z = c("A", "D")) # TRUE
#' d_separated(cg, "X", "Y", Z = NULL) # FALSE
#'
#' @family adjustment
#' @concept adjustment
#'
#' @export
d_separated <- function(
  cg,
  X = NULL,
  Y = NULL,
  Z = NULL,
  X_index = NULL,
  Y_index = NULL,
  Z_index = NULL
) {
  is_caugi(cg, throw_error = TRUE)
  if (
    length(X) > 1 || length(Y) > 1 || length(X_index) > 1 || length(Y_index) > 1
  ) {
    stop("Provide exactly one X and one Y.", call. = FALSE)
  }

  X_idx0 <- .resolve_idx0_get(cg@session, X, X_index)
  Y_idx0 <- .resolve_idx0_get(cg@session, Y, Y_index)
  Z_idx0 <- .resolve_idx0_mget(cg@session, Z, Z_index)

  rs_d_separated(cg@session, X_idx0, Y_idx0, Z_idx0)
}

# ──────────────────────────────────────────────────────────────────────────────
# ───────────────────────────── Adjustment sets ────────────────────────────────
# ──────────────────────────────────────────────────────────────────────────────

#' @title Compute an adjustment set
#'
#' @description Computes an adjustment set for `X -> Y` in a DAG.
#' @details
#' Types supported:
#' - `"parents"`: \eqn{\bigcup \mathrm{Pa}(X)} minus \eqn{X \cup Y}
#' - `"backdoor"`: an inclusion-minimal valid backdoor adjustment set, found in
#'   linear time as a minimal d-separator of `X` and `Y` in the proper backdoor
#'   graph (van der Zander, Liśkiewicz & Textor, 2019). Unlike `"parents"` (which
#'   always conditions on every parent of `X`), this drops nodes that are not
#'   needed to block a backdoor path. Note that minimal is not the same as
#'   statistically optimal; use `"optimal"` for the efficient O-set.
#' - `"optimal"`: O-set (only for single `x` and single `y`)
#'
#' @param cg A `caugi` object.
#' @param X,Y Node names.
#' @param X_index,Y_index Optional numeric 1-based indices.
#' @param type One of `"parents"`, `"backdoor"`, `"optimal"`.
#' The `optimal` option computes the O-set.
#'
#' @returns A character vector of node names representing the adjustment set.
#'
#' @examples
#' cg <- caugi(
#'   C %-->% X,
#'   X %-->% F,
#'   X %-->% D,
#'   A %-->% X,
#'   A %-->% K,
#'   K %-->% Y,
#'   D %-->% Y,
#'   D %-->% G,
#'   Y %-->% H,
#'   class = "DAG"
#' )
#'
#' adjustment_set(cg, "X", "Y", type = "parents") # C, A
#' adjustment_set(cg, "X", "Y", type = "backdoor") # A
#' adjustment_set(cg, "X", "Y", type = "optimal") # K
#'
#' @family adjustment
#' @concept adjustment
#'
#' @export
adjustment_set <- function(
  cg,
  X = NULL,
  Y = NULL,
  X_index = NULL,
  Y_index = NULL,
  type = c("optimal", "parents", "backdoor")
) {
  is_caugi(cg, throw_error = TRUE)
  if (
    length(X) > 1 || length(Y) > 1 || length(X_index) > 1 || length(Y_index) > 1
  ) {
    stop("Provide exactly one X and one Y.", call. = FALSE)
  }

  type <- match.arg(type)

  X_idx0 <- .resolve_idx0_get(cg@session, X, X_index)
  Y_idx0 <- .resolve_idx0_get(cg@session, Y, Y_index)

  idx0 <- switch(
    type,
    parents = rs_adjustment_set_parents(cg@session, X_idx0, Y_idx0),
    backdoor = rs_adjustment_set_backdoor(
      cg@session,
      X_idx0,
      Y_idx0
    ),
    optimal = rs_adjustment_set_optimal(cg@session, X_idx0, Y_idx0)
  )
  cg@nodes$name[idx0 + 1L]
}

#' @title Is a backdoor set valid?
#'
#' @description Checks whether `Z` is a valid backdoor adjustment set for
#' `X --> Y`.
#' @param cg A `caugi` object.
#'
#' @param X,Y Single node names.
#' @param Z Optional node set for conditioning
#' @param X_index,Y_index,Z_index Optional 1-based indices.
#'
#' @returns Logical value indicating if backdoor is valid or not.
#'
#' @examples
#' cg <- caugi(
#'   C %-->% X,
#'   X %-->% F,
#'   X %-->% D,
#'   A %-->% X,
#'   A %-->% K,
#'   K %-->% Y,
#'   D %-->% Y,
#'   D %-->% G,
#'   Y %-->% H,
#'   class = "DAG"
#' )
#'
#' is_valid_backdoor(cg, X = "X", Y = "Y", Z = NULL) # FALSE
#' is_valid_backdoor(cg, X = "X", Y = "Y", Z = "K") # TRUE
#' is_valid_backdoor(cg, X = "X", Y = "Y", Z = c("A", "C")) # TRUE
#'
#' @family adjustment
#' @concept adjustment
#'
#' @export
is_valid_backdoor <- function(
  cg,
  X = NULL,
  Y = NULL,
  Z = NULL,
  X_index = NULL,
  Y_index = NULL,
  Z_index = NULL
) {
  is_caugi(cg, throw_error = TRUE)
  if (
    length(X) > 1 || length(Y) > 1 || length(X_index) > 1 || length(Y_index) > 1
  ) {
    stop("Provide exactly one X and one Y.", call. = FALSE)
  }

  X_idx0 <- .resolve_idx0_get(cg@session, X, X_index)
  Y_idx0 <- .resolve_idx0_get(cg@session, Y, Y_index)
  Z_idx0 <- .resolve_idx0_mget(cg@session, Z, Z_index)

  rs_is_valid_backdoor_set(cg@session, X_idx0, Y_idx0, Z_idx0)
}

#' @title Get all backdoor sets up to a certain size.
#'
#' @description This function returns the backdoor sets up to size `max_size`,
#' which per default is set to 10.
#'
#' @param cg A `caugi`.
#' @param X,Y Single node name.
#' @param X_index,Y_index Optional 1-based indices (exclusive with name args).
#' @param minimal Logical; if `TRUE` (default), only minimal sets are returned.
#' @param max_size Integer; maximum size of sets to consider (default 3).
#'
#' @returns A list of character vectors, each an adjustment set
#' (possibly empty).
#'
#' @examples
#' cg <- caugi(
#'   C %-->% X,
#'   X %-->% F,
#'   X %-->% D,
#'   A %-->% X,
#'   A %-->% K,
#'   K %-->% Y,
#'   D %-->% Y,
#'   D %-->% G,
#'   Y %-->% H,
#'   class = "DAG"
#' )
#'
#' all_backdoor_sets(cg, X = "X", Y = "Y", max_size = 3L, minimal = FALSE)
#' #> [[1]]
#' #> [1] "A"
#' #>
#' #> [[2]]
#' #> [1] "K"
#' #>
#' #> [[3]]
#' #> [1] "C" "A"
#' #>
#' #> [[4]]
#' #> [1] "C" "K"
#' #>
#' #> [[5]]
#' #> [1] "A" "K"
#' #>
#' #> [[6]]
#' #> [1] "C" "A" "K"
#'
#' all_backdoor_sets(cg, X = "X", Y = "Y", max_size = 3L, minimal = TRUE)
#' #> [[1]]
#' #> [1] "A"
#' #>
#' #> [[2]]
#' #> [1] "K"
#'
#' @family adjustment
#' @concept adjustment
#'
#' @export
all_backdoor_sets <- function(
  cg,
  X = NULL,
  Y = NULL,
  X_index = NULL,
  Y_index = NULL,
  minimal = TRUE,
  max_size = 3L
) {
  is_caugi(cg, throw_error = TRUE)
  if (
    length(X) > 1 || length(Y) > 1 || length(X_index) > 1 || length(Y_index) > 1
  ) {
    stop("Provide exactly one X and one Y.", call. = FALSE)
  }

  X_idx0 <- .resolve_idx0_get(cg@session, X, X_index)
  Y_idx0 <- .resolve_idx0_get(cg@session, Y, Y_index)

  sets_idx0 <- rs_all_backdoor_sets(
    cg@session,
    X_idx0,
    Y_idx0,
    minimal,
    as.integer(max_size)
  )
  nm <- cg@nodes$name
  lapply(sets_idx0, \(idx0) nm[idx0 + 1L])
}

# ──────────────────────────────────────────────────────────────────────────────
# ────────────────────────── Minimal separator ─────────────────────────────────
# ──────────────────────────────────────────────────────────────────────────────

#' @title Compute a minimal separator
#'
#' @description Computes a minimal separator Z for sets X and Y, optionally with
#' mandatory inclusions and restrictions on the separator. Supports DAGs (where
#' the result is a d-separator), ADMGs, and ancestral graphs (where the result
#' is an m-separator).
#'
#' @details
#' A separator Z for X and Y is a set of nodes such that conditioning on Z
#' d- or m-separates X from Y in the graph. This function returns a minimal
#' separator: no proper subset of Z (excluding I) still separates X and Y.
#'
#' The algorithm runs in linear time O(n + m) and is unified across graph
#' classes via the Bayes-ball reachability of van der Zander & Liśkiewicz
#' (2020). For DAGs the algorithm specializes to FINDMINSEPINDAG; for ADMGs and
#' AGs (and their subclasses CPDAG, RCG, MAG) it uses the general FINDMINSEP
#' over mixed graphs.
#'
#' @param cg A `caugi` object (DAG, ADMG, or AG).
#' @param X,Y Character vectors of node names. Use `*_index` to pass 1-based
#'   indices.
#' @param I Nodes that must be included in the separator.
#' @param R Nodes allowed in the separator. If `NULL`, uses all nodes excluding
#'   X and Y.
#' @param X_index,Y_index,I_index,R_index Optional numeric 1-based indices
#'   (exclusive with corresponding name parameters).
#'
#' @returns A character vector of node names representing the minimal separator,
#' or `NULL` if no valid separator exists within the restriction R.
#'
#' @source van der Zander, B. & Liśkiewicz, M. (2020). Finding Minimal
#'   d-separators in Linear Time and Applications. Proceedings of The 35th
#'   Uncertainty in Artificial Intelligence Conference, in Proceedings of
#'   Machine Learning Research 115:637-647 Available from
#'   <https://proceedings.mlr.press/v115/van-der-zander20a.html>.
#'
#' @references
#' van der Zander, B. & Liśkiewicz, M. (2020). Finding Minimal d-separators in
#' Linear Time and Applications. In \emph{Proceedings of the 35th Conference on
#' Uncertainty in Artificial Intelligence (UAI 2020)}, PMLR 115:637--647.
#' \url{https://proceedings.mlr.press/v115/van-der-zander20a.html}.
#'
#' @examples
#' cg <- caugi(
#'   A %-->% X,
#'   X %-->% M,
#'   M %-->% Y,
#'   A %-->% Y,
#'   class = "DAG"
#' )
#'
#' # Find any minimal separator between X and Y
#' minimal_separator(cg, "X", "Y")
#'
#' # Force M to be in the separator
#' minimal_separator(cg, "X", "Y", I = "M")
#'
#' # Restrict separator to only {A, M}
#' minimal_separator(cg, "X", "Y", R = c("A", "M"))
#'
#' # Works on ADMGs (returns a minimal m-separator)
#' admg <- caugi(
#'   X %-->% Y,
#'   L %-->% X,
#'   L %-->% Y,
#'   class = "ADMG"
#' )
#' minimal_separator(admg, "X", "Y")
#'
#' @family adjustment
#' @concept adjustment
#'
#' @export
minimal_separator <- function(
  cg,
  X = NULL,
  Y = NULL,
  I = character(0),
  R = NULL,
  X_index = NULL,
  Y_index = NULL,
  I_index = NULL,
  R_index = NULL
) {
  is_caugi(cg, throw_error = TRUE)

  if (!is_dag(cg) && !is_admg(cg) && !is_ag(cg)) {
    stop(
      "`minimal_separator()` is only defined for DAGs, ADMGs, and AGs. ",
      "The graph has class \"",
      cg@graph_class,
      "\".",
      call. = FALSE
    )
  }

  # Resolve X and Y indices
  X_idx0 <- .resolve_idx0_mget(cg@session, X, X_index)
  Y_idx0 <- .resolve_idx0_mget(cg@session, Y, Y_index)

  # Validate that X and Y resolve to non-empty node sets
  if (length(X_idx0) == 0L) {
    stop(
      "`X` did not match any nodes in the graph. Check that the node names exist.",
      call. = FALSE
    )
  }
  if (length(Y_idx0) == 0L) {
    stop(
      "`Y` did not match any nodes in the graph. Check that the node names exist.",
      call. = FALSE
    )
  }
  # Resolve I indices (default to empty)
  I_idx0 <- .resolve_idx0_mget(cg@session, I, I_index)

  # Resolve R indices (default to all nodes except X and Y)
  if (is.null(R) && is.null(R_index)) {
    # Default: all nodes except X and Y
    all_idx0 <- seq.int(0L, nrow(cg@nodes) - 1L)
    R_idx0 <- setdiff(setdiff(all_idx0, X_idx0), Y_idx0)
  } else {
    R_idx0 <- .resolve_idx0_mget(cg@session, R, R_index)
  }

  # Call Rust function
  result_idx0 <- rs_minimal_separator(
    cg@session,
    as.integer(X_idx0),
    as.integer(Y_idx0),
    as.integer(I_idx0),
    as.integer(R_idx0)
  )

  # Convert result
  if (is.null(result_idx0)) {
    return(NULL)
  }

  # Convert 0-based indices to node names
  nm <- cg@nodes$name
  nm[result_idx0 + 1L]
}

#' @rdname minimal_separator
#' @export
minimal_d_separator <- function(
  cg,
  X = NULL,
  Y = NULL,
  I = character(0),
  R = NULL,
  X_index = NULL,
  Y_index = NULL,
  I_index = NULL,
  R_index = NULL
) {
  .Deprecated("minimal_separator")
  minimal_separator(
    cg = cg,
    X = X,
    Y = Y,
    I = I,
    R = R,
    X_index = X_index,
    Y_index = Y_index,
    I_index = I_index,
    R_index = R_index
  )
}

# ──────────────────────────────────────────────────────────────────────────────
# ──────────────────────── ADMG Adjustment Functions ───────────────────────────
# ──────────────────────────────────────────────────────────────────────────────

#' @title Is a set a valid adjustment set in an ADMG?
#'
#' @description Checks whether `Z` is a valid adjustment set for estimating
#' the causal effect of `X` on `Y` in an ADMG using the generalized adjustment
#' criterion.
#'
#' @param cg A `caugi` object of class ADMG.
#' @param X,Y Node names (can be vectors for multiple treatments/outcomes).
#' @param Z Conditioning set (character vector of node names).
#' @param X_index,Y_index,Z_index Optional 1-based indices.
#'
#' @returns Logical value indicating if the adjustment set is valid.
#'
#' @examples
#' # Classic confounding
#' cg <- caugi(
#'   L %-->% X,
#'   X %-->% Y,
#'   L %-->% Y,
#'   class = "ADMG"
#' )
#'
#' is_valid_adjustment_admg(cg, X = "X", Y = "Y", Z = NULL) # FALSE
#' is_valid_adjustment_admg(cg, X = "X", Y = "Y", Z = "L") # TRUE
#'
#' @family adjustment
#' @concept adjustment
#'
#' @export
is_valid_adjustment_admg <- function(
  cg,
  X = NULL,
  Y = NULL,
  Z = NULL,
  X_index = NULL,
  Y_index = NULL,
  Z_index = NULL
) {
  is_caugi(cg, throw_error = TRUE)
  # Validate that X and Y are provided

  if (is.null(X) && is.null(X_index)) {
    stop("X (or X_index) must be provided.", call. = FALSE)
  }
  if (is.null(Y) && is.null(Y_index)) {
    stop("Y (or Y_index) must be provided.", call. = FALSE)
  }

  X_idx0 <- .resolve_idx0_mget(cg@session, X, X_index)
  Y_idx0 <- .resolve_idx0_mget(cg@session, Y, Y_index)
  Z_idx0 <- .resolve_idx0_mget(cg@session, Z, Z_index)

  rs_is_valid_adjustment_set_admg(cg@session, X_idx0, Y_idx0, Z_idx0)
}

#' @title Get all valid adjustment sets in an ADMG
#'
#' @description Enumerates all valid adjustment sets for estimating the causal
#' effect of `X` on `Y` in an ADMG, up to a specified maximum size.
#'
#' @param cg A `caugi` object of class ADMG.
#' @param X,Y Node names (can be vectors for multiple treatments/outcomes).
#' @param X_index,Y_index Optional 1-based indices.
#' @param minimal Logical; if `TRUE` (default), only minimal sets are returned.
#' @param max_size Integer; maximum size of sets to consider (default 3).
#'
#' @returns A list of character vectors, each a valid adjustment set
#' (possibly empty list if none exist).
#'
#' @examples
#' cg <- caugi(
#'   L %-->% X,
#'   X %-->% Y,
#'   L %-->% Y,
#'   M %-->% Y,
#'   class = "ADMG"
#' )
#'
#' all_adjustment_sets_admg(cg, X = "X", Y = "Y", minimal = TRUE)
#' # Returns {L} as minimal adjustment set
#'
#' @family adjustment
#' @concept adjustment
#'
#' @export
all_adjustment_sets_admg <- function(
  cg,
  X = NULL,
  Y = NULL,
  X_index = NULL,
  Y_index = NULL,
  minimal = TRUE,
  max_size = 3L
) {
  is_caugi(cg, throw_error = TRUE)
  # Validate that X and Y are provided
  if (is.null(X) && is.null(X_index)) {
    stop("X (or X_index) must be provided.", call. = FALSE)
  }
  if (is.null(Y) && is.null(Y_index)) {
    stop("Y (or Y_index) must be provided.", call. = FALSE)
  }

  X_idx0 <- .resolve_idx0_mget(cg@session, X, X_index)
  Y_idx0 <- .resolve_idx0_mget(cg@session, Y, Y_index)

  sets_idx0 <- rs_all_adjustment_sets_admg(
    cg@session,
    X_idx0,
    Y_idx0,
    minimal,
    as.integer(max_size)
  )
  nm <- cg@nodes$name
  lapply(sets_idx0, \(idx0) nm[idx0 + 1L])
}

# ──────────────────────────────────────────────────────────────────────────────
# ───────────────────────────────── Helpers ────────────────────────────────────
# ──────────────────────────────────────────────────────────────────────────────

#' @title Resolve node name or index to 0-based index.
#'
#' @description Internal helper function to resolve either a node name or a
#' node index to a 0-based index using the Rust session.
#' `.resolve_idx0_get` expects a single value,
#' while `.resolve_idx0_mget` can return multiple values.
#'
#' @param session A GraphSession pointer.
#' @param node_name Optional character vector of node names.
#' @param node_index Optional numeric vector of 1-based node indices.
#'
#' @name .resolve_idx0_get
#'
#' @keywords internal
.resolve_idx0_get <- function(session, node_name = NULL, node_index = NULL) {
  if (!is.null(node_index)) {
    if (!is.null(node_name)) {
      stop("Provide either a node name or node index.", call. = FALSE)
    }
    as.integer(node_index - 1L)
  } else if (!is.null(node_name)) {
    rs_index_of(session, node_name)
  } else {
    stop(
      "Either the node name or the node index must be provided.",
      call. = FALSE
    )
  }
}

#' @name .resolve_idx0_get
#' @keywords internal
.resolve_idx0_mget <- function(session, node_name = NULL, node_index = NULL) {
  if (is.null(node_name) && is.null(node_index)) {
    integer(0)
  } else if (!is.null(node_index)) {
    if (!is.null(node_name)) {
      stop("Provide either a node name or node index.", call. = FALSE)
    }
    as.integer(node_index - 1L)
  } else if (!is.null(node_name)) {
    rs_indices_of(session, node_name)
  }
}