Repository: graspologic-org/graspologic

Issue:
[BUG] Using layout_umap and getting key error
## Expected Behavior
a graph layout using UMAP for dim reduction to 2d

## Actual Behavior
coming from layouts\auto.py
 KeyError: '2'

## Example Code
Please see [How to create a Minimal, Reproducible example](https://stackoverflow.com/help/minimal-reproducible-example) for some guidance on creating the best possible example of the problem
```python
import networkx as nx
from graspologic.layouts import layout_umap

g = nx.erdos_renyi_graph(10, 0.7, directed=False)
layout_umap(g)
```

## Full Traceback
```pytb
Paste the full traceback in case there is an exception

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-17-1c062a6ca19a> in <module>
      3 
      4 g = nx.erdos_renyi_graph(10, 0.7, directed=False)
----> 5 layout_umap(g)

~\miniconda3\lib\site-packages\graspologic\layouts\auto.py in layout_umap(graph, min_dist, n_neighbors, max_edges, random_seed)
    177         min_dist=min_dist, n_neighbors=n_neighbors, random_state=random_seed
    178     ).fit_transform(tensors)
--> 179     positions = _node_positions_from(lcc_graph, labels, points, random_seed=random_seed)
    180     return lcc_graph, positions
    181 

~\miniconda3\lib\site-packages\graspologic\layouts\auto.py in _node_positions_from(graph, labels, down_projection_2d, random_seed)
    241     scaled_points = _scale_points(down_projection_2d, covered_area)
    242     partitions = leiden(graph, random_seed=random_seed)
--> 243     positions = [
    244         NodePosition(
    245             node_id=key,

~\miniconda3\lib\site-packages\graspologic\layouts\auto.py in <listcomp>(.0)
    246             x=scaled_points[index][0],
    247             y=scaled_points[index][1],
--> 248             size=sizes[key],
    249             community=partitions[key],
    250         )

KeyError: '2'

```

## Your Environment
* Python version:  3.8
* graspologic version:





Discussion:
The core problem going on here is a bit complicated but also very simple at the same time - the TLDR is that we need to ensure we are always treating our keys as str(key) when we do these lookups because `leiden` is 100% of the time going to give you the string representation of your node back as the ID - which is what gets put into the partitions dictionary.  I had no interest trying to get Python's "pass any type of node id you want of any type" thing working in Rust so I just forced everything to encode as a string, which is great when the calling graph's nodes are strings too, but apparently has some issues in auto.py where it's using the original node ID (as an `int`, in this case) as the lookup to generate our `size` dictionary, but using str(node_id) as the key it's looking up with - which isn't going to appear in the sizes dict.

We just need to be careful that any dicts we create in this process are based on a `str(node_id)` or else we'll get things like this.

There's a similar typing issue in `leiden` and `modularity` where we implicitly cast your provided graph as `str(node_id)` for you to be nice and helpful when you call leiden, return `str(node_id) -> partition_id` mappings from leiden, and then when you try to call modularity on it to see how it turned out, it explodes because the graph you provide has `int(node_id)` keys and you're trying to look up `str(node_id)` mappings.

All of this is to say that maybe the right fix is for the python wrapper to our rust native impl should capture the node type in Python, let it cast to string for Rust and Leiden to process it, then before we hand the partition dictionary back we re-cast back to the original type.  And if we purposefully kept the original -> new mapping in memory in Python, we don't even have to have a lossy "cast back", since we'd just replace the new string representation of it back with the original as it was.

And that's the short explanation.
> All of this is to say that maybe the right fix is for the python wrapper to our rust native impl should capture the node type in Python, let it cast to string for Rust and Leiden to process it, then before we hand the partition dictionary back we re-cast back to the original type. And if we purposefully kept the original -> new mapping in memory in Python, we don't even have to have a lossy "cast back", since we'd just replace the new string representation of it back with the original as it was.

Big +1 for this, since I definitely think we should allow int or other datatypes as node keys, and this seems like the cleanest/simplest way to allow this.
It would also fix up our modularity bug calculation too, because the types would be what you'd have expected coming back.  So that's probably the right approach
@bdpedigo blech - we generate the edge list that leiden requires in `utils.py` where we generate a `List[Tuple[str, str, float]]`.  That's reasonable as a utility, but I think a `Tuple[Dict[str, Any], List[Tuple[str, str, float]]]` is not really something the common person is going to want - but it is what I need to do to fix this.

So the question is whether I can just axe this from utils or not.  There's benefits to having a "given a graph give me a basic edge list" but for all I know there's like 15 different versions of that behavior already.  So I figured I'd bring it up with you and see if you have a preference for what we do; axe it from utils or give everyone the `Tuple[Dict[str, Any], List[Tuple[str, str, float]]]` return type, which... doesn't seem super useful.  Or, I guess, replicate the function and have a private leiden specific version that does that.
hmmm i guess i dont see why the conversion to edgelist needs to be exposed to the user? am i missing something?

> Or, I guess, replicate the function and have a private leiden specific version that does that.

Is this you suggesting we don't need the user to see/care about what leiden is really eating under the hood? If so I agree with that.
We only exposed it before because having a quick way to write out an edge list regardless of what your current graph is was useful for us - sure it only saves like 2 lines of code, but a penny saved is a penny earned or something.  I don't know.

Regardless, it's currently in utils, and changing it to the new way is going to make it less obviously useful, so I'd rather remove it from utils entirely rather than say "here's this wacky function that almost nobody in the world is going to care about now but I guess you can use it to turn a graph in csr matrix form into an edge list if you want as long as you throw half of the top level tuple away" just doesn't sound particularly useful
totally fine with removing it from utils, personally!

Resolving patch:
diff --git a/graspologic/embed/n2v.py b/graspologic/embed/n2v.py
--- a/graspologic/embed/n2v.py
+++ b/graspologic/embed/n2v.py
@@ -9,6 +9,8 @@
 import networkx as nx
 import numpy as np
 
+from ..utils import remap_node_ids
+
 
 def node2vec_embed(
     graph: Union[nx.Graph, nx.DiGraph],
@@ -22,7 +24,7 @@ def node2vec_embed(
     iterations: int = 1,
     interpolate_walk_lengths_by_node_degree: bool = True,
     random_seed: Optional[int] = None,
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> Tuple[np.array, List[Any]]:
     """
     Generates a node2vec embedding from a given graph. Will follow the word2vec algorithm to create the embedding.
 
@@ -65,7 +67,7 @@ def node2vec_embed(
 
     Returns
     -------
-    Tuple[np.ndarray, np.ndarray]
+    Tuple[np.array, List[Any]]
         A tuple containing a matrix, with each row index corresponding to the embedding for each node. The tuple
         also contains a vector containing the corresponding vertex labels for each row in the matrix.
         The matrix and vector are positionally correlated.
@@ -128,7 +130,13 @@ def node2vec_embed(
         f"Completed. Ending time is {str(end)} Elapsed time is {str(start - end)}"
     )
 
-    return model.wv.vectors, model.wv.index2word
+    labels = node2vec_graph.original_graph.nodes()
+    remapped_labels = node2vec_graph.label_map_to_string
+
+    return (
+        np.array([model.wv.get_vector(remapped_labels[node]) for node in labels]),
+        labels,
+    )
 
 
 def _assert_is_positive_int(name: str, value: int):
@@ -233,7 +241,13 @@ def __init__(
         inout_hyperparameter: float,
         random_state: Optional[np.random.RandomState] = None,
     ):
-        self.graph: nx.Graph = graph
+        self.original_graph: nx.Graph = graph
+
+        graph_with_new_ids, new_id_map = remap_node_ids(graph=graph)
+
+        self.graph = graph_with_new_ids
+        self.label_map_to_string = new_id_map
+
         self.is_directed = self.graph.is_directed()
         self.p = return_hyperparameter
         self.q = inout_hyperparameter
@@ -299,7 +313,7 @@ def node2vec_walk(
 
     @staticmethod
     def _get_walk_length_interpolated(
-        degree: int, percentiles: np.ndarray, max_walk_length: int
+        degree: int, percentiles: list, max_walk_length: int
     ):
         """
         Given a node's degree, determine the length of a walk that should be used. If the degree is less than the
@@ -473,7 +487,7 @@ def _alias_setup(probabilities: List[float]):
     """
     number_of_outcomes = len(probabilities)
     alias = np.zeros(number_of_outcomes)
-    sampled_probabilities = np.zeros(number_of_outcomes, dtype=np.int)
+    sampled_probabilities = np.zeros(number_of_outcomes, dtype=int)
 
     smaller = []
     larger = []
diff --git a/graspologic/utils/__init__.py b/graspologic/utils/__init__.py
--- a/graspologic/utils/__init__.py
+++ b/graspologic/utils/__init__.py
@@ -21,6 +21,7 @@
     fit_plug_in_variance_estimator,
     remove_vertices,
     remap_labels,
+    remap_node_ids,
 )
 from .ptr import pass_to_ranks
 
@@ -45,5 +46,5 @@
     "fit_plug_in_variance_estimator",
     "remove_vertices",
     "remap_labels",
-    "to_weighted_edge_list",
+    "remap_node_ids",
 ]
diff --git a/graspologic/utils/utils.py b/graspologic/utils/utils.py
--- a/graspologic/utils/utils.py
+++ b/graspologic/utils/utils.py
@@ -5,7 +5,7 @@
 from collections import Iterable
 from functools import reduce
 from pathlib import Path
-from typing import Any, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Tuple, Union
 
 import networkx as nx
 import numpy as np
@@ -944,6 +944,52 @@ def remap_labels(
         return remapped_y_pred
 
 
+def remap_node_ids(
+    graph: nx.Graph, weight_attribute: str = "weight"
+) -> Tuple[nx.Graph, Dict[Any, str]]:
+    """
+    Given a graph with arbitrarily types node ids, return a new graph that contains the exact same edgelist
+    except the node ids are remapped to a string representation.
+
+    Parameters
+    ----------
+    graph : nx.Graph
+        A graph that has node ids of arbitrary types.
+    weight_attribute : str,
+        Default is ``weight``. An optional attribute to specify which column in your graph contains the weight value.
+
+    Returns
+    -------
+    Tuple[nx.Graph, Dict[Any, str]]
+        A new graph that contains the same edges except the node ids are remapped to strings. The keys in
+        the dictionary are the old node ids and the values are the newly remapped node ids.
+
+    Raises
+    ------
+    TypeError
+    """
+    if not isinstance(graph, nx.Graph):
+        raise TypeError("graph must be of type nx.Graph")
+
+    node_id_dict = dict()
+    graph_remapped = type(graph)()
+
+    for source, target, weight in graph.edges(data=weight_attribute):
+        if source not in node_id_dict:
+            node_id_dict[source] = str(len(node_id_dict.keys()))
+
+        if target not in node_id_dict:
+            node_id_dict[target] = str(len(node_id_dict.keys()))
+
+        graph_remapped.add_edge(node_id_dict[source], node_id_dict[target])
+
+        graph_remapped[node_id_dict[source]][node_id_dict[target]][
+            weight_attribute
+        ] = weight
+
+    return graph_remapped, node_id_dict
+
+
 def suppress_common_warnings():
     """
     Suppresses common warnings that occur when using graspologic.


Test patch:
diff --git a/tests/embed/test_n2v.py b/tests/embed/test_n2v.py
--- a/tests/embed/test_n2v.py
+++ b/tests/embed/test_n2v.py
@@ -1,25 +1,158 @@
 # Copyright (c) Microsoft Corporation and contributors.
 # Licensed under the MIT License.
 
-import unittest
+import io
 import networkx as nx
 import numpy as np
-from graspologic.embed import n2v
+import unittest
+
+from sklearn.metrics import adjusted_rand_score
+from sklearn.cluster import KMeans
+
+import graspologic as gc
+
+
+class Node2VecEmbedTest(unittest.TestCase):
+    def test_n2v_returns_same_labels_with_different_nodeid_types(self):
+        probability_matrix = np.array([[0.95, 0.01], [0.01, 0.95]])
+        number_of_nodes_per_community = [20, 20]
+
+        sbm_sample = gc.simulations.sbm(
+            number_of_nodes_per_community, probability_matrix
+        )
+        sbm_graph = nx.from_numpy_array(sbm_sample)
+
+        graph = nx.Graph()
+        graph_as_strings = nx.Graph()
+        for s, t in sbm_graph.edges():
+            graph.add_edge(s, t, weight=1)
+            graph_as_strings.add_edge(str(s), str(t), weight=1)
+
+        original_embedding = gc.embed.node2vec_embed(graph, random_seed=1)
+        string_embedding = gc.embed.node2vec_embed(graph_as_strings, random_seed=1)
+
+        k = KMeans(n_clusters=2)
+        original_labels = k.fit_predict(original_embedding[0])
+        string_labels = k.fit_predict(string_embedding[0])
+
+        expected_labels = np.zeros(40, dtype=int)
+        expected_labels[20:] = 1
+
+        original_ari = adjusted_rand_score(original_labels, expected_labels)
+        string_ari = adjusted_rand_score(string_labels, expected_labels)
+
+        self.assertEqual(original_ari, string_ari)
+
+    def test_n2v_directed_undirected_returns_same_clustering(self):
+        probability_matrix = np.array([[0.95, 0.01], [0.01, 0.95]])
+        number_of_nodes_per_community = [20, 20]
+
+        sbm_sample = gc.simulations.sbm(
+            number_of_nodes_per_community, probability_matrix
+        )
+        sbm_graph = nx.from_numpy_array(sbm_sample)
+
+        graph = nx.Graph()
+        graph_directed = nx.DiGraph()
+        for s, t in sbm_graph.edges():
+            graph.add_edge(s, t, weight=1)
+
+            graph_directed.add_edge(s, t, weight=1)
+            graph_directed.add_edge(t, s, weight=1)
+
+        undirected_embedding = gc.embed.node2vec_embed(graph, random_seed=1)
+        directed_embedding = gc.embed.node2vec_embed(graph_directed, random_seed=1)
 
+        k = KMeans(n_clusters=2)
+        undirected_labels = k.fit_predict(undirected_embedding[0])
+        directed_labels = k.fit_predict(directed_embedding[0])
+
+        expected_labels = np.zeros(40, dtype=int)
+        expected_labels[20:] = 1
+
+        undirected_ari = adjusted_rand_score(undirected_labels, expected_labels)
+        directed_ari = adjusted_rand_score(directed_labels, expected_labels)
+
+        self.assertEqual(undirected_ari, directed_ari)
 
-class TestN2V(unittest.TestCase):
     def test_node2vec_embed(self):
         g = nx.florentine_families_graph()
 
         for s, t in g.edges():
             g.add_edge(s, t, weight=1)
 
-        embedding = n2v.node2vec_embed(g, random_seed=1)
+        embedding = gc.embed.node2vec_embed(g, random_seed=1)
 
-        embedding2 = n2v.node2vec_embed(g, random_seed=1)
+        embedding2 = gc.embed.node2vec_embed(g, random_seed=1)
 
         np.testing.assert_array_equal(embedding[0], embedding2[0])
 
+    def test_node2vec_embedding_florentine_graph_correct_shape_is_returned(self):
+        graph = nx.florentine_families_graph()
+        for s, t in graph.edges():
+            graph.add_edge(s, t, weight=1)
+
+        model = gc.embed.node2vec_embed(graph)
+        model_matrix: np.ndarray = model[0]
+        vocab_list = model[1]
+        self.assertIsNotNone(model)
+        self.assertIsNotNone(model[0])
+        self.assertIsNotNone(model[1])
+
+        # model matrix should be 34 x 128
+        self.assertEqual(model_matrix.shape[0], 15)
+        self.assertEqual(model_matrix.shape[1], 128)
+
+        # vocab list should have exactly 34 elements
+        self.assertEqual(len(vocab_list), 15)
+
+    def test_node2vec_embedding_barbell_graph_correct_shape_is_returned(self):
+        graph = nx.barbell_graph(25, 2)
+        for s, t in graph.edges():
+            graph.add_edge(s, t, weight=1)
+
+        model = gc.embed.node2vec_embed(graph)
+        model_matrix: np.ndarray = model[0]
+        vocab_list = model[1]
+        self.assertIsNotNone(model)
+        self.assertIsNotNone(model[0])
+        self.assertIsNotNone(model[1])
+
+        # model matrix should be 34 x 128
+        self.assertEqual(model_matrix.shape[0], 52)
+        self.assertEqual(model_matrix.shape[1], 128)
+
+        # vocab list should have exactly 34 elements
+        self.assertEqual(len(vocab_list), 52)
+
+    def test_get_walk_length_lower_defaults_to_1(self):
+        expected_walk_length = 1
+
+        g = gc.embed.n2v._Node2VecGraph(nx.Graph(), 1, 1)
+        w = g._get_walk_length_interpolated(
+            degree=0, percentiles=[1, 2, 3, 4, 10, 100], max_walk_length=10
+        )
+
+        self.assertEqual(w, expected_walk_length)
+
+    def test_get_walk_length_higher_default_to_walk_length(self):
+        expected_walk_length = 100
+
+        g = gc.embed.n2v._Node2VecGraph(nx.Graph(), 1, 1)
+        w = g._get_walk_length_interpolated(
+            degree=10,
+            percentiles=[2, 3, 4, 5, 6, 7, 8, 9],
+            max_walk_length=expected_walk_length,
+        )
+
+        self.assertEqual(w, expected_walk_length)
+
+    def test_get_walk_length_in_middle_selects_interpolated_bucket(self):
+        expected_walk_length = 5
+
+        g = gc.embed.n2v._Node2VecGraph(nx.Graph(), 1, 1)
+        w = g._get_walk_length_interpolated(
+            degree=5, percentiles=[2, 3, 4, 5, 6, 7, 8, 9], max_walk_length=10
+        )
 
-if __name__ == "__main__":
-    unittest.main()
+        self.assertEqual(w, expected_walk_length)
diff --git a/tests/layouts/test_auto.py b/tests/layouts/test_auto.py
--- a/tests/layouts/test_auto.py
+++ b/tests/layouts/test_auto.py
@@ -2,8 +2,10 @@
 # Licensed under the MIT License.
 
 import unittest
+import networkx as nx
 import numpy
-from graspologic.layouts.auto import _get_bounds
+
+from graspologic.layouts.auto import _get_bounds, layout_umap
 
 
 class TestAuto(unittest.TestCase):
@@ -15,6 +17,34 @@ def test_get_bounds(self):
         self.assertEqual(10, maxx)
         self.assertEqual(5, maxy)
 
+    def test_layout_umap_string_node_ids(self):
+        graph = nx.florentine_families_graph()
+
+        for s, t in graph.edges():
+            graph.add_edge(s, t, weight=1)
+
+        _, node_positions = layout_umap(graph=graph)
+
+        self.assertEqual(len(node_positions), len(graph.nodes()))
+
+    def test_layout_umap_int_node_ids(self):
+        graph = nx.florentine_families_graph()
+        graph_int_node_ids = nx.Graph()
+        ids_as_ints = dict()
+
+        for s, t in graph.edges():
+            if s not in ids_as_ints:
+                ids_as_ints[s] = int(len(ids_as_ints.keys()))
+
+            if t not in ids_as_ints:
+                ids_as_ints[t] = int(len(ids_as_ints.keys()))
+
+            graph_int_node_ids.add_edge(ids_as_ints[s], ids_as_ints[t], weight=1)
+
+        _, node_positions = layout_umap(graph=graph_int_node_ids)
+
+        self.assertEqual(len(node_positions), len(graph.nodes()))
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_n2v.py b/tests/test_n2v.py
deleted file mode 100644
--- a/tests/test_n2v.py
+++ /dev/null
@@ -1,183 +0,0 @@
-# Copyright (c) Microsoft Corporation and contributors.
-# Licensed under the MIT License.
-
-import unittest
-import networkx as nx
-import numpy as np
-import graspologic.embed.n2v as n2v
-from graspologic.embed.n2v import _Node2VecGraph
-
-
-class Node2VecEmbedTest(unittest.TestCase):
-    def test_node2vec_embedding_correct_shape_is_returned(self):
-        import io
-
-        graph = nx.read_edgelist(
-            io.StringIO(_edge_list), nodetype=int, create_using=nx.DiGraph()
-        )
-
-        model = n2v.node2vec_embed(graph)
-        model_matrix: np.ndarray = model[0]
-        vocab_list = model[1]
-        self.assertIsNotNone(model)
-        self.assertIsNotNone(model[0])
-        self.assertIsNotNone(model[1])
-
-        # model matrix should be 34 x 128
-        self.assertEqual(model_matrix.shape[0], 34)
-        self.assertEqual(model_matrix.shape[1], 128)
-
-        # vocab list should have exactly 34 elements
-        self.assertEqual(len(vocab_list), 34)
-
-    def test_node2vec_embedding_florentine_graph_correct_shape_is_returned(self):
-        graph = nx.florentine_families_graph()
-        for s, t in graph.edges():
-            graph.add_edge(s, t, weight=1)
-
-        model = n2v.node2vec_embed(graph)
-        model_matrix: np.ndarray = model[0]
-        vocab_list = model[1]
-        self.assertIsNotNone(model)
-        self.assertIsNotNone(model[0])
-        self.assertIsNotNone(model[1])
-
-        # model matrix should be 34 x 128
-        self.assertEqual(model_matrix.shape[0], 15)
-        self.assertEqual(model_matrix.shape[1], 128)
-
-        # vocab list should have exactly 34 elements
-        self.assertEqual(len(vocab_list), 15)
-
-    def test_node2vec_embedding_barbell_graph_correct_shape_is_returned(self):
-        graph = nx.barbell_graph(25, 2)
-        for s, t in graph.edges():
-            graph.add_edge(s, t, weight=1)
-
-        model = n2v.node2vec_embed(graph)
-        model_matrix: np.ndarray = model[0]
-        vocab_list = model[1]
-        self.assertIsNotNone(model)
-        self.assertIsNotNone(model[0])
-        self.assertIsNotNone(model[1])
-
-        # model matrix should be 34 x 128
-        self.assertEqual(model_matrix.shape[0], 52)
-        self.assertEqual(model_matrix.shape[1], 128)
-
-        # vocab list should have exactly 34 elements
-        self.assertEqual(len(vocab_list), 52)
-
-    def test_get_walk_length_lower_defaults_to_1(self):
-        expected_walk_length = 1
-
-        g = _Node2VecGraph(nx.Graph(), 1, 1)
-        w = g._get_walk_length_interpolated(
-            degree=0, percentiles=[1, 2, 3, 4, 10, 100], max_walk_length=10
-        )
-
-        self.assertEqual(w, expected_walk_length)
-
-    def test_get_walk_length_higher_default_to_walk_length(self):
-        expected_walk_length = 100
-
-        g = _Node2VecGraph(nx.Graph(), 1, 1)
-        w = g._get_walk_length_interpolated(
-            degree=10,
-            percentiles=[2, 3, 4, 5, 6, 7, 8, 9],
-            max_walk_length=expected_walk_length,
-        )
-
-        self.assertEqual(w, expected_walk_length)
-
-    def test_get_walk_length_in_middle_selects_interpolated_bucket(self):
-        expected_walk_length = 5
-
-        g = _Node2VecGraph(nx.Graph(), 1, 1)
-        w = g._get_walk_length_interpolated(
-            degree=5, percentiles=[2, 3, 4, 5, 6, 7, 8, 9], max_walk_length=10
-        )
-
-        self.assertEqual(w, expected_walk_length)
-
-
-_edge_list = """
-1 32
-1 22
-1 20
-1 18
-1 14
-1 13
-1 12
-1 11
-1 9
-1 8
-1 7
-1 6
-1 5
-1 4
-1 3
-1 2
-2 31
-2 22
-2 20
-2 18
-2 14
-2 8
-2 4
-2 3
-3 14
-3 9
-3 10
-3 33
-3 29
-3 28
-3 8
-3 4
-4 14
-4 13
-4 8
-5 11
-5 7
-6 17
-6 11
-6 7
-7 17
-9 34
-9 33
-9 33
-10 34
-14 34
-15 34
-15 33
-16 34
-16 33
-19 34
-19 33
-20 34
-21 34
-21 33
-23 34
-23 33
-24 30
-24 34
-24 33
-24 28
-24 26
-25 32
-25 28
-25 26
-26 32
-27 34
-27 30
-28 34
-29 34
-29 32
-30 34
-30 33
-31 34
-31 33
-32 34
-32 33
-33 34
-"""
diff --git a/tests/test_utils.py b/tests/test_utils.py
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -547,3 +547,44 @@ def test_inputs(self):
 
         with pytest.raises(ValueError):
             remap_labels(self.y_pred, ["ant", "ant", "cat", "cat", "bird", "bird"])
+
+
+class TestRemapNodeIds(unittest.TestCase):
+    def test_remap_node_ids_invalid_typ_raises_typeerror(self):
+        invalid_types = [str, int, list]
+
+        for type in invalid_types:
+            with pytest.raises(TypeError):
+                gus.remap_node_ids(graph=type())
+
+    def _assert_graphs_are_equivalent(self, graph, new_graph, new_node_ids):
+        self.assertTrue(len(new_graph.nodes()) == len(graph.nodes()))
+        self.assertTrue(len(new_graph.edges()) == len(graph.edges()))
+
+        for source, target, weight in graph.edges(data="weight"):
+            self.assertTrue(source in new_node_ids.keys())
+            self.assertTrue(target in new_node_ids.keys())
+
+            new_weight = new_graph[new_node_ids[source]][new_node_ids[target]]["weight"]
+            self.assertEqual(weight, new_weight)
+
+    def test_remap_node_ids_graph_has_same_edges_but_remapped(self):
+        graph = nx.Graph()
+
+        graph.add_edge(0, 1, weight=10)
+        graph.add_edge(1, "someid", weight=100)
+
+        new_graph, new_node_ids = gus.remap_node_ids(graph)
+
+        self._assert_graphs_are_equivalent(graph, new_graph, new_node_ids)
+
+    def test_remap_node_ids_digraph_has_same_edges_but_remapped(self):
+        graph = nx.DiGraph()
+
+        graph.add_edge(0, 1, weight=10)
+        graph.add_edge(1, "someid", weight=100)
+        graph.add_edge("someid", 1, weight=21)
+
+        new_graph, new_node_ids = gus.remap_node_ids(graph)
+
+        self._assert_graphs_are_equivalent(graph, new_graph, new_node_ids)