Improved storage and fetching of Clebsch-Gordan coefficients

[GaffaSnobb]

Shell-model Hamiltonian

Recall the definition of the two-body part of the shell-model Hamiltonian:

$$H_\text{TB} = \sum\limits_{\alpha \leq \beta, \gamma \leq \delta, J, M} \langle \alpha \beta | V | \gamma \delta \rangle_J A^\dagger (\alpha \beta J M) A(\gamma \delta J M),$$ $$A^\dagger (\alpha \beta J M) = \frac{1}{\sqrt{1 + \delta_{\alpha \beta}}} \sum\limits_{m_\alpha, m_\beta} \langle j_\alpha, m_\alpha, j_\beta, m_\beta | J M \rangle c_{j_\alpha, m_\alpha}^\dagger c_{j_\beta, m_\beta}^\dagger.$$

There are of course many things to consider about these equations, but for now notice that the Clebsch-Gordan coefficients (CGCs) show up in the expressions for $A$ and $A^\dagger$. The CGCs are used in quantum mechanics to describe the coupling of two angular momenta to form a resultant total angular momentum state. They quantify the probability amplitude for a given combination of angular momenta and their z-components, taking the conservation of angular momentum into account. In the shell-model Hamiltonian, we can see that the CGCs are needed when we couple the two created single-particle states ($\alpha$ and $\beta$) and the two annihilated single-particle states ($\gamma$ and $\delta$).

Angular momentum conservation

The CGCs are identified by six quantum numbers which explain the total angular momentum and the z-component of the total angular momentum of the two single particles and the combined state that they couple to, namely $j_1, m_1, j_2, m_2, J, M$. We can immediately exploit the fact that angular momentum is conserved, which means that the CGCs are zero when the single-particle total angular momenta couple (well, they don't couple, thats the point) to a total angular momentum larger than their sum:

$$j_1 + j_2 < J \rightarrow \text{CGc} = 0,$$

and when the sum of the single-particle z-components do not sum exactly to the z-component of the total angular momentum they couple to:

$$m_1 + m_2 \neq M \rightarrow \text{CGc} = 0.$$

This is nice to know, because the fastest way to deal with the CGCs is probably to pre-calculate them in a lookup-table instead of calculating them every time they're needed. By angular momentum conservation we can reduce the number of iterations in the loops in the Hamiltonian by before doing any calculation, realising that the CGc will be zero and skip the calculation entirely.

An appropriate data structure

We still need a way to store the CGCs. One of the most readable ways of storing them is to use an std::unordered_map which more or less works exactly like a Python dict, except that you have to write the hashing function yourself if you want to use anything else but primitive data types (and a few others) as keys. We might for example create a struct to store the needed quantum numbers

struct Key6
{
    unsigned short j1;
    short m1;
    unsigned short j2;
    short m2;
    unsigned short J;
    short M;

    bool operator==(const Key6& other) const
    {
        return j1 == other.j1 && m1 == other.m1 && j2 == other.j2 && m2 == other.m2 && J == other.J && M == other.M;
    }
};

and then write a hashing function which produces a unique hash for every needed combination of quantum numbers

template <>
struct hash<Key6>
{
    size_t operator()(const Key6& k) const
    {
        size_t hash_j1 = std::hash<unsigned short>()(k.j1);
        size_t hash_m1 = std::hash<short>()(k.m1);
        size_t hash_j2 = std::hash<unsigned short>()(k.j2);
        size_t hash_m2 = std::hash<short>()(k.m2);
        size_t hash_J = std::hash<unsigned short>()(k.J);
        size_t hash_M = std::hash<short>()(k.M);

        return ((((hash_j1 ^ (hash_m1 << 1)) >> 1) ^ (hash_j2 << 1)) >> 1) ^
                ((((hash_m2 ^ (hash_J << 1)) >> 1) ^ (hash_M << 1)) >> 1);
    }
};

Disclaimer: ChatGPT wrote the hashing function. We can then create an std::unordered_map with Key6 as keys (using two times the angular momenta to remove half integers)

std::unordered_map<Key6, double> clebsch_gordan = {
    {{ 3, -3,  3,  3,  0,  0}, -0.5},
    ...
}

and then access the values by

Key6 key = {j1, m1, j2, m2, J, M};
double cg = clebsch_gordan.at(key);

This is a very readable way of doing it. There are however a few possible problems with the solution. Since I know little about hashing functions I do not know how well the Key6 hashing function performs. I have made sure that my program throws an error if a hash collision is detected, so that is no worry, but I believe that I can achieve a speed-up if I manage to convert the CGC map into a more primitive array. Further, if I want to run the two-body matrix element calculations on GPU, I cannot use std::unordered_map and I have to find a primitive way of storing the CGCs.

Storing Clebsch-Gordan coefficients in a (flattened) multidimensional array

I think it is inevitable to store the CGCs in an array, and to do that, I need to do some index gymnastics. Also, I probably have to store some zeros in the array to make the indices match with the quantum numbers of the CGCs. To illustrate the problem, let's imagine that the CGCs are identified by only two quantum numbers, $j_1$ and the corresponding $m_1$, for which we need only a 2D array. Let's say that we're dealing with the $sd$ major shell as the model space, meaning that we have the $1d5/2$, $2s1/2$, and $1d3/2$ orbitals. I always multiply the angular momenta by 2 as to avoid half-integer values, meaning that we have the following total angular momenta:

j = [1, 3, 5]
m_1 = [-1, 1]
m_3 = [-3, -1, 1, 3]
m_5 = [-5, -3, -1, 1, 3, 5]

Now, I want each of these values to directly correspond to an index in the array so we have to change the values a bit. The $m_i$ can be negative, but if we add $j_i$ to each of them they will be shifted into the positives:

m_5 + j_5 = [0, 2, 4, 6, 8, 10]

This is better, but with steps of two we would have to store a lot of unnecessary zeros so we'd better divide by two:

(m_5 + j_5)/2 = [0, 1, 2, 3, 4, 5]

We need to adjust the $j_i$ values too, and we can do that by adding 1, dividing by 2, subtracting 1:

(j + 1)/2 - 1 = [0, 1, 2]

As for now I will shift all the $m$ values by the largest $j$ value to keep the structure of the indices. This means that the $m$ values [-5, -3, -1, 1, 3, 5] will turn into indices [0, 1, 2, 3, 4, 5] while [-1, 1] turns into [2, 3] and not [0, 1]. This means that a lot of unnecessary zeros will be stored which is a potential for improvement later.

There will have to be a bit of arithmetic to convert the indices, but I hope that it will be less expensive than having to calculate a hash every time I want to fetch a CGC from the std::unordered_map.

To store the "CGCs" corresponding to the $j$ and $m$ values, we need a 2D array:

>>> np.zeros(shape=(3, 6), dtype=tuple)
[[0 0 (0, 2) (0, 3) 0 0]
 [0 (1, 1) (1, 2) (1, 3) (1, 4) 0]
 [(2, 0) (2, 1) (2, 2) (2, 3) (2, 4) (2, 5)]]

I have for illustrative purposes entered the $(j, m)$ index as the corresponding value and we can see that there will be some wasted space usage for storing zeros if we want to keep the dimensions uniform, 6 zeros in this case. Extending this to include all the needed quantum numbers, $j1, m1, j2, m2, J$, we will need a 5D array

np.zeros(shape=(3, 6, 3, 6, 6))

The 5th dimension of length 6 is because J = range(abs(j1 - j2), j1 + j2 + 2, 2). Visualising a 5D array is just messy so I won't explicitly show it here, but if you run the code yourself you'll see that there is a total of 1944 elements, of which 1384 are zeros. There is a lot of wasted space, so I am thinking about how to reduce the waste, but one step at the time. The total bytesize of the elements in the array is 15552 bytes, so hopefully it will still fit in CPU cache. Dont know how well that translates to GPU though.