# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:pyprobml]
#     language: python
#     name: conda-env-pyprobml-py
# ---

# # Discrete Probability Distribution Plot

# +
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import seaborn as sns

try:
    from probml_utils import savefig, latexify
except ModuleNotFoundError:
    # %pip install -qq git+https://github.com/probml/probml-utils.git
    from probml_utils import savefig, latexify

# + tags=["hide-input"]
latexify(width_scale_factor=2)

# +
# Bar graphs showing a uniform discrete distribution and another with full mass on one value.

N = 4


def make_graph(probs, N, save_name, fig=None, ax=None):
    x = jnp.arange(1, N + 1)
    if fig is None:
        fig, ax = plt.subplots()
    ax.bar(x, probs, align="center")
    ax.set_xlim([min(x) - 0.5, max(x) + 0.5])
    ax.set_xticks(x)
    ax.set_yticks(jnp.linspace(0, 1, N + 1))
    ax.set_xlabel("$x$")
    ax.set_ylabel("$Pr(X=x)$")
    sns.despine()
    if len(save_name) > 0:
        savefig(save_name)
    return fig, ax


uniform_probs = jnp.repeat(1.0 / N, N)
_, _ = make_graph(
    uniform_probs, N, "uniform_histogram_latexified"
)  # Do not add .pdf or .png as it is automatically added by savefig method

delta_probs = jnp.array([1, 0, 0, 0])
_, _ = make_graph(delta_probs, N, "delta_histogram_latexified");
# -

# ## Demo
#
# You can see different examples of discrete distributions by changing the seed in the following demo.

# +
from ipywidgets import interact


@interact(random_state=(1, 10), N=(2, 10))
def generate_random(random_state, N):
    key = jax.random.PRNGKey(random_state)
    probs = jax.random.uniform(key, shape=(N,))
    probs = probs / jnp.sum(probs)
    fig, ax = make_graph(probs, N, "")
    ax.set_yticks(jnp.linspace(0, 1, 11))

StarcoderdataJupyterScriptsDedupFiltered