TransferFunctionARRegression

class causalpy.pymc_models.TransferFunctionARRegression

Bayesian Transfer Function model with AR(1) errors for Graded Intervention Time Series.

This model extends the Transfer Function framework by explicitly modeling autocorrelation in the errors using an AR(1) process implemented via quasi-differencing. This approach properly accounts for temporal correlation in the residuals while jointly estimating transform parameters (adstock, saturation) and regression coefficients.

Mathematical Framework

The standard regression model with AR(1) errors is:

\[y[t] = \mu[t] + \epsilon[t] \epsilon[t] = \rho \cdot \epsilon[t-1] + \nu[t] \nu[t] \sim N(0, \sigma_\nu^2)\]

where \(\mu[t]\) is the regression mean (baseline + transformed treatment effects), \(\rho\) is the AR(1) coefficient (|ρ| < 1), and \(\nu[t]\) is white noise.

Quasi-Differencing Transformation

To enable Bayesian inference, we apply quasi-differencing:

\[y[t] - \rho \cdot y[t-1] = \mu[t] - \rho \cdot \mu[t-1] + \nu[t]\]

This transforms the model into one with independent errors \(\nu[t]\), which can be directly sampled in PyMC. The quasi-differenced likelihood is:

• For t=0: \(y[0] \sim N(\mu[0], \sigma_\nu / \sqrt{1-\rho^2})\) (stationary initial condition)

• For t>0: \(y[t] - \rho \cdot y[t-1] \sim N(\mu[t] - \rho \cdot \mu[t-1], \sigma_\nu)\)

Advantages Over Independent Errors

• Proper uncertainty quantification: Accounts for temporal correlation in credible intervals

• More efficient inference: Correctly models the error structure

• Better parameter recovery: Avoids bias from ignoring autocorrelation

• Diagnostic information: The \(\rho\) parameter indicates strength of temporal dependence

When to Use

Use this model when:

• Time series data exhibits autocorrelation in residuals

• Standard transfer function model shows diagnostic issues (e.g., correlated residuals)

• You need proper uncertainty propagation with temporal dependence

Use the standard TransferFunctionLinearRegression when:

• Residuals show minimal autocorrelation

• Computational efficiency is critical (AR model is slower)

• You want a simpler baseline model for comparison

type saturation_type:

Optional[str]

param saturation_type:

Type of saturation transform. Options: “hill”, “logistic”, “michaelis_menten”, None. If None, no saturation is applied.

type saturation_type:

str or None

type adstock_config:

Optional[Dict]

param adstock_config:

Configuration for adstock transform. Required keys: - “half_life_prior”: dict with prior specification (e.g., {“dist”: “Gamma”, “alpha”: 4, “beta”: 2}) - “l_max”: int, maximum lag - “normalize”: bool, whether to normalize weights If None, no adstock is applied.

type adstock_config:

dict or None

type saturation_config:

Optional[Dict]

param saturation_config:

Configuration for saturation transform. Structure depends on saturation_type: - For “hill”: {“slope_prior”: {…}, “kappa_prior”: {…}} - For “logistic”: {“lam_prior”: {…}} - For “michaelis_menten”: {“alpha_prior”: {…}, “lam_prior”: {…}}

type saturation_config:

dict or None

type coef_constraint:

str

param coef_constraint:

Constraint on treatment coefficients: “nonnegative” or “unconstrained”.

type coef_constraint:

str, default=”unconstrained”

type sample_kwargs:

Optional[Dict[str, Any]]

param sample_kwargs:

Additional kwargs passed to pm.sample().

type sample_kwargs:

dict, optional

Notes

• The AR(1) coefficient \(\rho\) has a Uniform(-0.99, 0.99) prior by default

• The quasi-differencing approach ensures the model remains computationally tractable

• Posterior predictive sampling requires forward simulation of the AR process

• Convergence can be slower than the independent errors model; consider increasing tune/draws

Prior Customization:

Priors are managed using the Prior class from pymc_extras and can be customized via the priors parameter:

from pymc_extras.prior import Prior

model = cp.pymc_models.TransferFunctionARRegression(
    saturation_type=None,
    adstock_config={...},
    priors={
        "beta": Prior(
            "Normal", mu=0, sigma=100, dims=["treated_units", "coeffs"]
        ),
        "rho": Prior(
            "Uniform", lower=-0.95, upper=0.95, dims=["treated_units"]
        ),
    },
)

By default, data-informed priors are set automatically via priors_from_data():

• Baseline coefficients (beta): Normal(0, 5 * std(y))

• Treatment coefficients (theta_treatment): Normal(0, 2 * std(y)) or HalfNormal(2 * std(y))

• Error std (sigma): HalfNormal(2 * std(y))

• AR(1) coefficient (rho): Uniform(-0.99, 0.99)

This adaptive approach ensures priors are reasonable regardless of data scale.

Examples

Basic usage:

import causalpy as cp

model = cp.pymc_models.TransferFunctionARRegression(
    saturation_type=None,
    adstock_config={
        "half_life_prior": {"dist": "Gamma", "alpha": 4, "beta": 2},
        "l_max": 8,
        "normalize": True,
    },
    sample_kwargs={"chains": 4, "draws": 2000, "tune": 1000},
)
result = cp.GradedInterventionTimeSeries(
    data=df,
    y_column="outcome",
    treatment_names=["treatment"],
    base_formula="1 + time + covariate",
    model=model,
)

References

[1]

Pankratz, A. (1991). “Forecasting with Dynamic Regression Models”. Wiley.

[2]

Box, G. E., & Jenkins, G. M. (2015). “Time Series Analysis: Forecasting and Control”. Wiley.

Methods

TransferFunctionARRegression.__init__([...])	Initialize TransferFunctionARRegression model.
TransferFunctionARRegression.add_coord(name)	Register a dimension coordinate with the model.
TransferFunctionARRegression.add_coords(...)	Vectorized version of Model.add_coord.
TransferFunctionARRegression.add_named_variable(var)	Add a random graph variable to the named variables of the model.
TransferFunctionARRegression.build_model(X, ...)	Build the PyMC model with transforms and AR(1) errors using quasi-differencing.
TransferFunctionARRegression.calculate_cumulative_impact(impact)
TransferFunctionARRegression.calculate_impact(...)	Calculate the causal impact as the difference between observed and predicted values.
TransferFunctionARRegression.check_start_vals(...)	Check that the logp is defined and finite at the starting point.
TransferFunctionARRegression.compile_d2logp([...])	Compiled log probability density hessian function.
TransferFunctionARRegression.compile_dlogp([...])	Compiled log probability density gradient function.
TransferFunctionARRegression.compile_fn(outs, *)	Compiles a PyTensor function.
TransferFunctionARRegression.compile_logp([...])	Compiled log probability density function.
TransferFunctionARRegression.copy()	Clone the model.
TransferFunctionARRegression.create_value_var(...)	Create a TensorVariable that will be used as the random variable's "value" in log-likelihood graphs.
TransferFunctionARRegression.d2logp([vars, ...])	Hessian of the models log-probability w.r.t.
TransferFunctionARRegression.debug([point, ...])	Debug model function at point.
TransferFunctionARRegression.dlogp([vars, ...])	Gradient of the models log-probability w.r.t.
TransferFunctionARRegression.eval_rv_shapes()	Evaluate shapes of untransformed AND transformed free variables.
TransferFunctionARRegression.fit(X, y, ...)	Fit the Transfer Function AR(1) model.
TransferFunctionARRegression.get_context([...])
TransferFunctionARRegression.initial_point([...])	Compute the initial point of the model.
TransferFunctionARRegression.logp([vars, ...])	Elemwise log-probability of the model.
TransferFunctionARRegression.logp_dlogp_function([...])	Compile a PyTensor function that computes logp and gradient.
TransferFunctionARRegression.make_obs_var(...)	Create a TensorVariable for an observed random variable.
TransferFunctionARRegression.name_for(name)	Check if name has prefix and adds if needed.
TransferFunctionARRegression.name_of(name)	Check if name has prefix and deletes if needed.
TransferFunctionARRegression.point_logps([...])	Compute the log probability of point for all random variables in the model.
TransferFunctionARRegression.predict(X)	Predict data given input data X
TransferFunctionARRegression.print_coefficients(labels)
TransferFunctionARRegression.priors_from_data(X, y)	Generate data-informed priors including AR(1) coefficient.
TransferFunctionARRegression.profile(outs, *)	Compile and profile a PyTensor function which returns outs and takes values of model vars as a dict as an argument.
TransferFunctionARRegression.register_data_var(data)	Register a data variable with the model.
TransferFunctionARRegression.register_rv(...)	Register an (un)observed random variable with the model.
TransferFunctionARRegression.replace_rvs_by_values(...)	Clone and replace random variables in graphs with their value variables.
TransferFunctionARRegression.score(X, y)	Score the Bayesian \(R^2\) given inputs X and outputs y.
TransferFunctionARRegression.set_data(name, ...)	Change the values of a data variable in the model.
TransferFunctionARRegression.set_dim(name, ...)	Update a mutable dimension.
TransferFunctionARRegression.set_initval(...)	Set an initial value (strategy) for a random variable.
TransferFunctionARRegression.shape_from_dims(dims)
TransferFunctionARRegression.to_graphviz(*)	Produce a graphviz Digraph from a PyMC model.

Attributes

basic_RVs	List of random variables the model is defined in terms of.
continuous_value_vars	All the continuous value variables in the model.
coords	Coordinate values for model dimensions.
datalogp	PyTensor scalar of log-probability of the observed variables and potential terms.
default_priors
dim_lengths	The symbolic lengths of dimensions in the model.
discrete_value_vars	All the discrete value variables in the model.
isroot
observedlogp	PyTensor scalar of log-probability of the observed variables.
parent
potentiallogp	PyTensor scalar of log-probability of the Potential terms.
prefix
root
unobserved_RVs	List of all random variables, including deterministic ones.
unobserved_value_vars	List of all random variables (including untransformed projections), as well as deterministics used as inputs and outputs of the model's log-likelihood graph.
value_vars	List of unobserved random variables used as inputs to the model's log-likelihood (which excludes deterministics).
varlogp	PyTensor scalar of log-probability of the unobserved random variables (excluding deterministic).
varlogp_nojac	PyTensor scalar of log-probability of the unobserved random variables (excluding deterministic) without jacobian term.

__init__(saturation_type=None, adstock_config=None, saturation_config=None, coef_constraint='unconstrained', sample_kwargs=None, priors=None)

Initialize TransferFunctionARRegression model.

Parameters:

• saturation_type (str | None)

• adstock_config (Dict | None)

• saturation_config (Dict | None)

• coef_constraint (str)

• sample_kwargs (Dict[str, Any] | None)

• priors (dict[str, Any] | None)

classmethod __new__(*args, **kwargs)