Source code for edstan.mcmc

from cmdstanpy import CmdStanMCMC
from numpy.typing import NDArray
import numpy as np
import pandas as pd




class EdStanMCMC:
    """
    A wrapper around :class:`pystan.CmdStanMCMC` that adds additional methods.

    This class delegates all unspecified attribute access to the underlying :class:`pystan.CmdStanMCMC` instance via
    :meth:`EdStanMCMC.__getattr__`. This allows it to behave like a :class:`pystan.CmdStanMCMC` object while also
    providing custom methods.
    """



    def __init__(
        self,
        mcmc: CmdStanMCMC,
        ii_labels: NDArray,
        jj_labels: NDArray,
        max_per_item: NDArray[np.integer],
    ):
        """
        Initializes an :class:`EdStanMCMC` instance.

        An instance of this class is generated by :meth:`ModelMCMC.sample_from_long` or
        :meth:`EdStanModel.sample_from_wide`. Though the class may be initialized directly, this is not the intended
        usage.

        :param mcmc: A fitted :mod:`edstan` model using MCMC.
        :param ii_labels: Labels associated with the items.
        :param jj_labels: Labels associated with the persons.
        :param max_per_item: The maximum score per item.
        """
        self.mcmc = mcmc
        self.ii_labels = ii_labels
        self.jj_labels = jj_labels
        self.max_per_item = max_per_item


    def __getattr__(self, name):
        return getattr(self.mcmc, name)



    def item_summary(self, **kwargs):
        """
        A wrapper around :meth:`pystan.CmdStanMCMC.summary` that provides posterior summaries grouped by item.

        :param kwargs: Additional optional arguments passed to :meth:`pystan.CmdStanMCMC.summary`, such as 'percentiles'
            and 'sig_figs'.
        :return: A summary :class:`pandas.DataFrame` filtered to include item and distribution parameters only, having a
            multi-index that associates parameters with their respective item labels (or the person distribution).
        """
        summary = self.mcmc.summary(**kwargs)
        summary.index.name = "parameter"

        expected = _get_expected_parameters_by_group(
            item_labels=self.ii_labels,
            max_per_item=self.max_per_item,
            rasch_family="sigma" in summary.index,
            ratings_model="kappa[1]" in summary.index,
        )

        return expected.merge(summary.reset_index(), on="parameter").set_index(
            ["parameter group", "parameter"]
        )




    def person_summary(self, **kwargs):
        """
        A wrapper around :meth:pystan.CmdStanMCMC.summary that provides posterior summaries grouped by person.

        :param kwargs: Additional optional arguments passed to :meth:`pystan.CmdStanMCMC.summary`, such as 'percentiles'
            and 'sig_figs'.
        :return: A summary :class:`pandas.DataFrame` filtered to include person parameters only, having a multi-index
            that associates parameters with their respective person labels.
        """
        summary = self.mcmc.summary(**kwargs)
        summary = summary.loc[summary.index.str.match("theta")]
        summary["person"] = self.jj_labels
        summary.index.name = "parameter"
        return summary.reset_index().set_index(["person", "parameter"])




def _get_expected_parameters_by_group(
    item_labels, max_per_item, rasch_family: bool, ratings_model: bool
):
    """Generate a DataFrame listing the expected item/distribution parameters and their groupings."""
    holder = []

    if ratings_model:
        betas_per_item = np.ones(len(max_per_item), dtype=int)
    else:
        betas_per_item = np.array(max_per_item, dtype=int)

    beta_counter = 0
    for item, item_max in enumerate(betas_per_item):
        if not rasch_family:
            holder.append((item_labels[item], f"alpha[{item + 1}]"))
        for _ in range(item_max):
            beta_counter += 1
            holder.append((item_labels[item], f"beta[{beta_counter}]"))

    if ratings_model:
        for j in range(int(max(max_per_item))):
            holder.append(("Rating scale steps", f"kappa[{j + 1}]"))

    holder.append(("Ability distribution", "lambda[1]"))

    if rasch_family:
        holder.append(("Ability distribution", "sigma"))

    df = pd.DataFrame(holder)
    df.columns = ["parameter group", "parameter"]

    return df