# Copyright 2022 - 2026 The PyMC Labs Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.
"""Utility functions."""
from __future__ import annotations
import re
from typing import TYPE_CHECKING, Any, Literal
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import xarray as xr
if TYPE_CHECKING:
from causalpy.experiments.synthetic_control import SyntheticControl
from causalpy.constants import HDI_PROB
def _as_scalar(value: Any) -> float:
"""Convert scalar-like values (including singleton arrays) to Python float.
Handles plain floats, 0-d and 1-element numpy arrays, and singleton
xarray DataArrays that arise when NumPy >= 2.4 enforces stricter
scalar-conversion rules.
Examples
--------
>>> _as_scalar(3.14)
3.14
>>> _as_scalar(np.array(2.5))
2.5
>>> _as_scalar(np.array([2.5]))
2.5
"""
return float(np.asarray(value).reshape(()))
def _bayesian_r2_score(y_true: np.ndarray, y_pred: np.ndarray) -> pd.Series:
"""Compute Bayesian R-squared across posterior predictive draws."""
var_y_est = np.var(y_pred, axis=1, ddof=0)
var_e = np.var(y_true - y_pred, axis=1, ddof=0)
r2_samples = var_y_est / (var_y_est + var_e)
return pd.Series(
[r2_samples.mean(), r2_samples.std(ddof=0)], index=["r2", "r2_std"]
)
def has_posterior_draws(Y: xr.DataArray) -> bool:
"""Whether *Y* carries genuine posterior uncertainty.
The canonical prediction container has ``chain`` and ``draw`` dimensions
on every backend; point-estimate backends emit singleton dimensions (a
point estimate is a posterior with one atom). Downstream code should key
statistical dispatch (HDI vs t-interval, ribbons, tail probabilities,
...) on this data property rather than on backend identity, so any
backend that emits many draws gets posterior summaries for free and a
degenerate single-draw run falls back to point summaries.
Parameters
----------
Y : xr.DataArray
A canonical prediction container with ``chain`` and ``draw``
dimensions.
"""
return Y.sizes.get("chain", 1) * Y.sizes.get("draw", 1) > 1
def _is_variable_dummy_coded(series: pd.Series) -> bool:
"""Check if a data in the provided Series is dummy coded. It should be 0 or 1
only."""
return len(set(series).difference({0, 1})) == 0
def _series_has_2_levels(series: pd.Series) -> bool:
"""Check that the variable in the provided Series has 2 levels"""
return len(pd.Categorical(series).categories) == 2
def round_num(n: float | xr.DataArray, round_to: int | None) -> str:
"""Return a string representing a number with significant figures.
Parameters
----------
n : float or xr.DataArray
Number to round. A zero-dimensional DataArray, as produced by
``.mean()`` on posterior samples, formats like the scalar it wraps.
round_to : int, optional
Number of significant figures. If None, defaults to 2.
Returns
-------
str
String representation of the number with specified significant
figures.
"""
sig_figs = _format_sig_figs(n, round_to)
return f"{n:.{sig_figs}g}"
def _format_sig_figs(value: float | xr.DataArray, default: int | None = None) -> int:
"""Get a default number of significant figures.
Gives the integer part or `default`, whichever is bigger.
Examples
--------
0.1234 --> 0.12
1.234 --> 1.2
12.34 --> 12
123.4 --> 123
"""
if default is None:
default = 2
if value == 0:
return 1
return max(int(np.log10(np.abs(value))) + 1, default)
def convert_to_string(x: float | xr.DataArray, round_to: int | None = 2) -> str:
"""Convert numeric inputs to a formatted string representation.
Parameters
----------
x : float or xr.DataArray
The numeric value or xarray DataArray to convert.
round_to : int, optional
Number of significant figures to round to. Defaults to 2.
Returns
-------
str
Formatted string representation. For floats, returns rounded
decimal. For DataArrays, returns mean with credible interval.
Raises
------
ValueError
If `x` is neither a float nor an xarray DataArray.
"""
if isinstance(x, float):
# In the case of a float, we return the number rounded to 2 decimal places
return f"{x:.2f}"
elif isinstance(x, xr.DataArray):
# In the case of an xarray object, we return the mean and CI
percentiles = x.quantile(
[(1 - HDI_PROB) / 2, 1 - (1 - HDI_PROB) / 2]
).to_numpy()
ci = (
rf"$CI_{{{HDI_PROB * 100:.0f}\%}}$"
+ f"[{round_num(percentiles[0], round_to)}, {round_num(percentiles[1], round_to)}]"
)
return f"{x.mean().to_numpy():.2f}" + ci
else:
raise ValueError(
"Type not supported. Please provide a float or an xarray object."
)
def get_interaction_terms(formula: str) -> list[str]:
"""
Extract interaction terms from a statistical model formula.
Parameters
----------
formula : str
A statistical model formula string (e.g., "y ~ x1 + x2*x3")
Returns
-------
list[str]
A list of interaction terms (those containing '*' or ':')
Examples
--------
>>> get_interaction_terms("y ~ 1 + x1 + x2*x3")
['x2*x3']
>>> get_interaction_terms("y ~ x1:x2 + x3")
['x1:x2']
>>> get_interaction_terms("y ~ x1 + x2 + x3")
[]
"""
# Define interaction indicators
INTERACTION_INDICATORS = ["*", ":"]
# Remove whitespace
formula_clean = formula.replace(" ", "")
# Extract right-hand side of the formula
rhs = formula_clean.split("~")[1]
# Split terms by '+' or '-' while keeping them intact
terms = re.split(r"(?=[+-])", rhs)
# Clean up terms and get interaction terms (those with '*' or ':')
interaction_terms = []
for term in terms:
# Remove leading + or - for processing
clean_term = term.lstrip("+-")
if any(indicator in clean_term for indicator in INTERACTION_INDICATORS):
interaction_terms.append(clean_term)
return interaction_terms
def check_convex_hull_violation(
treated_series: np.ndarray | xr.DataArray,
control_matrix: np.ndarray | xr.DataArray,
) -> dict:
"""
Check if treated series values fall within the range of control series.
For each time point, verify that:
min(controls) <= treated <= max(controls)
This is a necessary (but not sufficient) condition for the treated unit
to lie within the convex hull of control units.
Both arguments accept either ``np.ndarray`` or ``xr.DataArray`` inputs;
only positional (axis-based) operations are used internally.
Parameters
----------
treated_series : np.ndarray or xr.DataArray
1D array of treated unit values (shape: n_timepoints)
control_matrix : np.ndarray or xr.DataArray
2D array of control unit values (shape: n_timepoints x n_controls)
Returns
-------
dict
Dictionary with keys:
- 'passes': bool - whether the check passes
- 'n_violations': int - number of time points with violations
- 'pct_above': float - percentage of points where treated > max(controls)
- 'pct_below': float - percentage of points where treated < min(controls)
Examples
--------
>>> treated = np.array([1.0, 2.0, 3.0])
>>> controls = np.array([[0.5, 1.5], [1.5, 2.5], [2.5, 3.5]])
>>> result = check_convex_hull_violation(treated, controls)
>>> result["passes"]
True
"""
control_min = control_matrix.min(axis=1)
control_max = control_matrix.max(axis=1)
above = treated_series > control_max
below = treated_series < control_min
n_points = len(treated_series)
if n_points == 0:
return {
"passes": True,
"n_violations": 0,
"pct_above": 0.0,
"pct_below": 0.0,
}
return {
"passes": not (above.any() or below.any()),
"n_violations": int(above.sum() + below.sum()),
"pct_above": float(100 * above.sum() / n_points),
"pct_below": float(100 * below.sum() / n_points),
}
[docs]
def plot_correlations(
data: pd.DataFrame,
columns: list[str] | None = None,
method: Literal["pearson", "kendall", "spearman"] = "pearson",
figsize: tuple[float, float] | None = None,
ax: plt.Axes | None = None,
**kwargs: Any,
) -> tuple[pd.DataFrame, plt.Axes]:
"""Plot a pairwise correlation heatmap for panel data columns.
Computes the pairwise correlation matrix between the specified columns
(typically geographic units or time series) and displays it as a
lower-triangle heatmap. This is a pre-experiment diagnostic for
synthetic control analyses: markets that are highly correlated in the
pre-treatment period are more likely to produce reliable counterfactuals.
Parameters
----------
data : pd.DataFrame
Wide-format panel data with time as the index and locations/units as columns.
columns : list[str], optional
Subset of columns to include. If ``None``, all numeric columns are used.
method : {"pearson", "kendall", "spearman"}, default "pearson"
Correlation method passed to :meth:`pandas.DataFrame.corr`.
figsize : tuple[float, float], optional
Width and height in inches for the figure. Only used when ``ax`` is not provided. If ``None``, matplotlib's default is used.
ax : matplotlib.axes.Axes, optional
Axes on which to draw the heatmap. If ``None``, a new figure and axes are created (sized according to ``figsize``).
Returns
-------
tuple[pd.DataFrame, matplotlib.axes.Axes]
The correlation matrix and the axes containing the heatmap.
Other Parameters
----------------
**kwargs
Keyword arguments forwarded to :func:`seaborn.heatmap`: ``vmin``, ``vmax``, ``cmap``, ``center``, ``robust``, ``annot``, ``fmt``, ``annot_kws``, ``linewidths``, ``linecolor``, ``cbar``, ``cbar_kws``, ``cbar_ax``, ``square``, ``xticklabels``, ``yticklabels``, ``mask``, and the :meth:`matplotlib.axes.Axes.pcolormesh` keywords supported by the installed seaborn version. ``data`` and ``ax`` are supplied by CausalPy. This narrow third-party forwarder lets callers override CausalPy's heatmap defaults without duplicating seaborn's evolving forwarding surface; unknown keys are rejected by seaborn or matplotlib rather than ignored.
Examples
--------
.. code-block:: python
import causalpy as cp
df = cp.load_data("geolift1")
corr, ax = cp.plot_correlations(df)
# Larger figure with smaller annotation text
corr, ax = cp.plot_correlations(df, figsize=(10, 8), annot_kws={"size": 7})
"""
subset = data[columns] if columns is not None else data.select_dtypes("number")
corr = subset.corr(method=method)
mask = np.triu(np.ones_like(corr, dtype=bool))
if ax is None:
_, ax = plt.subplots(figsize=figsize)
defaults: dict[str, Any] = {
"mask": mask,
"cmap": sns.diverging_palette(230, 20, as_cmap=True),
"vmin": -1,
"vmax": 1,
"center": 0,
"square": True,
"linewidths": 0.5,
"cbar_kws": {"shrink": 0.8},
"annot": True,
"fmt": ".2f",
}
defaults.update(kwargs)
sns.heatmap(corr, ax=ax, **defaults)
return corr, ax
