Skip to content

Transforms

jimgw.core.transforms ¤

Transform ¤

Bases: ABC

Abstract base class for parameter-space transforms.

Transforms are responsible for mapping named parameter dictionaries between two coordinate spaces. Subclasses implement the actual numerical mapping; this base class handles name bookkeeping.

Attributes:

Name Type Description
name_mapping tuple[list[str], list[str]]

A pair (from_names, to_names) describing which input parameters are consumed and which output parameters are produced.
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
__init__(name_mapping: tuple[list[str], list[str]]) -> None ¤

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

Pair of (input_parameter_names, output_parameter_names).

 required
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤

NtoMTransform ¤

Bases: Transform

N-to-M transform: consumes N named parameters and produces M named parameters.

Only a forward pass is defined. This is the most general transform type.

Attributes:

Name Type Description
transform_func Callable

Callable that maps an input dict to an output dict.
transform_func: Callable[[dict[str, Float]], dict[str, Float]] instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
__repr__() ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
__init__(name_mapping: tuple[list[str], list[str]]) -> None ¤

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

Pair of (input_parameter_names, output_parameter_names).

 required
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤

NtoNTransform ¤

Bases: NtoMTransform

N-to-N transform with automatic Jacobian computation via forward-mode AD.

The number of input and output parameters must be equal. The log-Jacobian determinant of the transform is computed automatically using jax.jacfwd.

n_dim: int property ¤

Number of parameters consumed/produced by this transform.

name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
transform_func: Callable[[dict[str, Float]], dict[str, Float]] instance-attribute ¤
__repr__() ¤
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
__init__(name_mapping: tuple[list[str], list[str]]) -> None ¤

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

Pair of (input_parameter_names, output_parameter_names).

 required
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.

BijectiveTransform ¤

Bases: NtoNTransform

Bijective (invertible) N-to-N transform.

Extends :class:NtoNTransform with an inverse transform function and corresponding inverse and backward methods.

Attributes:

Name Type Description
inverse_transform_func Callable

Callable implementing the inverse map.
inverse_transform_func: Callable[[dict[str, Float]], dict[str, Float]] instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
transform_func: Callable[[dict[str, Float]], dict[str, Float]] instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.
__init__(name_mapping: tuple[list[str], list[str]]) -> None ¤

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

Pair of (input_parameter_names, output_parameter_names).

 required
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.

ConditionalBijectiveTransform ¤

Bases: BijectiveTransform

Bijective transform that depends on additional conditioning parameters.

Like :class:BijectiveTransform, but the transform and inverse functions also receive a set of fixed conditioning parameters (e.g. masses when converting spin angles). The Jacobian is computed only with respect to the primary (non-conditioning) parameters.

Attributes:

Name Type Description
conditional_names list[str]

Names of the conditioning parameters.
conditional_names: list[str] = conditional_names instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
transform_func: Callable[[dict[str, Float]], dict[str, Float]] instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

inverse_transform_func: Callable[[dict[str, Float]], dict[str, Float]] instance-attribute ¤
__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], conditional_names: list[str]) -> None ¤

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

Pair of (input_names, output_names) for the primary parameters.

 required
conditional_names list[str]

Names of the conditioning parameters that are passed through unchanged.

 required
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Apply the conditional forward transform and compute the log-Jacobian.

Parameters:

Name Type Description Default
x dict[str, Float]

Parameter dictionary containing both primary and conditioning parameters.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed parameter dictionary and the log absolute Jacobian determinant w.r.t. the primary parameters.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Apply the conditional inverse transform and compute the log-Jacobian.

Parameters:

Name Type Description Default
y dict[str, Float]

Parameter dictionary in the output space, containing both primary (output) and conditioning parameters.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The inverse-transformed parameter dictionary and the log absolute Jacobian determinant w.r.t. the primary output parameters.
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

ScaleTransform ¤

Bases: BijectiveTransform

Elementwise scaling transform: y = x * scale.

Attributes:

Name Type Description
scale Float

Multiplicative scale factor.
scale: Float = scale instance-attribute ¤
transform_func = lambda x: {(name_mapping[1][i]): (x[name_mapping[0][i]] * self.scale) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (x[name_mapping[1][i]] / self.scale) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], scale: Float) -> None ¤

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

(input_names, output_names).

 required
scale Float

Scale factor applied in the forward direction.

 required
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

OffsetTransform ¤

Bases: BijectiveTransform

Elementwise offset (translation) transform: y = x + offset.

Attributes:

Name Type Description
offset Float

Additive offset applied in the forward direction.
offset: Float = offset instance-attribute ¤
transform_func = lambda x: {(name_mapping[1][i]): (x[name_mapping[0][i]] + self.offset) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (x[name_mapping[1][i]] - self.offset) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], offset: Float) -> None ¤

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

(input_names, output_names).

 required
offset Float

Offset added in the forward direction.

 required
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

LogitTransform ¤

Bases: BijectiveTransform

Logit transform.

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

The name mapping between the input and output dictionary.

 required
transform_func = lambda x: {(name_mapping[1][i]): (1 / (1 + jnp.exp(-x[name_mapping[0][i]]))) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (logit(x[name_mapping[1][i]])) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]]) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

SineTransform ¤

Bases: BijectiveTransform

Sine transformation.

The original parameter is expected to be in [-pi/2, pi/2].

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

The name mapping between the input and output dictionary.

 required
transform_func = lambda x: {(name_mapping[1][i]): (jnp.sin(x[name_mapping[0][i]])) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (jnp.arcsin(x[name_mapping[1][i]])) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]]) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

CosineTransform ¤

Bases: BijectiveTransform

Cosine transformation.

The original parameter is expected to be in [0, pi].

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

The name mapping between the input and output dictionary.

 required
transform_func = lambda x: {(name_mapping[1][i]): (jnp.cos(x[name_mapping[0][i]])) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (jnp.arccos(x[name_mapping[1][i]])) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]]) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

BoundToBound ¤

Bases: BijectiveTransform

Linear rescaling from one bounded interval to another.

Maps x ∈ [original_lower, original_upper] to y ∈ [target_lower, target_upper] via a linear (affine) transform.

Attributes:

Name Type Description
original_lower_bound Float[Array, ' n_dim']

Lower bound(s) of the input interval.
original_upper_bound Float[Array, ' n_dim']

Upper bound(s) of the input interval.
target_lower_bound Float[Array, ' n_dim']

Lower bound(s) of the output interval.
target_upper_bound Float[Array, ' n_dim']

Upper bound(s) of the output interval.
original_lower_bound: Float[Array, ' n_dim'] = jnp.atleast_1d(original_lower_bound) instance-attribute ¤
original_upper_bound: Float[Array, ' n_dim'] = jnp.atleast_1d(original_upper_bound) instance-attribute ¤
target_lower_bound: Float[Array, ' n_dim'] = jnp.atleast_1d(target_lower_bound) instance-attribute ¤
target_upper_bound: Float[Array, ' n_dim'] = jnp.atleast_1d(target_upper_bound) instance-attribute ¤
transform_func = lambda x: {(name_mapping[1][i]): ((x[name_mapping[0][i]] - self.original_lower_bound[i]) * (self.target_upper_bound[i] - self.target_lower_bound[i]) / (self.original_upper_bound[i] - self.original_lower_bound[i]) + self.target_lower_bound[i]) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): ((x[name_mapping[1][i]] - self.target_lower_bound[i]) * (self.original_upper_bound[i] - self.original_lower_bound[i]) / (self.target_upper_bound[i] - self.target_lower_bound[i]) + self.original_lower_bound[i]) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], original_lower_bound: Float | Float[Array, ' n_dim'], original_upper_bound: Float | Float[Array, ' n_dim'], target_lower_bound: Float | Float[Array, ' n_dim'], target_upper_bound: Float | Float[Array, ' n_dim']) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

BoundToUnbound ¤

Bases: BijectiveTransform

Logit-based transform from a bounded interval to the real line.

Maps x ∈ (original_lower, original_upper) to y ∈ (-∞, +∞) via

.. math::

y = \text{logit}\!\left(\frac{x - x_{\min}}{x_{\max} - x_{\min}}\right).

The inverse maps back with the sigmoid function.

Attributes:

Name Type Description
original_lower_bound Float[Array, ' n_dim']

Lower bound(s) of the input interval.
original_upper_bound Float[Array, ' n_dim']

Upper bound(s) of the input interval.
original_lower_bound: Float[Array, ' n_dim'] = jnp.atleast_1d(original_lower_bound) instance-attribute ¤
original_upper_bound: Float[Array, ' n_dim'] = jnp.atleast_1d(original_upper_bound) instance-attribute ¤
transform_func = lambda x: {(name_mapping[1][i]): (logit((x[name_mapping[0][i]] - self.original_lower_bound[i]) / (self.original_upper_bound[i] - self.original_lower_bound[i]))) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): ((self.original_upper_bound[i] - self.original_lower_bound[i]) / (1 + jnp.exp(-x[name_mapping[1][i]])) + self.original_lower_bound[i]) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], original_lower_bound: Float | Float[Array, ' n_dim'], original_upper_bound: Float | Float[Array, ' n_dim']) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

SingleSidedUnboundTransform ¤

Bases: BijectiveTransform

Unbound upper limit transformation.

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

The name mapping between the input and output dictionary.

 required
original_lower_bound: Float[Array, ' n_dim'] = jnp.atleast_1d(original_lower_bound) instance-attribute ¤
transform_func = lambda x: {(name_mapping[1][i]): (jnp.log(x[name_mapping[0][i]] - self.original_lower_bound[i])) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (jnp.exp(x[name_mapping[1][i]]) + self.original_lower_bound[i]) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], original_lower_bound: Float | Float[Array, ' n_dim']) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

PowerLawTransform ¤

Bases: BijectiveTransform

PowerLaw transformation.

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

The name mapping between the input and output dictionary.

 required
xmin: Float = xmin instance-attribute ¤
xmax: Float = xmax instance-attribute ¤
alpha: Float = alpha instance-attribute ¤
transform_func = lambda x: {(name_mapping[1][i]): (self.xmin * jnp.exp(x[name_mapping[0][i]] * jnp.log(self.xmax / self.xmin))) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (jnp.log(x[name_mapping[1][i]] / self.xmin) / jnp.log(self.xmax / self.xmin)) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], xmin: Float, xmax: Float, alpha: Float) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

CartesianToPolarTransform ¤

Bases: BijectiveTransform

Transformation from (x, y) to (theta, r).

Parameters:

Name Type Description Default
parameter_name str

The name of the parameter to be transformed.

 required
transform_func = lambda x: {f'{parameter_name}_theta': jnp.arctan2(x[f'{parameter_name}_y'], x[f'{parameter_name}_x']) + jnp.pi, f'{parameter_name}_r': jnp.sqrt(x[f'{parameter_name}_x'] ** 2 + x[f'{parameter_name}_y'] ** 2)} instance-attribute ¤
inverse_transform_func = lambda x: {f'{parameter_name}_x': x[f'{parameter_name}_r'] * jnp.cos(x[f'{parameter_name}_theta']), f'{parameter_name}_y': x[f'{parameter_name}_r'] * jnp.sin(x[f'{parameter_name}_theta'])} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(parameter_name: str) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

PeriodicTransform ¤

Bases: BijectiveTransform

Transform a periodic parameter onto a 2D circle.

Maps (r, θ) — where θ ∈ [xmin, xmax] is the periodic angle and r is a scale — to Cartesian coordinates (x, y) = r(cos θ', sin θ') (with θ' = 2π(θ - xmin)/(xmax - xmin)). The inverse recovers (r, θ) from (x, y).

Attributes:

Name Type Description
xmin Float

Lower bound of the periodic parameter.
xmax Float

Upper bound of the periodic parameter.
xmin = xmin instance-attribute ¤
xmax = xmax instance-attribute ¤
transform_func = lambda x: {f'{name_mapping[1][0]}': x[name_mapping[0][0]] * jnp.cos(scaling * (x[name_mapping[0][1]] - self.xmin)), f'{name_mapping[1][1]}': x[name_mapping[0][0]] * jnp.sin(scaling * (x[name_mapping[0][1]] - self.xmin))} instance-attribute ¤
inverse_transform_func = lambda x: {name_mapping[0][1]: self.xmin + (jnp.pi + jnp.arctan2(x[name_mapping[1][1]], x[name_mapping[1][0]])) / scaling, name_mapping[0][0]: jnp.sqrt(x[name_mapping[1][0]] ** 2 + x[name_mapping[1][1]] ** 2)} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], xmin: Float, xmax: Float) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

RayleighTransform ¤

Bases: BijectiveTransform

Transformation from Uniform(0, 1) to Rayleigh distribution with scale parameter sigma.

Parameters:

Name Type Description Default
name_mapping tuple[list[str], list[str]]

The name mapping between the input and output dictionary.

 required
sigma = sigma instance-attribute ¤
transform_func = lambda x: {(name_mapping[1][i]): (sigma * jnp.sqrt(-2 * jnp.log(x[name_mapping[0][i]]))) for i in (range(len(name_mapping[0])))} instance-attribute ¤
inverse_transform_func = lambda x: {(name_mapping[0][i]): (jnp.exp(-(x[name_mapping[1][i]] / sigma) ** 2 / 2)) for i in (range(len(name_mapping[1])))} instance-attribute ¤
name_mapping: tuple[list[str], list[str]] = name_mapping instance-attribute ¤
n_dim: int property ¤

Number of parameters consumed/produced by this transform.

__repr__() ¤
__init__(name_mapping: tuple[list[str], list[str]], sigma: Float) ¤
propagate_name(x: Sequence[str]) -> tuple[str, ...] ¤
forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input x to transformed coordinate y.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The transformed dictionary.
transform(x: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Transform the input x to transformed coordinate y and return the log Jacobian determinant. This only works if the transform is a N -> N transform.

Parameters:

Name Type Description Default
x dict[str, Float]

The input dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The transformed dictionary and the log Jacobian determinant.
inverse(y: dict[str, Float]) -> tuple[dict[str, Float], Float] ¤

Inverse transform the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
tuple[dict[str, Float], Float]

tuple[dict[str, Float], Float]: The original dictionary and the log Jacobian determinant.
backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back the input y to original coordinate x.

Parameters:

Name Type Description Default
y dict[str, Float]

The transformed dictionary.

 required

Returns:

Type Description
dict[str, Float]

dict[str, Float]: The original dictionary.

reverse_bijective_transform(original_transform: BijectiveTransform) -> BijectiveTransform ¤

Construct the inverse of a bijective transform.

Swaps the forward and inverse functions and the name mapping of original_transform, returning a new :class:BijectiveTransform (or :class:ConditionalBijectiveTransform) whose forward direction is the original's inverse.

Parameters:

Name Type Description Default
original_transform BijectiveTransform

The transform to invert.

 required

Returns:

Name Type Description
BijectiveTransform BijectiveTransform

A new transform that is the reverse of the input.