{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a9463197",
   "metadata": {},
   "outputs": [],
   "source": [
    "from aare import Interpolator\n",
    "import numpy as np \n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fdc9ff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_rosenblatt_lookup_table(etacube, etabins_x, etabins_y, etabins_e):\n",
    "\n",
    "    \"\"\" \n",
    "        Rosenblatt transform: \n",
    "        transforms joint probability P_X,Y to uniform probability between [0,1] U\n",
    "        U_0 = F_X(x) = P(x' < x) marginal CDF of X\n",
    "        U_1 = F_Y|X(y|x) = P(y' < y | x) conditional CDF\n",
    "    \"\"\"\n",
    "\n",
    "    #marginal CDF for eta_x\n",
    "    pEtaX = etacube.sum(axis=1) # marginal distribution eta_x\n",
    "    CDF_EtaX = np.concatenate([np.zeros(shape=[1, len(etabins_e)]), np.cumsum(pEtaX, axis=0)], axis=0)\n",
    "    CDF_EtaX = CDF_EtaX / CDF_EtaX[-1] #normalize            \n",
    "\n",
    "    # conditional CDF F_Y|X (a CDF in y direction for each etaX bin)\n",
    "    CDF_EtaY_cond = np.cumsum(etacube, axis=1)\n",
    "    np.divide(CDF_EtaY_cond, np.repeat(pEtaX[:, None, :], len(etabins_y), axis=1), out=np.zeros_like(CDF_EtaY_cond, dtype=np.float64), where=np.repeat(pEtaX[:, None, :], len(etabins_y), axis=1)!=0)\n",
    "\n",
    "    CDF_EtaY_cond = np.concatenate([np.zeros(shape=[len(etabins_x), 1, len(etabins_e)]), CDF_EtaY_cond], axis=1)\n",
    "    #CDF_EtaY_cond = CDF_EtaY_cond /  np.where(CDF_EtaY_cond[:,-1:,:] == 0, 1, CDF_EtaY_cond[:,-1:,:]) #normalize\n",
    "\n",
    "    # calculate (u,v) at the center of each original bin\n",
    "    \"\"\" \n",
    "    Xc_1d = 0.5 * (F_EtaX_edges[:-1] + F_EtaX_edges[1:]) # why the average between two values?\n",
    "    Yc_2d = 0.5 * (F_EtaY_cond_edges[:, :-1] + F_EtaY_cond_edges[:, 1:])\n",
    "\n",
    "    X_tab = np.repeat(Xc_1d[:, None], nBinY, axis=1)        # (nBinX, nBinY)\n",
    "    Y_tab = Yc_2d                                           # (nBinX, nBinY)\n",
    "    \"\"\" \n",
    "    X_tab = np.repeat(CDF_EtaX[:, None, :], etabins_y.shape[0], axis=1) #TODO: actually does not depend on eta_y at all \n",
    "    Y_tab = CDF_EtaY_cond \n",
    "\n",
    "    print(\"shape X_tab: \", X_tab.shape)\n",
    "    print(\"shape Y_tab: \", Y_tab.shape)\n",
    "\n",
    "    return X_tab, Y_tab\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "616d04d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'y')"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 79696 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## joint conditional distribution \n",
    "rng = np.random.default_rng(42)\n",
    "n = 10_000_000  \n",
    " \n",
    "## circular distribution x,y dependant \n",
    "theta = rng.uniform(0, 4*np.pi, size=n)\n",
    "r = rng.uniform(0.3, 0.4) \n",
    "#r = 0.3 + 0.4 * (theta / (4*np.pi))  # 螺旋半径从 0.3 到 0.7\n",
    "x_raw = r * np.cos(theta) + 0.1 * rng.normal(size=n)\n",
    "y_raw = r * np.sin(theta) + 0.1 * rng.normal(size=n)\n",
    " \n",
    "# scale to [0, 1]\n",
    "x = (x_raw - x_raw.min()) / (x_raw.max() - x_raw.min())\n",
    "y = (y_raw - y_raw.min()) / (y_raw.max() - y_raw.min())\n",
    "    \n",
    "plt.scatter(x[:5000], y[:5000], s=1, alpha=0.5)\n",
    "plt.title(\"Original Distribution (x, y)\")\n",
    "plt.gca().set_aspect('equal')\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"y\")\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a713be97",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(x, y, CDF_X, CDF_Y, u, v): \n",
    "    #plot\n",
    "    fig, axs = plt.subplots(1, 4, figsize=(15, 4))\n",
    " \n",
    "    axs[0].scatter(x[:5000], y[:5000], s=1, alpha=0.5)\n",
    "    axs[0].set_title(\"Original (x, y)\")\n",
    "    axs[0].set_aspect('equal')\n",
    "    axs[0].set_xlabel(\"x\")\n",
    "    axs[0].set_ylabel(\"y\")\n",
    "    \n",
    "    axs[1].imshow(CDF_X[:,:,0], origin='lower', extent=[0,1,0,1], cmap='viridis')\n",
    "    axs[1].set_title(\"normalized CDF_X (u = F_X(x))\")\n",
    "    axs[1].set_xlabel(\"y\")\n",
    "    axs[1].set_ylabel(\"P_X(x < x')\")\n",
    "\n",
    "    axs[2].imshow(CDF_Y[:,:,0], origin='lower', extent=[0,1,0,1], cmap='viridis')\n",
    "    axs[2].set_title(\"normalized conditional CDF_Y (u = F_Y|X(y|x))\")\n",
    "    axs[2].set_xlabel(\"P_Y|X(y < y'|x)\")\n",
    "    axs[2].set_ylabel(\"x\")\n",
    "\n",
    "    axs[3].scatter(u[:5000], v[:5000], s=1, alpha=0.5, color='blue')\n",
    "    axs[3].set_title(\"Mapped (u, v)\")\n",
    "    axs[3].set_aspect('equal')\n",
    "    axs[3].set_xlim(0,1); axs[2].set_ylim(0,1)\n",
    "    axs[3].set_xlabel(\"u\")\n",
    "    axs[3].set_ylabel(\"v\")\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0e87346a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(300, 300, 2)\n"
     ]
    }
   ],
   "source": [
    "#calculate joint probability density function\n",
    "nBins = 300\n",
    "range = np.array([[x.min(),x.max()], [y.min(),y.max()]])\n",
    "bin_width = (range[:,0] - range[:,1])/nBins\n",
    "H, xedges, yedges = np.histogram2d(x, y, bins=nBins, range=range)\n",
    "H = H.astype(int) \n",
    "H = np.repeat(H[:,:,None], 2, axis=2) #add energy axis\n",
    "\n",
    "print(H.shape)\n",
    "ebins = [0,1,2] #dummy energy bins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d63e9295",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 296052 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## check that interpolator transforms to uniform distribution \n",
    "interpolator = Interpolator(H, xedges[:-1], yedges[:-1], ebins[:-1]) #dummy energy bins\n",
    "CDF_X = interpolator.get_ietax() #lookup table x \n",
    "CDF_Y = interpolator.get_ietay() #lookup table y \n",
    "\n",
    "x_indices = ((x - range[0,0])/bin_width[0]).astype(int)\n",
    "y_indices = ((y - range[1,0])/bin_width[1]).astype(int)\n",
    "\n",
    "u = CDF_X[x_indices, y_indices, 0] #transformed to uniform coordinates\n",
    "v = CDF_Y[x_indices, y_indices, 0] #transformed to uniform coordinates\n",
    "\n",
    "plot(x,y,CDF_X,CDF_Y, u,v)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "85823394",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape X_tab:  (301, 300, 2)\n",
      "shape Y_tab:  (300, 301, 2)\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 240092 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "x_indices = ((x - range[0,0])/bin_width[0]).astype(int)\n",
    "y_indices = ((y - range[1,0])/bin_width[1]).astype(int)\n",
    "\n",
    "CDF_X, CDF_Y = calculate_rosenblatt_lookup_table(H, xedges[:-1], yedges[:-1],ebins[:-1])\n",
    "\n",
    "u = CDF_X[x_indices, y_indices, 0] #transformed to uniform coordinates\n",
    "v = CDF_Y[x_indices, y_indices, 0] #transformed to uniform coordinates\n",
    " \n",
    "plot(x,y,CDF_X,CDF_Y, u,v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fe9e3306",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "worked: \n",
      "could normalize CDF eta_x\n",
      "construction worked\n",
      "calculation of conditional CDF etay worked\n",
      "assignment m_ietay worked\n",
      "assignment m_ietax worked\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 218732 characters]",
      "text/plain": [
       "<Figure size 1500x400 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "## check that cpp rosenblatt transforms to uniform distribution \n",
    "interpolator = Interpolator(xedges[:-1], yedges[:-1], ebins[:-1]) #dummy energy bins\n",
    "interpolator.rosenblatttransform(H)\n",
    "CDF_X = interpolator.get_ietax() #lookup table x \n",
    "CDF_Y = interpolator.get_ietay() #lookup table y \n",
    "\n",
    "x_indices = ((x - range[0,0])/bin_width[0]).astype(int)\n",
    "y_indices = ((y - range[1,0])/bin_width[1]).astype(int)\n",
    "\n",
    "u = CDF_X[x_indices, y_indices, 0] #transformed to uniform coordinates\n",
    "v = CDF_Y[x_indices, y_indices, 0] #transformed to uniform coordinates\n",
    "\n",
    "plot(x,y,CDF_X,CDF_Y, u,v)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6df08eb9",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a9d088d2",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "40a1ffff",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e0aa9dc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
