Queries

Probabilistic queries and inference computations on stochastic morphisms.

queries

Probability queries on stochastic morphisms.

Provides functions to extract probabilities, compute marginal probabilities, and compute expected values under stochastic morphisms.

prob

prob(f: Morphism, domain_indices: Tensor, codomain_indices: Tensor) -> Tensor

Extract probabilities at specific (domain, codomain) index pairs.

PARAMETER	DESCRIPTION
f	A (stochastic) morphism A → B. TYPE: Morphism
domain_indices	Integer indices into the domain. Shape (batch,) or (batch, n_dom_dims). TYPE: Tensor
codomain_indices	Integer indices into the codomain. Shape (batch,) or (batch, n_cod_dims). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Probability values at the specified index pairs. Shape (batch,).

Source code in src/quivers/stochastic/queries.py

def prob(
    f: Morphism,
    domain_indices: torch.Tensor,
    codomain_indices: torch.Tensor,
) -> torch.Tensor:
    """Extract probabilities at specific (domain, codomain) index pairs.

    Parameters
    ----------
    f : Morphism
        A (stochastic) morphism A → B.
    domain_indices : torch.Tensor
        Integer indices into the domain. Shape (batch,) or (batch, n_dom_dims).
    codomain_indices : torch.Tensor
        Integer indices into the codomain. Shape (batch,) or (batch, n_cod_dims).

    Returns
    -------
    torch.Tensor
        Probability values at the specified index pairs. Shape (batch,).
    """
    t = f.tensor

    if domain_indices.ndim == 1:
        domain_indices = domain_indices.unsqueeze(-1)

    if codomain_indices.ndim == 1:
        codomain_indices = codomain_indices.unsqueeze(-1)

    indices = torch.cat([domain_indices, codomain_indices], dim=-1)
    idx_tuple = tuple(indices[:, i] for i in range(indices.shape[1]))
    return t[idx_tuple]

marginal_prob

marginal_prob(f: Morphism, codomain_indices: Tensor) -> Tensor

Compute marginal probability over the domain for specific codomain indices.

Marginalizes (sums) over all domain elements: P(b) = Σ_a f(a, b) / |A|

This assumes a uniform prior over the domain.

PARAMETER	DESCRIPTION
f	A (stochastic) morphism A → B. TYPE: Morphism
codomain_indices	Integer indices into the codomain. Shape (batch,). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Marginal probabilities. Shape (batch,).

Source code in src/quivers/stochastic/queries.py

def marginal_prob(f: Morphism, codomain_indices: torch.Tensor) -> torch.Tensor:
    """Compute marginal probability over the domain for specific codomain indices.

    Marginalizes (sums) over all domain elements:
        P(b) = Σ_a f(a, b) / |A|

    This assumes a uniform prior over the domain.

    Parameters
    ----------
    f : Morphism
        A (stochastic) morphism A → B.
    codomain_indices : torch.Tensor
        Integer indices into the codomain. Shape (batch,).

    Returns
    -------
    torch.Tensor
        Marginal probabilities. Shape (batch,).
    """
    t = f.tensor
    n_dom = f.domain.ndim
    dom_dims = tuple(range(n_dom))

    # average over domain (uniform prior)
    marginal = t.mean(dim=dom_dims)

    if codomain_indices.ndim == 1:
        codomain_indices = codomain_indices.unsqueeze(-1)

    idx_tuple = tuple(codomain_indices[:, i] for i in range(codomain_indices.shape[1]))
    return marginal[idx_tuple]

expectation

expectation(f: Morphism, values: Tensor) -> Tensor

Compute expected value of a function under a stochastic morphism.

For each domain element a, computes: E_f[v | a] = Σ_b f(a, b) · values(b)

PARAMETER	DESCRIPTION
f	A (stochastic) morphism A → B. TYPE: Morphism
values	Real-valued function on the codomain. Shape (*codomain.shape). TYPE: Tensor

RETURNS	DESCRIPTION
Tensor	Expected values. Shape (*domain.shape).

Source code in src/quivers/stochastic/queries.py

def expectation(
    f: Morphism,
    values: torch.Tensor,
) -> torch.Tensor:
    """Compute expected value of a function under a stochastic morphism.

    For each domain element a, computes:
        E_f[v | a] = Σ_b f(a, b) · values(b)

    Parameters
    ----------
    f : Morphism
        A (stochastic) morphism A → B.
    values : torch.Tensor
        Real-valued function on the codomain. Shape (*codomain.shape).

    Returns
    -------
    torch.Tensor
        Expected values. Shape (*domain.shape).
    """
    t = f.tensor
    n_dom = f.domain.ndim
    n_cod = f.codomain.ndim

    # expand values to broadcast: (1, ..., 1, *cod_shape)
    v = values

    for _ in range(n_dom):
        v = v.unsqueeze(0)

    weighted = t * v
    cod_dims = tuple(range(n_dom, n_dom + n_cod))
    return weighted.sum(dim=cod_dims)