Repository: pyg-team/pytorch_geometric

Issue:
math term missed in dimenet_utils.associated_legendre_polynomials
### 🛠 Proposed Refactor

In order to get the function of 'spherical harmonics' basis used in DimeNet, we need to calc 'Associated Legendre Polynomial' in math. But I just found a term missed in torch_geometric.nn.models.dimenet_utils.associated_legendre_polynomials(). 

The author of DimeNet use `zero_m_only=True` by default, so there will be no error when running the code by default. 

But just a warning for those who also read or use this code, the author of the code (version of pyg: 2.3.1) missed a term when calculating $P_l^l$ which make all results except $P_l^0$ wrong.  [original code here](https://github.com/pyg-team/pytorch_geometric/blob/f71ead8ade8a67be23982114cfff649b7d074cfb/torch_geometric/nn/models/dimenet_utils.py#L86)

Take the result of $P_1^1$ as an example:

original code will output -1, but according to eq.14 at https://mathworld.wolfram.com/AssociatedLegendrePolynomial.html, the ground truth result is $-(1-x^2)^{0.5}$. 

### Suggest a potential alternative/fix

Beyond adding the missed term, I also add some comment to help understanding.
```python
def associated_legendre_polynomials(k, zero_m_only=True):
    '''
    helper function to calc Y_l^m
    '''
    z = sym.symbols('z')
    P_l_m = [[0] * (j + 1) for j in range(k)]

    P_l_m[0][0] = 1
    if k > 0:
        P_l_m[1][0] = z

        for j in range(2, k):
            P_l_m[j][0] = sym.simplify(((2 * j - 1) * z * P_l_m[j - 1][0] -
                                        (j - 1) * P_l_m[j - 2][0]) / j) # use the property of eq.7: https://mathworld.wolfram.com/AssociatedLegendrePolynomial.html
        if not zero_m_only:
            for i in range(1, k):
                P_l_m[i][i] = sym.simplify((1 - 2 * i) * P_l_m[i - 1][i - 1]*(1-z**2)**0.5) # add missed term (*(1-z**2)**0.5) here
                if i + 1 < k:
                    P_l_m[i + 1][i] = sym.simplify(
                        (2 * i + 1) * z * P_l_m[i][i]) # use the property of eq.11: https://mathworld.wolfram.com/AssociatedLegendrePolynomial.html
                for j in range(i + 2, k):
                    P_l_m[j][i] = sym.simplify(
                        ((2 * j - 1) * z * P_l_m[j - 1][i] -
                         (i + j - 1) * P_l_m[j - 2][i]) / (j - i)) # use the property of eq.7: https://mathworld.wolfram.com/AssociatedLegendrePolynomial.html

    return P_l_m
```
After fixing this tiny bug, all the results are the right


Discussion:

Resolving patch:
diff --git a/torch_geometric/nn/models/dimenet_utils.py b/torch_geometric/nn/models/dimenet_utils.py
--- a/torch_geometric/nn/models/dimenet_utils.py
+++ b/torch_geometric/nn/models/dimenet_utils.py
@@ -71,6 +71,7 @@ def sph_harm_prefactor(k, m):
 
 
 def associated_legendre_polynomials(k, zero_m_only=True):
+    r"""Helper function to calculate Y_l^m."""
     z = sym.symbols('z')
     P_l_m = [[0] * (j + 1) for j in range(k)]
 
@@ -79,15 +80,22 @@ def associated_legendre_polynomials(k, zero_m_only=True):
         P_l_m[1][0] = z
 
         for j in range(2, k):
+            # Use the property of Eq (7) in
+            # https://mathworld.wolfram.com/AssociatedLegendrePolynomial.html:
             P_l_m[j][0] = sym.simplify(((2 * j - 1) * z * P_l_m[j - 1][0] -
                                         (j - 1) * P_l_m[j - 2][0]) / j)
         if not zero_m_only:
             for i in range(1, k):
-                P_l_m[i][i] = sym.simplify((1 - 2 * i) * P_l_m[i - 1][i - 1])
+                P_l_m[i][i] = sym.simplify(
+                    (1 - 2 * i) * P_l_m[i - 1][i - 1] * (1 - z**2)**0.5)
                 if i + 1 < k:
+                    # Use the property of Eq (11) in
+                    # https://mathworld.wolfram.com/AssociatedLegendrePolynomial.html:
                     P_l_m[i + 1][i] = sym.simplify(
                         (2 * i + 1) * z * P_l_m[i][i])
                 for j in range(i + 2, k):
+                    # Use the property of Eq (7) in
+                    # https://mathworld.wolfram.com/AssociatedLegendrePolynomial.html:
                     P_l_m[j][i] = sym.simplify(
                         ((2 * j - 1) * z * P_l_m[j - 1][i] -
                          (i + j - 1) * P_l_m[j - 2][i]) / (j - i))


Test patch: