At PyMC Labs we actually had to write a probability distribution in an n-ball. The project will become open source later this year, but I can’t share the actual code until then.
We based our distribution from the methods described in this paper. Basically, a uniform distribution in a 5-ball, will have a piecewise pdf, like:
f(x) = \begin{cases}
\frac{\Gamma(\frac{5}{2} + 1)}{\pi^{5/2}} & \text{if }||x||\leq 1\\
0 & \text{if }||x||> 1\\
\end{cases}
Sampling here will be difficult for HMC because the gradient is 0 always until the points reach the bound of the sphere. If you want to be able to actually use HMC to draw samples in this space, you will need to implement a transform that maps points from \mathbb{R}^{5}\rightarrow \mathbb{B}^{5}_{2} (where \mathbb{B}^{5}_{2} stands for the unit ball in 5 dimensions under the L_2 norm).
The transform g(x), \mathbb{R}^{5}\rightarrow \mathbb{B}^{5}_{2} that we implemented was this:
g(x) = \frac{2x}{\pi||x||} \arctan(||x||)
We had a bit of trouble to write down the determinant of the jacobian transform (which you need to get the pdf in the unconstrained space, which enables sampling), but you basically end up with this expression:
J(g(x))= \frac{2 \vec{x}.\vec{x^t}}{\pi ||x||^{2} (1 + ||x||^{2})}
- \frac{2 \vec{x}.\vec{x^t} \arctan(||x||)}{\pi||x||^{3}}
+ \frac{2 \arctan(||x||) \mathbb{I}}{\pi ||x||}
With that transform in hand, we could implement a new distribution that used the transform during sampling. That way, HMC becomes very efficient as the unconstrained space and the jacobian information, allow it to navigate the uniform ball rather well. Our probability distribution actually also had a concentration parameter, that allowed it to concentrate points at the center or on the shell:
If you still want to use a density dist without a custom transform, you can. But keep in mind, that HMC might still struggle because of the sharp discontinuity in the pdf. What you would have to do to get a DensityDist working is this:
import numpy as np
import pymc as pm
from aesara import tensor as at
from matplotlib import pyplot as plt
def logp(value):
n = value.shape[-1]
r = at.sqrt(at.sum(value ** 2, axis=-1))
return at.switch(
r <= 1, at.log(at.gamma(n + 0.5) / np.pi ** (n / 2)), -np.inf
)
with pm.Model():
pm.DensityDist("x", logp=logp, shape=2)
step = pm.NUTS(target_accept=0.99, max_treedepth=4)
idata = pm.sample(step=step)
x = idata.posterior.x.stack(samples=["chain", "draw"]).values
plt.plot(x[0], x[1], "o", alpha=0.6)
You’ll see that the trace is plagued by divergences because HMC has trouble when it hits the edge of the ball. Nevertheless, the posterior kind of looks like what you want to get:
