Add TruncatedNormal#91
Conversation
Truncated Normal distribution with loc, scale, low, and high parameters. Every closed-form formula (sample, log_prob, cdf, icdf, survival function, entropy, mean, variance, mode, median) is cross-checked numerically against scipy.stats.truncnorm. KL divergence between arbitrary truncated normals isn't analytically tractable and raises NotImplementedError.
| low: Array | ||
| high: Array | ||
|
|
||
| def __init__(self, loc: Array, scale: Array, low: Array, high: Array): |
There was a problem hiding this comment.
can type better/more specific? + fields
|
|
||
| def __init__(self, loc: Array, scale: Array, low: Array, high: Array): | ||
| """Initializes a Truncated Normal distribution. | ||
|
|
There was a problem hiding this comment.
for docstrings with fields you can set the init.docs (I did this with uniform e.g.
| return self._standardize(self.high) | ||
|
|
||
| def _log_normalizer(self) -> Array: | ||
| """Calculates the log of the normalization constant Z.""" |
There was a problem hiding this comment.
we should include math in the description to define Z (of the class)
| """Calculates the log of the normalization constant Z.""" | ||
| cdf_high = jax.scipy.special.ndtr(self._std_high) | ||
| cdf_low = jax.scipy.special.ndtr(self._std_low) | ||
| return jnp.log(cdf_high - cdf_low) |
There was a problem hiding this comment.
also should this just be a property or member variable?
| """Calculates the log of the normalization constant Z.""" | ||
| cdf_high = jax.scipy.special.ndtr(self._std_high) | ||
| cdf_low = jax.scipy.special.ndtr(self._std_low) | ||
| return jnp.log(cdf_high - cdf_low) |
There was a problem hiding this comment.
we probably want to do this in log space? like with jax.scipy.special.log_ndtr right, since large distributions e.g. TruncatedNormal(0, 1, 8, 9).icdf(0.5) is inf but is stable with scipy.
| Z = jax.scipy.special.ndtr(beta) - jax.scipy.special.ndtr(alpha) | ||
|
|
||
| pdf_alpha = self._pdf(alpha) | ||
| pdf_beta = self._pdf(beta) |
There was a problem hiding this comment.
do we want to support single bounded distributions, e.g. [10, inf) truncation? this might hit a NaN here if so
2 participants
Truncated Normal distribution with loc, scale, low, and high parameters. Every closed-form formula (sample, log_prob, cdf, icdf, survival function, entropy, mean, variance, mode, median) is cross-checked numerically against scipy.stats.truncnorm. KL divergence between arbitrary truncated normals isn't analytically tractable and raises NotImplementedError.