Repository: qutip/qutip

Issue:
Add test for orbital
### Describe the Issue!

The function `orbital` is not tested and there seems to have an error in the spherical harmonics computation.
Adding a test and fixing any error, possibly using `scipy.special.sph_harm`, would be helpful.


Discussion:
Hello Eric,
I could have a look at it and add a test with the `scipy.special.sph_harm` .
Could you explain in more detail what errors you got?
Hi Anto, 
Could you please leave the issue for a GSoC candidate.

There are no error raised but looking are the equations, something looks wrong. For example:
For positive `m`: `SPlm[1]`, but the equivalent term for negative values: `SPlm[abs(m)]`.

So when adding the tests, I expect that some fixes will be needed.




OK, i will leave this issue for the GSoC.
I guess it is my turn then :) 
I looked at the `orbital` function and compared it  to the scipy implementation of `sph_harm` and the results don't seem to agree. There is one detail I need for clarification: The `orbital` function uses `ket`states as input. Which component of the `ket` state should refer to which configuration of `m`? 

For example, if I pass `basis(3,0)` as argument is the corresponding setting `l = 1, m = -1` or `l = 1, m = 0` or `l = 1, m = 1`? 
I am not sure what the desired assignment  is. For me it makes sense to start with the smallest `m` (in the example `m=-1`).
Thank you for taking this issue.

The contribution to the phase is `np.exp(1.0j * 1 * phi)`, ignoring the `m` and it's not documented. So I guess all ket are `m = 1`...
Since it's wrong and not documented, you are free to choose which ever way you prefer. But please update the docstring so that is is clear.


Resolving patch:
diff --git a/qutip/orbital.py b/qutip/orbital.py
--- a/qutip/orbital.py
+++ b/qutip/orbital.py
@@ -1,7 +1,7 @@
 __all__ = ['orbital']
 
 import numpy as np
-from scipy.special import factorial
+from scipy.special import factorial, sph_harm
 
 
 def orbital(theta, phi, *args):
@@ -11,94 +11,71 @@ def orbital(theta, phi, *args):
     defined by theta and phi which is
     :math:`\sum_{lm} c_{lm} Y_{lm}(theta,phi)` where :math:`C_{lm}` are the
     coefficients specified by the list of kets. Each ket has 2l+1 components
-    for some integer l.
+    for some integer l. The first entry of the ket defines the coefficient
+    c_{l,-l}, while the last entry of the ket defines the
+    coefficient c_{l, l}.
 
     Parameters
     ----------
-    theta : list/array
-        Polar angles
+    theta : int/float/list/array
+        Polar angles in [0, pi]
 
-    phi : list/array
-        Azimuthal angles
+    phi : int/float/list/array
+        Azimuthal angles in [0, 2*pi]
 
     args : list/array
         ``list`` of ket vectors.
 
     Returns
     -------
-    ``array`` for angular wave function
+    ``array`` for angular wave function evaluated at all
+              possible combinations of theta and phi
 
     """
-    psi = 0.0
     if isinstance(args[0], list):
         # use the list in args[0]
         args = args[0]
 
+    # convert to numpy array
+    theta = np.atleast_1d(theta)
+    phi = np.atleast_1d(phi)
+    # check that arrays are only 1D
+    if len(theta.shape) != 1:
+        raise ValueError('Polar angles theta must be 1D list')
+    if len(phi.shape) != 1:
+        raise ValueError('Azimuthal angles phi must be 1D list')
+
+    # make meshgrid
+    phi_mesh, theta_mesh = np.meshgrid(phi, theta)
+    # setup empty wavefunction
+    psi = np.zeros([theta.shape[0], phi.shape[0]], dtype=complex)
+    # iterate through provided kets
     for k in range(len(args)):
         ket = args[k]
-        if not ket.type == 'ket':
-            raise TypeError('Invalid input ket in orbital')
-        sk = ket.shape
-        nchk = (sk[0] - 1) / 2.0
-        if nchk != np.floor(nchk):
+        if ket.type == 'bra':
+            ket = ket.conj()
+        elif not ket.type == 'ket':
+            raise TypeError('Invalid type for input ket in orbital')
+        # Extract l value from the state
+        l = (ket.shape[0] - 1) / 2.0
+        if l != np.floor(l):
             raise ValueError(
                 'Kets must have odd number of components in orbital')
-        l = int((sk[0] - 1) / 2)
-        if l == 0:
-            SPlm = np.sqrt(2) * np.ones((np.size(theta), 1), dtype=complex)
-        else:
-            SPlm = _sch_lpmv(l, np.cos(theta))
-        fac = np.sqrt((2.0 * l + 1) / (8 * np.pi))
-        kf = ket.full()
-        psi += np.sqrt(2) * fac * kf[l, 0] * np.ones((np.size(phi),
-                                                      np.size(theta)),
-                                                     dtype=complex) * SPlm[0]
-        for m in range(1, l + 1):
-            psi += ((-1.0) ** m * fac * kf[l - m, 0]) * \
-                np.array([np.exp(1.0j * 1 * phi)]).T * \
-                np.ones((np.size(phi), np.size(theta)),
-                        dtype=complex) * SPlm[1]
-        for m in range(-l, 0):
-            psi = psi + (fac * kf[l - m, 0]) * \
-                np.array([np.exp(1.0j * 1 * phi)]).T * \
-                np.ones((np.size(phi), np.size(theta)), dtype=complex) * \
-                SPlm[abs(m)]
-    return psi
-
-
-# Schmidt Semi-normalized Associated Legendre Functions
-def _sch_lpmv(n, x):
-    '''
-    Outputs array of Schmidt Seminormalized Associated Legendre Functions
-    S_{n}^{m} for m<=n.
-
-    Parameters
-    ----------
-    n : int
-        Degree of polynomial.
+        l = int(l)
+        # get factors from ket
+        factors = ket.full()
+        # iterate through the possible m
 
-    x : float
-        Point at which to evaluate
+        for i in range(len(factors)):
+            # set correct m
+            m = i - l
+            # calculate spherical harmonics
+            # note that theta and phi are interchanged in scipy implementation
+            res = sph_harm(m, l, phi_mesh, theta_mesh)
+            psi += factors[i] * res
 
-    Returns
-    -------
-    array of values for Legendre functions.
+    # flatten output if only one row
+    if psi.shape[1] == 1:
+        psi = psi.flatten()
 
-    '''
-    from scipy.special import lpmv
-    n = int(n)
-    sch = np.array([1.0])
-    sch2 = np.array([(-1.0) ** m * np.sqrt(
-        (2.0 * factorial(n - m)) / factorial(n + m)) for m in range(1, n + 1)])
-    sch = np.append(sch, sch2)
-    if isinstance(x, float) or len(x) == 1:
-        leg = lpmv(np.arange(0, n + 1), n, x)
-        return np.array([sch * leg]).T
-    else:
-        for j in range(0, len(x)):
-            leg = lpmv(range(0, n + 1), n, x[j])
-            if j == 0:
-                out = np.array([sch * leg]).T
-            else:
-                out = np.append(out, np.array([sch * leg]).T, axis=1)
-    return out
+    return psi


Test patch:
diff --git a/qutip/tests/test_orbital.py b/qutip/tests/test_orbital.py
new file mode 100644
--- /dev/null
+++ b/qutip/tests/test_orbital.py
@@ -0,0 +1,55 @@
+import pytest
+from scipy.special import sph_harm
+from qutip.orbital import orbital
+import qutip
+import numpy as np
+
+
+def test_orbital_single_ket():
+    """ Checks output for a single ket as input"""
+    # setup mesh for theta and phi
+    theta_list = np.linspace(0, np.pi, num=50)
+    phi_list = np.linspace(0, 2 * np.pi, num=100)
+    for theta, phi in zip(theta_list, phi_list):
+        # set l and m
+        for l in range(0, 5):
+            for m in range(-l, l + 1):
+                q = qutip.basis(2 * l + 1, l + m)
+                # check that outputs are the same,
+                # note that theta and phi are interchanged for scipy
+                assert sph_harm(m, l, phi, theta) == orbital(theta, phi, q)
+
+
+def test_orbital_multiple_ket():
+    """ Checks if the combination of multiple kets works """
+    theta_list = np.linspace(0, np.pi, num=50)
+    phi_list = np.linspace(0, 2 * np.pi, num=100)
+    l, m = 5, 2
+    q1 = qutip.basis(2 * l + 1, l + m)
+    l, m = 3, -1
+    q2 = qutip.basis(2 * l + 1, l + m)
+    for theta, phi in zip(theta_list, phi_list):
+        exp = sph_harm(2, 5, phi, theta) + sph_harm(-1, 3, phi, theta)
+        assert orbital(theta, phi, q1, q2) ==  exp
+
+
+def test_orbital_explicit():
+    """ Checks explicit configurations of orbital functions"""
+    theta_list = np.linspace(0, np.pi, num=50)
+    phi_list = np.linspace(0, 2 * np.pi, num=100)
+    # Constant function
+    l, m = 0, 0
+    q = qutip.basis(2 * l + 1, l + m)
+    assert orbital(0, 0, q) == 0.5 * np.sqrt(1 / np.pi)
+    # cosine function
+    l, m = 1, 0
+    q = qutip.basis(2 * l + 1, l + m)
+    assert np.allclose(orbital(theta_list, 0, q),
+                       0.5 * np.sqrt(3 / np.pi) * np.cos(theta_list))
+    # cosine with phase
+    l, m = 1, 1
+    q = qutip.basis(2 * l + 1, l + m)
+    phi_mesh, theta_mesh = np.meshgrid(phi_list, theta_list)
+    assert np.allclose(orbital(theta_list, phi_list, q),
+                       -0.5 * np.sqrt(3 / (2 * np.pi)) * np.sin(
+                           theta_mesh) * np.exp(1j * phi_mesh))