From 3befd9e9cff555bba600ec808df16b1eaba3aa1b Mon Sep 17 00:00:00 2001
From: Gibraan Rahman <gibsramen@gmail.com>
Date: Tue, 12 Sep 2023 09:57:54 -0700
Subject: [PATCH] Fix custom model test

---
 tests/custom_model.stan    | 35 ++++++++++++++++++++---------------
 tests/test_custom_model.py | 10 ++++++----
 2 files changed, 26 insertions(+), 19 deletions(-)

diff --git a/tests/custom_model.stan b/tests/custom_model.stan
index e67156b..1bf2b62 100644
--- a/tests/custom_model.stan
+++ b/tests/custom_model.stan
@@ -1,42 +1,47 @@
-# Separate priors for intercept and host_common_name
 data {
   int<lower=0> N;       // number of samples
   int<lower=0> D;       // number of dimensions
+  real A;                                   // mean intercept
   int<lower=0> p;       // number of covariates
-  real depth[N];        // sequencing depths of microbes
+  vector[N] depth;      // sequencing depths of microbes
   matrix[N, p] x;       // covariate matrix
-  int y[N, D];          // observed microbe abundances
-  real<lower=0> B_p_1;  // stdev for Beta Normal prior
-  real<lower=0> B_p_2;  // stdev for Beta Normal prior
+  array[N, D] int y;    // observed microbe abundances
+  real<lower=0> B_p;    // stdev for intercept
   real<lower=0> phi_s;  // constant phi value
 }
 
 parameters {
   // parameters required for linear regression on the species means
-  matrix[p, D-1] beta_var;
+  row_vector<offset=A, multiplier=B_p>[D-1] beta_0;
+  matrix<multiplier=B_p>[p-1, D-1] beta_x;
+  real inv_disp;
 }
 
 transformed parameters {
+  matrix[p, D-1] beta_var = append_row(beta_0, beta_x);
   matrix[N, D-1] lam;
   matrix[N, D] lam_clr;
-  vector[N] z;
 
-  z = to_vector(rep_array(0, N));
-  lam = x * beta_var;
-  lam_clr = append_col(z, lam);
+  lam = x*beta_var;
+  for (n in 1:N){
+    lam[n] += depth[n];
+  }
+  lam_clr = append_col(to_vector(rep_array(0, N)), lam);
 }
 
 model {
-  // setting priors ...
+  inv_disp ~ lognormal(0, phi_s);
+
   for (i in 1:D-1){
-    beta_var[1, i] ~ normal(0., B_p_1); // uninformed prior
-    beta_var[2, i] ~ normal(0., B_p_2); // uninformed prior
+    for (j in 1:p){
+      beta_var[j, i] ~ normal(0., B_p); // uninformed prior
+    }
   }
+
   // generating counts
   for (n in 1:N){
     for (i in 1:D){
-      target += neg_binomial_2_log_lpmf(y[n, i] | depth[n] + lam_clr[n, i], phi_s);
+      target += neg_binomial_2_log_lpmf(y[n, i] | lam_clr[n, i], inv(inv_disp));
     }
   }
 }
-
diff --git a/tests/test_custom_model.py b/tests/test_custom_model.py
index 8a982d3..2da859a 100644
--- a/tests/test_custom_model.py
+++ b/tests/test_custom_model.py
@@ -7,8 +7,10 @@
 
 
 def test_custom_model(table_biom, metadata):
-    # Negative binomial model with separate prior values for intercept &
-    # host_common_name effect and constant overdispersion parameter.
+    # Negative binomial model with constant overdispersion parameter.
+    D = table_biom.shape[0]
+    A = np.log(1 / D)
+
     custom_model = TableModel(
         table=table_biom,
         model_path=resource_filename("tests", "custom_model.stan"),
@@ -19,10 +21,10 @@ def test_custom_model(table_biom, metadata):
     )
     custom_model.add_parameters(
         {
-            "B_p_1": 2.0,
-            "B_p_2": 5.0,
+            "B_p": 2.0,
             "phi_s": 0.2,
             "depth": np.log(table_biom.sum(axis="sample")),
+            "A": A
         }
     )
     custom_model.specify_model(