Repository: graspologic-org/graspologic

Issue:
add diag aug to ASE
ASE(A) operates on

A + diag(degree_vec / n-1)


Discussion:
on c. elegans, try diag augmentaion times some constant.
make pairs plots for each.

Resolving patch:
diff --git a/graspy/embed/ase.py b/graspy/embed/ase.py
--- a/graspy/embed/ase.py
+++ b/graspy/embed/ase.py
@@ -15,7 +15,13 @@
 import warnings
 
 from .base import BaseEmbed
-from ..utils import import_graph, is_fully_connected
+from ..utils import (
+    import_graph,
+    is_fully_connected,
+    augment_diagonal,
+    pass_to_ranks,
+    is_unweighted,
+)
 
 
 class AdjacencySpectralEmbed(BaseEmbed):
@@ -57,12 +63,19 @@ class AdjacencySpectralEmbed(BaseEmbed):
         'truncated'. The default is larger than the default in randomized_svd 
         to handle sparse matrices that may have large slowly decaying spectrum.
 
-    check_lcc : bool , optional (defult = True)
+    check_lcc : bool , optional (default = True)
         Whether to check if input graph is connected. May result in non-optimal 
         results if the graph is unconnected. If True and input is unconnected,
         a UserWarning is thrown. Not checking for connectedness may result in 
         faster computation.
 
+    diag_aug : bool, optional (default = True)
+        Whether to replace the main diagonal of the adjacency matrix with a vector 
+        corresponding to the degree (or sum of edge weights for a weighted network) 
+        before embedding. Empirically, this produces latent position estimates closer
+        to the ground truth. 
+
+
     Attributes
     ----------
     latent_left_ : array, shape (n_samples, n_components)
@@ -104,6 +117,7 @@ def __init__(
         algorithm="randomized",
         n_iter=5,
         check_lcc=True,
+        diag_aug=True,
     ):
         super().__init__(
             n_components=n_components,
@@ -113,6 +127,10 @@ def __init__(
             check_lcc=check_lcc,
         )
 
+        if not isinstance(diag_aug, bool):
+            raise TypeError("`diag_aug` must be of type bool")
+        self.diag_aug = diag_aug
+
     def fit(self, graph, y=None):
         """
         Fit ASE model to input graph
@@ -137,5 +155,8 @@ def fit(self, graph, y=None):
                 )
                 warnings.warn(msg, UserWarning)
 
+        if self.diag_aug:
+            A = augment_diagonal(A)
+
         self._reduce_dim(A)
         return self
diff --git a/graspy/models/rdpg.py b/graspy/models/rdpg.py
--- a/graspy/models/rdpg.py
+++ b/graspy/models/rdpg.py
@@ -114,7 +114,9 @@ def fit(self, graph, y=None):
             )
         graph = augment_diagonal(graph, weight=self.diag_aug_weight)
         graph += self.plus_c_weight / graph.size
-        ase = AdjacencySpectralEmbed(n_components=self.n_components, **self.ase_kws)
+        ase = AdjacencySpectralEmbed(
+            n_components=self.n_components, diag_aug=False, **self.ase_kws
+        )
         latent = ase.fit_transform(graph)
         self.latent_ = latent
         if type(self.latent_) == tuple:
diff --git a/graspy/utils/utils.py b/graspy/utils/utils.py
--- a/graspy/utils/utils.py
+++ b/graspy/utils/utils.py
@@ -563,19 +563,18 @@ def get_multigraph_intersect_lcc(graphs, return_inds=False):
 
 def augment_diagonal(graph, weight=1):
     r"""
-    Replaces the diagonal of adjacency matrix with the sum of the edge weight.
-    This is equivalent to the degree for a weighted network.
-
-    For directed graphs, the degree used is the out degree (number) of 
-    edges leaving the vertex. Ignores self-loops when calculating degree.
-
+    Replaces the diagonal of an adjacency matrix with :math:`\frac{d}{nverts - 1}` where
+    :math:`d` is the degree vector for an unweighted graph and the sum of magnitude of 
+    edge weights for each node for a weighted graph. For a directed graph the in/out 
+    :math:`d` is averaged.
+    
     Parameters
     ----------
     graph: nx.Graph, nx.DiGraph, nx.MultiDiGraph, nx.MultiGraph, np.ndarray
         Input graph in any of the above specified formats. If np.ndarray, 
         interpreted as an :math:`n \times n` adjacency matrix
-    weight: int
-         Weight to be applied to each index.
+    weight: float/int
+        scalar value to multiply the new diagonal vector by
     
     Returns
     -------
@@ -595,12 +594,16 @@ def augment_diagonal(graph, weight=1):
     """
     graph = import_graph(graph)
     graph = remove_loops(graph)
+
     divisor = graph.shape[0] - 1
-    # use out degree for directed graph
-    # ignore self loops in either case
-    degrees = np.sum(graph, axis=1)
+
+    in_degrees = np.sum(np.abs(graph), axis=0)
+    out_degrees = np.sum(np.abs(graph), axis=1)
+    degrees = (in_degrees + out_degrees) / 2
+
     diag = weight * degrees / divisor
     graph += np.diag(diag)
+
     return graph
 
 


Test patch:
diff --git a/tests/test_spectral_embed.py b/tests/test_spectral_embed.py
--- a/tests/test_spectral_embed.py
+++ b/tests/test_spectral_embed.py
@@ -1,5 +1,4 @@
 import unittest
-import graspy as gs
 import numpy as np
 from graspy.embed.ase import AdjacencySpectralEmbed
 from graspy.embed.lse import LaplacianSpectralEmbed
@@ -104,6 +103,11 @@ def test_unconnected_warning(self):
             ase = AdjacencySpectralEmbed()
             ase.fit(A)
 
+    def test_input_checks(self):
+        with self.assertRaises(TypeError):
+            ase = AdjacencySpectralEmbed(diag_aug="over 9000")
+            ase.fit()
+
 
 class TestLaplacianSpectralEmbed(unittest.TestCase):
     def test_output_dim(self):
diff --git a/tests/test_utils.py b/tests/test_utils.py
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -372,7 +372,7 @@ def test_augment_diagonal_undirected(self):
     def test_augment_diagonal_directed(self):
         A = np.array(
             [
-                [0, 1, 1, 0, 0],
+                [0, 1, -1, 0, 0],
                 [0, 0, 0, 2, 1],
                 [1, 0, 0, 1, 1],
                 [0, 2, 0, 0, 0],
@@ -380,11 +380,11 @@ def test_augment_diagonal_directed(self):
             ]
         )
         expected = A.copy().astype(float)
-        expected[0, 0] = 2.0 / 4
-        expected[1, 1] = 3.0 / 4
-        expected[2, 2] = 3.0 / 4
-        expected[3, 3] = 2.0 / 4
-        expected[4, 4] = 1.0 / 4
+        expected[0, 0] = 1.5 / 4
+        expected[1, 1] = 3 / 4
+        expected[2, 2] = 2.5 / 4
+        expected[3, 3] = 2.5 / 4
+        expected[4, 4] = 1.5 / 4
         A_aug = gus.augment_diagonal(A)
         np.testing.assert_array_equal(A_aug, expected)