Multivariate

Distributions over vectors, simplices, positive-definite matrices, or directions on the sphere.

MultivariateNormal(loc, scale_tril=None, *, cov=None, name)

Bases: TFPDistribution, FlatNumericRecordDistribution

Multivariate normal (Gaussian) distribution.

Parameters:

Name	Type	Description	Default
loc	array-like, shape ``(d,)``	Mean vector.	required
scale_tril	array-like, shape ``(d, d)``	Lower-triangular Cholesky factor of the covariance. Exactly one of scale_tril or cov must be provided.	None
cov	array-like, shape ``(d, d)``	Covariance matrix (Cholesky-decomposed internally).	None
name	str	Distribution name.	required

Source code in probpipe/distributions/multivariate.py

def __init__(
    self,
    loc: ArrayLike,
    scale_tril: ArrayLike | None = None,
    *,
    cov: ArrayLike | None = None,
    name: str,
):
    if scale_tril is not None and cov is not None:
        raise ValueError("Provide exactly one of scale_tril or cov, not both.")

    if scale_tril is not None:
        _, (loc, scale_tril) = _promote_floats(loc, scale_tril)
    elif cov is not None:
        _, (loc, cov) = _promote_floats(loc, cov)
    else:
        raise ValueError("One of scale_tril or cov must be provided.")

    if loc.ndim == 0:
        loc = loc.reshape(1)

    if cov is not None:
        if cov.shape != (loc.shape[0], loc.shape[0]):
            raise ValueError(
                f"cov shape {cov.shape} does not match loc length {loc.shape[0]}."
            )
        scale_tril = jnp.linalg.cholesky(cov)

    self._loc = loc
    self._scale_tril = scale_tril
    self._tfp_dist = tfd.MultivariateNormalTriL(
        loc=loc, scale_tril=scale_tril
    )
    super().__init__(name=name)

cov property

Full covariance matrix (computed from Cholesky factor).

Dirichlet(concentration, *, name)

Bases: TFPDistribution, FlatNumericRecordDistribution

Dirichlet distribution over the probability simplex.

Parameters:

Name	Type	Description	Default
concentration	array-like, shape ``(k,)``	Positive concentration (alpha) parameters.	required
name	str	Distribution name.	required

Source code in probpipe/distributions/multivariate.py

def __init__(
    self,
    concentration: ArrayLike,
    *,
    name: str,
):
    concentration = _as_float_array(concentration)
    if concentration.ndim == 0:
        raise ValueError("concentration must be at least 1-D.")

    self._concentration = concentration
    self._tfp_dist = tfd.Dirichlet(concentration=concentration)
    super().__init__(name=name)

Multinomial(total_count, probs=None, logits=None, *, name)

Bases: TFPDistribution, FlatNumericRecordDistribution

Multinomial distribution over count vectors.

Exactly one of probs or logits must be provided.

Parameters:

Name	Type	Description	Default
total_count	int or array - like	Number of trials.	required
probs	array-like, shape ``(k,)``	Event probabilities (need not be normalised).	None
logits	array-like, shape ``(k,)``	Log-odds of each event.	None
name	str	Distribution name.	required

Source code in probpipe/distributions/multivariate.py

def __init__(
    self,
    total_count: int | ArrayLike,
    probs: ArrayLike | None = None,
    logits: ArrayLike | None = None,
    *,
    name: str,
):
    if (probs is None) == (logits is None):
        raise ValueError("Exactly one of probs or logits must be provided.")

    if probs is not None:
        _, (total_count, probs) = _promote_floats(total_count, probs)
        self._probs = probs
        self._logits = None
        self._tfp_dist = tfd.Multinomial(
            total_count=total_count, probs=probs
        )
    else:
        _, (total_count, logits) = _promote_floats(total_count, logits)
        self._logits = logits
        self._probs = None
        self._tfp_dist = tfd.Multinomial(
            total_count=total_count, logits=logits
        )

    self._total_count = total_count
    super().__init__(name=name)

Wishart(df, scale_tril=None, *, scale=None, name)

Bases: TFPDistribution

Wishart distribution over positive-definite matrices.

Exactly one of scale_tril or scale must be provided.

Parameters:

Name	Type	Description	Default
df	float or array - like	Degrees of freedom (must be >= dimension).	required
scale_tril	array-like, shape ``(d, d)``	Lower-triangular Cholesky factor of the scale matrix.	None
scale	array-like, shape ``(d, d)``	Full scale matrix (Cholesky-decomposed internally).	None
name	str	Distribution name.	required

Source code in probpipe/distributions/multivariate.py

def __init__(
    self,
    df: float | ArrayLike,
    scale_tril: ArrayLike | None = None,
    *,
    scale: ArrayLike | None = None,
    name: str,
):
    if scale_tril is not None and scale is not None:
        raise ValueError(
            "Provide exactly one of scale_tril or scale, not both."
        )
    if scale_tril is None and scale is None:
        raise ValueError("One of scale_tril or scale must be provided.")

    if scale is not None:
        _, (df, scale) = _promote_floats(df, scale)
        scale_tril = jnp.linalg.cholesky(scale)
    else:
        _, (df, scale_tril) = _promote_floats(df, scale_tril)

    self._df = df
    self._scale_tril = scale_tril
    self._tfp_dist = tfd.WishartTriL(df=df, scale_tril=scale_tril)
    super().__init__(name=name)

scale property

Full scale matrix (computed from Cholesky factor).

VonMisesFisher(mean_direction, concentration, *, name)

Bases: TFPDistribution, FlatNumericRecordDistribution

Von Mises-Fisher distribution on the unit hypersphere.

Parameters:

Name	Type	Description	Default
mean_direction	array-like, shape ``(d,)``	Unit vector giving the mean direction.	required
concentration	float or array - like	Scalar concentration parameter (kappa >= 0).	required
name	str	Distribution name.	required

Source code in probpipe/distributions/multivariate.py

def __init__(
    self,
    mean_direction: ArrayLike,
    concentration: float | ArrayLike,
    *,
    name: str,
):
    _, (mean_direction, concentration) = _promote_floats(
        mean_direction, concentration
    )

    self._mean_direction = mean_direction
    self._concentration = concentration
    self._tfp_dist = tfd.VonMisesFisher(
        mean_direction=mean_direction, concentration=concentration
    )
    super().__init__(name=name)