Merge pull request #77 from uber/get-regression-coefficients

Public method to get regression coefficients

Author: steveyang90

orbit/constants/constants.py

@@ -1,3 +1,4 @@
+from collections import namedtuple
 from enum import Enum
 import os
 
@@ -129,4 +130,10 @@ class BacktestAnalyzeKeys(Enum):
 # Defaults Values
 DEFAULT_REGRESSOR_SIGN = '='
 DEFAULT_REGRESSOR_BETA = 0
-DEFAULT_REGRESSOR_SIGMA = 1.0
+DEFAULT_REGRESSOR_SIGMA = 1.0
+
+# beta coef columns
+COEFFICIENT_DF_COLS = namedtuple(
+    'coefficients_df_cols',
+    ['REGRESSOR', 'REGRESSOR_SIGN', 'COEFFICIENT']
+)('regressor', 'regressor_sign', 'coefficient')

orbit/lgt.py

@@ -10,7 +10,9 @@
 from orbit.constants.constants import (
     DEFAULT_REGRESSOR_SIGN,
     DEFAULT_REGRESSOR_BETA,
-    DEFAULT_REGRESSOR_SIGMA
+    DEFAULT_REGRESSOR_SIGMA,
+    COEFFICIENT_DF_COLS,
+    PredictMethod
 )
 from orbit.exceptions import (
     PredictionException,
@@ -362,6 +364,95 @@ def _set_model_param_names(self):
             self.model_param_names += [
                 lgt.RegressionStanSamplingParameters.REGULAR_REGRESSOR_BETA.value]
 
+    @staticmethod
+    def _concat_regression_coefs(pr_beta=None, rr_beta=None):
+        """Concatenates regression posterior matrix
+
+        In the case that `pr_beta` or `rr_beta` is a 1d tensor, transform to 2d tensor and
+        concatenate.
+
+        Args
+        ----
+        pr_beta : torch.tensor
+            postive-value constrainted regression betas
+        rr_beta : torch.tensor
+            regular regression betas
+
+        Returns
+        -------
+        torch.tensor
+            concatenated 2d tensor of shape (1, len(rr_beta) + len(pr_beta))
+
+        """
+        regressor_beta = None
+        if pr_beta is not None and rr_beta is not None:
+            pr_beta = pr_beta if len(pr_beta.shape) == 2 else pr_beta.reshape(1, -1)
+            rr_beta = rr_beta if len(rr_beta.shape) == 2 else rr_beta.reshape(1, -1)
+            regressor_beta = torch.cat((pr_beta, rr_beta), dim=1)
+        elif pr_beta is not None:
+            regressor_beta = pr_beta
+        elif rr_beta is not None:
+            regressor_beta = rr_beta
+
+        return regressor_beta
+
+    def get_regression_coefs(self, aggregation_method='mean'):
+        """Return DataFrame regression coefficients
+
+        Args
+        ----
+        aggregation_method : str
+            any PredictMethod except `full`
+        """
+        def _validate_args():
+            valid_args = set([x.value for x in PredictMethod])
+            valid_args = valid_args - set([PredictMethod.FULL_SAMPLING.value])
+
+            if aggregation_method not in valid_args:
+                raise IllegalArgument("aggregation_method must be one of {}".format(valid_args))
+
+        # init dataframe
+        reg_df = pd.DataFrame()
+
+        # end if no regressors
+        if self.num_of_regular_regressors + self.num_of_positive_regressors == 0:
+            return reg_df
+
+        _validate_args()
+
+        pr_beta = self.aggregated_posteriors\
+            .get(aggregation_method)\
+            .get(lgt.RegressionStanSamplingParameters.POSITIVE_REGRESSOR_BETA.value)
+
+        rr_beta = self.aggregated_posteriors\
+            .get(aggregation_method)\
+            .get(lgt.RegressionStanSamplingParameters.REGULAR_REGRESSOR_BETA.value)
+
+        # because `_conccat_regression_coefs` operates on torch tensors
+        pr_beta = torch.from_numpy(pr_beta) if pr_beta is not None else pr_beta
+        rr_beta = torch.from_numpy(rr_beta) if rr_beta is not None else rr_beta
+
+        regressor_betas = self._concat_regression_coefs(pr_beta, rr_beta)
+
+        # get column names
+        pr_cols = self.positive_regressor_col
+        rr_cols = self.regular_regressor_col
+
+        # note ordering here is not the same as `self.regressor_cols` because positive
+        # and negative do not have to be grouped on input
+        regressor_cols = pr_cols + rr_cols
+
+        # same note
+        regressor_signs \
+            = ["Positive"] * self.num_of_positive_regressors \
+            + ["Regular"] * self.num_of_regular_regressors
+
+        reg_df[COEFFICIENT_DF_COLS.REGRESSOR] = regressor_cols
+        reg_df[COEFFICIENT_DF_COLS.REGRESSOR_SIGN] = regressor_signs
+        reg_df[COEFFICIENT_DF_COLS.COEFFICIENT] = regressor_betas.flatten()
+
+        return reg_df
+
     def _predict(self, df=None, include_error=False, decompose=False):
         """Vectorized version of prediction math"""
 
@@ -407,16 +498,7 @@ def _predict(self, df=None, include_error=False, decompose=False):
         # regression components
         pr_beta = model.get(lgt.RegressionStanSamplingParameters.POSITIVE_REGRESSOR_BETA.value)
         rr_beta = model.get(lgt.RegressionStanSamplingParameters.REGULAR_REGRESSOR_BETA.value)
-        regressor_beta = None
-        if pr_beta is not None and rr_beta is not None:
-            pr_beta = pr_beta if len(pr_beta.shape) == 2 else pr_beta.reshape(1, -1)
-            rr_beta = rr_beta if len(rr_beta.shape) == 2 else rr_beta.reshape(1, -1)
-            regressor_beta = torch.cat((pr_beta, rr_beta), dim=1)
-        elif pr_beta is not None:
-            regressor_beta = pr_beta
-        elif rr_beta is not None:
-            regressor_beta = rr_beta
-
+        regressor_beta = self._concat_regression_coefs(pr_beta, rr_beta)
 
         ################################################################
         # Prediction Attributes

tests/test_dlt.py

@@ -4,6 +4,8 @@
 from orbit.dlt import DLT
 from orbit.exceptions import IllegalArgument
 
+from orbit.constants.constants import COEFFICIENT_DF_COLS
+
 
 def test_dlt_fit(iclaims_training_data):
     dlt = DLT(
@@ -269,6 +271,27 @@ def test_dlt_with_regressors_and_forecast(iclaims_training_data):
     assert list(predicted_df.columns) == expected_columns
 
 
+def test_get_regression_coefs(iclaims_training_data):
+    regressor_cols = ['trend.unemploy', 'trend.filling', 'trend.job']
+
+    dlt = DLT(
+        response_col='claims',
+        date_col='week',
+        seasonality=52,
+        chains=4,
+        prediction_percentiles=[5, 95, 30],
+        predict_method='full',
+        sample_method='mcmc',
+        regressor_col=regressor_cols
+    )
+
+    dlt.fit(df=iclaims_training_data)
+
+    reg_coefs = dlt.get_regression_coefs()
+
+    assert set(reg_coefs[COEFFICIENT_DF_COLS.REGRESSOR]) == set(regressor_cols)
+
+
 def test_dlt_multiple_fits(m3_monthly_data):
 
     dlt = DLT(response_col='value',

tests/test_lgt.py

@@ -5,6 +5,8 @@
 from orbit.lgt import LGT
 from orbit.exceptions import IllegalArgument
 
+from orbit.constants.constants import COEFFICIENT_DF_COLS
+
 
 def test_lgt_fit(iclaims_training_data):
     lgt = LGT(
@@ -270,6 +272,32 @@ def test_lgt_with_regressors_and_forecast(iclaims_training_data):
     assert list(predicted_df.columns) == expected_columns
 
 
+def test_get_regression_coefs(iclaims_training_data):
+    regressor_cols = ['trend.unemploy', 'trend.filling', 'trend.job']
+
+    lgt = LGT(
+        response_col='claims',
+        date_col='week',
+        seasonality=52,
+        chains=4,
+        prediction_percentiles=[5, 95, 30],
+        predict_method='full',
+        sample_method='mcmc',
+        regressor_col=regressor_cols,
+        regressor_sign=["=", "=", "+"]
+    )
+
+    lgt.fit(df=iclaims_training_data)
+
+    reg_coefs = lgt.get_regression_coefs()
+
+    assert set(reg_coefs[COEFFICIENT_DF_COLS.REGRESSOR]) == set(regressor_cols)
+
+    # negative case
+    with pytest.raises(IllegalArgument):
+        lgt.get_regression_coefs(aggregation_method='full')
+
+
 def test_lgt_multiple_fits(m3_monthly_data):
 
     lgt = LGT(response_col='value',

tests/test_pyro.py

@@ -32,6 +32,7 @@
 import pytest
 from orbit.lgt import LGT
 from orbit.exceptions import IllegalArgument
+from orbit.constants.constants import COEFFICIENT_DF_COLS
 
 
 def test_lgt_pyro_fit(iclaims_training_data):
@@ -113,3 +114,22 @@ def test_lgt_pyro_fit_and_full_predict(iclaims_training_data):
 
     assert predicted_out.shape == expected_shape
     assert list(predicted_out.columns) == expected_columns
+
+# todo: fix regression in pyro implementation
+# def test_get_regression_coefs(iclaims_training_data):
+#     regressor_cols = ['trend.unemploy', 'trend.filling', 'trend.job']
+#
+#     lgt = LGT(
+#         response_col='claims',
+#         date_col='week',
+#         seasonality=52,
+#         chains=4,
+#         predict_method='mean',
+#         inference_engine='pyro',
+#         regressor_col=regressor_cols,
+#         regressor_sign=["=", "=", "+"]
+#     )
+#
+#     lgt.fit(df=iclaims_training_data)
+#     reg_coefs = lgt.get_regression_coefs()
+#     assert set(reg_coefs[COEFFICIENT_DF_COLS.REGRESSOR]) == set(regressor_cols)