Simple Implementations of Freyja and Alcov models with unit tests

.gitignore

@@ -4,6 +4,8 @@ inst/doc
 working/
 *usher_barcodes*
 *pdf
+__pycache__/
 .Rproj.user
 README_cache/*
 provoc.Rproj
+.Rhistory

NAMESPACE

@@ -8,9 +8,11 @@ S3method(resid,provoc)
 S3method(residuals,provoc)
 S3method(summary,provoc)
 export(add_coverage)
+export(alcov)
 export(astronomize)
 export(coverage_at_aa)
 export(filter_varmat)
+export(freyja)
 export(fuse)
 export(get_canada_variants)
 export(get_convergence)
@@ -27,7 +29,16 @@ export(provoc_optim)
 export(simulate_coco)
 export(simulate_varmat)
 export(usher_barcodes)
+export(usher_barcodes)
 export(varmat_from_data)
 export(varmat_from_list)
 export(varmat_from_variants)
 importFrom(ggplot2,autoplot)
+import(glmnet)
+importFrom(ggplot2,aes)
+importFrom(ggplot2,geom_bar)
+importFrom(ggplot2,ggplot)
+importFrom(ggplot2,labs)
+importFrom(ggplot2,theme_minimal)
+importFrom(nnls,nnls)
+importFrom(scales,rescale)

R/alcov.R

@@ -0,0 +1,64 @@
+#' Simplified Alcov Model Using Non-Negative Linear Regression
+#'
+#' A simplified Alcov model for linear regression without intercepts,
+#' ensuring positivity of the coefficients. This function scales inputs and
+#' applies a non-negative linear regression to estimate variant proportions.
+#'
+#' @param Y Vector of frequencies for each sample.
+#' @param lmps Matrix or data frame of lineage definitions, similar to varmat.
+#' @param muts Char vector of mutation names, used to select and order columns in lmps if it is a data frame.
+#'
+#' @return An "Alcov" object with coefficients from the linear regression model, representing the estimated proportions of variants. Object can be plotted.
+#'
+#' @examples
+#' Y <- c(0.1, 0.2, 0.3, 0.4)
+#' lmps <- matrix(c(0, 1, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1), nrow = 4, byrow = TRUE)
+#' muts <- c("mut1", "mut2", "mut3")
+#'
+#' coef <- alcov(Y, lmps, muts)
+#' print(coef)
+#' 
+#' If you want to plot the coefficients:
+#' plot(coef, muts)
+#'
+#' @importFrom scales rescale
+#' @importFrom nnls nnls
+#' @export
+alcov <- function(Y, lmps, muts) {
+  # Ensure Y, lmps, and muts are properly aligned
+  if (is.data.frame(lmps)) {
+    lmps <- as.matrix(lmps[muts])
+  }
+
+  # Scale Y and lmps for better regression performance
+  Y_scaled <- rescale(Y)
+  lmps_scaled <- apply(lmps, 2, rescale)
+
+  # Linear Regression without intercept using non-negative least squares
+  model <- nnls(lmps_scaled, Y_scaled)
+
+  # Extract and return coeffs
+  alcov_coeffs <- coef(model)
+  class(alcov_coeffs) <- "alcov"
+  return(alcov_coeffs)
+}
+
+
+#' Plot Method for Alcov Coefficients
+#'
+#' @param coef Vector of coefficients returned by the alcov function.
+#' @param muts Vector of mutation names, which must match the length of `coef`.
+#' @importFrom ggplot2 ggplot geom_bar aes labs theme_minimal
+#' @export
+#' @method plot alcov
+plot.alcov <- function(coef, muts) {
+  if (!requireNamespace("ggplot2", quietly = TRUE)) {
+    stop("ggplot2 must be installed to use this function.")
+  }
+
+  df <- data.frame(Mutation = muts, Coefficient = coef)
+  ggplot(df, aes(x = Mutation, y = Coefficient, fill = Mutation)) +
+    geom_bar(stat = "identity") +
+    labs(title = "Alcov Model Coefficients", x = "Mutation", y = "Coefficient") +
+    theme_minimal()
+}

R/freyja.R

@@ -0,0 +1,68 @@
+#' Simplified Freyja Model Using Lasso Regression
+#'
+#' This function applies a Lasso regression to estimate proportions of variants
+#' based on frequency and adjusting for sequence depth.
+#'
+#' @param mix Vector of frequencies (count divided by coverage) for each sample.
+#' @param depths Vector of coverage for each sample.
+#' @param df_barcodes Numeric matrix varmat, with rows as samples and columns as mutations.
+#' @param muts Char vector of mutation names.
+#' @param eps Very small number representing the strength for Lasso regression. Default is 1e-4.
+#'
+#' @return A "Freyja" object with coefficients from the Lasso regression, representing the estimated proportions of variants. Object can be plotted.
+#'
+#' @examples
+#' mix <- c(0.1, 0.2, 0.3, 0.4)
+#' depths <- c(10, 20, 30, 40)
+#' df_barcodes <- matrix(c(0, 1, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1), nrow = 4, byrow = TRUE)
+#' muts <- c("mut1", "mut2", "mut3")
+#'
+#' coef <- freyja(mix, depths, df_barcodes, muts)
+#' print(coef)
+#' 
+#' If you want to plot the coefficients:
+#' plot(coef, muts)
+#'
+#' @import glmnet
+#' @export
+freyja <- function(mix, depths, df_barcodes, muts, eps=1e-4) {
+  # Adjust mutations based on coverage
+  depth_adjustment <- log(depths + 1) / max(log(depths + 1))
+  adjusted_mix <- mix * depth_adjustment
+
+  # Adjusted to replicate depth adjustment for each mutation
+  adjusted_barcodes <- t(t(df_barcodes) * depth_adjustment)
+
+  # Prepare for glmnet
+  x_matrix <- as.matrix(adjusted_barcodes)
+  y_vector <- as.vector(adjusted_mix)
+
+  # Initialize and fit the Lasso model
+  lasso_model <- glmnet(x_matrix, y_vector, alpha = 1, lambda = eps, intercept = FALSE, lower.limits = 0)
+
+  # Return the coeffs, excluding intercept
+  # The intercept is the first element in the glmnet coefficient matrix, so we skip it
+  freyja_coeffs <- coef(lasso_model)[-1]  # Removing intercept term which is included by default
+  class(freyja_coeffs) <- "freyja"
+  return(freyja_coeffs)
+}
+
+
+#' Plot Method for Freyja Coefficients
+#'
+#' @param coef Vector of coefficients returned by the freyja function.
+#' @param muts Vector of mutation names, which must match the length of `coef`.
+#' @importFrom ggplot2 ggplot geom_bar aes labs theme_minimal
+#' @export
+#' @method plot freyja
+plot.freyja <- function(coef, muts) {
+  if (!requireNamespace("ggplot2", quietly = TRUE)) {
+    stop("ggplot2 must be installed to use this function.")
+  }
+
+  df <- data.frame(Mutation = muts, Coefficient = coef)
+  ggplot(df, aes(x = Mutation, y = Coefficient, fill = Mutation)) +
+    geom_bar(stat = "identity") +
+    labs(title = "Freyja Model Coefficients", x = "Mutation", y = "Coefficient") +
+    theme_minimal()
+}

tests/testthat/test-other-models.R

@@ -0,0 +1,54 @@
+library(testthat)
+library(ggplot2)
+
+# Assemble
+mix <- c(0.1, 0.2, 0.3, 0.4)
+depths <- c(10, 20, 30, 40)
+df_barcodes <- matrix(c(0, 1, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1), nrow = 4, byrow = TRUE)
+muts <- c("mut1", "mut2", "mut3")
+Y <- c(0.1, 0.2, 0.3, 0.4)
+lmps <- matrix(c(0, 1, 0, 1, 0, 1, 1, 1, 0, 0, 0, 1), nrow = 4, byrow = TRUE)
+
+# Tests for Freyja model
+test_that("freyja model produces non-zero coefficients", {
+  # Act
+  freyja_coeffs <- provoc::freyja(mix, depths, df_barcodes, muts)
+
+  # Assert
+  expect_true(any(freyja_coeffs > 0), info = "Freyja model should produce some non-zero coefficients")
+})
+
+# Tests for Alcov model
+test_that("alcov model produces expected coefficients", {
+  # Act
+  alcov_coeffs <- provoc::alcov(Y, lmps, muts)
+
+  # Assert
+  expected_alcov_coeffs <- c(0.0, 0.3333333, 0.6666667)
+  expect_equal(alcov_coeffs, expected_alcov_coeffs, tolerance = 1e-5, 
+               info = "Alcov model coefficients should match expected values")
+})
+
+# Plotting coefficients for visual comparison
+test_that("plot coefficients for visual comparison", {
+
+  freyja_coeffs <- provoc::freyja(mix, depths, df_barcodes, muts)
+  alcov_coeffs <- provoc::alcov(Y, lmps, muts)
+
+  coefficients_df <- data.frame(
+    model = rep(c("Freyja", "Alcov"), each = 3),
+    mutation = rep(muts, 2),
+    coefficient = c(freyja_coeffs, alcov_coeffs)
+  )
+
+  p <- ggplot(coefficients_df, aes(x = mutation, y = coefficient, fill = model)) +
+    geom_bar(stat = "identity", position = position_dodge()) +
+    labs(title = "Coefficient Comparison between Freyja and Alcov Models",
+         y = "Coefficient",
+         x = "Mutation") +
+    scale_fill_manual(values = c("Freyja" = "blue", "Alcov" = "red")) +
+    theme_minimal()
+
+  print(p)
+})
+