Merge pull request #29 from hfoote/indexing

added nlm indexer

Author: NGC

EXPtools/utils/indexing.py

@@ -32,6 +32,37 @@ 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, lmax):
+    """
+    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.
+        lmax : int 
+            The maximum value of l in the basis
+
+    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 = int(total_terms(lmax)*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.