Fitting a linear model on part of the data and an exponential growth model on the rest with pymc3

Question

Let's generate some test data:

import numpy as np

observed = np.hstack([np.arange(10)*0.1 + 3, np.exp(np.arange(15)*.2 + .2)])

It looks like this:

import matplotlib.pyplot as plt

plt.plot(observed, marker='o', linestyle="None");

I would like to fit a linear model (y = a + bx) to the first part of the data, and an exponential growth model (y = exp(a + bx)) to the second part of the data. BUT let's pretend that I don't know, a-priori, where the switchpoint in.

I've tried writing this model in pymc3:

x = np.arange(len(observed))

with pm.Model() as model:
    sigma_1 = pm.HalfCauchy("sigma_1", beta=10)
    alpha_1 = pm.Normal("α_1", 0, sigma=20)
    beta_1 = pm.Normal("β_1", 0, sigma=20)
    sigma_0 = pm.HalfCauchy("sigma_0", beta=10)
    alpha_0 = pm.Normal("α_0", 0, sigma=20)
    beta_0 = pm.Normal("β_0", 0, sigma=20)

    switchpoint = pm.DiscreteUniform("switchpoint", lower=0, upper=len(x) - 1)

    exponential_growth = pm.Normal(
        "exponential_growth",
        mu=np.exp(alpha_1 + beta_1 * x),
        sigma=sigma_1,
        observed=observed,
    )

    linear_growth = pm.Normal(
        "linear_growth", mu=alpha_0 + beta_0 * x, sigma=sigma_0, observed=observed
    )

    likelihood = pm.math.switch(switchpoint >= x, linear_growth, exponential_growth)

    trace = sample(2000, cores=2)

though of course, this just fits two models to the whole data. I am not combining them in the correct way.

What is the correct way to specify that I would like to use linear_model before switchpoint, and exponential_model after it?

Answer 1

The switch can be used to compute the appropriate mu and sigma for the observation model with something like:

switchpoint = pm.DiscreteUniform("switchpoint", lower=0, upper=len(x) - 1)

mu, sigma = pm.math.switch(switchpoint >= x, 
                           (alpha_0 + beta_0 * x, sigma_0),
                           (np.exp(alpha_1 + beta_1 *x), sigma_1))


lik = pm.Normal('obs', mu=mu, sigma=sigma, observed=observed)