causal impact

import pandas as pd
import tensorflow_probability as tfp

from causalimpact import CausalImpact
from pandas import Timestamp
from pandas.tseries.offsets import DateOffset


def generate_data():
    f = DateOffset(months=1)
    d = {
            Timestamp("2012-01-01 00:00:00", freq=f): 20660.0,
            Timestamp("2012-02-01 00:00:00", freq=f): 21652.0,
            Timestamp("2012-03-01 00:00:00", freq=f): 28327.0,
            Timestamp("2012-04-01 00:00:00", freq=f): 20055.0,
            Timestamp("2012-05-01 00:00:00", freq=f): 8987.0,
            Timestamp("2012-06-01 00:00:00", freq=f): 8706.0,
            Timestamp("2012-07-01 00:00:00", freq=f): 8888.0,
            Timestamp("2012-08-01 00:00:00", freq=f): 8312.0,
            Timestamp("2012-09-01 00:00:00", freq=f): 6249.0,
            Timestamp("2012-10-01 00:00:00", freq=f): 7475.0,
            Timestamp("2012-11-01 00:00:00", freq=f): 13007.0,
            Timestamp("2012-12-01 00:00:00", freq=f): 18157.0,
            Timestamp("2013-01-01 00:00:00", freq=f): 21455.0,
            Timestamp("2013-02-01 00:00:00", freq=f): 23113.0,
            Timestamp("2013-03-01 00:00:00", freq=f): 29942.0,
            Timestamp("2013-04-01 00:00:00", freq=f): 20183.0,
            Timestamp("2013-05-01 00:00:00", freq=f): 9433.0,
            Timestamp("2013-06-01 00:00:00", freq=f): 8872.0,
            Timestamp("2013-07-01 00:00:00", freq=f): 9352.0,
            Timestamp("2013-08-01 00:00:00", freq=f): 7806.0,
            Timestamp("2013-09-01 00:00:00", freq=f): 6699.0,
            Timestamp("2013-10-01 00:00:00", freq=f): 7404.0,
            Timestamp("2013-11-01 00:00:00", freq=f): 12524.0,
            Timestamp("2013-12-01 00:00:00", freq=f): 17350.0,
            Timestamp("2014-01-01 00:00:00", freq=f): 21417.0,
            Timestamp("2014-02-01 00:00:00", freq=f): 23143.0,
            Timestamp("2014-03-01 00:00:00", freq=f): 31041.0,
            Timestamp("2014-04-01 00:00:00", freq=f): 22819.0,
            Timestamp("2014-05-01 00:00:00", freq=f): 9254.0,
            Timestamp("2014-06-01 00:00:00", freq=f): 9087.0,
            Timestamp("2014-07-01 00:00:00", freq=f): 9135.0,
            Timestamp("2014-08-01 00:00:00", freq=f): 8846.0,
            Timestamp("2014-09-01 00:00:00", freq=f): 5620.0,
            Timestamp("2014-10-01 00:00:00", freq=f): 9693.0,
            Timestamp("2014-11-01 00:00:00", freq=f): 13065.0,
            Timestamp("2014-12-01 00:00:00", freq=f): 19775.0,
            Timestamp("2015-01-01 00:00:00", freq=f): 24978.0,
            Timestamp("2015-02-01 00:00:00", freq=f): 26439.0,
            Timestamp("2015-03-01 00:00:00", freq=f): 34992.0,
            Timestamp("2015-04-01 00:00:00", freq=f): 21477.0,
            Timestamp("2015-05-01 00:00:00", freq=f): 10316.0,
            Timestamp("2015-06-01 00:00:00", freq=f): 9352.0,
            Timestamp("2015-07-01 00:00:00", freq=f): 9601.0,
            Timestamp("2015-08-01 00:00:00", freq=f): 9407.0,
            Timestamp("2015-09-01 00:00:00", freq=f): 6994.0,
            Timestamp("2015-10-01 00:00:00", freq=f): 10123.0,
            Timestamp("2015-11-01 00:00:00", freq=f): 13429.0,
            Timestamp("2015-12-01 00:00:00", freq=f): 20732.0,
            Timestamp("2016-01-01 00:00:00", freq=f): 25286.0,
            Timestamp("2016-02-01 00:00:00", freq=f): 27630.0,
            Timestamp("2016-03-01 00:00:00", freq=f): 42085.0,
            Timestamp("2016-04-01 00:00:00", freq=f): 32153.0,
            Timestamp("2016-05-01 00:00:00", freq=f): 11538.0,
            Timestamp("2016-06-01 00:00:00", freq=f): 9637.0,
            Timestamp("2016-07-01 00:00:00", freq=f): 9553.0,
            Timestamp("2016-08-01 00:00:00", freq=f): 9682.0,
            Timestamp("2016-09-01 00:00:00", freq=f): 8135.0,
            Timestamp("2016-10-01 00:00:00", freq=f): 13733.0,
            Timestamp("2016-11-01 00:00:00", freq=f): 15500.0,
        }
    return pd.DataFrame.from_dict(d, orient="index")


def main():
    data = generate_data()
    obs_series = data.iloc[:, 0].astype("float32")
    regular_data = tfp.sts.regularize_series(series=obs_series)
    local_linear = tfp.sts.LocalLinearTrend(observed_time_series=regular_data)
    seasonal = tfp.sts.Seasonal(num_seasons=12, observed_time_series=regular_data)
    model = tfp.sts.Sum([local_linear, seasonal], observed_time_series=regular_data)

    pre_period = ["2012-01-01", "2014-12-01"]
    post_period = ["2015-01-01", "2016-11-01"]
    ci = CausalImpact(regular_data, pre_period, post_period) # works
    ci = CausalImpact(regular_data, pre_period, post_period, model=model) # doesn't work

    print(ci.summary())


if __name__ == "__main__":
    main()