`causica.distributions.noise`¶

causica.distributions.noise.spline

Module Contents¶

Package Contents¶

Classes¶

`BernoulliNoise`	Creates a Bernoulli distribution parameterized by `probs`
`BernoulliNoiseModule`	Represents a BernoulliNoise distribution with learnable logits.
`CategoricalNoise`	Creates a one-hot categorical distribution parameterized by `probs` or
`CategoricalNoiseModule`	Represents a CategoricalNoise distribution with learnable logits.
`ContinuousNoiseDist`	Generic enumeration.
`JointNoise`	Represents an independent joint noise distribution of multiple variables types.
`JointNoiseModule`	Represents JointNoise with learnable parameters.
`IndependentNoise`	Like td.Idenpendent but also forwards Noise specific methods.
`Noise`	Extend Distributions to allow the noise (usually unparametrized) to be extracted from samples and vice versa.
`NoiseModule`	Module for producing Noise, where sampling from the distribution corresponds to adding noise to the input.
`SplineNoise`	A Spline Based Noise Distribution.
`SplineNoiseModule`	Implements Neural Spline Flow noise with learnable parameters.
`UnivariateCauchyNoise`	Samples from a Cauchy (Lorentz) distribution. The distribution of the ratio of
`UnivariateCauchyNoiseModule`	Represents a UnivariateCauchyNoise with learnable parameters for independent variables.
`UnivariateLaplaceNoise`	Creates a Laplace distribution parameterized by `loc` and `scale`.
`UnivariateLaplaceNoiseModule`	Represents a UnivariateLaplaceNoise with learnable parameters for independent variables.
`UnivariateNormalNoise`	Creates a normal (also called Gaussian) distribution parameterized by
`UnivariateNormalNoiseModule`	Represents a UnivariateNormalNoise with learnable parameters for independent variables.

Functions¶

`create_noise_modules`(→ dict[str, ...)	Create noise modules for each item of shapes and types.
`create_spline_dist_params`(→ list[SplineParams])	Create initial values for a spline distribution.

class causica.distributions.noise.BernoulliNoise(delta_logits: torch.Tensor, base_logits: torch.Tensor)[source]¶

Bases: torch.distributions.Bernoulli, causica.distributions.noise.noise.Noise

Creates a Bernoulli distribution parameterized by probs or logits (but not both).

Samples are binary (0 or 1). They take the value 1 with probability p and 0 with probability 1 - p.

Example:

>>> # xdoctest: +IGNORE_WANT("non-deterinistic")
>>> m = Bernoulli(torch.tensor([0.3]))
>>> m.sample()  # 30% chance 1; 70% chance 0
tensor([ 0.])

Parameters:¶

probs : Number, Tensor: the probability of sampling 1
logits : Number, Tensor: the log-odds of sampling 1

property mode¶

Override the default mode method to prevent it returning nan’s.

We favour sparseness, so if logit == 0, set the mode to be zero.

sample_to_noise(samples: torch.Tensor) → torch.Tensor[source]¶

Transform from the sample observations to corresponding noise variables.

This will draw from the noise posterior given the observations

A posterior sample of the Gumbel noise random variables given observation x and probabilities self.base_logits + logit_deltas.

This methodology is described in https://arxiv.org/pdf/1905.05824.pdf. See https://cmaddis.github.io/gumbel-machinery for derivation of Gumbel posteriors. For a derivation of this exact algorithm using softplus, see https://www.overleaf.com/8628339373sxjmtvyxcqnx.

Parameters:¶

samples: torch.Tensor¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: torch.Tensor) → torch.Tensor[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: torch.Tensor¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.BernoulliNoiseModule(dim: int, init_base_logits: float | torch.Tensor = 0.0)[source]¶

Bases: causica.distributions.noise.noise.NoiseModule[causica.distributions.noise.noise.IndependentNoise[BernoulliNoise]]

Represents a BernoulliNoise distribution with learnable logits.

forward(x: torch.Tensor | None = None) → causica.distributions.noise.noise.IndependentNoise[BernoulliNoise][source]¶

class causica.distributions.noise.CategoricalNoise(delta_logits: torch.Tensor, base_logits: torch.Tensor)[source]¶

Bases: torch.distributions.OneHotCategorical, causica.distributions.noise.noise.Noise

Creates a one-hot categorical distribution parameterized by probs or logits.

Samples are one-hot coded vectors of size probs.size(-1).

Note

The probs argument must be non-negative, finite and have a non-zero sum, and it will be normalized to sum to 1 along the last dimension. probs will return this normalized value. The logits argument will be interpreted as unnormalized log probabilities and can therefore be any real number. It will likewise be normalized so that the resulting probabilities sum to 1 along the last dimension. logits will return this normalized value.

See also: torch.distributions.Categorical() for specifications of probs and logits.

Example:

>>> # xdoctest: +IGNORE_WANT("non-deterinistic")
>>> m = OneHotCategorical(torch.tensor([ 0.25, 0.25, 0.25, 0.25 ]))
>>> m.sample()  # equal probability of 0, 1, 2, 3
tensor([ 0.,  0.,  0.,  1.])

Parameters:¶

probs : Tensor: event probabilities
logits : Tensor: event log probabilities (unnormalized)

sample_to_noise(samples: torch.Tensor) → torch.Tensor[source]¶

Transform from the sample observations to corresponding noise variables.

This will draw from the noise posterior given the observations

A posterior sample of the Gumbel noise random variables given observation x and probabilities self.base_logits + logit_deltas.

Parameters:¶

samples: torch.Tensor¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: torch.Tensor) → torch.Tensor[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: torch.Tensor¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.CategoricalNoiseModule(num_classes: int, init_base_logits: torch.Tensor | None = None)[source]¶

Bases: causica.distributions.noise.noise.NoiseModule[CategoricalNoise]

Represents a CategoricalNoise distribution with learnable logits.

forward(x: torch.Tensor | None = None) → CategoricalNoise[source]¶

class causica.distributions.noise.ContinuousNoiseDist[source]¶

Bases: enum.Enum

Generic enumeration.

Derive from this class to define new enumerations.

SPLINE = 'spline'¶

GAUSSIAN = 'gaussian'¶

class causica.distributions.noise.JointNoise(independent_noise_dists: Mapping[str, causica.distributions.noise.noise.Noise[torch.Tensor]])[source]¶

Bases: causica.distributions.noise.noise.Noise[tensordict.TensorDict]

Represents an independent joint noise distribution of multiple variables types.

Samples are TensorDicts containining independent variables.

property support : dict[str, Any | None]¶: Returns a Constraint object representing this distribution’s support.

property mode : tensordict.TensorDict¶: Returns the mode of the distribution.

property mean : tensordict.TensorDict¶: Returns the mean of the distribution.

arg_constraints¶

_apply_individually(value: tensordict.TensorDict, func: collections.abc.Callable[[causica.distributions.noise.noise.Noise, torch.Tensor], torch.Tensor]) → tensordict.TensorDict[source]¶

sample_to_noise(samples: tensordict.TensorDict) → tensordict.TensorDict[source]¶

Transform from the sample observations to corresponding noise variables.

Parameters:¶

samples: tensordict.TensorDict¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: tensordict.TensorDict) → tensordict.TensorDict[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: tensordict.TensorDict¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

sample(sample_shape: torch.Size = torch.Size()) → tensordict.TensorDict[source]¶: Generates a sample_shape shaped sample or sample_shape shaped batch of samples if the distribution parameters are batched.

log_prob(value: tensordict.TensorDict) → torch.Tensor[source]¶

Compute the log probs of the given values.

Note

Produces log_probs of the same shape as the inner noise distributions.

Parameters:¶

value: tensordict.TensorDict¶: Values matching the inner independent noise distributions by key. The batch shape of the TensorDict should be set to catch all dims except those associated with the event shape for the expected return shape to behave as specified.

Returns:¶

The log probability of the given value, of shape: torch.broadcast_shapes(value.batch_size, self.batch_shape).

Return type:¶

torch.Tensor

entropy() → torch.Tensor[source]¶

Returns entropy of distribution, batched over batch_shape.

Returns:¶: Tensor of shape batch_shape.

class causica.distributions.noise.JointNoiseModule(independent_noise_modules: Mapping[str, causica.distributions.noise.noise.NoiseModule[causica.distributions.noise.noise.Noise[torch.Tensor]]])[source]¶

Bases: causica.distributions.noise.noise.NoiseModule[JointNoise]

Represents JointNoise with learnable parameters.

Each noise module is used independently on their corresponding key of sample TensorDicts.

forward(x: tensordict.TensorDict | None = None) → JointNoise[source]¶

__getitem__(selection: Iterable[str]) → JointNoiseModule[source]¶: Return a JointNoiseModule representing the subset of variables specified in selection.

keys() → tuple[str, Ellipsis][source]¶: The keys for the different noise modules in order.

causica.distributions.noise.create_noise_modules(shapes: dict[str, torch.Size], types: dict[str, causica.datasets.variable_types.VariableTypeEnum], continuous_noise_dist: ContinuousNoiseDist) → dict[str, causica.distributions.noise.noise.NoiseModule[causica.distributions.noise.noise.Noise[torch.Tensor]]][source]¶

Create noise modules for each item of shapes and types.

Parameters:¶

shapes: dict[str, torch.Size]¶: The shape of each distribution. Currently only the last dimension is used.
types: dict[str, causica.datasets.variable_types.VariableTypeEnum]¶: Names of variables mapping to the variable type VariableTypeEnum.
continuous_noise_dist: ContinuousNoiseDist¶: The noise module to use for variable types of VariableTypeEnum.CONTINUOUS.

Raises:¶

ValueError – If any of the types or the continuous noise distribution is incorrectly specified.

Returns: Dict of independent noise modules following the shape and type specifications.

class causica.distributions.noise.IndependentNoise(base_distribution: BaseNoiseType_co, reinterpreted_batch_ndims: int, validate_args: bool | None = None)[source]¶

Bases: Generic[BaseNoiseType_co], torch.distributions.Independent, Noise[torch.Tensor]

Like td.Idenpendent but also forwards Noise specific methods.

base_dist : BaseNoiseType_co¶

sample_to_noise(samples: torch.Tensor) → torch.Tensor[source]¶

Transform from the sample observations to corresponding noise variables.

This just passes through to the underlying distribution.

Parameters:¶

samples: torch.Tensor¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: torch.Tensor) → torch.Tensor[source]¶

Generate samples using the given exogenous noise.

This just passes through to the underlying distribution.

Parameters:¶

noise: torch.Tensor¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.Noise(batch_shape: torch.Size = torch.Size(), event_shape: torch.Size = torch.Size(), validate_args: bool | None = None)[source]¶

Bases: Generic[SampleType], abc.ABC, torch.distributions.Distribution

Extend Distributions to allow the noise (usually unparametrized) to be extracted from samples and vice versa.

This class is used for Counterfactuals, where we want to keep the noise fixed but fix some values of the SEM.

abstract sample_to_noise(samples: SampleType) → SampleType[source]¶

Transform from the sample observations to corresponding noise variables.

Parameters:¶

samples: SampleType¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

abstract noise_to_sample(noise: SampleType) → SampleType[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: SampleType¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.NoiseModule(*args, **kwargs)[source]¶

Bases: abc.ABC, Generic[NoiseType_co], causica.distributions.distribution_module.DistributionModule[NoiseType_co]

Module for producing Noise, where sampling from the distribution corresponds to adding noise to the input.

All subclasses must allow sampling without input, which should correspond as closely as possible to adding the noise to a compatible zero vector.

abstract forward(x=None) → NoiseType_co[source]¶

class causica.distributions.noise.SplineNoise(base_loc: torch.Tensor, base_scale: torch.Tensor, spline_transforms: list[torch.distributions.AffineTransform | torch.distributions.ComposeTransform])[source]¶

Bases: torch.distributions.TransformedDistribution, causica.distributions.noise.noise.Noise

A Spline Based Noise Distribution.

Parametrized as in in [Neural Spline Flows](https://arxiv.org/pdf/1906.04032.pdf).

sample_to_noise(samples: torch.Tensor) → torch.Tensor[source]¶

Transform from the sample observations to corresponding noise variables.

Parameters:¶

samples: torch.Tensor¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: torch.Tensor) → torch.Tensor[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: torch.Tensor¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.SplineNoiseModule(dim: int, num_bins: int = 8, flow_steps: int = 1, init_knot_locations_scale: float = 0.01, init_derivatives_scale: float = 0.01)[source]¶

Bases: causica.distributions.noise.noise.NoiseModule[causica.distributions.noise.noise.IndependentNoise[SplineNoise]]

Implements Neural Spline Flow noise with learnable parameters.

See [Neural Spline Flows](https://arxiv.org/pdf/1906.04032.pdf).

forward(x: torch.Tensor | None = None) → causica.distributions.noise.noise.IndependentNoise[SplineNoise][source]¶

causica.distributions.noise.create_spline_dist_params(dim: int, num_bins: int, flow_steps: int, knot_locations_scale: float, derivatives_scale: float) → list[SplineParams][source]¶

Create initial values for a spline distribution.

Parameters:¶

dim: int¶: Number of dimensions of the represented variable.
num_bins: int¶: Number of spline bins.
flow_steps: int¶: Number of flow steps.
knot_locations_scale: float¶: Scale of random values used for knot_locations of PiecewiseRationalQuadraticTransform.
derivatives_scale: float¶: Scale of random values for derivatives of PiecewiseRationalQuadraticTransform.

Returns:¶

A list of parameters for `CompositeSplineLayer`s.

class causica.distributions.noise.UnivariateCauchyNoise(loc, scale, validate_args=None)[source]¶

Bases: torch.distributions.Cauchy, causica.distributions.noise.noise.Noise[torch.Tensor]

Samples from a Cauchy (Lorentz) distribution. The distribution of the ratio of independent normally distributed random variables with means 0 follows a Cauchy distribution.

Example:

>>> # xdoctest: +IGNORE_WANT("non-deterinistic")
>>> m = Cauchy(torch.tensor([0.0]), torch.tensor([1.0]))
>>> m.sample()  # sample from a Cauchy distribution with loc=0 and scale=1
tensor([ 2.3214])

Parameters:¶

loc : float or Tensor¶: mode or median of the distribution.
scale : float or Tensor¶: half width at half maximum.

sample_to_noise(samples: torch.Tensor) → torch.Tensor[source]¶

Transform from the sample observations to corresponding noise variables.

Parameters:¶

samples: torch.Tensor¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: torch.Tensor) → torch.Tensor[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: torch.Tensor¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.UnivariateCauchyNoiseModule(dim: int, init_log_scale: float | torch.Tensor = 0.0)[source]¶

Bases: causica.distributions.noise.noise.NoiseModule[causica.distributions.noise.noise.IndependentNoise[UnivariateCauchyNoise]]

Represents a UnivariateCauchyNoise with learnable parameters for independent variables.

forward(x: torch.Tensor | None = None) → causica.distributions.noise.noise.IndependentNoise[UnivariateCauchyNoise][source]¶

class causica.distributions.noise.UnivariateLaplaceNoise(loc, scale, validate_args=None)[source]¶

Bases: torch.distributions.Laplace, causica.distributions.noise.noise.Noise[torch.Tensor]

Creates a Laplace distribution parameterized by loc and scale.

Example:

>>> # xdoctest: +IGNORE_WANT("non-deterinistic")
>>> m = Laplace(torch.tensor([0.0]), torch.tensor([1.0]))
>>> m.sample()  # Laplace distributed with loc=0, scale=1
tensor([ 0.1046])

Parameters:¶

loc : float or Tensor¶: mean of the distribution
scale : float or Tensor¶: scale of the distribution

sample_to_noise(samples: torch.Tensor) → torch.Tensor[source]¶

Transform from the sample observations to corresponding noise variables.

Parameters:¶

samples: torch.Tensor¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: torch.Tensor) → torch.Tensor[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: torch.Tensor¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.UnivariateLaplaceNoiseModule(dim: int, init_log_scale: float | torch.Tensor = 0.0)[source]¶

Bases: causica.distributions.noise.noise.NoiseModule[causica.distributions.noise.noise.IndependentNoise[UnivariateLaplaceNoise]]

Represents a UnivariateLaplaceNoise with learnable parameters for independent variables.

forward(x: torch.Tensor | None = None) → causica.distributions.noise.noise.IndependentNoise[UnivariateLaplaceNoise][source]¶

class causica.distributions.noise.UnivariateNormalNoise(loc, scale, validate_args=None)[source]¶

Bases: torch.distributions.Normal, causica.distributions.noise.noise.Noise[torch.Tensor]

Creates a normal (also called Gaussian) distribution parameterized by loc and scale.

Example:

>>> # xdoctest: +IGNORE_WANT("non-deterinistic")
>>> m = Normal(torch.tensor([0.0]), torch.tensor([1.0]))
>>> m.sample()  # normally distributed with loc=0 and scale=1
tensor([ 0.1046])

Parameters:¶

loc : float or Tensor¶: mean of the distribution (often referred to as mu)
scale : float or Tensor¶: standard deviation of the distribution (often referred to as sigma)

sample_to_noise(samples: torch.Tensor) → torch.Tensor[source]¶

Transform from the sample observations to corresponding noise variables.

Parameters:¶

samples: torch.Tensor¶: Tensor of shape sample_shape + batch_shape + event_shape

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

noise_to_sample(noise: torch.Tensor) → torch.Tensor[source]¶

Generate samples using the given exogenous noise.

Parameters:¶

noise: torch.Tensor¶: noise variable with shape sample_shape + batch_shape.

Returns:¶

The generated samples with shape sample_shape + batch_shape + event_shape

class causica.distributions.noise.UnivariateNormalNoiseModule(dim: int, init_log_scale: float = 0.0)[source]¶

Bases: causica.distributions.noise.noise.NoiseModule[causica.distributions.noise.noise.IndependentNoise[UnivariateNormalNoise]]

Represents a UnivariateNormalNoise with learnable parameters for independent variables.

forward(x: torch.Tensor | None = None) → causica.distributions.noise.noise.IndependentNoise[UnivariateNormalNoise][source]¶

causica.distributions.noise¶

Module Contents¶

Package Contents¶

Classes¶

Functions¶

`causica.distributions.noise`¶