{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from tqdm import tqdm\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 7\n",
    "batch_size = 1024\n",
    "L = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def jackknife(data):\n",
    "    n = len(data)\n",
    "    jackknife_samples = np.empty((n, n-1))\n",
    "    \n",
    "    for i in range(n):\n",
    "        jackknife_samples[i] = np.delete(data, i)\n",
    "    \n",
    "    return jackknife_samples\n",
    "\n",
    "def jackknife_stats(data):\n",
    "    jackknife_samples = jackknife(data)\n",
    "    jackknife_means = np.mean(jackknife_samples, axis=1)\n",
    "    mean = np.mean(jackknife_means)\n",
    "    error = np.sqrt((len(data) - 1) * np.mean((jackknife_means - mean) ** 2))\n",
    "    \n",
    "    return mean, error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_autocorrelation(samples, max_lag=None):\n",
    "  \n",
    "    x = np.asarray(samples, dtype=float)\n",
    "    N = len(x)\n",
    "    if max_lag is None or max_lag >= N:\n",
    "        max_lag = N - 1\n",
    "\n",
    "    mean_x = np.mean(x)\n",
    "    var_x = np.var(x)\n",
    "    if np.isclose(var_x, 0.0):\n",
    "        return np.ones(max_lag + 1)\n",
    "\n",
    "    autocorr = np.zeros(max_lag + 1)\n",
    "    for lag in range(max_lag + 1):\n",
    "        x_front = x[:N - lag] - mean_x\n",
    "        x_back  = x[lag:] - mean_x\n",
    "        cov_k = np.sum(x_front * x_back) / N\n",
    "        autocorr[lag] = cov_k / var_x\n",
    "    \n",
    "    return autocorr\n",
    "\n",
    "def compute_ess(samples, max_lag=None):\n",
    "\n",
    "    x = np.asarray(samples, dtype=float)\n",
    "    N = len(x)\n",
    "    autocorr = compute_autocorrelation(x, max_lag=max_lag)\n",
    "    \n",
    "    tau_int = 1.0\n",
    "    for k in range(1, len(autocorr)):\n",
    "        if autocorr[k] < 0:\n",
    "            break\n",
    "        tau_int += 2.0 * autocorr[k]\n",
    "\n",
    "    ess = N / tau_int\n",
    "    return ess, tau_int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gauge_cooling(lattice, n_steps=1):\n",
    "    for _ in range(n_steps):\n",
    "        lattice = np.mod(lattice+np.pi, 2*np.pi) - np.pi\n",
    "    return lattice\n",
    "\n",
    "def metropolis_hastings_update(x, y, p, random_a):\n",
    "    accept_mask = random_a < p\n",
    "    updated_x = copy.deepcopy(x) \n",
    "    updated_x[accept_mask] = y[accept_mask]\n",
    "    acceptance_rate = np.sum(accept_mask) / accept_mask.size  \n",
    "    return updated_x, acceptance_rate\n",
    "\n",
    "def get_action(phi):\n",
    "\n",
    "    ph_w = np.cos(phi[:,:,:,0] - phi[:,:,:,1] + np.roll(phi[:,:,:,1], -1, 1) - np.roll(phi[:,:,:,0], -1, 2))\n",
    "    return -beta*np.sum(ph_w, axis=(1,2))\n",
    "\n",
    "def get_drift(phi):\n",
    "    phi_d = np.zeros((batch_size,L,L,2))\n",
    "    drift1 = beta * np.sin(phi[:,:,:,0] - phi[:,:,:,1] + np.roll(phi[:,:,:,1], -1, axis=1) - np.roll(phi[:,:,:,0], -1, axis=2))\n",
    "    phi_d[:,:,:,0] = -drift1 + np.roll(drift1, 1, axis=2)\n",
    "    phi_d[:,:,:,1] = drift1 - np.roll(drift1, 1, axis=1)\n",
    "    return phi_d\n",
    "\n",
    "\n",
    "\n",
    "def M_H(x,y,noise,h):\n",
    "    drift = beta*h*(get_drift(x) + get_drift(y))\n",
    "\n",
    "    p_1 = -(1/h)*0.25*np.sum((np.sqrt(2*h)*noise + drift)**2, axis=(1, 2, 3))\n",
    "    p_2 = 0.5*np.sum(noise ** 2, axis=(1, 2, 3))\n",
    "    delta_q = np.exp(p_1+p_2)\n",
    "    S_x = get_action(x)\n",
    "    S_y = get_action(y)\n",
    "    \n",
    "    delta_pi = np.exp(S_x-S_y)\n",
    "    accept_prob = delta_pi * delta_q\n",
    "    #print(accept_prob)\n",
    "    random_i =  np.random.rand(batch_size)\n",
    "    return metropolis_hastings_update(x,y,accept_prob,random_i)\n",
    "\n",
    "def wilson_loop(lattice, x, y):\n",
    "    U1 = np.exp(1j * lattice[x, y, 0])\n",
    "    U2 = np.exp(1j * lattice[x, (y + 1) % L, 0])\n",
    "    U3 = np.exp(1j * lattice[(x + 1) % L, y, 1])\n",
    "    U4 = np.exp(1j * lattice[x, y, 1])\n",
    "    return np.real(U1 * U3 * np.conj(U2) * np.conj(U4))\n",
    "\n",
    "def calculate_angle(array1,xp,yp):\n",
    "    var = L\n",
    "    right = array1[xp % var, yp % var][0]\n",
    "    up = array1[(xp + 1) % var, yp % var][1]\n",
    "    down = array1[xp % var, yp % var][1]\n",
    "    left = array1[xp % var, (yp + 1) % var][0]\n",
    "#    angel = right + up - left - down\n",
    "    return right + up - left - down\n",
    "\n",
    "\n",
    "def calculate_topological_charge(array):\n",
    "    top_charge = 0\n",
    "    for i in range(L):\n",
    "        for j in range(L):\n",
    "            angle = calculate_angle(array, i, j)\n",
    "            top_charge += project_angle(angle)\n",
    "\n",
    "    return (top_charge / (2 * np.pi))\n",
    "\n",
    "\n",
    "\n",
    "def project_angle(angle):\n",
    "    return np.remainder(angle+np.pi, (2 * np.pi)) - np.pi\n",
    "\n",
    "def topological_suspectiblity(array):\n",
    "    top_charge = 0\n",
    "    for i in range(L):\n",
    "        for j in range(L):\n",
    "            angle = calculate_angle(array, i, j)\n",
    "            top_charge += project_angle(angle)\n",
    "\n",
    "    return (top_charge / (2 * np.pi))**2\n",
    "\n",
    "def plaq(lat):\n",
    "    Wilson_loop_values=[]\n",
    "    lattice_size=lat.shape[2]\n",
    "    for conf in range(lat.shape[0]):\n",
    "        wilson_sum = 0\n",
    "        lattice = lat[conf]\n",
    "        for x in range(lattice_size):\n",
    "            for y in range(lattice_size):\n",
    "                wilson_sum += wilson_loop(lattice, x, y)\n",
    "        Wilson_loop_values.append(wilson_sum / (lattice_size ** 2))\n",
    "    return Wilson_loop_values\n",
    "\n",
    "def wilson_loopij(array, x_len, t_len):\n",
    "    Wilson_loop_values=[]\n",
    "    for conf in range(array.shape[0]):\n",
    "        init = np.zeros((L,L))\n",
    "        for x in range(x_len):\n",
    "            init += np.roll(array[conf,:,:,0],-x,axis=0)\n",
    "        for t in range(t_len):\n",
    "            init += np.roll(np.roll(array[conf,:,:,1],-x_len,axis=0),-t,axis=1)\n",
    "        for x in range(x_len):\n",
    "            init -= np.roll(np.roll(array[conf,:,:,0],(-x_len,-t_len),axis=(0,1)),x+1,axis=0)   \n",
    "        for t in range(t_len):\n",
    "            init -= np.roll(np.roll(array[conf,:,:,1],-t_len,axis=1),t+1,axis=1)\n",
    "        wil_l = np.mean(np.cos(init),axis=(0,1))\n",
    "        Wilson_loop_values.append(wil_l)\n",
    "    return Wilson_loop_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [10:55<00:00,  6.55s/it]\n"
     ]
    }
   ],
   "source": [
    "\n",
    "size =100\n",
    "\n",
    "    \n",
    "init_x = np.random.randn(batch_size,  L, L,2)  * np.pi\n",
    "MH_step = 1e-5\n",
    "rate = 0\n",
    "x = init_x\n",
    "rate_traj = []\n",
    "Q_list = []\n",
    "for i in tqdm(range(100)):\n",
    "    MH_step = 1e-9\n",
    "    rate = 0\n",
    "    step_size =  5*0.81**i\n",
    "    for _ in range(size):\n",
    "      score_t = get_drift(x)\n",
    "      mean_x = x + score_t * step_size\n",
    "      randn_noise = np.random.randn(batch_size,  L, L,2)\n",
    "      x_temp = mean_x + np.sqrt(2*step_size) * randn_noise\n",
    "      if i>90:\n",
    "        size = 1000\n",
    "        x, acc_rate = M_H(x,x_temp,randn_noise,step_size)\n",
    "        MH_step = MH_step+1\n",
    "        rate += acc_rate\n",
    "        Q_list.append(calculate_topological_charge(x[0]))\n",
    "      else:\n",
    "        x = x_temp\n",
    "      \n",
    "      x = gauge_cooling(x)\n",
    "    rate_traj.append(rate/MH_step)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(8100,)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.shape(Q_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " tau_int = 5001.00\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 32376 characters]",
      "text/plain": [
       "<Figure size 800x500 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "max_lag = 2500\n",
    "autocorr = compute_autocorrelation(Q_list, max_lag=max_lag)\n",
    "\n",
    "\n",
    "_, tau_int = compute_ess(Q_list, max_lag=max_lag)\n",
    "print(f\" tau_int = {tau_int:.2f}\")\n",
    "\n",
    "plt.figure(figsize=(8, 5))\n",
    "lags = np.arange(len(autocorr))\n",
    "plt.plot(lags, autocorr, marker='o', linestyle='-', label='Autocorrelation')\n",
    "plt.axhline(y=0, color='k', linestyle='--', linewidth=1)\n",
    "\n",
    "plt.title(f\"Autocorrelation vs. Lag\\n(tau_int ≈ {tau_int:.2f})\")\n",
    "plt.xlabel(\"Lag (MC)\")\n",
    "plt.ylabel(r\"Autocorrelation $\\rho_X(k)$\")\n",
    "plt.legend(loc='best')\n",
    "plt.grid(True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 28944 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\n",
    "\n",
    "time = range(len(rate_traj))  \n",
    "\n",
    "plt.figure(figsize=(8, 6))\n",
    "plt.plot(time, rate_traj, marker='x', linestyle='-', label='Rate Trajectory')\n",
    "plt.xlabel('Time')\n",
    "plt.ylabel('Rate')\n",
    "plt.title('Rate Trajectory over Time')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9816455078115184"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "rate/MH_step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.9231223575245571\n",
      "0.7221334939062264\n",
      "0.46638942047376364\n",
      "0.23935179962015088\n",
      "0.09353340279311201\n",
      "0.025372324000845776\n"
     ]
    }
   ],
   "source": [
    "for i in range(1,7):\n",
    "    print(np.mean(wilson_loopij(x[:1024],i,i)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0070794659534154005"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.std(wilson_loopij(x[:1024],1,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57600 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import scienceplots\n",
    "from matplotlib.ticker import MultipleLocator\n",
    "import matplotlib.pyplot as plt\n",
    "plt.style.use(['science', 'notebook','grid'])\n",
    "plt.rcParams['figure.dpi'] = 100\n",
    "plt.rcParams['savefig.dpi'] = 300\n",
    "plt.rcParams['font.size'] = 16\n",
    "plt.rcParams['xtick.labelsize'] = 14  \n",
    "plt.rcParams['ytick.labelsize'] = 14 \n",
    "\n",
    "plt.figure(figsize=(8, 6))\n",
    "\n",
    "\n",
    "\n",
    "Xrange=(0.8,1)\n",
    "plt.hist(wilson_loopij(x[:1024],1,1), bins=40,range=Xrange, histtype='step',linewidth=2,edgecolor='blue',color='blue',alpha=1,density=True,label=r'HMC $L = 16$')\n",
    "#plt.title('Wilson loop')\n",
    "ax = plt.gca()  \n",
    "ax.yaxis.set_major_locator(MultipleLocator(5))  \n",
    "ax.yaxis.set_minor_locator(MultipleLocator(2.5)) \n",
    "\n",
    "\n",
    "plt.legend()\n",
    "#plt.ylim()\n",
    "plt.xlabel(r'$\\langle W_{1 \\times 1} \\rangle$', fontsize=16)\n",
    "plt.ylabel(r'Frequency', fontsize=16)\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "top_charge_history=[]\n",
    "for conf in range(x.shape[0]):\n",
    "    top_charge_history.append(calculate_topological_charge(x[conf]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30556 characters]",
      "text/plain": [
       "<Figure size 800x600 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import scienceplots\n",
    "\n",
    "# 全局参数固定字体大小、dpi 等\n",
    "plt.style.use(['science', 'notebook','grid'])\n",
    "plt.rcParams['figure.dpi'] = 100\n",
    "plt.rcParams['savefig.dpi'] = 300\n",
    "plt.rcParams['font.size'] = 16\n",
    "plt.rcParams['xtick.labelsize'] = 14  \n",
    "plt.rcParams['ytick.labelsize'] = 14 \n",
    "\n",
    "fig_width = 8 \n",
    "fig_height = 6  \n",
    "\n",
    "fig, ax = plt.subplots(figsize=(fig_width, fig_height), dpi=100)\n",
    "\n",
    "bin_range = range(-5, 6)\n",
    "\n",
    "hist, bins = np.histogram(top_charge_history, bins=bin_range, density=True)\n",
    "\n",
    "\n",
    "ax.bar(bins[:-1] + 0.2, hist, width=0.5, edgecolor='black', color='orange', alpha=1, label=r'Langevin $\\beta=7$')\n",
    "\n",
    "\n",
    "ax.set_xlabel(r'Topological Charge $Q$', fontsize=16)\n",
    "ax.set_ylabel(r'Frequency', fontsize=16)\n",
    "\n",
    "\n",
    "plt.legend()\n",
    "\n",
    "\n",
    "ax = plt.gca()\n",
    "ax.yaxis.set_major_locator(MultipleLocator(0.2))  \n",
    "ax.yaxis.set_minor_locator(MultipleLocator(0.05))  \n",
    "\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "suspect1 = []\n",
    "for conf in range(x.shape[0]):\n",
    "    suspect1.append(topological_suspectiblity(x[conf]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.013526916503906252 0.0005832548539957425\n"
     ]
    }
   ],
   "source": [
    "mean, err = jackknife_stats(suspect1)\n",
    "print(mean/L**2,err/L**2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
