{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dicke-title",
   "metadata": {},
   "source": [
    "# Dicke Synchronization with a Collective Spin Ensemble\n",
    "\n",
    "This notebook uses the symmetric Dicke manifold to simulate a driven, collectively damped spin ensemble. The Hilbert-space dimension is only `N + 1`, so it is a compact way to explore finite-size collective dynamics before moving to larger trajectory studies."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dicke-install",
   "metadata": {},
   "source": [
    "If you are running this notebook in Colab, install the package first:\n",
    "\n",
    "```python\n",
    "%pip install openquantumsim\n",
    "```\n",
    "\n",
    "The first solver call on a new machine may spend a few minutes resolving and precompiling the Julia backend."
   ]
  },
  {
   "cell_type": "code",
   "id": "dicke-imports",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "import openquantumsim as oqs"
   ]
  },
  {
   "cell_type": "code",
   "id": "dicke-coherent-state",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def spin_coherent_dicke(space: oqs.DickeSpace, theta: float, phi: float = 0.0) -> np.ndarray:\n",
    "    \"\"\"Spin-coherent state in the Dicke basis ordered by descending excitations.\"\"\"\n",
    "    n_spins = space.n_spins\n",
    "    ket = np.zeros(space.dim, dtype=np.complex128)\n",
    "    for excitations in range(n_spins + 1):\n",
    "        amplitude = (\n",
    "            math.sqrt(math.comb(n_spins, excitations))\n",
    "            * np.sin(theta / 2) ** excitations\n",
    "            * np.cos(theta / 2) ** (n_spins - excitations)\n",
    "            * np.exp(1j * phi * excitations)\n",
    "        )\n",
    "        ket[n_spins - excitations] = amplitude\n",
    "    return ket / np.linalg.norm(ket)"
   ]
  },
  {
   "cell_type": "code",
   "id": "dicke-model",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_spins = 24\n",
    "space = oqs.DickeSpace(n_spins, label=\"ensemble\")\n",
    "\n",
    "omega = 1.0\n",
    "gamma = 0.35\n",
    "times = np.linspace(0.0, 12.0, 241)\n",
    "\n",
    "psi0 = spin_coherent_dicke(space, theta=np.pi / 3, phi=0.0)\n",
    "rho0 = oqs.ket2dm(psi0)\n",
    "\n",
    "H = omega * oqs.collective_x(space)\n",
    "c_ops = [np.sqrt(gamma / n_spins) * oqs.collective_lowering(space)]\n",
    "\n",
    "jx = (2.0 / n_spins) * oqs.collective_x(space)\n",
    "jz = (2.0 / n_spins) * oqs.collective_z(space)\n",
    "excitation = (1.0 / n_spins) * oqs.collective_excitation(space)"
   ]
  },
  {
   "cell_type": "code",
   "id": "dicke-solve",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = oqs.mesolve(\n",
    "    H,\n",
    "    rho0,\n",
    "    times,\n",
    "    c_ops=c_ops,\n",
    "    e_ops=[jx, jz, excitation],\n",
    "    options=oqs.Options(rtol=1e-8, atol=1e-10),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "id": "dicke-plot",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(7.0, 4.0))\n",
    "ax.plot(times, result.expect[0].real, label=r\"$2\\langle J_x\\rangle/N$\")\n",
    "ax.plot(times, result.expect[1].real, label=r\"$2\\langle J_z\\rangle/N$\")\n",
    "ax.plot(times, result.expect[2].real, label=r\"$\\langle n_e\\rangle/N$\")\n",
    "ax.set_xlabel(\"time\")\n",
    "ax.set_ylabel(\"normalized expectation\")\n",
    "ax.legend()\n",
    "ax.grid(alpha=0.3)\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dicke-next",
   "metadata": {},
   "source": [
    "Try changing `n_spins`, `omega`, and `gamma`. For trajectory studies, replace `mesolve` with `mcsolve` and keep the same Dicke-space operators."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}