latexdemo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.axisartist.axislines import AxesZero\n",
    "import matplotlib.gridspec as gridspec\n",
    "from matplotlib import cm, transforms\n",
    "import matplotlib.ticker as mtick\n",
    "from mpl_axes_aligner import align\n",
    "import numpy as np\n",
    "np.seterr(divide='ignore')\n",
    "from scipy.stats import poisson, norm, lognorm, gamma\n",
    "from scipy import optimize as opti\n",
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "from scipy import special\n",
    "from scipy.stats import norm\n",
    "from scipy.stats import multivariate_normal\n",
    "from scipy.signal import savgol_filter\n",
    "import h5py\n",
    "import torch\n",
    "import cupy as cp\n",
    "\n",
    "import wf_func as wff\n",
    "\n",
    "np.random.seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "Mu = 4\n",
    "Tau = 20\n",
    "Sigma = 5\n",
    "file = '4.0-20-5'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "with h5py.File('waveform/' + file + '.h5', 'r', libver='latest', swmr=True) as ipt:\n",
    "    ent = ipt['Readout/Waveform'][:]\n",
    "    tru = ipt['SimTriggerInfo/PEList'][:]\n",
    "    gmu = ipt['SimTriggerInfo/PEList'].attrs['gmu']\n",
    "    gsigma = ipt['SimTriggerInfo/PEList'].attrs['gsigma']\n",
    "    t0truth = ipt['SimTruth/T'][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normcombine(x, m, s, a):\n",
    "    return a[0] * norm.pdf((x - m[0]) / s[0]) + a[1] * norm.pdf((x - m[1]) / s[1])\n",
    "\n",
    "def normcombine2d(x, m, s, a, rho):\n",
    "    return a[0, 0] * multivariate_normal.pdf(x, mean=[m[0, 0], m[1, 0]], cov=matrix(s[0, 0], s[1, 0], rho[0, 0])) + a[0, 1] * multivariate_normal.pdf(x, mean=[m[0, 0], m[1, 1]], cov=matrix(s[0, 0], s[1, 1], rho[0, 1])) + a[1, 0] * multivariate_normal.pdf(x, mean=[m[0, 1], m[1, 0]], cov=matrix(s[0, 1], s[1, 0], rho[1, 0])) + a[1, 1] * multivariate_normal.pdf(x, mean=[m[0, 1], m[1, 1]], cov=matrix(s[0, 1], s[1, 1], rho[1, 1]))\n",
    "\n",
    "def matrix(sx, sy, rho):\n",
    "    return np.array([[sx ** 2, rho * sx * sy], [rho * sx * sy, sy ** 2]])\n",
    "\n",
    "def chargehist(t):\n",
    "    c = norm.pdf(t, loc=gmu, scale=gsigma)\n",
    "#     q1 = 150.8\n",
    "#     sigma = 37.59\n",
    "#     w = 2.433e-5\n",
    "#     alpha = 0.01335\n",
    "#     mu = 2.851e-5\n",
    "#     c = np.exp(-mu)*(w*alpha*np.exp(-alpha*t))\n",
    "#     c = c + mu*np.exp(-mu)*(\n",
    "#         (1-w)/(sigma*np.sqrt(2*np.pi))*np.exp(-(t-q1)**2/(2*sigma**2))+\n",
    "#         w*(alpha/2*np.exp(-alpha*(t-q1-alpha/2*sigma**2))*(1+special.erf(t-q1-alpha*sigma**2)/(np.sqrt(2)*sigma))))\n",
    "    return c"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "Thres = wff.Thres\n",
    "std = 1.\n",
    "spe_pre = wff.read_model('spe.h5', 1)\n",
    "p = spe_pre[0]['parameters']\n",
    "window = wff.window\n",
    "t_auto = np.arange(window).reshape(window, 1) - np.arange(window).reshape(1, window)\n",
    "mnecpu = wff.spe((t_auto + np.abs(t_auto)) / 2, p[0], p[1], p[2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Bbox([[0.125, 0.12], [0.9, 0.89]])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fig = plt.figure(figsize=(8, 6))\n",
    "t = np.arange(-4 * 5, 5 * 20, 0.1)\n",
    "# gs = gridspec.GridSpec(1, 1, figure=fig, left=0.15, right=0.95, top=0.95, bottom=0.15, wspace=0.4, hspace=0.5)\n",
    "# ax = fig.add_subplot(gs[0, 0])\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax.plot(t, wff.convolve_exp_norm(t, 20, 0), label=r'$(20,0)$', color='g')\n",
    "ax.plot(t, wff.convolve_exp_norm(t, 0, 5), label=r'$(0,5)$', color='r')\n",
    "ax.plot(t, wff.convolve_exp_norm(t, 20, 5), label=r'$(20,5)$', color='b')\n",
    "ax.annotate('', \n",
    "            xy=(17, 0.07), \n",
    "            xytext=(17, 0.04), \n",
    "            arrowprops=dict(facecolor='k', arrowstyle='<|-|>'))\n",
    "ax.text(x=20, y=0.055, s=r'$\\mu$')\n",
    "ax.annotate('', \n",
    "            xy=(-13, 0.09), \n",
    "            xytext=(13, 0.09), \n",
    "            arrowprops=dict(facecolor='k', arrowstyle='<|-|>'))\n",
    "ax.text(x=0, y=0.093, s=r'$t_0$')\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.grid()\n",
    "ax.set_xlim(xmin=-4 * int(5))\n",
    "ax.set_ylabel(r'$\\textrm{probability density}$')\n",
    "ax.legend(title=r'$(\\tau_\\ell, \\sigma_\\ell)/\\si{ns}$', loc='upper right')\n",
    "ax.set_ylim(0, ax.get_ylim()[1] * 1.2)\n",
    "# ax.annotate(r'$t_{0}$', xy=(0, 0), xytext=(5, 0.01), arrowprops=dict(facecolor='k', shrink=0.1, width=0.1, headwidth=2))\n",
    "fig.savefig('Note/figures/profile.pgf')\n",
    "fig.savefig('Note/figures/profile.pdf')\n",
    "plt.close()\n",
    "ax.get_position()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 4))\n",
    "# fig.tight_layout()\n",
    "gs = gridspec.GridSpec(1, 1, figure=fig, left=0.05, right=0.97, top=0.97, bottom=0.1, wspace=0.3, hspace=0.3)\n",
    "ax = fig.add_subplot(gs[0, 0])\n",
    "\n",
    "ax.spines['left'].set_position(('data', 0))\n",
    "ax.spines['bottom'].set_position(('data', 0))\n",
    "ax.plot(1, 0, '>k', transform=ax.get_yaxis_transform(), clip_on=False)\n",
    "ax.plot(0, 1, '^k', transform=ax.get_xaxis_transform(), clip_on=False)\n",
    "\n",
    "t = np.linspace(0, 6, 201)\n",
    "ax.plot(t, lognorm.pdf(t, loc=0, s=0.3), color='darkorange')\n",
    "ax.plot(t, lognorm.pdf(t, loc=3, s=0.3), color='darkblue')\n",
    "ax.fill_between(t, 0, lognorm.pdf(t, loc=0, s=0.3), color='darkorange', alpha=0.5)\n",
    "ax.fill_between(t, 0, lognorm.pdf(t, loc=3, s=0.3), color='darkblue', alpha=0.5)\n",
    "ax.set_xlim(0, 6)\n",
    "ax.set_ylim(bottom=1e-3)\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.set_xlabel(r'$\\mathrm{Time}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}$')\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "ax.annotate(text='', xy=(1.5, 1), xytext=(3.5, 1), arrowprops=dict(arrowstyle='<->'))\n",
    "ax.text(x=2.2, y=1.1, s=r'$\\sim D_\\mathrm{w}$')\n",
    "ax.text(x=0.7, y=0.3, s=r'$\\sim \\mathrm{RSS}$')\n",
    "ax.text(x=3.7, y=0.3, s=r'$\\sim \\mathrm{RSS}$')\n",
    "fig.savefig('Note/figures/tab.pgf')\n",
    "fig.savefig('Note/figures/tab.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>TriggerNo</th>\n",
       "      <th>PMTId</th>\n",
       "      <th>HitTime</th>\n",
       "      <th>Charge</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>469.16</td>\n",
       "      <td>96.56</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>471.43</td>\n",
       "      <td>200.71</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>487.39</td>\n",
       "      <td>220.02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>499.50</td>\n",
       "      <td>165.30</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>527.69</td>\n",
       "      <td>188.54</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  TriggerNo PMTId HitTime  Charge\n",
       "0         2     0  469.16   96.56\n",
       "1         2     0  471.43  200.71\n",
       "2         2     0  487.39  220.02\n",
       "3         2     0  499.50  165.30\n",
       "4         2     0  527.69  188.54"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "i = 2\n",
    "cid = ent[i]['ChannelID']\n",
    "eid = ent[i]['TriggerNo']\n",
    "truth = np.sort(tru[(tru['TriggerNo'] == eid) & (tru['PMTId'] == cid)], kind='stable', order=['TriggerNo', 'PMTId', 'HitPosInWindow'])\n",
    "wave = ent[i]['Waveform'].astype(np.float) * spe_pre[cid]['epulse']\n",
    "df = pd.DataFrame(truth)\n",
    "df = df.rename(columns={'HitPosInWindow':'HitTime'})\n",
    "charge = df['Charge'].copy()\n",
    "hittime = df['HitTime'].copy()\n",
    "df = df.astype({'Charge': 'float32'})\n",
    "df = df.astype({'TriggerNo' : 'str', 'PMTId' : 'str', 'HitTime' : 'str', 'Charge': 'str'})\n",
    "df['HitTime'] = ['{:.02f}'.format(s) for s in hittime]\n",
    "df['Charge'] = ['{:.02f}'.format(s) for s in charge]\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax1_lim = [min(0, wave.min() * 1.2), wave.max() * 1.3]\n",
    "ax2_lim = [ax1_lim[0] * 1.7 /  ax1_lim[-1], 1.7]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "ind = np.argwhere(wave > spe_pre[cid]['std'] * 5).flatten()\n",
    "xmin = ((ind.min() - spe_pre[cid]['mar_l']) // 20 - 1) * 20\n",
    "xmax = max(((ind.max() + spe_pre[cid]['mar_r']) // 20 + 1) * 20, xmin + 200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRIALS = 5000\n",
    "\n",
    "#profile\n",
    "def vcombine(A, cx, t, w_all):\n",
    "    '''\n",
    "    t is 2 x l_e\n",
    "    '''\n",
    "    frac, ti = cp.modf(t)\n",
    "    ti = cp.array(ti, np.int32)\n",
    "    A_vec = (1 - frac)[:, :, None] * A[w_all, :, ti] + frac[:, :, None] * A[w_all, :, ti+1]\n",
    "    c_vec = (1 - frac)[:, :, None] * cx[w_all, :, ti] + frac[:, :, None] * cx[w_all, :, ti+1]\n",
    "    return A_vec, c_vec\n",
    "\n",
    "vstep = cp.array((-1, 1), np.float32)\n",
    "\n",
    "#profile\n",
    "def vmove1(A_vec, c_vec, z, fmu, fsig2s_inv, det_fsig2s_inv, b0):\n",
    "    '''\n",
    "    A_vec: 行向量 x 2\n",
    "    c_vec: 行向量 x 2\n",
    "    vstep: 构造 (-1, 1) 则 np.einsum('ej,ejk,ek->e', fmu, fsig2s_inv, fmu) 恒等于 0\n",
    "\n",
    "    cp.einsum('eiw,ejw->eij', A_vec, c_vec)\n",
    "\n",
    "    cp.einsum('ejk,ek->ej', c_vec, z) + cp.einsum('ej,ejk->ek', fmu, fsig2s_inv)\n",
    "    cp.einsum('ej,ejk,ek->e', zc, beta, zc)\n",
    "    '''\n",
    "    ac = A_vec @ cp.transpose(c_vec, (0, 2, 1))\n",
    "    beta_inv = fsig2s_inv + (ac + cp.transpose(ac, (0, 2, 1))) / 2\n",
    "    beta = cp.linalg.inv(beta_inv)\n",
    "\n",
    "    zc = (c_vec @ z[:, :, None]).squeeze() + (fmu[:, None, :] @ fsig2s_inv).squeeze()\n",
    "    Δν = 0.5 * (zc[:, None, :] @ beta @ zc[:, :, None]).squeeze()\n",
    "\n",
    "    Δν += 0.5 * cp.log(cp.clip(cp.linalg.det(beta) * det_fsig2s_inv, 1/b0, b0))\n",
    "    return Δν, beta\n",
    "\n",
    "#profile\n",
    "def vmove2(A_vec, c_vec, fmu, A, beta):\n",
    "    '''\n",
    "    \"eij, eiv, ejw, ewt->evt\"\n",
    "    beta c_vec c_vec A\n",
    "    '''\n",
    "    Δcx = -cp.transpose(beta @ c_vec, (0, 2, 1)) @ (c_vec @ A)\n",
    "    Δz = -(fmu[:, None, :] @ A_vec).squeeze()\n",
    "    return Δcx, Δz\n",
    "\n",
    "def periodic(_h, _lt):\n",
    "    '''\n",
    "    enforce a periodic boundary at 0 and _lt - 1\n",
    "    '''\n",
    "    c_overflow = _h < 0\n",
    "    _h[c_overflow] += _lt[c_overflow] - 1\n",
    "    c_overflow = _h > _lt - 1\n",
    "    _h[c_overflow] -= _lt[c_overflow] - 1\n",
    "    return _h\n",
    "\n",
    "#profile\n",
    "def v_rt(_s, _ts, w_all):\n",
    "    frac, ti = np.modf(_s)\n",
    "    ti = np.array(ti, np.int32)\n",
    "    return (1 - frac) * _ts[w_all, ti] + frac * _ts[w_all, ti+1]\n",
    "    \n",
    "tau = Tau\n",
    "sigma = Sigma\n",
    "\n",
    "pi = np.pi\n",
    "b_t0 = [0., 600.]\n",
    "\n",
    "# def lc(t):\n",
    "#     return co + np.log(1.0 - erf((ass - t)/s2s)) - alpha*t\n",
    "\n",
    "def lc(t):\n",
    "    return wff.log_convolve_exp_norm(t, tau, sigma)\n",
    "\n",
    "if tau == 0.0:\n",
    "    sel_lc = cp.ElementwiseKernel(\n",
    "        'float32 t, bool sel',\n",
    "        'float32 lprob',\n",
    "        f\"lprob = sel ? -log({sigma}) - 0.5 * log(2.0 * {pi}) - 0.5 * (t / {sigma}) * (t / {sigma}) : 0\",\n",
    "        'sel_lc')\n",
    "elif sigma == 0.0:\n",
    "    sel_lc = cp.ElementwiseKernel(\n",
    "        'float32 t, bool sel',\n",
    "        'float32 lprob',\n",
    "        f\"lprob = sel ? (t > 0 ? -log({tau}) - t / {tau} : -1e4) : 0\",\n",
    "        'sel_lc')\n",
    "else:\n",
    "    alpha = 1 / tau\n",
    "    co = -np.log(2.0 * tau) + alpha * alpha * sigma * sigma / 2.0\n",
    "    ass = alpha * sigma * sigma\n",
    "    s2s = np.sqrt(2.0) * sigma\n",
    "    sel_lc = cp.ElementwiseKernel(\n",
    "        'float32 t, bool sel',\n",
    "        'float32 lprob',\n",
    "        f\"lprob = sel ? {co} + log(1.0 - erf(({ass} - t) / {s2s})) - {alpha} * t : 0\",\n",
    "        'sel_lc')\n",
    "\n",
    "sel_add = cp.ElementwiseKernel('float32 delta, bool sel', 'float32 dest', \"if(sel) dest += delta\", \"sel_add\")\n",
    "sel_assign = cp.ElementwiseKernel('float32 value, bool sel', 'float32 dest', \"if(sel) dest = value\", \"sel_assign\")\n",
    "\n",
    "#profile\n",
    "def batch(A, cx, index, tq, s, z, t0_min=100, t0_max=500):\n",
    "    \"\"\"\n",
    "    batch\n",
    "    ====\n",
    "    连续时间游走\n",
    "    cx: Cov^-1 * A, 详见 FBMP\n",
    "    s: list of PE locations\n",
    "    mu_t: LucyDDM 的估算 PE 数\n",
    "    z: residue waveform (raw - predicted)\n",
    "\n",
    "    home_s 是在指标意义上的连续变量\n",
    "\n",
    "    0   1   2   3   4   5   6   7   8   9   l_t = 10\n",
    "    +---+---+---+---+---+---+---+---+---+\n",
    "    0  0.1 0.2 0.2 0.5 0.4 0.3 0.2  0   0   q_s\n",
    "    0  0.1 0.3 0.5 1.0 1.4 1.7 1.9 1.9 1.9  cq\n",
    "     _h \\in [0, 9) 范围的测度是 9\n",
    "\n",
    "    int(_h) 是插值的下指标 占比为 1 - decimal(_h)\n",
    "    int(_h) + 1 是插值的上指标 占比为 decimal(_h)\n",
    "    \"\"\"\n",
    "    sig2w = cp.asarray(index[\"sig2w\"], np.float32)\n",
    "    sig2s = cp.asarray(index[\"sig2s\"], np.float32)\n",
    "    mus = cp.asarray(index[\"mus\"], np.float32)\n",
    "    NPE = index[\"NPE\"]\n",
    "\n",
    "    a0 = A[:, :, 0]\n",
    "    b0 = 1 + sig2s * (a0[:, None, :] @ a0[:, :, None]).squeeze() / sig2w\n",
    "\n",
    "    l_e = len(index)\n",
    "    l_s = s.shape[1]\n",
    "    w_all = cp.arange(l_e) # index of all the waveforms\n",
    "\n",
    "    l_t = tq.shape[1]\n",
    "    # istar [0, 1) 之间的随机数，用于点中 PE\n",
    "    istar = np.random.rand(l_e, TRIALS) # 同时可用于创生位置的选取\n",
    "\n",
    "    # 根据 p_cha 采样得到的 PE 序列。供以后的产生过程使用。这两行是使用了 InverseCDF 算法进行的MC采样。\n",
    "    fp = np.arange(l_t)\n",
    "    home_s = np.array([np.interp(_is, xp=_xp[:_lt], fp=fp[:_lt]) for _is, _xp, _lt in zip(istar, tq[\"cq\"], index[\"l_t\"])])\n",
    "\n",
    "    t0 = np.zeros(l_e, np.float32)\n",
    "    e_hit = NPE > 0\n",
    "\n",
    "    t0[e_hit] = v_rt(s[e_hit, 0], tq[\"t_s\"], e_hit)\n",
    "    e_nonhit = ~e_hit\n",
    "    t0[e_nonhit] = tq[\"t_s\"][e_nonhit, 0]\n",
    "    t0 = np.clip(t0, t0_min, t0_max)\n",
    "\n",
    "    # step: +1 创生一个 PE， -1 消灭一个 PE， +2 向左或向右移动\n",
    "    flip = np.random.choice((-1, 1, 2), (l_e, TRIALS), p=np.array((1, 1, 2)) / 4)\n",
    "    Δν_g = cp.zeros(l_e, dtype=np.float32)\n",
    "    Δν = np.zeros(l_e, dtype=np.float32)\n",
    "    ν = np.zeros(l_e, dtype=np.float32)\n",
    "    ν_max = np.zeros(l_e, dtype=np.float32)\n",
    "    Δν_history = np.zeros((l_e, TRIALS), dtype=np.float32) # list of Δν's\n",
    "    annihilations = np.zeros((l_e, TRIALS)) # float64\n",
    "    creations = np.zeros((l_e, TRIALS)) # float64\n",
    "    t0_history = np.zeros((l_e, TRIALS), dtype=np.float32)\n",
    "    s0_history = np.zeros((l_e, TRIALS), dtype=np.uint32) # 0-norm of s\n",
    "    s_history = np.zeros((l_e, TRIALS * l_s), dtype=np.float32)\n",
    "\n",
    "    # t0_accept = 0\n",
    "\n",
    "    log_mu = np.log(index[\"mu0\"])  # 猜测的 Poisson 流强度\n",
    "    loc = np.zeros((l_e, 2)) # float64\n",
    "    fsig2s_inv = cp.diag(vstep)[None, :, :] * (1/sig2s)[:, None, None]\n",
    "    det_fsig2s_inv = cp.linalg.det(fsig2s_inv)\n",
    "    fmu = vstep[None, :] * mus[:, None]\n",
    "\n",
    "    last_max_s = s.copy()\n",
    "    s_max_index = np.zeros(l_e, dtype=np.uint32)\n",
    "\n",
    "    for i, (t, step, home, wander, wt, accept, acct) in enumerate(\n",
    "        zip(\n",
    "            istar.T,\n",
    "            flip.T,\n",
    "            home_s.T,\n",
    "            np.random.normal(size=(TRIALS, l_e)),\n",
    "            np.random.normal(scale=10, size=(TRIALS, l_e)),\n",
    "            *np.log(np.random.rand(2, TRIALS, l_e)),\n",
    "        )\n",
    "    ):\n",
    "        if i % 1000 == 0:\n",
    "            print(i)\n",
    "\n",
    "        mNPE = np.max(NPE)\n",
    "        rt = v_rt(s[:, :mNPE], tq[\"t_s\"], np.arange(l_e)[:, None])\n",
    "        nt0 = t0 + wt.astype(np.float32)\n",
    "        sel = cp.asarray(np.arange(mNPE)[None, :] < NPE[:, None])\n",
    "        lc0 = cp.sum(sel_lc(cp.asarray(rt - t0[:, None]), sel), axis=1)\n",
    "        lc1 = cp.sum(sel_lc(cp.asarray(rt - nt0[:, None]), sel), axis=1)\n",
    "        t0_acceptance = np.logical_and(cp.asnumpy(lc1 - lc0) >= acct, np.logical_and(nt0 >= t0_min, nt0 <= t0_max))\n",
    "        # t0_accept += np.sum(t0_acceptance) / len(t0_acceptance)\n",
    "        np.putmask(t0, t0_acceptance, nt0)\n",
    "        t0_history[:, i] = t0\n",
    "\n",
    "        ### 光变曲线和移动计算\n",
    "        Δν[:] = 0\n",
    "        e_bounce = NPE == 0\n",
    "        step[e_bounce] = 1\n",
    "        accept[e_bounce] += np.log(4)  # 惩罚\n",
    "        ea_bounce = np.logical_and(NPE == 1, step == -1)\n",
    "        accept[ea_bounce] -= np.log(4) # 1 -> 0: 行动后从 0 脱出的几率大，需要鼓励\n",
    "\n",
    "        e_create = step == 1\n",
    "        e_minus = ~e_create\n",
    "        e_move = step == 2\n",
    "        e_pm = ~e_move\n",
    "        e_annihilate = step == -1\n",
    "        e_plus = ~e_annihilate\n",
    "        loc[e_create, 0] = l_t # 0 A_vec\n",
    "        op = np.array(t * NPE, dtype=np.int32)\n",
    "        loc[e_minus, 0] = s[e_minus, op[e_minus]] # annihilate + move\n",
    "        loc[e_move, 1] = periodic(loc[e_move, 0] + wander[e_move], index[\"l_t\"][e_move])\n",
    "        loc[e_create, 1] = periodic(home[e_create], index[\"l_t\"][e_create])\n",
    "        loc[e_annihilate, 1] = l_t\n",
    "\n",
    "        ### 矩阵 Δν 计算\n",
    "        vA, vc = vcombine(A, cx, cp.asarray(loc, np.float32), w_all[:, None])\n",
    "        Δν_g, beta = vmove1(vA, vc, z, fmu, fsig2s_inv, det_fsig2s_inv, b0)\n",
    "\n",
    "        ## -1 cases, step == 2, -1\n",
    "        Δν[e_minus] -= lc(v_rt(loc[e_minus, 0], tq[\"t_s\"], e_minus) - t0[e_minus])\n",
    "        ## +1 cases, step == 2, 1\n",
    "        Δν[e_plus] += lc(v_rt(loc[e_plus, 1], tq[\"t_s\"], e_plus) - t0[e_plus])\n",
    "        ## non-move cases, step == 1, -1\n",
    "        NPE[e_create] += 1\n",
    "        loc[e_annihilate, 1] = loc[e_annihilate, 0]\n",
    "        Δν[e_pm] += step[e_pm] * (log_mu[e_pm] - np.log(tq[\"q_s\"][e_pm, np.array(loc[e_pm, 1], dtype=np.int32) + 1]) - np.log(NPE[e_pm]))\n",
    "        loc[e_annihilate, 1] = l_t\n",
    "        NPE[e_create] -= 1\n",
    "        ########\n",
    "        Δν += cp.asnumpy(Δν_g)\n",
    "        #######\n",
    "\n",
    "        ### 计算 Δcx, Δz, 更新 cx 和 z。对 accept 进行特别处理\n",
    "        e_accept = Δν >= accept\n",
    "        _e_accept = cp.asarray(e_accept)\n",
    "        Δcx, Δz = vmove2(vA, vc, fmu, A, beta)\n",
    "        sel_add(Δcx, _e_accept[:, None, None], cx)\n",
    "        sel_add(Δz, _e_accept[:, None], z)\n",
    "        ########\n",
    "\n",
    "        # 增加\n",
    "        ea_create = np.logical_and(e_accept, e_create)\n",
    "        ea_plus = np.logical_and(e_accept, e_plus)\n",
    "        creations[ea_plus, i] = loc[ea_plus, 1]\n",
    "\n",
    "        s[ea_create, NPE[ea_create]] = loc[ea_create, 1]\n",
    "        NPE[ea_create] += 1\n",
    "        # 减少\n",
    "        ea_annihilate = np.logical_and(e_accept, step == -1)\n",
    "        ea_minus = np.logical_and(e_accept, e_minus)\n",
    "        annihilations[ea_minus, i] = loc[ea_minus, 0]\n",
    "\n",
    "        NPE[ea_annihilate] -= 1\n",
    "        s[ea_annihilate, op[ea_annihilate]] = s[ea_annihilate, NPE[ea_annihilate]]\n",
    "        # 移动\n",
    "        ea_move = np.logical_and(e_accept, e_move)\n",
    "        s[ea_move, op[ea_move]] = loc[ea_move, 1]\n",
    "\n",
    "        Δν[~e_accept] = 0\n",
    "        step[~e_accept] = 0\n",
    "\n",
    "        ν += Δν\n",
    "        breakthrough = ν > ν_max\n",
    "        ν_max[breakthrough] = ν[breakthrough]\n",
    "        s_max_index[breakthrough] = i\n",
    "        last_max_s[breakthrough] = s[breakthrough]\n",
    "        Δν_history[:, i] = Δν\n",
    "        flip[:, i] = step\n",
    "        s0_history[:, i] = NPE\n",
    "        s_history[:, i*l_s:(i+1)*l_s] = s\n",
    "    # print(\"t0 acceptance:\", t0_accept / TRIALS)\n",
    "    # Assign one PE to the waveform without any PE\n",
    "    s_init = np.full(l_s, 0.)\n",
    "    s_init[0] = 1.\n",
    "    last_max_s[last_max_s.sum(axis=1) == 0] = s_init\n",
    "    return flip, s0_history, t0_history, Δν_history, s_max_index, last_max_s, annihilations, creations, s_history\n",
    "\n",
    "def get_t0(a0, a1, s0_history, t0_history, loc, flip, index, tq, t00_l):\n",
    "    l_e = a1 - a0\n",
    "    l_s = loc.shape[1] // TRIALS\n",
    "    t0_l = np.empty(l_e)\n",
    "    mu_l = np.empty(l_e)\n",
    "    for i in range(a0, a1):\n",
    "        # accept = flip[i] != 0\n",
    "        accept = np.full(len(flip[i]), True)\n",
    "        NPE = s0_history[i]\n",
    "        number_sample_zero = np.sum(NPE == 0)\n",
    "        t00_list = np.repeat(t0_history[i][accept], NPE[accept])\n",
    "        step = np.repeat(np.arange(TRIALS)[accept], NPE[accept])\n",
    "        idx_base = np.arange(TRIALS)[accept] * l_s\n",
    "        idx = np.hstack([i_b + np.arange(npe) for i_b, npe in zip(idx_base, NPE[accept])])\n",
    "        mu_t = index[\"mu0\"][i]\n",
    "        b_mu = [max(1e-8, mu_t - 5 * np.sqrt(mu_t)), mu_t + 5 * np.sqrt(mu_t)]\n",
    "        l_t = index[\"l_t\"][i]\n",
    "        loc_i = loc[i][idx]\n",
    "        ilp_cha = np.log(1 / tq[\"q_s\"][i][:l_t])\n",
    "        guess = ilp_cha[loc[i][idx].astype(int)]\n",
    "        loc_i = np.interp(loc_i, xp=np.arange(0.5, l_t), fp=tq[\"t_s\"][i][:l_t])\n",
    "        # t00 = index[\"t0\"][i]\n",
    "        # t00 = loc_i.mean() + 1\n",
    "        t00 = t00_l[i]\n",
    "        t0_l[i - a0], mu_l[i - a0] = wff.fit_t0mu_gibbs(loc_i, t00_list, step, number_sample_zero, tau, sigma, mu_t, t00, b_mu, b_t0, TRIALS)\n",
    "    return t0_l, mu_l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "with h5py.File('result/fsmp_test/sparsify/' + file + '.h5', \"r\", libver=\"latest\", swmr=True) as ipt:\n",
    "    A = ipt[\"A\"][:]\n",
    "    cx = ipt[\"cx\"][:]\n",
    "    index = ipt[\"index\"][:]\n",
    "    s = ipt[\"s\"][:]\n",
    "    tq = ipt[\"tq\"][:]\n",
    "    z = ipt[\"z\"][:]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "42 173 6\n",
      "0\n",
      "1000\n",
      "2000\n",
      "3000\n",
      "4000\n"
     ]
    }
   ],
   "source": [
    "# Try using Metropolis\n",
    "# TRIALS = 5000\n",
    "# n = 1\n",
    "# b_t0 = [0., 600.]\n",
    "\n",
    "# # initialization\n",
    "# A, y, tlist, t0_t, t0_delta, cha, left_wave, right_wave = wff.initial_params(wave[::wff.nshannon], spe_pre[cid], Tau, Sigma, gmu, Thres['lucyddm'], p, is_t0=True, is_delta=False, n=n)\n",
    "# # assert len(np.unique(np.diff(tlist))) == 1\n",
    "# s_cha = np.cumsum(cha)\n",
    "# # moving average filter of size 2*n+1\n",
    "# cha = np.pad(s_cha[2*n+1:], (n+1, n), 'edge') - np.pad(s_cha[:-(2*n+1)], (n+1, n), 'edge')\n",
    "# cha += 1e-8 # for completeness of the random walk.\n",
    "# p_cha = cha / np.sum(cha)\n",
    "# mu_t = abs(y.sum() / gmu)\n",
    "\n",
    "# # Eq. (9) where the columns of A are taken to be unit-norm.\n",
    "# mus = np.sqrt(np.diag(np.matmul(A.T, A)))\n",
    "# assert np.std(mus) < 1e-4, 'mus must be equal'\n",
    "# mus = mus[0]\n",
    "# A = A / mus\n",
    "# p1 = mu_t * wff.convolve_exp_norm(tlist - t0_t, Tau, Sigma) / n + 1e-8\n",
    "# sig2w = spe_pre[cid]['std'] ** 2\n",
    "# sig2s = (gsigma * mus / gmu) ** 2\n",
    "\n",
    "# nu_star, T_star, c_star, es_history, NPE_evo, number_sample_zero = wff.metropolis_fsmp(y, A, sig2w, sig2s, mus, p1, p_cha, mu_t)\n",
    "\n",
    "# ilp_cha = np.log(cha.sum()) - np.log(cha)\n",
    "# guess = ilp_cha[es_history['loc'].astype(int)]\n",
    "# es_history['loc'] = np.interp(es_history['loc'], xp=np.arange(0.5, len(tlist)), fp=tlist)\n",
    "# ans = opti.fmin_l_bfgs_b(lambda x: -np.sum(wff.log_convolve_exp_norm(es_history['loc'] - x, Tau, Sigma)), x0=[t0_t], approx_grad=True, bounds=[b_t0], maxfun=500000)\n",
    "# t00 = ans[0].item() if ans[-1]['warnflag'] == 0 else t0_t\n",
    "# b_mu = [max(1e-8, mu_t - 5 * np.sqrt(mu_t)), mu_t + 5 * np.sqrt(mu_t)]\n",
    "# t0, mu = wff.fit_t0mu_guess(es_history['loc'], es_history['step'], \n",
    "#                             number_sample_zero, Tau, Sigma, guess, mu_t, t00, b_mu, b_t0, TRIALS)\n",
    "\n",
    "# j = 0\n",
    "# xmmse_most = np.zeros(len(tlist))\n",
    "# while np.all(xmmse_most <= 0):\n",
    "#     maxindex = nu_star.argsort()[::-1][j]\n",
    "#     zx = y - np.dot(A, mus * c_star[maxindex])\n",
    "#     Phi_s = wff.Phi(y, A, c_star[maxindex], mus, sig2s, sig2w)\n",
    "#     invPhi = np.linalg.inv(Phi_s)\n",
    "#     xmmse_most = mus * c_star[maxindex] + np.matmul(np.diagflat(sig2s * c_star[maxindex]), np.matmul(A.T, np.matmul(invPhi, zx)))\n",
    "#     j += 0\n",
    "# pet = np.repeat(tlist[xmmse_most > 0], c_star[maxindex][xmmse_most > 0])\n",
    "# cha = np.repeat(xmmse_most[xmmse_most > 0] / mus / c_star[maxindex][xmmse_most > 0], c_star[maxindex][xmmse_most > 0])\n",
    "\n",
    "# Trying to use wangyy's file\n",
    "# with h5py.File('result/fsmp/char/' + file + '.h5', 'r', libver='latest', swmr=True) as ipt:\n",
    "#     photoelectron = ipt['photoelectron'][:]\n",
    "\n",
    "# pet = photoelectron['HitPosInWindow'][(photoelectron['TriggerNo'] == eid) & (photoelectron['ChannelID'] == cid)]\n",
    "# cha = photoelectron['Charge'][(photoelectron['TriggerNo'] == eid) & (photoelectron['ChannelID'] == cid)] / gmu\n",
    "\n",
    "# with h5py.File('result/fsmp/solu/' + file + '.h5', 'r', libver='latest', swmr=True) as ipt:\n",
    "#     time = ipt['starttime'][:]\n",
    "\n",
    "# t0 = time['tswave'][(time['TriggerNo'] == eid) & (time['ChannelID'] == cid)]\n",
    "# mu = time['muwave'][(time['TriggerNo'] == eid) & (time['ChannelID'] == cid)]\n",
    "\n",
    "l_e = len(index)\n",
    "s_t = np.argsort(index[\"l_t\"])\n",
    "\n",
    "i_part = np.array([i, i+1])\n",
    "l_part = len(i_part)\n",
    "ind_part = index[i_part]\n",
    "lp_t = np.max(ind_part[\"l_t\"])\n",
    "lp_wave = np.max(ind_part[\"l_wave\"])\n",
    "lp_NPE = np.max(ind_part[\"NPE\"])\n",
    "print(lp_t, lp_wave, lp_NPE)\n",
    "\n",
    "null = np.zeros((l_part, lp_wave, 1), np.float32) # cx, A[:, :, -1] = 0  用于 +- 的空白维度\n",
    "s_null = np.zeros((l_part, lp_NPE * 2), np.float32) # 富余的 PE 活动空间\n",
    "(flip, s0_history, t0_history, Δν_history, s_max_index, last_max_s, annihilations, creations, loc\n",
    " ) = batch(cp.asarray(np.append(A[i_part, :lp_wave, :lp_t], null, axis=2), np.float32),\n",
    "           cp.asarray(np.append(cx[i_part, :lp_wave, :lp_t], null, axis=2), np.float32),\n",
    "           index[i_part], \n",
    "           tq[i_part, :lp_t], \n",
    "           np.append(s[i_part, :lp_NPE], s_null, axis=1),\n",
    "           cp.asarray(z[i_part, :lp_wave], np.float32))\n",
    "\n",
    "def get_t0(a0, a1, s0_history, t0_history, loc, flip, index, tq, t00_l):\n",
    "    l_e = a1 - a0\n",
    "    l_s = loc.shape[1] // TRIALS\n",
    "    t0_l = np.empty(l_e)\n",
    "    mu_l = np.empty(l_e)\n",
    "    for i in range(a0, a1):\n",
    "        # accept = flip[i] != 0\n",
    "        accept = np.full(len(flip[i]), True)\n",
    "        NPE = s0_history[i]\n",
    "        number_sample_zero = np.sum(NPE == 0)\n",
    "        t00_list = np.repeat(t0_history[i][accept], NPE[accept])\n",
    "        step = np.repeat(np.arange(TRIALS)[accept], NPE[accept])\n",
    "        idx_base = np.arange(TRIALS)[accept] * l_s\n",
    "        idx = np.hstack([i_b + np.arange(npe) for i_b, npe in zip(idx_base, NPE[accept])])\n",
    "        mu_t = index[\"mu0\"][i]\n",
    "        b_mu = [max(1e-8, mu_t - 5 * np.sqrt(mu_t)), mu_t + 5 * np.sqrt(mu_t)]\n",
    "        l_t = index[\"l_t\"][i]\n",
    "        loc_i = loc[i][idx]\n",
    "        ilp_cha = np.log(1 / tq[\"q_s\"][i][:l_t])\n",
    "        guess = ilp_cha[loc[i][idx].astype(int)]\n",
    "        loc_i = np.interp(loc_i, xp=np.arange(0.5, l_t), fp=tq[\"t_s\"][i][:l_t])\n",
    "        # t00 = index[\"t0\"][i]\n",
    "        # t00 = loc_i.mean() + 1\n",
    "        t00 = t00_l[i]\n",
    "        t0_l[i - a0], mu_l[i - a0] = wff.fit_t0mu_gibbs(loc_i, t00_list, step, number_sample_zero, tau, sigma, mu_t, t00, b_mu, b_t0, TRIALS)\n",
    "    return t0_l, mu_l"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "l_t = index['l_t'][i]\n",
    "tlist = tq['t_s'][i][:l_t]\n",
    "mus = index['mus'][i]\n",
    "sig2s = index['sig2s'][i]\n",
    "sig2w = index['sig2w'][i]\n",
    "y = ent[i]['Waveform'][index['a_wave'][i]:index['b_wave'][i]]\n",
    "A_i = A[i][:index['b_wave'][i]-index['a_wave'][i], :l_t]\n",
    "\n",
    "s_index = s0_history.flatten()[s_max_index + np.arange(2) * TRIALS]\n",
    "\n",
    "s = last_max_s[0][:s_index[0]]\n",
    "t, c = np.unique(np.sort(np.digitize(s, bins=np.arange(l_t)) - 1), return_counts=True)\n",
    "c_star = np.zeros(l_t, dtype=int)\n",
    "c_star[t] = c\n",
    "\n",
    "zx = y - np.dot(A_i, mus * c_star)\n",
    "Phi_s = wff.Phi(y, A_i, c_star, mus, sig2s, sig2w)\n",
    "invPhi = np.linalg.inv(Phi_s)\n",
    "xmmse_most = mus * c_star + np.matmul(np.diagflat(sig2s * c_star), np.matmul(A_i.T, np.matmul(invPhi, zx)))\n",
    "pet = np.repeat(tlist[xmmse_most > 0], c_star[xmmse_most > 0])\n",
    "cha = np.repeat(xmmse_most[xmmse_most > 0] / mus / c_star[xmmse_most > 0], c_star[xmmse_most > 0])\n",
    "mu_i = (c_star > 0).sum()\n",
    "t0_i, _ = wff.likelihoodt0(pet, char=cha, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='all')\n",
    "pet, cha = wff.clip(pet, cha, 0.0)\n",
    "# cha = cha * gmu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [469. 471. 487. 499. 527.], Weight = [0.33516652 1.51841157 1.3589372  1.03511034 1.14905059]\n",
      "wdist = 0.692857223254407, cdiff = -7.673615249627758\n",
      "RSS = 59.251812237554844\n",
      "1.935656595114267\n"
     ]
    }
   ],
   "source": [
    "t0, mu = get_t0(0, 1, s0_history, t0_history, loc, flip, index[i_part], tq[i_part, :lp_t], t0truth['T0'][i_part])\n",
    "\n",
    "pet, pwe = wff.clip(pet, cha, 0.0)\n",
    "pwe = pwe\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "wave_ylim = ax.get_ylim()\n",
    "fig.savefig('Note/figures/fsmp.pgf')\n",
    "fig.savefig('Note/figures/fsmp.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p)\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# fig = plt.figure(figsize=(8, 6))\n",
    "# # fig.tight_layout()\n",
    "# ax = fig.add_axes((.125, .12, .775, .77))\n",
    "# ax.plot(wave, label='Waveform')\n",
    "# ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "# ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "# ax.set_xlim(0, len(wave))\n",
    "# wave_ylim = ax.get_ylim()\n",
    "# ax.set_ylim(wave_ylim[0] * 1.05, wave_ylim[1] * 1.05)\n",
    "# ax.legend()\n",
    "# fig.savefig('Note/figures/wave.pgf')\n",
    "# fig.savefig('Note/figures/wave.pdf')\n",
    "# fig.savefig('Note/figures/wave.png')\n",
    "# fig.clf()\n",
    "# plt.close(fig)\n",
    "\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(truth['HitPosInWindow'], 0, truth['Charge'] / gmu, color='r', label='Charge', linewidth=1.0)\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax2.set_ylim(wave_ylim[0] / 30, wave_ylim[1] / 30)\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "ax.set_ylim(wave_ylim[0] * 0.7, wave_ylim[1] * 0.7)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/wave.pgf')\n",
    "fig.savefig('Note/figures/wave.pdf')\n",
    "fig.savefig('Note/figures/wave.png')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "# gs = gridspec.GridSpec(1, 1, figure=fig, left=0.15, right=0.85, top=0.95, bottom=0.15, wspace=0.4, hspace=0.5)\n",
    "# ax = fig.add_subplot(gs[0, 0])\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "t = np.arange(0, 100, 0.1)\n",
    "ax.plot(t, wff.spe(t, p[0], p[1], p[2]), color='b', label='Single PE response')\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.grid()\n",
    "ax.set_xlim(0, 80)\n",
    "ax.set_ylim(wave_ylim[0] * 0.7, wave_ylim[1] * 0.7)\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax.legend()\n",
    "fig.savefig('Note/figures/spe.pgf')\n",
    "fig.savefig('Note/figures/spe.pdf')\n",
    "fig.savefig('Note/figures/spe.png')\n",
    "plt.close()\n",
    "\n",
    "# fig = plt.figure(figsize=(8, 6))\n",
    "# # fig.tight_layout()\n",
    "# ax = fig.add_axes((.125, .12, .775, .77))\n",
    "# ax.vlines(truth['HitPosInWindow'], 0, truth['Charge'] / gmu, color='r', label='Charge')\n",
    "# ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "# ax.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "# ax.set_xlim(0, len(wave))\n",
    "# ax.set_ylim(wave_ylim[0] / 20, wave_ylim[1] / 20)\n",
    "# ax.axhline(y=0, color='k', linestyle='dashed', alpha=0.5)\n",
    "# ax.legend()\n",
    "# fig.savefig('Note/figures/charge.pgf')\n",
    "# fig.savefig('Note/figures/charge.pdf')\n",
    "# fig.savefig('Note/figures/charge.png')\n",
    "# fig.clf()\n",
    "# plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "# gs = gridspec.GridSpec(1, 1, figure=fig, left=0.15, right=0.85, top=0.95, bottom=0.15, wspace=0.4, hspace=0.5)\n",
    "# ax = fig.add_subplot(gs[0, 0])\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(truth['HitPosInWindow'], 0, truth['Charge'] / gmu, color='r', label='Charge')\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.hlines(2, 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/goal.pgf')\n",
    "fig.savefig('Note/figures/goal.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "869.1974059198186\n",
      "871.1250094058131\n"
     ]
    }
   ],
   "source": [
    "print(wave.sum())\n",
    "print(truth['Charge'][truth['Charge'] > 0].sum()) # made by noise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.load('result/takara/char/Channel00/cnn_testing_record_2022-04-09_06:20:24.npz')['arr_0']\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax.plot(np.arange(1, len(t)+1), t, label=r'$D_\\mathrm{w}$', color='C1')\n",
    "ax.set_xlabel(r'$\\mathrm{epoch}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Wasserstein\\ Distance}/\\si{ns}$')\n",
    "ax.legend()\n",
    "ax.grid()\n",
    "fig.savefig('Note/figures/epoch.pgf')\n",
    "fig.savefig('Note/figures/epoch.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481\n",
      " 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499\n",
      " 500 501 502 503 504 505 506 507 522 523 524 525 526 527 528 529 530 531\n",
      " 532 533 534 535], Weight = [0.04728835 0.0773302  0.12971889 0.14296212 0.17259763 0.18196951\n",
      " 0.16468631 0.16345413 0.15766181 0.12150366 0.10900383 0.10819887\n",
      " 0.08709687 0.05954987 0.05300041 0.04706557 0.04386505 0.05970238\n",
      " 0.10148353 0.11900193 0.14771352 0.15423084 0.1459499  0.13409241\n",
      " 0.12377091 0.11172638 0.10319991 0.08280433 0.08618624 0.07160304\n",
      " 0.10946468 0.10993917 0.13333807 0.12401741 0.12562733 0.11987713\n",
      " 0.10082539 0.09262077 0.07290046 0.06250504 0.05116336 0.03533935\n",
      " 0.04415651 0.04521806 0.05780158 0.07376266 0.10725069 0.10430699\n",
      " 0.1253408  0.11595931 0.09038139 0.09029217 0.08120139 0.06569599\n",
      " 0.05272566 0.04824543 0.05150447 0.04679789]\n",
      "wdist = 3.195640285348688, cdiff = 32.32581961192988\n",
      "RSS = 948.4877387238437\n",
      "2.695361674450112\n"
     ]
    }
   ],
   "source": [
    "pet, pwe = wff.threshold(wave, spe_pre[cid])\n",
    "pet, pwe = wff.clip(pet, pwe, Thres['threshold'])\n",
    "output = np.zeros(window)\n",
    "output[pet] = pwe\n",
    "alpha = opti.fmin_l_bfgs_b(lambda alpha: wff.rss_alpha(alpha, output, wave, mnecpu), x0=[0.01], approx_grad=True, bounds=[[1e-20, np.inf]], maxfun=50000)[0]\n",
    "pwe = pwe * alpha\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "ax2.annotate('', xy=(pet.mean(), pwe.max()*1.1), xytext=(pet.mean()+pet.ptp(), pwe.max()*1.1), arrowprops=dict(facecolor='k', shrink=0.01, width=2, headwidth=4))\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/threshold.pgf')\n",
    "fig.savefig('Note/figures/threshold.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p, fold='/tmp')\n",
    "t0 = wff.likelihoodt0(pet, char=pwe * gmu, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='charge')[0]\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [472 489 500 529], Weight = [1.60835669 1.30957107 1.22689949 1.13247991]\n",
      "wdist = 2.35038617939221, cdiff = -26.77229140635238\n",
      "RSS = 739.9283743731916\n",
      "4.852010259042913\n"
     ]
    }
   ],
   "source": [
    "pet, pwe = wff.findpeak(wave, spe_pre[cid])\n",
    "pet, pwe = wff.clip(pet, pwe, Thres['findpeak'])\n",
    "output = np.zeros(window)\n",
    "output[pet] = pwe\n",
    "alpha = opti.fmin_l_bfgs_b(lambda alpha: wff.rss_alpha(alpha, output, wave, mnecpu), x0=[0.01], approx_grad=True, bounds=[[1e-20, np.inf]], maxfun=50000)[0]\n",
    "pwe = pwe * alpha\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=1.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "loc = pet + spe_pre[cid]['peak_c']\n",
    "loc = loc[loc < window]\n",
    "amp = wave[loc]\n",
    "assert ax1_lim[-1] > amp.max() + 10\n",
    "for j in range(len(loc)):\n",
    "    ax.annotate('', xy=(loc[j], amp[j]+5), xytext=(loc[j], amp[j]+10), arrowprops=dict(facecolor='k', shrink=0.01, width=0.5, headwidth=2))\n",
    "ax2.annotate('', xy=(pet.mean(), pwe.max()*1.1), xytext=(pet.mean()+pet.ptp(), pwe.max()*1.1), arrowprops=dict(facecolor='k', shrink=0.01, width=2, headwidth=4))\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/findpeak.pgf')\n",
    "fig.savefig('Note/figures/findpeak.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p, fold='/tmp')\n",
    "t0 = wff.likelihoodt0(pet, char=pwe * gmu, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='charge')[0]\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [467 468 469 470 471 472 473 474 475 484 485 486 487 488 489 490 497 498\n",
      " 499 500 501 502 525 526 527 528 529 530 531], Weight = [0.13490864 0.20056036 0.25330008 0.28513933 0.29213436 0.27456125\n",
      " 0.23611504 0.18261099 0.12073645 0.11227218 0.18484422 0.24092033\n",
      " 0.26828681 0.2602826  0.21809638 0.15112532 0.1047817  0.16286432\n",
      " 0.19923032 0.20521867 0.18049444 0.13249467 0.12004205 0.16491299\n",
      " 0.19736295 0.20927639 0.19646852 0.1603138  0.10782198]\n",
      "wdist = 1.8696699565298651, cdiff = 18.00603457141207\n",
      "RSS = 153.70175366873292\n",
      "2.614635617851434\n"
     ]
    }
   ],
   "source": [
    "pet, pwe = wff.waveformfft(wave, spe_pre[cid])\n",
    "pet, pwe = wff.clip(pet, pwe, Thres['fftrans'])\n",
    "output = np.zeros(window)\n",
    "output[pet] = pwe\n",
    "alpha = opti.fmin_l_bfgs_b(lambda alpha: wff.rss_alpha(alpha, output, wave, mnecpu), x0=[0.01], approx_grad=True, bounds=[[1e-20, np.inf]], maxfun=50000)[0]\n",
    "pwe = pwe * alpha\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/fftrans.pgf')\n",
    "fig.savefig('Note/figures/fftrans.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p, fold='/tmp')\n",
    "t0 = wff.likelihoodt0(pet, char=pwe * gmu, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='charge')[0]\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [469 471 472 487 488 499 500 527 528], Weight = [0.45424007 1.17786671 0.21697766 0.57416642 0.85227729 0.50602988\n",
      " 0.48306009 0.45369149 0.7238494 ]\n",
      "wdist = 0.604852940125768, cdiff = -0.39650806898040825\n",
      "RSS = 10.01973168114615\n",
      "2.774549745000627\n"
     ]
    }
   ],
   "source": [
    "pet, pwe = wff.lucyddm(wave, spe_pre[cid]['spe'])\n",
    "pet, pwe = wff.clip(pet, pwe, Thres['lucyddm'])\n",
    "output = np.zeros(window)\n",
    "output[pet] = pwe\n",
    "alpha = opti.fmin_l_bfgs_b(lambda alpha: wff.rss_alpha(alpha, output, wave, mnecpu), x0=[0.01], approx_grad=True, bounds=[[1e-20, np.inf]], maxfun=50000)[0]\n",
    "pwe = pwe * alpha\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/lucyddm.pgf')\n",
    "fig.savefig('Note/figures/lucyddm.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p, fold='/tmp')\n",
    "t0 = wff.likelihoodt0(pet, char=pwe * gmu, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='charge')[0]\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [470 471 472 487 488 499 500 528], Weight = [0.8683057  0.25852837 0.75238555 0.73799834 0.6444972  0.49203569\n",
      " 0.50425877 1.15523982]\n",
      "wdist = 0.8129154294613066, cdiff = -5.021948382194073\n",
      "RSS = 19.534707730708483\n",
      "2.9598339189554395\n"
     ]
    }
   ],
   "source": [
    "with h5py.File('result/takara/char/' + file + '.h5', 'r', libver='latest', swmr=True) as ipt:\n",
    "    photoelec = ipt['photoelectron'][:]\n",
    "s = photoelec[(photoelec['TriggerNo'] == eid) & (photoelec['ChannelID'] == cid)]\n",
    "pet = s['HitPosInWindow']\n",
    "pwe = s['Charge']\n",
    "pwe = pwe / gmu\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "# ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/takara.pgf')\n",
    "fig.savefig('Note/figures/takara.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p, fold='/tmp')\n",
    "t0 = wff.likelihoodt0(pet, char=pwe * gmu, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='charge')[0]\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [465.  467.  468.5 469.5 471.5 472.  472.5 473.  474.  475.5 476.  482.5\n",
      " 483.5 486.  486.5 487.  487.5 488.5 489.  489.5 490.5 491.  492.  494.5\n",
      " 495.  496.  497.  498.5 500.  525.  527.  527.5 528.  529.  529.5 530.\n",
      " 530.5 531.  531.5 532. ], Weight = [0.00004056 0.00000161 0.50646792 0.00000606 1.3935919  0.00002333\n",
      " 0.00002546 0.0000265  0.0000705  0.00002203 0.00000904 0.00001444\n",
      " 0.00000447 0.00000808 0.00000127 0.000059   1.37724389 0.00001398\n",
      " 0.00000858 0.00001346 0.00002844 0.00002556 0.00001106 0.00000025\n",
      " 0.0000299  0.00003845 0.0000416  0.36676836 0.65237934 0.00000705\n",
      " 0.00002421 1.15792123 0.00000763 0.00003973 0.00002968 0.0000268\n",
      " 0.00001788 0.00000645 0.00002019 0.00004749]\n",
      "wdist = 0.6243421832926701, cdiff = 1.6777548751523437\n",
      "RSS = 8.834620517524197\n",
      "2.671449688188318\n"
     ]
    }
   ],
   "source": [
    "with h5py.File('result/mcmc/char/' + file + '.h5', 'r', libver='latest', swmr=True) as ipt:\n",
    "    photoelec = ipt['photoelectron'][:]\n",
    "s = photoelec[(photoelec['TriggerNo'] == eid) & (photoelec['ChannelID'] == cid)]\n",
    "pet = s['HitPosInWindow']\n",
    "pwe = s['Charge']\n",
    "pwe = pwe / gmu\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/mcmc.pgf')\n",
    "fig.savefig('Note/figures/mcmc.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p, fold='/tmp')\n",
    "t0 = wff.likelihoodt0(pet, char=pwe * gmu, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='charge')[0]\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PEnum is 5\n",
      "truth HitPosInWindow = [469.15547745 471.42601966 487.39357274 499.49537376 527.69280766], Weight = [0.60348736 1.25444866 1.37512431 1.0331502  1.1784267 ]\n",
      "truth RSS = 1003.0466121337955\n",
      "HitPosInWindow = [467. 468. 470. 471. 472. 487. 488. 499. 500. 527. 528.], Weight = [0.08310591 0.00032206 0.97774912 0.01575962 0.81853973 0.59466171\n",
      " 0.79408098 0.493306   0.51071231 0.32664292 0.83637863]\n",
      "wdist = 0.9068197801619707, cdiff = 1.0594602473590373\n",
      "RSS = 6.49430638214507\n",
      "2.7278145296716048\n"
     ]
    }
   ],
   "source": [
    "# pet, pwe = wff.xiaopeip(wave, spe_pre[cid], eta=0)\n",
    "pet, pwe = wff.xiaopeip(wave, spe_pre[cid], Tau, Sigma, Thres['lucyddm'], p, eta=0)\n",
    "pet, pwe = wff.clip(pet, pwe, Thres['xiaopeip'])\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "# fig.tight_layout()\n",
    "ax = fig.add_axes((.125, .12, .775, .77))\n",
    "ax2 = ax.twinx()\n",
    "ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "ax.plot(wave, label='Waveform')\n",
    "# ax.hlines(5 * spe_pre[cid]['std'], 0, window, color='g', label='Threshold')\n",
    "ax.set_xlim(xmin, xmax)\n",
    "lines, labels = ax.get_legend_handles_labels()\n",
    "lines2, labels2 = ax2.get_legend_handles_labels()\n",
    "ax2.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "ax2.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax2.legend(lines + lines2, labels + labels2)\n",
    "ax.set_xlabel(r'$\\mathrm{t}/\\si{ns}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Voltage}/\\si{mV}$')\n",
    "ax2.set_ylabel(r'$\\mathrm{Charge}$')\n",
    "ax.set_ylim(*ax1_lim)\n",
    "ax2.set_ylim(*ax2_lim)\n",
    "fig.savefig('Note/figures/xiaopeip.pgf')\n",
    "fig.savefig('Note/figures/xiaopeip.pdf')\n",
    "fig.clf()\n",
    "plt.close(fig)\n",
    "wff.demo(pet, pwe, truth, spe_pre[cid], window, wave, cid, p, fold='/tmp')\n",
    "t0 = wff.likelihoodt0(pet, char=pwe * gmu, gmu=gmu, Tau=Tau, Sigma=Sigma, mode='charge')[0]\n",
    "print((t0 - t0truth[i]['T0']).item())"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {},
   "source": [
    "with h5py.File('result/mcmc/solu/' + file + '.h5', 'r', libver='latest', swmr=True) as soluf, h5py.File('waveform/' + file + '.h5', 'r', libver='latest', swmr=True) as wavef:\n",
    "    start = wavef['SimTruth/T'][:]\n",
    "    time = soluf['starttime'][:]\n",
    "data = time['tswave'] - start['T0']\n",
    "vali = np.abs(data - np.mean(data)) < 5 * np.std(data, ddof=-1)\n",
    "data = data[vali]\n",
    "fig = plt.figure(figsize=(8, 6))\n",
    "ax = fig.add_axes((.1, .45, .85, .45))\n",
    "ax.hist(data, bins=100, density=1)\n",
    "ax.set_ylabel('arb. unit')\n",
    "ax.set_xticks([])\n",
    "ax.set_yticks([])\n",
    "ax.set_xlim(data.min() - 0.05, data.max() + 0.05)\n",
    "ax.set_ylim(0, ax.get_ylim()[1] * 1.05)\n",
    "axb = fig.add_axes((.1, .15, .85, .3))\n",
    "axb.boxplot(data, vert=False, sym='', patch_artist=True)\n",
    "axb.set_xlabel(r'$\\Delta t_{0}/\\si{ns}$')\n",
    "axb.set_yticks([])\n",
    "axb.set_xlim(ax.get_xlim())\n",
    "fig.savefig('Note/figures/mcmct0hist.pgf')\n",
    "fig.savefig('Note/figures/mcmct0hist.pdf')\n",
    "plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = ['lucyddm', 'xiaopeip', 'takara', 'fsmp', 'mcmc']\n",
    "\n",
    "for m in methods:\n",
    "    with h5py.File('result/' + m + '/dist/' + file + '.h5', 'r', libver='latest', swmr=True) as distfile:\n",
    "        dt = distfile['Record'][:]\n",
    "    N = np.percentile(dt['wdist'], 95)\n",
    "    M = 500\n",
    "\n",
    "    penum = np.unique(dt['NPE'])\n",
    "    l = min(50, penum.max())\n",
    "    wdist_stats = np.full((l, 6), np.nan)\n",
    "    edist_stats = np.full((l, 6), np.nan)\n",
    "    for i in range(l):\n",
    "        vali = dt['NPE'] == i+1\n",
    "        if np.sum(vali) == 0:\n",
    "            continue\n",
    "        dtwpi = dt['wdist'][vali]\n",
    "        dtepi = dt['RSS'][vali]\n",
    "        wdist_stats[i, 0] = np.median(dtwpi)\n",
    "        wdist_stats[i, 1] = np.median(np.abs(dtwpi - np.median(dtwpi)))\n",
    "        wdist_stats[i, 2] = np.mean(dtwpi)\n",
    "        wdist_stats[i, 3] = np.std(dtwpi)\n",
    "        wdist_stats[i, 4] = np.percentile(dtwpi, norm.cdf(-1) * 1e2)\n",
    "        wdist_stats[i, 5] = np.percentile(dtwpi, norm.cdf(1) * 1e2)\n",
    "        edist_stats[i, 0] = np.median(dtepi)\n",
    "        edist_stats[i, 1] = np.median(np.abs(dtepi - np.median(dtepi)))\n",
    "        edist_stats[i, 2] = np.mean(dtepi)\n",
    "        edist_stats[i, 3] = np.std(dtepi)\n",
    "        edist_stats[i, 4] = np.percentile(dtepi, norm.cdf(-1) * 1e2)\n",
    "        edist_stats[i, 5] = np.percentile(dtepi, norm.cdf(1) * 1e2)\n",
    "\n",
    "    L = len(dt)\n",
    "    data = dt['wdist']\n",
    "    fig = plt.figure(figsize=(8, 6))\n",
    "    ax1 = fig.add_axes((.125, .12, .6, .77))\n",
    "    # boxdict = ax1.boxplot(np.array([dt['wdist'][dt['NPE'] == i+1] for i in range(l)], dtype=np.object), sym='', patch_artist=True)\n",
    "    ax1.plot(np.arange(1, l + 1), wdist_stats[:, 0], color='b')\n",
    "    \n",
    "    ax1.errorbar(np.arange(1, l + 1), \n",
    "                wdist_stats[:, 0], \n",
    "                yerr=[wdist_stats[:, 0] - wdist_stats[:, 4], wdist_stats[:, 5] - wdist_stats[:, 0]], \n",
    "                fmt='o', elinewidth=1, capsize=3, color='b', label=r'$D_\\mathrm{w}$')\n",
    "    \n",
    "    ax1.fill_between(np.arange(1, l + 1), \n",
    "                    wdist_stats[:, 4], \n",
    "                    wdist_stats[:, 5], fc='k', alpha=0.1, color=None)\n",
    "\n",
    "    ax1.set_xticks(np.arange(1, 16, 2))\n",
    "    ax1.set_xticklabels(np.arange(1, 16, 2).astype(str))\n",
    "    ax1.set_xlim(0, l + 1)\n",
    "    # ax1.set_ylim(0, max([boxdict['whiskers'][2 * i + 1].get_xydata()[1, 1] for i in range(l)]) * 1.05)\n",
    "    # ax1.set_ylim(0, wdist_stats[:, 5].max() * 1.05)\n",
    "    ax1.set_xlabel(r'$N_{\\mathrm{PE}}$')\n",
    "    ax1.set_ylabel(r'$\\mathrm{Wasserstein\\ Distance}/\\si{ns}$')\n",
    "    ax1.legend()\n",
    "\n",
    "    ax2 = fig.add_axes((.725, .12, .175, .77))\n",
    "    bins = np.linspace(0, data[data < np.percentile(data, 90)].std() * 3 + np.mean(data), 41)\n",
    "    ax2.hist(data, bins=bins, density=1, orientation='horizontal')\n",
    "\n",
    "    ax2.set_xlabel(r'$\\mathrm{arb.\\ unit}$')\n",
    "    ax2.set_xlim(0, ax2.get_xlim()[1] * 1.05)\n",
    "    ax2.set_xticks([])\n",
    "    ax2.set_yticks([])\n",
    "    ax2.set_ylim(0, bins[-1])\n",
    "    ax1.set_ylim(0, bins[-1])\n",
    "    fig.savefig('Note/figures/' + m + 'chargestats.pgf')\n",
    "    fig.savefig('Note/figures/' + m + 'chargestats.pdf')\n",
    "    plt.close(fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 250/250 [00:02<00:00, 106.83it/s]\n"
     ]
    }
   ],
   "source": [
    "t = np.arange(0, 1000, 0.1) / gmu\n",
    "\n",
    "pdf = np.zeros_like(t)\n",
    "tlist = np.arange(-50, 200)\n",
    "\n",
    "for mu in tqdm(Mu * wff.convolve_exp_norm(tlist, Tau, Sigma)):\n",
    "    for i in range(1, 15):\n",
    "        # pdf += mu * poisson.pmf(i, mu) * norm.pdf(t, loc=1, scale=gsigma / gmu / np.sqrt(i))\n",
    "        pdf += mu * poisson.pmf(i, mu) * gamma.pdf(t, a=(gmu/gsigma)**2, loc=0, scale=(gsigma/gmu)**2/np.sqrt(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = ['lucyddm', 'xiaopeip', 'takara', 'fsmp']\n",
    "colors = {'truth':'k', 'lucyddm':'y', 'xiaopeip':'c', 'takara':'C0', 'fsmp':'r'}\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "fig.tight_layout()\n",
    "ax = fig.add_axes((.1, .12, .85, .80))\n",
    "t = np.arange(0, 1000, 0.1) / gmu\n",
    "# ax.plot(t, norm.pdf(t, loc=1, scale=gsigma / gmu) / (1 - norm.cdf(0, loc=1, scale=gsigma / gmu)), color=colors['truth'], alpha=0.2)\n",
    "ax.plot(t, pdf / pdf.sum() / np.diff(t)[0], label='$\\mathrm{ChargePDF}$', color=colors['truth'])\n",
    "# th = 160 * 5 * 1e-4\n",
    "th = 10 / gmu\n",
    "labels = {'truth':'\\mathrm{Truth}', 'lucyddm':'\\mathrm{LucyDDM}', 'xiaopeip':'\\mathrm{Fitting}', 'takara':'\\mathrm{CNN}', 'fsmp':'\\mathrm{FSMP}', 'fsmpwave':'\\mathrm{FSMP}'}\n",
    "for m in methods:\n",
    "    ch = h5py.File('result/' + m + '/char/' + file + '.h5', 'r', libver='latest', swmr=True)\n",
    "    cha = ch['photoelectron']['Charge'] / gmu\n",
    "    ax.hist(cha[cha > th], bins=np.linspace(th, 400 / gmu, 101), label='$'+labels[m]+'$', histtype='step', density=True, color=colors[m], linewidth=2.)\n",
    "ax.set_xlim(10 / gmu, 310 / gmu)\n",
    "ax.set_yticks([])\n",
    "# ax.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "# ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax.legend()\n",
    "# ax.set_xlabel('$\\mathrm{Charge}/\\si{mV\\cdot ns}$')\n",
    "ax.set_xlabel('$\\mathrm{Charge}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Normalized\\ Count}$')\n",
    "plt.savefig('Note/figures/recchargehist.png')\n",
    "plt.savefig('Note/figures/recchargehist.pdf')\n",
    "plt.savefig('Note/figures/recchargehist.pgf')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 500/500 [00:04<00:00, 108.78it/s]\n"
     ]
    }
   ],
   "source": [
    "t = np.arange(0, 1000, 0.1) / gmu\n",
    "\n",
    "pdf = np.zeros_like(t)\n",
    "b = 0.5\n",
    "tlist = np.arange(-50, 200, b)\n",
    "\n",
    "for mu in tqdm(25 * wff.convolve_exp_norm(tlist, Tau, Sigma) * b):\n",
    "    for i in range(1, 15):\n",
    "        # pdf += mu * poisson.pmf(i, mu) * norm.pdf(t, loc=1, scale=gsigma / gmu / np.sqrt(i))\n",
    "        pdf += mu * poisson.pmf(i, mu) * gamma.pdf(t, a=(gmu/gsigma)**2, loc=0, scale=(gsigma/gmu)**2/np.sqrt(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = ['lucyddm', 'xiaopeip', 'takara', 'fsmp']\n",
    "colors = {'truth':'k', 'lucyddm':'y', 'xiaopeip':'c', 'takara':'C0', 'fsmp':'r'}\n",
    "fig = plt.figure(figsize=(10, 6))\n",
    "fig.tight_layout()\n",
    "ax = fig.add_axes((.1, .12, .85, .80))\n",
    "# ax.plot(t, norm.pdf(t, loc=1, scale=gsigma / gmu) / (1 - norm.cdf(0, loc=1, scale=gsigma / gmu)), color=colors['truth'], alpha=0.2)\n",
    "ax.plot(t, pdf / pdf.sum() / np.diff(t)[0], label='$\\mathrm{ChargePDF}$', color=colors['truth'])\n",
    "# th = 160 * 5 * 1e-4\n",
    "th = 10 / gmu\n",
    "labels = {'truth':'\\mathrm{Truth}', 'lucyddm':'\\mathrm{LucyDDM}', 'xiaopeip':'\\mathrm{Fitting}', 'takara':'\\mathrm{CNN}', 'fsmp':'\\mathrm{FSMP}', 'fsmpwave':'\\mathrm{FSMP}'}\n",
    "for m in methods:\n",
    "    ch = h5py.File('result/' + m + '/char/15.0-20-5' + '.h5', 'r', libver='latest', swmr=True)\n",
    "    cha = ch['photoelectron']['Charge'] / gmu\n",
    "    ax.hist(cha[cha > th], bins=np.linspace(th, 400 / gmu, 101), label='$'+labels[m]+'$', histtype='step', density=True, color=colors[m], linewidth=2.)\n",
    "ax.set_xlim(10 / gmu, 310 / gmu)\n",
    "ax.set_yticks([])\n",
    "# ax.yaxis.get_major_formatter().set_powerlimits((0, 1))\n",
    "# ax.yaxis.set_major_formatter(mtick.FormatStrFormatter('%.1f'))\n",
    "ax.legend()\n",
    "# ax.set_xlabel('$\\mathrm{Charge}/\\si{mV\\cdot ns}$')\n",
    "ax.set_xlabel('$\\mathrm{Charge}$')\n",
    "ax.set_ylabel(r'$\\mathrm{Normalized\\ Count}$')\n",
    "plt.savefig('Note/figures/recchargehist25.png')\n",
    "# plt.savefig('Note/figures/recchargehist.pdf')\n",
    "# plt.savefig('Note/figures/recchargehist.pgf')\n",
    "plt.show()"
   ]
  }
 ],
 "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.9.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}