Ben Letham 7 лет назад
Родитель
Сommit
aa37fb48ac
1 измененных файлов с 53 добавлено и 44 удалено
  1. 53 44
      python/stan/win/prophet.stan

+ 53 - 44
python/stan/win/prophet.stan

@@ -1,13 +1,13 @@
 functions {
 functions {
-  matrix get_changepoint_matrix(vector t, vector t_change, int T, int S) {
+  real[ , ] get_changepoint_matrix(real[] t, real[] t_change, int T, int S) {
     // Assumes t and t_change are sorted.
     // Assumes t and t_change are sorted.
-    matrix[T, S] A;
-    row_vector[S] a_row;
+    real A[T, S];
+    real a_row[S];
     int cp_idx;
     int cp_idx;
 
 
     // Start with an empty matrix.
     // Start with an empty matrix.
-    A = rep_matrix(0, T, S);
-    a_row = rep_row_vector(0, S);
+    A = rep_array(0, T, S);
+    a_row = rep_array(0, S);
     cp_idx = 1;
     cp_idx = 1;
 
 
     // Fill in each row of A.
     // Fill in each row of A.
@@ -23,9 +23,9 @@ functions {
 
 
   // Logistic trend functions
   // Logistic trend functions
 
 
-  vector logistic_gamma(real k, real m, vector delta, vector t_change, int S) {
-    vector[S] gamma;  // adjusted offsets, for piecewise continuity
-    vector[S + 1] k_s;  // actual rate in each segment
+  real[] logistic_gamma(real k, real m, real[] delta, real[] t_change, int S) {
+    real gamma[S];  // adjusted offsets, for piecewise continuity
+    real k_s[S + 1];  // actual rate in each segment
     real m_pr;
     real m_pr;
 
 
     // Compute the rate in each segment
     // Compute the rate in each segment
@@ -43,45 +43,49 @@ functions {
     return gamma;
     return gamma;
   }
   }
   
   
-  vector logistic_trend(
+  real[] logistic_trend(
     real k,
     real k,
     real m,
     real m,
-    vector delta,
-    vector t,
-    vector cap,
-    matrix A,
-    vector t_change,
+    real[] delta,
+    real[] t,
+    real[] cap,
+    real[ , ] A,
+    real[] t_change,
     int S,
     int S,
     int T
     int T
   ) {
   ) {
-    vector[S] gamma;
-    vector[T] Y;
+    real gamma[S];
+    real Y[T];
 
 
     gamma = logistic_gamma(k, m, delta, t_change, S);
     gamma = logistic_gamma(k, m, delta, t_change, S);
     for (i in 1:T) {
     for (i in 1:T) {
-      Y[i] = cap[i] / (1 + exp(-(k + dot_product(A[i], delta)) * (t[i] - (m + dot_product(A[i], gamma)))));
+      Y[i] = cap[i] / (1 + exp(-(k + dot_product(A[i], delta))
+        * (t[i] - (m + dot_product(A[i], gamma)))));
     }
     }
     return Y;
     return Y;
   }
   }
 
 
   // Linear trend function
   // Linear trend function
 
 
-  vector linear_trend(
+  real[] linear_trend(
     real k,
     real k,
     real m,
     real m,
-    vector delta,
-    vector t,
-    matrix A,
-    vector t_change,
+    real[] delta,
+    real[] t,
+    real[ , ] A,
+    real[] t_change,
     int S,
     int S,
     int T
     int T
   ) {
   ) {
-    vector[S] gamma;
-    vector[T] Y;
-    
-    gamma = (-t_change .* delta);
+    real gamma[S];
+    real Y[T];
+
+    for (i in 1:S) {
+      gamma[i] = -t_change[i] * delta[i];
+    }
     for (i in 1:T) {
     for (i in 1:T) {
-      Y[i] = (k + dot_product(A[i], delta)) * t[i] + (m + dot_product(A[i], gamma));
+      Y[i] = (k + dot_product(A[i], delta)) * t[i] + (
+        m + dot_product(A[i], gamma));
     }
     }
     return Y;
     return Y;
   }
   }
@@ -90,37 +94,37 @@ functions {
 data {
 data {
   int T;                // Number of time periods
   int T;                // Number of time periods
   int<lower=1> K;       // Number of regressors
   int<lower=1> K;       // Number of regressors
-  vector[T] t;          // Time
-  vector[T] cap;        // Capacities for logistic trend
-  vector[T] y;          // Time series
+  real t[T];            // Time
+  real cap[T];          // Capacities for logistic trend
+  real y[T];            // Time series
   int S;                // Number of changepoints
   int S;                // Number of changepoints
-  vector[S] t_change;   // Times of trend changepoints
-  matrix[T,K] X;        // Regressors
+  real t_change[S];     // Times of trend changepoints
+  real X[T,K];         // Regressors
   vector[K] sigmas;     // Scale on seasonality prior
   vector[K] sigmas;     // Scale on seasonality prior
   real<lower=0> tau;    // Scale on changepoints prior
   real<lower=0> tau;    // Scale on changepoints prior
   int trend_indicator;  // 0 for linear, 1 for logistic
   int trend_indicator;  // 0 for linear, 1 for logistic
-  vector[K] s_a;        // Indicator of additive features
-  vector[K] s_m;        // Indicator of multiplicative features
+  real s_a[K];          // Indicator of additive features
+  real s_m[K];          // Indicator of multiplicative features
 }
 }
 
 
 transformed data {
 transformed data {
-  matrix[T, S] A;
+  real A[T, S];
   A = get_changepoint_matrix(t, t_change, T, S);
   A = get_changepoint_matrix(t, t_change, T, S);
 }
 }
 
 
 parameters {
 parameters {
   real k;                   // Base trend growth rate
   real k;                   // Base trend growth rate
   real m;                   // Trend offset
   real m;                   // Trend offset
-  vector[S] delta;          // Trend rate adjustments
+  real delta[S];            // Trend rate adjustments
   real<lower=0> sigma_obs;  // Observation noise
   real<lower=0> sigma_obs;  // Observation noise
-  vector[K] beta;           // Regressor coefficients
+  real beta[K];             // Regressor coefficients
 }
 }
 
 
 transformed parameters {
 transformed parameters {
-  vector[T] trend;
-  vector[T] Y;
-  vector[T] Xb_a;
-  vector[T] Xb_m;
+  real trend[T];
+  real Y[T];
+  real beta_m[K];
+  real beta_a[K];
 
 
   if (trend_indicator == 0) {
   if (trend_indicator == 0) {
     trend = linear_trend(k, m, delta, t, A, t_change, S, T);
     trend = linear_trend(k, m, delta, t, A, t_change, S, T);
@@ -128,10 +132,15 @@ transformed parameters {
     trend = logistic_trend(k, m, delta, t, cap, A, t_change, S, T);
     trend = logistic_trend(k, m, delta, t, cap, A, t_change, S, T);
   }
   }
 
 
+  for (i in 1:K) {
+    beta_m[i] = beta[i] * s_m[i];
+    beta_a[i] = beta[i] * s_a[i];
+  }
+
   for (i in 1:T) {
   for (i in 1:T) {
-    Xb_a[i] = dot_product(X[i], beta .* s_a);
-    Xb_m[i] = dot_product(X[i], beta .* s_m);
-    Y[i] = trend[i] * (1 + Xb_m[i]) + Xb_a[i];
+    Y[i] = (
+      trend[i] * (1 + dot_product(X[i], beta_m)) + dot_product(X[i], beta_a)
+    );
   }
   }
 }
 }