From 16472a57004a587054d94bdbd87e3c05ac6a2ee4 Mon Sep 17 00:00:00 2001
From: Cuong Duong <cuong@canva.com>
Date: Sun, 21 Mar 2021 04:59:58 +1100
Subject: [PATCH] Upgrade CmdStanPy interface (#1849)

* cmdstan variable extraction update

* add backwards compatibility

* Fix bug

* Upgrade stan version in requirements and in CI testing

Co-authored-by: Ben Letham <bletham@gmail.com>
---
 .github/workflows/build-and-test.yml | 6 +++---
 python/fbprophet/forecaster.py       | 1 +
 python/fbprophet/models.py           | 7 +++++--
 python/requirements.txt              | 2 +-
 4 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/build-and-test.yml b/.github/workflows/build-and-test.yml
index a457fdc..c9a838b 100644
--- a/.github/workflows/build-and-test.yml
+++ b/.github/workflows/build-and-test.yml
@@ -27,7 +27,7 @@ jobs:
         cd python && python setup.py develop test
         python setup.py clean
         rm -rf fbprophet/stan_model
-        wget https://github.com/stan-dev/cmdstan/releases/download/v2.22.1/cmdstan-2.22.1.tar.gz -O /tmp/cmdstan.tar.gz > /dev/null
+        wget https://github.com/stan-dev/cmdstan/releases/download/v2.26.1/cmdstan-2.26.1.tar.gz -O /tmp/cmdstan.tar.gz > /dev/null
         tar -xvf /tmp/cmdstan.tar.gz -C /tmp > /dev/null
-        make -C /tmp/cmdstan-2.22.1/ build > /dev/null
-        CMDSTAN=/tmp/cmdstan-2.22.1 STAN_BACKEND=CMDSTANPY python setup.py develop test
+        make -C /tmp/cmdstan-2.26.1/ build > /dev/null
+        CMDSTAN=/tmp/cmdstan-2.26.1 STAN_BACKEND=CMDSTANPY python setup.py develop test
diff --git a/python/fbprophet/forecaster.py b/python/fbprophet/forecaster.py
index 56d062a..8b28e95 100644
--- a/python/fbprophet/forecaster.py
+++ b/python/fbprophet/forecaster.py
@@ -1170,6 +1170,7 @@ class Prophet(object):
         else:
             self.params = self.stan_backend.fit(stan_init, dat, **kwargs)
 
+        self.stan_fit = self.stan_backend.stan_fit
         # If no changepoints were requested, replace delta with 0s
         if len(self.changepoints) == 0:
             # Fold delta into the base rate k
diff --git a/python/fbprophet/models.py b/python/fbprophet/models.py
index 8b62275..224bce3 100644
--- a/python/fbprophet/models.py
+++ b/python/fbprophet/models.py
@@ -128,7 +128,7 @@ class CmdStanPyBackend(IStanBackend):
                                           inits=stan_init,
                                           iter_sampling=iter_half,
                                           **kwargs)
-        res = self.stan_fit.sample
+        res = self.stan_fit.draws()
         (samples, c, columns) = res.shape
         res = res.reshape((samples * c, columns))
         params = self.stan_to_dict_numpy(self.stan_fit.column_names, res)
@@ -182,7 +182,10 @@ class CmdStanPyBackend(IStanBackend):
         end = 0
         two_dims = len(data.shape) > 1
         for cname in column_names:
-            parsed = cname.split(".")
+            if "." in cname:
+                parsed = cname.split(".")
+            else:
+                parsed = cname.split("[")
 
             curr = parsed[0]
             if prev is None:
diff --git a/python/requirements.txt b/python/requirements.txt
index 0d60db8..f85562c 100644
--- a/python/requirements.txt
+++ b/python/requirements.txt
@@ -1,5 +1,5 @@
 Cython>=0.22
-cmdstanpy==0.9.5
+cmdstanpy==0.9.68
 pystan>=2.14
 numpy>=1.15.4
 pandas>=1.0.4