{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Simulate 8 TF detectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "## ACME autophase routines, keep collapsed\n",
    "def autops(data, fn, p0=0.0, p1=0.0):\n",
    "    \"\"\"\n",
    "    Automated phase correction from NMRglue by https://github.com/jjhelmus\n",
    "    These functions provide support for automatic phasing of NMR data. \n",
    "\n",
    "\n",
    "    Automatic linear phase correction\n",
    "\n",
    "    Parameters\n",
    "\n",
    "        data : ndarray\n",
    "\n",
    "             Array of NMR data.\n",
    "\n",
    "        fn : str or function\n",
    "\n",
    "             Algorithm to use for phase scoring. Built in functions can be\n",
    "             specified by one of the following strings: \"acme\", \"peak_minima\"\n",
    "\n",
    "        p0 : float\n",
    "\n",
    "            Initial zero order phase in degrees.\n",
    "\n",
    "        p1 : float\n",
    "\n",
    "            Initial first order phase in degrees.\n",
    "\n",
    "    Returns\n",
    "\n",
    "        ndata : ndarray\n",
    "\n",
    "            Phased NMR data.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    import numpy as np\n",
    "    import scipy.optimize\n",
    "    from io import StringIO # Python3 use: from io import StringIO\n",
    "    from contextlib import redirect_stdout\n",
    "\n",
    "    \n",
    "    if not callable(fn):\n",
    "        fn = {\n",
    "            'peak_minima': _ps_peak_minima_score,\n",
    "            'acme': _ps_acme_score,\n",
    "        }[fn]\n",
    "    \n",
    "    opt = [p0, p1]\n",
    "    with StringIO() as buf, redirect_stdout(buf):   \n",
    "        opt = scipy.optimize.fmin(fn, x0=opt, args=(data, ))\n",
    "        mystdout = buf.getvalue()\n",
    "    return ps(data, p0=opt[0], p1=opt[1]), opt[0], opt[1], mystdout\n",
    "\n",
    "\n",
    "def _ps_acme_score(ph, data):\n",
    "    \"\"\"\n",
    "    Phase correction using ACME algorithm by Chen Li et al.\n",
    "    Journal of Magnetic Resonance 158 (2002) 164-168\n",
    "\n",
    "    Parameters\n",
    "    * pd : tuple, current p0 and p1 values\n",
    "    * data : ndarray, array of NMR data.\n",
    "\n",
    "    Returns\n",
    "    * score : float, value of the objective function (phase score)\n",
    "\n",
    "    \"\"\"\n",
    "    import numpy as np\n",
    "\n",
    "    stepsize = 1\n",
    "\n",
    "    phc0, phc1 = ph\n",
    "\n",
    "    s0 = ps(data, p0=phc0, p1=phc1)\n",
    "    data = np.real(s0)\n",
    "\n",
    "    # Calculation of first derivatives\n",
    "    ds1 = np.abs((data[1:]-data[:-1]) / (stepsize*2))\n",
    "    p1 = ds1 / np.sum(ds1)\n",
    "\n",
    "    # Calculation of entropy\n",
    "    p1[p1 == 0] = 1\n",
    "\n",
    "    h1 = -p1 * np.log(p1)\n",
    "    h1s = np.sum(h1)\n",
    "\n",
    "    # Calculation of penalty\n",
    "    pfun = 0.0\n",
    "    as_ = data - np.abs(data)\n",
    "    sumas = np.sum(as_)\n",
    "\n",
    "    if sumas < 0:\n",
    "        pfun = pfun + np.sum((as_/2) ** 2)\n",
    "\n",
    "    p = 1000 * pfun\n",
    "\n",
    "    return h1s + p\n",
    "\n",
    "\n",
    "def _ps_peak_minima_score(ph, data):\n",
    "    \"\"\"\n",
    "    Phase correction using simple minima-minimisation around highest peak\n",
    "    This is a naive approach but is quick and often achieves reasonable\n",
    "    results.  The optimisation is performed by finding the highest peak in the\n",
    "    spectra (e.g. TMSP) and then attempting to reduce minima surrounding it.\n",
    "    Parameters\n",
    "    * pd : tuple, current p0 and p1 values\n",
    "    * data : ndarray, array of NMR data.\n",
    "\n",
    "    Returns\n",
    "    * score : float, value of the objective function (phase score)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    phc0, phc1 = ph\n",
    "\n",
    "    s0 = ps(data, p0=phc0, p1=phc1)\n",
    "    data = np.real(s0)\n",
    "\n",
    "    i = np.argmax(data)\n",
    "    mina = np.min(data[i-100:i])\n",
    "    minb = np.min(data[i:i+100])\n",
    "\n",
    "    return np.abs(mina - minb)\n",
    "\n",
    "def ps(data, p0=0.0, p1=0.0, inv=False):\n",
    "    \"\"\"\n",
    "    Linear phase correction\n",
    "\n",
    "    Parameters\n",
    "\n",
    "        data : ndarray\n",
    "\n",
    "            Array of NMR data.\n",
    "\n",
    "        p0 : float\n",
    "\n",
    "            Zero order phase in degrees.\n",
    "\n",
    "        p1 : float\n",
    "\n",
    "            First order phase in degrees.\n",
    "\n",
    "        inv : bool, optional\n",
    "\n",
    "            True for inverse phase correction\n",
    "\n",
    "    Returns\n",
    "\n",
    "        ndata : ndarray\n",
    "\n",
    "            Phased NMR data.\n",
    "\n",
    "    \"\"\"\n",
    "    import numpy as np\n",
    "\n",
    "    p0 = p0 * np.pi / 180.  # convert to radians\n",
    "    p1 = p1 * np.pi / 180.\n",
    "    size = data.shape[-1]\n",
    "    apod = np.exp(1.0j * (p0 + (p1 * np.arange(size) / size))).astype(data.dtype)\n",
    "    if inv:\n",
    "        apod = 1 / apod\n",
    "    return apod * data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test multidetector asymmetry with a mix of 6 components with beatings average field ~0.35 T mean decay 1 mus-1 no noise\n",
    "# tested separately\n",
    "## effect of small (~1.6 ns) t0 discrepancies, seems ok\n",
    "## effect of theta deviations from n*pi/4, sum =0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "$\\sum_i \\cos\\theta_i =$ 1.33e-01\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 35852 characters]\n",
      "text/plain": [
       "<Figure size 1008x504 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# %matplotlib -l  lists all the possible options\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as P\n",
    "\n",
    "# generate 8 detector count rates, GPS like\n",
    "N0,dt,taumu,gamu = 1000, 0.5e-3, 2.197.,2*np.pi*135.5 # ref detector rate in counts/bin, bin width in mus, muon lifetime in mus, gyromagnetic ratio in mus-1/T \n",
    "ag,wg,sg = np.array([0.15,0.2,0.3]),gamu*np.array([.35,0.31,0.37]),np.array([0.8,3.2,1.6]) # Gaussian amplitudes sum to 0.6,omega as gamu*field in Tesla, decays in us-1\n",
    "al,wl,ll = np.array([0.18,0.12,0.05]), gamu*np.array([0.3,0.36,0.39]),np.array([0.3,1.5,2.5]) # Lorentzian amplitudes sum to 0.4\n",
    "t0 =  np.zeros(8) # 0.0015 + np.array([0.0006,-0.0008,-0.0004,-0.0006,-0.0007,0.0008,0.0004,0.0007])  # errors in t0 placement of each detector, average is 0.22 # np.zeros(8) # \n",
    "alpha = np.array([1.,0.95,0.93,0.97,1.03,1.05,1.04,0.99])# ratios of N0/Ni\n",
    "theta = np.linspace(0,7/8,8)*2*np.pi # exact angles! check np.cos(theta).sum()=-4-4e-16\n",
    "dtheta = np.array([0,-1,-1,1,-1,0,1,1,])*np.pi/64\n",
    "theta = theta + dtheta # sum = 0\n",
    "# detectors along columns:\n",
    "t0 = np.reshape(t0,(8,1))\n",
    "alpha = np.reshape(alpha,(8,1))\n",
    "theta = np.reshape(theta,(8,1))\n",
    "\n",
    "print(r'$\\sum_i \\cos\\theta_i =$ {:.2e}'.format(np.cos(theta).sum()))\n",
    "t = np.arange(0.,10.,dt)  # nominal-time vector along rows in musec\n",
    "\n",
    "tt, tt0 = np.meshgrid(t,t0)\n",
    "tt, aalpha = np.meshgrid(t,alpha)\n",
    "tt, ttheta = np.meshgrid(t,theta)\n",
    "\n",
    "N = sum([N0/aalpha*np.exp(-(tt-tt0)/taumu)*(1.+ag[k]*np.cos(wg[k]*(tt-tt0)+ttheta)*np.exp(-0.5*(sg[k]*(tt-tt0))**2)+al[k]*np.cos(wl[k]*(tt-tt0)+ttheta)*np.exp(-ll[k]*(tt-tt0))) for k in range(ag.shape[0])],0) # unpolarized counts for all detectors\n",
    "# 8 detectors counts\n",
    "Px = sum([ag[k]*np.cos(wg[k]*(t-t0[0]))*np.exp(-0.5*(sg[k]*(t-t0[0]))**2)+al[k]*np.cos(wl[k]*(t-t0[0]))*np.exp(-ll[k]*(t-t0[0])) for k in range(ag.shape[0])],0)\n",
    "# pure cos polarization\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "last = 20000\n",
    "fix,ax = P.subplots(2,4,figsize=(14,7),sharex=True,sharey=True)\n",
    "for det in range(theta.shape[0]):\n",
    "    ax[divmod(det,4)].plot(t[:last],N[det,:last],label='# {}'.format(det))\n",
    "    ax[divmod(det,4)].legend()\n",
    "    if det>3:\n",
    "        ax[divmod(det,4)].set_xlabel(r't [$\\mu$s]')\n",
    "    if det % 4 == 0:\n",
    "        ax[divmod(det,4)].set_ylabel('Polarization')\n",
    "P.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "7.336459499900938 -7.376554261091433\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 89572 characters]\n",
      "text/plain": [
       "<Figure size 720x360 with 3 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# fake analysis \n",
    "Ntot=8*N0*.np.exp(-tt/taumu)\n",
    "Pc = 2*sum(aalpha*N*np.cos(ttheta),0)/Ntot\n",
    "Ps = 2*sum(aalpha*N*np.sin(ttheta),0)/Ntot\n",
    "fig,ax = P.subplots(1,3,figsize=(10,5),sharey=True)\n",
    "\n",
    "\n",
    "fftP = np.fft.fft(Pc)+1.0j*np.fft.fft(Ps)\n",
    "fftPrephase,p0,p1,dum = autops(fftP,'acme')\n",
    "print(p0,p1) #'p0 = {:.1f} deg, p1 = {:.1f} deg',format(p0,p1)\n",
    "Prephase = np.fft.ifft(fftPrephase)\n",
    "ax[0].plot(t[:last],Pc[:last],label='cos')\n",
    "ax[0].plot(t[:last],Ps[:last],label='sin')\n",
    "ax[1].plot(t[:last],Px[:last],label='original')\n",
    "ax[2].plot(t[:last],Prephase.real[0:last],label='acme cos')\n",
    "ax[2].plot(t[:last],Prephase.imag[0:last],label='acme sin')\n",
    "ax[0].legend()\n",
    "ax[1].legend()\n",
    "ax[2].legend()\n",
    "ax[0].set_xlabel(r't [$\\mu$s]')\n",
    "ax[0].set_ylabel('Polarization')\n",
    "ax[1].set_xlabel(r't [$\\mu$s]')\n",
    "ax[2].set_xlabel(r't [$\\mu$s]')\n",
    "ax[0].set_ylim(-1.05,1.05)\n",
    "P.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
