Build Your Own Model

One important feature of orbit is to allow users to build and customize some prototype models promptly to serve their own purpose. Users just need to code up the core model structure part, then orbit will facilitate and streamline the downstream functionalities, such as fit-predict, diagnostics, etc.

In this section, we give a demo on how to build up a simple linear model using Pyro with orbit through multi-inheritance from abstract templates.

[1]:
import pandas as pd
import numpy as np
import torch
import pyro
import pyro.distributions as dist
from copy import deepcopy
import matplotlib.pyplot as plt

import orbit
from orbit.models.template import BaseTemplate
from orbit.models.template import FullBayesianTemplate
from orbit.estimators.pyro_estimator import PyroEstimatorVI

from orbit.utils.simulation import make_regression
from orbit.diagnostics.plot import plot_predicted_data
from orbit.utils.plot import get_orbit_style
plt.style.use(get_orbit_style())

%matplotlib inline
[2]:
print(orbit.__version__)
1.0.17

The ingredents to construct a new model: 1. fitter (MyFitter) - the place you define the new model structure including the parameters and how they are tied to priors and data evidence. 2. data mapper (MyDataMapper) - the way to define labels interfacing from python to sampling api 3. base model (BaseRegression) - additional process need to be added across the init-fit-predict steps and how we utilize posteriors to perform in-sample / oo-sample prediction

Define a fitter

[3]:
class MyFitter:
    max_plate_nesting = 1  # max number of plates nested in model

    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.float)
            self.__dict__[key] = value

    def __call__(self):
        extra_out = {}

        p = self.regressor.shape[1]
        bias = pyro.sample("bias", dist.Normal(0, 1))
        weight = pyro.sample("weight", dist.Normal(0, 1).expand([p]).to_event(1))
        yhat = bias + weight @ self.regressor.transpose(-1, -2)
        obs_sigma = pyro.sample("obs_sigma", dist.HalfCauchy(self.response_sd))

        with pyro.plate("response_plate", self.num_of_obs):
            pyro.sample("response", dist.Normal(yhat, obs_sigma), obs=self.response)
        return extra_out

Define data mapper

[4]:
from enum import Enum

class MyDataMapper(Enum):
    NUM_OF_OBSERVATIONS = 'NUM_OF_OBS'
    RESPONSE = 'RESPONSE'
    RESPONSE_SD = 'RESPONSE_SD'
    REGRESSOR = 'REGRESSOR'

Put it into the template

Note that here we use MyFitter and MyDataMapper to create the base model - BaseRegression.

[5]:
class BaseRegression(BaseTemplate):
    _fitter = MyFitter
    _data_input_mapper = MyDataMapper
    def __init__(self, regressor_col, **kwargs):
        super().__init__(**kwargs)  # create estimator in base class
        self.regressor_col = regressor_col
        self.regressor = None

    def _set_model_param_names(self):
        self._model_param_names = ['bias', 'weight', 'obs_sigma']

    def _set_dynamic_attributes(self, df):
        super()._validate_training_df(df)
        super()._set_training_df_meta(df)

        self.regressor = df[self.regressor_col].values

        super()._set_model_data_input()
        self._set_init_values()

    def _predict(self, posterior_estimates, df, include_error=False, **kwargs):
        model = deepcopy(posterior_estimates)
        new_regressor = df[self.regressor_col].values.T
        bias =  np.expand_dims(model.get('bias'),-1)
        obs_sigma = np.expand_dims(model.get('obs_sigma'), -1)
        weight = model.get('weight')

        pred_len = df.shape[0]
        batch_size = weight.shape[0]

        prediction = bias + np.matmul(weight, new_regressor) + \
            np.random.normal(0, obs_sigma, size=(batch_size, pred_len))
        return {'prediction': prediction}

Package a concrete model

Then we mixed BaseRegression with the FullBayesianTemplate from orbit. One additional step is to declare the right supported_estimator_types i.e. PyroEstimatorVI in this case. This is due to our default uses StanEstimatorMCMC and the estimation method we use here is svi in pyro. Please see our docs for types of estimators we support.

[6]:
class PyroVIRegression(FullBayesianTemplate, BaseRegression):

    _supported_estimator_types = [PyroEstimatorVI]

    def __init__(self, estimator_type=PyroEstimatorVI, **kwargs):
        super().__init__(estimator_type=estimator_type, **kwargs)

Test out the new model

Prepare the input data.

[7]:
x, y, coefs = make_regression(120, [3.0, -1.0], bias=1.0, scale=1.0)
[8]:
df = pd.DataFrame(
    np.concatenate([y.reshape(-1, 1), x], axis=1), columns=['y', 'x1', 'x2']
)
df['week'] = pd.date_range(start='2016-01-04', periods=len(y), freq='7D')
[9]:
df.head(5)
[9]:
y x1 x2 week
0 2.382337 0.345584 0.000000 2016-01-04
1 2.812929 0.330437 -0.000000 2016-01-11
2 3.600130 0.905356 0.446375 2016-01-18
3 -0.884275 -0.000000 0.581118 2016-01-25
4 2.704941 0.364572 0.294132 2016-02-01
[10]:
test_size = 20
train_df = df[:-test_size]
test_df = df[-test_size:]

Instantiate the new model object.

[11]:
mod = PyroVIRegression(
    response_col='y',
    date_col='week',
    regressor_col=['x1','x2'],
    verbose=True,
    num_steps=501,
    seed=2021,
)
[12]:
mod.fit(df=train_df)
INFO:root:Guessed max_plate_nesting = 2
step    0 loss = 27333, scale = 0.077552
step  100 loss = 12590, scale = 0.0092793
step  200 loss = 12597, scale = 0.0098217
step  300 loss = 12591, scale = 0.0095262
step  400 loss = 12593, scale = 0.0092962
step  500 loss = 12591, scale = 0.0095438
[13]:
estimated_weights = mod._posterior_samples['weight']

We can check the cofficients with truth.

[14]:
print("True Coef: {:.3f}, {:.3f}".format(coefs[0], coefs[1]) )
estimated_coef = np.median(estimated_weights, axis=0)
print("Estimated Coef: {:.3f}, {:.3f}".format(estimated_coef[0], estimated_coef[1]))
True Coef: 3.000, -1.000
Estimated Coef: 2.951, -0.964

We can also check the future forecast accuracy.

[15]:
predicted_df = mod.predict(df)
[16]:
_ = plot_predicted_data(train_df, predicted_df, 'week', 'y', test_actual_df=test_df, prediction_percentiles=[5, 95])
../_images/tutorials_build_your_own_model_28_0.png