Repository: neurodata/graspy (python)

Issue:
update requirements to scipy>=1.4
Scipy 1.4
- has much faster linear assignment problem, making FAQ way faster
- has MGC, which we eventually want for new nonpar, signal subgraph

Pull request description:
Update requirements to scipy >= 1.4
<!--
Thanks for contributing a pull request! Please ensure you have taken a look at
the contribution guidelines: https://github.com/neurodata/graspy/blob/master/CONTRIBUTING.md#pull-request-checklist
-->

#### Reference Issues/PRs
<!--
Example: Fixes #1234. See also #3456.
Please use keywords (e.g., Fixes) to create link to the issues or pull requests
you resolved, so that they will automatically be closed when your pull request
is merged. See https://github.com/blog/1506-closing-issues-via-pull-requests
-->
Fixes #302.

#### What does this implement/fix? Explain your changes.

Updated files [setup.py](https://github.com/neurodata/graspy/blob/master/setup.py#L13-22) and [requirement.txt](https://github.com/neurodata/graspy/blob/master/requirements.txt) so that new scipy version requirement is: scipy>=1.4.0

#### Any other comments?



Patch:
diff --git a/docs/reference/index.rst b/docs/reference/index.rst
index 6ebac002..7add0ef2 100644
--- a/docs/reference/index.rst
+++ b/docs/reference/index.rst
@@ -16,3 +16,4 @@ Reference
    utils
    pipeline
    match
+   subgraph
\ No newline at end of file
diff --git a/docs/reference/subgraph.rst b/docs/reference/subgraph.rst
new file mode 100644
index 00000000..803127b8
--- /dev/null
+++ b/docs/reference/subgraph.rst
@@ -0,0 +1,9 @@
+Subgraph
+========
+
+.. currentmodule:: graspy.subgraph
+
+Signal-Subgraph Estimators
+------------------
+
+.. autoclass:: SignalSubgraph
diff --git a/docs/tutorial.rst b/docs/tutorial.rst
index df48fca7..dc335e28 100644
--- a/docs/tutorial.rst
+++ b/docs/tutorial.rst
@@ -81,3 +81,14 @@ including an introduction to the module, and how to utilize the seeding feature.
 
    tutorials/matching/faq
    tutorials/matching/sgm
+
+.. _subgraph_tutorials:
+
+Subgraph
+========
+The following tutorial demonstrates how to estimate the signal-subgraph of samples of a graph/class model according to either the coherent or incoherent estimator models.
+
+.. toctree::
+   :maxdepth: 1
+
+   tutorials/subgraph/subgraph
diff --git a/docs/tutorials/subgraph/subgraph.ipynb b/docs/tutorials/subgraph/subgraph.ipynb
new file mode 100644
index 00000000..3bf7a5d1
--- /dev/null
+++ b/docs/tutorials/subgraph/subgraph.ipynb
@@ -0,0 +1,231 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Signal Subgraph Estimators"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### This tutorial will introduce the following signal-subgraph estimators:\n",
+    "- Incoherent subgraph estimator\n",
+    "- Coherent subgraph estimator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import graspy.subgraph as sg\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Preliminaries\n",
+    "\n",
+    "The general graph model is characterized by $M_V(m,s;\\pi,p,q)$, where: \n",
+    "\n",
+    "$V$ - the number of vertices in the graph\n",
+    "\n",
+    "$n$ - the number of graph samples\n",
+    "\n",
+    "$s$ - the number of edges that must be present in the subgraph\n",
+    "\n",
+    "$m$ - the number of vertices that each edge in the subgraph must be incident to\n",
+    "\n",
+    "$\\pi$ - the probability of a graph sample being of class 1\n",
+    "\n",
+    "$p$ - the probability of edges in the signal-subgraph, conditioned on Class 0\n",
+    "\n",
+    "$q$ - the probability of edges in the signal-subgraph, conditioned on Class 1\n",
+    "\n",
+    "The signal-subgraph is a subset of edges with distinct class-conditional likelihood parameters. The signal-subgraph estimator evaluates a test statistic for each edge in the graph and selects $s$ edges with the lowest test statistic, where $s$ is the desired size of the signal-subgraph. \n",
+    "\n",
+    "The estimator that is used to find the signal-subgraph determines certain properties of the resulting subgraph. Both estimators use $s$ to determine the size of the resulting subgraph. $m$ is only used for the coherent estimator, which constrains the subgraph to $m$ vertices. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Incoherent Signal-Subgraph Estimator\n",
+    "\n",
+    "For this example we will randomly select 20 edges from a graph with 70 vertices. These edges will have distinct class-conditional edge probabilities, and the graphs will be sampled from the model $M_{70}(20; 0.5, 0.8, 0.1)$, with $n = 100$."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from graspy.plot import heatmap\n",
+    "\n",
+    "verts = 70\n",
+    "sedges = 20\n",
+    "pi = 0.5\n",
+    "p = 0.8\n",
+    "q = 0.1\n",
+    "nsamples = 100\n",
+    "\n",
+    "np.random.seed(8888)\n",
+    "classlabels = np.zeros(nsamples, dtype=int)\n",
+    "classlabels[1::2] = 1\n",
+    "\n",
+    "sigsubindex = np.random.choice(verts ** 2, sedges, replace=False)\n",
+    "vect = p * np.ones(verts ** 2)\n",
+    "vect[sigsubindex] = q\n",
+    "vect = np.reshape(vect, (verts, verts))\n",
+    "expected = np.where(vect == q, 1, 0)\n",
+    "\n",
+    "blank = vect[:, :, None] + np.zeros(int(nsamples / 2))\n",
+    "A = p * np.ones((verts, verts, nsamples))\n",
+    "A[:, :, 1::2] = blank\n",
+    "A = np.random.binomial(1, A)\n",
+    "\n",
+    "sigsub = sg.SignalSubgraph()\n",
+    "sigsub.fit_transform(graphs=A, labels=classlabels, constraints=sedges)\n",
+    "\n",
+    "estimatesigsub = np.zeros((verts, verts))\n",
+    "estimatesigsub[sigsub.sigsub_] = 1\n",
+    "\n",
+    "fig, ax = plt.subplots(ncols=2, figsize=(10, 5), constrained_layout=True)\n",
+    "heatmap(expected, ax=ax[0], cbar=False, title=\"Expected Signal-Subgraph\")\n",
+    "heatmap(estimatesigsub, ax=ax[1], cbar=False, title=\"Estimated Signal-Subgraph\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Note that because $p$ and $q$ are sufficiently distinct and $n$ is sufficiently large, the Expected and Estimated signal-subgraphs should match exactly. "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Incoherent Signal-Subgraph Estimator\n",
+    "\n",
+    "Once again, we will randomly select 20 edges from a graph with 70 vertices. These edges will have distinct class-conditional edge probabilities, but the graphs will be sampled from the model $M_{70}(1, 20; 0.5, 0.8, 0.1)$, with $n = 100$.\n",
+    "\n",
+    "The estimated signal-subgraph will have 20 edges, constrained so that each edge must be incident to the same vertex. First, we will use the same expected signal-subgraph as the previous example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "np.random.seed(8888)\n",
+    "classlabels = np.zeros(nsamples, dtype=int)\n",
+    "classlabels[1::2] = 1\n",
+    "\n",
+    "sigsubindex = np.random.choice(verts ** 2, sedges, replace=False)\n",
+    "vect = p * np.ones(verts ** 2)\n",
+    "vect[sigsubindex] = q\n",
+    "vect = np.reshape(vect, (verts, verts))\n",
+    "expected = np.where(vect == q, 1, 0)\n",
+    "\n",
+    "blank = vect[:, :, None] + np.zeros(int(nsamples / 2))\n",
+    "A = p * np.ones((verts, verts, nsamples))\n",
+    "A[:, :, 1::2] = blank\n",
+    "A = np.random.binomial(1, A)\n",
+    "\n",
+    "sigsub = sg.SignalSubgraph()\n",
+    "sigsub.fit_transform(A, classlabels, [20, 1])\n",
+    "\n",
+    "estimatesigsub = np.zeros((verts, verts))\n",
+    "estimatesigsub[sigsub.sigsub_] = 1\n",
+    "\n",
+    "fig, ax = plt.subplots(ncols=2, figsize=(10, 5), constrained_layout=True)\n",
+    "heatmap(expected, ax=ax[0], cbar=False, title=\"Expected Signal-Subgraph\")\n",
+    "heatmap(estimatesigsub, ax=ax[1], cbar=False, title=\"Estimated Signal-Subgraph\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Notice how the coherent estimator constrains the estimated signal-subgraph to 20 edges that are incident to 1 vertex with the best total significance values. Now, we will try an expected signal-subgraph that is also limited to one vertex."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "mverts = 1\n",
+    "\n",
+    "np.random.seed(7777)\n",
+    "classlabels = np.zeros(nsamples, dtype=int)\n",
+    "classlabels[1::2] = 1\n",
+    "\n",
+    "m = np.random.choice(verts, mverts)\n",
+    "vect = p * np.ones(2 * verts * mverts - (mverts ** 2))\n",
+    "vect[np.random.choice(len(vect), sedges, replace=False)] = q\n",
+    "\n",
+    "blank = p * np.ones((verts, verts))\n",
+    "blank[m, :] = np.nan\n",
+    "blank[:, m] = np.nan\n",
+    "blank[np.isnan(blank)] = vect\n",
+    "expected = np.where(blank == q, 1, 0)\n",
+    "\n",
+    "blank = blank[:, :, None] + np.zeros(int(nsamples / 2))\n",
+    "A = p * np.ones((verts, verts, nsamples))\n",
+    "A[:, :, 1::2] = blank\n",
+    "A = np.random.binomial(1, A)\n",
+    "\n",
+    "sigsub = sg.SignalSubgraph()\n",
+    "sigsub.fit_transform(graphs=A, labels=classlabels, constraints=sedges)\n",
+    "\n",
+    "estimatesigsub = np.zeros((verts, verts))\n",
+    "estimatesigsub[sigsub.sigsub_] = 1\n",
+    "\n",
+    "fig, ax = plt.subplots(ncols=2, figsize=(10, 5), constrained_layout=True)\n",
+    "heatmap(expected, ax=ax[0], cbar=False, title=\"Expected Signal-Subgraph\")\n",
+    "heatmap(estimatesigsub, ax=ax[1], cbar=False, title=\"Estimated Signal-Subgraph\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now that the expected signal-subgraph is constrained to the coherent signal-subgraph model, the Expected and Estimated signal-subgraphs are exactly equal. "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
diff --git a/graspy/embed/base.py b/graspy/embed/base.py
index b4f91f50..d94c1b35 100644
--- a/graspy/embed/base.py
+++ b/graspy/embed/base.py
@@ -17,7 +17,7 @@ from abc import abstractmethod
 import numpy as np
 from sklearn.base import BaseEstimator
 
-from ..utils import import_graph, is_almost_symmetric
+from ..utils import augment_diagonal, import_graph, is_almost_symmetric
 from .svd import selectSVD
 
 
@@ -175,6 +175,7 @@ class BaseEmbedMulti(BaseEmbed):
         algorithm="randomized",
         n_iter=5,
         check_lcc=True,
+        diag_aug=True,
     ):
         super().__init__(
             n_components=n_components,
@@ -184,6 +185,10 @@ class BaseEmbedMulti(BaseEmbed):
             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 _check_input_graphs(self, graphs):
         """
         Checks if all graphs in list have same shapes.
@@ -238,3 +243,30 @@ class BaseEmbedMulti(BaseEmbed):
         self.n_vertices_ = out[0].shape[0]
 
         return out
+
+    def _diag_aug(self, graphs):
+        """
+        Augments the diagonal off each input graph. Returns the original
+        input object type.
+
+        Parameters
+        ----------
+        graphs : list of nx.Graph or ndarray, or ndarray
+            If list of nx.Graph, each Graph must contain same number of nodes.
+            If list of ndarray, each array must have shape (n_vertices, n_vertices).
+            If ndarray, then array must have shape (n_graphs, n_vertices, n_vertices).
+
+
+        Returns
+        -------
+        out : list of ndarray, or ndarray
+            If input is list of ndarray, then list is returned.
+            If input is ndarray, then ndarray is returned.
+        """
+        if isinstance(graphs, list):
+            out = [augment_diagonal(g) for g in graphs]
+        elif isinstance(graphs, np.ndarray):
+            # Copying is necessary to not overwrite input array
+            out = np.array([augment_diagonal(graphs[i]) for i in range(self.n_graphs_)])
+
+        return out
diff --git a/graspy/embed/mase.py b/graspy/embed/mase.py
index 44fa5aeb..e31aa16b 100644
--- a/graspy/embed/mase.py
+++ b/graspy/embed/mase.py
@@ -70,6 +70,11 @@ class MultipleASE(BaseEmbedMulti):
         Whether to scale individual eigenvectors with eigenvalues in first embedding
         stage.
 
+    diag_aug : bool, optional (default = True)
+        Whether to replace the main diagonal of each adjacency matrices with
+        a vector corresponding to the degree (or sum of edge weights for a
+        weighted network) before embedding.
+
     Attributes
     ----------
     n_graphs_ : int
@@ -101,6 +106,7 @@ class MultipleASE(BaseEmbedMulti):
         algorithm="randomized",
         n_iter=5,
         scaled=True,
+        diag_aug=True,
     ):
         if not isinstance(scaled, bool):
             msg = "scaled must be a boolean, not {}".format(scaled)
@@ -111,6 +117,7 @@ class MultipleASE(BaseEmbedMulti):
             n_elbows=n_elbows,
             algorithm=algorithm,
             n_iter=n_iter,
+            diag_aug=diag_aug,
         )
         self.scaled = scaled
 
@@ -200,6 +207,10 @@ class MultipleASE(BaseEmbedMulti):
         # Check if undirected
         undirected = all(is_almost_symmetric(g) for g in graphs)
 
+        # Diag augment
+        if self.diag_aug:
+            graphs = self._diag_aug(graphs)
+
         # embed
         Uhat, Vhat = self._reduce_dim(graphs)
         self.latent_left_ = Uhat
diff --git a/graspy/embed/omni.py b/graspy/embed/omni.py
index 995e60e7..706762c0 100644
--- a/graspy/embed/omni.py
+++ b/graspy/embed/omni.py
@@ -97,6 +97,11 @@ class OmnibusEmbed(BaseEmbedMulti):
         in non-optimal results if the average graph is unconnected. If True and average
         graph is unconnected, a UserWarning is thrown.
 
+    diag_aug : bool, optional (default = True)
+        Whether to replace the main diagonal of each adjacency matrices with
+        a vector corresponding to the degree (or sum of edge weights for a
+        weighted network) before embedding.
+
     Attributes
     ----------
     n_graphs_ : int
@@ -136,6 +141,7 @@ class OmnibusEmbed(BaseEmbedMulti):
         algorithm="randomized",
         n_iter=5,
         check_lcc=True,
+        diag_aug=True,
     ):
         super().__init__(
             n_components=n_components,
@@ -143,6 +149,7 @@ class OmnibusEmbed(BaseEmbedMulti):
             algorithm=algorithm,
             n_iter=n_iter,
             check_lcc=check_lcc,
+            diag_aug=diag_aug,
         )
 
     def fit(self, graphs, y=None):
@@ -173,6 +180,10 @@ class OmnibusEmbed(BaseEmbedMulti):
                 )
                 warnings.warn(msg, UserWarning)
 
+        # Diag augment
+        if self.diag_aug:
+            graphs = self._diag_aug(graphs)
+
         # Create omni matrix
         omni_matrix = _get_omni_matrix(graphs)
 
diff --git a/requirements.txt b/requirements.txt
index 55d70431..68b2d9af 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,7 @@
 networkx>=2.1
 numpy>=1.8.1
 scikit-learn>=0.19.1
-scipy>=1.1.0
+scipy>=1.4.0
 seaborn>=0.9.0
 matplotlib>=3.0.0,<=3.3.0
 hyppo>=0.1.2
diff --git a/setup.py b/setup.py
index 7378a02c..8f7696ef 100644
--- a/setup.py
+++ b/setup.py
@@ -15,7 +15,7 @@ REQUIRED_PACKAGES = [
     "networkx>=2.1",
     "numpy>=1.8.1",
     "scikit-learn>=0.19.1",
-    "scipy>=1.1.0",
+    "scipy>=1.4.0",
     "seaborn>=0.9.0",
     "matplotlib>=3.0.0",
     "hyppo>=0.1.3",


Test patch:
diff --git a/tests/test_mase.py b/tests/test_mase.py
index b5b1a9cb..6ea6fc69 100644
--- a/tests/test_mase.py
+++ b/tests/test_mase.py
@@ -1,9 +1,9 @@
 import numpy as np
 import pytest
-
 from graspy.cluster.gclust import GaussianCluster
 from graspy.embed.mase import MultipleASE
 from graspy.simulations.simulations import er_np, sbm
+from numpy import array_equal
 
 
 def make_train_undirected(n=[128, 128], m=10, alpha=1):
@@ -46,6 +46,10 @@ def test_bad_inputs():
         "Invalid unscaled"
         mase = MultipleASE(scaled="1")
 
+    with pytest.raises(TypeError):
+        wrong_diag_aug = "True"
+        mase = MultipleASE(diag_aug=wrong_diag_aug)
+
     with pytest.raises(ValueError):
         "Test single graph input"
         MultipleASE().fit(single_graph)
@@ -64,84 +68,82 @@ def test_bad_inputs():
         MultipleASE().fit(different_size_graphs)
 
 
+def test_diag_aug():
+    np.random.seed(5)
+    n = 100
+    p = 0.25
+
+    graphs_list = [er_np(n, p) for _ in range(2)]
+    graphs_arr = np.array(graphs_list)
+
+    # Test that array and list inputs results in same embeddings
+    mase_arr = MultipleASE(diag_aug=True).fit_transform(graphs_arr)
+    mase_list = MultipleASE(diag_aug=True).fit_transform(graphs_list)
+
+    assert array_equal(mase_list, mase_arr)
+
+
 def test_graph_clustering():
     """
     There should be 4 total clusters since 4 class problem.
     n_components = 2
     """
-    # undirected case
-    np.random.seed(2)
     n = [128, 128]
     m = 10
-    X = make_train_undirected(n, m)
-
-    res = MultipleASE(2).fit(X).scores_.reshape((m * 4, -1))
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 4
 
-    # directed case
-    np.random.seed(3)
-    X = make_train_directed(n, m)
+    def run(diag_aug, scaled):
+        # undirected case
+        np.random.seed(2 + diag_aug + scaled)
+        X = make_train_undirected(n, m)
 
-    res = MultipleASE(2).fit(X).scores_.reshape((m * 4, -1))
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 4
+        res = (
+            MultipleASE(2, diag_aug=diag_aug, scaled=scaled)
+            .fit(X)
+            .scores_.reshape((m * 4, -1))
+        )
+        gmm = GaussianCluster(10, covariance_type="all").fit(res)
+        assert gmm.n_components_ == 4
 
-    # Scaled cases
-    # undirected case
-    np.random.seed(12)
-    X = make_train_undirected(n, m)
+        # directed case
+        np.random.seed(3 + diag_aug + scaled)
+        X = make_train_directed(n, m)
 
-    res = MultipleASE(2, scaled=True).fit(X).scores_.reshape((m * 4, -1))
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 4
+        res = MultipleASE(2, diag_aug=diag_aug).fit(X).scores_.reshape((m * 4, -1))
+        gmm = GaussianCluster(10, covariance_type="all").fit(res)
+        assert gmm.n_components_ == 4
 
-    # directed case
-    np.random.seed(13)
-    X = make_train_directed(n, m)
-
-    res = MultipleASE(2, scaled=True).fit(X).scores_.reshape((m * 4, -1))
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 4
+    run(diag_aug=False, scaled=False)
+    run(diag_aug=True, scaled=False)
+    run(diag_aug=False, scaled=True)
+    run(diag_aug=True, scaled=True)
 
 
 def test_vertex():
     """
     There should be 2 clusters since each graph is a 2 block model
     """
-    # undirected case
-    np.random.seed(4)
     n = [128, 128]
     m = 10
-    X = make_train_undirected(n, m)
 
-    # undirected case
-    res = MultipleASE(n_components=2).fit(X).latent_left_
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 2
+    def run(diag_aug, scaled):
+        # undirected case
+        np.random.seed(4 + diag_aug + scaled)
+        X = make_train_undirected(n, m)
 
-    # directed case
-    np.random.seed(5)
-    X = make_train_directed(n, m)
+        res = MultipleASE(n_components=2).fit(X).latent_left_
+        gmm = GaussianCluster(10, covariance_type="all").fit(res)
+        assert gmm.n_components_ == 2
 
-    mase = MultipleASE(n_components=2).fit(X)
-    res = np.hstack([mase.latent_left_, mase.latent_right_])
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 2
+        # directed case
+        np.random.seed(5 + diag_aug + scaled)
+        X = make_train_directed(n, m)
 
-    # Scaled and undirected case
-    np.random.seed(4)
-    X = make_train_undirected(n, m)
-
-    res = MultipleASE(n_components=2, scaled=True).fit_transform(X)
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 2
-
-    # Scaled and directed case
-    np.random.seed(5)
-    X = make_train_directed(n, m)
+        mase = MultipleASE(n_components=2).fit(X)
+        res = np.hstack([mase.latent_left_, mase.latent_right_])
+        gmm = GaussianCluster(10, covariance_type="all").fit(res)
+        assert gmm.n_components_ == 2
 
-    left, right = MultipleASE(n_components=2, scaled=True).fit_transform(X)
-    res = np.hstack([left, right])
-    gmm = GaussianCluster(10, covariance_type="all").fit(res)
-    assert gmm.n_components_ == 2
+    run(diag_aug=False, scaled=False)
+    run(diag_aug=True, scaled=False)
+    run(diag_aug=False, scaled=True)
+    run(diag_aug=True, scaled=True)
diff --git a/tests/test_omni.py b/tests/test_omni.py
index 27e0a14a..15b46ee9 100755
--- a/tests/test_omni.py
+++ b/tests/test_omni.py
@@ -69,7 +69,11 @@ def test_omni_matrix_random():
     assert_allclose(A, expected_output)
 
 
-def test_omni_matrix_invalid_inputs():
+def test_invalid_inputs():
+    with pytest.raises(TypeError):
+        wrong_diag_aug = "True"
+        omni = OmnibusEmbed(diag_aug=wrong_diag_aug)
+
     with pytest.raises(ValueError):
         empty_list = []
         omni = OmnibusEmbed(n_components=2)
@@ -110,6 +114,21 @@ def test_omni_unconnected():
         omni.fit(graphs)
 
 
+def test_diag_aug():
+    np.random.seed(5)
+    n = 100
+    p = 0.25
+
+    graphs_list = [er_np(n, p) for _ in range(2)]
+    graphs_arr = np.array(graphs_list)
+
+    # Test that array and list inputs results in same embeddings
+    omni_arr = OmnibusEmbed(diag_aug=True).fit_transform(graphs_arr)
+    omni_list = OmnibusEmbed(diag_aug=True).fit_transform(graphs_list)
+
+    assert array_equal(omni_list, omni_arr)
+
+
 def test_omni_embed():
     """
     We compare the difference of norms of OmniBar and ABar.
@@ -121,15 +140,18 @@ def test_omni_embed():
         n = arr.shape[0] // 2
         return (arr[:n] + arr[n:]) / 2
 
-    X, A1, A2 = generate_data(1000)
-    Abar = (A1 + A2) / 2
+    def run(diag_aug):
+        X, A1, A2 = generate_data(1000, seed=2)
+        Abar = (A1 + A2) / 2
+
+        omni = OmnibusEmbed(n_components=3, diag_aug=diag_aug)
+        OmniBar = compute_bar(omni.fit_transform([A1, A2]))
 
-    np.random.seed(11)
-    omni = OmnibusEmbed(n_components=3)
-    OmniBar = compute_bar(omni.fit_transform([A1, A2]))
+        omni = OmnibusEmbed(n_components=3, diag_aug=diag_aug)
+        ABar = compute_bar(omni.fit_transform([Abar, Abar]))
 
-    omni = OmnibusEmbed(n_components=3)
-    ABar = compute_bar(omni.fit_transform([Abar, Abar]))
+        tol = 1.0e-2
+        assert allclose(norm(OmniBar, axis=1), norm(ABar, axis=1), rtol=tol, atol=tol)
 
-    tol = 1.0e-2
-    assert allclose(norm(OmniBar, axis=1), norm(ABar, axis=1), rtol=tol, atol=tol)
+    run(diag_aug=True)
+    run(diag_aug=False)