hide arviz style change in BaseExperiment.plot

Author: drbenvincent

causalpy/__init__.py

@@ -11,7 +11,6 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import arviz as az
 
 import causalpy.pymc_models as pymc_models
 import causalpy.skl_models as skl_models
@@ -28,8 +27,6 @@
 from .experiments.regression_kink import RegressionKink
 from .experiments.synthetic_control import SyntheticControl
 
-az.style.use("arviz-darkgrid")
-
 __all__ = [
     "__version__",
     "DifferenceInDifferences",

causalpy/experiments/base.py

@@ -63,12 +63,17 @@ def plot(self, *args, **kwargs) -> tuple:
         Internally, this function dispatches to either `_bayesian_plot` or `_ols_plot`
         depending on the model type.
         """
-        if isinstance(self.model, PyMCModel):
-            return self._bayesian_plot(*args, **kwargs)
-        elif isinstance(self.model, RegressorMixin):
-            return self._ols_plot(*args, **kwargs)
-        else:
-            raise ValueError("Unsupported model type")
+        import arviz as az
+        import matplotlib.pyplot as plt
+
+        # Apply arviz-darkgrid style only during plotting, then revert
+        with plt.style.context(az.style.library["arviz-darkgrid"]):
+            if isinstance(self.model, PyMCModel):
+                return self._bayesian_plot(*args, **kwargs)
+            elif isinstance(self.model, RegressorMixin):
+                return self._ols_plot(*args, **kwargs)
+            else:
+                raise ValueError("Unsupported model type")
 
     @abstractmethod
     def _bayesian_plot(self, *args, **kwargs):