{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Parameters estimation\n\nPlot obtained values of the parameters versus the actual ones\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nfrom matplotlib import pyplot as plt\nfrom GLE_analysisEM import GLE_Estimator, GLE_BasisTransform\nfrom GLE_analysisEM.post_processing import memory_kernel, forcefield, forcefield_plot2D, correlation, memory_timescales\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 = 2\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))\ngenerator = GLE_Estimator(verbose=2, dim_x=dim_x, dim_h=dim_h, basis=pot_gen, force_init=force, init_params=\"random\", random_state=random_state)\nX, idx, Xh = generator.sample(n_samples=20000, n_trajs=25, x0=0.0, v0=0.0)\n\n# ------ Estimation ------#\n# basis = GLE_BasisTransform(basis_type=\"linear\")\nbasis = GLE_BasisTransform(basis_type=\"free_energy_kde\")\nestimator = GLE_Estimator(verbose=2, verbose_interval=10, dim_x=dim_x, dim_h=dim_h, basis=basis, n_init=1, random_state=7, tol=1e-7, no_stop=False, max_iter=100, multiprocessing=8)\nestimator.fit(X, idx_trajs=idx)\n\n# ------ Plotting ------#\nfig, axs = plt.subplots(2, 2)\n\n\n# ------ Force field ------#\naxs[0, 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, 0].plot(xfx_true[:, 0], xfx_true[:, 1], label=\"True force field\")\n    axs[0, 0].plot(xfx[:, 0], xfx[:, 1], label=\"Fitted force field\")\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, 0].quiver(x_true, y_true, fx_true, fy_true, width=0.001, color=\"green\", label=\"True force field\")\n    axs[0, 0].quiver(x, y, fx, fy, width=0.001, color=\"blue\", label=\"Fitted force field\")\n\naxs[0, 0].legend(loc=\"upper right\")\n\n# ------ Memory kernel ------#\naxs[0, 1].set_title(\"Memory kernel\")\ntime, kernel = memory_kernel(1000, estimator.dt, estimator.get_coefficients(), dim_x)\ntime_true, kernel_true = memory_kernel(1000, generator.dt, generator.get_coefficients(), dim_x)\n\n\naxs[0, 1].plot(time, kernel[:, 0, 0], label=\"Fitted memory kernel\")\naxs[0, 1].plot(time_true, kernel_true[:, 0, 0], label=\"True memory kernel\")\naxs[0, 1].legend(loc=\"upper right\")\n\n# ------ Memory eigenvalues ------#\naxs[1, 1].set_title(\"Kernel Eigenvalues\")\nmem_ev = memory_timescales(estimator.get_coefficients(), dim_x=dim_x)\nmem_ev_true = memory_timescales(generator.get_coefficients(), dim_x=dim_x)\naxs[1, 1].scatter(np.real(mem_ev), np.imag(mem_ev), label=\"Ev fitted\")\naxs[1, 1].scatter(np.real(mem_ev_true), np.imag(mem_ev_true), label=\"Ev true\")\n# axs[1, 1].set_aspect(1)\n\n\ndef simulated_vacf(estimator):\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 /= vacf[0]\n    return time, vacf\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]\nvacf_num /= vacf_num[0]\ntime_sim, vacf_sim = simulated_vacf(estimator)\naxs[1, 0].plot(time[: len(time_sim) // 2], vacf_sim, label=\"Fitted VACF\")\naxs[1, 0].set_title(\"Velocity autocorrelation function\")\naxs[1, 0].plot(time[: len(time) // 2], vacf_num, label=\"Numerical VACF\")\naxs[1, 0].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
}