Samplers

The samplers module provides various sampling algorithms to generate samples from a given probability distribution. These algorithms can be used to sample a explicit (graph) or implicit (\(f(z)=0\)) distribution. The samplers can be used for various tasks such as: * build a dataset or mesh for a manifold * visualize a distribution * build a dataset of geodesics using one of the methods in geodesics.generate

Sample from an explicit manifold

The following code samples from an explicit manifold defined by a potential function \(\phi(x):\mathbb{R}^i\rightarrow \mathbb{R}^m\), where \(i\) is the number of (independent) inputs and \(m\) the number of constraints. The dimensionality of ambient space is \(n = i+m\). * Using randinput_expl we can obtain a random sample of the manifold by sampling uniformly via latin-hypercube in the input space and projecting the samples onto the manifold using \(\phi\). This method can be used to obtain a uniform sample of the input space, but it does not guarantee a uniform sample of the manifold itself, but it's fast. * Using volume_expl we can obtain a non-repeated uniform sample of the manifold. The function perform an over-sampling of \(phi\), which gets then re-sampled \(w_i \sim \sqrt{\text{det} [G(x_i)]}\) where \(G = J^T J\) is the metric tensor induced by \(\phi\).

import os
PLOT_W = 400
PLOT_H = 400
os.environ['JAX_PLATFORMS'] = "cpu"  # use "cuda" if you have a GPU
os.environ['JAX_ENABLE_X64'] = "1"

from jnlr.utils.samplers import volume_expl, randinput_expl
from jnlr.utils.manifolds import f_ackley as phi
from jnlr.utils.plot_utils import plot_3d_projection


Y_vol = volume_expl(phi, n_samples=500, oversample=5, roi_R=5)
fig = plot_3d_projection(Y_vol, phi, width=PLOT_W, height=PLOT_H)

# set view from the top
fig.update_layout(scene_camera=dict(eye=dict(x=0., y=0., z=2.5)))

Y_vol = randinput_expl(phi, n_samples=500)
fig = plot_3d_projection(Y_vol, phi, width=PLOT_W, height=PLOT_H)
# set view from the top
fig.update_layout(scene_camera=dict(eye=dict(x=0., y=0., z=2.5)))

from jnlr.utils.samplers import langevin_implicit

def f(xyz):
    x, y, z = xyz
    return x**4 - x**2 + y**2 + z**2 - 2*x*z
Y_lang = langevin_implicit(f, n_samples=3000, burn=100, thin=2, sigma=0.1, lam=0.01, kappa=0.02, R=3.0)
fig = plot_3d_projection(Y_lang, width=PLOT_W, height=PLOT_H)
fig.update_layout(scene_camera=dict(eye=dict(x=0., y=0., z=2.5)))

from jnlr.utils.implicit_hypersurfaces import f_dodecahedron
Y_lang = langevin_implicit(f_dodecahedron, n_samples=3000, burn=100, thin=2, sigma=0.1, lam=0.01, kappa=0.02, R=3.0)
fig = plot_3d_projection(Y_lang, width=PLOT_W, height=PLOT_H)
fig.update_layout(scene_camera=dict(eye=dict(x=0., y=0., z=2.5)))

from jnlr.utils.plot_utils import plot_3d_projection
from jnlr.utils.implicit_hypersurfaces import surface7, surface_b, surface_c

Y_lang = langevin_implicit(surface7, n_samples=5000, burn=100, thin=1, sigma=0.3, lam=0, kappa=0, R=5.0, tol=1e-2)
plot_3d_projection(Y_lang, width=PLOT_W, height=PLOT_H)

Y_lang = langevin_implicit(surface_c, n_samples=5000, burn=100, thin=1, sigma=0.5, lam=0.1, kappa=0, R=5.0, tol=1e-3)
plot_3d_projection(Y_lang, width=PLOT_W, height=PLOT_H)

Y_lang = langevin_implicit(surface_b, n_samples=5000, burn=100, thin=1, sigma=0.3, lam=0, kappa=0, R=5.0, tol=1e-2)
plot_3d_projection(Y_lang, width=PLOT_W, height=PLOT_H)