Show how to efficiently store and access Wigner 3j coefficients when one of the j values is <=2

README.md

The notebook is viewable as such here. The storage isn't strictly as efficient as it could be, but it's pretty darn good. And the size of that storage only goes as jmax**2.

Wigner3j.ipynb

{
 "metadata": {
  "name": "",
  "signature": "sha256:3037e357808ca7f291f6a56117b6c1eb434d52d273d3d3b634d2ef6ed2817cfc"
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Given that you know that one $j_i$ value (let's say $j_3$) will be at most 2, you also know that $\\lvert m_3 \\rvert$ will be at most $2$, which means that $m_2$ must be within 2 of $-m_1$.  Also, the triangle inequality tells you that $j_2$ will be within 2 of $j_1$.\n",
      "\n",
      "Because of symmetries, we are free to stipulate that $j_2 \\leq j_1$ and $j_3 \\in \\{0,1,2\\}$.  And let's suppose that $j_1$ can go up to $j_{max} = 80$.  We can also cut off about half of the values by using symmetry to stipulate that $m_1 \\geq 0$.  The naive way to do this (shown below) includes some invalid indices.  But that should be okay, since you'll just never access them."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "from __future__ import division, print_function\n",
      "import numpy as np\n",
      "from sympy import init_printing, symbols, summation\n",
      "from sympy.physics.wigner import wigner_3j\n",
      "init_printing()"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "heading",
     "level": 1,
     "metadata": {},
     "source": [
      "Implementation"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Construct an array of Wigner 3j values as follows:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "w3j = np.empty((149445,), dtype=float)\n",
      "\n",
      "i=0\n",
      "for j1 in range(80+1):\n",
      "    for m1 in range(j1+1):\n",
      "        for j2 in range(j1-2,j1+1):\n",
      "            for m2 in range(-m1-2,-m1+2+1):\n",
      "                for j3 in [0,1,2]:\n",
      "                    for m3 in [-m1-m2]:\n",
      "                        tmp = wigner_3j(j1,j2,j3,m1,m2,m3)\n",
      "                        try:\n",
      "                            w3j[i] = float(tmp.evalf(n=18))\n",
      "                        except AttributeError:\n",
      "                            w3j[i] = float(tmp)\n",
      "                        i += 1\n",
      "if i!= w3j.size:\n",
      "    raise ValueError(\"Something went wrong in assigning the values!\")"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 2
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "You can save that array to a text file with something like"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "np.savetxt('Wigner3j_j1_80_j3_2.txt', w3j)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 3
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Then, load that array from the text file into whatever program you're using, and express the 3j values as a function with something like the following code:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "w3j = np.loadtxt('Wigner3j_j1_80_j3_2.txt')\n",
      "\n",
      "def wigner3j(j1,j2,j3,m1,m2,m3):\n",
      "    if abs(m1)>j1 or abs(m2)>j2 or abs(m3)>j3 or abs(j1-j2)>j3 or j1+j2<j3 or m1+m2+m3!=0:\n",
      "        return 0.0\n",
      "    if j3>2: # Re-order so that j3 is <=2, if possible\n",
      "        if j1<=2:\n",
      "            return wigner3j(j3,j2,j1,m3,m2,m1)\n",
      "        elif j2<=2:\n",
      "            return wigner3j(j1,j3,j2,m1,m3,m2) * (-1)**(j1+j2+j3)\n",
      "        else:\n",
      "            raise ValueError(\"wigner3j was not designed to handle indices {0}\".format((j1,j2,j3,m1,m2,m3)))\n",
      "    if j2>j1: # Re-order so that j1>=j2\n",
      "        return wigner3j(j2,j1,j3,m2,m1,m3) * (-1)**(j1+j2+j3)\n",
      "    if m1<0: # Use symmetry to eliminate ~1/2 the storage\n",
      "        return wigner3j(j1,j2,j3,-m1,-m2,-m3) * (-1)**(j1+j2+j3)\n",
      "    # The // division by 2 is standard c-style integer division\n",
      "    return w3j[(45*j1**2 + 15*j1 + 30*j2 + 2*j3 + 96*m1 + 6*m2 + 72) // 2]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 4
    },
    {
     "cell_type": "heading",
     "level": 1,
     "metadata": {},
     "source": [
      "Explanation of the index"
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Here, you can do a simple loop to get the indices.  It will include some invalid sets of indices, but that should be fine; you'll just never use them."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "indices = [[j1,j2,j3,m1,m2,m3]\n",
      "           for j1 in range(80+1) for m1 in range(j1+1)\n",
      "           for j2 in range(j1-2,j1+1) for m2 in range(-m1-2,-m1+2+1)\n",
      "           for j3 in [0,1,2] for m3 in [-m1-m2]]\n",
      "print(\"Total number of elements:\",len(indices))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "Total number of elements: 149445\n"
       ]
      }
     ],
     "prompt_number": 5
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "The array is created by a nested loop, where the innermost loop goes over values of $j_3$, and the outermost loop goes over values of $j_1$.  [Actually, the innermost loop as written is a trivial loop over one value of $m_3$, so we ignore it.]  So to index it, you have to start with the number of iterations required to get the innermost loop from the initial value of $j_3$ to the desired value of $j_3$.  For the second-to-innermost loop, you have to increase the index by the number of iterations of both loops required to increase the value of $m_2$ from its initial value to the desired value.  And so on.\n",
      "\n",
      "This is achieved by the following sum:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "j1,j2,j3,m1,m2,m3 = symbols('j1,j2,j3,m1,m2,m3', integer=True)\n",
      "\n",
      "index_value = (\n",
      " # no increment needed to get to m3\n",
      " + 0\n",
      "\n",
      " # get j3 from its initial value (0) to j3\n",
      " + 1 * (j3-(0))\n",
      "\n",
      " # 3 iterations for each step in m2, and its initial value is (-m1-2)\n",
      " + 3 * (m2-(-m1-2))\n",
      "\n",
      " # 15 iterations for each step in j2, and its initial value is (j1-2)\n",
      " + 15 * (j2-(j1-2))\n",
      "\n",
      " # 45 iterations for each step in m1, and its initial value is (0)\n",
      " + 45 * (m1-(0))\n",
      "\n",
      " # 45*(j1+1) iterations for each step in j1, and j1 needs to step from 2 to j1-1\n",
      " + summation( 45*(j1+1), (j1,0,j1-1))\n",
      "\n",
      ").expand().simplify()\n",
      "\n",
      "display(index_value)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "latex": [
        "$$\\frac{45 j_{1}^{2}}{2} + \\frac{15 j_{1}}{2} + 15 j_{2} + j_{3} + 48 m_{1} + 3 m_{2} + 36$$"
       ],
       "metadata": {},
       "output_type": "display_data",
       "png": "iVBORw0KGgoAAAANSUhEUgAAAZEAAAAvBAMAAADQuOrVAAAAMFBMVEX///8AAAAAAAAAAAAAAAAA\nAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAv3aB7AAAAD3RSTlMAMpndu3bvImbNiRBU\nq0Qb3U6NAAAACXBIWXMAAA7EAAAOxAGVKw4bAAAGDUlEQVRoBdVZW2gcVRj+9pKdzc5muhRfWhXH\nNhStpF1KVSxI11r1odKGohQUbcV6qZV2iEJFhCwFtS+SUNuHtWLXtPESCM6DrVCQBC99qiSokaiU\nrj6I+NJK0ya92PU/M3POnNk5k6Fakt0DmTn/93/zzf/NOTNzZgM0tFznXQ1Iq4Yf4+dWLb2h7ntx\nm9UAtWi4C721Fi09VPaYHYJaBVgIHLQX1qjck8zE3eG62wmXKWFGcyCrgd31C6yWd+gvZ7JeoGmb\nyIlECSSbKGjbDrw3YfGK9vKO2OvrzpATiSIyTdZ57QAwLmoyqkZBBLyzlZz4FI423b4acLLo1DdU\ndkNrDSdZi5xMvnk7oA+W2O1Qb7BBIXPiU8L55kCOgZxsYG+RJF6MKIk58Sk3RbDmG64yJ0C6iLdx\nNaIY5kRQFv4QwZpnOF9wnSSnYRvTgPKCu0445et5Ljni9CdATjr6kbtGr5Ii1BecnPgUNKmT706f\nvjiR6UfyCtBuQl0mOYmjRFyouYW3072OVD8wMhrtJI4ytyVHnG0Kmom3LICtHZVTh8YkjhKhPafw\n0fokeobZl+JS+lM56dp9fzmGMqcVx59sVYQT6cgoyiG6y55MrCBmfqdEF101SqvwY0+wg6Lb60OP\nRicBbejoKD2qBj+USTdrf1OoGhPBciknRSw6r9LzIlOvdxOQHReo1FGjqwETHaMSr7FLk3qs3AhK\ncQ/yM9A/wJ0Spl1OmxRKTr6Qsk7XpfSsneAJnW4u1j7vYk++4S1u5G953kdYT+jSKjxfhtYfTLuR\nR2q7igVFRZ4rbx3FebRXsV7mTNzBIumCfytnnb5L8eF8t9dvIycdPs57Is8BZy90aRWeK0FXVQqP\nZOzCiMopV/6jpl9CnxXQdwPpgnMxBYtDXA//zUn1AIwpO1fmcvJe2FXPLnFmNrselg9U9X0xVZZh\nQs9xsm9zATjxmEQW+QDKdZ1V+JmZT2AM73l8whjcKB0pXcb9qrxQZq/zS4s31+RjQ31+xlBCAEKP\nOUnbGq14zEwZA+tslyLyDB0c9lCu66zCk/VDyKUfwsjTbGInSlybkxJHTCkfVl68CvqlMl7mxyn3\nXEyZdEBRKXNC7UFkrQWjOYsWaU7jeYYmy8mSA4rL7azC9y66WLuFFhi9VfQh0VV0KRAk4KOCyIeV\n6Qn8jF63scxGXdE8Ne5EwTjnUnilzn1C0EpLx1krZSWmg3mGJqtZ7yhP11mFZ8fRt9MeqbH1xY80\nXRVOMn4+rEzned4+D/TV3DMqtp9WKvdUKvsVGQ7lKpV3d1QqJovZmNCvsJ/VgCPI1Ng0g5wnFEj2\n08bXdVbhGRuJqxhz1hfPCSc+SbOQ9vNh5T3AptEd5KRM2pGNj0kkIXjHk95KG2Dm8+5k8/MOerzg\nKnm6zip8hKBJ9l5bA2NaD4/JgitIz/j5kDJNrE1lenbNMiZ0hut0YgL30Y8B7Lsz1U0bamL2MfTY\nQReUdLcjYwGP4CXo12jI3gg7yZSQmvbzIWX6hWutRTNzmaet3l2nkx5kX/CGY8ATFE6cQUqWXNjX\nnYJxCNlumox0D7VXq2En2TLGpDwalfchfwHpqn6jnl3Jn85PQlv+PU2uZJGWXqZbsj8mDGU3hNOE\nE7YKv3X5ChhFtI0j/0sh7AS/L9ko5UPKxpLOGjDQVXCl2VZbsl6KHFyc0WP1dD7g9fhOXHMO0MQy\ngb/wmwuIPKHHC9qMizbq+gcrnEhJ6oaVg3kWfel8yQfwPwMR9HHn2SSDRlmOqJ+wewtIvH9qg4t7\neQdNWdmdLtqgK0vkt/FISQopc7a8/5UezHIc7nfQx/y2MBxAeqvrnOX9uTCaGBy2Aqgi0Jb+062A\nBUQfDkFlkZE6zwJnClIc7maK6GBfMLO1js1lRVqNKog3BDpsxzlpvxLv5IaU8v9F1tpxGinvjRfH\nm+e8cTm2gLPVWEozEFKl2Crou7sVWmdskclSLKUZCPSWjWtPxRGaI/8VaAU9a8uboEV007eEibYY\nJ/Q725am9wG8MjQQ8x8ebc2Q/3HaxI4O1+v0ETlba6ev4OJshPnN/QsqJ/4TvJUGVgAAAABJRU5E\nrkJggg==\n",
       "text": [
        "     2                                         \n",
        "45\u22c5j\u2081    15\u22c5j\u2081                                 \n",
        "\u2500\u2500\u2500\u2500\u2500\u2500 + \u2500\u2500\u2500\u2500\u2500 + 15\u22c5j\u2082 + j\u2083 + 48\u22c5m\u2081 + 3\u22c5m\u2082 + 36\n",
        "  2        2                                   "
       ]
      }
     ],
     "prompt_number": 6
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "print(2*index_value)"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "45*j1**2 + 15*j1 + 30*j2 + 2*j3 + 96*m1 + 6*m2 + 72\n"
       ]
      }
     ],
     "prompt_number": 7
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "We can check that this actually correctly indexes the `indices_2` array like this:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "def index(j1,j2,j3,m1,m2,m3):\n",
      "    return (45*j1**2 + 15*j1 + 30*j2 + 2*j3 + 96*m1 + 6*m2 + 72) // 2\n",
      "\n",
      "for i in range(len(indices)):\n",
      "    assert i == index(*indices[i])\n",
      "print(\"That seems to have worked\")"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "That seems to have worked\n"
       ]
      }
     ],
     "prompt_number": 8
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "As an additional check, we can calculate the total size of the array, as the index with $j_1 = j_{max}+1$ and everything else as the lowest possible indices for that value of $j_1$."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "jmax=80\n",
      "n_indices = index((jmax+1),(jmax+1)-2,0,0,-2,2)\n",
      "assert n_indices == len(indices)\n",
      "print(\"The size {0} seems to be correct.\".format(n_indices))"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "The size 149445 seems to be correct.\n"
       ]
      }
     ],
     "prompt_number": 9
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Just for fun, let's actually test that the result is the same as sympy's for every value you'll need up to $j_{1,2}=10$, and for every such value with $j_{1,2}=80$.  (That limited range should be convincing enough, and the rest just take too long.)"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "for j1 in range(10) + [80]:\n",
      "    for m1 in range(-j1,j1+1):\n",
      "        for j2 in range(10) + [80]:\n",
      "            for m2 in range(-j2,j2+1):\n",
      "                for j3 in [0,1,2]:\n",
      "                    for m3 in range(-j3,j3+1):\n",
      "                        a = wigner_3j(j1,j2,j3,m1,m2,m3)\n",
      "                        try:\n",
      "                            a = float(a.evalf(n=18))\n",
      "                        except AttributeError:\n",
      "                            a = float(a)\n",
      "                        b = wigner3j(j1,j2,j3,m1,m2,m3)\n",
      "                        assert a==b, 'sympy...wigner_3j(*{0})={1}; wigner3j(*{0})={2}'.format([j1,j2,j3,m1,m2,m3], a, b)\n",
      "print(\"That seems to have worked\")"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "That seems to have worked\n"
       ]
      }
     ],
     "prompt_number": 10
    }
   ],
   "metadata": {}
  }
 ]
}