{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aefba47a",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "outputs": [],
   "source": [
    "# ruff: noqa: E402"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b90e8cb",
   "metadata": {},
   "source": [
    "# Estimating parameters of an anharmonic oscillator\n",
    "\n",
    "The anharnomic oscillator can be modelled by a non-linear partial differential\n",
    "equation as described in section 6.3.4 of the book [Fundamentals of Algorithms\n",
    "and Data Assimilation](https://www.amazon.com/Data-Assimilation-Methods-Algorithms-Applications/dp/1611974534)\n",
    "by Mark Asch, Marc Bocquet and Maëlle Nodet.\n",
    "\n",
    "-------------\n",
    "\n",
    "The discrete two-parameter model is:\n",
    "\n",
    "$$x_{k+1} - 2 x_{k} + x_{k-1} = \\omega^2 x_{k} - \\lambda^2 x_{k}^3$$\n",
    "\n",
    "with initial conditions $x_0 = 0$ and $x_1 = 1$.\n",
    "\n",
    "In other words we have the following recurrence relationship:\n",
    "\n",
    "$x_{k+1} = \\omega^2 x_{k} - \\lambda^2 x_{k}^3 + 2 x_{k} - x_{k-1}$\n",
    "\n",
    "\n",
    "-------------\n",
    "The state vector can be written as:\n",
    "\n",
    "$$\n",
    "\\mathbf{u}_k = \\begin{bmatrix}\n",
    "x_{k} \\\\\n",
    "x_{k-1}\n",
    "\\end{bmatrix}\n",
    "\\quad\n",
    "\\mathcal{M}_{k+1} =\n",
    " \\begin{bmatrix}\n",
    "2 + \\omega^2 - \\lambda^2 x_k^2 & -1 \\\\\n",
    "1 & 0\n",
    "\\end{bmatrix}\n",
    "$$\n",
    "\n",
    "so that $\\mathbf{u}_{k+1} = \\mathcal{M}_{k+1}(\\mathbf{u}_k)$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce9e725",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib import pyplot as plt\n",
    "from scipy import stats\n",
    "\n",
    "import iterative_ensemble_smoother as ies\n",
    "\n",
    "rng = np.random.default_rng(12345)\n",
    "\n",
    "plt.rcParams[\"figure.figsize\"] = [8, 3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73052f52",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def plot_result(A, response_x_axis, title=None):\n",
    "    \"\"\"Plot the anharmonic oscillator, as well as marginal\n",
    "    and joint distributions for the parameters.\"\"\"\n",
    "\n",
    "    responses = forward_model(A, response_x_axis)\n",
    "\n",
    "    fig, axes = plt.subplots(1, 2 + len(A), figsize=(9, 2.25))\n",
    "    if title:\n",
    "        fig.suptitle(title)\n",
    "\n",
    "    axes[0].plot(response_x_axis, responses, color=\"black\", alpha=0.1)\n",
    "    for ax, param, label in zip(axes[1:], A, [\"omega\", \"lambda\"]):\n",
    "        ax.hist(param, bins=\"fd\")\n",
    "        ax.set_xlabel(label)\n",
    "\n",
    "    axes[-1].scatter(A[0, :], A[1, :], s=5)\n",
    "\n",
    "    fig.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e64dd78",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd8ba791",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "def generate_observations(K):\n",
    "    \"\"\"Run the model with true parameters, then generate observations.\"\"\"\n",
    "    # Evaluate using true parameter values on a fine grid with K sample points\n",
    "    x = simulate_anharmonic(omega=3.5e-2, lmbda=3e-4, K=K)\n",
    "\n",
    "    # Generate observations every `nobs` steps\n",
    "    nobs = 50\n",
    "    # Generate observation points [0, 50, 100, 150, ...] when `nobs` = 50\n",
    "    observation_x_axis = np.arange(K // nobs) * nobs\n",
    "\n",
    "    # Generate noisy observations, with value at least 5\n",
    "    observation_errors = np.maximum(5, 0.2 * np.abs(x[observation_x_axis]))\n",
    "    observation_values = rng.normal(loc=x[observation_x_axis], scale=observation_errors)\n",
    "\n",
    "    return observation_values, observation_errors, observation_x_axis\n",
    "\n",
    "\n",
    "def simulate_anharmonic(omega, lmbda, K):\n",
    "    \"\"\"Evaluate the anharmonic oscillator.\"\"\"\n",
    "    x = np.zeros(K)\n",
    "    x[0] = 0\n",
    "    x[1] = 1\n",
    "\n",
    "    # Looping from 2 because we have initial conditions at k=0 and k=1.\n",
    "    for k in range(2, K):\n",
    "        x[k] = x[k - 1] * (2 + omega**2 - lmbda**2 * x[k - 1] ** 2) - x[k - 2]\n",
    "\n",
    "    return x\n",
    "\n",
    "\n",
    "def forward_model(A, response_x_axis):\n",
    "    \"\"\"Evaluate on each column (ensemble realization).\"\"\"\n",
    "    return np.vstack(\n",
    "        [simulate_anharmonic(*params, K=len(response_x_axis)) for params in A.T]\n",
    "    ).T\n",
    "\n",
    "\n",
    "response_x_axis = range(2500)\n",
    "observation_values, observation_errors, observation_x_axis = generate_observations(\n",
    "    len(response_x_axis)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69cc3195",
   "metadata": {},
   "source": [
    "## Plot observations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eb80878",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 3))\n",
    "ax.set_title(\"Anharmonic Oscillator\")\n",
    "ax.plot(\n",
    "    np.arange(2500),\n",
    "    simulate_anharmonic(omega=3.5e-2, lmbda=3e-4, K=2500),\n",
    "    label=\"Truth\",\n",
    "    color=\"black\",\n",
    ")\n",
    "\n",
    "ax.scatter(observation_x_axis, observation_values, label=\"Observations\")\n",
    "\n",
    "fig.tight_layout()\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69aad59c",
   "metadata": {},
   "source": [
    "## Create and plot prior\n",
    "\n",
    "\n",
    "This section shows a \"trick\" for working with non-normal priors.\n",
    "\n",
    "Let $g$ be the forward model, so that $y = g(x)$.\n",
    "Assume for instance that $x$ must be a positive parameter.\n",
    "If so, then we can for instance use an exponential prior on $x$.\n",
    "But if we use samples from an exponential prior directly in an ensemble smoother,\n",
    "there is no guarantee that the posterior samples will be positive.\n",
    "\n",
    "The trick is to sample $x \\sim \\mathcal{N}(0, 1)$,\n",
    "then define a function $f$ that maps from standard normal to the\n",
    "exponential distribution.\n",
    "This function can be constructed by first mapping from standard normal to\n",
    "the interval $[0, 1)$ using the CDF, then mapping to exponential using the\n",
    "quantile function (inverse CDF) of the exponential distribution.\n",
    "\n",
    "In summary:\n",
    "\n",
    "- We send $x \\sim \\mathcal{N}(0, 1)$ into the ensemble smoother as before\n",
    "- We use the composite function $g(f(x))$ as our forward model, instead of $g(x)$\n",
    "- When we plot the prior and posterior, we plot $f(x)$ instead of $x$ directly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d97f407",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create Uniform prior distributions\n",
    "realizations = 100\n",
    "PRIORS = [stats.uniform(loc=0.025, scale=0.02), stats.uniform(loc=0.0002, scale=0.0002)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4adf856",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prior_to_standard_normal(A):\n",
    "    \"\"\"Map prior to U(0, 1), then use inverse of CDF to map to N(0, 1).\"\"\"\n",
    "    return np.vstack(\n",
    "        [stats.norm().ppf(prior.cdf(A_param)) for (A_param, prior) in zip(A, PRIORS)]\n",
    "    )\n",
    "\n",
    "\n",
    "def standard_normal_to_prior(A):\n",
    "    \"\"\"Reverse mapping.\"\"\"\n",
    "    return np.vstack(\n",
    "        [prior.ppf(stats.norm().cdf(A_param)) for (A_param, prior) in zip(A, PRIORS)]\n",
    "    )\n",
    "\n",
    "\n",
    "# verify that mappings are invertible\n",
    "A_uniform = np.vstack([prior.rvs(realizations) for prior in PRIORS])\n",
    "assert np.allclose(\n",
    "    standard_normal_to_prior(prior_to_standard_normal(A_uniform)), A_uniform\n",
    ")\n",
    "assert np.allclose(\n",
    "    prior_to_standard_normal(standard_normal_to_prior(A_uniform)), A_uniform\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e34efc8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a standard normal prior and plot it in transformed space\n",
    "A = rng.standard_normal(size=(2, realizations))\n",
    "plot_result(standard_normal_to_prior(A), response_x_axis)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71da149e",
   "metadata": {},
   "source": [
    "## A single update step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a936f188",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_result(standard_normal_to_prior(A), response_x_axis)\n",
    "\n",
    "# The forward model is composed with the transformation from N(0, 1) to uniform priors\n",
    "responses_before = forward_model(standard_normal_to_prior(A), response_x_axis)\n",
    "Y = responses_before[observation_x_axis]\n",
    "\n",
    "# Run smoother for one step\n",
    "smoother = ies.SIES(\n",
    "    parameters=A,\n",
    "    covariance=observation_errors**2,\n",
    "    observations=observation_values,\n",
    "    seed=42,\n",
    ")\n",
    "\n",
    "new_A = smoother.sies_iteration(Y, step_length=1.0)\n",
    "\n",
    "\n",
    "plot_result(standard_normal_to_prior(new_A), response_x_axis)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9237a5b",
   "metadata": {},
   "source": [
    "## Iterative smoother"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "937ff571",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "from iterative_ensemble_smoother.utils import steplength_exponential\n",
    "\n",
    "\n",
    "def iterative_smoother(A):\n",
    "    A_current = np.copy(A)\n",
    "    iterations = 4\n",
    "    smoother = ies.SIES(\n",
    "        parameters=A,\n",
    "        covariance=observation_errors**2,\n",
    "        observations=observation_values,\n",
    "        seed=42,\n",
    "    )\n",
    "\n",
    "    plot_result(\n",
    "        standard_normal_to_prior(A_current),\n",
    "        response_x_axis,\n",
    "        title=\"Prior\",\n",
    "    )\n",
    "\n",
    "    for iteration in range(iterations):\n",
    "        responses_before = forward_model(\n",
    "            standard_normal_to_prior(A_current), response_x_axis\n",
    "        )\n",
    "\n",
    "        Y = responses_before[observation_x_axis]\n",
    "\n",
    "        A_current = smoother.sies_iteration(\n",
    "            Y, step_length=steplength_exponential(iteration + 1)\n",
    "        )\n",
    "\n",
    "        plot_result(\n",
    "            standard_normal_to_prior(A_current),\n",
    "            response_x_axis,\n",
    "            title=f\"SIES iteration {iteration+1}\",\n",
    "        )\n",
    "\n",
    "\n",
    "iterative_smoother(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "408c2e7d",
   "metadata": {},
   "source": [
    "## ES-MDA (Ensemble Smoother - Multiple Data Assimilation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b35cd78",
   "metadata": {},
   "outputs": [],
   "source": [
    "smoother = ies.ESMDA(\n",
    "    # If a 1D array is passed, it is interpreted\n",
    "    # as the diagonal of the covariance matrix.\n",
    "    covariance=observation_errors**2,\n",
    "    observations=observation_values,\n",
    "    # The inflation factors used in ESMDA\n",
    "    # They are scaled so that sum_i alpha_i^-1 = 1\n",
    "    alpha=np.array([8, 4, 2, 1]),\n",
    "    seed=42,\n",
    ")\n",
    "\n",
    "\n",
    "A_current = np.copy(A)\n",
    "\n",
    "plot_result(standard_normal_to_prior(A_current), response_x_axis, title=\"Prior\")\n",
    "\n",
    "for iteration in range(smoother.num_assimilations()):\n",
    "    # Evaluate the model\n",
    "    responses_before = forward_model(\n",
    "        standard_normal_to_prior(A_current), response_x_axis\n",
    "    )\n",
    "    Y = responses_before[observation_x_axis]\n",
    "\n",
    "    # Assimilate data\n",
    "    A_current = smoother.assimilate(A_current, Y)\n",
    "\n",
    "    # Plot\n",
    "    plot_result(\n",
    "        standard_normal_to_prior(A_current),\n",
    "        response_x_axis,\n",
    "        title=f\"ESMDA iteration {iteration+1}\",\n",
    "    )"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}