{
"cells": [
{
"cell_type": "markdown",
"id": "992fe18e",
"metadata": {},
"source": [
"# Real-time dynamics\n",
"\n",
"In this tutorial, we will discuss how to use Quantax to simulate real-time dynamics $\\ket{\\dot \\psi} = -i H \\ket{\\psi}$ in a $4 \\times 4$ transverse-field Ising model.\n",
"\n",
"We will start from a product state $\\ket{\\psi_0} = (\\ket{\\uparrow} + \\ket{\\downarrow})^{\\otimes N}$, and simulate quench given Ising Hamiltonian."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "af0046c4",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import jax\n",
"import jax.numpy as jnp\n",
"import equinox as eqx\n",
"import quantax as qtx\n",
"from quantax.symmetry import TransND, C4v, SpinInverse\n",
"from IPython.display import clear_output\n",
"%config InlineBackend.figure_format = 'svg'\n",
"\n",
"qtx.set_default_dtype(jnp.complex128)\n",
"\n",
"lattice = qtx.sites.Square(4)\n",
"N = lattice.Nsites\n",
"\n",
"H = qtx.operator.Ising(h=6)"
]
},
{
"cell_type": "markdown",
"id": "c2a48118",
"metadata": {},
"source": [
"## Exact evolution\n",
"\n",
"In small systems, one can utilize QuSpin to obtain exact evolution trajectory.\n",
"\n",
"We start with the initial product state $\\ket{\\psi_0}$. To reduce the cost of evolution, it's usually important to restrict the state into a specific symmetry sector."
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "3675de5e",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/aochen/quantax_env/lib/python3.12/site-packages/quantax/symmetry/symmetry.py:288: GeneralBasisWarning: using non-commuting symmetries can lead to unwanted behaviour of general basis, make sure that quantum numbers are invariant under non-commuting symmetries!\n",
" basis = spin_basis_general(\n"
]
}
],
"source": [
"psi0 = np.ones((1 << N,), dtype=jnp.complex128)\n",
"psi0 = qtx.state.DenseState(psi0)\n",
"\n",
"symm = TransND() @ C4v() @ SpinInverse()\n",
"psi0 = psi0.todense(symm).normalize()"
]
},
{
"cell_type": "markdown",
"id": "018da761",
"metadata": {},
"source": [
"Then we perform the evolution using a QuSpin operator. It will take some time."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "59e7b61b",
"metadata": {},
"outputs": [],
"source": [
"op = H.get_quspin_op(symm)\n",
"t0 = 0.0\n",
"times = np.linspace(0, 1, 1001)\n",
"\n",
"psit = op.evolve(psi0.psi, t0, times)"
]
},
{
"cell_type": "markdown",
"id": "c1f9d64e",
"metadata": {},
"source": [
"Now we can measure the evolution of $\\braket{\\sigma_i^x}$."
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "fb3ad6ef",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"from quantax.operator import sigma_x\n",
"\n",
"opx = sum(sigma_x(i) for i in range(N)) / N\n",
"psit = [qtx.state.DenseState(psi, symm=symm) for psi in psit.T]\n",
"sx_exact = [psi @ opx @ psi for psi in psit]\n",
"sx_exact = np.asarray(sx_exact).real\n",
"plt.plot(times, sx_exact)\n",
"plt.xlabel('Time')\n",
"plt.ylabel(r'$\\langle \\sigma_i^x \\rangle$')\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"id": "9e60d951",
"metadata": {},
"source": [
"## NQS preparation\n",
"\n",
"The first step of neural-network unitary dynamics is also preparing an initial state. It's hard to set the wavefunction amplitudes directly in an NQS. Here, we prepare the initial state by searching for the ground state of $H_0 = -\\sum_i \\sigma_i^x$.\n",
"\n",
"We will use {py:class}`~quantax.model.ResConv` as the variational wavefunction, which has been used in the tutorials of {doc}`square ` and {doc}`triangular ` J1-J2 models."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "e90d1e56",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"E0106 09:48:23.025522 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n",
"E0106 09:48:23.254585 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n",
"E0106 09:48:23.577108 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n",
"E0106 09:48:23.912475 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n",
"E0106 09:48:24.170000 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n",
"E0106 09:48:24.491987 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n",
"E0106 09:48:24.827526 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n",
"E0106 09:48:25.164962 1531416 cuda_timer.cc:87] Delay kernel timed out: measured time has sub-optimal accuracy. There may be a missing warmup execution, please investigate in Nsight Systems.\n"
]
},
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"model = qtx.model.ResConv(\n",
" nblocks=2, channels=8, kernel_size=3, dtype=jnp.float32, out_dtype=jnp.complex128\n",
")\n",
"# Translation symmetry is encoded in ResConv, so it's not needed here\n",
"symm1 = C4v() @ SpinInverse()\n",
"\n",
"state = qtx.state.Variational(model, symm=symm1)\n",
"sampler0 = qtx.sampler.LocalFlip(state, 1024)\n",
"optimizer = qtx.optimizer.SR(state, -opx)\n",
"\n",
"sx = qtx.utils.DataTracer()\n",
"\n",
"for i in range(100):\n",
" samples = sampler0.sweep()\n",
" step = optimizer.get_step(samples)\n",
" state.update(step * 0.1)\n",
" sx.append(-optimizer.energy)\n",
"\n",
"sx.plot()\n",
"plt.show()\n",
"\n",
"state.save(\"/tmp/dynamics.eqx\")"
]
},
{
"cell_type": "markdown",
"id": "c5d73f6e",
"metadata": {},
"source": [
"In practice, the choice of the initial state $\\ket{\\psi_0}$ depends on the physical background. For instance, $\\ket{\\psi_0}$ can be an excitation of a ground state if one needs to study low-energy spectrum."
]
},
{
"cell_type": "markdown",
"id": "5a6319de",
"metadata": {},
"source": [
"## Non-stochastic TDVP\n",
"\n",
"The method simulating unitary dynamics with variational wavefunctions is called time-dependent variational principle (TDVP). We start with non-stochastic TDVP supported by {py:class}`~quantax.optimizer.ER`. As explained in {doc}`Build your network `, it's restricted to small systems and helps us to check the expressivity of neural networks by excluding stochastic errors. \n",
"\n",
"Assume the evolution trajectory is determined by\n",
"\n",
"$$\n",
"\\dot \\theta = f(\\theta),\n",
"$$\n",
"\n",
"where $\\theta$ is variational parameters. In ground-state search, one update $\\theta$ up to the first order, i.e.\n",
"\n",
"$$\n",
"\\theta(t_0 + \\Delta t) = \\theta(t_0) + f(\\theta(t_0)).\n",
"$$\n",
"\n",
"In unitary dynamics, one has to employ higher order updates to ensure accuracy in every step. We will utilize the Heun's method given by\n",
"\n",
"$$\n",
"\\begin{aligned}\n",
" k_1 &= f(\\theta(t_0)), \\\\\n",
" k_2 &= f(\\theta(t_0) + k_1 \\Delta t), \\\\\n",
" \\theta(t_0 + \\Delta t) &= \\theta(t_0) + \\frac{k_1 + k_2}{2} \\Delta t.\n",
"\\end{aligned}\n",
"$$"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "5d198b80",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"state = qtx.state.Variational(model, \"/tmp/dynamics.eqx\", symm=symm1)\n",
"optimizer = qtx.optimizer.ER(state, H, imag_time=False, symm=symm)\n",
"\n",
"dt = 0.001\n",
"\n",
"fidelity = qtx.utils.DataTracer()\n",
"sx_etdvp = qtx.utils.DataTracer()\n",
"\n",
"for i in range(1, 1001):\n",
" step1 = optimizer.get_step()\n",
" state.update(step1 * dt)\n",
" step2 = optimizer.get_step()\n",
" state.update((step2 - step1) * dt / 2)\n",
"\n",
" if i % 10 == 0:\n",
" clear_output(wait=True)\n",
" dense = state.todense(symm).normalize()\n",
" fidelity.append(abs(dense @ psit[i]) ** 2, time=i * dt)\n",
" sx_etdvp.append((dense @ opx @ dense).real, time=i * dt)\n",
"\n",
" fidelity.plot()\n",
" plt.xlabel('Time')\n",
" plt.ylabel('Fidelity')\n",
" plt.show()\n",
"\n",
" plt.plot(times, sx_exact, 'k--')\n",
" sx_etdvp.plot()\n",
" plt.xlabel('Time')\n",
" plt.ylabel(r'$\\langle \\sigma_i^x \\rangle$')\n",
" plt.show()"
]
},
{
"cell_type": "markdown",
"id": "871899f2",
"metadata": {},
"source": [
"The fidelity is high during the evolution, and $\\braket{\\sigma_i^x}$ matches well with the exact values. These results show that `ResConv` has enough expressivity for the dynamics.\n",
"\n",
"## Stochastic TDVP\n",
"\n",
"Next, we will introduce stochastic evolution supported by `TimeEvol`, which is equivalent to real-time {py:class}`~quantax.optimizer.SR`. We need to specify more samples and sampling steps to reduce stochastic errors."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "dc67b563",
"metadata": {},
"outputs": [
{
"data": {
"image/svg+xml": [
"\n",
"\n",
"\n"
],
"text/plain": [
""
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"state = qtx.state.Variational(model, \"/tmp/dynamics.eqx\", symm=symm1)\n",
"sampler = qtx.sampler.LocalFlip(\n",
" state, nsamples=8192, thermal_steps=50 * N, sweep_steps=5 * N\n",
")\n",
"optimizer = qtx.optimizer.TimeEvol(state, H)\n",
"\n",
"dt = 0.001\n",
"\n",
"sx_etdvp = qtx.utils.DataTracer()\n",
"\n",
"for i in range(1, 1001):\n",
" samples = sampler.sweep()\n",
" step1 = optimizer.get_step(samples)\n",
" state.update(step1 * dt)\n",
"\n",
" samples = sampler.sweep()\n",
" step2 = optimizer.get_step(samples)\n",
" state.update((step2 - step1) * dt / 2)\n",
"\n",
" if i % 10 == 0:\n",
" clear_output(wait=True)\n",
" sx_etdvp.append(opx.expectation(state, samples), time=i * dt)\n",
"\n",
" plt.plot(times, sx_exact, \"k--\")\n",
" sx_etdvp.plot()\n",
" plt.xlabel(\"Time\")\n",
" plt.ylabel(r\"$\\langle \\sigma_i^x \\rangle$\")\n",
" plt.show()\n",
"\n",
" if i % 50 == 0:\n",
" # Reset all Markov chains to improve sampling quality\n",
" sampler.reset()"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "quantax_env",
"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.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}