Param

class pmrf.parameters.Param(*, value: ArrayLike | None = None, distribution: AbstractDistribution | None = None, constraint: AbstractConstraint | None = None, name: str | None = None, scale: Any = 1.0, fixed: bool = False, metadata: Any = None, raw_value: Any = None)

Bases: AbstractVariable, AbstractWrappable[Array], AbstractAnnotated[Any]

The canonical parameter container for ParamRF.

Parameters can be created by instantiating this class, or using factories in pmrf.parameters, most of which are re-exported at root (e.g. pmrf.Unconstrained(), pmrf.Fixed(), pmrf.Bounded()).

Wraps a Parax variable, applying an optional scale, name and metadata.

Creates a generic parameter.

The incoming value can be any ArrayLike object.

Parameters:
  • value (ArrayLike, optional) – The unscaled value of the parameter.

  • distribution (Optional[AbstractDistribution], optional) – The unscaled probability distribution for the parameter. See pmrf.distributions.

  • constraint (Optional[AbstractConstraint], optional) – The unscaled constraint to apply to the parameter. See pmrf.constraints.

  • name (str, optional) – A name for the parameter, by default None.

  • scale (float, optional) – The scaling factor to apply, by default 1.0.

  • fixed (bool, optional) – Initializes the parameter as fixed. Defaults to False.

  • metadata (Any, optional) – Arbitrary metadata for the parameter, by default None.

  • raw_value (Optional[prx.AbstractVariable], optional) – The raw Parax variable to wrap. Mutually exclusive with value.

as_fixed() Param

Returns a fixed version of this parameter.

Returns:

A new parameter instance wrapped as fixed.

Return type:

Param

as_free() Param

Returns a free (variable) version of this parameter.

Returns:

A new parameter instance wrapped as free.

Return type:

Param

at(where: Callable[[Self], T] | str | tuple[str, ...] | list[str]) Lens[Self, T]

(experimental) A functional interface for parameter manipulation.

This is a wrapper around equinox.tree_at via the jax-optix library.

Similar to pmrf.Model.at() but only accepts callables. See the documentation for that method for more details.

Returns:

A lens object focused on the root of the current instance.

Return type:

Lens

wrap(value: Array) Self

Updates the internal state of the parameter using a physical value.

Parameters:

value (Array) – The physical value to wrap.

Returns:

A new instance of the parameter with the updated state.

Return type:

Self

Raises:

ValueError – If the underlying Parax variable is not wrappable.

property bijector: AbstractBijector | None

The full bijector mapping the raw value to the scaled physical value.

Composes raw_to_constrained_bijector with constrained_to_physical_bijector.

Returns:

The bijector if a constraint exists, otherwise None.

Return type:

AbstractBijector | None

property bounds: tuple[Array | ndarray | bool | number | bool | int | float | complex, Array | ndarray | bool | number | bool | int | float | complex] | None

The unscaled lower and upper bounds of the parameter.

Returns:

A tuple of (lower_bound, upper_bound) if bounds exist, otherwise None.

Return type:

tuple[ArrayLike, ArrayLike] | None

property constrained_to_physical_bijector: AbstractBijector | None

The bijector mapping the constrained value to the scaled physical value.

This is the parameter’s scale. Distributions and bounds are authored in the constrained space, so this is the step needed to compare them against a value that has been unwrapped.

Returns:

The bijector if the parameter is scaled, otherwise None.

Return type:

AbstractBijector | None

property constraint: AbstractConstraint | None

The unscaled constraint associated with the parameter.

Returns:

The constraint if one exists, otherwise None.

Return type:

AbstractConstraint | None

property distribution: AbstractDistribution | None

The unscaled probability distribution associated with the parameter.

Returns:

The distribution if one exists, otherwise None.

Return type:

AbstractDistribution | None

property fixed: bool

Indicates whether the parameter is fixed (constant).

Returns:

True if the parameter is fixed, False otherwise.

Return type:

bool

metadata: Any = None

Arbitrary metadata to store alongside the parameter.

name: str | None = None

A name for the parameter.

property raw_leaf: Array | None

Returns this parameter’s single remaining raw/whitened array leaf, without relying on the internal structure of the wrapped Parax variable (e.g. Random, Fixed). Intended for use after masking or partitioning has already reduced this parameter’s own metadata (its distribution, constraint, etc.) to None, leaving only its own raw value behind.

Returns:

The one remaining leaf, or None if none remain (e.g. this parameter was masked out entirely).

Return type:

jax.Array | None

Raises:

ValueError – If more than one leaf remains, e.g. called before masking.

property raw_to_constrained_bijector: AbstractBijector | None

The bijector mapping the raw value to the constrained value.

The raw value is the latent one held in raw_value, which Parax refers to as the unconstrained space. It is called raw here to avoid confusion with pmrf.Unconstrained(), which creates a parameter without bounds.

Returns:

The bijector if a constraint exists, otherwise None.

Return type:

AbstractBijector | None

raw_value: AbstractVariable

The raw value of the parameter.

scale: float = 1.0

The scale of the parameter.

property unscaled_value: Array

Returns the original unscaled value.

Returns:

The computed array value.

Return type:

jax.Array

property value: Array

Returns the scaled physical value.

Returns:

The computed array value.

Return type:

jax.Array