nifty.re.likelihood_impl module#

class Categorical(data, axis=-1, sampling_dtype=<class 'float'>)[source]#

Bases: Likelihood

Categorical likelihood of the data, equivalent to cross entropy

Parameters:
  • data (tree-like structure of jnp.ndarray and int) – Which of the categories is the realized in the data. Must agree with the input shape except for the shape[axis] of the leafs

  • axis (int) – Leaf-axis over which the categories are formed

  • sampling_dtype (dtype, optional) – Data-type for sampling.

See Likelihood for details on the properties.

__init__(data, axis=-1, sampling_dtype=<class 'float'>)[source]#
axis: int = -1#
data: Any#
energy(primals)[source]#

Applies the energy to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

energy – Energy at the position primals.

Return type:

float

left_sqrt_metric(primals, tangents)[source]#

Applies the left-square-root (LSM) of the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric. Must be of shape lsm_tangents_shape.

  • **primals_kw (Any) – Additional arguments passed on to the LSM.

Returns:

metric_sample – Tree-like structure of the same type as primals to which the left-square-root of the metric has been applied to.

Return type:

tree-like structure

metric(primals, tangents)[source]#

Applies the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric.

  • **primals_kw (Any) – Additional arguments passed on to the metric.

Returns:

naturally_curved – Tree-like structure of the same type as primals to which the metric has been applied to.

Return type:

tree-like structure

class Gaussian(data, noise_cov_inv: Callable | None = None, noise_std_inv: Callable | None = None)[source]#

Bases: Likelihood

Gaussian likelihood of the data

Parameters:
  • data (tree-like structure of jnp.ndarray and float) – Data with additive noise following a Gaussian distribution.

  • noise_cov_inv (callable acting on type of data) – Function applying the inverse noise covariance of the Gaussian.

  • noise_std_inv (callable acting on type of data) – Function applying the square root of the inverse noise covariance.

Notes

If noise_std_inv is None it is inferred by assuming a diagonal noise covariance, i.e. by applying it to a vector of ones and taking the square root. If both noise_cov_inv and noise_std_inv are None, a unit covariance is assumed.

See Likelihood for details on the properties.

__init__(data, noise_cov_inv: Callable | None = None, noise_std_inv: Callable | None = None)[source]#
data: Any#
energy(primals)[source]#

Applies the energy to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

energy – Energy at the position primals.

Return type:

float

left_sqrt_metric(primals, tangents)[source]#

Applies the left-square-root (LSM) of the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric. Must be of shape lsm_tangents_shape.

  • **primals_kw (Any) – Additional arguments passed on to the LSM.

Returns:

metric_sample – Tree-like structure of the same type as primals to which the left-square-root of the metric has been applied to.

Return type:

tree-like structure

metric(primals, tangents)[source]#

Applies the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric.

  • **primals_kw (Any) – Additional arguments passed on to the metric.

Returns:

naturally_curved – Tree-like structure of the same type as primals to which the metric has been applied to.

Return type:

tree-like structure

noise_cov_inv: Callable#
noise_std_inv: Callable#
normalized_residual(primals)[source]#

Applies the normalized_residual to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

normalized_residual – Structure of the same type as lsm_tangents_shape for which the normalized_residual is computed.

Return type:

tree-like structure

transformation(primals)[source]#

Applies the coordinate transformation that maps into a coordinate system in which the metric of the likelihood is the Euclidean metric.

Parameters:
  • primals (tree-like structure) – Position at which to transform.

  • **primals_kw (Any) – Additional arguments passed on to the transformation.

Returns:

transformed_sample – Structure of the same type as lsm_tangents_shape to which the geometric transformation has been applied to.

Return type:

tree-like structure

class NDVariableCovarianceGaussian(data, covariance=True)[source]#

Bases: Likelihood

Gaussian likelihood of the data with a variable covariance/precision matrix.

Parameters:
  • data (tree-like structure of jnp.ndarray and float) – Data with additive noise following a multivariate Gaussian distribution. Every leaf must have shape (…, d) with d being the dimension of the multivariate Gaussian distribution.

  • covariance (bool) – If True (default), mat is assumed to be a covariance matrix. If False, mat is assumed to be a precision matrix.

Notes

The likelihood acts on a tuple of (mean, mat) Every leaf of mean must have shape (…, d). Every leaf of mat must have shape (…, d, d).

Please ensure that your forward model guarantees a valid (symmetric and positive-definite) covariance or precision matrix.

See Likelihood for details on the properties.

__init__(data, covariance=True)[source]#
covariance: bool = True#
data: Any#
energy(primals)[source]#

Applies the energy to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

energy – Energy at the position primals.

Return type:

float

left_sqrt_metric(primals, tangents)[source]#

Applies the left-square-root (LSM) of the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric. Must be of shape lsm_tangents_shape.

  • **primals_kw (Any) – Additional arguments passed on to the LSM.

Returns:

metric_sample – Tree-like structure of the same type as primals to which the left-square-root of the metric has been applied to.

Return type:

tree-like structure

metric(primals, tangents)[source]#

Applies the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric.

  • **primals_kw (Any) – Additional arguments passed on to the metric.

Returns:

naturally_curved – Tree-like structure of the same type as primals to which the metric has been applied to.

Return type:

tree-like structure

normalized_residual(primals)[source]#

Applies the normalized_residual to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

normalized_residual – Structure of the same type as lsm_tangents_shape for which the normalized_residual is computed.

Return type:

tree-like structure

transformation(primals)[source]#

Notes

A global transformation to Euclidean space does not exist. A local approximation invoking the residual is used instead.

class Poissonian(data, sampling_dtype=<class 'float'>)[source]#

Bases: Likelihood

Computes the negative log-likelihood, i.e. the Hamiltonians of an expected count Vector constrained by Poissonian count data.

Represents up to an f-independent term \(log(d!)\):

\[E(f) = -\log \text{Poisson}(d|f) = \sum f - d^\dagger \log(f),\]

where f is a Vector in data space of the expectation values for the counts.

Parameters:
  • data (jnp.ndarray or tree-like structure of jnp.ndarray and float) – Data Vector with counts. Needs to have integer dtype and all values need to be non-negative.

  • sampling_dtype (dtype, optional) – Data-type for sampling.

See Likelihood for details on the properties.

__init__(data, sampling_dtype=<class 'float'>)[source]#
data: Any#
energy(primals)[source]#

Applies the energy to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

energy – Energy at the position primals.

Return type:

float

left_sqrt_metric(primals, tangents)[source]#

Applies the left-square-root (LSM) of the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric. Must be of shape lsm_tangents_shape.

  • **primals_kw (Any) – Additional arguments passed on to the LSM.

Returns:

metric_sample – Tree-like structure of the same type as primals to which the left-square-root of the metric has been applied to.

Return type:

tree-like structure

metric(primals, tangents)[source]#

Applies the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric.

  • **primals_kw (Any) – Additional arguments passed on to the metric.

Returns:

naturally_curved – Tree-like structure of the same type as primals to which the metric has been applied to.

Return type:

tree-like structure

normalized_residual(primals)[source]#

Applies the normalized_residual to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

normalized_residual – Structure of the same type as lsm_tangents_shape for which the normalized_residual is computed.

Return type:

tree-like structure

transformation(primals)[source]#

Applies the coordinate transformation that maps into a coordinate system in which the metric of the likelihood is the Euclidean metric.

Parameters:
  • primals (tree-like structure) – Position at which to transform.

  • **primals_kw (Any) – Additional arguments passed on to the transformation.

Returns:

transformed_sample – Structure of the same type as lsm_tangents_shape to which the geometric transformation has been applied to.

Return type:

tree-like structure

class StudentT(data, dof, noise_cov_inv: Callable | None = None, noise_std_inv: Callable | None = None)[source]#

Bases: Likelihood

Student’s t likelihood of the data

Parameters:
  • data (tree-like structure of jnp.ndarray and float) – Data with additive noise following a Gaussian distribution.

  • dof (tree-like structure of jnp.ndarray and float) – Degree-of-freedom parameter of Student’s t distribution.

  • noise_cov_inv (callable acting on type of data) – Function applying the inverse noise covariance of the Gaussian.

  • noise_std_inv (callable acting on type of data) – Function applying the square root of the inverse noise covariance.

Notes

If noise_std_inv is None it is inferred by assuming a diagonal noise covariance, i.e. by applying it to a vector of ones and taking the square root. If both noise_cov_inv and noise_std_inv are None, a unit covariance is assumed.

See Likelihood for details on the properties.

__init__(data, dof, noise_cov_inv: Callable | None = None, noise_std_inv: Callable | None = None)[source]#
data: Any#
dof: Any#
energy(primals)[source]#

Applies the energy to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

energy – Energy at the position primals.

Return type:

float

left_sqrt_metric(primals, tangents)[source]#

Applies the left-square-root (LSM) of the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric. Must be of shape lsm_tangents_shape.

  • **primals_kw (Any) – Additional arguments passed on to the LSM.

Returns:

metric_sample – Tree-like structure of the same type as primals to which the left-square-root of the metric has been applied to.

Return type:

tree-like structure

metric(primals, tangents)[source]#

Applies the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric.

  • **primals_kw (Any) – Additional arguments passed on to the metric.

Returns:

naturally_curved – Tree-like structure of the same type as primals to which the metric has been applied to.

Return type:

tree-like structure

noise_cov_inv: Callable#
noise_std_inv: Callable#
normalized_residual(primals)[source]#

Applies the normalized_residual to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

normalized_residual – Structure of the same type as lsm_tangents_shape for which the normalized_residual is computed.

Return type:

tree-like structure

transformation(primals)[source]#

Applies the coordinate transformation that maps into a coordinate system in which the metric of the likelihood is the Euclidean metric.

Parameters:
  • primals (tree-like structure) – Position at which to transform.

  • **primals_kw (Any) – Additional arguments passed on to the transformation.

Returns:

transformed_sample – Structure of the same type as lsm_tangents_shape to which the geometric transformation has been applied to.

Return type:

tree-like structure

class VariableCovarianceGaussian(data)[source]#

Bases: Likelihood

Gaussian likelihood of the data with a variable covariance

Parameters:

data (jnp.ndarray of float or jft.Vector of tree-like structure) – of jnp.ndarray. Data with additive noise following a Gaussian distribution.

Notes

The likelihood acts on a tuple of (mean, std_inv).

See Likelihood for details on the properties.

__init__(data)[source]#
data: Any#
energy(primals)[source]#

Applies the energy to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

energy – Energy at the position primals.

Return type:

float

left_sqrt_metric(primals, tangents)[source]#

Applies the left-square-root (LSM) of the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric. Must be of shape lsm_tangents_shape.

  • **primals_kw (Any) – Additional arguments passed on to the LSM.

Returns:

metric_sample – Tree-like structure of the same type as primals to which the left-square-root of the metric has been applied to.

Return type:

tree-like structure

metric(primals, tangents)[source]#

Applies the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric.

  • **primals_kw (Any) – Additional arguments passed on to the metric.

Returns:

naturally_curved – Tree-like structure of the same type as primals to which the metric has been applied to.

Return type:

tree-like structure

normalized_residual(primals)[source]#

Applies the normalized_residual to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

normalized_residual – Structure of the same type as lsm_tangents_shape for which the normalized_residual is computed.

Return type:

tree-like structure

transformation(primals)[source]#

Notes

A global transformation to Euclidean space does not exist. A local approximation invoking the residual is used instead.

class VariableCovarianceStudentT(data, dof)[source]#

Bases: Likelihood

Student’s t likelihood of the data with a variable covariance

Parameters:
  • data (tree-like structure of jnp.ndarray and float) – Data with additive noise following a Gaussian distribution.

  • dof (tree-like structure of jnp.ndarray and float) – Degree-of-freedom parameter of Student’s t distribution.

Notes

The likelihood acts on a tuple of (mean, std).

See Likelihood for details on the properties.

__init__(data, dof)[source]#
data: Any#
dof: Any#
energy(primals)[source]#

Applies the energy to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

energy – Energy at the position primals.

Return type:

float

left_sqrt_metric(primals, tangents)[source]#

Applies the left-square-root (LSM) of the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric. Must be of shape lsm_tangents_shape.

  • **primals_kw (Any) – Additional arguments passed on to the LSM.

Returns:

metric_sample – Tree-like structure of the same type as primals to which the left-square-root of the metric has been applied to.

Return type:

tree-like structure

metric(primals, tangent)[source]#

Applies the metric at primals to tangents.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the metric.

  • tangents (tree-like structure) – Instance to which to apply the metric.

  • **primals_kw (Any) – Additional arguments passed on to the metric.

Returns:

naturally_curved – Tree-like structure of the same type as primals to which the metric has been applied to.

Return type:

tree-like structure

normalized_residual(primals)[source]#

Applies the normalized_residual to primals.

Parameters:
  • primals (tree-like structure) – Position at which to evaluate the energy.

  • **primals_kw (Any) – Additional arguments passed on to the energy.

Returns:

normalized_residual – Structure of the same type as lsm_tangents_shape for which the normalized_residual is computed.

Return type:

tree-like structure