diff --git a/R/diversity.R b/R/diversity.R
index 434909b9..57a8b721 100644
--- a/R/diversity.R
+++ b/R/diversity.R
@@ -1,5 +1,5 @@
 `diversity` <-
-  function (x, index = "shannon", MARGIN = 1, base = exp(1)) 
+    function (x, index = "shannon", MARGIN = 1, base = exp(1))
 {
     x <- drop(as.matrix(x))
     if (any(x < 0, na.rm = TRUE))
@@ -16,16 +16,16 @@
         x <- -x * log(x, base)
     else
         x <- x*x
-    if (length(dim(x)) > 1) 
+    if (length(dim(x)) > 1)
         H <- apply(x, MARGIN, sum, na.rm = TRUE)
     else
         H <- sum(x, na.rm = TRUE)
-    if (index == "simpson") 
+    if (index == "simpson")
         H <- 1 - H
-    else if (index == "invsimpson") 
+    else if (index == "invsimpson")
         H <- 1/H
     ## check NA in data
-    if (any(is.na(total)))
-        is.na(H) <- is.na(total)
+    if (any(NAS <- is.na(total)))
+        H[NAS] <- NA
     H
 }