{
    "nbformat": 4,
    "nbformat_minor": 5,
    "metadata": {
        "kernelspec": {
            "display_name": "Python",
            "language": "python",
            "name": "python"
        },
        "language_info": {
            "name": "python"
        }
    },
    "cells": [
        {
            "id": "0a443425",
            "cell_type": "markdown",
            "source": "<div class='alert alert-warning'>\n\nSciPy's interactive examples with Jupyterlite are experimental and may not always work as expected. Execution of cells containing imports may result in large downloads (up to 60MB of content for the first import from SciPy). Load times when importing from SciPy may take roughly 10-20 seconds. If you notice any problems, feel free to open an [issue](https://github.com/scipy/scipy/issues/new/choose).\n\n</div>",
            "metadata": {}
        },
        {
            "id": "d67f4901",
            "cell_type": "markdown",
            "source": "**Visualization of the probability density**\n\nPlot the probability density in three dimensions for increasing\nconcentration parameter. The density is calculated by the ``pdf``\nmethod.\n",
            "metadata": {}
        },
        {
            "id": "751797f1",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "import numpy as np\nimport matplotlib.pyplot as plt\nfrom scipy.stats import vonmises_fisher\nfrom matplotlib.colors import Normalize\nn_grid = 100\nu = np.linspace(0, np.pi, n_grid)\nv = np.linspace(0, 2 * np.pi, n_grid)\nu_grid, v_grid = np.meshgrid(u, v)\nvertices = np.stack([np.cos(v_grid) * np.sin(u_grid),\n                     np.sin(v_grid) * np.sin(u_grid),\n                     np.cos(u_grid)],\n                    axis=2)\nx = np.outer(np.cos(v), np.sin(u))\ny = np.outer(np.sin(v), np.sin(u))\nz = np.outer(np.ones_like(u), np.cos(u))\ndef plot_vmf_density(ax, x, y, z, vertices, mu, kappa):\n    vmf = vonmises_fisher(mu, kappa)\n    pdf_values = vmf.pdf(vertices)\n    pdfnorm = Normalize(vmin=pdf_values.min(), vmax=pdf_values.max())\n    ax.plot_surface(x, y, z, rstride=1, cstride=1,\n                    facecolors=plt.cm.viridis(pdfnorm(pdf_values)),\n                    linewidth=0)\n    ax.set_aspect('equal')\n    ax.view_init(azim=-130, elev=0)\n    ax.axis('off')\n    ax.set_title(rf\"$\\kappa={kappa}$\")\nfig, axes = plt.subplots(nrows=1, ncols=3, figsize=(9, 4),\n                         subplot_kw={\"projection\": \"3d\"})\nleft, middle, right = axes\nmu = np.array([-np.sqrt(0.5), -np.sqrt(0.5), 0])\nplot_vmf_density(left, x, y, z, vertices, mu, 5)\nplot_vmf_density(middle, x, y, z, vertices, mu, 20)\nplot_vmf_density(right, x, y, z, vertices, mu, 100)\nplt.subplots_adjust(top=1, bottom=0.0, left=0.0, right=1.0, wspace=0.)\nplt.show()",
            "outputs": []
        },
        {
            "id": "96dc5dd1",
            "cell_type": "markdown",
            "source": "As we increase the concentration parameter, the points are getting more\nclustered together around the mean direction.\n\n**Sampling**\n\nDraw 5 samples from the distribution using the ``rvs`` method resulting\nin a 5x3 array.\n",
            "metadata": {}
        },
        {
            "id": "52cece95",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "rng = np.random.default_rng()\nmu = np.array([0, 0, 1])\nsamples = vonmises_fisher(mu, 20).rvs(5, random_state=rng)\nsamples",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([[ 0.3884594 , -0.32482588,  0.86231516],\n       [ 0.00611366, -0.09878289,  0.99509023],\n       [-0.04154772, -0.01637135,  0.99900239],\n       [-0.14613735,  0.12553507,  0.98126695],\n       [-0.04429884, -0.23474054,  0.97104814]])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "454cbdb6",
            "cell_type": "markdown",
            "source": "These samples are unit vectors on the sphere $S^2$. To verify,\nlet us calculate their euclidean norms:\n",
            "metadata": {}
        },
        {
            "id": "02881043",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "np.linalg.norm(samples, axis=1)",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "array([1., 1., 1., 1., 1.])"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "d2797889",
            "cell_type": "markdown",
            "source": "Plot 20 observations drawn from the von Mises-Fisher distribution for\nincreasing concentration parameter $\\kappa$. The red dot highlights\nthe mean direction $\\mu$.\n",
            "metadata": {}
        },
        {
            "id": "4103dfe4",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "def plot_vmf_samples(ax, x, y, z, mu, kappa):\n    vmf = vonmises_fisher(mu, kappa)\n    samples = vmf.rvs(20)\n    ax.plot_surface(x, y, z, rstride=1, cstride=1, linewidth=0,\n                    alpha=0.2)\n    ax.scatter(samples[:, 0], samples[:, 1], samples[:, 2], c='k', s=5)\n    ax.scatter(mu[0], mu[1], mu[2], c='r', s=30)\n    ax.set_aspect('equal')\n    ax.view_init(azim=-130, elev=0)\n    ax.axis('off')\n    ax.set_title(rf\"$\\kappa={kappa}$\")\nmu = np.array([-np.sqrt(0.5), -np.sqrt(0.5), 0])\nfig, axes = plt.subplots(nrows=1, ncols=3,\n                         subplot_kw={\"projection\": \"3d\"},\n                         figsize=(9, 4))\nleft, middle, right = axes\nplot_vmf_samples(left, x, y, z, mu, 5)\nplot_vmf_samples(middle, x, y, z, mu, 20)\nplot_vmf_samples(right, x, y, z, mu, 100)\nplt.subplots_adjust(top=1, bottom=0.0, left=0.0,\n                    right=1.0, wspace=0.)\nplt.show()",
            "outputs": []
        },
        {
            "id": "903b58d1",
            "cell_type": "markdown",
            "source": "The plots show that with increasing concentration $\\kappa$ the\nresulting samples are centered more closely around the mean direction.\n\n**Fitting the distribution parameters**\n\nThe distribution can be fitted to data using the ``fit`` method returning\nthe estimated parameters. As a toy example let's fit the distribution to\nsamples drawn from a known von Mises-Fisher distribution.\n",
            "metadata": {}
        },
        {
            "id": "6a9a6819",
            "cell_type": "code",
            "metadata": {},
            "execution_count": null,
            "source": "mu, kappa = np.array([0, 0, 1]), 20\nsamples = vonmises_fisher(mu, kappa).rvs(1000, random_state=rng)\nmu_fit, kappa_fit = vonmises_fisher.fit(samples)\nmu_fit, kappa_fit",
            "outputs": [
                {
                    "output_type": "execute_result",
                    "metadata": {},
                    "data": {
                        "text/plain": "(array([0.01126519, 0.01044501, 0.99988199]), 19.306398751730995)"
                    },
                    "execution_count": null
                }
            ]
        },
        {
            "id": "26809b0f",
            "cell_type": "markdown",
            "source": "We see that the estimated parameters `mu_fit` and `kappa_fit` are\nvery close to the ground truth parameters.\n",
            "metadata": {}
        }
    ]
}