prediction

Prediction helpers for brms models.

This module wraps brms prediction utilities and returns typed result objects that contain both an ArviZ InferenceData view and the underlying R result.

Notes

Executed inside the worker process that hosts the embedded R session.

Attributes

FitResult = IDResult[IDBrm] module-attribute

ProxyListSexpVector = Union[SexpWrapper, ListSexpVector, None] module-attribute

Classes

IDLogLikelihoodInsample

Bases: IDConstantData

Typed .log_likelihood extension to idata

Source code in brmspy/types/brms_results.py

class IDLogLikelihoodInsample(IDConstantData):
    """Typed .log_likelihood extension to idata"""

    log_likelihood: xr.Dataset

Attributes

log_likelihood instance-attribute

IDLogLikelihoodOutsample

Bases: IDPredictionsConstantData

Typed .log_likelihood extension to idata

Source code in brmspy/types/brms_results.py

class IDLogLikelihoodOutsample(IDPredictionsConstantData):
    """Typed .log_likelihood extension to idata"""

    log_likelihood: xr.Dataset

Attributes

log_likelihood instance-attribute

IDObservedData

Bases: IDConstantData

Typed .posterior extension to idata

Source code in brmspy/types/brms_results.py

class IDObservedData(IDConstantData):
    """Typed .posterior extension to idata"""

    observed_data: xr.Dataset

Attributes

observed_data instance-attribute

IDPosteriorPredictive

Bases: IDConstantData

Typed .posterior_predictive extension to idata

Source code in brmspy/types/brms_results.py

class IDPosteriorPredictive(IDConstantData):
    """Typed .posterior_predictive extension to idata"""

    posterior_predictive: xr.Dataset

Attributes

posterior_predictive instance-attribute

IDPredictions

Bases: IDPredictionsConstantData

Typed .predictions extension to idata

Source code in brmspy/types/brms_results.py

class IDPredictions(IDPredictionsConstantData):
    """Typed .predictions extension to idata"""

    predictions: xr.Dataset

Attributes

predictions instance-attribute

IDResult dataclass

Bases: Generic[T_idata], RListVectorExtension

Generic result container with arviz and R objects.

Attributes:

Name	Type	Description
idata	InferenceData	arviz InferenceData object
r	ListVector	R object from brms

Source code in brmspy/types/brms_results.py

@dataclass
class IDResult(Generic[T_idata], RListVectorExtension):
    """Generic result container with arviz and R objects.

    Attributes
    ----------
    idata : arviz.InferenceData
        arviz InferenceData object
    r : robjects.ListVector
        R object from brms
    """

    idata: T_idata

    def __repr__(self) -> str:
        return repr(self.idata)

Attributes

idata instance-attribute

Functions

__repr__()

Source code in brmspy/types/brms_results.py

def __repr__(self) -> str:
    return repr(self.idata)

__init__(r, idata)

IDPosterior

Bases: IDConstantData

Typed .posterior extension to idata

Source code in brmspy/types/brms_results.py

class IDPosterior(IDConstantData):
    """Typed .posterior extension to idata"""

    posterior: xr.Dataset

Attributes

posterior instance-attribute

Functions

_arviz_add_constant_data(idata, constant_data_dict, group_name='constant_data', obs_id=None)

Add a non-draw group (constant_data or predictions_constant_data) to an idata.

Extracts obs_id coords directly from the existing idata. This avoids ArviZ's auto (chain, draw) dims and keeps the group purely 1D along obs_id.

Source code in brmspy/helpers/_rpy2/_conversion.py

def _arviz_add_constant_data(
    idata: az.InferenceData,
    constant_data_dict: dict[str, np.ndarray],
    group_name: Literal["constant_data", "predictions_constant_data"] = "constant_data",
    obs_id: None | list[str] | np.ndarray = None,
) -> az.InferenceData:
    """
    Add a non-draw group (constant_data or predictions_constant_data) to an idata.

    Extracts obs_id coords directly from the existing idata. This avoids ArviZ's
    auto (chain, draw) dims and keeps the group purely 1D along obs_id.
    """
    if not constant_data_dict:
        return idata

    # ---- 1) Extract obs_id coords from any existing group ----
    if obs_id is None:
        for group in idata.groups():
            ds = idata[group]
            if ds is not None and "obs_id" in ds.coords:
                obs_id = ds.coords["obs_id"].values
                break

        if obs_id is None:
            raise ValueError(
                "Could not locate 'obs_id' in any existing idata group; "
                "cannot attach constant_data."
            )

    # ---- 2) Build dims & coords for the new constant group ----
    const_dims = {name: ["obs_id"] for name in constant_data_dict.keys()}
    const_coords = {"obs_id": obs_id}

    # ---- 3) Build a small InferenceData and extend ----
    if group_name == "constant_data":
        const_idata = az.from_dict(
            constant_data=constant_data_dict, coords=const_coords, dims=const_dims
        )
    else:
        const_idata = az.from_dict(
            predictions_constant_data=constant_data_dict,
            coords=const_coords,
            dims=const_dims,
        )

    idata.extend(const_idata)
    return idata

_brmsfit_get_constant_data(brmsfit_obj, newdata=None, resp_names=None)

Extract constant_data for ArviZ.

• If newdata is None: use brmsfit$data.
• Else: use the provided newdata.
• Drop response columns and 'obs_id' (responses go to observed_data, obs_id is handled as a coord).
• Return a dict[var_name -> np.ndarray] with length N (N = number of rows).

Source code in brmspy/helpers/_rpy2/_conversion.py

def _brmsfit_get_constant_data(
    brmsfit_obj,
    newdata: None | pd.DataFrame = None,
    resp_names: None | list[str] = None,
) -> dict[str, np.ndarray]:
    """
    Extract constant_data for ArviZ.

    - If newdata is None: use brmsfit$data.
    - Else: use the provided newdata.
    - Drop response columns and 'obs_id' (responses go to observed_data,
      obs_id is handled as a coord).
    - Return a dict[var_name -> np.ndarray] with length N (N = number of rows).
    """
    import rpy2.robjects as ro
    from rpy2.robjects import pandas2ri
    from rpy2.robjects.conversion import localconverter

    if resp_names is None:
        resp_names = _brmsfit_get_response_names(brmsfit_obj)

    if newdata is None:
        # in-sample: use brmsfit$data (R data.frame) -> pandas.DataFrame
        r_data = brmsfit_obj.rx2("data")
        with localconverter(ro.default_converter + pandas2ri.converter):
            df = pandas2ri.rpy2py(r_data)
        if not isinstance(df, pd.DataFrame):
            df = pd.DataFrame(df)
    else:
        # out-of-sample: use newdata as given
        df = newdata.copy()

    # Ensure we don't accidentally mutate caller's frame
    df = df.copy()

    # Drop response variables if present
    drop_cols: set[str] = set(resp_names or [])

    # Drop obs_id column if present; obs_id is handled as a coord
    if "_obs_id_" in df.columns:
        df = df.set_index("_obs_id_", drop=True)

    keep_cols = [c for c in df.columns if c not in drop_cols]

    constant_data: dict[str, np.ndarray] = {}
    for col in keep_cols:
        # Just pass through whatever dtype it has; xarray can handle object too
        constant_data[col] = df[col].to_numpy()

    return constant_data

_brmsfit_get_dims_and_coords(brmsfit_obj, newdata=None, resp_names=None)

Infer dims/coords for ArviZ from a brmsfit object and optional newdata.

Rules for obs_id: - If newdata is None: 1) If obs_id column exists in fit$data and is unique: use that. 2) Else, if rownames of fit$data are unique: use those. 3) Else: use a sequential integer range [0, N). - If newdata is not None: 1) If obs_id column exists in newdata and is unique: use that. 2) Else: use newdata.index (with a warning if not unique).

Source code in brmspy/helpers/_rpy2/_conversion.py

def _brmsfit_get_dims_and_coords(
    brmsfit_obj,
    newdata: None | pd.DataFrame = None,
    resp_names: None | list[str] = None,
) -> tuple[TypeDims, TypeCoords]:
    """
    Infer dims/coords for ArviZ from a brmsfit object and optional newdata.

    Rules for obs_id:
    - If newdata is None:
        1) If `obs_id` column exists in `fit$data` and is unique: use that.
        2) Else, if rownames of `fit$data` are unique: use those.
        3) Else: use a sequential integer range [0, N).
    - If newdata is not None:
        1) If `obs_id` column exists in newdata and is unique: use that.
        2) Else: use newdata.index (with a warning if not unique).
    """
    import rpy2.robjects as ro

    fun_nrow = cast(Callable, ro.r("nrow"))

    if resp_names is None:
        resp_names = _brmsfit_get_response_names(brmsfit_obj)

    if newdata is None:
        # in-sample: look at brmsfit$data
        r_data = brmsfit_obj.rx2("data")
        n_obs = int(fun_nrow(r_data)[0])

        obs_id = _get_obs_id_from_r_data(r_data, n_obs)

    else:
        # out-of-sample: look at newdata
        n_obs = int(len(newdata))
        obs_id = _get_obs_id_from_newdata(newdata, n_obs)

    obs_id_arr = np.asarray(obs_id)

    dims: TypeDims = {}
    coords: TypeCoords = {}

    coords["obs_id"] = obs_id_arr

    for resp in resp_names:
        dims[resp] = ["obs_id"]

    # you can add more dims/coords for responses here later
    # e.g. multi-response mapping, etc.

    return dims, coords

_brmsfit_get_observed_data(brmsfit_obj, resp_names=None)

Source code in brmspy/helpers/_rpy2/_conversion.py

def _brmsfit_get_observed_data(
    brmsfit_obj, resp_names: None | list[str] = None
) -> dict[str, np.ndarray]:
    import rpy2.robjects as ro
    from rpy2.robjects import pandas2ri
    from rpy2.robjects.conversion import localconverter

    if resp_names is None:
        resp_names = _brmsfit_get_response_names(brmsfit_obj)

    observed_data_dict: dict[str, np.ndarray] = {}
    n_obs = 0

    try:
        r_get_y = cast(Callable, ro.r("brms::get_y"))
        y_r = r_get_y(brmsfit_obj)

        with localconverter(ro.default_converter + pandas2ri.converter):
            y_py = pandas2ri.rpy2py(y_r)

        if isinstance(y_py, pd.DataFrame):
            n_obs = y_py.shape[0]
            if not resp_names:
                resp_names = list(y_py.columns)
            for resp in resp_names:
                if resp in y_py.columns:
                    observed_data_dict[resp] = y_py[resp].to_numpy()

        elif isinstance(y_py, pd.Series):
            n_obs = y_py.shape[0]
            if not resp_names:
                resp_names = [str(y_py.name) or "y"]
            observed_data_dict[resp_names[0]] = y_py.to_numpy()

        else:
            arr = np.asarray(y_py)
            if arr.ndim == 1:
                n_obs = arr.shape[0]
                if not resp_names:
                    resp_names = ["y"]
                observed_data_dict[resp_names[0]] = arr
            elif arr.ndim == 2:
                n_obs, k = arr.shape
                if not resp_names:
                    resp_names = [f"y_{j}" for j in range(k)]
                for j, resp in enumerate(resp_names):
                    observed_data_dict[resp] = arr[:, j]

    except Exception as e:
        log_warning(f"[brmsfit_to_idata] Could not extract observed data: {e}")

    return observed_data_dict

_brmsfit_get_posterior(brmsfit_obj, **kwargs)

Source code in brmspy/helpers/_rpy2/_conversion.py

def _brmsfit_get_posterior(
    brmsfit_obj: Sexp, **kwargs
) -> tuple[dict[str, np.ndarray], Sexp]:
    import rpy2.robjects as ro
    from rpy2.robjects import pandas2ri
    from rpy2.robjects.conversion import localconverter

    # -------------------------------------------------
    # POSTERIOR (parameters) via posterior::as_draws_df
    # -------------------------------------------------
    as_draws_df = cast(Callable, ro.r("posterior::as_draws_df"))
    draws_r = as_draws_df(brmsfit_obj, **kwargs)

    with localconverter(ro.default_converter + pandas2ri.converter):
        df = pandas2ri.rpy2py(draws_r)

    chain_col = ".chain" if ".chain" in df.columns else "chain"
    draw_col = ".draw" if ".draw" in df.columns else "draw"

    df["draw_idx"] = df.groupby(chain_col)[draw_col].transform(
        lambda x: np.arange(len(x), dtype=int)
    )

    chains = np.sort(df[chain_col].unique())
    n_chains = len(chains)
    n_draws = int(df["draw_idx"].max()) + 1

    posterior_dict: dict[str, np.ndarray] = {}

    for col in df.columns:
        if col in (chain_col, draw_col, ".iteration", "draw_idx"):
            continue

        mat = (
            df.pivot(index="draw_idx", columns=chain_col, values=col)
            .sort_index(axis=0)
            .reindex(columns=chains)
            .to_numpy()
            .T
        )
        posterior_dict[col] = mat

    return posterior_dict, draws_r

_brmsfit_get_predict_generic(brmsfit_obj, function='brms::posterior_predict', resp_names=None, **kwargs)

Source code in brmspy/helpers/_rpy2/_conversion.py

def _brmsfit_get_predict_generic(
    brmsfit_obj,
    function: Literal[
        "brms::posterior_predict",
        "brms::log_lik",
        "brms::posterior_linpred",
        "brms::posterior_epred",
    ] = "brms::posterior_predict",
    resp_names: None | list[str] = None,
    **kwargs,
) -> tuple[dict[str, np.ndarray], Sexp | dict[str, Sexp]]:
    import rpy2.robjects as ro

    if resp_names is None:
        resp_names = _brmsfit_get_response_names(brmsfit_obj)

    nchains, ndraws = _brmsfit_get_counts(brmsfit_obj)

    post_pred_dict: dict[str, np.ndarray] = {}

    r: dict[str, Sexp] | Sexp = ro.NULL
    try:
        r_pp_wrapper = cast(Callable, ro.r(function))

        if not resp_names:
            # No response names found - univariate default
            pp_r = r_pp_wrapper(brmsfit_obj, **kwargs)
            r = pp_r
            post_pred_dict["y"] = _reshape_to_arviz(np.asarray(pp_r), nchains, ndraws)

        elif len(resp_names) == 1:
            # Single response
            resp = resp_names[0]
            pp_r = r_pp_wrapper(
                brmsfit_obj, **kwargs, resp=resp
            )  # Pass as plain string
            r = pp_r
            post_pred_dict[resp] = _reshape_to_arviz(np.asarray(pp_r), nchains, ndraws)

        else:
            # Multivariate: loop over response names
            r = {}
            for resp in resp_names:
                pp_r = r_pp_wrapper(
                    brmsfit_obj, **kwargs, resp=resp
                )  # Pass as plain string!
                post_pred_dict[resp] = _reshape_to_arviz(
                    np.asarray(pp_r), nchains, ndraws
                )
                r[resp] = pp_r

    except Exception as e:
        log_warning(
            f"[brmsfit_to_idata] Could not extract posterior predictive/log_lik: {e}"
        )
        import traceback

        traceback.print_exc()

    return post_pred_dict, r

_brmsfit_get_response_names(brmsfit_obj)

Source code in brmspy/helpers/_rpy2/_conversion.py

def _brmsfit_get_response_names(brmsfit_obj) -> list[str]:
    import rpy2.robjects as ro

    # ------------------------------
    # RESPONSE NAMES via brmsterms()
    # ------------------------------
    resp_names: list[str] = []
    try:
        # Method 1: Use brmsterms to get response variable names
        r_code = """
        function(fit) {
            bterms <- brms::brmsterms(fit$formula)
            if (inherits(bterms, "mvbrmsterms")) {
                # Multivariate: extract response names from each term
                names(bterms$terms)
            } else {
                # Univariate: get the single response
                resp <- bterms$respform
                if (!is.null(resp)) {
                    all.vars(resp)[1]
                } else {
                    # Fallback: parse from formula
                    all.vars(fit$formula$formula)[1]
                }
            }
        }
        """
        get_resp_names = cast(Callable, ro.r(r_code))
        resp_r = get_resp_names(brmsfit_obj)
        resp_names = list(resp_r)
    except Exception as e:
        log_warning(
            f"[brmsfit_to_idata] Could not get response names via brmsterms: {e}"
        )

        # Fallback: try to extract from model formula directly
        try:
            r_fallback = """
            function(fit) {
                # Try to get response names from the model's data
                y <- brms::get_y(fit)
                if (is.matrix(y) || is.data.frame(y)) {
                    colnames(y)
                } else if (!is.null(names(y))) {
                    unique(names(y))[1]
                } else {
                    "y"
                }
            }
            """
            get_resp_fallback = cast(Callable, ro.r(r_fallback))
            resp_r = get_resp_fallback(brmsfit_obj)
            if hasattr(resp_r, "__iter__") and not isinstance(resp_r, str):
                resp_names = [str(r) for r in resp_r if r is not None]
            else:
                resp_names = [str(resp_r)]
        except Exception as e2:
            log_warning(f"[brmsfit_to_idata] Fallback also failed: {e2}")

    return resp_names

_idata_add_resp_names_suffix(idata, suffix, resp_names)

In-place: append suffix to all variables in resp_names across all applicable InferenceData groups.

Mutates idata directly.

Source code in brmspy/helpers/_rpy2/_conversion.py

def _idata_add_resp_names_suffix(
    idata: az.InferenceData,
    suffix: str,
    resp_names: list[str],
) -> None:
    """
    In-place: append `suffix` to all variables in `resp_names` across all
    applicable InferenceData groups.

    Mutates `idata` directly.
    """
    if not suffix or not resp_names:
        return

    for group in idata.groups():
        ds = getattr(idata, group, None)
        if ds is None:
            continue

        rename_map = {
            resp: f"{resp}{suffix}" for resp in resp_names if resp in ds.data_vars
        }

        if rename_map:
            ds = ds.rename(rename_map)
            setattr(idata, group, ds)

kwargs_r(kwargs)

Convert Python keyword arguments to R-compatible format.

Convenience function that applies py_to_r() to all values in a keyword arguments dictionary, preparing them for R function calls.

Parameters:

Name	Type	Description	Default
kwargs	dict or None	Dictionary of keyword arguments where values may be Python objects (dicts, lists, DataFrames, arrays, etc.)	required

Returns:

Type	Description
dict	Dictionary with same keys but R-compatible values, or empty dict if None

Notes

This is a thin wrapper around py_to_r() that operates on dictionaries. It's commonly used to prepare keyword arguments for R function calls via rpy2.

Examples:

from brmspy.helpers.conversion import kwargs_r
import pandas as pd
import numpy as np

# Prepare kwargs for R function
py_kwargs = {
    'data': pd.DataFrame({'y': [1, 2], 'x': [1, 2]}),
    'prior': {'b': [0, 1]},
    'chains': 4,
    'iter': 2000
}

r_kwargs = kwargs_r(py_kwargs)
# All values converted to R objects
# Can now call: r_function(**r_kwargs)

See Also

py_to_r : Underlying conversion function for individual values brmspy.brms.fit : Uses this to prepare user kwargs for R

Source code in brmspy/helpers/_rpy2/_conversion.py

def kwargs_r(kwargs: dict | None) -> dict:
    """
    Convert Python keyword arguments to R-compatible format.

    Convenience function that applies py_to_r() to all values in a
    keyword arguments dictionary, preparing them for R function calls.

    Parameters
    ----------
    kwargs : dict or None
        Dictionary of keyword arguments where values may be Python objects
        (dicts, lists, DataFrames, arrays, etc.)

    Returns
    -------
    dict
        Dictionary with same keys but R-compatible values, or empty dict if None

    Notes
    -----
    This is a thin wrapper around `py_to_r()` that operates on dictionaries.
    It's commonly used to prepare keyword arguments for R function calls via rpy2.

    Examples
    --------

    ```python
    from brmspy.helpers.conversion import kwargs_r
    import pandas as pd
    import numpy as np

    # Prepare kwargs for R function
    py_kwargs = {
        'data': pd.DataFrame({'y': [1, 2], 'x': [1, 2]}),
        'prior': {'b': [0, 1]},
        'chains': 4,
        'iter': 2000
    }

    r_kwargs = kwargs_r(py_kwargs)
    # All values converted to R objects
    # Can now call: r_function(**r_kwargs)
    ```

    See Also
    --------
    py_to_r : Underlying conversion function for individual values
    brmspy.brms.fit : Uses this to prepare user kwargs for R
    """
    if kwargs is None:
        return {}
    return {k: py_to_r(v) for k, v in kwargs.items()}

py_to_r(obj)

Convert arbitrary Python objects to R objects via rpy2.

Comprehensive converter that handles nested structures (dicts, lists), DataFrames, arrays, and scalars. Uses rpy2's converters with special handling for dictionaries (→ R named lists) and lists of dicts.

Parameters:

Name	Type	Description	Default
obj	any	Python object to convert. Supported types: - None → R NULL - dict → R named list (ListVector), recursively - list/tuple of dicts → R list of named lists - list/tuple (other) → R vector or list - pd.DataFrame → R data.frame - np.ndarray → R vector/matrix - scalars (int, float, str, bool) → R atomic types	required

Returns:

Type	Description
rpy2 R object	R representation of the Python object

Notes

Conversion Rules:

1. None: → R NULL
2. DataFrames: → R data.frame (via pandas2ri)
3. Dictionaries: → R named list (ListVector), recursively converting values
4. Lists of dicts: → R list with 1-based indexed names containing named lists
5. Other lists/tuples: → R vectors or lists (via rpy2 default)
6. NumPy arrays: → R vectors/matrices (via numpy2ri)
7. Scalars: → R atomic values

Recursive Conversion:

Dictionary values are recursively converted, allowing nested structures:

{'a': {'b': [1, 2, 3]}}  →  list(a = list(b = c(1, 2, 3)))

List of Dicts:

Lists containing only dicts are converted to R lists with 1-based indexing:

[{'x': 1}, {'x': 2}]  →  list("1" = list(x = 1), "2" = list(x = 2))

Examples:

from brmspy.helpers.conversion import py_to_r
import numpy as np
import pandas as pd

# Scalars
py_to_r(5)        # R: 5
py_to_r("hello")  # R: "hello"
py_to_r(None)     # R: NULL

# Arrays
py_to_r(np.array([1, 2, 3]))  # R: c(1, 2, 3)

# DataFrames
df = pd.DataFrame({'x': [1, 2], 'y': [3, 4]})
py_to_r(df)  # R: data.frame(x = c(1, 2), y = c(3, 4))

See Also

r_to_py : Convert R objects back to Python kwargs_r : Convert keyword arguments dict for R function calls brmspy.brms.fit : Uses this for converting data to R

Source code in brmspy/helpers/_rpy2/_converters/_dispatch.py

def py_to_r(obj: PyObject) -> Sexp:
    """
    Convert arbitrary Python objects to R objects via rpy2.

    Comprehensive converter that handles nested structures (dicts, lists),
    DataFrames, arrays, and scalars. Uses rpy2's converters with special
    handling for dictionaries (→ R named lists) and lists of dicts.

    Parameters
    ----------
    obj : any
        Python object to convert. Supported types:
        - None → R NULL
        - dict → R named list (ListVector), recursively
        - list/tuple of dicts → R list of named lists
        - list/tuple (other) → R vector or list
        - pd.DataFrame → R data.frame
        - np.ndarray → R vector/matrix
        - scalars (int, float, str, bool) → R atomic types

    Returns
    -------
    rpy2 R object
        R representation of the Python object

    Notes
    -----
    **Conversion Rules:**

    1. **None**: → R NULL
    2. **DataFrames**: → R data.frame (via pandas2ri)
    3. **Dictionaries**: → R named list (ListVector), recursively converting values
    4. **Lists of dicts**: → R list with 1-based indexed names containing named lists
    5. **Other lists/tuples**: → R vectors or lists (via rpy2 default)
    6. **NumPy arrays**: → R vectors/matrices (via numpy2ri)
    7. **Scalars**: → R atomic values

    **Recursive Conversion:**

    Dictionary values are recursively converted, allowing nested structures:
    ```python
    {'a': {'b': [1, 2, 3]}}  →  list(a = list(b = c(1, 2, 3)))
    ```

    **List of Dicts:**

    Lists containing only dicts are converted to R lists with 1-based indexing:
    ```python
    [{'x': 1}, {'x': 2}]  →  list("1" = list(x = 1), "2" = list(x = 2))
    ```

    Examples
    --------

    ```python
    from brmspy.helpers.conversion import py_to_r
    import numpy as np
    import pandas as pd

    # Scalars
    py_to_r(5)        # R: 5
    py_to_r("hello")  # R: "hello"
    py_to_r(None)     # R: NULL

    # Arrays
    py_to_r(np.array([1, 2, 3]))  # R: c(1, 2, 3)

    # DataFrames
    df = pd.DataFrame({'x': [1, 2], 'y': [3, 4]})
    py_to_r(df)  # R: data.frame(x = c(1, 2), y = c(3, 4))
    ```

    See Also
    --------
    r_to_py : Convert R objects back to Python
    kwargs_r : Convert keyword arguments dict for R function calls
    brmspy.brms.fit : Uses this for converting data to R
    """
    import rpy2.robjects as ro

    if obj is None:
        return ro.NULL

    if isinstance(obj, ro.Sexp):
        return obj

    if isinstance(obj, RListVectorExtension) and isinstance(obj.r, ro.Sexp):
        return obj.r

    _type = type(obj)
    converter = None

    if _type in _registry._PY2R_CONVERTERS:
        # O(1) lookup first
        converter = _registry._PY2R_CONVERTERS[_type]
    else:
        for _type, _con in _registry._PY2R_CONVERTERS.items():
            if isinstance(obj, _type):
                converter = _con
                break

    assert len(_registry._PY2R_CONVERTERS) > 0, "NO PY2R CONVERTERS"
    assert (
        converter
    ), "object fallback must be in place in __init__.py! This is an issue with the library, not the user!"
    return converter(obj)

posterior(model, **kwargs)

Return posterior draws as idata.

Wrapper around R posterior::as_draws_df().

Parameters:

Name	Type	Description	Default
model	FitResult	Fitted model.	required
**kwargs		Forwarded to posterior::as_draws_df(). e.g inc_warmup, regex, variable	{}

Returns:

Type	Description
PosteriorEpredResult	Result containing idata (ArviZ InferenceData) and an underlying R handle.

Examples:

from brmspy import brms

fit = brms.brm("y ~ x", data=df, chains=4)
ep = brms.posterior(fit)

ep.idata.posterior

Source code in brmspy/_brms_functions/prediction.py

def posterior(
    model: FitResult | ProxyListSexpVector, **kwargs
) -> IDResult[IDPosterior]:
    """
    Return posterior draws as idata.

    Wrapper around R ``posterior::as_draws_df()``.

    Parameters
    ----------
    model : FitResult
        Fitted model.
    **kwargs
        Forwarded to ``posterior::as_draws_df()``. e.g inc_warmup, regex, variable

    Returns
    -------
    PosteriorEpredResult
        Result containing `idata` (ArviZ `InferenceData`) and an underlying R handle.

    Examples
    --------
    ```python
    from brmspy import brms

    fit = brms.brm("y ~ x", data=df, chains=4)
    ep = brms.posterior(fit)

    ep.idata.posterior
    ```
    """
    model_r = py_to_r(model)
    kwargs = kwargs_r(kwargs)

    resp_names = _brmsfit_get_response_names(model_r)
    dims, coords = _brmsfit_get_dims_and_coords(model_r, resp_names=resp_names)

    result, r = _brmsfit_get_posterior(model_r, **kwargs)
    idata = az.from_dict(posterior=result, dims=dims, coords=coords)

    # Add constant data
    constant_data_dict = _brmsfit_get_constant_data(
        model_r, newdata=None, resp_names=resp_names
    )
    _arviz_add_constant_data(
        idata, constant_data_dict, "constant_data", obs_id=coords["obs_id"]
    )

    return IDResult(r=cast(ProxyListSexpVector, r), idata=cast(IDPosterior, idata))

observed_data(model)

Source code in brmspy/_brms_functions/prediction.py

def observed_data(model: FitResult | ProxyListSexpVector) -> IDResult[IDObservedData]:
    import rpy2.robjects as ro

    model_r = py_to_r(model)

    resp_names = _brmsfit_get_response_names(model_r)
    dims, coords = _brmsfit_get_dims_and_coords(model_r, resp_names=resp_names)

    result = _brmsfit_get_observed_data(model_r, resp_names=resp_names)
    r = cast(Any, ro.NULL)

    idata = az.from_dict(observed_data=result, coords=coords, dims=dims)
    idata = cast(IDObservedData, idata)

    # Add constant data
    constant_data_dict = _brmsfit_get_constant_data(
        model_r, newdata=None, resp_names=resp_names
    )
    _arviz_add_constant_data(idata, constant_data_dict, "constant_data")

    return IDResult(r=r, idata=idata)

posterior_epred(model, newdata=None, **kwargs)

posterior_epred(model: FitResult | ProxyListSexpVector, newdata: Literal[None] = None) -> IDResult[IDPosterior]

posterior_epred(model: FitResult | ProxyListSexpVector, newdata: pd.DataFrame) -> IDResult[IDPredictions]

Compute expected posterior predictions (noise-free).

Wrapper around R brms::posterior_epred(). This returns draws of the expected value (typically on the response scale), without observation noise.

Parameters:

Name	Type	Description	Default
model	FitResult	Fitted model.	required
newdata	DataFrame or None	New data for predictions. If None, uses the training data.	None
**kwargs		Forwarded to brms::posterior_epred().	{}

Returns:

Type	Description
PosteriorEpredResult	Result containing idata (ArviZ InferenceData) and an underlying R handle.

See Also

brms::posterior_epred : R documentation

Examples:

from brmspy import brms

fit = brms.brm("y ~ x", data=df, chains=4)
ep = brms.posterior_epred(fit)

ep.idata.posterior

Source code in brmspy/_brms_functions/prediction.py

def posterior_epred(
    model: FitResult | ProxyListSexpVector,
    newdata: pd.DataFrame | None = None,
    **kwargs,
) -> IDResult:
    """
    Compute expected posterior predictions (noise-free).

    Wrapper around R ``brms::posterior_epred()``. This returns draws of the
    expected value (typically on the response scale), without observation noise.

    Parameters
    ----------
    model : FitResult
        Fitted model.
    newdata : pandas.DataFrame or None, default=None
        New data for predictions. If ``None``, uses the training data.
    **kwargs
        Forwarded to ``brms::posterior_epred()``.

    Returns
    -------
    PosteriorEpredResult
        Result containing `idata` (ArviZ `InferenceData`) and an underlying R handle.

    See Also
    --------
    brms::posterior_epred : [R documentation](https://paulbuerkner.com/brms/reference/posterior_epred.brmsfit.html)

    Examples
    --------
    ```python
    from brmspy import brms

    fit = brms.brm("y ~ x", data=df, chains=4)
    ep = brms.posterior_epred(fit)

    ep.idata.posterior
    ```
    """
    model_r = py_to_r(model)
    data_r = py_to_r(newdata)
    kwargs = kwargs_r(kwargs)

    resp_names = _brmsfit_get_response_names(model_r)
    dims, coords = _brmsfit_get_dims_and_coords(
        model_r, resp_names=resp_names, newdata=newdata
    )

    result, r = _brmsfit_get_predict_generic(
        model_r,
        newdata=data_r,
        function="brms::posterior_epred",
        resp_names=resp_names,
        **kwargs,
    )

    if newdata is None:
        idata = az.from_dict(posterior=result, coords=coords, dims=dims)
        idata = cast(IDPosterior, idata)
    else:
        idata = az.from_dict(predictions=result, coords=coords, dims=dims)
        idata = cast(IDPredictions, idata)

    _idata_add_resp_names_suffix(idata, "_mean", resp_names)

    # Add constant data
    constant_data_dict = _brmsfit_get_constant_data(
        model_r, newdata=newdata, resp_names=resp_names
    )
    group_name = "constant_data" if newdata is None else "predictions_constant_data"
    _arviz_add_constant_data(idata, constant_data_dict, group_name)

    return IDResult(r=cast(ProxyListSexpVector, r), idata=idata)

posterior_predict(model, newdata=None, **kwargs)

posterior_predict(model: FitResult | ProxyListSexpVector, newdata: Literal[None] = None, **kwargs) -> IDResult[IDPosteriorPredictive]

posterior_predict(model: FitResult | ProxyListSexpVector, newdata: pd.DataFrame, **kwargs) -> IDResult[IDPredictions]

Draw from the posterior predictive distribution (includes observation noise).

Wrapper around R brms::posterior_predict().

Parameters:

Name	Type	Description	Default
model	FitResult	Fitted model.	required
newdata	DataFrame or None	New data for predictions. If None, uses the training data.	None
**kwargs		Forwarded to brms::posterior_predict().	{}

Returns:

Type	Description
PosteriorPredictResult	Result containing idata (ArviZ InferenceData) and an underlying R handle.

See Also

brms::posterior_predict : R documentation

Examples:

from brmspy import brms

fit = brms.brm("y ~ x", data=df, chains=4)
pp = brms.posterior_predict(fit)

pp.idata.posterior_predictive

Source code in brmspy/_brms_functions/prediction.py

def posterior_predict(
    model: FitResult | ProxyListSexpVector,
    newdata: pd.DataFrame | None = None,
    **kwargs,
) -> IDResult:
    """
    Draw from the posterior predictive distribution (includes observation noise).

    Wrapper around R ``brms::posterior_predict()``.

    Parameters
    ----------
    model : FitResult
        Fitted model.
    newdata : pandas.DataFrame or None, default=None
        New data for predictions. If ``None``, uses the training data.
    **kwargs
        Forwarded to ``brms::posterior_predict()``.

    Returns
    -------
    PosteriorPredictResult
        Result containing `idata` (ArviZ `InferenceData`) and an underlying R handle.

    See Also
    --------
    brms::posterior_predict : [R documentation](https://paulbuerkner.com/brms/reference/posterior_predict.brmsfit.html)

    Examples
    --------
    ```python
    from brmspy import brms

    fit = brms.brm("y ~ x", data=df, chains=4)
    pp = brms.posterior_predict(fit)

    pp.idata.posterior_predictive
    ```
    """
    model_r = py_to_r(model)
    data_r = py_to_r(newdata)
    kwargs = kwargs_r(kwargs)

    resp_names = _brmsfit_get_response_names(model_r)
    dims, coords = _brmsfit_get_dims_and_coords(
        model_r, resp_names=resp_names, newdata=newdata
    )

    result, r = _brmsfit_get_predict_generic(
        model_r,
        newdata=data_r,
        function="brms::posterior_predict",
        resp_names=resp_names,
        **kwargs,
    )

    if newdata is None:
        idata = az.from_dict(
            posterior_predictive=result,
            dims=dims,
            coords=coords,
        )
        idata = cast(IDPosteriorPredictive, idata)
    else:
        idata = az.from_dict(
            predictions=result,
            dims=dims,
            coords=coords,
        )
        idata = cast(IDPredictions, idata)

    # Add constant data
    constant_data_dict = _brmsfit_get_constant_data(
        model_r, newdata=newdata, resp_names=resp_names
    )
    group_name = "constant_data" if newdata is None else "predictions_constant_data"
    _arviz_add_constant_data(idata, constant_data_dict, group_name)

    return IDResult(r=cast(ProxyListSexpVector, r), idata=idata)

posterior_linpred(model, newdata=None, **kwargs)

posterior_linpred(model: FitResult | ProxyListSexpVector, newdata: Literal[None] = None, **kwargs) -> IDResult[IDPosterior]

posterior_linpred(model: FitResult | ProxyListSexpVector, newdata: pd.DataFrame, **kwargs) -> IDResult[IDPredictions]

Draw from the linear predictor.

Wrapper around R brms::posterior_linpred(). This typically returns draws on the link scale (before applying the inverse link), unless you pass transform=True.

Parameters:

Name	Type	Description	Default
model	FitResult	Fitted model.	required
newdata	DataFrame or None	New data for predictions. If None, uses the training data.	None
**kwargs		Forwarded to brms::posterior_linpred() (commonly transform or ndraws).	{}

Returns:

Type	Description
PosteriorLinpredResult	Result containing idata (ArviZ InferenceData) and an underlying R handle.

See Also

brms::posterior_linpred : R documentation

Examples:

from brmspy import brms

fit = brms.brm("y ~ x", data=df, chains=4)
lp = brms.posterior_linpred(fit, transform=False)

lp.idata.predictions

Source code in brmspy/_brms_functions/prediction.py

def posterior_linpred(
    model: FitResult | ProxyListSexpVector,
    newdata: pd.DataFrame | None = None,
    **kwargs,
) -> IDResult:
    """
    Draw from the linear predictor.

    Wrapper around R ``brms::posterior_linpred()``. This typically returns draws
    on the link scale (before applying the inverse link), unless you pass
    ``transform=True``.

    Parameters
    ----------
    model : FitResult
        Fitted model.
    newdata : pandas.DataFrame or None, default=None
        New data for predictions. If ``None``, uses the training data.
    **kwargs
        Forwarded to ``brms::posterior_linpred()`` (commonly ``transform`` or ``ndraws``).

    Returns
    -------
    PosteriorLinpredResult
        Result containing `idata` (ArviZ `InferenceData`) and an underlying R handle.

    See Also
    --------
    brms::posterior_linpred : [R documentation](https://paulbuerkner.com/brms/reference/posterior_linpred.brmsfit.html)

    Examples
    --------
    ```python
    from brmspy import brms

    fit = brms.brm("y ~ x", data=df, chains=4)
    lp = brms.posterior_linpred(fit, transform=False)

    lp.idata.predictions
    ```
    """
    import rpy2.robjects as ro

    model_r = py_to_r(model)
    data_r = py_to_r(newdata)
    kwargs = kwargs_r(kwargs)

    resp_names = _brmsfit_get_response_names(model_r)
    dims, coords = _brmsfit_get_dims_and_coords(
        model_r, resp_names=resp_names, newdata=newdata
    )

    result, r = _brmsfit_get_predict_generic(
        model_r,
        newdata=data_r,
        function="brms::posterior_linpred",
        resp_names=resp_names,
        **kwargs,
    )

    if newdata is None:
        idata = az.from_dict(
            posterior=result,
            dims=dims,
            coords=coords,
        )
        idata = cast(IDPosterior, idata)
    else:
        idata = az.from_dict(
            predictions=result,
            dims=dims,
            coords=coords,
        )
        idata = cast(IDPredictions, idata)

    _idata_add_resp_names_suffix(idata, "_linpred", resp_names)

    # Add constant data
    constant_data_dict = _brmsfit_get_constant_data(
        model_r, newdata=newdata, resp_names=resp_names
    )
    group_name = "constant_data" if newdata is None else "predictions_constant_data"
    _arviz_add_constant_data(idata, constant_data_dict, group_name)

    return IDResult(r=cast(ProxyListSexpVector, r), idata=idata)

log_lik(model, newdata=None, **kwargs)

log_lik(model: FitResult | ProxyListSexpVector, newdata: Literal[None] = None, **kwargs) -> IDResult[IDLogLikelihoodInsample]

log_lik(model: FitResult | ProxyListSexpVector, newdata: pd.DataFrame, **kwargs) -> IDResult[IDLogLikelihoodOutsample]

Compute pointwise log-likelihood draws.

Wrapper around R brms::log_lik(). The result is useful for LOO/WAIC via ArviZ.

Parameters:

Name	Type	Description	Default
model	FitResult	Fitted model.	required
newdata	DataFrame or None	New data. If None, uses the training data.	None
**kwargs		Forwarded to brms::log_lik().	{}

Returns:

Type	Description
LogLikResult	Result containing idata (ArviZ InferenceData) and an underlying R handle.

See Also

brms::log_lik : R documentation

Examples:

from brmspy import brms
import arviz as az

fit = brms.brm("y ~ x", data=df, chains=4)
ll = brms.log_lik(fit)

az.loo(ll.idata)

Source code in brmspy/_brms_functions/prediction.py

def log_lik(
    model: FitResult | ProxyListSexpVector,
    newdata: pd.DataFrame | None = None,
    **kwargs,
) -> IDResult:
    """
    Compute pointwise log-likelihood draws.

    Wrapper around R ``brms::log_lik()``. The result is useful for LOO/WAIC via ArviZ.

    Parameters
    ----------
    model : FitResult
        Fitted model.
    newdata : pandas.DataFrame or None, default=None
        New data. If ``None``, uses the training data.
    **kwargs
        Forwarded to ``brms::log_lik()``.

    Returns
    -------
    LogLikResult
        Result containing `idata` (ArviZ `InferenceData`) and an underlying R handle.

    See Also
    --------
    brms::log_lik : [R documentation](https://paulbuerkner.com/brms/reference/log_lik.brmsfit.html)

    Examples
    --------
    ```python
    from brmspy import brms
    import arviz as az

    fit = brms.brm("y ~ x", data=df, chains=4)
    ll = brms.log_lik(fit)

    az.loo(ll.idata)
    ```
    """
    import rpy2.robjects as ro

    model_r = py_to_r(model)
    data_r = py_to_r(newdata)
    kwargs = kwargs_r(kwargs)

    resp_names = _brmsfit_get_response_names(model_r)
    dims, coords = _brmsfit_get_dims_and_coords(
        model_r, resp_names=resp_names, newdata=newdata
    )

    result, r = _brmsfit_get_predict_generic(
        model_r,
        newdata=data_r,
        function="brms::log_lik",
        resp_names=resp_names,
        **kwargs,
    )
    if newdata is None:
        idata = az.from_dict(log_likelihood=result, dims=dims, coords=coords)
        idata = cast(IDLogLikelihoodInsample, idata)
    else:
        idata = az.from_dict(log_likelihood=result, dims=dims, coords=coords)
        idata = cast(IDLogLikelihoodOutsample, idata)

    # Add constant data
    constant_data_dict = _brmsfit_get_constant_data(
        model_r, newdata=newdata, resp_names=resp_names
    )
    group_name = "constant_data" if newdata is None else "predictions_constant_data"
    _arviz_add_constant_data(idata, constant_data_dict, group_name)

    return IDResult(r=cast(ProxyListSexpVector, r), idata=idata)