SoftplusT#

class brainstate.nn.SoftplusT(lower)#

Softplus transformation mapping unbounded values to positive semi-infinite interval.

This transformation uses the softplus function to map any real value to the interval [lower, ∞). It provides a smooth, differentiable alternative to ReLU activation and is commonly used to constrain parameters to be positive, such as variance parameters or rate constants.

The transformation is defined by:

\[\text{forward}(x) = \log(1 + e^x) + \text{lower}\]

The inverse transformation is:

\[\text{inverse}(y) = \log(e^{y - \text{lower}} - 1)\]
Parameters:

lower (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Lower bound of the target interval.

lower#

Lower bound of the interval.

Type:

array_like

unit#

Physical unit of the lower bound.

Type:

brainunit.Unit

Notes

The softplus function is the smooth approximation to the ReLU function: \(\lim_{\beta \to \infty} \frac{1}{\beta} \log(1 + e^{\beta x}) = \max(0, x)\)

For large positive x, softplus(x) ≈ x, and for large negative x, softplus(x) ≈ 0. The function is strictly positive and has a well-defined inverse.

Examples

>>> # Map to positive reals [0, ∞)
>>> transform = SoftplusT(0.0)
>>> x = jnp.array([-5.0, 0.0, 5.0])
>>> y = transform.forward(x)
>>> # y ≈ [0.007, 0.693, 5.007]
>>> # Map to interval [2, ∞) for positive-definite parameters
>>> transform = SoftplusT(2.0)
>>> x = jnp.array([0.0])
>>> y = transform.forward(x)
>>> # y ≈ [2.693]
forward(x)[source]#

Transform unbounded input to positive semi-infinite interval.

Parameters:

x (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input values in unbounded domain (-∞, ∞).

Returns:

Transformed values in interval [lower, ∞).

Return type:

Array

Notes

Uses jax.nn.softplus for numerical stability: softplus(x) = log(1 + exp(x)), which is exact for all x (the previous log1p(exp(x)) form saturated for x beyond ~20, breaking the forward map and round-trip).

inverse(y)[source]#

Transform positive semi-infinite input back to unbounded domain.

Parameters:

y (Array | ndarray | bool | number | bool | int | float | complex | Quantity) – Input values in interval [lower, ∞).

Returns:

Transformed values in unbounded domain (-∞, ∞).

Return type:

Array

Notes

Input must be strictly greater than lower bound to avoid numerical issues. Uses the numerically stable softplus inverse z + log(-expm1(-z)) which is accurate across the whole range (the previous log(exp(z) - 1) form clipped z at ~20 and so failed to invert large constrained values). z is dimensionless after dividing out self.unit, so its mantissa is taken to keep the bare-jnp log/expm1 calls valid.

log_abs_det_jacobian(x, y)[source]#

Compute log absolute determinant of the Jacobian.

For softplus: d/dx[log(1 + exp(x))] = sigmoid(x) log|det J| = sum(log(sigmoid(x))) = sum(x - softplus(x))

Return type:

Array