# 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 hashlib
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 _design_fingerprint(*inputs: Any) -> tuple:
"""Structural hash of build-time inputs (shapes, dtypes, raw bytes).
Inputs may be mappings of named arrays or array-likes. A second
``build()`` whose fingerprint differs from the recorded one means the
caller is trying to reuse an immutable graph with different data, which
must fail loudly instead of being silently ignored.
"""
def _digest(value: Any) -> Any:
if isinstance(value, dict):
return tuple(sorted((key, _digest(item)) for key, item in value.items()))
arr = np.ascontiguousarray(np.asarray(value))
return (
arr.shape,
str(arr.dtype),
hashlib.blake2b(arr.tobytes(), digest_size=16).hexdigest(),
)
return tuple(_digest(value) for value in inputs)
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, round_to: int | None) -> str:
"""Return a string representing a number with significant figures.
Parameters
----------
n : float
Number to round.
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, 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
