{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7341a367",
   "metadata": {},
   "outputs": [],
   "source": [
    "# ruff: noqa: E402"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1f7b359",
   "metadata": {},
   "source": [
    "# Adaptive Localization\n",
    "\n",
    "In this example we run adaptive localization\n",
    "on a sparse Gauss-linear problem.\n",
    "\n",
    "- Each response is only affected by $3$ parameters.\n",
    "- This is represented by a tridiagonal matrix $A$ in the forward model $g(x) = Ax$.\n",
    "- The problem is Gauss-Linear, so in this case ESMDA will sample\n",
    "  the true posterior when the number of ensemble members (realizations) is large.\n",
    "- The sparse correlation structure will lead to spurious correlations, in the sense\n",
    "  that a parameter and response might appear correlated when in fact they are not."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b782f5",
   "metadata": {},
   "source": [
    "## Import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30d4ee53",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy as sp\n",
    "from matplotlib import pyplot as plt\n",
    "\n",
    "from iterative_ensemble_smoother import ESMDA\n",
    "from iterative_ensemble_smoother.experimental import AdaptiveESMDA"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd9c64d0",
   "metadata": {},
   "source": [
    "## Create problem data\n",
    "\n",
    "Some settings worth experimenting with:\n",
    "\n",
    "- Decreasing `prior_std=1` will pull the posterior solution toward zero.\n",
    "- Increasing `num_ensemble` will increase the quality of the solution.\n",
    "- Increasing `num_observations / num_parameters`\n",
    "  will increase the quality of the solution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9afa268c",
   "metadata": {},
   "outputs": [],
   "source": [
    "rng = np.random.default_rng(42)\n",
    "\n",
    "# Dimensionality of the problem\n",
    "num_parameters = 100\n",
    "num_observations = 50\n",
    "num_ensemble = 25\n",
    "prior_std = 1\n",
    "\n",
    "# Number of iterations to use in ESMDA\n",
    "alpha = 5"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f7a9b56",
   "metadata": {},
   "source": [
    "## Create problem data - sparse tridiagonal matrix $A$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d2e259",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "diagonal = np.ones(min(num_parameters, num_observations))\n",
    "\n",
    "# Create a tridiagonal matrix (easiest with scipy)\n",
    "A = sp.sparse.diags(\n",
    "    [diagonal, diagonal, diagonal],\n",
    "    offsets=[-1, 0, 1],\n",
    "    shape=(num_observations, num_parameters),\n",
    "    dtype=float,\n",
    ").toarray()\n",
    "\n",
    "# We add some noise that is insignificant compared to the\n",
    "# actual local structure in the forward model\n",
    "A = A + rng.standard_normal(size=A.shape) * 0.01\n",
    "\n",
    "plt.title(\"Linear mapping $A$ in forward model $g(x) = Ax$\")\n",
    "plt.imshow(A)\n",
    "plt.xlabel(\"Parameters (inputs)\")\n",
    "plt.ylabel(\"Responses (outputs)\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3371f5d",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "- Below we draw prior realizations $X \\sim N(0, \\sigma)$.\n",
    "- The true parameter values used to generate observations are in the range $[-1, 1]$.\n",
    "- As the number of realizations (ensemble members) goes to infinity,\n",
    "  the correlation between the prior and the true parameter values converges to zero.\n",
    "- The correlation is zero for a finite number of realizations too,\n",
    "  but statistical noise might induce some spurious correlations between\n",
    "  the prior and the true parameter values.\n",
    "\n",
    "**In summary the baseline correlation is zero.**\n",
    "Anything we can do to increase the correlation above zero beats the baseline,\n",
    "which is using the prior (no update).\n",
    "\n",
    "The correlation coefficient does not take into account the uncertainty\n",
    "represeted in the posterior, only the mean posterior value is compared with\n",
    "the true parameter values.\n",
    "To compare distributions we could use the\n",
    "[Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence),\n",
    "but we do not pursue this here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3d08ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_prior(num_parameters, num_ensemble, prior_std):\n",
    "    \"\"\"Sample prior from N(0, prior_std).\"\"\"\n",
    "    return rng.normal(size=(num_parameters, num_ensemble)) * prior_std\n",
    "\n",
    "\n",
    "def g(X):\n",
    "    \"\"\"Apply the forward model.\"\"\"\n",
    "    return A @ X\n",
    "\n",
    "\n",
    "# Create observations: obs = g(x) + N(0, 1)\n",
    "x_true = np.linspace(-1, 1, num=num_parameters)\n",
    "observation_noise = rng.standard_normal(size=num_observations)  # N(0, 1) noise\n",
    "observations = g(x_true) + observation_noise\n",
    "\n",
    "# Initial ensemble X ~ N(0, prior_std) and diagonal covariance with ones\n",
    "X = get_prior(num_parameters, num_ensemble, prior_std)\n",
    "\n",
    "# Covariance matches the noise added to observations above\n",
    "covariance = np.ones(num_observations)  # N(0, 1) covariance"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1663291",
   "metadata": {},
   "source": [
    "## Solve the maximum likelihood problem\n",
    "\n",
    "We can solve $Ax = b$, where $b$ are the observations,\n",
    "for the maximum likelihood estimate.\n",
    "\n",
    "Notice that unlike using a Ridge model,\n",
    "solving $Ax = b$ directly does not use any prior information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ca5b10e",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_ml, *_ = np.linalg.lstsq(A, observations, rcond=None)\n",
    "\n",
    "plt.figure(figsize=(7, 3))\n",
    "plt.scatter(np.arange(len(x_true)), x_true, label=\"True parameter values\")\n",
    "plt.scatter(np.arange(len(x_true)), x_ml, label=\"ML estimate (no prior)\")\n",
    "plt.xlabel(\"Parameter index\")\n",
    "plt.ylabel(\"Parameter value\")\n",
    "plt.grid(True, ls=\"--\", zorder=0, alpha=0.33)\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c73cefc",
   "metadata": {},
   "source": [
    "## Solve using ESMDA\n",
    "\n",
    "We create an `ESMDA` instance and solve the Guass-linear problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19174a10",
   "metadata": {},
   "outputs": [],
   "source": [
    "smoother = ESMDA(\n",
    "    covariance=covariance,\n",
    "    observations=observations,\n",
    "    alpha=alpha,\n",
    "    seed=1,\n",
    ")\n",
    "\n",
    "X_i = np.copy(X)\n",
    "for i, alpha_i in enumerate(smoother.alpha, 1):\n",
    "    print(\n",
    "        f\"ESMDA iteration {i}/{smoother.num_assimilations()}\"\n",
    "        + f\" with inflation factor alpha_i={alpha_i}\"\n",
    "    )\n",
    "    X_i = smoother.assimilate(X_i, Y=g(X_i))\n",
    "\n",
    "\n",
    "X_posterior = np.copy(X_i)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c477ce6",
   "metadata": {},
   "source": [
    "## Plot and compare solutions\n",
    "\n",
    "Compare the true parameters with both the ML estimate\n",
    "from linear regression and the posterior means obtained using `ESMDA`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc746bee",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 3))\n",
    "plt.scatter(np.arange(len(x_true)), x_true, label=\"True parameter values\")\n",
    "plt.scatter(np.arange(len(x_true)), x_ml, label=\"ML estimate (no prior)\")\n",
    "plt.scatter(\n",
    "    np.arange(len(x_true)), np.mean(X_posterior, axis=1), label=\"Posterior mean\"\n",
    ")\n",
    "plt.xlabel(\"Parameter index\")\n",
    "plt.ylabel(\"Parameter value\")\n",
    "plt.grid(True, ls=\"--\", zorder=0, alpha=0.33)\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1dece12",
   "metadata": {},
   "source": [
    "## Solve using AdaptiveESMDA\n",
    "\n",
    "We create an `AdaptiveESMDA` instance and solve the Guass-linear problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc2ea87c",
   "metadata": {},
   "outputs": [],
   "source": [
    "adaptive_smoother = AdaptiveESMDA(\n",
    "    covariance=covariance,\n",
    "    observations=observations,\n",
    "    seed=1,\n",
    ")\n",
    "\n",
    "X_i = np.copy(X)\n",
    "\n",
    "# Loop over alpha defined in ESMDA instance,\n",
    "# the vector of inflation values alpha is then\n",
    "# of the same size in AdaptiveESMDA and ESMDA,\n",
    "# and it's correctly scaled\n",
    "assert np.isclose(np.sum(1 / smoother.alpha), 1), \"Incorrect scaling\"\n",
    "for i, alpha_i in enumerate(smoother.alpha, 1):\n",
    "    print(\n",
    "        f\"AdaptiveESMDA iteration {i}/{len(smoother.alpha)}\"\n",
    "        + f\" with inflation factor alpha_i={alpha_i}\"\n",
    "    )\n",
    "\n",
    "    # Run forward model\n",
    "    Y_i = g(X_i)\n",
    "\n",
    "    # Perturb observations\n",
    "    D_i = adaptive_smoother.perturb_observations(\n",
    "        ensemble_size=X_i.shape[1], alpha=alpha_i\n",
    "    )\n",
    "\n",
    "    # Assimilate data\n",
    "    X_i = adaptive_smoother.assimilate(\n",
    "        X=X_i, Y=Y_i, D=D_i, alpha=alpha_i, verbose=False\n",
    "    )\n",
    "\n",
    "\n",
    "X_adaptive_posterior = np.copy(X_i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73102263",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 3))\n",
    "plt.scatter(np.arange(len(x_true)), x_true, label=\"True parameter values\")\n",
    "plt.scatter(np.arange(len(x_true)), x_ml, label=\"ML estimate (no prior)\")\n",
    "plt.scatter(\n",
    "    np.arange(len(x_true)), np.mean(X_posterior, axis=1), label=\"Posterior mean\"\n",
    ")\n",
    "plt.scatter(\n",
    "    np.arange(len(x_true)),\n",
    "    np.mean(X_adaptive_posterior, axis=1),\n",
    "    label=\"Posterior mean (adaptive)\",\n",
    ")\n",
    "plt.xlabel(\"Parameter index\")\n",
    "plt.ylabel(\"Parameter value\")\n",
    "plt.grid(True, ls=\"--\", zorder=0, alpha=0.33)\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84eecf27",
   "metadata": {},
   "source": [
    "## Correlations between true parameters and solution means\n",
    "\n",
    "- A more sophisticated way to measure fit might be to use [Kullback–Leibler divergence](https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence).\n",
    "- Here we simply look at the correlations of the means."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4516f962",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "for arr, label in zip(\n",
    "    [\n",
    "        np.mean(X, axis=1),\n",
    "        x_ml,\n",
    "        np.mean(X_posterior, axis=1),\n",
    "        np.mean(X_adaptive_posterior, axis=1),\n",
    "    ],\n",
    "    [\"Prior\", \"ML estimate\", \"Posterior mean\", \"Posterior mean (adaptive)\"],\n",
    "):\n",
    "    corr = sp.stats.pearsonr(x_true, arr).statistic\n",
    "    print(label, corr)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8922579a",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "source": [
    "## Run on several ensemble sizes and seeds\n",
    "\n",
    "To get a more complete picture of how the number of realizations (ensemble size)\n",
    "affects the result, we loop over various ensemble sizes and compute posteriors.\n",
    "To capture average behavior and reduce the influence of sampling noise,\n",
    "we also loop over several seeds.\n",
    "\n",
    "Recall that there are two sources of randomness in the result:\n",
    "\n",
    "- The draw of prior realizations $X \\sim N(0, \\sigma)$.\n",
    "- The noise (given by the `covariance` argument) that ESMDA adds to the observations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6752299a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def corr_true(array):\n",
    "    return sp.stats.pearsonr(x_true, array).statistic\n",
    "\n",
    "\n",
    "ENSEMBLE_SIZES = list(range(2, 51))\n",
    "NUM_SEEDS = 25\n",
    "\n",
    "# Store average correlation coefficients\n",
    "ESMDA_corrs = []\n",
    "AdaptiveESMDA_corrs = []\n",
    "prior_corrs = []  # Baseline comparison - no update\n",
    "\n",
    "# Loop over increasingly large ensemble sizes\n",
    "for ensemble_size in ENSEMBLE_SIZES:\n",
    "    # Posteriors means for this size\n",
    "    ESMDA_means = []\n",
    "    AdaptiveESMDA_means = []\n",
    "    prior_means = []\n",
    "\n",
    "    # Loop over seeds used in ESMDA/AdaptiveESMDA,\n",
    "    # which determine the perturbations of the observations.\n",
    "    for seed in range(NUM_SEEDS):\n",
    "        # Prior\n",
    "        X = get_prior(num_parameters, ensemble_size, prior_std)\n",
    "        prior_means.append(np.mean(X, axis=1))\n",
    "\n",
    "        # ================ ESMDA ==============\n",
    "        smoother = ESMDA(\n",
    "            covariance=covariance,\n",
    "            observations=observations,\n",
    "            alpha=5,\n",
    "            seed=seed,\n",
    "        )\n",
    "\n",
    "        X_i = np.copy(X)\n",
    "        for i, alpha_i in enumerate(smoother.alpha, 1):\n",
    "            X_i = smoother.assimilate(X_i, Y=g(X_i))\n",
    "\n",
    "        ESMDA_means.append(np.mean(X_i, axis=1))\n",
    "\n",
    "        # ============ AdaptiveESMDA ============\n",
    "        adaptive_smoother = AdaptiveESMDA(\n",
    "            covariance=covariance,\n",
    "            observations=observations,\n",
    "            seed=seed,\n",
    "        )\n",
    "\n",
    "        X_i = np.copy(X)\n",
    "\n",
    "        for i, alpha_i in enumerate(smoother.alpha, 1):\n",
    "            # Perturb observations\n",
    "            D_i = adaptive_smoother.perturb_observations(\n",
    "                ensemble_size=X_i.shape[1], alpha=alpha_i\n",
    "            )\n",
    "\n",
    "            # Assimilate data\n",
    "            X_i = adaptive_smoother.assimilate(X=X_i, Y=g(X_i), D=D_i, alpha=alpha_i)\n",
    "\n",
    "        AdaptiveESMDA_means.append(np.mean(X_i, axis=1))\n",
    "\n",
    "    # Collect results for all runs of this size\n",
    "    ESMDA_corr = np.mean([corr_true(arr) for arr in ESMDA_means])\n",
    "    AdaptiveESMDA_corr = np.mean([corr_true(arr) for arr in AdaptiveESMDA_means])\n",
    "    prior_corr = np.mean([corr_true(arr) for arr in prior_means])\n",
    "\n",
    "    # Add to respective lists\n",
    "    ESMDA_corrs.append(ESMDA_corr)\n",
    "    AdaptiveESMDA_corrs.append(AdaptiveESMDA_corr)\n",
    "    prior_corrs.append(prior_corr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49c79153",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create figure title with a lot of information\n",
    "title = \"ESMDA vs AdaptiveESMDA on different ensemble sizes.\\n\"\n",
    "title += (\n",
    "    f\"Each point represents the average of {NUM_SEEDS} runs with different seeds.\\n\"\n",
    ")\n",
    "title += f\"Parameters = {num_parameters}, Observations =\"\n",
    "title += f\" {num_observations}, len(alpha) = {len(smoother.alpha)}\"\n",
    "\n",
    "\n",
    "# Create the figure\n",
    "plt.figure(figsize=(7, 3.5))\n",
    "plt.title(title)\n",
    "\n",
    "plt.plot(ENSEMBLE_SIZES, ESMDA_corrs, label=\"ESMDA\")\n",
    "plt.plot(ENSEMBLE_SIZES, AdaptiveESMDA_corrs, label=\"AdaptiveESMDA\")\n",
    "plt.plot(ENSEMBLE_SIZES, prior_corrs, label=\"Prior\")\n",
    "\n",
    "plt.xlabel(\"Ensemble size\")\n",
    "plt.ylabel(\"Correlation coeff of ensemble\\nmeans against true parameters\")\n",
    "\n",
    "plt.grid(True)\n",
    "plt.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,py:percent"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}