Merge branch 'release/1.0.4'

Cupy iplementation of KPM

Author: devmessias

README.md

@@ -1,12 +1,14 @@
 # ![eMaTe](emate.png)
 
-eMaTe is a python package implemented in tensorflow which the main goal is provide useful methods capable of estimate spectral densities and trace functions of large sparse matrices. 
+eMaTe it is a python package which the main goal is to provide  methods capable of estimating the spectral densities and trace 
+functions of large sparse matrices. eMaTe can run in both CPU and GPU and can estimate the spectral density and related trace functions, such as entropy and Estrada index, even in directed or undirected networks with million of nodes.
 
 ## Install                                                                                                              
 ```
 pip install emate
 ```
 
+If you a have a GPU you should also install cupy.
 ## Kernel Polynomial Method (KPM)
 
 The Kernel Polynomial Method can estimate the spectral density of large sparse Hermitan matrices with a computational cost almost linear. This method combines three key ingredients: the Chebyshev expansion + the stochastic trace estimator + kernel smoothing.
@@ -26,15 +28,15 @@ vals = np.linalg.eigvalsh(W).real
 ```
 
 ```python
-from emate.hermitian import pykpm
+from emate.hermitian import tfkpm
 from stdog.utils.misc import ig2sparse 
 
 W = ig2sparse(G)
 
-num_moments = 300
-num_vecs = 200
+num_moments = 40
+num_vecs = 40
 extra_points = 10
-ek, rho = pykpm(W, num_moments, num_vecs, extra_points)
+ek, rho = tfkpm(W, num_moments, num_vecs, extra_points)
 ```
 
 ```python
@@ -43,6 +45,18 @@ plt.hist(vals, density=True, bins=100, alpha=.9, color="steelblue")
 plt.scatter(ek, rho, c="tomato", zorder=999, alpha=0.9, marker="d")
 
 ```
+If the CUPY package it is available in your machine, you can also use the cupy implementation. When compared to tf-kpm, the
+Cupy-kpm is slower for median matrices (100k) and faster for larger matrices (> 10^6). The main reason it's because the tf-kpm was implemented in order to calc all te moments in a single step. 
+
+```python
+from emate.hermitian import cupykpm
+
+num_moments = 40
+num_vecs = 40
+extra_points = 10
+ek, rho = cupykpm(W, num_moments, num_vecs, extra_points)
+```
+
 
 ![](docs/source/imgs/kpm.png)
 
@@ -113,8 +127,5 @@ approximated_entropy, exact_entropy
 [[2] Ubaru, S., Chen, J., & Saad, Y. (2017). Fast Estimation of tr(f(A)) via Stochastic Lanczos Quadrature. SIAM Journal on Matrix Analysis and Applications, 38(4), 1075-1099.](https://epubs.siam.org/doi/abs/10.1137/16M1104974)
 
 
-## Acknowledgements
-
-This work has been supported also by FAPESP grants  11/50761-2  and  2015/22308-2.   Research  carriedout using the computational resources of the Center forMathematical  Sciences  Applied  to  Industry  (CeMEAI)funded by FAPESP (grant 2013/07375-0).
 
 

emate.png

@@ -3,7 +3,13 @@
 
 from emate import linalg
 from emate import utils
-__version__ = "1.0.3"
+
+try:
+    import cupy as cp
+except:
+    print("Warning: Cupy package not found")
+
+__version__ = "1.0.4"
 __license__ = ""
 __author__ = "Bruno Messias; Thomas K Peron"
 __author_email__ = "[email protected]"

emate/__init__.py

@@ -1,4 +1,4 @@
-from emate.hermitian.kpm import pykpm
-from emate.hermitian import tfops
+from emate.hermitian.kpm import pykpm, cupykpm, tfkpm
+from emate.hermitian import tfops, cupyops
 
-__all__ = ["pykpm", "tfops"]
+__all__ = ["pykpm", "tfops", "cupykpm", "cupykpm", "tfkpm"]

emate/hermitian/__init__.py

@@ -0,0 +1,107 @@
+"""
+Kernel Polynomial Method
+========================
+
+The kernel polynomial method is an algorithm to obtain an approximation
+for the spectral density of a Hermitian matrix. This algorithm combines
+expansion in polynomials of Chebyshev with the stochastic trace in order
+to obtain such approximation.
+
+Applications
+------------
+
+    - Hamiltonian matrices associated with quantum mechanics
+    - Magnetic Laplacian associated with directed graphs
+    - etc
+
+Available functions
+-------------------
+
+
+"""
+import numpy as np
+try:
+    import cupy as cp
+except:
+    cp = None
+
+from emate.utils.cupyops.signal import dctIII
+
+
+def get_moments(
+    H_rescaled,
+    num_moments,
+    dimension,
+    precision=32
+):
+    """
+    Parameters
+    ----------
+        H: sparse cupy of rank 2
+        num_moments: (uint) number of cheby. moments
+        dimension: (uint) size of the matrix
+
+        alpha0: Tensor(shape=(H.shape[0], num_vecs), dtype=tf_complex)
+        alpha1: Tensor(shape=(H.shape[0], num_vecs), dtype=tf_complex)
+
+
+    Returns
+    -------
+    """
+    cp_complex = cp.complex64
+    if precision == 64:
+        cp_complex = cp.complex128
+        
+    alpha0 = cp.exp(1j*2*cp.pi*cp.random.rand(dimension))
+    alpha1 = H_rescaled.dot(alpha0)
+    mu = cp.zeros(num_moments, dtype=cp_complex)
+    mu[0] = (alpha0.T.conj()).dot(alpha0)
+    mu[1] = (alpha0.T.conj()).dot(alpha1)
+
+    for i_moment in range(1, num_moments//2):
+        alpha2 = 2*H_rescaled.dot(alpha1)-alpha0
+        mu[2*i_moment] = 2*(alpha1.T.conj()).dot(alpha1) - mu[0]
+        mu[2*i_moment+1] = 2*(alpha2.T.conj()).dot(alpha1) - mu[1]
+
+        alpha0 = alpha1
+        alpha1 = alpha2
+
+    return mu
+
+
+def apply_kernel(
+    moments,
+    kernel,
+    dimension,
+    num_moments,
+    num_vecs,
+    extra_points=1,
+):
+    """
+    Parameters
+    ----------
+
+    """
+
+    moments = cp.sum(moments.real, axis=0)
+    moments = moments/num_vecs/dimension
+
+    num_points = extra_points+num_moments
+
+    if kernel is not None:
+        moments = moments*kernel
+
+    mu_ext = cp.zeros(num_points)
+    mu_ext[0:num_moments] = moments
+
+    smooth_moments = dctIII(mu_ext)
+    points = cp.arange(0, num_points)
+    ek = cp.cos(cp.pi*(points+0.5)/num_points)
+    gk = cp.pi*cp.sqrt(1.-ek**2)
+   
+    rho = cp.divide(smooth_moments, gk)
+
+    return ek, rho
+
+
+__all__ = ["apply_kernel", "get_moments"]

emate/hermitian/cupyops/__init__.py

@@ -32,20 +32,34 @@
 """
 import numpy as np
 import tensorflow as tf
+try:
+    import cupy as cp
+except:
+    cp = None
 
-from emate.linalg import rescale_matrix
+from emate.linalg import rescale_matrix, get_bounds
+from emate.linalg.misc import rescale_cupy
 
 from emate.utils.tfops.vector_factories import normal_complex
 
-from emate.utils.tfops.kernels import jackson as jackson_kernel
-from emate.hermitian.tfops.kpm import get_moments, apply_kernel, rescale_kpm
+from emate.utils.tfops.kernels import jackson as tf_jackson
+from emate.utils.cupyops.kernels import jackson as cupy_jackson
 
+from emate.hermitian.tfops.kpm import get_moments, apply_kernel
+from emate.hermitian.cupyops import kpm as cupyops
+
+
+def rescale_kpm(ek, rho, scale_fact_a, scale_fact_b):
+
+    ek = ek*scale_fact_a + scale_fact_b
+    rho = rho/scale_fact_a
+    return ek, rho
 
 def pykpm(
     H,
-    num_moments,
-    num_vecs,
-    extra_points,
+    num_moments=10,
+    num_vecs=10,
+    extra_points=2,
     precision=32,
     lmin=None,
     lmax=None,
@@ -100,6 +114,9 @@ def pykpm(
 
     """
 
+    if (lmin is None) or (lmax is None):
+        lmin, lmax = get_bounds(H)
+
     H, scale_fact_a, scale_fact_b = rescale_matrix(H, lmin, lmax,)
 
     coo = H.tocoo()
@@ -128,10 +145,10 @@ def pykpm(
 
     dimension = H.shape[0]
 
-    tf.reset_default_graph()
+    tf.compat.v1.reset_default_graph()
     with tf.device(device):
-        sp_indices = tf.placeholder(dtype=tf.int64, name="sp_indices")
-        sp_values = tf.placeholder(
+        sp_indices = tf.compat.v1.placeholder(dtype=tf.int64, name="sp_indices")
+        sp_values = tf.compat.v1.placeholder(
             dtype=tf_type,
             name="sp_values"
         )
@@ -146,7 +163,7 @@ def pykpm(
             precision=precision
         )
         moments = get_moments(H, num_vecs, num_moments, alpha0)
-        kernel0 = jackson_kernel(num_moments, precision=32)
+        kernel0 = tf_jackson(num_moments, precision=32)
         if precision == 64:
             moments = tf.cast(moments, tf.float32)
         ek, rho = apply_kernel(
@@ -159,10 +176,55 @@ def pykpm(
         )
         ek, rho = rescale_kpm(ek, rho, scale_fact_a, scale_fact_b)
 
-    with tf.Session() as sess:
+    with tf.compat.v1.Session() as sess:
         ek, rho = sess.run([ek, rho], feed_dict)
 
     return ek, rho
 
+def cupykpm(
+    H,
+    num_moments=10,
+    num_vecs=10,
+    extra_points=12,
+    precision=32,
+    lmin=None,
+    lmax=None,
+    epsilon=0.01
+):
+
+    dimension = H.shape[0]
+
+    if (lmin is None) or (lmax is None):
+        lmin, lmax = get_bounds(H)
+
+    H  = cp.sparse.csr_matrix(
+        (
+            cp.array(H.data.astype("complex64")), 
+            cp.array(H.indices),
+            cp.array( H.indptr)
+        ), 
+        shape=H.shape, dtype="complex64"
+    )
+
+    H, scale_fact_a, scale_fact_b = rescale_cupy(H, lmin, lmax, epsilon)
+    
+    moments = cp.array([
+        cupyops.get_moments(H, num_moments, dimension, precision=precision)
+        for i in range(num_vecs)
+    ])
+    kernel0 = cupy_jackson(num_moments, precision=precision)
+ 
+    ek, rho = cupyops.apply_kernel(
+        moments,
+        kernel0,
+        dimension,
+        num_moments,
+        num_vecs,
+        extra_points
+    )
+    ek, rho = rescale_kpm(ek, rho, scale_fact_a, scale_fact_b)
+
+    return ek, rho
 
-__all__ = ["pykpm"]
+tfkpm = pykpm
+__all__ = ["pykpm", "cupykpm", "tfkpm"]

emate/hermitian/cupyops/kpm.py

@@ -69,7 +69,7 @@ def get_moments(
 
         # first_moment.shape = (num_vecs, )
         if complex_eval:
-            alpha0conj = tf.conj(alpha0)
+            alpha0conj = tf.math.conj(alpha0)
         else:
             alpha0conj = alpha0
 
@@ -128,8 +128,8 @@ def body(
             alpha2 = matrix_mul-alpha0
 
             if complex_eval:
-                alpha1conj = tf.conj(alpha1)
-                alpha2conj = tf.conj(alpha2)
+                alpha1conj = tf.math.conj(alpha1)
+                alpha2conj = tf.math.conj(alpha2)
 
             else:
                 alpha1conj = alpha1
@@ -256,22 +256,15 @@ def apply_kernel(
         num_points = num_moments+extra_points
 
         smooth_moments = tf.reshape(smooth_moments, [num_points])
-        smooth_moments = tf.spectral.dct(smooth_moments, type=3)
+        smooth_moments = tf.signal.dct(smooth_moments, type=3)
 
         points = tf.range(num_points, dtype=tf.float32)
 
         ek = tf.cos(np.pi*(points+0.5)/(num_points))
         gk = np.pi*tf.sqrt(1.-ek**2)
-        rho = tf.divide(smooth_moments, gk)
+        rho = tf.math.divide(smooth_moments, gk)
 
     return ek, rho
 
 
-def rescale_kpm(ek, rho, scale_fact_a, scale_fact_b):
-
-    ek = ek*scale_fact_a + scale_fact_b
-    rho = rho/scale_fact_a
-    return ek, rho
-
-
-__all__ = ["rescale_kpm", "apply_kernel", "get_moments"]
+__all__ = [ "apply_kernel", "get_moments"]

emate/hermitian/kpm.py

@@ -14,6 +14,6 @@
 """
 
 from emate.linalg import tfops
-from emate.linalg.misc import rescale_matrix, get_bounds
+from emate.linalg.misc import rescale_matrix, get_bounds, rescale_cupy
 
-__all__ = ["rescale_matrix", "get_bounds", "tfops"]
+__all__ = ["rescale_matrix", "get_bounds", "tfops", "rescale_cupy"]