{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multivariate inference\n",
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Problem definition\n",
    "from probeye.definition.inverse_problem import InverseProblem\n",
    "from probeye.definition.forward_model import ForwardModelBase\n",
    "from probeye.definition.distribution import Normal, Uniform\n",
    "from probeye.definition.sensor import Sensor\n",
    "from probeye.definition.likelihood_model import GaussianLikelihoodModel\n",
    "from probeye.definition.correlation_model import ExpModel\n",
    "\n",
    "# Inference\n",
    "from probeye.inference.emcee.solver import EmceeSolver\n",
    "\n",
    "# Postprocessing\n",
    "from probeye.postprocessing.sampling_plots import create_pair_plot\n",
    "from probeye.postprocessing.sampling_plots import create_posterior_plot\n",
    "from probeye.postprocessing.sampling_plots import create_trace_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fixed parameters\n",
    "I = 1e9  # mm^4\n",
    "L = 10_000  # mm\n",
    "\n",
    "# Measurements\n",
    "x_sensors = [2500, 5000]  # mm  (always from lower to higher, bug in inv_cov_vec_1D in tripy)\n",
    "d_sensors = [35, 50]  # mm\n",
    "sigma_model = 2.5  # mm\n",
    "pearson = 0.5\n",
    "l_corr = -np.abs(x_sensors[1] - x_sensors[0]) / np.log(pearson)  # mm (assuming exponential correlation)\n",
    "\n",
    "# Prior\n",
    "E_mean = 60  # GPa\n",
    "E_std = 20  # GPa\n",
    "Q_mean = 60  # kN\n",
    "Q_std = 30  # kN\n",
    "Q_loc_low = 0  # mm\n",
    "Q_loc_high = 10000  # mm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define forward model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def beam_deflection(E, Q, a, x):  # a is load position, x is sensor position\n",
    "    if x < a:\n",
    "        b = L - a\n",
    "        return Q * b * x * (L ** 2 - b ** 2 - x ** 2) / (6 * E * I * L)\n",
    "    \n",
    "    return Q * a * (L - x) * (2 * L * x - x ** 2 - a ** 2) / (6 * E * I * L)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BeamModel(ForwardModelBase):\n",
    "    def interface(self):\n",
    "        self.parameters = ...\n",
    "        self.input_sensors = Sensor(\"x\")\n",
    "        self.output_sensors = Sensor(\"y\", std_model=\"sigma\")\n",
    "\n",
    "    def response(self, inp: dict) -> dict:\n",
    "        return {...}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define inverse problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Instantiate the inverse problem\n",
    "problem = ...\n",
    "\n",
    "..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Solve with MCMC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "emcee_solver = EmceeSolver(problem, show_progress=True)\n",
    "inference_data = emcee_solver.run(n_steps=2000, n_initial_steps=2000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Posterior plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "post_plot_array = create_posterior_plot(\n",
    "    inference_data,\n",
    "    emcee_solver.problem,\n",
    "    kind=\"kde\",\n",
    "    title=\"Kernel density estimate of the posterior distribution\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pair plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pair_plot_array = create_pair_plot(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trace plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace_plot_array = create_trace_plot(...)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Posterior predictives"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract samples\n",
    "posterior_samples = inference_data.posterior.to_array()  # parameters, chains, samples\n",
    "posterior_samples = np.array(posterior_samples)\n",
    "posterior_samples = posterior_samples.reshape(posterior_samples.shape[0], -1).T  # samples, parameters\n",
    "\n",
    "# Make predictions for x in the range of the beam\n",
    "x_range = np.linspace(0, 10000, 100)\n",
    "predictions = np.zeros((len(posterior_samples), len(x_range)))\n",
    "\n",
    "for i, sample in enumerate(posterior_samples):\n",
    "    ...\n",
    "\n",
    "# Calculate mean and 95% intervals\n",
    "mean_pred = ...\n",
    "lower_bound = ...  # 2.5th percentile\n",
    "upper_bound = ...  # 97.5th percentile"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot posterior predictive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(12, 4))\n",
    "plt.plot(x_range, -mean_pred, label='Mean Prediction')\n",
    "plt.fill_between(x_range, -lower_bound, -upper_bound, color='lightblue', alpha=0.5, label='95% Interval')\n",
    "plt.xlabel('x (mm)')\n",
    "plt.ylabel('Deflection')\n",
    "plt.title('Posterior Predictive Deflections along the Beam')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}