diff --git a/EXPtools/utils/indexing.py b/EXPtools/utils/indexing.py index 13980ee..a5261d3 100644 --- a/EXPtools/utils/indexing.py +++ b/EXPtools/utils/indexing.py @@ -32,6 +32,38 @@ def I(l, m): assert abs(m) <= l, "m must be less than or equal to l" return int(l * (l + 1) / 2) + abs(m) +def I_nlm(n, l, m): + """ + Calculate the index of a spherical harmonic element given the angular numbers n, l, and m . + + Parameters + ---------- + n : int + The radial number + + l : int + The angular number + + m : int + The magnetic quantum number, ranging from 0 to l. + + Returns: + -------- + idx : int + The index corresponding to the specified angular numbers. + """ + assert isinstance(n, int) and isinstance(l, int) and isinstance(m, int), "all args must be integers" + assert n >= 0, "n must be greater than 0" + assert l >= 0, "l must be greater than 0" + assert abs(m) <= l, "m must be less than or equal to l" + + # determine the index of the n=n, l=0, m=0 term + n_idx = 0 + for j in range(n+1): + n_idx += int(j * (j + 1) / 2) # number of terms for a given n + # add the number of indices to reach the desired l,m + return I(l, m) + n_idx + def inverse_I(I): """ Calculate the angular numbers l and m given the index of a spherical harmonic element.