Repository: dmlc/dgl

Issue:
`dgl.lap_pe` cannot scale to large graphs due to materialization of dense adjacency matrix
## 🐛 Bug

`dgl.lap_pe` cannot scale to large graphs due to materialization of dense adjacency matrix

## To Reproduce

@vijaydwivedi75 told me that he would like to scale `dgl.lap_pe` to larger graphs like OGB products, but current implementation throws a memory error:

```python
import dgl 
from ogb.nodeproppred import DglNodePropPredDataset
dataset = DglNodePropPredDataset(name='ogbn-products')
graph = dataset[0][0]

dgl.lap_pe(graph, 5)
```

```
Traceback (most recent call last): 
  File "<stdin>", line 1, in <module>
  File "/home/vdwivedi/miniconda3/envs/transformers/lib/python3.10/site-packages/dgl/transforms/functional.py", line 3674, in lap_pe
    EigVal, EigVec = np.linalg.eig(L.toarray())
  File "/home/vdwivedi/miniconda3/envs/transformers/lib/python3.10/site-packages/scipy/sparse/_compressed.py", line 1051, in toarray
    out = self._process_toarray_args(order, out)
  File "/home/vdwivedi/miniconda3/envs/transformers/lib/python3.10/site-packages/scipy/sparse/_base.py", line 1298, in _process_toarray_args
    return np.zeros(self.shape, dtype=self.dtype, order=order)
numpy.core._exceptions._ArrayMemoryError: Unable to allocate 43.6 TiB for an array with shape (2449029, 2449029) and data type float64
```

## Expected behavior

We could make it work on graphs like OGB products (see fix below)

## Environment

 - DGL Version (e.g., 1.0): master
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3):
 - OS (e.g., Linux):
 - How you installed DGL (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version (if applicable):
 - GPU models and configuration (e.g. V100):
 - Any other relevant information:

## Additional context

@vijaydwivedi75 also suggested a fix.  Instead of materializing the dense adjacency matrix in
https://github.com/dmlc/dgl/blob/df97f2e8eea35636ab299e9fca988e269ad747ac/python/dgl/transforms/functional.py#L3675
We could as well use scipy's implementation which can work on sparse matrices:
```python
EigVal, EigVec = sp.linalg.eigs(L, k=pos_enc_dim+1, which='SR', tol=1e-2) # works fine for ogbn-products
```


Discussion:

Resolving patch:
diff --git a/python/dgl/transforms/functional.py b/python/dgl/transforms/functional.py
--- a/python/dgl/transforms/functional.py
+++ b/python/dgl/transforms/functional.py
@@ -3672,13 +3672,26 @@ def lap_pe(g, k, padding=False, return_eigval=False):
     L = sparse.eye(g.num_nodes()) - N * A * N
 
     # select eigenvectors with smaller eigenvalues O(n + klogk)
-    EigVal, EigVec = np.linalg.eig(L.toarray())
-    max_freqs = min(n - 1, k)
-    kpartition_indices = np.argpartition(EigVal, max_freqs)[: max_freqs + 1]
-    topk_eigvals = EigVal[kpartition_indices]
-    topk_indices = kpartition_indices[topk_eigvals.argsort()][1:]
-    topk_EigVec = EigVec[:, topk_indices]
-    eigvals = F.tensor(EigVal[topk_indices], dtype=F.float32)
+    if k + 1 < n - 1:
+        # Use scipy if k + 1 < n - 1 for memory efficiency.
+        EigVal, EigVec = scipy.sparse.linalg.eigs(
+            L, k=k + 1, which="SR", tol=1e-2
+        )
+        topk_indices = EigVal.argsort()[1:]
+        # Since scipy may return complex value, to avoid crashing in NN code,
+        # convert them to real number.
+        topk_eigvals = EigVal[topk_indices].real
+        topk_EigVec = EigVec[:, topk_indices].real
+    else:
+        # Fallback to numpy since scipy.sparse do not support this case.
+        EigVal, EigVec = np.linalg.eig(L.toarray())
+        max_freqs = min(n - 1, k)
+        kpartition_indices = np.argpartition(EigVal, max_freqs)[: max_freqs + 1]
+        topk_eigvals = EigVal[kpartition_indices]
+        topk_indices = kpartition_indices[topk_eigvals.argsort()][1:]
+        topk_EigVec = EigVec[:, topk_indices]
+        topk_EigVal = EigVal[topk_indices]
+    eigvals = F.tensor(topk_EigVal, dtype=F.float32)
 
     # get random flip signs
     rand_sign = 2 * (np.random.rand(max_freqs) > 0.5) - 1.0


Test patch: