Plotting log-normal probability distributions with varying parameters (Python 3.6+)

log-normal.py

# MIT License
# (c) Yuki Koyama

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize
import seaborn as sns
import math


def log_normal_distribution(x, mu, sigma_2):
    r = math.log(x) - mu
    return (1.0 / (x * math.sqrt(2.0 * math.pi * sigma_2))) * math.exp(-0.5 * (r * r) / sigma_2)


def create_varying_mu():
    # Define the resolution of mu variation
    num_samples_per_graph = 11

    # Define the list of sigma^2 values
    sigma_2_list = [0.05, 0.25, 1.00]

    # Define the x range (note: we cannot include x = 0 here)
    x = np.arange(0.001, 3.001, 0.001)

    # Prepare for graph plotting
    sns.set()
    sns.set_palette("inferno", num_samples_per_graph)
    sns.set_context("paper")
    fig = plt.figure(figsize=(4 * len(sigma_2_list), 4), dpi=300)

    # Iterate over sigma^2
    for i, sigma_2 in enumerate(sigma_2_list):
        ax = fig.add_subplot(1, len(sigma_2_list), i + 1)

        for j in range(num_samples_per_graph):
            # Define the mu value (from 0.5 to 1.5 when num_samples_per_graph = 11)
            mu = math.log(0.1 * (j + 5))

            # Calculate probability density function values
            y = x.copy()
            for y_elem in np.nditer(y, op_flags=['readwrite']):
                y_elem[...] = log_normal_distribution(y_elem, mu, sigma_2)

            # Create a subplot
            ax.set_title("Log-Normal Distributions ($\sigma^{2} = " + f"{sigma_2:.2f}" + "$)")
            ax.plot(x, y, label="$\mu = \log(" + f"{math.exp(mu):.1f}" + ")$")
            ax.legend(loc='best')

    # Save to a file
    fig.savefig('log-normal-varying-mu.png', bbox_inches='tight')


def create_varying_sigma_2():
    # Define the resolution of sigma^2 variation
    num_samples_per_graph = 44

    # Define the mu value
    mu = math.log(1.0)

    # Define the x range (note: we cannot include x = 0 here)
    x = np.arange(0.001, 2.001, 0.001)

    # Prepare for graph plotting
    sns.set()
    sns.set_palette("inferno", num_samples_per_graph)
    sns.set_context("paper")
    fig = plt.figure(figsize=(6, 4), dpi=300)
    ax = fig.add_subplot(1, 1, 1)

    for j in range(num_samples_per_graph):
        # Define the sigma^2 value
        sigma_2 = 0.10 * (j + 1)

        # Calculate probability density function values
        y = x.copy()
        for y_elem in np.nditer(y, op_flags=['readwrite']):
            y_elem[...] = log_normal_distribution(y_elem, mu, sigma_2)

        # Create a subplot
        ax.set_title("Log-Normal Distributions ($\mu = \log(" + f"{math.exp(mu):.1f}" + ")$)")
        ax.plot(x, y, label="$\sigma^{2} = " + f"{sigma_2:.2f}" + "$")

    # Set a colorbar legend
    norm = Normalize(0.10 * 1, 0.10 * (num_samples_per_graph + 1))
    mappable = ScalarMappable(cmap="inferno", norm=norm)
    mappable.set_array([])
    ticks = np.linspace(norm.vmin, norm.vmax, 5) # Need to tweak depending on num_samples_per_graph
    tick_labels = []
    for ticks_elem in np.nditer(ticks):
        tick_labels.append("$\sigma^{2} = " + f"{ticks_elem:.1f}" + "$")
    cbar = fig.colorbar(mappable, orientation='vertical')
    cbar.set_ticks(ticks)
    cbar.set_ticklabels(tick_labels)
    cbar.ax.tick_params(width=0, pad=0)

    # Save to a file
    fig.savefig('log-normal-varying-sigma-2.png', bbox_inches='tight')


if __name__ == "__main__":
    create_varying_mu()
    create_varying_sigma_2()

Its a beautiful thing. Thank you for helping me understand visually what the parameters do to the distribution