diff --git a/NAMESPACE b/NAMESPACE
index 485f9f0..b02be59 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -17,7 +17,6 @@ export(min_norm)
 export(plot_factors)
 export(plot_factors_bars)
 export(snp_ldsc)
-import(data.table)
 import(dplyr)
 import(flashier)
 import(ggplot2)
@@ -25,3 +24,13 @@ import(lpSolve)
 import(purrr)
 import(readr)
 import(reshape2)
+importFrom(data.table,"%chin%")
+importFrom(data.table,":=")
+importFrom(data.table,.I)
+importFrom(data.table,data.table)
+importFrom(data.table,fifelse)
+importFrom(data.table,fwrite)
+importFrom(data.table,is.data.table)
+importFrom(data.table,setDT)
+importFrom(data.table,setcolorder)
+importFrom(data.table,setnames)
diff --git a/R/estimate_R.R b/R/estimate_R.R
index dd0e07c..02c43a8 100644
--- a/R/estimate_R.R
+++ b/R/estimate_R.R
@@ -67,7 +67,7 @@ R_ldsc <- function(Z_hat,
     stopifnot(nrow(N) == J & identical(colnames(N),colnames(Z_hat)))
   }else if(class(N) == "numeric"){
     stopifnot(identical(names(N),colnames(Z_hat)))
-    N <- matrix(rep(N, each = J), nrow = J)
+    N <- matrix(rep(N, each = J), nrow = J, ncol = M, dimnames = dimnames(Z_hat))
   }
 
   if(is.null(comparisons)){
diff --git a/R/gwas_format.R b/R/gwas_format.R
index 138243a..e43d308 100644
--- a/R/gwas_format.R
+++ b/R/gwas_format.R
@@ -10,132 +10,162 @@
 #'@param pos Position column (optional)
 #'@param p_value p-value column (optional)
 #'@param sample_size  Sample size column (optional) or an integer
+#'@param allele_freq Effect allele frequency column (optional)
 #'@param compute_pval Logical, compute the p-value using a normal approximation if missing? Defaults to TRUE.
 #'@param output_file File to write out formatted data. If missing formatted data will be returned.
+#'@param return_og_snps Option to return all input SNPs in original order. SNPs filtered out will have info set to NA
 #'@details This function will try to merge data sets X1 and X2 on the specified columns. Where
 #'necessary, it will flip the sign of effects so that the effect allele is the same in both
 #'data sets. It will remove variants with ambiguous alleles or where the alleles (G/C or A/T) or
 #'with alleles that do not match between data sets (e.g A/G in one data set and A/C in the other).
 #'It will not remove variants that are simply strand flipped between the two data sets (e. g. A/C in one data set, T/G in the other).
-#'@return A data frame with columns chrom, pos, snp, A1, A2, beta_hat, se, p_value, and sample_size with all SNPs
+#'@return A data frame with columns chrom, pos, snp, A1, A2, beta_hat, se, p_value, sample_size, and allele_freq with all SNPs
 #'aligned so that A is the effect allele. This is ready to be used with gwas_merge with formatted = TRUE.
 #'@export
 gwas_format <- function(X, snp, beta_hat, se, A1, A2,
                         chrom, pos, p_value,
                         sample_size, allele_freq,
-                        output_file, compute_pval = TRUE){
+                        output_file, compute_pval = TRUE, return_og_snps = FALSE){
 
+  # --- make data.table ---
+  setDT(X)
+
+  # --- check for missing inputs ---
   if(missing(snp) | missing(beta_hat) | missing(se) | missing(A1) | missing(A2)){
     stop("snp, beta_hat, se, A1, and A2 are required.\n")
   }
-  if(missing(chrom)){
-    X <- mutate(X, chrom = NA)
-    chrom <- "chrom"
-  }else if(is.na(chrom)){
-    X <- mutate(X, chrom = NA)
+ 
+  if(missing(chrom) || is.na(chrom)){
+    X[, chrom := NA]
     chrom <- "chrom"
   }
-  if(missing(pos)){
-    X <- mutate(X, pos = NA_integer_)
-    pos <- "pos"
-  }else if(is.na(pos)){
-    X <- mutate(X, pos = NA_integer_)
+  if(missing(pos) || is.na(pos)){
+    X[, pos := NA_integer_]
     pos <- "pos"
   }
 
-  if(missing(p_value)){
-    X <- mutate(X, p_value = NA_real_)
-    p_value <- "p_value"
-    p_val_missing <- TRUE
-  }else if(is.na(p_value)){
-    X <- mutate(X, p_value = NA_real_)
+  if(missing(p_value) || is.na(p_value)){
+    X[, p_value := NA_real_]
     p_value <- "p_value"
     p_val_missing <- TRUE
   }else{
     p_val_missing <- FALSE
   }
 
-  if(missing(sample_size)){
-    X <- mutate(X, sample_size = NA_real_)
-    sample_size <- "sample_size"
-  }else if(is.na(sample_size)){
-    X <- mutate(X, sample_size = NA_real_)
+  if(missing(sample_size) || is.na(sample_size)){
+    X[, sample_size := NA_real_]
     sample_size <- "sample_size"
   }else if(is.numeric(sample_size)){
-    X <- mutate(X, sample_size = sample_size)
+    X[, sample_size := NA_real_]
     sample_size <- "sample_size"
   }
 
-  if(missing(allele_freq)){
-    X <- mutate(X, af = NA_real_)
-    allele_freq <- "af"
-  }else if(is.na(allele_freq)){
-    X <- mutate(X, af = NA_real_)
+  if(missing(allele_freq) || is.na(allele_freq)){
+    X[, af := NA_real_]
     allele_freq <- "af"
   }
+  
+  # --- keep columns we want and rename ---
+  old_cols <- c(chrom, pos, snp, A1, A2, beta_hat, se, p_value, sample_size, allele_freq)
 
-  keep_cols <- c(chrom, pos, snp, A1, A2, beta_hat, se, p_value, sample_size, allele_freq)
+  new_cols <- c(
+    "chrom", "pos", "snp", "A1", "A2",
+    "beta_hat", "se", "p_value", "sample_size", "allele_freq"
+  )
 
-  X <- X %>%
-      select(all_of(keep_cols))%>%
-      rename(snp = snp,
-            beta_hat =beta_hat,
-            se = se,
-            A1 = A1,
-            A2 = A2,
-            chrom = chrom,
-            pos = pos,
-            p_value = p_value,
-            sample_size = sample_size,
-            allele_freq = allele_freq) %>%
-      mutate(A1 = toupper(A1),
-             A2 = toupper(A2))
+  # checks
+  stopifnot(all(old_cols %chin% names(X)))
 
-  if(p_val_missing & compute_pval){
-    X <- X %>% mutate(p_value = 2*pnorm(-abs(beta_hat/se)))
+  # drop columns not needed, by reference
+  drop_cols <- setdiff(names(X), old_cols)
+  if (length(drop_cols)) {
+    X[, (drop_cols) := NULL]
   }
 
+  # reorder columns
+  setcolorder(X, old_cols)
+
+  # rename columns, by reference
+  setnames(X, old = old_cols, new = new_cols)
+
+  # uppercase alleles, by reference
+  X[, `:=`(
+    A1 = toupper(A1),
+    A2 = toupper(A2)
+  )]
+
   cat("There are ", nrow(X), " variants.\n")
 
+  # --- remove invalid snps ---
+  # before filtering begins, label rows with indices, so we can join back to get all snps if they want
+  # use ids instead of snp names bc may not be unique
+  X[, row_id := .I]
+  og_snp_index <- X[, .(row_id, snp)]
+ 
   #Duplicated variants
-  dup_vars <- X$snp[which(duplicated(X$snp))]
-  X <- X %>% filter(!snp %in% dup_vars)
-  cat("Removing ", length(dup_vars), " duplicated variants leaving ", nrow(X), "variants.\n")
+  # drop ALL instances of a variants which ever appears > 1x
+  dup_vars <- X[duplicated(snp), unique(snp)]
+  X <- X[!(snp %in% dup_vars)]
+  cat("Removing", length(dup_vars), "duplicated variants leaving", nrow(X), "variants.\n")
 
   #Illegal alleles
-  illegal_vars <- X %>%
-                  filter((!A1 %in% c("A", "C", "T", "G") | !A2 %in% c("A", "C", "T", "G") )) %>%
-                  select(snp)
-  if(length(illegal_vars) > 0){
-    X <- X %>% filter(!snp %in% illegal_vars$snp)
-    cat("Removing ", length(illegal_vars), " variants with illegal alleles leaving ", nrow(X), "variants.\n")
+  valid_alleles <- c("A", "C", "T", "G")
+  illegal_vars <- X[
+    !A1 %chin% valid_alleles | !A2 %chin% valid_alleles,
+    snp
+  ]
+  if(length(illegal_vars)){
+    X <- X[!(snp %in% illegal_vars)]
+    cat("Removing", length(illegal_vars), "variants with illegal alleles leaving", nrow(X), "variants.\n")
   }else{
     cat("No variants have illegal alleles.\n")
   }
 
   #Ambiguous alleles
+  # when filtering rows, need to assign to var to keep result
   n <- nrow(X)
-  X <- remove_ambiguous(X, upper = TRUE)
-  cat("Removed ", n-nrow(X), " variants with ambiguous strand.\n")
+  X <- remove_ambiguous(X)
+  cat("Removing", n-nrow(X), " variants with ambiguous strand leaving", nrow(X), "variants.\n")
 
-  # make X into a data.table for I can do the new align_beta on it.  ideally this would be for everything
-  setDT(X)
-  
+  # --- compute pval ----
+  # ask jean do we always want to compute even if not missing
+  if(p_val_missing & compute_pval){
+    X[, beta_hat := as.numeric(beta_hat)]
+    X[, se       := as.numeric(se)]
+    
+    X[, p_value := fifelse(
+      !is.na(beta_hat) & !is.na(se) & se != 0,
+      2 * pnorm(-abs(beta_hat / se)),
+      NA_real_
+    )]
+   cat("Computed p-value\n")
+  }
+
+  # --- harmonize alleles ---
   cat("Flipping strand and effect allele so A1 is always A\n")
   align_beta(X)
 
-  # data table syntax
-  X <- X[, .(chrom, pos, snp, A1, A2, beta_hat, se, p_value, sample_size, allele_freq)]
-
+  # --- write out results ---
+  if (return_og_snps){
+    # label all snps that made it thru filtering as kept
+    #X[, pass_filt := TRUE]
+    # make matrix with all snps given to function
+    X_full <- X[
+      og_snp_index,
+      on = .(row_id, snp)
+    ]
+    # make NAs in kept from dropped rows into false
+    #X_full[is.na(pass_filt), pass_filt := FALSE]
+
+    X <- X_full
+  }
   if(!missing(output_file)){
     cat("Writing out ", nrow(X), " variants to file.\n")
-    # changed from path= to file=
-    write_tsv(X, file = output_file)
+    fwrite(X, file = output_file, sep="\t", na = "NA")  
     return(NULL)
   }
   cat("Returning ", nrow(X), " variants.\n")
   return(X)
-
 }
 
 
@@ -145,76 +175,77 @@ gwas_format <- function(X, snp, beta_hat, se, A1, A2,
 #   return(dat)
 # }
 
-remove_ambiguous <- function(X, upper=TRUE){
-  if(upper){
-    X <- X %>% dplyr::filter(!(A1 == "G" & A2 == "C") &
-                               !(A1 == "C" & A2 == "G") &
-                               !(A1 == "A" & A2 == "T") &
-                               !(A1 == "T" & A2 == "A"))
-    return(X)
-  }
-  X <- X %>% filter(!(A1 == "g" & A2 == "c") &
-                      !(A1 == "c" & A2 == "g") &
-                      !(A1 == "a" & A2 == "t") &
-                      !(A1 == "t" & A2 == "a"))
-  return(X)
+remove_ambiguous <- function(X) {
+  stopifnot(data.table::is.data.table(X))
+  stopifnot(all(c("A1", "A2") %chin% names(X)))
+
+  ambig_pairs <- data.table(A1 = c("G", "C", "A", "T", "g", "c", "a", "t"),
+                            A2 = c("C", "G", "T", "A", "c", "g", "t", "a"))
+
+  idx_ambig <- X[ambig_pairs, on = .(A1, A2), which = TRUE, nomatch = NULL]
+
+  if (!length(idx_ambig)) {
+    return(invisible(X))
+  } else {
+    return(invisible(X[-idx_ambig]))
+  }
 }
 
 # flip signs and strands so that allele 1 is always A
 # now modifies X in-place w/ data table for speed and memory savings
 # believe beta_hat and af are always the names assigned in gwas_format.  but kept for consistency w/ old func
-align_beta <- function(X, upper = TRUE,
-                          beta_col = "beta_hat",
-                          af_col   = "af") {
+align_beta <- function(X, beta_col = "beta_hat", af_col = "allele_freq") {
 
+  # --- checks ---
   stopifnot(is.data.table(X))
-  stopifnot(all(c("A1", "A2") %in% names(X)))
-  stopifnot(beta_col %in% names(X))
-
+  stopifnot(all(c("A1", "A2") %chin% names(X)))
+  if (!is.character(X[["A1"]])) X[, A1 := as.character(A1)]
+  if (!is.character(X[["A2"]])) X[, A2 := as.character(A2)]
+  stopifnot(beta_col %chin% names(X))
+  
+  # --- setup ---
   flp <- c("A"="T","G"="C","T"="A","C"="G",
            "a"="t","t"="a","c"="g","g"="c")
 
-  af_present <- af_col %in% names(X)
-  if (!af_present) {
-    X[, (af_col) := NA_real_]  # create af col temporarily so code can be uniform
-  }
+  af_present <- af_col %chin% names(X)
 
-  # flip strands if we have Ts to get As
-  X[, flip_strands_flag :=
-       if (upper) (A1 == "T" | A2 == "T") else (A1 == "t" | A2 == "t")]
-  X[, `:=`(
-    flipped_A1 = fifelse(flip_strands_flag, flp[A1], A1),
-    flipped_A2 = fifelse(flip_strands_flag, flp[A2], A2)
-  )]
+  # make beta and af numeric if needed
+  if (!is.numeric(X[[beta_col]])) {
+    X[, (beta_col) := as.numeric(get(beta_col))]
+  }
 
-  # flag that is true if the A1 is A (we want)
-  X[, A1_A_flag := flipped_A1 %chin% c("A","a")]
+  if (af_present && !is.numeric(X[[af_col]])) {
+    X[, (af_col) := as.numeric(get(af_col))]
+  }
 
-  # swap A1 and A2 if A1 is not A
-  X[, `:=`(
-    A1 = fifelse(A1_A_flag, flipped_A1, flipped_A2),
-    A2 = fifelse(A1_A_flag, flipped_A2, flipped_A1)
-  )]
+  # --- flipping ---
+  # flip strands if we have Ts to get As
+  idx_flip_strands <- X[, which(A1 %chin% c("T", "t") | A2 %chin% c("T", "t"))]
 
-  # flip beta hats if the A1 is not A
-  # sd means subset of data (select just beta column)
-  X[, (beta_col) := {
-    b <- .SD[[1L]]
-    if (!is.numeric(b)) b <- as.numeric(b)
-    fifelse(A1_A_flag, b, -b)
-  }, .SDcols = beta_col]
-
-  # flip af if the A1 is not A
-  X[, (af_col) := {
-    p <- .SD[[1L]]
-    if (!is.numeric(p)) p <- as.numeric(p)
-    fifelse(A1_A_flag, p, 1 - p)
-  }, .SDcols = af_col]
-
-  # delete columns we don't need anymore
-  X[, c("flip_strands_flag", "flipped_A1", "flipped_A2", "A1_A_flag") := NULL]
-  if (!af_present) X[, (af_col) := NULL]
+  if (length(idx_flip_strands)) {
+    X[idx_flip_strands, `:=`(
+      A1 = unname(flp[A1]),
+      A2 = unname(flp[A2])
+    )]
+  }
+  
+  # we want A1 as A 
+  idx_swap_alleles <- X[, which(!(A1 %chin% c("A", "a")))]
+  
+  if (length(idx_swap_alleles)) {
+    X[idx_swap_alleles, `:=`(
+      A1 = A2,
+      A2 = A1
+    )]
+
+    # flip beta and af if A1 was not A
+    X[idx_swap_alleles, (beta_col) := -get(beta_col)]
+    if (af_present){
+      X[idx_swap_alleles, (af_col)   := 1 - get(af_col)]
+    }
+  }
 
+  # --- return ---
   # since these are in-place mods, we can call func w/o assignment to new var.  but nobody wants to see the whole table
   return(invisible(X))
 }
diff --git a/R/package.R b/R/package.R
index d121965..af36761 100644
--- a/R/package.R
+++ b/R/package.R
@@ -7,6 +7,6 @@
 #' @import flashier
 #' @import dplyr readr reshape2 lpSolve
 #' @import purrr ggplot2
-#' @import data.table
+#' @importFrom data.table setDT setcolorder setnames data.table is.data.table fifelse fwrite %chin% .I :=
 #' @name GFA
 NULL
diff --git a/man/gwas_format.Rd b/man/gwas_format.Rd
index 3f802ba..6759e7d 100644
--- a/man/gwas_format.Rd
+++ b/man/gwas_format.Rd
@@ -17,7 +17,8 @@ gwas_format(
   sample_size,
   allele_freq,
   output_file,
-  compute_pval = TRUE
+  compute_pval = TRUE,
+  return_og_snps = FALSE
 )
 }
 \arguments{
@@ -41,12 +42,16 @@ gwas_format(
 
 \item{sample_size}{Sample size column (optional) or an integer}
 
+\item{allele_freq}{Effect allele frequency column (optional)}
+
 \item{output_file}{File to write out formatted data. If missing formatted data will be returned.}
 
 \item{compute_pval}{Logical, compute the p-value using a normal approximation if missing? Defaults to TRUE.}
+
+\item{return_og_snps}{Option to return all input SNPs in original order. SNPs filtered out will have info set to NA}
 }
 \value{
-A data frame with columns chrom, pos, snp, A1, A2, beta_hat, se, p_value, and sample_size with all SNPs
+A data frame with columns chrom, pos, snp, A1, A2, beta_hat, se, p_value, sample_size, and allele_freq with all SNPs
 aligned so that A is the effect allele. This is ready to be used with gwas_merge with formatted = TRUE.
 }
 \description{