Optimization

class Transform

class SigmoidTransform(lower, upper)

Sigmoid transformation.

Parameters:

• lower (Array | ndarray | bool | number | bool | int | float | complex)

• upper (Array | ndarray | bool | number | bool | int | float | complex)

class SoftplusTransform(lower)

Softplus transformation.

Parameters:

lower (Array | ndarray | bool | number | bool | int | float | complex)

class NegSoftplusTransform(upper)

Negative softplus transformation.

Parameters:

upper (Array | ndarray | bool | number | bool | int | float | complex)

class AffineTransform(scale, shift)

Parameters:

• scale (Array | ndarray | bool | number | bool | int | float | complex)

• shift (Array | ndarray | bool | number | bool | int | float | complex)

class ChainTransform(transforms)

Chaining together multiple transformations

Parameters:

transforms (Sequence[Transform])

class MaskedTransform(mask, transform)

Parameters:

• mask (Array | ndarray | bool | number | bool | int | float | complex)

• transform (Transform)

class CustomTransform(forward_fn, inverse_fn)

Custom transformation

Parameters:

• forward_fn (Callable)

• inverse_fn (Callable)

class ParamTransform(tf_dict)

Parameter transformation utility.

This class is used to transform parameters usually from an unconstrained space to a constrained space and back (bacause most biophysical parameter are bounded). The user can specify a PyTree of transforms that are applied to the parameters.

Parameters:

tf_dict (List[Dict[str, Transform]] | Transform)

tf_dict

A PyTree of transforms for each parameter.

forward(params)

Pushes unconstrained parameters through a tf such that they fit the interval.

Parameters:

params (List[Dict[str, Array | ndarray | bool | number | bool | int | float | complex]] | Array | ndarray | bool | number | bool | int | float | complex) – A list of dictionaries (or any PyTree) with unconstrained parameters.

Returns:

A list of dictionaries (or any PyTree) with transformed parameters.

Return type:

Dict[str, Array]

inverse(params)

Takes parameters from within the interval and makes them unconstrained.

Parameters:

params (List[Dict[str, Array | ndarray | bool | number | bool | int | float | complex]] | Array | ndarray | bool | number | bool | int | float | complex) – A list of dictionaries (or any PyTree) with transformed parameters.

Returns:

A list of dictionaries (or any PyTree) with unconstrained parameters.

Return type:

Dict[str, Array]

class TypeOptimizer(optimizer, optimizer_args, opt_params)

optax wrapper which allows different argument values for different params.

Parameters:

• optimizer (Callable)

• optimizer_args (Dict[str, Any])

• opt_params (List[Dict[str, Array]])

init(opt_params)

Initialize the optimizers. Equivalent to optax.optimizers.init().

Parameters:

opt_params (List[Dict[str, Array]])

Return type:

List

update(gradient, opt_state)

Update the optimizers. Equivalent to optax.optimizers.update().

Parameters:

• gradient (Array)

• opt_state (List)

Return type:

Tuple[List, List]