R/vectorfield_landscape.R

In fitlandr: Fit Vector Fields and Potential Landscapes from Intensive Longitudinal Data

#' Estimate a 3D potential landscape from a vector field
#'
#' Two methods are available: `method = "pathB"` and `method = "simlandr"`. See *Details* section.
#'
#' For `method = "simlandr"`, the landscape is constructed based on the generalized potential landscape by Wang et al. (2008), implemented by the `simlandr` package. This function is a wrapper of [sim_vf()] and [simlandr::make_3d_static()]. Use those two functions separately for more customization.
#'
#' For `method = "pathB"`, the landscape is constructed based on the deterministic path-integral quasi-potential defined by Bhattacharya et al. (2011).
#'
#' We recommend the `simlandr` method for psychological data because it is more stable.
#'
#' Parallel computing based on `future` is supported for both methods. Use `future::plan("multisession")` to enable this and speed up computation.
#'
#' @param vf A `vectorfield` object estimated by [fit_2d_vf()].
#' @param method The method used for landscape construction. Can be `pathB` or `simlandr`.
#' @param .pathB_options Only for `method = "pathB"`. Options controlling the path-integral algorithm. Should be generated by [sim_vf_options()].
#' @param .sim_vf_options Only for `method = "simlandr"`. Options controlling the vector field simulation. Should be generated by [sim_vf_options()].
#' @param .simlandr_options Only for `method = "simlandr"`. Options controlling the landscape construction. Should be generated by [simlandr_options()].
#' @inheritParams predict.vectorfield
#'
#' @return A `landscape` object as described in [simlandr::make_3d_static()], or a `3d_static_landscape_B` object, which inherits from the `landscape` class and contains the following elements: `dist`, the distribution estimation for landscapes; `plot`, a 3D plot using `plotly`; plot_2, a 2D plot using `ggplot2`; x, y, from `vf`.
#' @examplesIf interactive()
#' # generate data
#' single_output_grad <- simlandr::sim_fun_grad(length = 200, seed = 1614)
#' # fit the vector field
#' v2 <- fit_2d_vf(single_output_grad, x = "x", y = "y", method = "MVKE")
#' plot(v2)
#' # fit the landscape
#' future::plan("multisession")
#' set.seed(1614)
#' l2 <- fit_3d_vfld(v2,
#' .sim_vf_options = sim_vf_options(chains = 16, stepsize = 1, forbid_overflow = TRUE),
#' .simlandr_options = simlandr_options(adjust = 5, Umax = 4))
#' plot(l2, 2)
#' future::plan("sequential")
#' @export
fit_3d_vfld <- function(vf, method = c("simlandr", "pathB"), .pathB_options = pathB_options(vf), .sim_vf_options = sim_vf_options(vf), .simlandr_options = simlandr_options(vf), linear_interp = FALSE) {
  method <- match.arg(method[1], c("pathB", "simlandr"))
  if (method == "pathB") {
    all_pars <- .pathB_options %>% lapply(eval_pass_missing, list(vf = vf))
    all_pars$f <- function(x) {
    	stats::predict(all_pars$vf, pos = x, linear_interp = linear_interp, calculate_a = FALSE)$v
    }
    cli::cli_progress_step("Calculating path integrals")
    resultB <- do.call(path_integral_B, all_pars)
    cli::cli_progress_step("Aligning potentials")
    out_B <- do.call(align_pot_B, c(list(resultB = resultB), all_pars))
    out_B$d <- out_B$z

    p <- plotly::plot_ly(x = out_B$x, y = out_B$y, z = out_B$z, type = "surface")
    p <- plotly::layout(p, scene = list(
      xaxis = list(title = vf$x),
      yaxis = list(title = vf$y), zaxis = list(title = "U")
    )) %>%
      plotly::colorbar(title = "U")
    p2 <- ggplot2::ggplot(simlandr::make_2d_tidy_dist(out_B), ggplot2::aes(
      x = x,
      y = y
    )) +
      ggplot2::geom_raster(ggplot2::aes(fill = d)) +
      ggplot2::scale_fill_viridis_c() +
      ggplot2::labs(x = vf$x, y = vf$y, fill = "U") +
      ggplot2::theme_bw()
    result <- c(list(
      dist = out_B, plot = p, plot_2 = p2, x = vf$x,
      y = vf$y
    ), all_pars)
    class(result) <- c(
      "3d_static_landscape_B", "3d_static_landscape", "3d_landscape",
      "landscape"
    )
    return(result)
  } else if (method == "simlandr") {
    cli::cli_progress_step("Simulating the model")
    simulation_output <- do.call(sim_vf, c(.sim_vf_options, list(linear_interp = linear_interp)) %>% lapply(eval_pass_missing, list(vf = vf, chains = .$chains)))
    cli::cli_progress_step("Constructing the landscape")
    return(do.call(simlandr::make_3d_static, c(list(output = simulation_output), .simlandr_options %>% lapply(eval_pass_missing, list(vf = vf)))))
  }
}