From a1513305ff8a0946e0f75a02b470960863e900bf Mon Sep 17 00:00:00 2001 From: Juan Orduz Date: Thu, 11 Apr 2024 12:33:59 +0200 Subject: [PATCH] improve coords matching (#623) --- tests/clv/models/test_shifted_beta_geo.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/clv/models/test_shifted_beta_geo.py b/tests/clv/models/test_shifted_beta_geo.py index 4657b1084..0563069a6 100644 --- a/tests/clv/models/test_shifted_beta_geo.py +++ b/tests/clv/models/test_shifted_beta_geo.py @@ -183,7 +183,7 @@ def test_model_convergence(self, data, model_config): def test_distribution_customer_churn_time(self): dataset = pd.DataFrame( { - "customer_id": [1, 2, 3], + "customer_id": [0, 1, 2], "t_churn": [10, 10, 10], "T": 10, } @@ -195,11 +195,13 @@ def test_distribution_customer_churn_time(self): model.fit(fit_method="map") customer_thetas = np.array([0.1, 0.5, 0.9]) model.idata = az.from_dict( - { + posterior={ "alpha": np.ones((2, 500)), # Two chains, 500 draws each "beta": np.ones((2, 500)), "theta": np.full((2, 500, 3), customer_thetas), - } + }, + coords={"customer_id": [0, 1, 2]}, + dims={"theta": ["customer_id"]}, ) res = model.distribution_customer_churn_time(