add more options for selection

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: coala
-Version: 0.2.2.9001
-Date: 2015-11-24
+Version: 0.2.2.9002
+Date: 2015-11-25
 License: MIT + file LICENSE
 Title: A Framework for Coalescent Simulation
 Authors@R: c(

NEWS.md

@@ -2,6 +2,10 @@ coala 0.3.0
 ===========
 In development
 
+## Major improvements
+* Support for more selection models, including for local adaptation.
+
+## Small Changes
 * Adds `feat_ignore_singletons`, which is a feature that makes coala ignore 
   singletons when calculating the summary statistics.
 

R/feature_selection.R

@@ -1,22 +1,57 @@
 selection_class <- R6Class("selection", inherit = feature_class,
-  private = list(strength_AA = NA, strength_Aa = NA, additive = FALSE),
+  private = list(
+    strength_AA = 0,
+    strength_Aa = 0,
+    strength_aa = 0,
+    additive = FALSE,
+    force_keep = TRUE,
+    start_frequency = 0.05,
+    Ne = 10000,
+    position = 0.5,
+    start = TRUE
+  ),
   public = list(
-    initialize = function(strength_AA, strength_Aa, population, time,
-                          additive = FALSE) {
-      if (additive) {
-        private$strength_AA <- private$add_parameter(strength_AA)
-        private$additive <- TRUE
-      } else {
+    initialize = function(strength_AA, strength_Aa, strength_aa,
+                          population, time, additive = FALSE,
+                          start, start_frequency, Ne, position, force_keep) {
+
+      assert_that(is.logical(additive) && length(additive) == 1)
+      private$additive <- additive
+
+      assert_that(is.logical(start) && length(start) == 1)
+      private$start <- start
+
+      assert_that(is.numeric(start_frequency))
+      private$start_frequency <- start_frequency
+
+      assert_that(is.number(Ne))
+      private$Ne <- Ne
+
+      assert_that(is.numeric(position) && length(position) == 1)
+      private$position <- position
+
+      assert_that(is.logical(force_keep) && length(force_keep) == 1)
+      private$force_keep <- force_keep
+
+      private$strength_Aa <- private$add_parameter(strength_Aa)
+      if (!additive) {
         private$strength_AA <- private$add_parameter(strength_AA)
-        private$strength_Aa <- private$add_parameter(strength_Aa)
+        private$strength_aa <- private$add_parameter(strength_aa)
       }
 
-      private$set_population(population)
       private$time <- private$add_parameter(time)
+      private$set_population(population)
     },
     get_strength_AA = function() private$strength_AA,
     get_strength_Aa = function() private$strength_Aa,
+    get_strength_aa = function() private$strength_aa,
     is_additive = function() private$additive,
+    get_start = function() private$start,
+    get_force_keep = function() private$force_keep,
+    get_Ne = function() private$Ne,
+    get_start_freq = function() private$start_frequency,
+    get_position = function() private$position,
+
     print = function() {
       if (private$additive) {
         cat("Additive selection with strength ", print_par(private$strength_AA))
@@ -32,13 +67,38 @@ selection_class <- R6Class("selection", inherit = feature_class,
 
 #' Adds positive selection to a model
 #'
-#' @param population The population in which the allele is selected.
+#' @param population The population in which the allele is selected. Can either
+#'   be \code{all} for all population, or the number of a population.
 #' @param time The time at which the selection starts.
 #' @param strength_AA The selection strength for the selected homozygote
 #' @param strength_Aa The selection strength for the heterozygote.
+#' @param strength_aa The selection strength for the recessive homoygote.
 #' @param strength_A This sets the strength for the selected allele in an
-#'   haploid model. \code{strength_AA} and \code{strength_Aa} are ignored
-#'   when this is argument is given.
+#'   haploid model or a diploid model with additive selection.
+#'   \code{strength_AA}, \code{strength_Aa}, \code{strength_aa}
+#'   are ignored when this is argument is given.
+#' @param start Whether selection should start at this time point. At this time
+#'   the selected allele is introduced in the population with an initial
+#'   starting frequency. This must be set to \code{TRUE} for exactly one
+#'   selection feature in the model. The values of \code{start_frequency},
+#'   \code{Ne}, \code{position} and \code{force_keep} are used for the model.
+#'   You can add aditional selection feature to the model to set the
+#'   selection strength for more demes or change it at different time points,
+#'   but these need to have \code{start = FALSE}.
+#' @param start_frequency The start frequency at which the selected allele is
+#'   introduced at time \code{time}. If the model has multiple population, this
+#'   can either be a numeric vector that contains the initial frequency for each
+#'   population or a single number. In the latter case, the value is used for
+#'   all population specified with \code{populations}, and 0 is used for all
+#'   other populations.
+#' @param Ne The effective population size that is used for the forward
+#'   simulations.
+#' @param position The position of the selected site, relative to the
+#'   simulated sequence. Values between 0 and 1 are within the simulated area,
+#'   while smaller values are to the left of it and larger ones to the right.
+#' @param force_keep Whether to restart simulatin in which the selected goes to
+#'   extinction or not.
+#'
 #' @export
 #' @examples
 #' # Positive selection in population 2:
@@ -48,16 +108,32 @@ selection_class <- R6Class("selection", inherit = feature_class,
 #'                  strength_Aa=par_range('s', 100, 2000),
 #'                  population = 2,
 #'                  time=par_expr(tau))
-#'
-feat_selection <- function(strength_AA, strength_Aa, strength_A,
-                           population = 1, time) {
-  if (!missing(strength_A)) {
-    return(selection_class$new(strength_A,
+feat_selection <- function(strength_AA = 0,
+                           strength_Aa = 0,
+                           strength_aa = 0,
+                           strength_A = NULL,
+                           population = "all",
+                           time,
+                           start = TRUE,
+                           start_frequency = 0.0005,
+                           Ne = 10000,
+                           position = 0.5,
+                           force_keep = TRUE) {
+
+  if (!is.null(strength_A)) {
+    return(selection_class$new(strength_Aa = strength_A,
                                population = population,
                                time = time,
-                               additive = TRUE))
+                               additive = TRUE,
+                               start = start,
+                               start_frequency = start_frequency,
+                               Ne = Ne,
+                               position = position,
+                               force_keep = force_keep))
   }
-  selection_class$new(strength_AA, strength_Aa, population, time)
+
+  selection_class$new(strength_AA, strength_Aa, strength_aa, population, time,
+                      FALSE, start, start_frequency, Ne, position, force_keep)
 }
 
 
@@ -74,19 +150,56 @@ conv_to_scrm_arg.selection <- conv_to_ms_arg.selection
 #' @describeIn conv_to_ms_arg Feature conversion
 #' @export
 conv_to_msms_arg.selection <- function(feature, model) {
-  n_pop <- length(get_populations(model))
-  start_freq <- rep(0, n_pop)
-  start_freq[feature$get_population()] <- 0.0005
+
+  # Determine general options if we are starting selection
+  if (feature$get_start()) {
+    n_pop <- length(get_populations(model))
+
+    # Determine initial frequency of selected allele
+    if (length(feature$get_start_freq()) == 1) {
+      if (feature$get_population() == "all") {
+        start_freq <- rep(feature$get_start_freq(), n_pop)
+      } else {
+        start_freq <- rep(0, n_pop)
+        start_freq[feature$get_population()] <- feature$get_start_freq()
+      }
+    } else {
+      start_freq <- feature$get_start_freq()
+    }
+
+    start_cmd <- paste0("-SI', ", feature$get_time(), ", '",
+                        n_pop, " " , paste(start_freq, collapse = " "), " ",
+                        "-Sp', ", feature$get_position(), ", '",
+                        "-N', ", feature$get_Ne(), ", '",
+                        ifelse(feature$get_force_keep(), "-SForceKeep ", ""))
+  } else {
+    start_cmd <- ""
+  }
+
   if (feature$is_additive()) {
-    strength <- paste0("-SA',", feature$get_strength_AA(), ", '")
+    if (feature$get_population() == "all") {
+      strength <- paste0("-SA',", feature$get_strength_Aa(), ", '")
+    } else {
+      strength <- paste0("-Sc',",
+                         feature$get_time(), ", ",
+                         feature$get_population(), ", ",
+                         feature$get_strength_Aa(), ", '")
+    }
   } else {
-    strength <- paste0("-SAA',", feature$get_strength_AA(), ", '",
-                       "-SAa',", feature$get_strength_Aa(), ", '")
+    if (feature$get_population() == "all") {
+      strength <- paste0("-SAA',", feature$get_strength_AA(), ", '",
+                         "-SAa',", feature$get_strength_Aa(), ", '",
+                         "-Saa',", feature$get_strength_aa(), ", '")
+    } else {
+      strength <- paste0("-Sc',",
+                         feature$get_time(), ", ",
+                         feature$get_population(), ", ",
+                         feature$get_strength_AA(), ", ",
+                         feature$get_strength_Aa(), ", ",
+                         feature$get_strength_aa(), ", '")
+    }
   }
-  paste0("-SI', ", feature$get_time(), ", '",
-         n_pop, " " , paste(start_freq, collapse = " "), " ",
-         strength,
-         "-Sp 0.5 -SForceKeep -N 10000 ")
+  paste0(start_cmd, strength)
 }
 
 #' @describeIn conv_to_ms_arg Feature conversion

tests/testthat/test-feature-selection.R

@@ -1,23 +1,50 @@
 context("Feature Selection")
 
-test_that("generation of selection cmd works", {
+test_that("generation of initial selection cmd works", {
   if (!has_msms()) skip("msms not installed")
   msms <- get_simulator("msms")
   model  <- model_theta_tau() +
-    feat_selection(111, 222, population = 1, time = 5)
+    feat_selection(111, 222, 333, population = "all", time = 5)
   cmd <- msms$get_cmd(model)
-  expect_true(grepl("-N 10000", cmd))
-  expect_true(grepl("-SI 5 2 5e-04 0", cmd))
-  expect_true(grepl("-SAA 111", cmd))
-  expect_true(grepl("-SAa 222", cmd))
+  expect_true(grepl(" -N 10000 ", cmd))
+  expect_true(grepl(" -SForceKeep ", cmd))
+  expect_true(grepl(" -SI 5 2 5e-04 5e-04 ", cmd))
+  expect_true(grepl(" -Sp 0.5 ", cmd))
+  expect_true(grepl(" -SAA 111 ", cmd))
+  expect_true(grepl(" -SAa 222 ", cmd))
+  expect_true(grepl(" -Saa 333 ", cmd))
   expect_true(grepl(" $", cmd))
 
   model <- model_theta_tau() +
-    feat_selection(strength_A = 123, population = 1, time = 5)
+    feat_selection(strength_A = 123, population = 1, time = 5, Ne = 1122,
+                   start_frequency = c(0.1), position = 0.3)
   cmd <- msms$get_cmd(model)
-  expect_true(grepl("-N 10000", cmd))
-  expect_true(grepl("-SI 5 2 5e-04 0", cmd))
-  expect_true(grepl("-SA 123", cmd))
+  expect_true(grepl(" -Sc 5 1 123 ", cmd))
+  expect_true(grepl(" -SI 5 2 0.1 0 ", cmd))
+  expect_true(grepl("-N 1122", cmd))
+  expect_true(grepl(" -Sp 0.3 ", cmd))
+  expect_true(grepl(" $", cmd))
+
+  model <- model_theta_tau() +
+    feat_selection(strength_A = 123, population = 1, time = 5,
+                   start_frequency = c(0.1, 0.2))
+  expect_true(grepl(" -SI 5 2 0.1 0.2 ", msms$get_cmd(model)))
+})
+
+
+test_that("generation of selection cmd with multiple demes works", {
+  if (!has_msms()) skip("msms not installed")
+  msms <- get_simulator("msms")
+  model  <- model_theta_tau() +
+    feat_selection(111, 222, 333, population = 1, time = 5) +
+    feat_selection(444, 555, 666, population = 2, time = 5, start = FALSE)
+  cmd <- msms$get_cmd(model)
+  expect_true(grepl(" -N 10000 ", cmd))
+  expect_true(grepl(" -SForceKeep ", cmd))
+  expect_true(grepl(" -SI 5 2 5e-04 0 ", cmd))
+  expect_true(grepl(" -Sp 0.5 ", cmd))
+  expect_true(grepl(" -Sc 5 1 111 222 333 ", cmd))
+  expect_true(grepl(" -Sc 5 2 444 555 666 ", cmd))
   expect_true(grepl(" $", cmd))
 })