{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5260bbb0",
   "metadata": {},
   "source": [
    "# Using AngularKets and AngularStates in different spin coupling schemes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d147bcd",
   "metadata": {},
   "source": [
    "We define two kind of objets:\n",
    "- `AngularKet` objects represent canonical ketstates in a given coupling scheme.\n",
    "For each coupling scheme, the `AngularKet` objects form a complete orthonormal basis. \n",
    "These objects are immutable and hashable.\n",
    "- `AngularState` objects represent a statevector, which is defined by a coefficient vector and a list of `AngularKet` objects, which must all be of the same coupling scheme."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "803cf823",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rydstate.angular import AngularKetFJ, AngularKetJJ, AngularKetLS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ef6f51f",
   "metadata": {},
   "source": [
    "For AngularKets we always store the following basic angular momentum quantum numbers (qn):\n",
    "- the nuclear spin qn `i_c` (0 if no hyperfine splitting should be considered)\n",
    "- the core electron spin qn `s_c` (0 for Alkali atoms, 1/2 for Alkaline earth atoms)\n",
    "- the core electron orbital qn `l_c` (0 by default)\n",
    "- the Rydberg electron spin qn `s_r = 0.5`  and\n",
    "- the Rydberg electron orbital qn `l_r`.\n",
    "\n",
    "As well as the total combined qn:\n",
    "- the total atom angular momentum qn `f_tot`\n",
    "- Optionally: the magnetic quantum number `m`, which is the projection of `f_tot` onto the quantization axis.\n",
    "\n",
    "These basic angular momentum qns can couple via different coupling schemes to the total atom angular momentum `f_tot`.\n",
    "The different coupling schemes define in addition the following quantum numbers:\n",
    "- `AngularKetLS`:\n",
    "  - the total orbital momentum `(l_c, l_r)l_tot`,\n",
    "  - the total spin `(s_c, s_r)s_tot`, and\n",
    "  - the total angular momentum `(l_tot, s_tot)j_tot`\n",
    "- `AngularKetJJ`:\n",
    "  - the angular momentum of the core electron `(l_c, s_c)j_c`,\n",
    "  - the angular momentum of the Rydberg electron `(s_r, l_r)j_r`, and\n",
    "  - the total angular momentum `(j_c, j_r)j_tot`\n",
    "- `AngularKetFJ`:\n",
    "  - the angular momentum of the core electron `(l_c, s_c)j_c`,\n",
    "  - the angular momentum of the Rydberg electron `(s_r, l_r)j_r`, and \n",
    "  - the total spin momentum of the core electron `(j_c, i_c)f_c`\n",
    "\n",
    "Note the notation `(a, b)c` means that the angular momenta `a` and `b` couple to the combined angular momentum `c`.\n",
    "\n",
    "To create an angular ket, you can simply call the respective class and specify the needed quantum numbers.\n",
    "Note, that you only have to specify as many quantum numbers as needed to uniquely define the ket, all other quantum numbers will then be determined automatically.\n",
    "Furthermore, you can specify the atomic species to automatically set the nuclear spin `i_c` and the core electron spin `s_c` (the core electron orbital `l_c` will always be set to 0 by default).\n",
    "\n",
    "The magnetic quantum number `m` is optional, and only needed if you want to calculate concrete matrix elements.\n",
    "If you don't specify `m`, you still can calculate reduced matrix elements as well as reduced overlaps between two `AngularKet` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "edddbef8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ket1=AngularKetLS(i_c=2.5, s_c=0.5, l_c=0, s_r=0.5, l_r=0, s_tot=0.0, l_tot=0, j_tot=0.0, f_tot=2.5)\n",
      "ket2=AngularKetJJ(i_c=2.5, s_c=0.5, l_c=0, s_r=0.5, l_r=0, j_c=0.5, j_r=0.5, j_tot=0.0, f_tot=2.5)\n",
      "ket3=AngularKetFJ(i_c=2.5, s_c=0.5, l_c=0, s_r=0.5, l_r=0, j_c=0.5, f_c=2.0, j_r=0.5, f_tot=2.5)\n"
     ]
    }
   ],
   "source": [
    "ket1 = AngularKetLS(s_tot=0, l_r=0, j_tot=0, species=\"Yb173\")\n",
    "print(f\"{ket1=}\")\n",
    "ket2 = AngularKetJJ(j_tot=0, l_r=0, f_tot=2.5, species=\"Yb173\")\n",
    "print(f\"{ket2=}\")\n",
    "ket3 = AngularKetFJ(f_c=2, l_r=0, f_tot=2.5, species=\"Yb173\")\n",
    "print(f\"{ket3=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8327c524",
   "metadata": {},
   "source": [
    "Calculating overlaps between two angular states of different coupling schemes is as simply as calling the `calc_reduced_overlap` method.\n",
    "Note, that this method ignores any given magnetic quantum numbers `m` (if specified).\n",
    "This method will automatically check the coupling schemes of the two states and calculate the needed Wigner6j and Wigner9j symbols."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e8f13db5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ket1.calc_reduced_overlap(ket2)=1.0\n",
      "ket1.calc_reduced_overlap(ket3)=np.float64(-0.6454972243679028)\n",
      "ket2.calc_reduced_overlap(ket3)=-0.6454972243679028\n"
     ]
    }
   ],
   "source": [
    "print(f\"{ket1.calc_reduced_overlap(ket2)=}\")\n",
    "print(f\"{ket1.calc_reduced_overlap(ket3)=}\")\n",
    "print(f\"{ket2.calc_reduced_overlap(ket3)=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b58dd2",
   "metadata": {},
   "source": [
    "You can also convert the ket objects to state objects of different coupling schemes by using the `to_state` method with the desired coupling scheme as an argument.\n",
    "This will return a `AngularState` object, which contains a list of `AngularKet` objects and the respective coefficients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd8c1782",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.0*JJ(i_c=2.5, s_c=0.5, l_c=0, s_r=0.5, l_r=0, j_c=0.5, j_r=0.5, j_tot=0.0, f_tot=2.5)\n",
      "-0.6454972243679028*FJ(i_c=2.5, s_c=0.5, l_c=0, s_r=0.5, l_r=0, j_c=0.5, f_c=2.0, j_r=0.5, f_tot=2.5), 0.7637626158259734*FJ(i_c=2.5, s_c=0.5, l_c=0, s_r=0.5, l_r=0, j_c=0.5, f_c=3.0, j_r=0.5, f_tot=2.5)\n"
     ]
    }
   ],
   "source": [
    "print(ket1.to_state(\"JJ\"))\n",
    "print(ket1.to_state(\"FJ\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ff43497",
   "metadata": {},
   "source": [
    "The `AngularState` class also provides methods to calculate expectation and standard deviations of quantum numbers: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5b536e6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FJ(i_c=2.5, s_c=0.5, l_c=0, s_r=0.5, l_r=1, j_c=0.5, f_c=2.0, j_r=1.5, f_tot=2.5)\n",
      "exp_f_c=np.float64(2.0), std_f_c=0\n",
      "exp_s_tot=np.float64(0.7777777777777777), std_s_tot=0.41573970964154916\n"
     ]
    }
   ],
   "source": [
    "ket = AngularKetFJ(f_c=2, l_r=1, j_r=1.5, f_tot=2.5, species=\"Yb173\")\n",
    "\n",
    "ket_as_statefj = ket.to_state(\"FJ\")\n",
    "exp_f_c = ket_as_statefj.calc_exp_qn(\"f_c\")\n",
    "std_f_c = ket_as_statefj.calc_std_qn(\"f_c\")\n",
    "\n",
    "ket_as_statels = ket.to_state(\"LS\")\n",
    "exp_s_tot = ket_as_statels.calc_exp_qn(\"s_tot\")\n",
    "std_s_tot = ket_as_statels.calc_std_qn(\"s_tot\")\n",
    "\n",
    "print(ket)\n",
    "print(f\"{exp_f_c=}, {std_f_c=}\")\n",
    "print(f\"{exp_s_tot=}, {std_s_tot=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b0b3b72",
   "metadata": {},
   "source": [
    "And we can calculate matrix elements of operators between two `AngularState` objects by using the `calc_matrix_element` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51756a1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.1753272381496312\n",
      "0.13801311186847098\n"
     ]
    }
   ],
   "source": [
    "ket1 = AngularKetFJ(f_c=2, l_r=1, j_r=1.5, f_tot=2.5, species=\"Yb173\")\n",
    "ket2 = AngularKetFJ(f_c=2, l_r=2, j_r=1.5, f_tot=2.5, species=\"Yb173\")\n",
    "\n",
    "print(ket1.calc_reduced_matrix_element(ket2, operator=\"spherical\", kappa=1))\n",
    "print(ket1.calc_reduced_matrix_element(ket1, operator=\"s_tot\", kappa=1))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rydstate",
   "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.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}