setup universal MPI - does not parallelise!

docs/NB_filtering_cell_types.Rmd

@@ -7,6 +7,7 @@ output: html_document
 
 ```{r setup, include=FALSE}
 # Import libraries
+setwd("~/PostDoc/RNAseq-noise-model/docs")
 
 #setwd("/wehisan/bioinf/bioinf-data/Papenfuss_lab/projects/mangiola.s/PostDoc/RNAseq-noise-model")
 set.seed(13254)
@@ -67,6 +68,8 @@ counts =
 		read_csv(f) %>% mutate(sample = sample %>% as.character, isoform = isoform %>% as.character)
 	} %>%
 
+  filter(symbol %>% is.na %>% `!`) %>%
+
 	# Median redundant
   do_parallel_start(40, "symbol") %>%
   do({
@@ -116,7 +119,7 @@ counts =
   		  `anova p-value` = 
   		      ifelse(
   		        (.) %>% distinct(`Data base`) %>% nrow > 1,
-  		        (.) %>% aov(`read count normalised log` ~ `Cell type formatted` + `Data base`, .) %>% anova %$%   		    `Pr(>F)` %>%  `[` (1),
+  		        (.) %>% aov(`read count normalised log` ~ `Cell type formatted` + `Data base`, .) %>% anova %$% `Pr(>F)` %>% `[` (1),
   		        (.) %>% aov(`read count normalised log` ~ `Cell type formatted`, .) %>% anova %$% `Pr(>F)` %>% `[` (1)
   		      ) 
 		    )
@@ -133,10 +136,12 @@ counts =
     `hard bimodality` =
       (`bimodal p-value` < 0.05) + 
       (`bimodal coefficient` > 0.6666667) + 
-      (`anova p-value` < 0.05 ) >= 
-      2
+      (`anova p-value` < 0.05 ) >= 2
   ) %>%
-  mutate(`soft bimodality` = `anova p-value` < 0.0001)
+  mutate(`soft bimodality` = `anova p-value` < 0.0001) %>%
+
+  mutate(symbol = symbol %>% as.factor) %>%
+  mutate(sample = sample %>% as.factor) 
 ```
 
 Inference 
@@ -150,44 +155,99 @@ counts_stan =
   counts %>%
   filter(!filt_for_calc) %>%
   filter(!`hard bimodality`) %>%
-  #inner_join( (.) %>% distinct(symbol) %>% sample_n(500) ) %>%
+  inner_join( (.) %>% distinct(symbol) %>% sample_n(66) ) %>%
 
   left_join(
     (.) %>% 
     distinct(sample, symbol, `read count normalised`) %>% 
+    arrange(symbol) %>%
     count(symbol) %>% 
     mutate(end = cumsum(n)) %>%
     mutate(start = c(1, .$end %>% rev() %>% `[` (-1) %>% rev %>% `+` (1)))
   ) %>% 
-  arrange(symbol) 
+  arrange(symbol) %>%
+
+  # Add sample idx
+  mutate(sample_idx = as.integer(sample)) 
 
-data_for_stan = list(
-  N = counts_stan %>% nrow,
-  G = counts_stan %>% distinct(symbol) %>% nrow,
-  S = counts_stan %>% distinct(sample) %>% nrow,
-  symbol_start_end = counts_stan %>% distinct(start, end) %>% select(start, end),
-  counts = counts_stan %>% select(`read count normalised`) %>% pull(1),
-  sample_idx = 
-    counts_stan %>%
-    mutate(sample = factor(sample)) %>% 
-    mutate(sample_idx = as.integer(sample)) %>% 
-    select(sample_idx) %>%
-    pull(1),
+shards = 5
+counts_stan_MPI =
+  counts_stan %>%
+
+  # Add index for MPI
+  left_join(
+    (.) %>% 
+      distinct(symbol) %>%
+      mutate( idx_MPI = head( rep(1:shards, (.) %>% nrow %>% `/` (shards) %>% ceiling ) %>% sort, n=(.) %>% nrow) )
+  )
+
+data_for_stan_MPI = list(
+
+  N = counts_stan_MPI %>% distinct(sample, symbol, idx_MPI) %>% count(idx_MPI) %>% pull(n) %>% max,
+  M = counts_stan_MPI %>% distinct(start, idx_MPI) %>% count(idx_MPI) %>% pull(n) %>% max,
+  G = counts_stan_MPI %>% distinct(symbol) %>% nrow(),
+  S = counts_stan_MPI %>% distinct(sample) %>% nrow(),
+  n_shards = shards,
+  G_per_shard = counts_stan_MPI %>% distinct(symbol, idx_MPI) %>% count(idx_MPI) %>% pull(n) ,
+  G_per_shard_idx = c(0, counts_stan_MPI %>% distinct(symbol, idx_MPI) %>% count(idx_MPI) %>% pull(n) %>% cumsum),
+
+  counts = foreach(idx = counts_stan_MPI %>% distinct(idx_MPI) %>% pull(idx_MPI), .combine = full_join) %do% {
+    counts_stan_MPI %>%
+    filter(idx_MPI == idx) %>%
+    select(`read count`) %>% 
+    setNames(idx) %>%
+    mutate(i = 1:n())
+  } %>%
+  select(-i) %>%
+    replace(is.na(.), 0) %>%
+    mutate_if(is.numeric, as.integer) %>%
+    as_matrix() %>% t,
+
+  symbol_start =   foreach(idx = counts_stan_MPI %>% distinct(idx_MPI) %>% pull(idx_MPI), .combine = full_join) %do% {
+    counts_stan_MPI %>%
+    filter(idx_MPI == idx) %>%
+    distinct(start) %>% 
+    rbind(  counts_stan_MPI %>% filter(idx_MPI == idx) %>% pull(end) %>% max) %>%
+    mutate(start = start - ( (.) %>% head(n=1) %>% pull(start)) + 1) %>%
+    setNames(idx) %>%
+    mutate(i = 1:n())
+  } %>%
+  select(-i) %>%
+    replace(is.na(.), 0) %>%
+    mutate_if(is.numeric, as.integer) %>%
+    as_matrix() %>% t,
+
+  sample_idx =  foreach(idx = counts_stan_MPI %>% distinct(idx_MPI) %>% pull(idx_MPI), .combine = full_join) %do% {
+    counts_stan_MPI %>%
+    filter(idx_MPI == idx) %>%
+    select(sample_idx) %>% 
+    setNames(idx) %>%
+    mutate(i = 1:n())
+  } %>%
+  select(-i) %>%
+  replace(is.na(.), 0) %>%
+    mutate_if(is.numeric, as.integer) %>%
+    as_matrix() %>% t
 
-  # Info on data set
-  omit_data = 0,
-  generate_quantities = 1,
-  is_prior_asymetric = 1
 )
 
-# fit =
-#   sampling(
-#       nb_model,
-#     data = data_for_stan, 
-#     chains=4, iter=500
-# )
-# fit_parsed = fit %>% spread_draws(lambda[G], sigma[G], sigma_raw[G]) 
-#fit %>% summary() %$% summary %>% as_tibble(rownames="par")
+file("~/.R/Makevars") %>% {
+  writeLines(c("CXX14FLAGS += -O3", "CXXFLAGS += -DSTAN_THREADS", "CXXFLAGS += -pthread"), (.)); 
+  (.)
+} %>%
+close(.)
+nb_model = stan_model(here("stan",sprintf("%s.stan", "negBinomial_tidy_MPI")))
+Sys.setenv("STAN_NUM_THREADS" = 8)
+
+
+fit =
+  sampling(
+      nb_model,
+    data = data_for_stan_MPI,
+    chains=1, iter=1
+)
+fit_parsed = fit %>% spread_draws(lambda[G], sigma[G], sigma_raw[G])
+fit %>% summary() %$% summary %>% as_tibble(rownames="par")
 
 
 load("fit_parsed_all_genes_t_cell_18022019.RData")

stan/negBinomial_tidy_MPI.stan

@@ -1,136 +1,116 @@
 functions{
-	real gamma_log_lpdf(vector x_log, real a, real b){
 
-		// This function is the  probability of the log gamma funnction
-		// in case you have data that is aleady in log form
+	 vector lp_reduce( vector global_parameters , vector local_parameters , real[] xr , int[] xi ) {
+	 	int M = xi[1];
+	 	int N = xi[2];
+	 	int G_per_shard = xi[3];
+	 	int symbol_start[M+1] = xi[(3+1):(3+1+M)];
+	 	int sample_idx[N] = xi[(3+1+M+1):(3+1+M+1+N-1)];
+	 	int counts[N] = xi[(3+1+M+1+N):size(xi)];
 
-		vector[rows(x_log)] jacob = x_log; //jacobian
-		real norm_constant = a * log(b) -lgamma(a);
-		real a_minus_1 = a-1;
-		return sum( jacob ) + norm_constant * rows(x_log) + sum(  x_log * a_minus_1 - exp(x_log) * b ) ;
 
-	}
-
-
-	vector gen_inv_logit(vector b1s, vector X_no_intercept, real inflection, real y_cross) {
-
-		real b0 = -inflection * b1s[1];
-		vector[num_elements(b1s) + 1] b = append_row(b0, b1s);
-		matrix[rows(X_no_intercept), num_elements(b1s) + 1] X = append_col(rep_vector(1, rows(X_no_intercept)), X_no_intercept);
-
-		return  y_cross * (1 + exp(-b0)) ./ to_vector(1 + exp(- (  X * b ))) ;
-	}
+	 	vector[G_per_shard] lambda_MPI = local_parameters[1:G_per_shard];
+	 	vector[G_per_shard] sigma_MPI = local_parameters[(G_per_shard+1):(G_per_shard*2)];
+	 	vector[G_per_shard] lp;
 
-	 vector lp_reduce( vector beta , vector theta , real[] xr , int[] xi ) {
-	 	int M = rows(theta)/2;
-	 	int S = rows(to_vector(xi))/M;
-	 	vector[M] lambda_MPI = theta[1:M];
-	 	vector[M] sigma_MPI = theta[(M+1):(M*2)];
-	 	vector[M] lp;
-		int counts[M, S];
 
-		for(m in 1:M){
-			int j = 1 + (m-1)*S;
-	    int k = m*S;
-			counts[m] = xi[j:k];
-		}
-
-	 for(m in 1:M)
-	 	lp[m] =  neg_binomial_2_log_lpmf(counts[m] | lambda_MPI[m], sigma_MPI[m]);
+	 for(g in 1:G_per_shard){
+	 	lp[g] =  neg_binomial_2_log_lpmf(counts[symbol_start[g]:symbol_start[g+1]] | lambda_MPI[g], sigma_MPI[g]);
+	 }
 
 
     return [sum(lp)]';
   }
 
-
 }
 data {
+  int<lower=0> N;
+  int<lower=0> M;
 	int<lower=0> G;
 	int<lower=0> S;
-	int<lower=0> counts[G, S];
-
+  int n_shards;
+	int<lower=0> counts[n_shards, N];
+	int<lower=0> symbol_start[n_shards, M+1];
+	int<lower=0> sample_idx[n_shards, N];
+	int<lower=0> G_per_shard[n_shards];
+	int<lower=0> G_per_shard_idx[n_shards + 1];
 
 }
 transformed data {
-  // 7 shards
-  // M = N/7 = 124621/7 = 17803
-  int n_shards = 10;
-  int M = G/n_shards;
-  int counts_MPI[n_shards, M*S];  // 2M because two variables, and they get stacked in array
-  real xr[n_shards, 0];
   vector[0] global_parameters;
-  // an empty set of per-shard parameters
-  //vector[M] theta[n_shards];
-  // split into shards
+  real xr[n_shards, 0];
+
+  int<lower=0> int_MPI[n_shards, 3+(M+1)+N+N];
+
+  // Shards - MPI
   for ( i in 1:n_shards ) {
-    int j = 1 + (i-1)*M;
-    int k = i*M;
-    counts_MPI[i,] = to_array_1d(counts[ j:k ]);
+  int M_N_Gps[3];
+  M_N_Gps[1] = M;
+  M_N_Gps[2] = N;
+  M_N_Gps[3] = G_per_shard[i];
+
+  int_MPI[i,] = append_array(append_array(append_array(M_N_Gps, symbol_start[i]), sample_idx[i]), counts[i]);
   }
 }
 parameters {
+  // Overall properties of the data
+  real lambda_mu; // So is compatible with logGamma prior
+  real<lower=0> lambda_sigma;
+  real lambda_skew;
+  vector[S] exposure_rate;
 
-	// Overall properties of the data
-	real<lower=0> lambda_mu; // So is compatible with logGamma prior
-	real<lower=0> lambda_sigma_raw;
-	vector[S] exposure_rate;
+  // Gene-wise properties of the data
+  vector[G] lambda;
+  vector<lower=0>[G] sigma_raw;
 
-	// Gene-wise properties of the data
-	vector[G] lambda;
-	vector<lower=0>[G] sigma_raw;
-
-	// Signa linear model
+  // Signa linear model
 
-	real sigma_inflection;
-	real<lower=0> sigma_y_cross;
-	vector<upper=0>[1] sigma_slope;
-	real<lower=0>sigma_sigma;
+  real<upper=0> sigma_slope;
+  real sigma_intercept;
+  real<lower=0>sigma_sigma;
 
 }
 transformed parameters {
-	real<lower=0> lambda_sigma = lambda_sigma_raw / 1000;
-	vector[G] sigma = 1.0 ./ sigma_raw;
-
-	// Sigma linear model
-	//vector[G] sigma_mu = exp(lambda * sigma_slope + sigma_0); //the expected values (linear predictor)
-	vector[G] sigma_mu = gen_inv_logit(sigma_slope, lambda, sigma_inflection, sigma_y_cross) ;
-	vector[G] sigma_alpha = sigma_mu .* sigma_mu / sigma_sigma; //shape parameter for the gamma distribution
-	vector[G] sigma_beta = sigma_mu / sigma_sigma; //rate parameter for the gamma distribution
+  // Sigma
+  vector[G] sigma = 1.0 ./ sigma_raw;
 
+// Shards - MPI
 	vector[2*M] lambda_sigma_MPI[n_shards];
+	for( i in 1:(n_shards) ) {
+
+	  vector[ (M*2) - (G_per_shard[i]*2) ] buffer = rep_vector(0.0,(M*2) - (G_per_shard[i]*2));
+
+		lambda_sigma_MPI[i] = append_row(append_row(
+		  lambda[(G_per_shard_idx[i]+1):(G_per_shard_idx[i+1])],
+		  sigma[(G_per_shard_idx[i]+1):(G_per_shard_idx[i+1])]
+		), buffer);
 
-	for( i in 1:n_shards ) {
-		int j = 1 + (i-1)*M;
-    int k = i*M;
-		lambda_sigma_MPI[i] = append_row(lambda[j:k], sigma[j:k]);
 	}
+
 }
 model {
 
-	// Overall properties of the data
-	lambda_mu ~ gamma(3,2);
-	lambda_sigma_raw ~ normal(0,1);
-	//sigma_raw ~ normal(0,1);
-	exposure_rate ~ normal(0,1);
-	sum(exposure_rate) ~ normal(0, 0.001 * S);
+  // Overall properties of the data
+int i = 1;
+  lambda_mu ~ normal(0,2);
+  lambda_sigma ~ normal(0,2);
+  lambda_skew ~ normal(0,2);
 
-	sigma_inflection ~ normal(0,2);
-	sigma_y_cross ~ cauchy(0,2);
-	sigma_slope ~ normal(0,1);
-	sigma_sigma ~ cauchy(0,2.5);
+  //sigma_raw ~ normal(0,1);
+  exposure_rate ~ normal(0,1);
+  sum(exposure_rate) ~ normal(0, 0.001 * S);
 
-	// Gene-wise properties of the data
-	target += sum( map_rect( lp_reduce , global_parameters , lambda_sigma_MPI , xr , counts_MPI ) );
+  sigma_intercept ~ normal(0,2);
+  sigma_slope ~ normal(0,2);
+  sigma_sigma ~ normal(0,2);
 
-  lambda ~ gamma_log_lpdf(lambda_mu, lambda_sigma);
-	sigma_raw ~ gamma(sigma_alpha,sigma_beta);
+  // Gene-wise properties of the data
+  // lambda ~ normal_or_gammaLog(lambda_mu, lambda_sigma, is_prior_asymetric);
+  lambda ~  skew_normal(lambda_mu,lambda_sigma, lambda_skew);
+  sigma_raw ~ lognormal(sigma_slope * lambda + sigma_intercept,sigma_sigma);
 
-	// Sample from data
-	// if(omit_data==0) for(g in 1:G)
-	// 	counts[symbol_start_end[g,1]:symbol_start_end[g,2]] ~ neg_binomial_2_log(
-	// 		exposure_rate[sample_idx[symbol_start_end[g,1]:symbol_start_end[g,2]]] + lambda[g],
-	// 		sigma[g]
-	// 	);
+	// Gene-wise properties of the data
+	target += sum( map_rect( lp_reduce , global_parameters , lambda_sigma_MPI , xr , int_MPI ) );
 
 }