R/zzz.R

In mlr3cluster: Cluster Extension for 'mlr3'

#' @import checkmate
#' @import data.table
#' @import mlr3
#' @import mlr3misc
#' @import paradox
#' @importFrom graphics plot
#' @importFrom R6 R6Class
#' @importFrom stats model.frame predict rnorm runif terms
"_PACKAGE"

mlr3cluster_tasks = new.env(parent = emptyenv())
mlr3cluster_learners = new.env(parent = emptyenv())
mlr3cluster_task_generators = new.env(parent = emptyenv())

register_item = function(env, type) {
  function(name, constructor) {
    if (name %chin% names(env)) {
      stopf("%s %s registered twice.", type, name)
    }
    env[[name]] = constructor
  }
}

register_task = register_item(mlr3cluster_tasks, "task")
register_learner = register_item(mlr3cluster_learners, "learner")
register_task_generator = register_item(mlr3cluster_task_generators, "task generator")

register_mlr3 = function(...) {
  # reflections
  mlr_reflections = utils::getFromNamespace("mlr_reflections", ns = "mlr3")
  mlr_reflections$task_types = mlr_reflections$task_types[!"clust"]
  # fmt: skip
  mlr_reflections$task_types = setkeyv(rbind(mlr_reflections$task_types, rowwise_table(
    ~type,    ~package,       ~task,        ~learner,       ~prediction,        ~prediction_data,       ~measure,
    "clust",  "mlr3cluster",  "TaskClust",  "LearnerClust", "PredictionClust",  "PredictionDataClust",  "MeasureClust"
  ), fill = TRUE), "type")

  mlr_reflections$task_col_roles$clust = mlr_reflections$task_col_roles$regr
  mlr_reflections$task_properties$clust = mlr_reflections$task_properties$regr
  mlr_reflections$learner_properties$clust = union(
    mlr_reflections$learner_properties$regr,
    c("partitional", "hierarchical", "exclusive", "overlapping", "fuzzy", "complete", "partial", "density")
  )
  mlr_reflections$learner_predict_types$clust = list(partition = "partition", prob = c("partition", "prob"))
  mlr_reflections$measure_properties$clust = mlr_reflections$measure_properties$regr
  mlr_reflections$default_measures$clust = "clust.dunn"
  mlr_reflections$loaded_packages = union(mlr_reflections$loaded_packages, "mlr3cluster")

  # tasks
  mlr_tasks = utils::getFromNamespace("mlr_tasks", ns = "mlr3")
  iwalk(as.list(mlr3cluster_tasks), function(task, id) mlr_tasks$add(id, task))

  # task generators
  mlr_task_generators = utils::getFromNamespace("mlr_task_generators", ns = "mlr3")
  iwalk(as.list(mlr3cluster_task_generators), function(generator, id) mlr_task_generators$add(id, generator))

  # learners
  mlr_learners = utils::getFromNamespace("mlr_learners", ns = "mlr3")
  iwalk(as.list(mlr3cluster_learners), function(learner, id) mlr_learners$add(id, learner))

  # measures
  mlr_measures = utils::getFromNamespace("mlr_measures", ns = "mlr3")
  iwalk(as.list(measures), function(info, name) {
    mlr_measures$add(paste0("clust.", name), MeasureClustSimple, name = name)
  })
}

.onLoad = function(libname, pkgname) {
  backports::import(pkgname)

  register_namespace_callback(pkgname, "mlr3", register_mlr3)
}

.onUnload = function(libpath) {
  walk(names(mlr3cluster_tasks), function(id) mlr_tasks$remove(id))
  walk(names(mlr3cluster_task_generators), function(id) mlr_task_generators$remove(id))
  walk(names(mlr3cluster_learners), function(id) mlr_learners$remove(id))
  walk(names(measures), function(id) mlr_measures$remove(paste0("clust.", id)))

  mlr_reflections$task_types = mlr_reflections$task_types[!"clust"]
  reflections = c(
    "measure_properties",
    "default_measures",
    "learner_properties",
    "learner_predict_types",
    "task_properties",
    "task_col_roles"
  )
  walk(reflections, function(x) mlr_reflections[[x]] = remove_named(mlr_reflections[[x]], "clust"))
  mlr_reflections$loaded_packages = setdiff(mlr_reflections$loaded_packages, "mlr3cluster")
}

leanify_package()