Adding estimated group methylation levels and dispersion to output

src/radmeth/radmeth.cpp

@@ -176,40 +176,11 @@ enum class row_status : std::uint8_t {
   na_extreme_cnt,
 };
 
-[[nodiscard]] static std::vector<double>
-drop_idx(const std::vector<double> &v, const std::size_t idx_to_drop) {
-  std::vector<double> u;
-  u.reserve(std::size(v) - 1);
-  for (auto i = 0u; i < std::size(v); ++i)
-    if (i != idx_to_drop)
-      u.push_back(v[i]);
-  return u;
-}
-
-/// ADS: this function is not currently used, as the threads do not operate in
-/// "chunks"
-/*
-static inline void
-get_chunk_bounds(const std::uint32_t n_elements, const std::uint32_t n_chunks,
-                 std::vector<std::pair<std::uint32_t, std::uint32_t>> &chunks) {
-  const std::uint32_t q = n_elements / n_chunks;
-  const std::uint32_t r = n_elements - q * n_chunks;
-  std::uint32_t block_start{};
-  for (std::size_t i = 0; i < n_chunks; ++i) {
-    const auto sz = (i < r) ? q + 1 : q;
-    const auto block_end = block_start + sz;
-    chunks[i] = {block_start, block_end};
-    block_start = block_end;
-  }
-}
-*/
-
 static void
 radmeth(const bool show_progress, const bool more_na_info,
         const std::uint32_t n_threads, const std::string &table_filename,
         const std::string &outfile, const Regression &alt_model,
         const Regression &null_model, const std::uint32_t test_factor_idx) {
-  static constexpr auto suffix_fmt = "\t%ld\t%ld\t%ld\t%ld\n";
   static constexpr auto buf_size = 1024;
   static constexpr auto max_lines = 16384;
 
@@ -251,7 +222,7 @@ that the design matrix and the proportion table are correctly formatted.
   std::vector<Regression> alt_models(n_threads, alt_model);
   std::vector<Regression> null_models(n_threads, null_model);
 
-  // std::vector<std::pair<std::uint32_t, std::uint32_t>> chunks(n_threads);
+  const auto n_groups = alt_model.n_groups();
 
   // Iterate over rows in the file and do llr test on proportions from
   // each. Do this in sets of rows to avoid having to spawn too many threads.
@@ -262,19 +233,20 @@ that the design matrix and the proportion table are correctly formatted.
     if (n_lines == 0)
       break;
 
-    /// ADS: chunks not used
-    // get_chunk_bounds(n_lines, n_threads, chunks);
-
     std::vector<std::thread> threads;
     for (auto thread_id = 0u; thread_id < n_threads; ++thread_id) {
       threads.emplace_back([&, thread_id] {
-        /// ADS: chunks not used
-        // const auto &[chunk_beg, chunk_end] = chunks[thread_id];
+        std::vector<double> p_estim_alt;
+        std::vector<double> p_estim_null;
+        double phi_estim_alt{};
+        double phi_estim_null{};
+
         auto &t_alt_model = alt_models[thread_id];
         auto &t_null_model = null_models[thread_id];
-        /// ADS: chunks not used
-        // for (auto b = chunk_beg; b != chunk_end; ++b) {
         for (auto b = 0u; b < n_lines; ++b) {
+          // ADS: rows done by different threads are interleaved because the
+          // difficult (e.g., high-coverage) rows can be consecutive and this
+          // balances work better.
           if (b % n_threads != thread_id)
             continue;
           t_alt_model.props.parse(lines[b]);
@@ -291,37 +263,19 @@ that the design matrix and the proportion table are correctly formatted.
             if (has_extreme_counts(t_alt_model))
               return std::tuple{1.0, row_status::na_extreme_cnt};
 
-            std::vector<double> alternate_params;
-            fit_regression_model(t_alt_model, alternate_params);
+            fit_regression_model(t_alt_model, p_estim_alt, phi_estim_alt);
+
             t_null_model.props = t_alt_model.props;
 
-            auto null_params = drop_idx(alternate_params, test_factor_idx);
+            fit_regression_model(t_null_model, p_estim_null, phi_estim_null);
 
-            fit_regression_model(t_null_model, null_params);
             const double p_value =
               llr_test(t_null_model.max_loglik, t_alt_model.max_loglik);
 
             return (p_value != p_value) ? std::tuple{1.0, row_status::na}
                                         : std::tuple{p_value, row_status::ok};
           }();
 
-          std::size_t n_reads_factor = 0;
-          std::size_t n_reads_others = 0;
-          std::size_t n_meth_factor = 0;
-          std::size_t n_meth_others = 0;
-
-          const auto &mc = t_alt_model.props.mc;
-          const auto &vec = t_alt_model.design.tmatrix[test_factor_idx];
-          for (std::size_t s = 0; s < n_samples; ++s)
-            if (vec[s] != 0) {
-              n_reads_factor += mc[s].n_reads;
-              n_meth_factor += mc[s].n_meth;
-            }
-            else {
-              n_reads_others += mc[s].n_reads;
-              n_meth_others += mc[s].n_meth;
-            }
-
           n_bytes[b] = [&] {
             // clang-format off
             const int n_prefix_bytes =
@@ -349,17 +303,26 @@ that the design matrix and the proportion table are correctly formatted.
 
             cursor += n_pval_bytes;
 
-            // clang-format off
-            const int n_suffix_bytes = std::sprintf(cursor, suffix_fmt,
-                                                    n_reads_factor,
-                                                    n_meth_factor,
-                                                    n_reads_others,
-                                                    n_meth_others);
-            // clang-format on
-            if (n_suffix_bytes < 0)
-              return n_suffix_bytes;
+            const int n_param_bytes = [&] {
+              std::int32_t n_bytes = 0;
+              for (auto g_idx = 0u; g_idx < n_groups; ++g_idx) {
+                const int n = std::sprintf(cursor, "\t%f", p_estim_alt[g_idx]);
+                cursor += n;
+                if (n < 0)
+                  return n;
+                n_bytes += n;
+              }
+              const int n = std::sprintf(cursor, "\t%f\n", phi_estim_alt);
+              if (n < 0)
+                return n;
+              n_bytes += n;
+              return n_bytes;
+            }();
+
+            if (n_param_bytes < 0)
+              return n_param_bytes;
 
-            return n_prefix_bytes + n_pval_bytes + n_suffix_bytes;
+            return n_prefix_bytes + n_pval_bytes + n_param_bytes;
           }();
         }
       });
@@ -450,6 +413,7 @@ main_radmeth(int argc, char *argv[]) {
     // initialize full design matrix from file
     Regression alt_model;
     alt_model.design = read_design(design_filename);
+
     if (verbose)
       std::cerr << "Alternate model:\n" << alt_model.design << '\n';
 

src/radmeth/radmeth_model.cpp

@@ -169,18 +169,37 @@ Design::drop_factor(const std::size_t factor_idx) {
 
 std::ostream &
 operator<<(std::ostream &out, const Design &design) {
-  for (std::size_t factor = 0; factor < design.factor_names.size(); ++factor) {
-    if (factor != 0)
-      out << '\t';
-    out << design.factor_names[factor];
+  static constexpr std::uint32_t max_samples_to_report = 200;
+  const auto n_samples = design.n_samples();
+  const auto n_factors = design.n_factors();
+  if (n_samples <= max_samples_to_report) {
+    for (std::size_t factor = 0; factor < n_factors; ++factor) {
+      if (factor != 0)
+        out << '\t';
+      out << design.factor_names[factor];
+    }
+    out << '\n';
+
+    for (std::size_t i = 0; i < n_samples; ++i) {
+      out << design.sample_names[i];
+      for (std::size_t j = 0; j < n_factors; ++j)
+        out << '\t' << static_cast<std::uint32_t>(design.matrix[i][j]);
+      out << '\n';
+    }
   }
-  out << '\n';
 
-  for (std::size_t i = 0; i < design.n_samples(); ++i) {
-    out << design.sample_names[i];
-    for (std::size_t j = 0; j < design.n_factors(); ++j)
-      out << "\t" << static_cast<std::uint32_t>(design.matrix[i][j]);
-    out << "\n";
+  // compute the number of samples per group
+  const auto n_groups = design.n_groups();
+  std::vector<std::uint32_t> n_samples_per_group(n_groups, 0);
+  for (auto s_idx = 0u; s_idx < n_samples; ++s_idx)
+    ++n_samples_per_group[design.group_id[s_idx]];
+
+  out << "group_id | factor_levels (" << n_factors << " factors) | n_samples\n";
+  for (std::size_t g_idx = 0; g_idx < n_groups; ++g_idx) {
+    out << g_idx;
+    for (std::size_t f_idx = 0; f_idx < n_factors; ++f_idx)
+      out << '\t' << static_cast<std::uint32_t>(design.groups[g_idx][f_idx]);
+    out << '\t' << n_samples_per_group[g_idx] << '\n';
   }
 
   return out;

src/radmeth/radmeth_model.hpp

@@ -107,6 +107,11 @@ struct Regression {
     return design.n_groups();
   }
 
+  [[nodiscard]] std::size_t
+  n_params() const {
+    return n_factors() + 1;
+  }
+
   [[nodiscard]] std::size_t
   props_size() const {
     return std::size(props.mc);

src/radmeth/radmeth_optimize.cpp

@@ -212,7 +212,7 @@ get_cumulative(const std::vector<std::uint32_t> &group_id,
       const auto val = get_value(mc[c_idx]);
       auto &vec = get_vector(cumul[g_idx]);
       for (auto i = 0u; i < val; ++i)
-        vec[i]++;
+        ++vec[i];
     }
   };
   // call the lambda 3 times for m_counts, r_counts, d_counts
@@ -229,24 +229,20 @@ get_cumulative(const std::vector<std::uint32_t> &group_id,
     [](cumul_counts &c) -> std::vector<std::uint32_t> & { return c.d_counts; });
 }
 
-[[nodiscard]] bool
-fit_regression_model(Regression &r, std::vector<double> &params_init) {
+void
+fit_regression_model(Regression &r, std::vector<double> &p_estimates,
+                     double &dispersion_estimate) {
   static constexpr auto init_dispersion_param = -2.5;
   const auto stepsize = Regression::stepsize;
   const auto max_iter = Regression::max_iter;
 
-  get_cumulative(r.design.group_id, r.design.n_groups(), r.props.mc, r.cumul);
+  const auto n_groups = r.n_groups();
+  get_cumulative(r.design.group_id, n_groups, r.props.mc, r.cumul);
   set_max_r_count(r);
 
-  // one more than the number of factors
-  const std::size_t n_params = r.n_factors() + 1;
-  if (params_init.empty()) {
-    params_init.resize(n_params, 0.0);
-    params_init.back() = init_dispersion_param;
-  }
-  if (std::size(params_init) != n_params)
-    throw std::runtime_error("Wrong number of initial parameters.");
-  r.p_v.resize(r.n_groups());
+  r.p_v.resize(n_groups);
+
+  const std::size_t n_params = r.n_params();
   const auto tol = std::sqrt(n_params) * r.n_samples() * Regression::tolerance;
   // clang-format off
   auto loglik_bundle = gsl_multimin_function_fdf{
@@ -258,6 +254,12 @@ fit_regression_model(Regression &r, std::vector<double> &params_init) {
   };
   // clang-format on
 
+  // set the parameters: zero for "p" parameters and the final one for
+  // dispersion using the constant
+  auto params = gsl_vector_alloc(n_params);
+  gsl_vector_set_all(params, 0.0);
+  gsl_vector_set(params, n_params - 1, init_dispersion_param);
+
   // Alternatives:
   // - gsl_multimin_fdfminimizer_conjugate_pr
   // - gsl_multimin_fdfminimizer_conjugate_fr
@@ -266,10 +268,6 @@ fit_regression_model(Regression &r, std::vector<double> &params_init) {
   const auto minimizer = gsl_multimin_fdfminimizer_conjugate_pr;
   auto s = gsl_multimin_fdfminimizer_alloc(minimizer, n_params);
 
-  auto params = gsl_vector_alloc(n_params);
-  for (auto i = 0u; i < n_params; ++i)
-    gsl_vector_set(params, i, params_init[i]);
-
   gsl_multimin_fdfminimizer_set(s, &loglik_bundle, params, stepsize, tol);
 
   int status = 0;
@@ -282,11 +280,23 @@ fit_regression_model(Regression &r, std::vector<double> &params_init) {
     status = gsl_multimin_test_gradient(s->gradient, tol);
   } while (status == GSL_CONTINUE && ++iter < max_iter);
 
+  /// ADS: the condition below might not always pass even if we are doing
+  /// ok. It's not clear how to check for failure.
+
+  // if (status != GSL_SUCCESS)
+  //   throw std::runtime_error("failed to fit model parameters");
+
   const auto param_estimates = gsl_multimin_fdfminimizer_x(s);
+
+  const auto &groups = r.design.groups;
+  p_estimates.clear();
+  for (auto g_idx = 0u; g_idx < n_groups; ++g_idx)
+    p_estimates.push_back(get_p(groups[g_idx], param_estimates));
+  const auto disp_param = gsl_vector_get(param_estimates, n_params - 1);
+  dispersion_estimate = 1.0 / std::exp(disp_param);
+
   r.max_loglik = log_likelihood(param_estimates, r);
 
   gsl_multimin_fdfminimizer_free(s);
   gsl_vector_free(params);
-
-  return status == GSL_SUCCESS;
 }

src/radmeth/radmeth_optimize.hpp

@@ -23,7 +23,8 @@
 
 struct Regression;
 
-bool
-fit_regression_model(Regression &r, std::vector<double> &initial_params);
+void
+fit_regression_model(Regression &r, std::vector<double> &p_estimates,
+                     double &dispersion_estimate);
 
 #endif