tramdag.transforms¶

tramdag.transforms ¶

Univariate transforms for CausalFlowDAG nodes.

Each continuous node carries a monotone 1-D transform h (zuko-backed) that maps the observed value to the latent scale; ordinal nodes carry a cutpoint ("ordered logit") transform. Together with the additive shift terms they form one triangular flow from the standard-logistic latent to the observed variables.

continuous: \(u = h(x) + s(\mathrm{pa})\) with \(h\) Bernstein / RQ-spline / affine on the value range scaled from the train range_q quantiles to \([-B, B]\).
ordinal: \(P(Y \le k) = \sigma(\vartheta_k - s(\mathrm{pa}))\) with increasing cutpoints \(\vartheta\).

docs/model.md is the guide to both.

Notes

BOUND is the half-width \(B\) of the pre-scaled domain. It is fixed, because the quantile pre-map makes one domain work for any data scale.

RANGE_Q is the quantile level of the range pre-map. The flow calibration scales the train q/1 - q quantiles onto \([-B, B]\). marginal_init_theta maps that domain onto the latent q/1 - q quantiles. Both must use the same level, because they calibrate each other only then.

_CDF_EPS has two uses in the marginal starts. As a floor on an empirical CDF value, it bounds the initial cutpoints at about \(\pm\operatorname{logit}(10^{-3}) \approx \pm 6.9\). An empty class or a saturated tail then starts implausible but finite. As a floor on the latent-scale gap between two neighbouring cutpoints or Bernstein control points, it keeps the sequence strictly increasing. The softplus inverse needs this.

_U_EPS is the clamp margin of the uniform draw in StandardLogistic. The value 1e-7 keeps u off 0 and 1 at float32 resolution. It caps \(|z|\) at about 16.1, which is past any latent that a fitted shift reaches.

StandardLogistic ¶

Standard logistic base distribution (the TRAM latent).

Source code in src/tramdag/transforms.py

class StandardLogistic:
    """Standard logistic base distribution (the TRAM latent)."""

    @staticmethod
    def log_prob(z: Tensor) -> Tensor:
        """Give the log density at ``z``.

        Parameters
        ----------
        z : Tensor
            Evaluation points.

        Returns
        -------
        Tensor
            The log density, same shape as ``z``.
        """
        return -z - 2.0 * torch.nn.functional.softplus(-z)

    @staticmethod
    def sample(shape, device=None, generator=None) -> Tensor:
        """Draw samples of the given ``shape``.

        Parameters
        ----------
        shape : tuple[int, ...]
            Shape of the sample tensor.
        device : torch.device | str | None, optional
            Target device, by default ``None``.
        generator : torch.Generator | None, optional
            Random source for a reproducible draw, by default ``None``.

        Returns
        -------
        Tensor
            The samples.
        """
        u = torch.rand(shape, device=device, generator=generator)
        return StandardLogistic.icdf(u)

    @staticmethod
    def icdf(u: Tensor) -> Tensor:
        """Give the quantile at probability ``u``.

        Parameters
        ----------
        u : Tensor
            Probabilities in (0, 1). Clamped off 0 and 1 by ``_U_EPS``.

        Returns
        -------
        Tensor
            The quantiles, same shape as ``u``.
        """
        return torch.logit(u, eps=_U_EPS)

log_prob `staticmethod` ¶

log_prob(z)

Give the log density at z.

Parameters:

Name	Type	Description	Default
`z`	`Tensor`	Evaluation points.	required

Returns:

Type	Description
`Tensor`	The log density, same shape as `z`.

Source code in src/tramdag/transforms.py

@staticmethod
def log_prob(z: Tensor) -> Tensor:
    """Give the log density at ``z``.

    Parameters
    ----------
    z : Tensor
        Evaluation points.

    Returns
    -------
    Tensor
        The log density, same shape as ``z``.
    """
    return -z - 2.0 * torch.nn.functional.softplus(-z)

sample `staticmethod` ¶

sample(shape, device=None, generator=None)

Draw samples of the given shape.

Parameters:

Name	Type	Description	Default
`shape`	`tuple[int, ...]`	Shape of the sample tensor.	required
`device`	`device \| str \| None`	Target device, by default `None`.	`None`
`generator`	`Generator \| None`	Random source for a reproducible draw, by default `None`.	`None`

Returns:

Type	Description
`Tensor`	The samples.

Source code in src/tramdag/transforms.py

@staticmethod
def sample(shape, device=None, generator=None) -> Tensor:
    """Draw samples of the given ``shape``.

    Parameters
    ----------
    shape : tuple[int, ...]
        Shape of the sample tensor.
    device : torch.device | str | None, optional
        Target device, by default ``None``.
    generator : torch.Generator | None, optional
        Random source for a reproducible draw, by default ``None``.

    Returns
    -------
    Tensor
        The samples.
    """
    u = torch.rand(shape, device=device, generator=generator)
    return StandardLogistic.icdf(u)

icdf `staticmethod` ¶

icdf(u)

Give the quantile at probability u.

Parameters:

Name	Type	Description	Default
`u`	`Tensor`	Probabilities in (0, 1). Clamped off 0 and 1 by `_U_EPS`.	required

Returns:

Type	Description
`Tensor`	The quantiles, same shape as `u`.

Source code in src/tramdag/transforms.py

@staticmethod
def icdf(u: Tensor) -> Tensor:
    """Give the quantile at probability ``u``.

    Parameters
    ----------
    u : Tensor
        Probabilities in (0, 1). Clamped off 0 and 1 by ``_U_EPS``.

    Returns
    -------
    Tensor
        The quantiles, same shape as ``u``.
    """
    return torch.logit(u, eps=_U_EPS)

BernsteinUT ¶

Bases: _ScaledUT

TRAM-style Bernstein polynomial transform (zuko BernsteinTransform).

Parameters:

Name	Type	Description	Default
`n_coeffs`	`int`	Number of unconstrained coefficients, by default 20. The polynomial has degree `n_coeffs + 1` (tram's `order`).	`20`
`range_q`	`float`	Domain quantile level, see `_ScaledUT`.	`RANGE_Q`

Source code in src/tramdag/transforms.py

class BernsteinUT(_ScaledUT):
    """TRAM-style Bernstein polynomial transform (zuko ``BernsteinTransform``).

    Parameters
    ----------
    n_coeffs : int, optional
        Number of unconstrained coefficients, by default 20. The polynomial
        has degree ``n_coeffs + 1`` (tram's ``order``).
    range_q : float, optional
        Domain quantile level, see ``_ScaledUT``.
    """

    def __init__(self, n_coeffs: int = 20, range_q: float = RANGE_Q):
        super().__init__(range_q)
        self._n = n_coeffs

    @property
    def n_params(self) -> int:
        """int: number of transform parameters this transform needs."""
        return self._n

    def _build(self, theta: Tensor):
        return BernsteinTransform(theta, bound=self.bound)

    def marginal_init_theta(self, column: np.ndarray | None) -> Tensor:
        r"""Give the unconstrained Bernstein coefficients of the marginal start.

        With ``column`` the control points follow the node's **empirical
        marginal**: control point $k$ is $\operatorname{logit} \hat F(y_k)$ at the
        value $y_k$ sitting at $k/\text{order}$ of the pre-scaled domain, so
        the polynomial starts as the Bernstein approximation of
        $\operatorname{logit} \hat F(y)$. Without it the control points are
        equally spaced, the plain linear map from the pre-scaled domain
        $[-B, B]$ onto $[\operatorname{logit} q, \operatorname{logit}(1-q)]$ with
        $q$ = ``range_q``.

        Parameters
        ----------
        column : numpy.ndarray | None, optional
            The node's raw training column. ``None`` gives the linear map,
            which needs no data.

        Raises
        ------
        ValueError
            With ``range_q=0``: the domain ends are the data min/max, whose
            latent quantile target $\operatorname{logit} 0$ is undefined. Skip
            ``marginal_init`` for a min-max-domain model.

        Returns
        -------
        Tensor
            The coefficients, shape ``(n_params,)``.

        Notes
        -----
        The unconstrained coefficients come from inverting zuko's
        ``BernsteinTransform._constrain_theta`` (a cumsum of softplus
        differences with the first two and the last two control points tied,
        hence the two averaged pairs below).
        """
        n = self._n
        if self.range_q == 0:
            raise ValueError(
                "the marginal start maps the domain ends onto the latent "
                "range_q quantiles, and logit(0) is undefined; a "
                "range_q=0 (min-max domain) model has no marginal start; "
                "fit it with marginal_init=False"
            )
        order = n + 1
        points = self._init_control_points(column if n >= 3 else None, order)
        diffs = np.maximum(np.diff(points), _CDF_EPS)
        first, last = diffs[:2].mean(), diffs[-2:].mean()
        diffs[:2], diffs[-2:] = first, last
        shift = math.log(2.0) * n / 2.0
        theta = np.empty(n)
        theta[0] = points[0] + shift
        theta[1:] = np.log(np.expm1(diffs[1:n]))
        return torch.as_tensor(theta, dtype=self.xmin.dtype, device=self.xmin.device)

    def _init_control_points(self, column: np.ndarray | None, order: int) -> np.ndarray:
        r"""Give the ``order + 1`` target control points of the marginal start.

        $\operatorname{logit}$ of the empirical CDF at the equally spaced domain values
        when a column is given, an equally spaced ramp from ``a`` to ``-a``
        otherwise.
        """
        q = self.range_q
        a = math.log(q) - math.log(1.0 - q)
        lo, hi = float(self.xmin), float(self.xmax)
        u = np.arange(order + 1) / order
        if column is None or hi <= lo:
            return a + u * (-2.0 * a)
        values = np.sort(np.asarray(column, dtype=float))
        cdf = np.searchsorted(values, lo + u * (hi - lo), side="right") / values.size
        cdf = np.clip(cdf, _CDF_EPS, 1 - _CDF_EPS)
        return np.log(cdf) - np.log1p(-cdf)

n_params `property` ¶

n_params

int: number of transform parameters this transform needs.

marginal_init_theta ¶

marginal_init_theta(column)

Give the unconstrained Bernstein coefficients of the marginal start.

With column the control points follow the node's empirical marginal: control point \(k\) is \(\operatorname{logit} \hat F(y_k)\) at the value \(y_k\) sitting at \(k/\text{order}\) of the pre-scaled domain, so the polynomial starts as the Bernstein approximation of \(\operatorname{logit} \hat F(y)\). Without it the control points are equally spaced, the plain linear map from the pre-scaled domain \([-B, B]\) onto \([\operatorname{logit} q, \operatorname{logit}(1-q)]\) with \(q\) = range_q.

Parameters:

Name	Type	Description	Default
`column`	`ndarray \| None`	The node's raw training column. `None` gives the linear map, which needs no data.	required

Raises:

Type	Description
`ValueError`	With `range_q=0`: the domain ends are the data min/max, whose latent quantile target \(\operatorname{logit} 0\) is undefined. Skip `marginal_init` for a min-max-domain model.

Returns:

Type	Description
`Tensor`	The coefficients, shape `(n_params,)`.

Notes

The unconstrained coefficients come from inverting zuko's BernsteinTransform._constrain_theta (a cumsum of softplus differences with the first two and the last two control points tied, hence the two averaged pairs below).

Source code in src/tramdag/transforms.py

def marginal_init_theta(self, column: np.ndarray | None) -> Tensor:
    r"""Give the unconstrained Bernstein coefficients of the marginal start.

    With ``column`` the control points follow the node's **empirical
    marginal**: control point $k$ is $\operatorname{logit} \hat F(y_k)$ at the
    value $y_k$ sitting at $k/\text{order}$ of the pre-scaled domain, so
    the polynomial starts as the Bernstein approximation of
    $\operatorname{logit} \hat F(y)$. Without it the control points are
    equally spaced, the plain linear map from the pre-scaled domain
    $[-B, B]$ onto $[\operatorname{logit} q, \operatorname{logit}(1-q)]$ with
    $q$ = ``range_q``.

    Parameters
    ----------
    column : numpy.ndarray | None, optional
        The node's raw training column. ``None`` gives the linear map,
        which needs no data.

    Raises
    ------
    ValueError
        With ``range_q=0``: the domain ends are the data min/max, whose
        latent quantile target $\operatorname{logit} 0$ is undefined. Skip
        ``marginal_init`` for a min-max-domain model.

    Returns
    -------
    Tensor
        The coefficients, shape ``(n_params,)``.

    Notes
    -----
    The unconstrained coefficients come from inverting zuko's
    ``BernsteinTransform._constrain_theta`` (a cumsum of softplus
    differences with the first two and the last two control points tied,
    hence the two averaged pairs below).
    """
    n = self._n
    if self.range_q == 0:
        raise ValueError(
            "the marginal start maps the domain ends onto the latent "
            "range_q quantiles, and logit(0) is undefined; a "
            "range_q=0 (min-max domain) model has no marginal start; "
            "fit it with marginal_init=False"
        )
    order = n + 1
    points = self._init_control_points(column if n >= 3 else None, order)
    diffs = np.maximum(np.diff(points), _CDF_EPS)
    first, last = diffs[:2].mean(), diffs[-2:].mean()
    diffs[:2], diffs[-2:] = first, last
    shift = math.log(2.0) * n / 2.0
    theta = np.empty(n)
    theta[0] = points[0] + shift
    theta[1:] = np.log(np.expm1(diffs[1:n]))
    return torch.as_tensor(theta, dtype=self.xmin.dtype, device=self.xmin.device)

SplineUT ¶

Bases: _ScaledUT

Monotone rational-quadratic spline (zuko MonotonicRQSTransform).

Parameters:

Name	Type	Description	Default
`bins`	`int`	Number of spline bins, by default 8.	`8`
`range_q`	`float`	Domain quantile level, see `_ScaledUT`.	`RANGE_Q`

Source code in src/tramdag/transforms.py

class SplineUT(_ScaledUT):
    """Monotone rational-quadratic spline (zuko ``MonotonicRQSTransform``).

    Parameters
    ----------
    bins : int, optional
        Number of spline bins, by default 8.
    range_q : float, optional
        Domain quantile level, see ``_ScaledUT``.
    """

    def __init__(self, bins: int = 8, range_q: float = RANGE_Q):
        super().__init__(range_q)
        self.bins = bins

    @property
    def n_params(self) -> int:
        """int: number of transform parameters this transform needs."""
        return 3 * self.bins - 1

    def _build(self, theta: Tensor):
        K = self.bins
        widths, heights, derivs = (
            theta[..., :K],
            theta[..., K : 2 * K],
            theta[..., 2 * K :],
        )
        return MonotonicRQSTransform(widths, heights, derivs, bound=self.bound)

n_params `property` ¶

n_params

int: number of transform parameters this transform needs.

AffineUT ¶

Bases: _ScaledUT

Monotone affine transform: the node-conditional is a logistic GLM.

Parameters:

Name	Type	Description	Default
`range_q`	`float`	Domain quantile level, see `_ScaledUT`.	`RANGE_Q`

Source code in src/tramdag/transforms.py

class AffineUT(_ScaledUT):
    """Monotone affine transform: the node-conditional is a logistic GLM.

    Parameters
    ----------
    range_q : float, optional
        Domain quantile level, see ``_ScaledUT``.
    """

    @property
    def n_params(self) -> int:
        """int: number of transform parameters this transform needs."""
        return 2

    def _build(self, theta: Tensor):
        return MonotonicAffineTransform(theta[..., 0], theta[..., 1])

n_params `property` ¶

n_params

int: number of transform parameters this transform needs.

ordinal_bounds ¶

ordinal_bounds(theta_tilde, shift, y)

Give the shifted cutpoint interval of each observed level.

scores.py reads these bounds to form the latent-scale derivative.

Parameters:

Name	Type	Description	Default
`theta_tilde`	`Tensor`	Unconstrained cutpoint parameters, shape `(n, levels - 1)`.	required
`shift`	`Tensor`	The node's shift, shape `(n,)`.	required
`y`	`Tensor`	Observed levels `0..levels-1`, shape `(n,)`.	required

Returns:

Type	Description
`tuple[Tensor, Tensor]`	The lower and upper shifted cutpoint of each row's level, `(n,)` each.

Source code in src/tramdag/transforms.py

def ordinal_bounds(
    theta_tilde: Tensor, shift: Tensor, y: Tensor
) -> tuple[Tensor, Tensor]:
    """Give the shifted cutpoint interval of each observed level.

    ``scores.py`` reads these bounds to form the latent-scale derivative.

    Parameters
    ----------
    theta_tilde : Tensor
        Unconstrained cutpoint parameters, shape ``(n, levels - 1)``.
    shift : Tensor
        The node's shift, shape ``(n,)``.
    y : Tensor
        Observed levels ``0..levels-1``, shape ``(n,)``.

    Returns
    -------
    tuple[Tensor, Tensor]
        The lower and upper shifted cutpoint of each row's level, ``(n,)`` each.
    """
    cut = ordinal_cutpoints(theta_tilde) - shift.view(-1, 1)
    idx = torch.arange(theta_tilde.shape[0], device=theta_tilde.device)
    y = y.long()
    return cut[idx, y], cut[idx, y + 1]

ordinal_cutpoints ¶

ordinal_cutpoints(theta_tilde)

Constrain unconstrained parameters to increasing cutpoints.

The cutpoints are \(\tilde\vartheta_0\) followed by \(\tilde\vartheta_0 + \mathrm{cumsum}(\exp \tilde\vartheta_{1:})\), with \(\pm\infty\) at both ends.

Parameters:

Name	Type	Description	Default
`theta_tilde`	`Tensor`	Unconstrained cutpoint parameters, shape `(n, K-1)`.	required

Returns:

Type	Description
`Tensor`	Increasing cutpoints with `-inf`/`+inf` ends, shape `(n, K+1)`.

Notes

For K == 2, theta_tilde[:, 1:] is empty and its cumsum has shape (n, 0). This case needs no special treatment.

Source code in src/tramdag/transforms.py

def ordinal_cutpoints(theta_tilde: Tensor) -> Tensor:
    r"""Constrain unconstrained parameters to increasing cutpoints.

    The cutpoints are $\tilde\vartheta_0$ followed by
    $\tilde\vartheta_0 + \mathrm{cumsum}(\exp \tilde\vartheta_{1:})$, with
    $\pm\infty$ at both ends.

    Parameters
    ----------
    theta_tilde : Tensor
        Unconstrained cutpoint parameters, shape ``(n, K-1)``.

    Returns
    -------
    Tensor
        Increasing cutpoints with ``-inf``/``+inf`` ends, shape
        ``(n, K+1)``.

    Notes
    -----
    For ``K == 2``, ``theta_tilde[:, 1:]`` is empty and its ``cumsum`` has shape
    ``(n, 0)``. This case needs no special treatment.
    """
    n = theta_tilde.shape[0]
    neg_inf = torch.full(
        (n, 1), -torch.inf, device=theta_tilde.device, dtype=theta_tilde.dtype
    )
    pos_inf = torch.full(
        (n, 1), torch.inf, device=theta_tilde.device, dtype=theta_tilde.dtype
    )
    first = theta_tilde[:, :1]
    rest = first + torch.cumsum(torch.exp(theta_tilde[:, 1:]), dim=1)
    return torch.cat([neg_inf, first, rest, pos_inf], dim=1)

ordinal_marginal_init_theta ¶

ordinal_marginal_init_theta(counts)

Give the unconstrained cutpoint parameters that match class counts.

The marginal \(P(Y \le k) = \sigma(\vartheta_k)\) of the result matches the empirical class frequencies.

Parameters:

Name	Type	Description	Default
`counts`	`array - like`	Per-class count vector, length K. An empty class is handled by `_CDF_EPS`.	required

Returns:

Type	Description
`Tensor`	The unconstrained parameters `theta_tilde`, shape `(K-1,)`.

Notes

Inverts ordinal_cutpoints: the targets are \(c_k = \operatorname{logit} \hat F(k)\) from the empirical CDF, clamped off 0 and 1.

Source code in src/tramdag/transforms.py

def ordinal_marginal_init_theta(counts) -> Tensor:
    r"""Give the unconstrained cutpoint parameters that match class counts.

    The marginal $P(Y \le k) = \sigma(\vartheta_k)$ of the result matches
    the empirical class frequencies.

    Parameters
    ----------
    counts : array-like
        Per-class count vector, length K. An empty class is handled by
        ``_CDF_EPS``.

    Returns
    -------
    Tensor
        The unconstrained parameters ``theta_tilde``, shape ``(K-1,)``.

    Notes
    -----
    Inverts ``ordinal_cutpoints``: the targets are
    $c_k = \operatorname{logit} \hat F(k)$ from the empirical CDF, clamped off 0
    and 1.
    """
    counts = np.asarray(counts, dtype=np.float64)
    p = counts / counts.sum()
    F = np.clip(np.cumsum(p)[:-1], _CDF_EPS, 1 - _CDF_EPS)
    c = np.log(F) - np.log1p(-F)
    diffs = np.maximum(np.diff(c), _CDF_EPS)
    tt = np.empty_like(c)
    tt[0] = c[0]
    tt[1:] = np.log(diffs)
    return torch.as_tensor(tt)

ordinal_log_prob ¶

ordinal_log_prob(theta_tilde, shift, y)

Give \(\log P(Y = y \mid \vartheta, s)\).

The model is \(P(Y \le k) = \sigma(\vartheta_k - s)\).

Parameters:

Name	Type	Description	Default
`theta_tilde`	`Tensor`	Unconstrained cutpoint parameters, shape `(n, K-1)`.	required
`shift`	`Tensor`	Total shift, shape `(n,)`.	required
`y`	`Tensor`	Observed levels in `0..K-1`, shape `(n,)`.	required

Returns:

Type	Description
`Tensor`	The log probabilities, shape `(n,)`.

Notes

Computed in log-space as logsigmoid(b) + log1mexp(logsigmoid(a) - logsigmoid(b)), taking per element the better-conditioned side (the CDF or the survival side). Do not replace this with the difference of two sigmoids: that form has exactly zero gradient once the sigmoids saturate in float32, and a node stuck there never recovers.

The survival side is the CDF side of the negated, swapped bounds. One evaluation with flipped inputs replaces the computation of both sides and the selection between them.

Source code in src/tramdag/transforms.py

def ordinal_log_prob(theta_tilde: Tensor, shift: Tensor, y: Tensor) -> Tensor:
    r"""Give $\log P(Y = y \mid \vartheta, s)$.

    The model is $P(Y \le k) = \sigma(\vartheta_k - s)$.

    Parameters
    ----------
    theta_tilde : Tensor
        Unconstrained cutpoint parameters, shape ``(n, K-1)``.
    shift : Tensor
        Total shift, shape ``(n,)``.
    y : Tensor
        Observed levels in ``0..K-1``, shape ``(n,)``.

    Returns
    -------
    Tensor
        The log probabilities, shape ``(n,)``.

    Notes
    -----
    Computed in log-space as ``logsigmoid(b) + log1mexp(logsigmoid(a) -
    logsigmoid(b))``, taking per element the better-conditioned side (the CDF
    or the survival side). Do not replace this with the difference of two
    sigmoids: that form has exactly zero gradient once the sigmoids saturate
    in float32, and a node stuck there never recovers.

    The survival side is the CDF side of the negated, swapped bounds. One
    evaluation with flipped inputs replaces the computation of both sides and the
    selection between them.
    """
    lower, upper = ordinal_bounds(theta_tilde, shift, y)
    ls = torch.nn.functional.logsigmoid
    flip = upper + lower > 0
    a = torch.where(flip, -upper, lower)
    b = torch.where(flip, -lower, upper)
    return ls(b) + _log1mexp((ls(a) - ls(b)).clamp(max=-1e-7))

ordinal_pmf ¶

ordinal_pmf(theta_tilde, shift)

Give the class probabilities of every row.

Parameters:

Name	Type	Description	Default
`theta_tilde`	`Tensor`	Unconstrained cutpoint parameters, shape `(n, K-1)`.	required
`shift`	`Tensor`	Total shift, shape `(n,)`.	required

Returns:

Type	Description
`Tensor`	The class probabilities, shape `(n, K)`.

Source code in src/tramdag/transforms.py

def ordinal_pmf(theta_tilde: Tensor, shift: Tensor) -> Tensor:
    """Give the class probabilities of every row.

    Parameters
    ----------
    theta_tilde : Tensor
        Unconstrained cutpoint parameters, shape ``(n, K-1)``.
    shift : Tensor
        Total shift, shape ``(n,)``.

    Returns
    -------
    Tensor
        The class probabilities, shape ``(n, K)``.
    """
    cdf = torch.sigmoid(ordinal_cutpoints(theta_tilde) - shift.view(-1, 1))
    return cdf[:, 1:] - cdf[:, :-1]

ordinal_sample ¶

ordinal_sample(theta_tilde, shift, z)

Map latents to ordinal levels.

The rule is \(x = \#\{j : \vartheta_j - s < z\}\) over the finite cutpoints.

Parameters:

Name	Type	Description	Default
`theta_tilde`	`Tensor`	Unconstrained cutpoint parameters, shape `(n, K-1)`.	required
`shift`	`Tensor`	Total shift, shape `(n,)`.	required
`z`	`Tensor`	Latent values, shape `(n,)`.	required

Returns:

Type	Description
`Tensor`	The levels as floats, shape `(n,)`.

Source code in src/tramdag/transforms.py

def ordinal_sample(theta_tilde: Tensor, shift: Tensor, z: Tensor) -> Tensor:
    r"""Map latents to ordinal levels.

    The rule is $x = \#\{j : \vartheta_j - s < z\}$ over the finite cutpoints.

    Parameters
    ----------
    theta_tilde : Tensor
        Unconstrained cutpoint parameters, shape ``(n, K-1)``.
    shift : Tensor
        Total shift, shape ``(n,)``.
    z : Tensor
        Latent values, shape ``(n,)``.

    Returns
    -------
    Tensor
        The levels as floats, shape ``(n,)``.
    """
    finite = ordinal_cutpoints(theta_tilde)[:, 1:-1] - shift.view(-1, 1)
    return (z.view(-1, 1) > finite).sum(dim=1).to(shift.dtype)

ordinal_abduct ¶

ordinal_abduct(theta_tilde, shift, y, generator=None)

Abduct the latent of an ordinal node (the abduction step).

The latent is drawn from the standard logistic, truncated to the interval that is consistent with the observed level.

Parameters:

Name	Type	Description	Default
`theta_tilde`	`Tensor`	Unconstrained cutpoint parameters, shape `(n, K-1)`.	required
`shift`	`Tensor`	Total shift, shape `(n,)`.	required
`y`	`Tensor`	Observed levels in `0..K-1`, shape `(n,)`.	required
`generator`	`Generator \| None`	Random source for the truncated draw, by default `None`.	`None`

Returns:

Type	Description
`Tensor`	The latents, shape `(n,)`.

Source code in src/tramdag/transforms.py

def ordinal_abduct(
    theta_tilde: Tensor,
    shift: Tensor,
    y: Tensor,
    generator: torch.Generator | None = None,
) -> Tensor:
    """Abduct the latent of an ordinal node (the abduction step).

    The latent is drawn from the standard logistic, truncated to the
    interval that is consistent with the observed level.

    Parameters
    ----------
    theta_tilde : Tensor
        Unconstrained cutpoint parameters, shape ``(n, K-1)``.
    shift : Tensor
        Total shift, shape ``(n,)``.
    y : Tensor
        Observed levels in ``0..K-1``, shape ``(n,)``.
    generator : torch.Generator | None, optional
        Random source for the truncated draw, by default ``None``.

    Returns
    -------
    Tensor
        The latents, shape ``(n,)``.
    """
    lower, upper = ordinal_bounds(theta_tilde, shift, y)
    u_lo, u_hi = torch.sigmoid(lower), torch.sigmoid(upper)
    u = u_lo + (u_hi - u_lo) * torch.rand(
        lower.shape, device=lower.device, generator=generator
    )
    return StandardLogistic.icdf(u)

make_univariate_transform ¶

make_univariate_transform(name, **kwargs)

Build a scaled univariate transform by name.

Parameters:

Name	Type	Description	Default
`name`	`str`	One of `"bernstein"`, `"spline"`, `"affine"`.	required
`**kwargs`		Passed to the transform class.	`{}`

Returns:

Type	Description
`_ScaledUT`	The transform.

Raises:

Type	Description
`ValueError`	If `name` is not one of the three.

Source code in src/tramdag/transforms.py

def make_univariate_transform(name: str, **kwargs) -> _ScaledUT:
    """Build a scaled univariate transform by name.

    Parameters
    ----------
    name : str
        One of ``"bernstein"``, ``"spline"``, ``"affine"``.
    **kwargs
        Passed to the transform class.

    Returns
    -------
    _ScaledUT
        The transform.

    Raises
    ------
    ValueError
        If ``name`` is not one of the three.
    """
    if name not in _TRANSFORMS:
        raise ValueError(
            f"unknown transform {name!r}; choose one of {sorted(_TRANSFORMS)}"
        )
    return _TRANSFORMS[name](**kwargs)

tramdag.transforms¶