geposan/R/method_neural.R

#' Find genes by training and applying a neural network.
#'
#' @param seed The seed will be used to make the results reproducible.
#' @param n_models This number specifies how many sets of training data should
#'   be created. For each set, there will be a model trained on the remaining
#'   training data and validated using this set. For non-training genes, the
#'   final score will be the mean of the result of applying the different
#'   models. There should be at least two training sets. The analysis will only
#'   work, if there is at least one reference gene per training set.
#'
#' @return An object of class `geposan_method`.
#'
#' @export
neural <- function(seed = 180199, n_models = 5) {
  method(
    id = "neural",
    name = "Neural",
    description = "Assessment by neural network",
    function(preset, progress) {
      species_ids <- preset$species_ids
      gene_ids <- preset$gene_ids
      reference_gene_ids <- preset$reference_gene_ids

      cached(
        "neural",
        c(species_ids, gene_ids, reference_gene_ids, seed, n_models),
        { # nolint
          reference_count <- length(reference_gene_ids)
          stopifnot(n_models %in% 2:reference_count)

          # Make results reproducible.
          tensorflow::set_random_seed(seed)

          # Step 1: Prepare input data.
          # ---------------------------

          # Prefilter distances by species.
          distances <- geposan::distances[species %chin% species_ids]

          # Input data for the network. This contains the gene ID as
          # an identifier as well as the per-species gene distances as
          # input variables.
          data <- data.table(gene = gene_ids)

          # Buffer to keep track of the names of the input variables.
          input_vars <- NULL

          # Make a columns containing positions and distances for each
          # species.
          for (species_id in species_ids) {
            species_data <- distances[
              species == species_id,
              .(gene, distance)
            ]

            # Only include species with at least 25% known values.
            # As positions and distances always coexist, we don't
            # loose any data here.

            species_data <- stats::na.omit(species_data)

            if (nrow(species_data) >= 0.25 * length(gene_ids)) {
              data <- merge(data, species_data, all.x = TRUE)

              # Replace missing data with mean values. The neural
              # network can't handle NAs in a meaningful way.
              # Choosing extreme values here would result in
              # heavily biased results. Therefore, the mean value
              # is chosen as a compromise. However, this will of
              # course lessen the significance of the results.

              mean_distance <- round(
                species_data[, mean(distance)]
              )

              data[is.na(distance), distance := mean_distance]

              # Name the new column after the species.
              setnames(data, "distance", species_id)

              # Add the input variable to the buffer.
              input_vars <- c(input_vars, species_id)
            }
          }

          progress(0.1)

          # Step 2: Prepare training data.
          # ------------------------------

          # Take out the reference data.

          reference_data <- data[gene %chin% reference_gene_ids]
          reference_data[, score := 1.0]

          # Take out random samples from the remaining genes. This is
          # another compromise with a negative impact on
          # significance. We assume that a random gene is not likely
          # to match the reference genes.

          without_reference_data <- data[
            !gene %chin% reference_gene_ids
          ]

          control_data <- without_reference_data[
            sample(
              nrow(without_reference_data),
              reference_count
            )
          ]

          control_data[, score := 0.0]

          # Split the training data into random sets to have
          # validation data for each model.

          # Scramble the source tables.
          reference_data <- reference_data[sample(reference_count)]
          control_data <- control_data[sample(reference_count)]

          networks <- list()

          indices <- seq_len(reference_count)
          indices_split <- split(indices, indices %% n_models)

          for (i in seq_len(n_models)) {
            training_data <- rbindlist(list(
              reference_data[!indices_split[[i]]],
              control_data[!indices_split[[i]]]
            ))

            validation_data <- rbindlist(list(
              reference_data[indices_split[[i]]],
              control_data[indices_split[[i]]]
            ))

            networks[[i]] <- list(
              training_data = training_data,
              validation_data = validation_data
            )
          }

          # Step 3: Create, train and apply neural network.
          # -----------------------------------------------

          # Layers for the neural network.
          input_layer <- length(input_vars)
          layer1 <- input_layer
          layer2 <- 0.5 * input_layer
          layer3 <- 0.5 * layer2

          # Convert data to matrix and normalize it.
          to_matrix <- function(data) {
            data_matrix <- as.matrix(data[, ..input_vars])
            colnames(data_matrix) <- NULL
            keras::normalize(data_matrix)
          }

          data_matrix <- to_matrix(data)
          output_vars <- NULL

          for (i in seq_along(networks)) {
            # Create a new model for each training session, because
            # the model would keep its state across training
            # sessions otherwise.
            model <- keras::keras_model_sequential() |>
              keras::layer_dense(
                units = layer1,
                activation = "relu",
                input_shape = input_layer,
              ) |>
              keras::layer_dense(
                units = layer2,
                activation = "relu",
                kernel_regularizer = keras::regularizer_l2()
              ) |>
              keras::layer_dense(
                units = layer3,
                activation = "relu",
                kernel_regularizer = keras::regularizer_l2()
              ) |>
              keras::layer_dense(
                units = 1,
                activation = "sigmoid"
              ) |>
              keras::compile(
                loss = keras::loss_mean_absolute_error(),
                optimizer = keras::optimizer_adam()
              )

            # Train the model.

            network <- networks[[i]]

            training_data <- network$training_data
            training_matrix <- to_matrix(training_data)
            validation_data <- network$validation_data
            validation_matrix <- to_matrix(validation_data)

            fit <- keras::fit(
              model,
              x = training_matrix,
              y = training_data$score,
              validation_data = list(
                x_val = validation_matrix,
                y_val = validation_data$score
              ),
              epochs = 500,
              verbose = FALSE
            )

            # Apply the model.

            data[, new_score := stats::predict(model, data_matrix)]

            # Remove the values of the training data itself.
            data[gene %chin% training_data$gene, new_score := NA]

            output_var <- sprintf("score%i", i)
            setnames(data, "new_score", output_var)
            output_vars <- c(output_vars, output_var)


            # Store the details.

            networks[[i]]$model <- keras::serialize_model(model)
            networks[[i]]$fit <- fit

            progress(0.1 + i * (0.9 / n_models))
          }

          # Compute the final score as the mean score.
          data[,
            score := mean(as.numeric(.SD), na.rm = TRUE),
            .SDcols = output_vars,
            by = gene
          ]

          progress(1.0)

          result(
            method = "neural",
            scores = data[, .(gene, score)],
            details = list(
              seed = seed,
              n_models = n_models,
              all_results = data[, !..input_vars],
              networks = networks
            )
          )
        }
      )
    }
  )
}