diff --git a/R/simuloiAlleeli.R b/R/simuloiAlleeli.R
index 9f26d21..955684a 100644
--- a/R/simuloiAlleeli.R
+++ b/R/simuloiAlleeli.R
@@ -4,11 +4,22 @@
 #' @export
 
 simuloiAlleeli <- function(allfreqs, pop, loc) {
-    if (length(dim(allfreqs)) == 3) { # distinguish between arrays and matrices
-        freqs <- allfreqs[, loc, pop]
+    if (length(dim(allfreqs)) == 0) {
+        freqs <- 1
     } else {
-        freqs <- allfreqs[, loc]
+        if (length(dim(allfreqs)) == 3) { # distinguish between array and matrix
+            freqs <- allfreqs[, loc, pop]
+        } else {
+            freqs <- allfreqs[, loc]
+        }
     }
+    # freqs <- ifelse(is.null(length(dim(allfreqs)), allfreqs[loc], 0)
+    # freqs <- switch() + 1,
+    #     allfreqs[, loc],
+    #     allfreqs[, loc, pop]
+    # )
+    
+
     cumsumma <- cumsum(freqs)
     arvo <- runif(1)
     isommat <- which(cumsumma > arvo)
diff --git a/tests/testthat/test-admix1.R b/tests/testthat/test-admix1.R
index 962dcdf..04e004a 100644
--- a/tests/testthat/test-admix1.R
+++ b/tests/testthat/test-admix1.R
@@ -183,9 +183,11 @@ test_that("computePersonalAllFreqs works like on Matlab", {
 
 test_that("simuloiAlleeli works like on Matlab", {
     # TODO: test on vector
+    sk1 <- 2
     ra1 <- array(1:12, c(2, 2, 3))
     mx1 <- matrix(c(3, 5, 0, 9), 2)
     mx2 <- matrix(c(3, 5, 0, 9, 5, 8), 2)
+    expect_equal(simuloiAlleeli(sk1, 1, 1), 1)
     expect_equal(simuloiAlleeli(ra1, 2, 1), 1)
     expect_equal(simuloiAlleeli(mx1, 1, 2), 2)
     expect_equal(simuloiAlleeli(mx2, 1, 3), 1)