Merge branch 'main' into sc-api-change

Author: drbenvincent

causalpy/pymc_models.py

@@ -87,10 +87,20 @@ def _data_setter(self, X) -> None:
 
         This method is used internally to register new data for the model for
         prediction.
+
+        NOTE: We are actively changing the `X`. Often, this matrix will have a different
+        number of rows than the original data. So to make the shapes work, we need to
+        update all data nodes in the model to have the correct shape. The values are not
+        used, so we set them to 0. In our case, we just have data nodes X and y, but if
+        in the future we get more complex models with more data nodes, then we'll need
+        to update all of them - ideally programmatically.
         """
+        new_no_of_observations = X.shape[0]
         with self:
-            # TODO: update coords
-            pm.set_data({"X": X})
+            pm.set_data(
+                {"X": X, "y": np.zeros(new_no_of_observations)},
+                coords={"obs_ind": np.arange(new_no_of_observations)},
+            )
 
     def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None:
         """Draw samples from posterior, prior predictive, and posterior predictive
@@ -118,7 +128,6 @@ def predict(self, X):
 
         .. caution::
             Results in KeyError if model hasn't been fit.
-
         """
 
         # Ensure random_seed is used in sample_prior_predictive() and