Refinement of KTR

orbit/constants/ktr.py

@@ -10,7 +10,13 @@ class DataInputMapper(Enum):
     # observation related
     _NUM_OF_OBSERVATIONS = 'N_OBS'
     _RESPONSE = 'RESPONSE'
+    # used for pyro
+    # _IS_VALID_RESPONSE = 'IS_VALID_RES'
+    _NUM_OF_VALID_RESPONSE = 'N_VALID_RES'
+    # mainly used for stan
+    _WHICH_VALID_RESPONSE = 'WHICH_VALID_RES'
     _RESPONSE_SD = 'SDY'
+    _RESPONSE_MEAN = 'MEAN_Y'
     _DEGREE_OF_FREEDOM = 'DOF'
     # ----------  Level  ---------- #
     _NUM_KNOTS_LEVEL = 'N_KNOTS_LEV'
@@ -35,6 +41,8 @@ class DataInputMapper(Enum):
     _INSERT_PRIOR_SD = 'PRIOR_SD'
     _INSERT_PRIOR_TP_IDX = 'PRIOR_TP_IDX'
     _INSERT_PRIOR_IDX = 'PRIOR_IDX'
+    _LEVEL_KNOTS = 'LEV_KNOT_LOC'
+    _SEAS_TERM = 'SEAS_TERM'
 
 
 class BaseSamplingParameters(Enum):

orbit/diagnostics/plot.py

@@ -112,13 +112,14 @@ def plot_predicted_data(training_actual_df, predicted_df, date_col, actual_col,
         fig.savefig(path)
     if is_visible:
         plt.show()
+    else:
+        plt.close()
 
     return ax
 
 
 def plot_predicted_components(predicted_df, date_col, prediction_percentiles=None, plot_components=None,
-                              title="", figsize=None, path=None):
-
+                              title="", is_visible=True, figsize=None, path=None):
     """ Plot predicted componenets with the data frame of decomposed prediction where components
     has been pre-defined as `trend`, `seasonality` and `regression`.
     Parameters
@@ -138,6 +139,8 @@ def plot_predicted_components(predicted_df, date_col, prediction_percentiles=Non
         title of the plot
     figsize: tuple
         figsize pass through to `matplotlib.pyplot.figure()`
+    is_visible: boolean
+        whether we want to show the plot. If called from unittest, is_visible might = False.
     path: str
         path to save the figure
    Returns
@@ -183,6 +186,11 @@ def plot_predicted_components(predicted_df, date_col, prediction_percentiles=Non
     if path:
         plt.savefig(path)
 
+    if is_visible:
+        plt.show()
+    else:
+        plt.close()
+
     return axes
 
 

orbit/pyro/ktr.py

@@ -4,6 +4,7 @@
 import pyro
 import pyro.distributions as dist
 
+# FIXME: this is sort of dangerous; consider better implementation later
 torch.set_default_tensor_type('torch.DoubleTensor')
 
 
@@ -14,7 +15,13 @@ def __init__(self, data):
         for key, value in data.items():
             key = key.lower()
             if isinstance(value, (list, np.ndarray)):
-                value = torch.tensor(value, dtype=torch.double)
+                if key in ['which_valid_res']:
+                    # to use as index, tensor type has to be long or int
+                    value = torch.tensor(value)
+                else:
+                    # loc/scale cannot be in long format
+                    # sometimes they may be supplied as int, so dtype conversion is needed
+                    value = torch.tensor(value, dtype=torch.double)
             self.__dict__[key] = value
 
     def __call__(self):
@@ -30,9 +37,15 @@ def __call__(self):
         use _coef, _weight etc. instead of _beta, use _scale instead of _sigma
         """
         response = self.response
+        which_valid = self.which_valid_res
+
         n_obs = self.n_obs
+        n_valid = self.n_valid_res
         sdy = self.sdy
+        meany = self.mean_y
         dof = self.dof
+        lev_knot_loc = self.lev_knot_loc
+        seas_term = self.seas_term
 
         pr = self.pr
         rr = self.rr
@@ -43,13 +56,6 @@ def __call__(self):
         k_coef = self.k_coef
         n_knots_lev = self.n_knots_lev
         n_knots_coef = self.n_knots_coef
-        regressors = torch.zeros(n_obs)
-        if n_pr > 0 and n_rr > 0:
-            regressors = torch.cat([rr, pr], dim=-1)
-        elif n_pr > 0:
-            regressors = pr
-        elif n_rr > 0:
-            regressors = rr
 
         lev_knot_scale = self.lev_knot_scale
 
@@ -63,11 +69,33 @@ def __call__(self):
         pr_knot_pool_scale = self.pr_knot_pool_scale
         pr_knot_scale = self.pr_knot_scale.unsqueeze(-1)
 
+        # transformation of data
+        regressors = torch.zeros(n_obs)
+        if n_pr > 0 and n_rr > 0:
+            regressors = torch.cat([rr, pr], dim=-1)
+        elif n_pr > 0:
+            regressors = pr
+        elif n_rr > 0:
+            regressors = rr
+
+        response_tran = response - meany - seas_term
+
+        # sampling begins here
         extra_out = {}
 
         # levels sampling
-        lev_knot = pyro.sample("lev_knot", dist.Laplace(0, lev_knot_scale).expand([n_knots_lev]))
-        lev = (lev_knot @ k_lev.transpose(-2, -1))
+        # with pyro.plate("lev_plate", n_knots_lev):
+        #     lev_drift = pyro.sample("lev_drift", dist.Laplace(0, lev_knot_scale))
+        # lev_knot_tran = lev_drift.cumsum(-1)
+        # lev = (lev_knot_tran @ k_lev.transpose(-2, -1))
+
+        # levels sampling
+        if len(lev_knot_loc) > 0:
+            lev_knot_tran = pyro.sample("lev_knot", dist.Normal(lev_knot_loc - meany, lev_knot_scale).expand([n_knots_lev]))
+            lev = (lev_knot_tran @ k_lev.transpose(-2, -1))
+        else:
+            lev_knot_tran = pyro.sample("lev_knot", dist.Laplace(0, lev_knot_scale).expand([n_knots_lev]))
+            lev = (lev_knot_tran @ k_lev.transpose(-2, -1))
 
         # regular regressor sampling
         if n_rr > 0:
@@ -129,11 +157,19 @@ def __call__(self):
                 pyro.sample("prior_{}_{}".format(tp, idx), dist.Normal(m, sd),
                             obs=coef[..., tp, idx])
 
-        pyro.sample("init_lev", dist.Normal(response[0], sdy), obs=lev[..., 0])
-
         obs_scale = pyro.sample("obs_scale", dist.HalfCauchy(sdy))
-        with pyro.plate("response_plate", n_obs):
-            pyro.sample("response", dist.StudentT(dof, yhat[..., :], obs_scale), obs=response)
-
-        extra_out.update({'yhat': yhat, 'lev': lev, 'coef': coef, 'coef_knot': coef_knot})
+        with pyro.plate("response_plate", n_valid):
+            pyro.sample("response",
+                        dist.StudentT(dof, yhat[..., which_valid], obs_scale),
+                        obs=response_tran[which_valid])
+
+        lev_knot = lev_knot_tran + meany
+
+        extra_out.update({
+            'yhat': yhat + seas_term + meany,
+            'lev': lev + meany,
+            'lev_knot': lev_knot,
+            'coef': coef,
+            'coef_knot': coef_knot
+        })
         return extra_out

orbit/stan/ktr.stan

@@ -0,0 +1,104 @@
+data {
+  // response related
+  int<lower=0> N_OBS;
+  int<lower=0> N_VALID_RES;
+  real<lower=0> SDY;
+  real MEAN_Y;
+  int<lower=0> DOF;
+  vector[N_OBS] RESPONSE; 
+  int WHICH_VALID_RES[N_VALID_RES];
+  // trend related
+  int<lower=0> N_KNOTS_LEV;
+  matrix[N_OBS, N_KNOTS_LEV] K_LEV;
+  real<lower=0> LEV_KNOT_SCALE;
+  // regrression related
+  int<lower=0> N_RR;
+  int<lower=0> N_PR;
+  vector[N_RR] RR_KNOT_POOL_LOC;
+  vector<lower=0>[N_RR] RR_KNOT_POOL_SCALE;
+  vector<lower=0>[N_RR] RR_KNOT_SCALE;
+  matrix[N_OBS, N_RR] RR;
+  matrix[N_OBS, N_PR] PR;
+  vector[N_PR] PR_KNOT_POOL_LOC;
+  vector<lower=0>[N_PR] PR_KNOT_POOL_SCALE;
+  vector<lower=0>[N_PR] PR_KNOT_SCALE;
+  // kernel
+  int<lower=0> N_KNOTS_COEF;
+  matrix[N_OBS, N_KNOTS_COEF] K_COEF;
+}
+transformed data {
+  matrix[N_OBS, N_RR + N_PR] REGRESSORS;
+  vector[N_OBS] RESPONSE_TRAN;
+  // convert numpy index to stan
+  int WHICH_VALID_RES2[N_VALID_RES];
+  for (n in 1:N_VALID_RES) {
+    WHICH_VALID_RES2[n] = WHICH_VALID_RES[n] + 1;
+  }
+  REGRESSORS = append_col(RR, PR);
+  RESPONSE_TRAN = RESPONSE - MEAN_Y;
+}
+
+parameters {
+  // vector[N_KNOTS_LEV] lev_knot;
+  vector[N_KNOTS_LEV] lev_knot_drift;
+  vector<lower=0>[N_PR] pr_knot_loc;
+  vector[N_RR] rr_knot_loc;
+  matrix<lower=0>[N_KNOTS_COEF, N_PR] pr_knot;
+  matrix[N_KNOTS_COEF, N_RR] rr_knot;
+  real<lower=0, upper=SDY> obs_scale;
+}
+transformed parameters {
+  vector[N_OBS] lev;
+  vector[N_OBS] regression;
+  vector[N_OBS] yhat;  
+  vector[N_KNOTS_LEV] lev_knot_tran;
+  matrix[N_OBS, N_PR + N_RR] coef;
+  matrix[N_OBS, N_RR] rr_coef;
+  matrix<lower=0>[N_OBS, N_PR] pr_coef;
+
+  lev_knot_tran = cumulative_sum(lev_knot_drift);
+  lev = K_LEV * lev_knot_tran;
+
+  rr_coef = rep_matrix(0, N_OBS, N_RR);
+  pr_coef = rep_matrix(0, N_OBS, N_PR);
+
+  if (N_RR > 0) rr_coef = K_COEF * rr_knot;
+  if (N_PR > 0) pr_coef = K_COEF * pr_knot;
+
+  coef = append_col(rr_coef, pr_coef);
+  if (N_RR + N_PR  > 0) {
+    for (n in 1:N_OBS) {
+      regression[n] = sum(REGRESSORS[n, :] .* coef[n, :]);
+    }
+  } else {
+    regression = rep_vector(0, N_OBS);
+  }
+  yhat = lev + regression;
+}
+
+
+model {
+  // lev_knot ~ double_exponential(0, LEV_KNOT_SCALE);
+  lev_knot_drift ~ double_exponential(0, LEV_KNOT_SCALE);
+  rr_knot_loc ~ normal(RR_KNOT_POOL_LOC, RR_KNOT_POOL_SCALE);
+  for (n in 1:N_KNOTS_COEF) {
+    rr_knot[n,:] ~ normal(rr_knot_loc, RR_KNOT_SCALE);
+  }
+  pr_knot_loc ~ normal(PR_KNOT_POOL_LOC, PR_KNOT_POOL_SCALE);
+  for (n in 1:N_KNOTS_COEF) {
+    pr_knot[n,:] ~ normal(pr_knot_loc, PR_KNOT_SCALE);
+  }
+  obs_scale ~ cauchy(0, SDY)T[0, SDY];
+  RESPONSE_TRAN[WHICH_VALID_RES2] ~ student_t(DOF, yhat[WHICH_VALID_RES2], obs_scale);
+}
+
+generated quantities {
+  matrix[N_RR + N_PR, N_KNOTS_COEF] coef_knot;
+  vector[N_KNOTS_LEV] lev_knot;
+  lev_knot = lev_knot_tran + MEAN_Y;
+  if (N_RR + N_PR > 0) {
+    coef_knot = append_col(rr_knot, pr_knot)';
+  } else {
+    coef_knot = rep_matrix(0, N_RR + N_PR, N_KNOTS_COEF);
+  }
+}