lucyddm.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import re\n",
    "import copy\n",
    "import time\n",
    "import math\n",
    "import argparse\n",
    "import pickle\n",
    "from functools import partial\n",
    "from multiprocessing import Pool, cpu_count\n",
    "\n",
    "import numpy as np\n",
    "np.seterr(divide='ignore')\n",
    "# np.seterr(all='raise')\n",
    "from celluloid import Camera\n",
    "from mpl_axes_aligner import align\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.animation import FuncAnimation, ArtistAnimation\n",
    "from scipy.signal import savgol_filter\n",
    "from tqdm import tqdm\n",
    "import h5py\n",
    "\n",
    "import wf_func as wff\n",
    "spe_pre = wff.read_model('spe.h5', 1)\n",
    "\n",
    "import matplotlib\n",
    "matplotlib.use('Agg')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opdt = np.dtype([('TriggerNo', np.uint32), ('ChannelID', np.uint32), ('HitPosInWindow', np.float64), ('Charge', np.float64)])\n",
    "with h5py.File('waveform/4.0-20-5.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'][:]\n",
    "    l = len(ipt['Readout/Waveform'])\n",
    "    print('{} waveforms will be computed'.format(l))\n",
    "    window = int(len(ipt['Readout/Waveform'][0]['Waveform']) / wff.nshannon)\n",
    "    assert window >= len(spe_pre[0]['spe']), 'Single PE too long which is {}'.format(len(spe_pre[0]['spe']))\n",
    "    Mu = ipt['Readout/Waveform'].attrs['mu']\n",
    "    Tau = ipt['Readout/Waveform'].attrs['tau']\n",
    "    Sigma = ipt['Readout/Waveform'].attrs['sigma']\n",
    "    gmu = ipt['SimTriggerInfo/PEList'].attrs['gmu']\n",
    "    gsigma = ipt['SimTriggerInfo/PEList'].attrs['gsigma']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\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[ent[i]['ChannelID']]['epulse']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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 = min(((ind.max() + spe_pre[cid]['mar_r']) // 20 + 1) * 20, xmin + 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def savegif(waveform, spe_pre):\n",
    "    waveform = savgol_filter(waveform, 11, 4)\n",
    "    n = 0\n",
    "    spe = np.append(np.zeros(len(spe_pre) - 2 * n - 1), np.abs(spe_pre))\n",
    "    waveform = np.clip(waveform, 1e-6, np.inf)\n",
    "    spe = np.clip(spe, 1e-6, np.inf)\n",
    "    waveform = waveform / np.sum(spe)\n",
    "    wave_deconv = waveform.copy()\n",
    "    spe_mirror = spe[::-1]\n",
    "    fig = plt.figure(figsize=(8, 6))\n",
    "    camera = Camera(fig)\n",
    "    ax = fig.add_subplot()\n",
    "    ax2 = ax.twinx()\n",
    "    ax2.set_ylim(top=1.0)\n",
    "    ax.set_xlim(xmin, xmax)\n",
    "    ax.set_xlabel('t/ns')\n",
    "    ax.set_ylabel('Voltage/mV')\n",
    "    ax2.set_ylabel('Charge')\n",
    "    for i in tqdm(range(100)):\n",
    "        relative_blur = waveform / np.convolve(wave_deconv, spe, mode='same')\n",
    "        new_wave_deconv = wave_deconv * np.convolve(relative_blur, spe_mirror, mode='same')\n",
    "        wave_deconv = new_wave_deconv\n",
    "        pet = np.arange(0, len(waveform) - n)\n",
    "        pwe = wave_deconv[n:] / wave_deconv[n:].sum() * waveform.sum()\n",
    "        \n",
    "        ax.plot(waveform * np.sum(spe), label='Waveform', color='b')\n",
    "        ax.hlines(5, 0, window, color='k', label='WaveThres')\n",
    "        ax2.vlines(truth['HitPosInWindow'], 0, truth['Charge'], color='g', label='Truth', linewidth=1.0, alpha=0.5)\n",
    "        ax2.vlines(pet, 0, pwe, color='r', label='Charge', linewidth=0.5)\n",
    "        ax2.hlines(0.1, 0, window, color='k', linestyle='--', label='ChaThres')\n",
    "        \n",
    "        align.yaxes(ax, 0, ax2, 0)\n",
    "        camera.snap()\n",
    "    animation = camera.animate(interval=100)\n",
    "    animation.save('lucyddm.gif')\n",
    "    plt.close(fig)\n",
    "    return"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "savegif(wave, spe_pre[i]['spe'])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.7rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}