Solvers in header, package header, WIP #146

inst/include/finalsize.h

@@ -0,0 +1,13 @@
+// Written manually to allow header export
+// copied from https://github.com/r-pkg-examples/rcpp-shared-cpp-functions
+
+#ifndef finalsize_finalsize_H_GEN_
+#define finalsize_finalsize_H_GEN_
+
+// Include the Rcpp Header
+#include <Rcpp.h>
+
+#include "iterative_solver.h"
+#include "newton_solver.h"
+
+#endif // finalsize_finalsize_H_GEN_

inst/include/iterative_solver.h

@@ -0,0 +1,130 @@
+// Copyright 2023 'finalsize' authors. See repository licence in LICENSE.md.
+#ifndef INST_INCLUDE_ITERATIVE_SOLVER_H_
+#define INST_INCLUDE_ITERATIVE_SOLVER_H_
+
+// clang-format off
+#include <Rcpp.h>
+#include <RcppEigen.h>
+// clang-format on
+
+// Enable C++14 via this plugin (Rcpp 0.10.3 or later)
+// [[Rcpp::plugins(cpp14)]]
+// [[Rcpp::depends(RcppEigen)]]
+
+// add to namespace finalsize
+namespace finalsize {
+
+/// function for iterative solver
+// taken from Edwin van Leeuwen at
+// https://gitlab.com/epidemics-r/code_snippets/feature/newton_solver/include/finalsize.hpp
+
+/// @brief Solve an equation using the Newton method with tolerance checking
+/// @param contact_matrix Social contact matrix. Entry \eqn{mm_{ij}} gives
+/// the average number of contacts in group \eqn{i} reported by participants
+/// in group \eqn{j}.
+/// @param demography_vector Demography vector. Entry \eqn{pp_{i}} gives
+/// proportion of total population in group \eqn{i}. Normalised if needed.
+/// @param susceptibility A matrix giving the susceptibility of individuals
+/// in demographic group \eqn{i} and risk group \eqn{j}.
+/// @param iterations Number of solver iterations. Defaults to 10,000.
+/// @param tolerance Solver error tolerance. Solving ends when the error
+/// drops below this tolerance. Defaults to set `1e-6`. Larger tolerance
+/// values are likely to lead to inaccurate final size estimates.
+/// @param step_rate The solver step rate. Defaults to 1.9 as a value found to
+/// work well.
+/// @param adapt_step Boolean, whether the solver step rate should be changed
+/// based on the solver error. Defaults to TRUE.
+/// @return A vector of final sizes.
+inline Eigen::ArrayXd solve_final_size_iterative(
+    const Eigen::MatrixXd &contact_matrix,
+    const Eigen::VectorXd &demography_vector,
+    const Eigen::VectorXd &susceptibility, const int iterations = 10000,
+    const double tolerance = 1e-6, double step_rate = 1.9,
+    const bool adapt_step = true) {
+  // count number of demography groups
+  int nDim = demography_vector.size();
+
+  Eigen::ArrayXi zeros(nDim);
+  zeros.fill(0);
+
+  Eigen::ArrayXd epi_final_size(nDim);  // prev in settings struct
+  epi_final_size.fill(0.5);
+
+  Eigen::MatrixXd contact_matrix_ = contact_matrix;
+  for (int i = 0; i < contact_matrix.rows(); ++i) {
+    // Check if value should be 0 for (limited) performance increase
+    if (demography_vector(i) == 0 || susceptibility(i) == 0 ||
+        contact_matrix.row(i).sum() == 0) {
+      zeros[i] = 1;
+      epi_final_size[i] = 0;
+    }
+    for (int j = 0; j < contact_matrix.cols(); ++j) {
+      if (zeros[j]) {
+        contact_matrix_(i, j) = 0;
+      } else {
+        // Scale contacts appropriately
+        contact_matrix_(i, j) = contact_matrix(i, j) * demography_vector(j);
+      }
+    }
+  }
+
+  Eigen::VectorXd epi_final_size_return(nDim);
+  // define functions to minimise error in final size estimate
+  auto f = [&contact_matrix_, &susceptibility, &zeros](
+               const Eigen::VectorXd &epi_final_size,
+               Eigen::VectorXd &epi_final_size_return) {
+    Eigen::VectorXd s = contact_matrix_ * (-epi_final_size);
+    for (int i = 0; i < contact_matrix_.rows(); ++i) {
+      if (zeros[i] == 1) {
+        epi_final_size_return[i] = 0;
+        continue;
+      }
+      epi_final_size_return[i] = 1;
+
+      epi_final_size_return(i) -= exp(susceptibility(i) * s(i));
+    }
+  };
+
+  double current_error = step_rate * nDim;
+  double error = NAN;
+  const double step_change = 1.1;
+
+  for (auto i = 0; i < iterations; ++i) {
+    f(epi_final_size, epi_final_size_return);
+
+    Eigen::ArrayXd dpi = epi_final_size - epi_final_size_return.array();
+    error = dpi.abs().sum();
+    if (error < tolerance) {
+      epi_final_size -= dpi;
+      break;
+    }
+
+    double change = (current_error - error);
+    if (change > 0) {
+      epi_final_size -= step_rate * dpi;
+      if (adapt_step) {
+        step_rate *= step_change;
+      }
+    } else {
+      epi_final_size -= dpi;
+      if (adapt_step) step_rate = std::max(0.9 * step_rate, 1.0);
+    }
+    current_error = error;
+  }
+  if (current_error / tolerance > 100.0) {
+    Rcpp::warning(
+        "Solver error > 100x solver tolerance, try increasing iterations");
+  }
+
+  // Adjust numerical errors;
+  for (auto i = 0; i < epi_final_size.size(); ++i) {
+    if (zeros[i] ||
+        ((epi_final_size(i) < 0) && (epi_final_size(i) > -tolerance)))
+      epi_final_size(i) = 0;
+  }
+  return epi_final_size;
+}
+
+}  // namespace finalsize
+
+#endif  // INST_INCLUDE_ITERATIVE_SOLVER_H_

inst/include/newton_solver.h

@@ -0,0 +1,119 @@
+// Copyright 2023 'finalsize' authors. See repository licence in LICENSE.md.
+#ifndef INST_INCLUDE_NEWTON_SOLVER_H_
+#define INST_INCLUDE_NEWTON_SOLVER_H_
+
+// clang-format off
+#include <Rcpp.h>
+#include <RcppEigen.h>
+// clang-format on
+
+// Enable C++14 via this plugin (Rcpp 0.10.3 or later)
+// [[Rcpp::plugins(cpp14)]]
+// [[Rcpp::depends(RcppEigen)]]
+
+// add to namespace finalsize
+namespace finalsize {
+
+/// function for newton solver
+// taken from Edwin van Leeuwen at
+// https://gitlab.com/epidemics-r/code_snippets/feature/newton_solver/include/finalsize.hpp
+
+/// @brief Solve an equation using the Newton method with tolerance checking
+/// @param contact_matrix Social contact matrix. Entry \eqn{mm_{ij}} gives
+/// the average number of contacts in group \eqn{i} reported by participants
+/// in group \eqn{j}.
+/// @param demography_vector Demography vector. Entry \eqn{pp_{i}} gives
+/// proportion of total population in group \eqn{i}. Normalised if needed.
+/// @param susceptibility A matrix giving the susceptibility of individuals
+/// in demographic group \eqn{i} and risk group \eqn{j}.
+/// @param iterations Number of solver iterations. Defaults to 10,000.
+/// @param tolerance Solver error tolerance. Solving ends when the error
+/// drops below this tolerance. Defaults to set `1e-6`. Larger tolerance
+/// values are likely to lead to inaccurate final size estimates.
+/// @return A vector of final sizes.
+inline Eigen::ArrayXd solve_final_size_newton(const Eigen::MatrixXd &contact_matrix,
+                                       const Eigen::VectorXd &demography_vector,
+                                       const Eigen::VectorXd &susceptibility,
+                                       const int iterations = 10000,
+                                       const double tolerance = 1e-6) {
+  // count number of demography groups
+  int nDim = demography_vector.size();
+
+  Eigen::ArrayXi zeros(nDim);
+  zeros.fill(0);
+
+  Eigen::ArrayXd epi_final_size(nDim);  // prev in settings struct
+  epi_final_size.fill(0.5);
+
+  Eigen::MatrixXd contact_matrix_ = contact_matrix;
+  for (int i = 0; i < contact_matrix.rows(); ++i) {
+    // Check if value should be 0 for (limited) performance increase
+    if (demography_vector(i) == 0 || susceptibility(i) == 0 ||
+        contact_matrix.row(i).sum() == 0) {
+      zeros[i] = 1;
+      epi_final_size[i] = 0;
+    }
+    for (int j = 0; j < contact_matrix.cols(); ++j) {
+      if (zeros[j] == 1) {
+        contact_matrix_(i, j) = 0;
+      } else {
+        // Scale contacts appropriately
+        // Could add transmissibility (j)?
+        contact_matrix_(i, j) =
+            susceptibility(i) * contact_matrix(i, j) * demography_vector(j);
+      }
+    }
+  }
+
+  Eigen::VectorXd cache_v = epi_final_size;
+  // a function f1 that multiplies the contact matrix by the final size guess
+  // + the log of the guess
+  auto f1 = [&contact_matrix_](const Eigen::VectorXd &x,
+                               Eigen::VectorXd &cache) {
+    cache = -(contact_matrix_ * (1 - x.array()).matrix() +
+              x.array().log().matrix());
+  };
+
+  Eigen::MatrixXd cache_m = contact_matrix_;
+  // a function f2 which adds the negative of the contact matrix
+  // to a diagonal matrix of the current final size guess
+  auto f2 = [&contact_matrix_](const Eigen::VectorXd &x,
+                               Eigen::MatrixXd &cache) {
+    cache = (1.0 / x.array()).matrix().asDiagonal();
+    cache = -contact_matrix_ + cache;
+  };
+
+  // a function dx_f that wraps f1, f2, and performs a matrix solve
+  auto dx_f = [&f1, &f2](const Eigen::VectorXd &x, Eigen::VectorXd &cache,
+                         Eigen::MatrixXd &cache_m) {
+    f2(x, cache_m);
+    f1(x, cache);
+    cache = cache_m.partialPivLu().solve(cache).array();
+  };
+
+  // iterate over n-iterations or until the solver tolerance is met
+  Eigen::VectorXd x(nDim);
+  x.fill(1e-6);
+  double error = 0.0;
+  for (auto i = 0; i < iterations; ++i) {
+    dx_f(x, cache_v, cache_m);
+
+    error = cache_v.array().abs().sum();
+    x += cache_v;
+    if (error < tolerance) {
+      break;
+    }
+  }
+  if (error / tolerance > 100.0) {
+    Rcpp::warning(
+        "Solver error > 100x solver tolerance, try increasing iterations");
+  }
+
+  epi_final_size = 1 - x.array();
+
+  return epi_final_size;
+}
+
+}  // namespace finalsize
+
+#endif  // INST_INCLUDE_NEWTON_SOLVER_H_