final_pole_module.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bb0c2a4c-628e-467e-884b-025a80031f2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import cmath as cm\n",
    "import pandas as pd\n",
    "import random\n",
    "import pickle\n",
    "import os\n",
    "import datetime\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.pyplot import figure\n",
    "import chainer\n",
    "from chainer import configuration\n",
    "from chainer.dataset import convert\n",
    "import sklearn\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.model_selection import train_test_split\n",
    "from tabulate import tabulate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "735cb3cd-8d91-418e-8d2b-78ed6dd2692c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#constants\n",
    "hbarc = 197.3\n",
    "\n",
    "#Heavy quark sector\n",
    "JPsi = 3096.90\n",
    "# Pion = 138.0394\n",
    "\n",
    "Proton = 938.27\n",
    "\n",
    "Dplus = 1869.66\n",
    "Dstar0 = 1864.83 #This value is for D0; I just happened to name it Dstar0\n",
    "#and found it too much a hassle to rename the variable\n",
    "#https://pdglive.lbl.gov/Particle.action?init=0&node=S032&home=MXXX035\n",
    "Sigmaplus = 1189.37\n",
    "Sigmaplus_C = 2452.9 #https://pdglive.lbl.gov/Particle.action?init=0&node=B104&home=BXXX040\n",
    "\n",
    "\n",
    "T1 = (JPsi + Proton)/hbarc\n",
    "T2 = (Sigmaplus_C + Dstar0)/hbarc\n",
    "\n",
    "T4 = 4500/hbarc\n",
    "\n",
    "mu_JPsiP   = 1/(1/JPsi + 1/Proton)\n",
    "mu_Sigmaplus_CDstar0  = 1/(1/Sigmaplus_C + 1/Dstar0)\n",
    "\n",
    "mu1 = mu_JPsiP/hbarc\n",
    "mu2 = mu_Sigmaplus_CDstar0/hbarc\n",
    "\n",
    "NEpoints = 100 #37\n",
    "delE = 200/(NEpoints)\n",
    "\n",
    "#Generate energy axis\n",
    "Einput = np.zeros([NEpoints,],dtype = 'float64')\n",
    "for ndx in range(NEpoints):\n",
    "    rand_points = np.random.uniform(low=4212+(ndx)*delE, high=4212+(ndx+1)*delE) #np.random.uniform(low=T1*hbarc+(ndx)*delE, high=T1*hbarc+(ndx+1)*delE)\n",
    "    Einput[ndx] = rand_points\n",
    "\n",
    "k1 = np.sqrt(Einput**2.0-(T1*hbarc)**2.0)\n",
    "k2 = np.zeros([NEpoints,],dtype = 'complex_')\n",
    "for kndx in range(len(Einput)):\n",
    "    k2pts = cm.sqrt(Einput[kndx]**2.0-(T2*hbarc)**2.0)\n",
    "    k2[kndx] = k2pts\n",
    "    \n",
    "Nreal = 3000 #2000 #2000 #300 #1000 \n",
    "Nimag = 3000 #2000 #2000 #300 #1000\n",
    "\n",
    "#Generate Npole poles within the counting region\n",
    "#units in MeV\n",
    "Erealbelow = np.random.uniform(low=T2*hbarc-100, high=T2*hbarc, size=int(Nreal/2))\n",
    "Erealabove = np.random.uniform(low=T2*hbarc, high=T2*hbarc+100, size=int(Nreal/2))\n",
    "Ereal = np.concatenate((Erealbelow, Erealabove))\n",
    "Eimag = np.random.uniform(low=0.5, high=50, size=Nimag)\n",
    "    \n",
    "#Generate poles beyond the counting region\n",
    "#units in MeV\n",
    "Erealbelow = np.random.uniform(low=T2*hbarc-2500, high=T2*hbarc-2000, size=int(Nreal/2))#np.random.uniform(low=T1*hbarc-2000, high=T1*hbarc-100, size=int(Nreal/2))\n",
    "Erealabove = np.random.uniform(low=T2*hbarc+1000, high=T2*hbarc+1500, size=int(Nreal/2))#np.random.uniform(low=T2*hbarc+500, high=T2*hbarc+600, size=int(Nreal/2))\n",
    "\n",
    "Erealfar = np.concatenate((Erealbelow, Erealabove))\n",
    "Eimagfar = np.random.uniform(low=700.0, high=2000.0, size=Nimag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "95291506-7958-4e37-a56e-d2561ab55c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "class unif_pole:\n",
    "    def __init__(self, RS, Ereal, Eimag):\n",
    "        self.Ereal = Ereal\n",
    "        self.Eimag = Eimag\n",
    "        self.RS  = RS\n",
    "        Epole = Ereal - (1j)*Eimag\n",
    "        self.pos = Epole\n",
    "        \n",
    "        #compute uniformized momentum pole\n",
    "        #for channel 1 and channel 2\n",
    "        k1pole = cm.sqrt((Epole/hbarc)**2-T1**2)\n",
    "        k2pole = cm.sqrt((Epole/hbarc)**2-T2**2)\n",
    "        \n",
    "        #Riemann sheet assignment\n",
    "        beta1  = RS[0]*abs(k1pole.imag)\n",
    "        beta2  = RS[1]*abs(k2pole.imag)\n",
    "        \n",
    "        \n",
    "        #Get the real part, we need to be consistent with signs\n",
    "        alpha1 = -np.sign(beta1)*np.abs(k1pole.real)\n",
    "        alpha2 = -np.sign(beta2)*np.abs(k2pole.real)\n",
    "\n",
    "        #Just for counterchecking\n",
    "        #signs of beta1 and beta2 should agree \n",
    "        #with RS[0] and RS[1], respectively\n",
    "        self.alpha1 = alpha1\n",
    "        self.alpha2 = alpha2\n",
    "        self.beta1 = beta1\n",
    "        self.beta2 = beta2\n",
    "        \n",
    "        #Construct pole channel momenta\n",
    "        polep1 = (1j)*beta1 + alpha1\n",
    "        polep2 = (1j)*beta2 + alpha2\n",
    "        \n",
    "        Delta = cm.sqrt(T2**2 - T1**2)\n",
    "        self.Delta = Delta\n",
    "        \n",
    "        #Uniformization of the assigned pole\n",
    "        omega_pole = (polep1 + polep2)/Delta\n",
    "        recip_omg_pol = 1/omega_pole\n",
    "        self.omega_pole = omega_pole\n",
    "        self.recip_omg_pol = recip_omg_pol\n",
    "        \n",
    "        #Get pole regulator\n",
    "        #[bt] sheet\n",
    "        if RS==[-1,1]: \n",
    "            omega_reg = np.abs(recip_omg_pol)*cm.exp(-0.5*np.pi*(1j))\n",
    "        #[bb] sheet\n",
    "        elif RS==[-1,-1]:\n",
    "            omega_reg = np.abs(recip_omg_pol)*cm.exp(-0.5*np.pi*(1j))\n",
    "        #[tb] sheet\n",
    "        elif RS==[1,-1]:\n",
    "            omega_reg = np.abs(recip_omg_pol)*cm.exp(-0.5*np.pi*(1j))\n",
    "            \n",
    "        self.omega_reg = omega_reg\n",
    "        self.recip_omg_reg = 1/omega_reg\n",
    "        \n",
    "        p1_reg = omega_reg + 1/omega_reg\n",
    "        p2_reg = omega_reg - 1/omega_reg\n",
    "        \n",
    "        #Riemann sheet identifier for pole regulator\n",
    "        def RSlabel(pimag1, pimag2):\n",
    "            if pimag1>0 and pimag2>0:\n",
    "                RS = 'tt' #sheet 1\n",
    "            elif pimag1<0 and pimag2>0:\n",
    "                RS = 'bt' #sheet 2\n",
    "            elif pimag1<0 and pimag2<0:\n",
    "                RS = 'bb' #sheet 3\n",
    "            elif pimag1>0 and pimag2<0:\n",
    "                RS = 'tb' #sheet 4\n",
    "            return RS\n",
    "        #If you want to check the Riemann sheet of pole regulator\n",
    "        self.regulator = ['{:.2f}'.format(np.sqrt(p1_reg**2.0+T1**2.0)*hbarc), RSlabel(p1_reg.imag, p2_reg.imag)]\n",
    "        self.assignedpole = ['{:.2f}'.format(np.sqrt(polep1**2.0+T1**2.0)*hbarc), RSlabel(polep1.imag, polep2.imag)]\n",
    "    \n",
    "    #The indent on this part is very important\n",
    "    #Calculate the S-matrix contribution of the uniformized pole\n",
    "    def smat11(self, Ecm):\n",
    "        #Get channel momenta of Ecm\n",
    "        p1 = np.sqrt((Ecm/hbarc)**2.0-T1**2.0)\n",
    "        p2 = np.zeros([NEpoints,],dtype = 'complex_')\n",
    "        for pndx in range(len(Ecm)):\n",
    "            p2_pts = cm.sqrt((Ecm[pndx]/hbarc)**2.0-T2**2.0)\n",
    "            p2[pndx] = p2_pts\n",
    "\n",
    "        #Get uniformized parameter\n",
    "        omega = (p1 + p2)/self.Delta\n",
    "        \n",
    "        #Numerator of S-matrix\n",
    "        Numpol = (omega-np.conj(self.recip_omg_pol))*(omega+self.recip_omg_pol)\n",
    "        Numreg = (omega-np.conj(self.recip_omg_reg))*(omega+self.recip_omg_reg)\n",
    "        Num = Numpol*Numreg\n",
    "        \n",
    "        #Denominator of S-matrix\n",
    "        Denpol = (omega-self.omega_pole)*(omega+np.conj(self.omega_pole))\n",
    "        Denreg = (omega-self.omega_reg)*(omega+np.conj(self.omega_reg))\n",
    "        Den = Denpol*Denreg\n",
    "            \n",
    "        return np.abs((self.omega_pole*self.omega_reg))**2.0*Num/Den\n",
    "    \n",
    "    def smat22(self, Ecm):\n",
    "        #Get channel momenta of Ecm\n",
    "        p1 = np.sqrt((Ecm/hbarc)**2.0-T1**2.0)\n",
    "        p2 = np.zeros([NEpoints,],dtype = 'complex_')\n",
    "        for pndx in range(len(Ecm)):\n",
    "            p2_pts = cm.sqrt((Ecm[pndx]/hbarc)**2.0-T2**2.0)\n",
    "            p2[pndx] = p2_pts\n",
    "\n",
    "        #Get uniformized parameter\n",
    "        omega = (p1 + p2)/self.Delta\n",
    "        \n",
    "        #Numerator of S-matrix\n",
    "        Numpol = (omega+np.conj(self.recip_omg_pol))*(omega-self.recip_omg_pol)\n",
    "        Numreg = (omega+np.conj(self.recip_omg_reg))*(omega-self.recip_omg_reg)\n",
    "        Num = Numpol*Numreg\n",
    "        #Denominator of S-matrix\n",
    "        \n",
    "        Denpol = (omega-self.omega_pole)*(omega+np.conj(self.omega_pole))\n",
    "        Denreg = (omega-self.omega_reg)*(omega+np.conj(self.omega_reg))\n",
    "        Den = Denpol*Denreg\n",
    "            \n",
    "        return np.abs((self.omega_pole*self.omega_reg))**2.0*Num/Den\n",
    "    \n",
    "    def smatdet(self, Ecm):\n",
    "        #Get channel momenta of Ecm\n",
    "        p1 = np.sqrt((Ecm/hbarc)**2.0-T1**2.0)\n",
    "        p2 = np.zeros([NEpoints,],dtype = 'complex_')\n",
    "        for pndx in range(len(Ecm)):\n",
    "            p2_pts = cm.sqrt((Ecm[pndx]/197.3)**2.0-T2**2.0)\n",
    "            p2[pndx] = p2_pts\n",
    "\n",
    "        #Get uniformized parameter\n",
    "        omega = (p1 + p2)/self.Delta\n",
    "        \n",
    "        #Numerator of S-matrix\n",
    "        Numpol = (omega-np.conj(self.omega_pole))*(omega+self.omega_pole)\n",
    "        Numreg = (omega-np.conj(self.omega_reg))*(omega+self.omega_reg)\n",
    "        Num = Numpol*Numreg\n",
    "        \n",
    "        #Denominator of S-matrix\n",
    "        Denpol = (omega-self.omega_pole)*(omega+np.conj(self.omega_pole))\n",
    "        Denreg = (omega-self.omega_reg)*(omega+np.conj(self.omega_reg))\n",
    "        Den = Denpol*Denreg\n",
    "            \n",
    "        return np.abs((self.omega_pole*self.omega_reg))**2.0*Num/Den"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "635a024f-be74-4bba-a4bf-f4b9566a6ad4",
   "metadata": {
    "tags": []
   },
   "source": [
    "### On cusps\n",
    "\n",
    "The off-diagonal term of the $S$-matrix reads as\n",
    "\\begin{equation}\n",
    "S_{12}^2 = S_{11}S_{22} - \\text{det}(S).\n",
    "\\end{equation}\n",
    "\n",
    "When plotting $S_{12}$, we will encounter discontinuity in our plots. The ```np.angle(z)``` has a left branch cut. At best, we could rotate the cut and shift the discontinuity. However, this poses a problem when we want to consider a wide range of $E$. To circumvent this problem, we introduce the ```detect_cusp``` function. What this does is to detect the discontinuities of the input and changes the sign of the element from there onwards until it detects another discontinuity. For example, given an array with $20$ elements, e.g. ```Z[0], Z[1],..., Z[19]``` and suppose there are discontinuities at\n",
    "```Z[9], Z[14], Z[18]```. The output of the function is an array with:\n",
    "\n",
    "1) The original ```Z[0]``` to ```Z[8]``` of the input array\n",
    "2) Negated ( $Z \\to -Z$) ```Z[9]``` to ```Z[13]``` of the input array \n",
    "3) Original ```Z[14]``` to ```Z[17]``` of the input array\n",
    "4) Negated ```Z[18]``` and ```Z[19]``` of the input array \n",
    "\n",
    "These 4 then are concatenated such that the len(output) = len(input). \n",
    "\n",
    "We use this function as follows:\n",
    "\n",
    "1) Compute $\\text{det}(S)$:\n",
    "```smatdet = np.prod(smatdet, axis = 0)```\n",
    "2) Construct $S_{12}^2$:\n",
    "```pwat12sqr = (-1.0)*(smat11 * smat22 - smatdet)/4.0```\n",
    "3) Construct $S_{12}$ using polar representation. Get the modulo first:\n",
    "```mod12sqr = np.abs(pwat12sqr)```\n",
    "4) Get the argument of $S_{12}^2$ which we will divide by 2 upon construction of the polar representation:\n",
    "```arg12sqrBC = np.array([np.angle(z) for z in pwat12sqr])```\n",
    "5) Express in polar representation:\n",
    "```pwat12withdisc = np.sqrt(mod12sqr)*np.exp(1j* arg12sqrBC/ 2)```\n",
    "6) Express final pwat12 by using detect_cusp function multiplied by heaviside function:\n",
    "```detect_cusp(pwat12withdisc)*np.heaviside(Einput-T2*197.3,0)```\n",
    "\n",
    "The multiplication of heaviside should only be done after using ```detect_cusp```. \n",
    "If we put heaviside before ```detect_cusp```, the sign of the real and imag part of our plots will be inverted for some cases.\n",
    "The reason is python considers ```0. +(-) 0.j``` to have a positive (negative) imaginary part.\n",
    "It immediately detects a cusp at the threshold if we put the heaviside function inside ```detect_cusp```.\n",
    "\n",
    "Note that there is no physics behind this. This is purely motivated by the numerical output of python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e820369b-9238-4d34-9de0-f0df8e7054fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect_cusp(input_array, threshold=0.1):\n",
    "    # Calculate the derivative of the real and imaginary parts\n",
    "    real_diff = np.diff(np.real(input_array))\n",
    "    imag_diff = np.diff(np.imag(input_array))\n",
    "\n",
    "    # Detect points where the real part has a cusp\n",
    "    real_cusp_indices = np.where(np.abs(real_diff) > threshold)[0] + 1\n",
    "\n",
    "    # Detect points where the imaginary part has a cusp\n",
    "    imag_cusp_indices = np.where(np.abs(imag_diff) > threshold)[0] + 1\n",
    "\n",
    "    # Combine the indices and remove duplicates\n",
    "    all_cusp_indices = np.unique(np.concatenate((real_cusp_indices, imag_cusp_indices)))\n",
    "\n",
    "    # Modify the array by negating the sign for both real and imaginary parts\n",
    "    for cusp_index in all_cusp_indices:\n",
    "        input_array[cusp_index:] = -input_array[cusp_index:]\n",
    "\n",
    "    return input_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "1e82aea1-299b-42ef-8aa0-75407d1269f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of poles to be generated per class: 4000000\n",
      "Ndata to be generated= 16000000\n",
      "Your directory is: curriculumXX_training\n"
     ]
    }
   ],
   "source": [
    "# inspect = True\n",
    "inspect = True\n",
    "directory = 'curriculum02_training'\n",
    "#curr01 datasets: 00, 01, 11, 21\n",
    "\n",
    "#directory = 'curriculum02_training'\n",
    "#curr02 datasets: 00, 01, 11, 21, 02\n",
    "#directory = 'curriculum03_training'\n",
    "#curr03 datasets: 00, 01, 11, 21, 02, 12\n",
    "#directory = 'curriculum04_training'\n",
    "#curr04 datasets: 00, 01, 11, 21, 02, 12, 22\n",
    "#directory = 'curriculum05_training'\n",
    "#curr05 datasets: 00, 01, 11, 21, 02, 12, 22, 10\n",
    "#directory = 'curriculum06_training'\n",
    "#curr06 datasets: 00, 01, 11, 21, 02, 12, 22, 10, 20\n",
    "#directory = 'curriculum07_training'\n",
    "#curr07 datasets: 00, 01, 11, 21, 02, 12, 22, 10, 20, 30\n",
    "\n",
    "\n",
    "#directory = 'curriculum08_training'\n",
    "#curr08 datasets: 00, 01, 11, 21, 02, 12, 22, 10, 20, 30, 03\n",
    "\n",
    "#directory = 'curriculum32_training'\n",
    "#all datasets\n",
    "\n",
    "\n",
    "#directory = 'sample_plot'\n",
    "\n",
    "if not os.path.isdir(directory):\n",
    "    os.makedirs(directory)\n",
    "print('Number of poles to be generated per class:', Nreal*Nimag)\n",
    "print('Ndata to be generated=', 4*Nreal*Nimag)\n",
    "print('Your directory is:', directory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e1f5479c-1c30-47f0-8bac-1aa3df74a124",
   "metadata": {},
   "outputs": [],
   "source": [
    "#descriptive labels of network output\n",
    "#at most 4 poles in all RS\n",
    "labelz = [\n",
    "#default no pole\n",
    "    'no nearby pole',                          #00\n",
    "#poles in [bt]    \n",
    "    '1 pole  in [bt]',                          #01\n",
    "    '2 poles in [bt]',                         #02\n",
    "    '3 poles in [bt]',                         #03\n",
    "    '4 poles in [bt]',                         #04\n",
    "#[bt] and [bb] no shadow pair    \n",
    "    '3 poles in [bt] and 1 pole  in [bb]',      #05\n",
    "    '2 poles in [bt] and 1 pole  in [bb]',      #06\n",
    "    '2 poles in [bt] and 2 poles in [bb]',     #07 \n",
    "    '1 pole  in [bt] and 2 poles in [bb]',      #08\n",
    "    '1 pole  in [bt] and 3 poles in [bb]',     #09\n",
    "    '1 pole  in [bt] and 1 pole  in [bb]',      #10\n",
    "#poles in [bb] only    \n",
    "    '1 pole  in [bb]',                          #11\n",
    "    '2 poles in [bb]',                         #12\n",
    "    '3 poles in [bb]',                         #13\n",
    "    '4 poles in [bb]',                         #14\n",
    "#[bb] and [tb] no shadow pair    \n",
    "    '3 poles in [bb] and 1 pole  in [tb]',      #15\n",
    "    '2 poles in [bb] and 1 pole  in [tb]',      #16\n",
    "    '2 poles in [bb] and 2 poles in [tb]',     #17   \n",
    "    '1 pole  in [bb] and 2 poles in [tb]',      #18\n",
    "    '1 pole  in [bb] and 3 poles in [tb]',     #19\n",
    "    '1 pole  in [bb] and 1 pole  in [tb]',      #20\n",
    "#poles in [tb] only    \n",
    "    '1 pole  in [tb]',                          #21\n",
    "    '2 poles in [tb]',                         #22\n",
    "    '3 poles in [tb]',                         #23\n",
    "    '4 poles in [tb]',                         #24    \n",
    "#[tb] and [bt]\n",
    "    '3 poles in [tb] and 1 pole  in [bt]',      #25\n",
    "    '2 poles in [tb] and 1 pole  in [bt]',      #26\n",
    "    '2 poles in [tb] and 2 poles in [bt]',     #27\n",
    "    '1 pole  in [tb] and 2 poles in [bt]',      #28\n",
    "    '1 pole  in [tb] and 3 poles in [bt]',     #29 \n",
    "    '1 pole  in [tb] and 1 pole  in [bt]',      #30      \n",
    "#poles in all three\n",
    "    '2 poles in [bt] and 1 pole  in [bb] and 1 pole  in [tb]',    #31\n",
    "    '1 pole  in [bt] and 2 poles in [bb] and 1 pole  in [tb]',    #32\n",
    "    '1 pole  in [bt] and 1 pole  in [bb] and 2 poles in [tb]',    #33\n",
    "    '1 pole  in [bt] and 1 pole  in [bb] and 1 pole  in [tb]'      #34\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "08166eb5-b69a-4364-84dc-6da98d3b09f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "hepdata = pd.read_csv(\"hep.csv\", usecols=[0])\n",
    "\n",
    "hep_data = hepdata.loc[(hepdata.MEV > 4260) & (hepdata.MEV < 4380)].copy()\n",
    "\n",
    "E_exp = hep_data[\"MEV\"].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a22abc70-5b49-42b0-b163-aeec241c3e71",
   "metadata": {},
   "outputs": [],
   "source": [
    "def skip_duplicate(real1, imag1, Nreal, Nimag):\n",
    "    # Create lists of available real and imag values\n",
    "    real_list = [entry for entry in range(1, Nreal) if entry != real1]\n",
    "    imag_list = [entry for entry in range(1, Nimag) if entry != imag1]\n",
    "\n",
    "    # Randomly choose real values without duplication\n",
    "    real_choices = np.random.choice(real_list, 10, replace=False)\n",
    "    # Randomly choose imag values without duplication\n",
    "    imag_choices = np.random.choice(imag_list, 10, replace=False)\n",
    "\n",
    "    # Combine real and imag values into two lists\n",
    "    real_values = [real1] + list(real_choices)\n",
    "    imag_values = [imag1] + list(imag_choices)\n",
    "\n",
    "    indx = [real_values, imag_values]\n",
    "    return indx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d729faae-ae03-45ff-b2c2-81ed08d08afc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def export_data(Einput, ReT11, ImT11, labelout, data_info, output_directory):\n",
    "    # Create the specified output directory if it doesn't exist\n",
    "    os.makedirs(output_directory, exist_ok=True)\n",
    "    \n",
    "    # Define the data to export\n",
    "    data_to_export = {\n",
    "        'Einput.pkl': Einput,\n",
    "        'ModEsq.pkl': ModEsq,\n",
    "        'labelout.pkl': labelout,\n",
    "        'data_info.pkl': data_info,\n",
    "    }\n",
    "    \n",
    "    # Export each piece of data\n",
    "    for file_name, data in data_to_export.items():\n",
    "        with open(os.path.join(output_directory, file_name), 'wb') as file:\n",
    "            pickle.dump(data, file, protocol=4)\n",
    "    \n",
    "    # Collect data for input layer\n",
    "    EinputMod = np.concatenate((Einput,ModEsq), axis=1)\n",
    "    \n",
    "#     # Alternatively, you can design a DNN with Einput, ReT11, and ImT11 in the input layer\n",
    "#     T11 = np.concatenate((Einput, ReT11, ImT11), axis=1)\n",
    "    \n",
    "    # Export the collected data for the input layer\n",
    "    data_to_export = {\n",
    "        'EinputMod.pkl': EinputMod\n",
    "    }\n",
    "    \n",
    "    # Export each piece of data for the input layer\n",
    "    for file_name, data in data_to_export.items():\n",
    "        with open(os.path.join(output_directory, file_name), 'wb') as file:\n",
    "            pickle.dump(data, file, protocol=4)\n",
    "    \n",
    "    print('Export completed.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "890ea354-092d-4710-abf4-8aabbe5526cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def import_data(directory):\n",
    "    # Construct file paths\n",
    "    file_paths = {\n",
    "        'Einput': os.path.join(directory, 'Einput.pkl'),\n",
    "        'ModEsq': os.path.join(directory, 'ModEsq.pkl'),\n",
    "        'labelout': os.path.join(directory, 'labelout.pkl'),\n",
    "        'data_info': os.path.join(directory, 'data_info.pkl')\n",
    "    }\n",
    "\n",
    "    # Initialize empty dictionaries for data\n",
    "    data = {}\n",
    "\n",
    "    # Load data from files\n",
    "    for key, file_path in file_paths.items():\n",
    "        with open(file_path, 'rb') as file:\n",
    "            data[key] = pickle.load(file)\n",
    "\n",
    "    return data['Einput'], data['ModEsq'], data['labelout'], data['data_info']\n",
    "\n",
    "# Example usage:\n",
    "# Einput, ReT11, ImT11, labelout, data_info = import_data('input_directory')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e2340461-8b8c-44a8-9bc8-553995074f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def seerealimagpart(Einput, ModEsq, RePWAT, ImPWAT, labelout, data_info):\n",
    "#     # Randomly select an index for data sample\n",
    "#     chckind = np.random.randint(0, len(labelout))\n",
    "    \n",
    "#     # Define threshold values and corresponding points for plotting\n",
    "# #     thresholds = [\n",
    "# #         ('T_JPsiP', T1),\n",
    "# #         ('T_Sigmaplus_CDstar0', T2),\n",
    "# #         ('T_4', T4)\n",
    "# #     ]\n",
    "    \n",
    "# #     threshold_values = [\n",
    "# #     [-max(max(RePWAT[chckind]**2.0+ImPWAT[chckind]**2.0), max(np.abs(RePWAT[chckind])), max(np.abs(ImPWAT[chckind]))),\n",
    "# #      max(max(RePWAT[chckind]**2.0+ImPWAT[chckind]**2.0), max(np.abs(RePWAT[chckind])), max(np.abs(ImPWAT[chckind])))\n",
    "# #     ] for _ in range(3)\n",
    "# #                         ]\n",
    "\n",
    "#     # Set up the plot\n",
    "#     # max_value = max(max(RePWAT[chckind]**2.0 + ImPWAT[chckind]**2.0), max(np.abs(RePWAT[chckind])), max(np.abs(ImPWAT[chckind])))\n",
    "#     # plt.ylim(-max_value - 0.05, max_value + 0.05)\n",
    "#     # plt.plot([T1, T1], [0, 0], 'red')  # Horizontal line\n",
    "#     # plt.plot([T2, T2], [0, 0], 'red')  # Horizontal line\n",
    "#     # plt.plot([T4, T4], [0, 0], 'red')  # Horizontal line\n",
    "#     # for threshold, values in zip(thresholds, threshold_values):\n",
    "#     #     plt.plot(threshold[1], values, 'red')\n",
    "    \n",
    "#     # Plot data points\n",
    "#     plt.axhline(y = 0, color = 'r', linestyle = '-')\n",
    "#     plt.axvline(x = T1*hbarc, color = 'r' )\n",
    "#     plt.axvline(x = T2*hbarc, color = 'r')\n",
    "    \n",
    "#     plt.plot(Einput[chckind], RePWAT[chckind], '+', label='RePWAT')\n",
    "#     plt.plot(Einput[chckind], ImPWAT[chckind], '*', label='ImPWAT')\n",
    "#     plt.plot(Einput[chckind], ModEsq[chckind], 'o', label='ModSq')\n",
    "#     plt.legend(loc = 'upper left',frameon=True)\n",
    "    \n",
    "#     # Set plot labels and title\n",
    "#     plt.title('Input Data', fontsize=15)\n",
    "#     plt.xlabel('$E_{cm}$ (MeV)', fontsize=15)\n",
    "#     plt.xticks(fontsize=15)\n",
    "#     plt.ylabel('$Re T_{11}$, $Im T_{11}$', fontsize=15)\n",
    "#     plt.yticks(fontsize=15)\n",
    "#     plt.tight_layout()\n",
    "\n",
    "#     isactive = ['No', 'Yes']\n",
    "    \n",
    "\n",
    "    \n",
    "#     # Display class information in a table\n",
    "#     print('class', '{:02d}'.format(labelout[chckind]), ':', labelz[labelout[chckind]])\n",
    "#     table_data = [\n",
    "#         ['n', 'Energy Pole (MeV)', 'RS', 'Active']\n",
    "#     ]\n",
    "#     for i in range(7):\n",
    "#         table_data.append([str(i+1), '{:.2f}'.format(data_info[chckind][0][i]),\n",
    "#                            data_info[chckind][1][i], isactive[data_info[chckind][2][i]]])\n",
    "    \n",
    "#     print(tabulate(table_data))\n",
    "    \n",
    "#     table = plt.table(cellText=table_data, loc='upper left', colWidths=[0.1] * 4, cellLoc='center', edges='open')\n",
    "#     table.auto_set_font_size(False)\n",
    "#     table.set_fontsize(10)\n",
    "#     table.scale(1, 1.5)  # Adjust the scaling as needed\n",
    "    \n",
    "#     # Show the plot and return the selected index\n",
    "#     return plt.show(), chckind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f1983f94-c558-4983-ae64-7f5bb37878ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_timestamp():\n",
    "    now = datetime.datetime.now()\n",
    "    return now.strftime(\"%Y%m%d%H%M%S\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ea419c05-f549-442e-9d5b-85cd3036fc6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def seerealimagpart(Einput, ModEsq, Re11, Im11, Re22, Im22, labelout, data_info, \n",
    "                    save_dir='path/to/save'):\n",
    "    \n",
    "    #assert isinstance(Einput, np.ndarray) and Einput.dtype == np.complex128, \"Einput must be a complex numpy array\"\n",
    "    #assert isinstance(ModEsq, np.ndarray) and ModEsq.dtype == np.complex128, \"ModEsq must be a complex numpy array\"\n",
    "    #assert isinstance(Re11, np.ndarray) and Re11.dtype == np.complex128, \"Re11 must be a complex numpy array\"\n",
    "    #assert isinstance(Im11, np.ndarray) and Im11.dtype == np.complex128, \"Im11 must be a complex numpy array\"\n",
    "    #assert isinstance(Re22, np.ndarray) and Re22.dtype == np.complex128, \"Re22 must be a complex numpy array\"\n",
    "    #assert isinstance(Im22, np.ndarray) and Im22.dtype == np.complex128, \"Im22 must be a complex numpy array\"\n",
    "    #assert isinstance(labelout, np.ndarray) and labelout.dtype == np.int, \"labelout must be an integer numpy array\"\n",
    "\n",
    "    # Randomly select an index for data sample\n",
    "    chckind = np.random.randint(0, len(labelout))\n",
    "    \n",
    "    # Create the save directory if it doesn't exist\n",
    "    if not os.path.exists(save_dir):\n",
    "        os.makedirs(save_dir)\n",
    "        \n",
    "    # Set up the figure with a 2x2 grid layout\n",
    "    fig, axs = plt.subplots(2, 2, figsize=(15, 10))\n",
    "\n",
    "    # Plot data points for the first set\n",
    "    axs[0, 0].axhline(y=0, color='r', linestyle='-')\n",
    "    # axs[0, 0].axvline(x=T1*hbarc, color='r')\n",
    "    axs[0, 0].axvline(x=T2*hbarc, color='r')\n",
    "    axs[0, 0].plot(Einput[chckind], Re11[chckind], '+', label='Re11')\n",
    "    axs[0, 0].plot(Einput[chckind], Im11[chckind], '*', label='Im11')\n",
    "    axs[0, 0].plot(Einput[chckind], ModEsq[chckind], 'o', label='ModSq')\n",
    "    axs[0, 0].legend(loc='upper left', frameon=True)\n",
    "    axs[0, 0].set_title('$T_{11}$')\n",
    "    axs[0, 0].set_xlabel('$E_{cm}$ (MeV)')\n",
    "    axs[0, 0].set_ylabel('$Re T_{11}$, $Im T_{11}$, $|T_{11}+T_{12}|^2$')\n",
    "\n",
    "    isactive = ['No', 'Yes']\n",
    "    # Display class information in a table for the first set\n",
    "    axs[0, 1].axis('off')  # Turn off axis for the table subplot\n",
    "    table_data = [\n",
    "        ['n', 'Energy Pole (MeV)', 'RS', 'Active']\n",
    "    ]\n",
    "    for i in range(10):\n",
    "        table_data.append([str(i+1), '{:.2f}'.format(data_info[chckind][0][i]),\n",
    "                           data_info[chckind][1][i], isactive[data_info[chckind][2][i]]])\n",
    "\n",
    "    table = axs[0, 1].table(cellText=table_data, cellLoc='center', colWidths=[0.2] * 4, loc='center')\n",
    "    table.auto_set_font_size(False)\n",
    "    table.set_fontsize(10)\n",
    "    table.scale(1.5, 1.5)  # Adjust the scaling as needed\n",
    "\n",
    "    # Plot data points for the second set\n",
    "    axs[1, 0].axhline(y=0, color='r', linestyle='-')\n",
    "    # axs[1, 0].axvline(x=T1*hbarc, color='r')\n",
    "    axs[1, 0].axvline(x=T2*hbarc, color='r')\n",
    "    axs[1, 0].plot(Einput[chckind], Re22[chckind], '+', label='Re22')\n",
    "    axs[1, 0].plot(Einput[chckind], Im22[chckind], '*', label='Im22')\n",
    "    axs[1, 0].plot(Einput[chckind], ModEsq[chckind], 'o', label='ModSq2')\n",
    "    axs[1, 0].legend(loc='upper left', frameon=True)\n",
    "    axs[1, 0].set_title('$T_{22}$')\n",
    "    axs[1, 0].set_xlabel('$E_{cm}$ (MeV)')\n",
    "    axs[1, 0].set_ylabel('$Re T_{22}$, $Im T_{22}$, $|T_{11}+T_{12}|^2$')\n",
    "\n",
    "    # Display class information in a table for the second set\n",
    "    axs[1, 1].axis('off')  # Turn off axis for the table subplot\n",
    "    table_data = [\n",
    "        ['n', 'Energy Pole (MeV)', 'RS', 'Active']\n",
    "    ]\n",
    "    for i in range(10):\n",
    "        table_data.append([str(i+1), '{:.2f}'.format(data_info[chckind][0][i]),\n",
    "                           data_info[chckind][1][i], isactive[data_info[chckind][2][i]]])\n",
    "\n",
    "    table = axs[1, 1].table(cellText=table_data, cellLoc='center', colWidths=[0.2] * 4, loc='center')\n",
    "    table.auto_set_font_size(False)\n",
    "    table.set_fontsize(10)\n",
    "    table.scale(1.5, 1.5)  # Adjust the scaling as needed\n",
    "\n",
    "    # Save the figure with a timestamp\n",
    "    run_identifier = generate_timestamp()\n",
    "\n",
    "    # Modify the file names in the code accordingly\n",
    "    figure_filename1 = f\"21bt_{run_identifier}.png\"\n",
    "    \n",
    "    # Save the figures with the updated file names\n",
    "    fig.savefig(os.path.join(save_dir, figure_filename1))\n",
    "\n",
    "    # Show the plot and return the selected index\n",
    "    plt.tight_layout()\n",
    "    return plt.show(), chckind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ea68b050-060a-4fb1-a76a-c0721979d614",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_traintest(directory,curriculum):\n",
    "    # If Nshuffle=0, it is understood that the testing dataset is to be prepared.\n",
    "    # Otherwise, it is the training dataset, and Nshuffle determines shuffling times.\n",
    "    out = directory\n",
    "    # import prepared dataset\n",
    "    inputtraining = pickle.load(open(os.path.join(out, 'labelin_curr{:02d}.pkl'.format(curriculum)), 'rb'))\n",
    "    outputtraining = pickle.load(open(os.path.join(out, 'labelout_curr{:02d}.pkl'.format(curriculum)), 'rb'))\n",
    "    inputtraining = np.float32(np.asarray(inputtraining))\n",
    "    \n",
    "    # shuffle the imported data\n",
    "    X_train, X_test, y_train, y_test = train_test_split(inputtraining, outputtraining, test_size=0.2, random_state=42, stratify=outputtraining)\n",
    "    train = chainer.datasets.TupleDataset(X_train, y_train)\n",
    "    # split training set with testing set\n",
    "    pickle.dump(train, open(os.path.join(out, 'chainer_train_curr{:02d}.pkl'.format(curriculum)), 'wb'), protocol=4) #pickle.dump(train, open(os.path.join(out, 'chainer_train.pkl'), 'wb'), protocol=4)\n",
    "        \n",
    "    test = chainer.datasets.TupleDataset(X_test, y_test)\n",
    "    # split training set with the testing set\n",
    "    pickle.dump(test, open(os.path.join(out, 'chainer_test_curr{:02d}.pkl'.format(curriculum)), 'wb'), protocol=4)#pickle.dump(test, open(os.path.join(out, 'chainer_test.pkl'), 'wb'), protocol=4)\n",
    "\n",
    "    print(f'Size of training dataset: {len(X_train)}')\n",
    "    print(f'Size of testing dataset: {len(X_test)}')\n",
    "    print(f\"Test output values: {np.unique(y_test, return_counts=True)[0]}\")\n",
    "    print(f\"Test output value counts: {np.unique(y_test, return_counts=True)[1]}\")\n",
    "        \n",
    "    return\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "951fdfa2-72e0-4b62-ad56-0d9d596b4eb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 100 #49\n",
    "\n",
    "Einput = np.linspace(4212,4412, N)\n",
    "\n",
    "M = 5619.6 #mass of the parent particle (\\Lambda_b^0)\n",
    "m1 = 3096.9 #mass of particle 1 (J/\\psi)\n",
    "m2 = 493.677 #mass of particle 2 (K^-)\n",
    "m3 = 938.27208816 #mass of particle 3 (p)\n",
    "proj_axis = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3ba88d18-3098-4872-8de4-838332287dc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Dalitz(M, m1, m2, m3, proj_axis):\n",
    "    \n",
    "    \"\"\"\n",
    "    input: \n",
    "        M: mass of the parent particle in GeV\n",
    "        m1, m2, m3: mass of the final states in GeV\n",
    "\n",
    "    output:\n",
    "        Dalitz plot in array and figure\n",
    "    \"\"\"\n",
    "\n",
    "    m23sq = np.linspace((m2 + m3)**2, (M - m1)**2, 10000)\n",
    "    #m23sq = np.linspace(2.5e6, 6.5e6, 10000) #invariant mass of the second and third final states \n",
    "    #m13sq = np.linspace((m1 + m3)**2, (M - m2)**2, N) #invariant mass of the first and third final states\n",
    "    m13sq = Einput**2\n",
    "    \n",
    "    E_1 = (M**2 + m1**2 - m23sq)/(2*M)\n",
    "    E_2 = (M**2 + m2**2 - m13sq)/(2*M)\n",
    "\n",
    "    X, Y = np.meshgrid(E_1, E_2)\n",
    "    condition = (4*(X**2 - m1**2)*(Y**2 - m2**2) - (M**2 + m1**2 + m2**2 - m3**2 - 2*M*(X + Y) + 2*X*Y)**2) >=0\n",
    "    E_3 = (M - X - Y) >= 0\n",
    "\n",
    "    plot = condition*E_3 > 0\n",
    "\n",
    "    #cp = plt.pcolormesh(m23sq, m13sq, plot, cmap='Blues')\n",
    "\n",
    "    phase_space = np.sum(plot, axis=proj_axis)\n",
    "    \n",
    "    return phase_space\n",
    "\n",
    "phase_space0 = Dalitz(M, m1, m2, m3, proj_axis)\n",
    "norm = np.linalg.norm(phase_space0)\n",
    "phase_space = phase_space0/norm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "509ef7c3-a47e-40f9-958e-7668ed99bbbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def polynomial(coeff, x):\n",
    "    total = np.sum([coeff[i]*(x**i) for i in range(len(coeff))], axis=0)\n",
    "    norm = np.linalg.norm([total])\n",
    "    poly_bg = total/norm\n",
    "    return poly_bg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f37868-e835-4437-8a80-8021bdce3e03",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}