Skip to content

Transforms

Functions:

Name Description
reverse_bijective_transform

Construct the inverse of a bijective transform.

BijectiveTransform ¤

Bases: NtoNTransform

Bijective (invertible) N-to-N transform.

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

Attributes:

Name Type Description
inverse_transform_func Callable[[dict[str, Float]], dict[str, Float]]

Callable implementing the inverse map.

Methods:

Name Description
backward

Pull back from the output coordinate to the input coordinate (no Jacobian).
inverse

Inverse transform from the output coordinate back to the input coordinate.

backward(y: dict[str, Float]) -> dict[str, Float] ¤

Pull back from the output coordinate to the input coordinate (no Jacobian).

Parameters:

Name Type Description Default
y dict[str, Float]

Parameter dictionary in the transformed (output) space.

 required

Returns:

Type Description
dict[str, Float]

Parameter dictionary in the original (input) space.

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

Inverse transform from the output coordinate back to the input coordinate.

Parameters:

Name Type Description Default
y dict[str, Float]

Parameter dictionary in the transformed (output) space.

 required

Returns:

Type Description
dict[str, Float]

Tuple of the inverse-transformed parameter dictionary and the log
FloatScalar

absolute Jacobian determinant.

BoundToBound ¤

Bases: BijectiveTransform

Linear rescaling from one bounded interval to another.

Maps x in [original_lower, original_upper] to y in [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.

Methods:

Name Description
__init__

Args:

__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']) ¤

Parameters:

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

Pair of (input_names, output_names).

 required
original_lower_bound float | Float[Array, ' n_dim']

Lower bound(s) of the source interval.

 required
original_upper_bound float | Float[Array, ' n_dim']

Upper bound(s) of the source interval.

 required
target_lower_bound float | Float[Array, ' n_dim']

Lower bound(s) of the target interval.

 required
target_upper_bound float | Float[Array, ' n_dim']

Upper bound(s) of the target interval.

 required

BoundToUnbound ¤

Bases: BijectiveTransform

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

Maps values from (original_lower, original_upper) to the real line via

$$

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.

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]], original_lower_bound: float | Float[Array, ' n_dim'], original_upper_bound: float | Float[Array, ' n_dim']) ¤

Parameters:

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

Pair of (input_names, output_names).

 required
original_lower_bound float | Float[Array, ' n_dim']

Lower bound(s) of the source interval.

 required
original_upper_bound float | Float[Array, ' n_dim']

Upper bound(s) of the source interval.

 required

CartesianToPolarTransform ¤

Bases: BijectiveTransform

Cartesian-to-polar transform: (x, y) -> (theta, r) with theta in [0, 2*pi].

Methods:

Name Description
__init__

Args:

__init__(parameter_name: str) ¤

Parameters:

Name Type Description Default
parameter_name str

Base name; consumes <name>_x and <name>_y, produces <name>_theta and <name>_r.

 required

ConditionalBijectiveTransform ¤

Bases: BijectiveTransform

Bijective transform that depends on additional conditioning parameters.

Like 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.

Methods:

Name Description
__init__

Args:
inverse

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

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

__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

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

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
dict[str, Float]

The inverse-transformed parameter dictionary and the log absolute
FloatScalar

Jacobian determinant w.r.t. the primary output parameters.

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

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
dict[str, Float]

The transformed parameter dictionary and the log absolute Jacobian
FloatScalar

determinant w.r.t. the primary parameters.

CosineTransform ¤

Bases: BijectiveTransform

Cosine transformation: y = cos(x), with x in [0, pi].

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]]) ¤

Parameters:

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

Pair of (input_names, output_names).

 required

GaussianTransform ¤

Bases: BijectiveTransform

CDF transform from Uniform(0, 1) to a Gaussian (normal) distribution.

Maps u in (0, 1) to x = mu + sigma * ndtri(u) using the probit (quantile) function. The inverse maps x -> ndtr((x - mu) / sigma) in (0, 1) using the normal CDF.

Attributes:

Name Type Description
mu float

Mean of the Gaussian distribution.
sigma float

Standard deviation of the Gaussian distribution.

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]], mu: float, sigma: float) ¤

Parameters:

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

Pair of (input_names, output_names).

 required
mu float

Mean of the Gaussian distribution.

 required
sigma float

Standard deviation of the Gaussian distribution.

 required

LogitTransform ¤

Bases: BijectiveTransform

Logit transform: y = sigmoid(x) (forward) / x = logit(y) (backward).

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]]) ¤

Parameters:

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

Pair of (input_names, output_names).

 required

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.

Methods:

Name Description
forward

Push forward the input to the transformed coordinate.

forward(x: dict[str, Float]) -> dict[str, Float] ¤

Push forward the input to the transformed coordinate.

Parameters:

Name Type Description Default
x dict[str, Float]

Parameter dictionary in the input space.

 required

Returns:

Type Description
dict[str, Float]

Parameter dictionary in the transformed space.

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.

Methods:

Name Description
transform

Transform to the output coordinate and return the log-Jacobian determinant.

Attributes:

Name Type Description
n_dim int

Number of parameters consumed/produced by this transform.

n_dim: int property ¤

Number of parameters consumed/produced by this transform.

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

Transform to the output coordinate and return the log-Jacobian determinant.

Parameters:

Name Type Description Default
x dict[str, Float]

Parameter dictionary in the input space.

 required

Returns:

Type Description
dict[str, Float]

Tuple of the transformed parameter dictionary and the log absolute
FloatScalar

Jacobian determinant.

OffsetTransform ¤

Bases: BijectiveTransform

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

Attributes:

Name Type Description
offset float

Additive offset applied in the forward direction.

Methods:

Name Description
__init__

Args:

__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

PeriodicTransform ¤

Bases: BijectiveTransform

Transform a periodic parameter onto a 2D circle.

Maps (r, theta) - where theta in [xmin, xmax] is the periodic angle and r is a scale - to Cartesian coordinates (x, y) = r(cos theta_prime, sin theta_prime) (with theta_prime = 2*pi*(theta - xmin)/(xmax - xmin)). The inverse recovers (r, theta) from (x, y).

Attributes:

Name Type Description
xmin Float

Lower bound of the periodic parameter.
xmax Float

Upper bound of the periodic parameter.

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]], xmin: float, xmax: float) ¤

Parameters:

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

Pair of ([scale_name, angle_name], [x_name, y_name]).

 required
xmin float

Lower bound of the periodic parameter.

 required
xmax float

Upper bound of the periodic parameter.

 required

PowerLawTransform ¤

Bases: BijectiveTransform

Power-law CDF transform from Uniform(0, 1) to a power-law distribution.

Maps the unit interval to [xmin, xmax] with the CDF of p(x) ∝ x^alpha. The alpha = -1 (log-uniform) case is handled separately.

Attributes:

Name Type Description
xmin float

Lower bound of the output domain.
xmax float

Upper bound of the output domain.
alpha float

Power-law exponent.

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]], xmin: float, xmax: float, alpha: float) ¤

Parameters:

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

Pair of (input_names, output_names).

 required
xmin float

Lower bound of the output domain.

 required
xmax float

Upper bound of the output domain.

 required
alpha float

Power-law exponent (alpha = -1 gives log-uniform).

 required

RayleighTransform ¤

Bases: BijectiveTransform

CDF transform from Uniform(0, 1) to a Rayleigh distribution.

Maps u in [0, 1] to x = sigma * sqrt(-2 * log(u)).

Attributes:

Name Type Description
sigma

Scale parameter of the Rayleigh distribution.

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]], sigma: float) ¤

Parameters:

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

Pair of (input_names, output_names).

 required
sigma float

Scale parameter of the Rayleigh distribution.

 required

ScaleTransform ¤

Bases: BijectiveTransform

Elementwise scaling transform: y = x * scale.

Attributes:

Name Type Description
scale float

Multiplicative scale factor.

Methods:

Name Description
__init__

Args:

__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

SineTransform ¤

Bases: BijectiveTransform

Sine transformation: y = sin(x), with x in [-pi/2, pi/2].

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]]) ¤

Parameters:

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

Pair of (input_names, output_names).

 required

SingleSidedUnboundTransform ¤

Bases: BijectiveTransform

Log transform from a lower-bounded domain to the real line: y = log(x - lower).

Attributes:

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

Lower bound(s) of the input domain.

Methods:

Name Description
__init__

Args:

__init__(name_mapping: tuple[list[str], list[str]], original_lower_bound: float | Float[Array, ' n_dim']) ¤

Parameters:

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

Pair of (input_names, output_names).

 required
original_lower_bound float | Float[Array, ' n_dim']

Lower bound(s) of the input domain.

 required

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.

Methods:

Name Description
__init__

Args:
propagate_name

Propagate a parameter-name tuple through this transform.

__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, ...] ¤

Propagate a parameter-name tuple through this transform.

Removes the names consumed by this transform and appends the names it produces, preserving the order of unchanged names.

Parameters:

Name Type Description Default
x Sequence[str]

Sequence of parameter names in the current space.

 required

Returns:

Type Description
tuple[str, ...]

New tuple of parameter names after applying this transform.

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 BijectiveTransform (or ConditionalBijectiveTransform) whose forward direction is the original's inverse.

Parameters:

Name Type Description Default
original_transform BijectiveTransform

The transform to invert.

 required

Returns:

Type Description
BijectiveTransform

A new transform that is the reverse of the input.