src/pifun.cpp


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

#include <RcppArmadillo.h>

using namespace Rcpp;
using namespace arma;

vec removalPiFun ( vec p ){
  int J = p.size();
  vec pi(J);
  pi(0) = p(0);
  for(int j=1; j<J; j++){
    pi(j) = pi(j-1) / p(j-1) * (1-p(j-1)) * p(j);
  }
  return(pi);
}

vec removalPiFun(vec p, uvec times){
  int J = p.size();
  for (int j=0; j<J; j++){
    p(j) = 1 - pow(1 - p(j), times(j));
  }
  return removalPiFun(p);
}

vec doublePiFun( vec p ){
  //p must have 2 columns
  vec pi(3);
  pi(0) = p(0) * (1 - p(1));
  pi(1) = p(1) * (1 - p(0));
  pi(2) = p(0) * p(1);
  return(pi);
}

vec depDoublePiFun( vec p ){
  //p must have 2 columns
  vec pi(2);
  pi(0) = p(0);
  pi(1) = p(1) * (1 - p(0));
  return(pi);
}

vec piFun( vec p , std::string pi_fun ){
  if(pi_fun == "removalPiFun"){
    return(removalPiFun(p));
  } else if(pi_fun == "doublePiFun"){
    return(doublePiFun(p));
  } else if(pi_fun == "depDoublePiFun"){
    return(depDoublePiFun(p));
  } else {
    stop("Invalid pifun type");
  }
}