leaspy.variables.distributions

This module defines the distributions used for sampling variables.

Attributes

Normal
Bernoulli
WeibullRightCensored
WeibullRightCensoredWithSources

Classes

StatelessDistributionFamily	Abstract base class defining a stateless interface for distribution families.
StatelessDistributionFamilyFromTorchDistribution	Wrapper to build a StatelessDistributionFamily class from an existing torch distribution class.
BernoulliFamily	Bernoulli family (stateless).
NormalFamily	Normal / Gaussian family (stateless).
MultivariateNormalFamily	Multivariate Normal family with diagonal covariance (stateless).
MixtureNormalFamily	A stateless mixture of univariate normal distributions.
AbstractWeibullRightCensoredFamily	Abstract base class defining a stateless interface for distribution families.
WeibullRightCensoredFamily	Abstract base class defining a stateless interface for distribution families.
WeibullRightCensoredWithSourcesFamily	Abstract base class defining a stateless interface for distribution families.
SymbolicDistribution	Class providing symbolic methods for distribution families.

Module Contents

class StatelessDistributionFamily

Bases: abc.ABC

Abstract base class defining a stateless interface for distribution families.

This class represents a family of probability distributions in a stateless manner: no parameters are stored in the instance. All methods operate purely via classmethods, using explicitly passed distribution parameters.

Notes

• Subclasses must define the parameters class variable, listing parameter names in order.

• Each method operates solely on the passed tensors; no state or caching is assumed.

parameters: ClassVar[tuple[str, Ellipsis]]

classmethod validate_parameters(*params)

Abstractmethod:

Parameters:

params (Any)

Return type:

tuple[Tensor, Ellipsis]

Validate consistency of distribution parameters, returning them with out-of-place modifications if needed.

classmethod shape(*params_shapes)

Shape of distribution samples (without any additional expansion), given shapes of distribution parameters.

Parameters:

*params_shapestuple of int

The shapes of the distribution parameters, passed in the order defined by cls.parameters.

Returns:

tuple of int

The shape of the distribution samples.

Raises:

LeaspyInputError

If the number of provided shapes does not match the expected number of parameters.

NotImplementedError

If the distribution has no parameters, and sample shape cannot be inferred.

Parameters:

params_shapes (tuple[int, Ellipsis])

Return type:

tuple[int, Ellipsis]

classmethod sample(*params, sample_shape=())

Abstractmethod:

Parameters:

• params (Tensor)

• sample_shape (tuple[int, Ellipsis])

Return type:

Tensor

Sample values, given distribution parameters (sample_shape is prepended to shape of distribution parameters).

classmethod mode(*params)

Abstractmethod:

Parameters:

params (Tensor)

Return type:

Tensor

Mode of distribution (returning first value if discrete ties), given distribution parameters.

classmethod mean(*params)

Abstractmethod:

Parameters:

params (Tensor)

Return type:

Tensor

Mean of distribution (if defined), given distribution parameters.

classmethod stddev(*params)

Abstractmethod:

Parameters:

params (Tensor)

Return type:

Tensor

Standard-deviation of distribution (if defined), given distribution parameters.

classmethod nll(x, *params)

Negative log-likelihood of value, given distribution parameters.

Parameters:

• x (WeightedTensor[float])

• params (Tensor)

Return type:

WeightedTensor[float]

classmethod regularization(x, *params)

Negative log-likelihood of value, given distribution parameters.

Parameters:

• x (Tensor)

• params (Tensor)

Return type:

WeightedTensor[float]

classmethod nll_jacobian(x, *params)

Jacobian w.r.t. value of negative log-likelihood, given distribution parameters.

Parameters:

• x (WeightedTensor[float])

• params (Tensor)

Return type:

WeightedTensor[float]

classmethod nll_and_jacobian(x, *params)

Negative log-likelihood of value and its jacobian w.r.t. value, given distribution parameters.

Parameters:

• x (WeightedTensor[float])

• params (Tensor)

Return type:

tuple[WeightedTensor[float], WeightedTensor[float]]

class StatelessDistributionFamilyFromTorchDistribution

Bases: StatelessDistributionFamily

Wrapper to build a StatelessDistributionFamily class from an existing torch distribution class.

Attributes:

dist_factoryCallable […, torch.distributions.Distribution]

A class variable that points to a factory function or class used to instantiate the corresponding PyTorch distribution.

dist_factory: ClassVar[Callable[Ellipsis, torch.distributions.Distribution]]

classmethod validate_parameters(*params)

Validate consistency of distribution parameters, returning them with out-of-place modifications if needed.

Parameters:

paramsAny

The parameters to pass to the distribution factory.

Returns:

tuple [ torch.Tensor, …]

The validated parameters.

Parameters:

params (Any)

Return type:

tuple[Tensor, Ellipsis]

classmethod sample(*params, sample_shape=())

Sample values, given distribution parameters (sample_shape is prepended to shape of distribution parameters).

Parameters:

paramstorch.Tensor

The distribution parameters.

sample_shapetuple of int, optional

Desired shape of the sample batch (prepended to the shape of a single sample). Defaults to an empty tuple, i.e., a single sample.

Returns:

torch.Tensor

Sampled tensor

Parameters:

• params (Tensor)

• sample_shape (tuple[int, Ellipsis])

Return type:

Tensor

classmethod mode(*params)

Abstractmethod:

Parameters:

params (Tensor)

Return type:

Tensor

Mode of distribution (returning first value if discrete ties), given distribution parameters.

This method should be overridden in subclasses that wrap torch distributions which explicitly define a mode.

Parameters:

paramstorch.Tensor

The distribution parameters.

Returns:

torch.Tensor

The value of the distribution’s mode.

Raises:

NotImplementedError

If the mode is not explicitly implemented for the specific distribution subclass.

Parameters:

params (Tensor)

Return type:

Tensor

classmethod mean(*params)

Return the mean of the distribution, if defined.

Parameters:

paramstorch.Tensor

The distribution parameters.

Returns:

torch.Tensor

The value of the distribution’s mean.

Parameters:

params (Tensor)

Return type:

Tensor

classmethod stddev(*params)

Return the standard-deviation of the distribution.

Parameters:

paramstorch.Tensor

The distribution parameters.

Returns:

torch.Tensor

The value of the distribution’s standard deviation.

Parameters:

params (Tensor)

Return type:

Tensor

class BernoulliFamily

Bases: StatelessDistributionFamilyFromTorchDistribution

Bernoulli family (stateless).

Inherits from StatelessDistributionFamilyFromTorchDistribution.

parameters: ClassVar = ('loc',)

dist_factory: ClassVar

class NormalFamily

Bases: StatelessDistributionFamilyFromTorchDistribution

Normal / Gaussian family (stateless).

Inherits from StatelessDistributionFamilyFromTorchDistribution.

parameters: ClassVar = ('loc', 'scale')

dist_factory: ClassVar

nll_constant_standard: ClassVar

classmethod mode(loc, scale)

Return the mode of the distribution given the distribution’s loc and scale parameters.

Parameters:

loctorch.Tensor

The distribution loc.

scaletorch.Tensor

The distribution scale.

Returns:

torch.Tensor

The value of the distribution’s mode.

Parameters:

• loc (Tensor)

• scale (Tensor)

Return type:

Tensor

classmethod mean(loc, scale)

Return the mean of the distribution, given the distribution loc and scale parameters.

Parameters:

loctorch.Tensor

The distribution loc parameters.

scaletorch.Tensor

The distribution scale parameters.

Returns:

torch.Tensor

The value of the distribution’s mean.

Parameters:

• loc (Tensor)

• scale (Tensor)

Return type:

Tensor

classmethod stddev(loc, scale)

Return the standard-deviation of the distribution, given loc and scale of the distribution.

Parameters:

loctorch.Tensor

The distribution loc parameter.

scaletorch.Tensor

The distribution scale parameter.

Returns:

torch.Tensor

The value of the distribution’s standard deviation.

Parameters:

• loc (Tensor)

• scale (Tensor)

Return type:

Tensor

class MultivariateNormalFamily

Bases: StatelessDistributionFamily

Multivariate Normal family with diagonal covariance (stateless).

parameters: ClassVar = ('loc', 'scale')

dist_factory: ClassVar

nll_constant_standard: ClassVar

classmethod multi_dist_factory(loc, scale)

Parameters:

• loc (Tensor)

• scale (Tensor)

Return type:

MultivariateNormal

classmethod mode(loc, scale)

Mode of distribution (returning first value if discrete ties), given distribution parameters.

Parameters:

• loc (Tensor)

• scale (Tensor)

Return type:

Tensor

classmethod mean(loc, scale)

Mean of distribution (if defined), given distribution parameters.

Parameters:

• loc (Tensor)

• scale (Tensor)

Return type:

Tensor

classmethod stddev(loc, scale)

Standard-deviation of distribution (if defined), given distribution parameters.

Parameters:

• loc (Tensor)

• scale (Tensor)

Return type:

Tensor

class MixtureNormalFamily

Bases: StatelessDistributionFamily

A stateless mixture of univariate normal distributions.

This class defines a mixture distribution where each component is a univariate normal distribution, and the mixture weights are defined by a categorical distribution.

Parameters:

• loc (Tensor): Mean of each normal component, one for each cluster.

• scale (Tensor): Standard deviation of each normal component.

• probs (Tensor): Probabilities associated with each cluster; must sum to 1.

The mixture is modeled using torch.distributions.MixtureSameFamily, where:

• The mixing distribution is a Categorical defined by probs.

• The component distribution is Normal(loc, scale).

parameters: ClassVar = ('loc', 'scale', 'probs')

nll_constant_standard: ClassVar

classmethod dist_factory(loc, scale, probs)

Construct a MixtureSameFamily distribution of univariate normal components.

Parameters:

loctorch.Tensor

Mean(s) of the normal components. Shape should be broadcastable with scale and probs.

scaletorch.Tensor

Standard deviation(s) of the normal components.

probstorch.Tensor

Probabilities of each mixture component. Must sum to 1 along the component axis.

Returns:

class:

.MixtureSameFamily A mixture distribution with categorical mixing and normal components.

Parameters:

• loc (Tensor)

• scale (Tensor)

• probs (Tensor)

Return type:

MixtureSameFamily

classmethod sample(*params, sample_shape=())

Draw samples from the mixture of normal distributions.

Parameters:

*paramstorch.Tensor

Distribution parameters in the order (loc, scale, probs). These should be broadcastable to define a valid MixtureSameFamily distribution.

sample_shapetuple of int, optional

The desired sample shape. Defaults to an empty tuple for a single sample.

Returns:

torch.Tensor

A tensor of samples drawn from the specified mixture distribution. The shape is sample_shape + batch_shape.

Parameters:

• params (Tensor)

• sample_shape (tuple[int, Ellipsis])

Return type:

Tensor

classmethod set_component_distribution(component_distribution, loc, scale)

Ensure that the component distribution is an instance of the torch.distributions.Normal.

Parameters:

component_distributiontorch.distributions.Distribution

The distribution object to validate. Must be an instance of torch.distributions.Normal.

loctorch.Tensor

Mean(s) for the normal distribution.

scaletorch.Tensor

Standard deviation(s) for the normal distribution.

Returns:

torch.distributions.Distribution

The newly set Normal distribution instance.

Raises:

ValueError

If component_distribution is not an instance of torch.distributions.Normal.

Parameters:

• component_distribution (torch.distributions)

• loc (Tensor)

• scale (Tensor)

Return type:

torch.distributions

classmethod extract_probs(*params)

Extract the mixture probabilities from the distribution parameters.

Parameters:

*paramsAny

Distribution parameters (expected: loc, scale, probs), passed to dist_factory.

Returns:

torch.Tensor

A 1D tensor of probabilities for each component in the mixture.

Parameters:

params (Any)

Return type:

Tensor

classmethod extract_n_clusters(*params)

Return the number of mixture components (i.e., clusters) in the distribution.

Parameters:

*paramsAny

Distribution parameters (expected: loc, scale, probs), passed to extract_probs.

Returns:

int

The number of clusters/components in the mixture distribution.

Parameters:

params (Any)

Return type:

int

classmethod extract_cluster_parameters(which_cluster)

Extract the parameters (probability, mean, standard deviation) for a specific cluster.

Parameters:

which_clusterint

The index of the cluster to extract parameters for.

Returns:

tuple of torch.Tensor

A tuple containing: - prob: Probability of the selected cluster. - loc: Mean of the selected normal component. - scale: Standard deviation of the selected normal component.

Parameters:

which_cluster (int)

Return type:

tuple[Tensor, Tensor, Tensor]

classmethod mode(*params)

Return the mode of the mixture distribution.

For a mixture of normal distributions, this returns the mean of the overall mixture distribution, which serves as a proxy for the mode.

Parameters:

*paramsAny

Distribution parameters (expected: loc, scale, probs), passed to dist_factory.

Returns:

torch.Tensor

The mode of the mixture distribution.

Parameters:

params (Any)

Return type:

Tensor

classmethod mean(*params)

Return the mean of the mixture distribution.

For a mixture of normal distributions, this returns the mean of the overall mixture distribution.

Parameters:

*paramsAny

Distribution parameters (expected: loc, scale, probs), passed to dist_factory.

Returns:

torch.Tensor

The mean of the mixture distribution.

Parameters:

params (Any)

Return type:

Tensor

classmethod stddev(*params)

Return the standard deviation of the mixture distribution.

For a mixture of normal distributions, this returns the standard deviation of the overall mixture distribution.

Parameters:

*paramsAny

Distribution parameters (expected: loc, scale, probs), passed to dist_factory.

Returns:

torch.Tensor

The standard deviation of the mixture distribution.

Parameters:

params (Any)

Return type:

Tensor

class AbstractWeibullRightCensoredFamily

Bases: StatelessDistributionFamily

Abstract base class defining a stateless interface for distribution families.

This class represents a family of probability distributions in a stateless manner: no parameters are stored in the instance. All methods operate purely via classmethods, using explicitly passed distribution parameters.

Notes

• Subclasses must define the parameters class variable, listing parameter names in order.

• Each method operates solely on the passed tensors; no state or caching is assumed.

dist_weibull: ClassVar

precision: float = 0.0001

classmethod validate_parameters(*params)

Abstractmethod:

Parameters:

params (Any)

Return type:

tuple[Tensor, Ellipsis]

Validate consistency of distribution parameters, returning them with out-of-place modifications if needed.

Parameters:

paramsAny

The parameters to pass to the distribution factory.

Returns:

tuple [ torch.Tensor, …]

The validated parameters.

Parameters:

params (Any)

Return type:

tuple[Tensor, Ellipsis]

classmethod sample(nu, rho, xi, tau, sample_shape=())

Sample values from a Weibull distribution.

Parameters:

nutorch.Tensor

rhotorch.Tensor

xitorch.Tensor

tautorch.Tensor

sample_shapetuple of int, optional

Shape of the samples to draw

Returns:

torch.Tensor

Samples drawn from the transformed Weibull distribution, shaped according to sample_shape combined with the distribution parameter shapes.

Parameters:

• nu (Tensor)

• rho (Tensor)

• xi (Tensor)

• tau (Tensor)

• sample_shape (tuple[int, Ellipsis])

Return type:

Tensor

classmethod mode(*params)

Abstractmethod:

Parameters:

params (Tensor)

Return type:

Tensor

Return the mode of the distribution (returning first value if discrete ties).

Parameters:

paramstorch.Tensor

The distribution parameters.

Returns:

torch.Tensor

The value of the distribution’s mode.

Parameters:

params (Tensor)

Return type:

Tensor

classmethod mean(nu, rho, xi, tau)

Return the mean of the distribution, if defined.

Parameters:

nutorch.Tensor

rhotorch.Tensor

xitorch.Tensor

tautorch.Tensor

Returns:

torch.Tensor

The value of the distribution’s mean.

Parameters:

• nu (Tensor)

• rho (Tensor)

• xi (Tensor)

• tau (Tensor)

Return type:

Tensor

classmethod stddev(nu, rho, xi, tau)

Return the standard-deviation of the distribution, given distribution parameters.

Parameters:

nutorch.Tensor

rhotorch.Tensor

xitorch.Tensor

tautorch.Tensor

Returns:

torch.Tensor

The value of the distribution’s standard deviation.

Parameters:

• nu (Tensor)

• rho (Tensor)

• xi (Tensor)

• tau (Tensor)

Return type:

Tensor

classmethod compute_log_likelihood_hazard(x, nu, rho, xi, tau, *params)

Compute the log hazard component of the likelihood for given data and parameters.

Parameters:

xWeightedTensor

The observed data values with associated weights.

nutorch.Tensor

rhotorch.Tensor

xitorch.Tensor

tautorch.Tensor

params :class:`torch.Tensor`

Additional parameters for reparametrization if needed.

Returns:

torch.Tensor

The log hazard values for each observation, zeroed out for censored data.

Parameters:

• x (WeightedTensor)

• nu (Tensor)

• rho (Tensor)

• xi (Tensor)

• tau (Tensor)

• params (Tensor)

Return type:

Tensor

classmethod compute_hazard(x, nu, rho, xi, tau, *params)

Compute the hazard function values for given observations and parameters.

Parameters:

xWeightedTensor

The observed data values with associated weights.

nutorch.Tensor

rhotorch.Tensor

xitorch.Tensor

tautorch.Tensor

params :class:`torch.Tensor`

Additional parameters for reparametrization if needed.

Returns:

torch.Tensor

Hazard values computed for each observation.

Parameters:

• x (WeightedTensor)

• nu (Tensor)

• rho (Tensor)

• xi (Tensor)

• tau (Tensor)

• params (Tensor)

Return type:

Tensor

classmethod compute_log_survival(x, nu, rho, xi, tau, *params)

Compute the log survival function for the Weibull distribution given observations and parameters.

Parameters:

xtorch.Tensor

nutorch.Tensor

rhotorch.Tensor

xitorch.Tensor

tau ::class:`torch.Tensor`

paramstorch.Tensor

Additional optional parameters used in reparametrization.

Returns:

torch.Tensor

Log survival values for each observation.

Parameters:

• x (Tensor)

• nu (Tensor)

• rho (Tensor)

• xi (Tensor)

• tau (Tensor)

• params (Tensor)

Return type:

Tensor

classmethod compute_predictions(x, nu, rho, xi, tau, *params)

Compute predicted survival or cumulative incidence probabilities for time-to-event data using a reparametrized Weibull model.

Parameters:

xtorch.Tensor

nutorch.Tensor

rhotorch.Tensor

xitorch.Tensor

tau ::class:`torch.Tensor`

paramstorch.Tensor

Additional optional parameters used in reparametrization.

Returns:

torch.Tensor

Predicted survival probabilities or cumulative incidence values (depending on event type count), normalized by baseline survival at the last visit.

Parameters:

• x (Tensor)

• nu (Tensor)

• rho (Tensor)

• xi (Tensor)

• tau (Tensor)

• params (Tensor)

Return type:

Tensor

class WeibullRightCensoredFamily

Bases: AbstractWeibullRightCensoredFamily

Abstract base class defining a stateless interface for distribution families.

This class represents a family of probability distributions in a stateless manner: no parameters are stored in the instance. All methods operate purely via classmethods, using explicitly passed distribution parameters.

Notes

• Subclasses must define the parameters class variable, listing parameter names in order.

• Each method operates solely on the passed tensors; no state or caching is assumed.

parameters: ClassVar = ('nu', 'rho', 'xi', 'tau')

class WeibullRightCensoredWithSourcesFamily

Bases: AbstractWeibullRightCensoredFamily

Abstract base class defining a stateless interface for distribution families.

This class represents a family of probability distributions in a stateless manner: no parameters are stored in the instance. All methods operate purely via classmethods, using explicitly passed distribution parameters.

Notes

• Subclasses must define the parameters class variable, listing parameter names in order.

• Each method operates solely on the passed tensors; no state or caching is assumed.

parameters: ClassVar = ('nu', 'rho', 'xi', 'tau', 'survival_shifts')

class SymbolicDistribution

Class providing symbolic methods for distribution families.

parameters_names: tuple[str, Ellipsis]

dist_family: Type[StatelessDistributionFamily]

validate_parameters: Callable[Ellipsis, tuple[Tensor, Ellipsis]]

Function of named distribution parameters, to validate these parameters.

shape: Callable[Ellipsis, tuple[int, Ellipsis]]

Function of named shapes of distribution parameters, to get shape of distribution samples.

mode: Callable[Ellipsis, Tensor]

Function of named distribution parameters, to get mode of distribution.

mean: Callable[Ellipsis, Tensor]

Function of named distribution parameters, to get mean of distribution.

stddev: Callable[Ellipsis, Tensor]

Function of named distribution parameters, to get std-deviation of distribution.

get_func(func, *extra_args_names, **kws)

Retrieve a function (e.g., ‘sample’, ‘mode’, ‘mean’) from the associated stateless distribution family, wrapped as a NamedInputFunction.

Parameters:

funcstr

Name of the method to retrieve from the stateless distribution family.

*extra_args_namesstr

Additional parameter names (e.g., ‘sample_shape’) to include before standard parameters.

**kwsdict

Optional keyword arguments passed to the NamedInputFunction.

Returns:

NamedInputFunction

A callable wrapper over the distribution method with named inputs.

Parameters:

• func (str)

• extra_args_names (str)

get_func_sample(sample_shape=())

Factory of symbolic sampling function.

Parameters:

sample_shapetuple of int, optional

The shape of the sample. Default=().

Returns:

NamedInputFunction

The sample function.

Parameters:

sample_shape (tuple[int, Ellipsis])

Return type:

NamedInputFunction[Tensor]

get_func_sample_multivariate(sample_shape=())

Factory of symbolic sampling function.

Parameters:

sample_shapetuple of int, optional

The shape of the sample. Default=().

Returns:

NamedInputFunction

The sample function.

Parameters:

sample_shape (tuple[int, Ellipsis])

Return type:

NamedInputFunction[Tensor]

get_func_regularization(value_name)

Factory method to return a symbolic function computing the negative log-likelihood (used for regularization) from a given value.

Parameters:

value_namestr

Returns:

WeightedTensor]

The named input function to use to compute negative log likelihood.

Parameters:

value_name (str)

Return type:

NamedInputFunction[WeightedTensor[float]]

get_func_nll(value_name)

Factory method to return a symbolic function computing the negative log-likelihood (NLL) from a given value.

Parameters:

value_namestr

Returns:

WeightedTensor]

The named input function to use to compute negative log likelihood.

Parameters:

value_name (str)

Return type:

NamedInputFunction[WeightedTensor[float]]

get_func_nll_jacobian(value_name)

Factory of symbolic function: state -> jacobian w.r.t. value of negative log-likelihood.

Parameters:

value_namestr

Returns:

NamedInputFunction

The named input function to use to compute negative log likelihood jacobian.

Parameters:

value_name (str)

Return type:

NamedInputFunction[WeightedTensor[float]]

get_func_nll_and_jacobian(value_name)

Factory of symbolic function: state -> (negative log-likelihood, its jacobian w.r.t. value).

Parameters:

value_namestr

Returns:

tuple`[:class:`~leaspy.utils.functional._named_input_function.NamedInputFunction`[:class:`~leaspy.utils.weighted_tensor._weighted_tensor.WeightedTensor]]

The named input functions to use to compute negative log likelihood and its jacobian.

Parameters:

value_name (str)

Return type:

NamedInputFunction[tuple[WeightedTensor[float], WeightedTensor[float]]]

classmethod bound_to(dist_family)

Return a factory to create SymbolicDistribution bound to the provided distribution family.

Parameters:

dist_familyStatelessDistributionFamily

The distribution family to use to create a SymbolicDistribution.

Returns:

factoryCallable`[..., :class:`~leaspy.variables.distributions.SymbolicDistribution]

The factory.

Parameters:

dist_family (Type[StatelessDistributionFamily])

Return type:

Callable[Ellipsis, SymbolicDistribution]

Normal

Bernoulli

WeibullRightCensored

WeibullRightCensoredWithSources