{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Parameters estimation from Markovian\n\nPlot obtained values of the parameters versus the actual ones when the parameters are estimated from a Markov model\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nfrom matplotlib import pyplot as plt\n\nfrom GLE_analysisEM import GLE_BasisTransform, GLE_Estimator, GLE_PotentialTransform\n\nfrom GLE_analysisEM import Markov_Estimator\nfrom GLE_analysisEM.post_processing import forcefield, forcefield_plot2D, correlation\n\nfrom sklearn.preprocessing import FunctionTransformer\n\na = 0.025\nb = 1.0\n\n\ndef dV(X):\n    \"\"\"\n    Compute the force field\n    \"\"\"\n    return 4 * a * np.power(X, 3) - 2 * b * X\n\n\ndim_x = 1\ndim_h = 1\nrandom_state = None\nforce = -np.identity(dim_x)\n# force = [[-0.25, -1], [1, -0.25]]\nA = np.array([[5e-2, -1.0], [1.0, 0.1]])\n\n# ------ Generation ------#\n# pot_gen = GLE_BasisTransform(basis_type=\"linear\")\npot_gen = GLE_BasisTransform(transformer=FunctionTransformer(dV))\n# pot_gen_polynom = GLE_BasisTransform(basis_type=\"polynomial\", degree=3)\ngenerator = GLE_Estimator(verbose=2, dim_x=dim_x, dim_h=2, basis=pot_gen, init_params=\"random\", force_init=force, random_state=random_state)\nX, idx, Xh = generator.sample(n_samples=20000, n_trajs=25, x0=0.0, v0=0.0)\nprint(\"---- Real ones ----\")\nprint(generator.get_coefficients())\n\n# for n in range(dim_x):\n#     plt.plot(X[:, 0], X[:, n * 2 + 2], label=\"v{}\".format(n + 1))\n#     plt.plot(X[:, 0], X[:, n * 2 + 1], label=\"x{}\".format(n + 1))\n#\n# plt.show()\n# ------ Estimation ------#\n# basis = GLE_BasisTransform(basis_type=\"linear\")\nbasis = GLE_BasisTransform(basis_type=\"polynomial\", degree=3).fit(X[1 : 1 + dim_x])\nestimator = Markov_Estimator(init_params=\"random\", verbose=2, verbose_interval=1, dim_x=dim_x, basis=basis, n_init=1, OptimizeForce=True, random_state=7)\nestimator.fit(X, idx_trajs=idx)\n# print(estimator.get_coefficients())\n\n# Free energy\n\npotential = GLE_PotentialTransform(estimator=\"histogram\", dim_x=dim_x)\npotential.fit(X)\n\n# ------ Plotting ------#\nfig, axs = plt.subplots(2)\n\n\n# ------ Force field ------#\naxs[0].set_title(\"Force field\")\nforce_true = generator.get_coefficients()[\"force\"]\nforce_fitted = estimator.get_coefficients()[\"force\"]\nif dim_x == 1:\n    x_lims = [[-10, 10, 25]]\n    xfx_true = forcefield(x_lims, pot_gen, force_true)\n    xfx = forcefield(x_lims, basis, force_fitted)\n    axs[0].plot(xfx_true[:, 0], xfx_true[:, 1], label=\"True force field\")\n    axs[0].plot(xfx[:, 0], xfx[:, 1], label=\"Fitted force field\")\n\n    x_lims = [[-8, 8, 150]]\n    x_val = np.linspace(x_lims[0][0], x_lims[0][1], x_lims[0][2]).reshape(-1, 1)\n    pot_val = potential.predict(x_val)\n    axs[0].plot(x_val[:, 0], pot_val[:, 0], label=\"Fitted potential\")\n\nif dim_x == 2:\n    x_lims = [[-2, 2, 10], [-2, 2, 10]]\n    x_true, y_true, fx_true, fy_true = forcefield_plot2D(x_lims, basis, force_true)\n    x, y, fx, fy = forcefield_plot2D(x_lims, basis, force_fitted)\n    axs[0].quiver(x_true, y_true, fx_true, fy_true, width=0.001, color=\"green\", label=\"True force field\")\n    axs[0].quiver(x, y, fx, fy, width=0.001, color=\"blue\", label=\"Fitted force field\")\n\naxs[0].legend(loc=\"upper right\")\n\n\ndef simulated_vacf(estimator, basis):\n    \"\"\"\n    Get vacf via numericall simulation of the model\n    \"\"\"\n    Ntrajs = 5\n    X, idx, Xh = estimator.sample(n_samples=10000, n_trajs=Ntrajs)\n    traj_list = np.split(X, idx)\n    vacf = 0.0\n    for n, trj in enumerate(traj_list):\n        vacf += correlation(trj[:, 1 + estimator.dim_x])\n        time = trj[:, 0]\n    # vacf /= Ntrajs\n    return time, vacf / vacf[0]\n\n\n# ------ Diffusion ------#\ntraj_list = np.split(X, idx)\nvacf_num = 0.0\nfor n, trj in enumerate(traj_list):\n    vacf_num += correlation(trj[:, 2])\n    time = trj[:, 0]\n# vacf_num /= len(traj_list)\nvacf_num /= vacf_num[0]\ntime_sim, vacf_sim = simulated_vacf(estimator, basis)\naxs[1].plot(time[: len(time_sim) // 2], vacf_sim, label=\"Fitted VACF\")\naxs[1].set_title(\"Velocity autocorrelation function\")\naxs[1].plot(time[: len(time) // 2], vacf_num, label=\"Numerical VACF\")\n\n\naxs[1].legend(loc=\"upper right\")\n\nplt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.9.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}