{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "8d31666e",
      "metadata": {
        "id": "8d31666e"
      },
      "source": [
        "# Exercise Session 3: Statistical Inference"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7e30f093",
      "metadata": {
        "id": "7e30f093"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from scipy.stats import norm\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f9a57bbd",
      "metadata": {
        "id": "f9a57bbd"
      },
      "source": [
        "## Exercise: Mass Estimation from Position Measurements in 1D Brownian Motion"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f01b1ee6",
      "metadata": {
        "id": "f01b1ee6"
      },
      "source": [
        "A Brownian particle of mass $m_\\star$ moves in one dimension in a fluid characterized by a friction coefficient $\\gamma$ and a temperature $T$. Its velocity $v(t)$ can be modelled as an Ornstein-Uhlenbeck process, that is as a random variable varying in time according to the stochastic differential equation ([SDE](https://en.wikipedia.org/wiki/Stochastic_differential_equation)):\n",
        "$$\n",
        "m_\\star \\, dv(t) = - \\gamma v(t) \\, dt + \\sqrt{2 \\gamma k_B T} \\, dW(t),\n",
        "$$\n",
        "where $k_B = 1.380649×10^{−23} \\mathrm{J/K}$ is the Boltzmann constant and $W(t)$ is a [Wiener process](https://en.wikipedia.org/wiki/Wiener_process). The solution of the above SDE is a random function of time following a Gaussian distribution, which at stationarity (i.e. $t \\to \\infty$) becomes the well-known [Maxwell-Boltzmann distribution](https://en.wikipedia.org/wiki/Maxwell%E2%80%93Boltzmann_distribution).\n",
        "\n",
        "In practice, we are interested in a model for the position of the particle in time, since the position of the Brownian particle is what we measure directly. Assuming that the velocity is stationary, for the stochastic process above the displacement $\\Delta x$ occurring in the time interval $\\Delta t$ is also a random Gaussian variable, with zero mean and variance given by the formula:\n",
        "$$\n",
        "\\mathrm{Var}\\{ \\Delta x \\} = \\sigma^2(m_\\star, \\gamma, T, \\Delta t) =  \\frac{2 k_B T}{\\gamma} \\biggl[ \\Delta t - \\frac{m_\\star}{\\gamma} \\Bigl( 1 - e^{-\\gamma \\Delta t / m_\\star} \\Bigr) \\biggr].\n",
        "$$\n",
        "\n",
        "Given these premises, in this exercise we will consider the problem of inferring the true mass $m_\\star$ from a series of position measurements realized when the system is stationary, given $\\gamma$ and $T$."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f550938d",
      "metadata": {
        "id": "f550938d"
      },
      "source": [
        "### Task 1: simulating the Brownian motion, i.e. generating the dataset of displacements"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "459f1642",
      "metadata": {
        "id": "459f1642"
      },
      "source": [
        "You observe the particle's position at equally spaced times $t_0, t_1, \\dots, t_n$, with $t_i - t_{i-1} = \\Delta t$. Using the following values for the parameters of the problem"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "94269bf3",
      "metadata": {
        "id": "94269bf3"
      },
      "outputs": [],
      "source": [
        "# Parameters\n",
        "m_star = 5.0e-10\n",
        "gamma = 1.0e-8\n",
        "temp = 300.0 # Temperature T\n",
        "dt = 0.01\n",
        "n = 200"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "49b8b189",
      "metadata": {
        "id": "49b8b189"
      },
      "source": [
        "simulate $n$ displacements $\\Delta x_1, \\dots, \\Delta x_n$ as i.i.d. (independent and identically distributed) Gaussian variables with zero mean and variance $\\sigma^2(m_\\star, \\gamma, T, \\Delta t)$. Then, construct the corresponding trajectory. For simplicity, assume that $x_0 = 0$."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "61640e1f",
      "metadata": {
        "id": "61640e1f"
      },
      "source": [
        "**1.1** Code the function $\\sigma^2(m, \\gamma, T, \\Delta t)$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "537e0451",
      "metadata": {
        "id": "537e0451"
      },
      "outputs": [],
      "source": [
        "def displacement_var(m, gamma, temp, dt, kB = 1.380649e-23):\n",
        "    ..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6044fee4",
      "metadata": {
        "id": "6044fee4"
      },
      "source": [
        "**1.2** Write a function that takes $m$, $\\gamma$, $T$, $\\Delta t$ and $n$ as inputs and generates the dataset of displacements as a vector of length $n$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "cd66cfe2",
      "metadata": {
        "id": "cd66cfe2"
      },
      "outputs": [],
      "source": [
        "def generate_dataset():\n",
        "    ..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "64343bd6",
      "metadata": {
        "id": "64343bd6"
      },
      "source": [
        "**1.3** Generate a dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "60b4499c",
      "metadata": {
        "id": "60b4499c"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "233b7292",
      "metadata": {
        "id": "233b7292"
      },
      "source": [
        "**1.4** As a sanity check, plot its histogram and the Gaussian p.d.f. $\\mathcal{N}(0, \\sigma^2(m_\\star, \\gamma, T, \\Delta t))$ from which the dataset should be drawn. Use reasonable limits for the $x$-axis.\n",
        "\n",
        "*Hint*: To plot the histogram you can use ```hist``` from ```matplotlib.pyplot``` (check [documentation](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.hist.html)). Set ```density = True``` to have the histogram represent a p.d.f. (remember that large $y$-values are expected when your p.d.f. is narrow). To plot the theoretical Gaussian you can employ the function ```norm``` from ```scipy.stats```, which has already been imported at the beginning of this notebook (check [documentation](https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.norm.html))."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9762e989",
      "metadata": {
        "id": "9762e989"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "6b16022b",
      "metadata": {
        "id": "6b16022b"
      },
      "source": [
        "**1.5** Plot the trajectory corresponding to your dataset.\n",
        "\n",
        "*Remark*: The plot will display **one realization** of the particle trajectory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7aa2f146",
      "metadata": {
        "id": "7aa2f146"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "10c57fe5",
      "metadata": {
        "id": "10c57fe5"
      },
      "source": [
        "### Task 2: estimation of $m_\\star$ through maximum likelihood"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e864b3cc",
      "metadata": {
        "id": "e864b3cc"
      },
      "source": [
        "In this section, we implement the maximum likelihood estimator (MLE) for the mass, and use it to estimate $m_\\star$ for the dataset that you generated."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3733201d",
      "metadata": {
        "id": "3733201d"
      },
      "source": [
        "**2.1** Prove that MLE $\\hat{m}$ is the solution to the following equation\n",
        "$$\n",
        "\\sigma^2(\\hat{m}, \\gamma, T, \\Delta t) = S_n^2\n",
        "$$\n",
        "where $S_n^2 = \\frac{1}{n} \\sum_{i = 1}^n (\\Delta x_i)^2$ is the (biased) estimator for the variance.\n",
        "\n",
        "*Hint*: It is useful to maximise the log-likelihood instead of the likelihood to convert products into sums (recall that logs do not alter the position of maxima)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "15c53b75",
      "metadata": {
        "id": "15c53b75"
      },
      "source": [
        "**2.2** Show that the previous equation can be cast as a fixed-point equation in $\\hat{m}$, i.e. an equation of the form $\\hat{m} = F(\\hat{m})$, with $F$ to be determined."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c451f1ea",
      "metadata": {
        "id": "c451f1ea"
      },
      "source": [
        "*Remark*: **What is the fixed-point iteration method?**\n",
        "\n",
        "The fixed-point iteration method is a simple numerical technique to solve equations of the form $x = F(x)$. The idea is:\n",
        "1. Start with an initial guess $x_0$.\n",
        "2. Compute a new value by plugging into the function: $x_{k+1} = F(x_k)$.\n",
        "3. Repeat until the sequence $\\{ x_k \\}$ stops changing (converges).\n",
        "\n",
        "In particular, regarding point 3., one usually needs to define a convergence criterion. A standard choice is to compute at each step the absolute difference $|x_{k+1} - x_{k}|$ and to claim convergence when this is smaller than a conventionally small enough tolerance. Then, $x_{k+1}$ is taken as the approximate solution to the fixed-point equation.\n",
        "\n",
        "However, convergence is not always guaranteed. If the function $F(x)$ is well-behaved (specifically, if $|F'(x)| < 1$ near the solution), then the iteration will converge to the fixed point $x^*$ that satisfies $x^* = F(x^*)$.\n",
        "\n",
        "In our problem, the MLE equation for $\\hat{m}$ cannot be solved in closed form, but we were able to rewrite it as a fixed-point equation. This allows us to approximate $\\hat{m}$ by repeatedly applying $F$ until the value stabilizes."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "39629847",
      "metadata": {
        "id": "39629847"
      },
      "source": [
        "**2.3** Assuming that $F(m)$ is well-behaved, write a function that computes the fixed point of the above equation using the fixed-point iteration method.\n",
        "\n",
        "*Hint*: A proper function must have an optional initial value for the mass, an optional value for the tolerance and an optional maximum number of iterations (for safety). Additionally, in the case of no convergence, it should raise an error or simply return ```None``` after printing a warning message."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1d275e7e",
      "metadata": {
        "id": "1d275e7e"
      },
      "outputs": [],
      "source": [
        "def fixed_point_estimator():\n",
        "    ..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bdcd227e",
      "metadata": {
        "id": "bdcd227e"
      },
      "source": [
        "**2.4** Compute $\\hat{m}$ (via fixed-point iteration) on the dataset that you generated. Then plot the log-likelihood of your dataset as a function of the parameter $m$ together with the two vertical lines $m = \\hat{m}$ and $m = m_\\star$. Add a legend for clarity.\n",
        "\n",
        "*Hint*: Take a reasonable initial value for $m$ (that means positive). Play with the tolerance value and the maximum number of iterations in your fixed-point iteration in order to get the best approximation of the MLE. Use a log scale for the $x$-axis. It is helpful to use different line-styles for the two vertical lines.\n",
        "\n",
        "*Remark*: Again, the point that you obtain via fixed-point iteration is an **approximation** for the point maximizing the (log-)likelihood, namely the MLE. The latter in turn is what we use to estimate the ground-truth value of the mass. In general, there is no reason to expect the MLE to give a \"very good\" estimate of the true parameter (we will come back to this point in the last task of the exercise)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "9677f05a",
      "metadata": {
        "id": "9677f05a"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ec689a9a",
      "metadata": {
        "id": "ec689a9a"
      },
      "source": [
        "### Task 3: likelihood flattening in the diffusive regime"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "caec01df",
      "metadata": {
        "id": "caec01df"
      },
      "source": [
        "For an Ornstein-Uhlenbeck velocity process we can distinguish two interesting regimes:\n",
        "- If $\\gamma \\Delta t / m_\\star << 1$, the motion is *ballistic* (dominated by inertia).\n",
        "- If $\\gamma \\Delta t / m_\\star >> 1$, the motion is *diffusive* (dominated by diffusion)."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c30841f5",
      "metadata": {
        "id": "c30841f5"
      },
      "source": [
        "**3.1** Check that the values for $m_\\star$, $\\gamma$ and $\\Delta t$ considered until now do <u>not</u> correspond to a diffusive regime."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "152c64d7",
      "metadata": {
        "id": "152c64d7"
      },
      "source": [
        "**3.2** Replace the previously employed value for $m_\\star$ with"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "60657ab1",
      "metadata": {
        "id": "60657ab1"
      },
      "outputs": [],
      "source": [
        "m_star2 = 5.0e-15"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a41b93f6",
      "metadata": {
        "id": "a41b93f6"
      },
      "source": [
        "and repeat the same check. Are you still out of the diffusive regime with this new value of $m_\\star$?"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "e790adca",
      "metadata": {
        "id": "e790adca"
      },
      "source": [
        "**3.3** Using the same $\\gamma$, $T$, $\\Delta t$ and $n$ as before, generate a dataset of displacements using the new value for the ground-truth mass. Then plot the corresponding log-likelihood as a function of $m$. Do you observe any difference w.r.t. the last plot? Is maximizing the (log-)likelihood still a reasonable approach in this situation?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "80abe161",
      "metadata": {
        "id": "80abe161"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "2a48dd50",
      "metadata": {
        "id": "2a48dd50"
      },
      "source": [
        "**3.4** Compute (via fixed-point iteration) the MLE on the dataset generated with the new value of $m_\\star$. Play with the tolerance value and the maximum number of iterations. Are you able to obtain a reasonable estimate, that is one of the same order of magnitude of the true value?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "41cb68f4",
      "metadata": {
        "id": "41cb68f4"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b05bc245",
      "metadata": {
        "id": "b05bc245"
      },
      "source": [
        "**3.5** Find an explanation for the results that you have just found.\n",
        "\n",
        "*Hint*: One way is to rewrite $\\sigma^2(m_\\star, \\gamma, T, \\Delta t)$ for $\\gamma \\Delta t / m_\\star >> 1$."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "3270760b",
      "metadata": {
        "id": "3270760b"
      },
      "source": [
        "### Task 4: dependence on the dataset size"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a130a34e",
      "metadata": {
        "id": "a130a34e"
      },
      "source": [
        "Let's return to the initial setting of parameters. What happens if we have a small number of observations?"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "99c18c73",
      "metadata": {
        "id": "99c18c73"
      },
      "source": [
        "**4.1** Generate a dataset of size $n = 20$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ee949773",
      "metadata": {
        "id": "ee949773"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a6895b23",
      "metadata": {
        "id": "a6895b23"
      },
      "source": [
        "**4.2** Repeat task 2.4 on this new dataset. What happens to the MLE?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "5541d9fb",
      "metadata": {
        "id": "5541d9fb"
      },
      "outputs": [],
      "source": [
        "..."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "eaaedaf2",
      "metadata": {
        "id": "eaaedaf2"
      },
      "source": [
        "**4.3** Find an explanation for the results that you have just found.\n",
        "\n",
        "*Hint*: The law of large numbers is involved."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "base",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.12.2"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}