diff --git a/docs/conf.py b/docs/conf.py
index c9bb725..5d669ad 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -1,5 +1,4 @@
 #!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 #
 # queueuing-tool documentation build configuration file, created by
 # sphinx-quickstart on Sun Feb  1 14:13:20 2015.
@@ -13,27 +12,25 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
-import sys
 import os
+import sys
 
 try:
-    import unittest.mock as mock
+    from unittest import mock
 except ImportError:
-    import mock
+    from unittest import mock
 
 import alabaster
-import numpydoc
-
 
 MOCK_MODULES = [
-    'choice',
-    'queueing_tool.queues.choice',
-    'matplotlib',
-    'numpy',
-    'numpy.random',
-    'priority_queue',
-    'queueing_tool.network.priority_queue',
-    'pygraphviz'
+    "choice",
+    "queueing_tool.queues.choice",
+    "matplotlib",
+    "numpy",
+    "numpy.random",
+    "priority_queue",
+    "queueing_tool.network.priority_queue",
+    "pygraphviz",
 ]
 sys.modules.update((mod_name, mock.Mock()) for mod_name in MOCK_MODULES)
 
@@ -41,58 +38,58 @@
 # If extensions (or modules to document with autodoc) are in another directory,
 # add these directories to sys.path here. If the directory is relative to the
 # documentation root, use os.path.abspath to make it absolute, like shown here.
-sys.path.insert(0, os.path.abspath('..'))
+sys.path.insert(0, os.path.abspath(".."))
 
 # -- General configuration ------------------------------------------------
 
 # If your documentation needs a minimal Sphinx version, state it here.
-#needs_sphinx = '1.0'
+# needs_sphinx = '1.0'
 
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
 extensions = [
-    'sphinx.ext.autodoc',
-    'sphinx.ext.autosummary',
-    'sphinx.ext.doctest',
-    'sphinx.ext.extlinks',
-    'sphinx.ext.intersphinx',
-    'sphinx.ext.mathjax',
-    'sphinx.ext.todo',
-    'sphinx.ext.viewcode',
-    'numpydoc',
-    'alabaster',
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.doctest",
+    "sphinx.ext.extlinks",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.todo",
+    "sphinx.ext.viewcode",
+    "numpydoc",
+    "alabaster",
 ]
 
 intersphinx_mapping = {
-    'python': ('http://docs.python.org/3', None),
-    'matplotlib': ('http://matplotlib.sourceforge.net', None),
-    'numpy': ('http://docs.scipy.org/doc/numpy', None),
-    'networkx': ('http://networkx.readthedocs.org/en/networkx-1.11/', None)
+    "python": ("http://docs.python.org/3", None),
+    "matplotlib": ("http://matplotlib.sourceforge.net", None),
+    "numpy": ("http://docs.scipy.org/doc/numpy", None),
+    "networkx": ("http://networkx.readthedocs.org/en/networkx-1.11/", None),
 }
 
 extlinks = {
-    'doi': ('http://dx.doi.org/%s', 'DOI: '),
-    'arxiv': ('http://arxiv.org/abs/%s', 'arXiv: ')
+    "doi": ("http://dx.doi.org/%s", "DOI: "),
+    "arxiv": ("http://arxiv.org/abs/%s", "arXiv: "),
 }
 
 numpydoc_show_class_members = False
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # The suffix of source filenames.
-source_suffix = '.rst'
+source_suffix = ".rst"
 
 # The encoding of source files.
-#source_encoding = 'utf-8-sig'
+# source_encoding = 'utf-8-sig'
 
 # The master toctree document.
-master_doc = 'index'
+master_doc = "index"
 
 # General information about the project.
-project = 'Queueing-tool'
-copyright = '2024, Daniel Jordon'
+project = "Queueing-tool"
+copyright = "2024, Daniel Jordon"
 
 # The version info for the project you're documenting, acts as replacement for
 # |version| and |release|, also used in various other places throughout the
@@ -102,6 +99,7 @@
 
 # The short X.Y version.
 import queueing_tool as qt
+
 version = qt.__version__
 
 # The full version, including alpha/beta/rc tags.
@@ -109,137 +107,140 @@
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
-exclude_patterns = ['_build']
+exclude_patterns = ["_build"]
 
 # If true, '()' will be appended to :func: etc. cross-reference text.
-#add_function_parentheses = True
+# add_function_parentheses = True
 
 # If true, the current module name will be prepended to all description
 # unit titles (such as .. function::).
-#add_module_names = True
+# add_module_names = True
 
 # The name of the Pygments (syntax highlighting) style to use.
-pygments_style = 'sphinx'
+pygments_style = "sphinx"
 
 # A list of ignored prefixes for module index sorting.
-#modindex_common_prefix = []
+# modindex_common_prefix = []
 
 # -- Options for HTML output ----------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
 # a list of builtin themes. http://sphinx-doc.org/theming.html
-html_theme = 'alabaster'  # 'qt_sphinx13'#
-#html_style = 'adjust_alabaster.css'
+html_theme = "alabaster"  # 'qt_sphinx13'#
+# html_style = 'adjust_alabaster.css'
 # Theme options are theme-specific and customize the look and feel of a theme
 # further.  For a list of options available for each theme, see the
 # documentation.
 html_theme_options = {
-    'code_font_size': '0.81em',
-    'github_user': 'djordon',
-    'github_repo': 'queueing-tool',
+    "code_font_size": "0.81em",
+    "github_user": "djordon",
+    "github_repo": "queueing-tool",
 }
-#html_context = {'css_files': ['_static/adjust_alabaster.css',]}
+# html_context = {'css_files': ['_static/adjust_alabaster.css',]}
 
 # Add any paths that contain custom themes here, relative to this directory.
-html_theme_path = [alabaster.get_path()]#'_themes']
+html_theme_path = [alabaster.get_path()]  #'_themes']
 
 html_sidebars = {
-    '**': [
-        'about.html',
-        'navigation.html',
-        'relations.html',
-        'searchbox.html',
-        'donate.html',
+    "**": [
+        "about.html",
+        "navigation.html",
+        "relations.html",
+        "searchbox.html",
+        "donate.html",
     ]
 }
 
 # The name for this set of Sphinx documents.  If None, it defaults to
 # "<project> v<release> documentation".
-#html_title = None
+# html_title = None
 
 # A shorter title for the navigation bar.  Default is the same as html_title.
-#html_short_title = None
+# html_short_title = None
 
 # The name of an image file (relative to this directory) to place at the top
 # of the sidebar.
-#html_logo = None
+# html_logo = None
 
 # The name of an image file (within the static path) to use as favicon of the
 # docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
 # pixels large.
-#html_favicon = None
+# html_favicon = None
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ['_static']
+html_static_path = ["_static"]
 
 # Add any extra paths that contain custom files (such as robots.txt or
 # .htaccess) here, relative to this directory. These files are copied
 # directly to the root of the documentation.
-#html_extra_path = []
+# html_extra_path = []
 
 # If true, links to the reST sources are added to the pages.
-#html_show_sourcelink = True
+# html_show_sourcelink = True
 
 # If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
-#html_show_sphinx = True
+# html_show_sphinx = True
 
 # If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
-#html_show_copyright = True
+# html_show_copyright = True
 
 # If true, an OpenSearch description file will be output, and all pages will
 # contain a <link> tag referring to it.  The value of this option must be the
 # base URL from which the finished HTML is served.
-#html_use_opensearch = ''
+# html_use_opensearch = ''
 
 # This is the file name suffix for HTML files (e.g. ".xhtml").
-#html_file_suffix = None
+# html_file_suffix = None
 
 # Output file base name for HTML help builder.
-htmlhelp_basename = 'queueuing-tooldoc'
+htmlhelp_basename = "queueuing-tooldoc"
 
 
 # -- Options for LaTeX output ---------------------------------------------
 
 latex_elements = {
-# The paper size ('letterpaper' or 'a4paper').
-#'papersize': 'letterpaper',
-
-# The font size ('10pt', '11pt' or '12pt').
-#'pointsize': '10pt',
-
-# Additional stuff for the LaTeX preamble.
-#'preamble': '',
+    # The paper size ('letterpaper' or 'a4paper').
+    #'papersize': 'letterpaper',
+    # The font size ('10pt', '11pt' or '12pt').
+    #'pointsize': '10pt',
+    # Additional stuff for the LaTeX preamble.
+    #'preamble': '',
 }
 
 # Grouping the document tree into LaTeX files. List of tuples
 # (source start file, target name, title,
 #  author, documentclass [howto, manual, or own class]).
 latex_documents = [
-  ('index', 'queueuing-tool.tex', 'queueuing-tool Documentation',
-   'Daniel Jordon', 'manual'),
+    (
+        "index",
+        "queueuing-tool.tex",
+        "queueuing-tool Documentation",
+        "Daniel Jordon",
+        "manual",
+    ),
 ]
 
 # The name of an image file (relative to this directory) to place at the top of
 # the title page.
-#latex_logo = None
+# latex_logo = None
 
 # For "manual" documents, if this is true, then toplevel headings are parts,
 # not chapters.
-#latex_use_parts = False
+# latex_use_parts = False
 
 # If true, show page references after internal links.
-#latex_show_pagerefs = False
+# latex_show_pagerefs = False
 
 # If true, show URL addresses after external links.
-#latex_show_urls = False
+# latex_show_urls = False
 
 # Documents to append as an appendix to all manuals.
-#latex_appendices = []
+# latex_appendices = []
 
 # If false, no module index is generated.
-#latex_domain_indices = True
+# latex_domain_indices = True
 
 
 # -- Options for manual page output ---------------------------------------
@@ -247,12 +248,11 @@
 # One entry per manual page. List of tuples
 # (source start file, name, description, authors, manual section).
 man_pages = [
-    ('index', 'queueuing-tool', 'queueuing-tool Documentation',
-     ['Daniel Jordon'], 1)
+    ("index", "queueuing-tool", "queueuing-tool Documentation", ["Daniel Jordon"], 1)
 ]
 
 # If true, show URL addresses after external links.
-#man_show_urls = False
+# man_show_urls = False
 
 
 # -- Options for Texinfo output -------------------------------------------
@@ -261,22 +261,29 @@
 # (source start file, target name, title, author,
 #  dir menu entry, description, category)
 texinfo_documents = [
-  ('index', 'queueuing-tool', 'queueuing-tool Documentation',
-   'Daniel Jordon', 'queueuing-tool', 'One line description of project.',
-   'Miscellaneous'),
+    (
+        "index",
+        "queueuing-tool",
+        "queueuing-tool Documentation",
+        "Daniel Jordon",
+        "queueuing-tool",
+        "One line description of project.",
+        "Miscellaneous",
+    ),
 ]
 
 # Documents to append as an appendix to all manuals.
-#texinfo_appendices = []
+# texinfo_appendices = []
 
 # If false, no module index is generated.
-#texinfo_domain_indices = True
+# texinfo_domain_indices = True
 
 # How to display URL addresses: 'footnote', 'no', or 'inline'.
-#texinfo_show_urls = 'footnote'
+# texinfo_show_urls = 'footnote'
 
 # If true, do not generate a @detailmenu in the "Top" node's menu.
-#texinfo_no_detailmenu = False
+# texinfo_no_detailmenu = False
+
 
 def setup(app):
-    app.add_js_file('copybutton.js')
+    app.add_js_file("copybutton.js")
diff --git a/examples/example_grocery_store.py b/examples/example_grocery_store.py
index 91551dc..9b40de4 100644
--- a/examples/example_grocery_store.py
+++ b/examples/example_grocery_store.py
@@ -1,15 +1,15 @@
 import functools
 
 import numpy as np
-import queueing_tool as qt
 
+import queueing_tool as qt
 
 # Make an adjacency list
 adja_list = {0: {1: {}}, 1: {k: {} for k in range(2, 22)}}
 
 # Make an object that has the same dimensions as your adjacency list that
 # specifies the type of queue along each edge.
-edge_list = {0: {1: 1}, 1: {k: 2 for k in range(2, 22)}}
+edge_list = {0: {1: 1}, 1: dict.fromkeys(range(2, 22), 2)}
 
 # Creates a networkx directed graph using the adjacency list and edge list
 g = qt.adjacency2graph(adjacency=adja_list, edge_type=edge_list)
@@ -21,7 +21,7 @@
 
 # Define the parameters for each of the queues
 def rate(t):
-    return 25 + 350 * np.sin(np.pi * t / 2)**2
+    return 25 + 350 * np.sin(np.pi * t / 2) ** 2
 
 
 def ser_f(t):
@@ -31,17 +31,15 @@ def ser_f(t):
 def identity(t):
     return t
 
+
 arr_f = functools.partial(qt.poisson_random_measure, rate=rate, rate_max=375)
 
 
 # Make a mapping between the edge types and the parameters used to make those
 # queues. If a particular parameter is not given then th defaults are used.
 q_args = {
-    1: {'arrival_f': arr_f,
-        'service_f': identity,
-        'AgentFactory': qt.GreedyAgent},
-    2: {'num_servers': 1,
-        'service_f': ser_f}
+    1: {"arrival_f": arr_f, "service_f": identity, "AgentFactory": qt.GreedyAgent},
+    2: {"num_servers": 1, "service_f": ser_f},
 }
 
 # Put it all together to create the network
@@ -49,7 +47,7 @@ def identity(t):
 
 # The default layout is spring_layout, which doesn't look good for this network.
 # This makes a new one
-qn.g.new_vertex_property('pos')
+qn.g.new_vertex_property("pos")
 pos = {}
 for v in qn.g.nodes():
     if v == 0:
@@ -57,14 +55,14 @@ def identity(t):
     elif v == 1:
         pos[v] = [0, 0.5]
     else:
-        pos[v] = [-5. + (v - 2.0) / 2, 0]
+        pos[v] = [-5.0 + (v - 2.0) / 2, 0]
 
 qn.g.set_pos(pos)
 # qn.draw(fname="store1.png", transparent=True, figsize=(12, 3),
 #        bgcolor=[0,0,0,0], bbox_inches='tight')
 
 # List the maximum number of agents from the default of 1000 to infinity
-qn.max_agents = np.infty
+qn.max_agents = np.inf
 
 # Before any simulations can take place the network must be initialized to
 # allow arrivals from outside the network. This specifies that only type 1
diff --git a/queueing_tool/__init__.py b/queueing_tool/__init__.py
index 8ab261e..36a2699 100644
--- a/queueing_tool/__init__.py
+++ b/queueing_tool/__init__.py
@@ -1,23 +1,18 @@
-from __future__ import absolute_import
 
 try:
     from importlib.metadata import version
 except ModuleNotFoundError:
     from importlib_metadata import version
 
-from queueing_tool.queues import *
-import queueing_tool.queues as queues
-
-from queueing_tool.network import *
-import queueing_tool.network as network
-
+from queueing_tool import graph, network, queues
 from queueing_tool.graph import *
-import queueing_tool.graph as graph
+from queueing_tool.network import *
+from queueing_tool.queues import *
 
 __version__ = version(__package__ or __name__)
 
 __all__ = []
-__all__.extend(['__version__'])
+__all__.extend(["__version__"])
 __all__.extend(queues.__all__)
 __all__.extend(network.__all__)
 __all__.extend(graph.__all__)
diff --git a/queueing_tool/graph/__init__.py b/queueing_tool/graph/__init__.py
index 4d8c013..e30144c 100644
--- a/queueing_tool/graph/__init__.py
+++ b/queueing_tool/graph/__init__.py
@@ -14,36 +14,28 @@
     ~queueing_tool.network.QueueingToolError
 """
 
-from queueing_tool.graph.graph_functions import (
-    graph2dict
-)
+from queueing_tool.graph.graph_functions import graph2dict
 from queueing_tool.graph.graph_generation import (
-    generate_random_graph,
     generate_pagerank_graph,
+    generate_random_graph,
+    generate_transition_matrix,
     minimal_random_graph,
     set_types_random,
     set_types_rank,
-    generate_transition_matrix
-)
-from queueing_tool.graph.graph_preparation import (
-    add_edge_lengths,
-    _prepare_graph
-)
-from queueing_tool.graph.graph_wrapper import (
-    adjacency2graph,
-    QueueNetworkDiGraph
 )
+from queueing_tool.graph.graph_preparation import _prepare_graph, add_edge_lengths
+from queueing_tool.graph.graph_wrapper import QueueNetworkDiGraph, adjacency2graph
 
 __all__ = [
-    '_prepare_graph',
-    'add_edge_lengths',
-    'adjacency2graph',
-    'generate_random_graph',
-    'generate_pagerank_graph',
-    'generate_transition_matrix',
-    'graph2dict',
-    'minimal_random_graph',
-    'set_types_rank',
-    'set_types_random',
-    'QueueNetworkDiGraph'
+    "QueueNetworkDiGraph",
+    "_prepare_graph",
+    "add_edge_lengths",
+    "adjacency2graph",
+    "generate_pagerank_graph",
+    "generate_random_graph",
+    "generate_transition_matrix",
+    "graph2dict",
+    "minimal_random_graph",
+    "set_types_random",
+    "set_types_rank",
 ]
diff --git a/queueing_tool/graph/graph_functions.py b/queueing_tool/graph/graph_functions.py
index f45eed0..88739dc 100644
--- a/queueing_tool/graph/graph_functions.py
+++ b/queueing_tool/graph/graph_functions.py
@@ -5,20 +5,20 @@
 
 
 def _calculate_distance(latlon1, latlon2):
-    """Calculates the distance between two points on earth.
-    """
+    """Calculates the distance between two points on earth."""
     lat1, lon1 = latlon1
     lat2, lon2 = latlon2
     dlon = lon2 - lon1
     dlat = lat2 - lat1
     R = 6371  # radius of the earth in kilometers
-    a = np.sin(dlat / 2)**2 + np.cos(lat1) * np.cos(lat2) * (np.sin(dlon / 2))**2
+    a = np.sin(dlat / 2) ** 2 + np.cos(lat1) * np.cos(lat2) * (np.sin(dlon / 2)) ** 2
     c = 2 * np.pi * R * np.arctan2(np.sqrt(a), np.sqrt(1 - a)) / 180
     return c
 
 
 def _test_graph(graph):
-    """A function that makes sure ``graph`` is either a
+    """
+    A function that makes sure ``graph`` is either a
     :any:`networkx.DiGraph` or a string or file object to one.
 
     Parameters
@@ -34,6 +34,7 @@ def _test_graph(graph):
     TypeError
         Raises a :exc:`~TypeError` if ``graph`` cannot be turned into a
         :any:`networkx.DiGraph`.
+
     """
     if not isinstance(graph, QueueNetworkDiGraph):
         try:
@@ -44,7 +45,8 @@ def _test_graph(graph):
 
 
 def graph2dict(g, return_dict_of_dict=True):
-    """Takes a graph and returns an adjacency list.
+    """
+    Takes a graph and returns an adjacency list.
 
     Parameters
     ----------
@@ -77,6 +79,7 @@ def graph2dict(g, return_dict_of_dict=True):
     3: {2: {}}}
     >>> qt.graph2dict(g, return_dict_of_dict=False)
     {0: [1, 2], 1: [0], 2: [0, 3], 3: [2]}
+
     """
     if not isinstance(g, nx.DiGraph):
         g = QueueNetworkDiGraph(g)
@@ -84,5 +87,4 @@ def graph2dict(g, return_dict_of_dict=True):
     dict_of_dicts = nx.to_dict_of_dicts(g)
     if return_dict_of_dict:
         return dict_of_dicts
-    else:
-        return {k: list(val.keys()) for k, val in dict_of_dicts.items()}
+    return {k: list(val.keys()) for k, val in dict_of_dicts.items()}
diff --git a/queueing_tool/graph/graph_generation.py b/queueing_tool/graph/graph_generation.py
index 87b6cdc..30abc35 100644
--- a/queueing_tool/graph/graph_generation.py
+++ b/queueing_tool/graph/graph_generation.py
@@ -4,13 +4,14 @@
 import networkx as nx
 import numpy as np
 
-from queueing_tool.graph.graph_functions import _test_graph, _calculate_distance
+from queueing_tool.graph.graph_functions import _calculate_distance, _test_graph
 from queueing_tool.graph.graph_wrapper import QueueNetworkDiGraph
 from queueing_tool.union_find import UnionFind
 
 
 def generate_transition_matrix(g, seed=None):
-    """Generates a random transition matrix for the graph ``g``.
+    """
+    Generates a random transition matrix for the graph ``g``.
 
     Parameters
     ----------
@@ -27,6 +28,7 @@ def generate_transition_matrix(g, seed=None):
         probability of transitioning from vertex ``i`` to vertex ``j``.
         If there is no edge connecting vertex ``i`` to vertex ``j``
         then ``mat[i, j] = 0``.
+
     """
     g = _test_graph(g)
 
@@ -42,7 +44,7 @@ def generate_transition_matrix(g, seed=None):
         if deg == 1:
             mat[v, ind] = 1
         elif deg > 1:
-            probs = np.ceil(np.random.rand(deg) * 100) / 100.
+            probs = np.ceil(np.random.rand(deg) * 100) / 100.0
             if np.isclose(np.sum(probs), 0):
                 probs[np.random.randint(deg)] = 1
 
@@ -52,7 +54,8 @@ def generate_transition_matrix(g, seed=None):
 
 
 def generate_random_graph(num_vertices=250, prob_loop=0.5, **kwargs):
-    """Creates a random graph where the edges have different types.
+    """
+    Creates a random graph where the edges have different types.
 
     This method calls :func:`.minimal_random_graph`, and then adds
     a loop to each vertex with ``prob_loop`` probability. It then
@@ -105,6 +108,7 @@ def generate_random_graph(num_vertices=250, prob_loop=0.5, **kwargs):
     Note that none of the edge types in the above example are 0. It is
     recommended use edge type indices starting at 1, since 0 is
     typically used for terminal edges.
+
     """
     g = minimal_random_graph(num_vertices, **kwargs)
     for v in g.nodes():
@@ -117,7 +121,8 @@ def generate_random_graph(num_vertices=250, prob_loop=0.5, **kwargs):
 
 
 def generate_pagerank_graph(num_vertices=250, **kwargs):
-    """Creates a random graph where the vertex types are
+    """
+    Creates a random graph where the vertex types are
     selected using their pagerank.
 
     Calls :func:`.minimal_random_graph` and then
@@ -149,6 +154,7 @@ def generate_pagerank_graph(num_vertices=250, **kwargs):
     vertices as type 3, and adds loops to these vertices as well. These
     loops then have edge types that correspond to the vertices type.
     The rest of the edges are set to type 1.
+
     """
     g = minimal_random_graph(num_vertices, **kwargs)
     r = np.zeros(num_vertices)
@@ -171,7 +177,7 @@ def generate_pagerank_graph(num_vertices=250, **kwargs):
                 page_rank = nx.pagerank_numpy(g)
             except:
                 raise exe
- 
+
     for k, pr in page_rank.items():
         r[k] = pr
     g = set_types_rank(g, rank=r, **kwargs)
@@ -179,7 +185,8 @@ def generate_pagerank_graph(num_vertices=250, **kwargs):
 
 
 def minimal_random_graph(num_vertices, seed=None, **kwargs):
-    """Creates a connected graph with random vertex locations.
+    """
+    Creates a connected graph with random vertex locations.
 
     Parameters
     ----------
@@ -204,6 +211,7 @@ def minimal_random_graph(num_vertices, seed=None, **kwargs):
     ``v``, all other vertices with Euclidean distance less or equal to
     ``r`` are connect by an edge --- where ``r`` is the smallest number
     such that the graph ends up connected.
+
     """
     if isinstance(seed, numbers.Integral):
         np.random.seed(seed)
@@ -214,11 +222,11 @@ def minimal_random_graph(num_vertices, seed=None, **kwargs):
     for k in range(num_vertices - 1):
         for j in range(k + 1, num_vertices):
             v = points[k] - points[j]
-            edges.append((k, j, v[0]**2 + v[1]**2))
+            edges.append((k, j, v[0] ** 2 + v[1] ** 2))
 
-    mytype = [('n1', int), ('n2', int), ('distance', float)]
+    mytype = [("n1", int), ("n2", int), ("distance", float)]
     edges = np.array(edges, dtype=mytype)
-    edges = np.sort(edges, order='distance')
+    edges = np.sort(edges, order="distance")
     unionF = UnionFind([k for k in range(num_vertices)])
 
     g = nx.Graph()
@@ -235,9 +243,9 @@ def minimal_random_graph(num_vertices, seed=None, **kwargs):
     return g
 
 
-def set_types_random(g, proportions=None, loop_proportions=None, seed=None,
-                     **kwargs):
-    """Randomly sets ``edge_type`` (edge type) properties of the graph.
+def set_types_random(g, proportions=None, loop_proportions=None, seed=None, **kwargs):
+    """
+    Randomly sets ``edge_type`` (edge type) properties of the graph.
 
     This function randomly assigns each edge a type. The probability of
     an edge being a specific type is proscribed in the
@@ -284,6 +292,7 @@ def set_types_random(g, proportions=None, loop_proportions=None, seed=None,
     defaults to four types in the graph (types 0, 1, 2, and 3). It sets
     non-loop edges to be either 1, 2, or 3 33% chance, and loops are
     types 0, 1, 2, 3 with 25% chance.
+
     """
     g = _test_graph(g)
 
@@ -291,10 +300,10 @@ def set_types_random(g, proportions=None, loop_proportions=None, seed=None,
         np.random.seed(seed)
 
     if proportions is None:
-        proportions = {k: 1. / 3 for k in range(1, 4)}
+        proportions = dict.fromkeys(range(1, 4), 1.0 / 3)
 
     if loop_proportions is None:
-        loop_proportions = {k: 1. / 4 for k in range(4)}
+        loop_proportions = dict.fromkeys(range(4), 1.0 / 4)
 
     edges = [e for e in g.edges() if e[0] != e[1]]
     loops = [e for e in g.edges() if e[0] == e[1]]
@@ -319,15 +328,16 @@ def set_types_random(g, proportions=None, loop_proportions=None, seed=None,
     for k, e in enumerate(loops):
         eTypes[e] = values[k]
 
-    g.new_edge_property('edge_type')
+    g.new_edge_property("edge_type")
     for e in g.edges():
-        g.set_ep(e, 'edge_type', eTypes[e])
+        g.set_ep(e, "edge_type", eTypes[e])
 
     return g
 
 
 def set_types_rank(g, rank, pType2=0.1, pType3=0.1, seed=None, **kwargs):
-    """Creates a stylized graph. Sets edge and types using `pagerank`_.
+    """
+    Creates a stylized graph. Sets edge and types using `pagerank`_.
 
     This function sets the edge types of a graph to be either 1, 2, or
     3. It sets the vertices to type 2 by selecting the top
@@ -368,6 +378,7 @@ def set_types_rank(g, rank, pType2=0.1, pType3=0.1, seed=None, **kwargs):
     TypeError
         Raised when the parameter ``g`` is not of a type that can be
         made into a :any:`DiGraph<networkx.DiGraph>`.
+
     """
     g = _test_graph(g)
 
@@ -378,49 +389,51 @@ def set_types_rank(g, rank, pType2=0.1, pType3=0.1, seed=None, **kwargs):
     nDests = int(np.ceil(g.number_of_nodes() * pType2))
     dests = np.where(rank >= tmp[-nDests])[0]
 
-    if 'pos' not in g.vertex_properties():
+    if "pos" not in g.vertex_properties():
         g.set_pos()
 
-    dest_pos = np.array([g.vp(v, 'pos') for v in dests])
+    dest_pos = np.array([g.vp(v, "pos") for v in dests])
     nFCQ = int(pType3 * g.number_of_nodes())
-    min_g_dist = np.ones(nFCQ) * np.infty
+    min_g_dist = np.ones(nFCQ) * np.inf
     ind_g_dist = np.ones(nFCQ, int)
 
-    r, theta = np.random.random(nFCQ) / 500., np.random.random(nFCQ) * 360.
+    r, theta = np.random.random(nFCQ) / 500.0, np.random.random(nFCQ) * 360.0
     xy_pos = np.array([r * np.cos(theta), r * np.sin(theta)]).transpose()
     g_pos = xy_pos + dest_pos[np.array(np.mod(np.arange(nFCQ), nDests), int)]
 
     for v in g.nodes():
         if v not in dests:
-            tmp = np.array([_calculate_distance(g.vp(v, 'pos'), g_pos[k, :]) for k in range(nFCQ)])
+            tmp = np.array(
+                [_calculate_distance(g.vp(v, "pos"), g_pos[k, :]) for k in range(nFCQ)]
+            )
             min_g_dist = np.min((tmp, min_g_dist), 0)
             ind_g_dist[min_g_dist == tmp] = v
 
     ind_g_dist = np.unique(ind_g_dist)
-    fcqs = set(ind_g_dist[:min(nFCQ, len(ind_g_dist))])
+    fcqs = set(ind_g_dist[: min(nFCQ, len(ind_g_dist))])
     dests = set(dests)
-    g.new_vertex_property('loop_type')
+    g.new_vertex_property("loop_type")
 
     for v in g.nodes():
         if v in dests:
-            g.set_vp(v, 'loop_type', 3)
+            g.set_vp(v, "loop_type", 3)
             if not g.is_edge((v, v)):
                 g.add_edge(v, v)
         elif v in fcqs:
-            g.set_vp(v, 'loop_type', 2)
+            g.set_vp(v, "loop_type", 2)
             if not g.is_edge((v, v)):
                 g.add_edge(v, v)
 
-    g.new_edge_property('edge_type')
+    g.new_edge_property("edge_type")
     for e in g.edges():
-        g.set_ep(e, 'edge_type', 1)
+        g.set_ep(e, "edge_type", 1)
 
     for v in g.nodes():
-        if g.vp(v, 'loop_type') in [2, 3]:
+        if g.vp(v, "loop_type") in [2, 3]:
             e = (v, v)
-            if g.vp(v, 'loop_type') == 2:
-                g.set_ep(e, 'edge_type', 2)
+            if g.vp(v, "loop_type") == 2:
+                g.set_ep(e, "edge_type", 2)
             else:
-                g.set_ep(e, 'edge_type', 3)
+                g.set_ep(e, "edge_type", 3)
 
     return g
diff --git a/queueing_tool/graph/graph_preparation.py b/queueing_tool/graph/graph_preparation.py
index e751d00..5493a62 100644
--- a/queueing_tool/graph/graph_preparation.py
+++ b/queueing_tool/graph/graph_preparation.py
@@ -1,15 +1,13 @@
 import networkx as nx
 import numpy as np
 
-from queueing_tool.graph.graph_functions import _test_graph, _calculate_distance
-from queueing_tool.graph.graph_wrapper import (
-    adjacency2graph,
-    QueueNetworkDiGraph
-)
+from queueing_tool.graph.graph_functions import _calculate_distance, _test_graph
+from queueing_tool.graph.graph_wrapper import QueueNetworkDiGraph, adjacency2graph
 
 
 def add_edge_lengths(g):
-    """Add add the edge lengths as a :any:`DiGraph<networkx.DiGraph>`
+    """
+    Add add the edge lengths as a :any:`DiGraph<networkx.DiGraph>`
     for the graph.
 
     Uses the ``pos`` vertex property to get the location of each
@@ -36,18 +34,19 @@ def add_edge_lengths(g):
 
     """
     g = _test_graph(g)
-    g.new_edge_property('edge_length')
+    g.new_edge_property("edge_length")
 
     for e in g.edges():
-        latlon1 = g.vp(e[1], 'pos')
-        latlon2 = g.vp(e[0], 'pos')
-        g.set_ep(e, 'edge_length', np.round(_calculate_distance(latlon1, latlon2), 3))
+        latlon1 = g.vp(e[1], "pos")
+        latlon2 = g.vp(e[0], "pos")
+        g.set_ep(e, "edge_length", np.round(_calculate_distance(latlon1, latlon2), 3))
 
     return g
 
 
 def _prepare_graph(g, g_colors, q_cls, q_arg, adjust_graph):
-    """Prepares a graph for use in :class:`.QueueNetwork`.
+    """
+    Prepares a graph for use in :class:`.QueueNetwork`.
 
     This function is called by ``__init__`` in the
     :class:`.QueueNetwork` class. It creates the :class:`.QueueServer`
@@ -108,50 +107,51 @@ def _prepare_graph(g, g_colors, q_cls, q_arg, adjust_graph):
     TypeError
         Raised when the parameter ``g`` is not of a type that can be
         made into a :any:`networkx.DiGraph`.
+
     """
     g = _test_graph(g)
 
     if adjust_graph:
-        pos = nx.get_node_attributes(g, 'pos')
+        pos = nx.get_node_attributes(g, "pos")
         ans = nx.to_dict_of_dicts(g)
         g = adjacency2graph(ans, adjust=2, is_directed=g.is_directed())
         g = QueueNetworkDiGraph(g)
         if len(pos) > 0:
             g.set_pos(pos)
 
-    g.new_vertex_property('vertex_color')
-    g.new_vertex_property('vertex_fill_color')
-    g.new_vertex_property('vertex_pen_width')
-    g.new_vertex_property('vertex_size')
+    g.new_vertex_property("vertex_color")
+    g.new_vertex_property("vertex_fill_color")
+    g.new_vertex_property("vertex_pen_width")
+    g.new_vertex_property("vertex_size")
 
-    g.new_edge_property('edge_control_points')
-    g.new_edge_property('edge_color')
-    g.new_edge_property('edge_marker_size')
-    g.new_edge_property('edge_pen_width')
+    g.new_edge_property("edge_control_points")
+    g.new_edge_property("edge_color")
+    g.new_edge_property("edge_marker_size")
+    g.new_edge_property("edge_pen_width")
 
-    queues = _set_queues(g, q_cls, q_arg, 'cap' in g.vertex_properties())
+    queues = _set_queues(g, q_cls, q_arg, "cap" in g.vertex_properties())
 
-    if 'pos' not in g.vertex_properties():
+    if "pos" not in g.vertex_properties():
         g.set_pos()
 
     for k, e in enumerate(g.edges()):
-        g.set_ep(e, 'edge_pen_width', 1.25)
-        g.set_ep(e, 'edge_marker_size', 8)
+        g.set_ep(e, "edge_pen_width", 1.25)
+        g.set_ep(e, "edge_marker_size", 8)
         if e[0] == e[1]:
-            g.set_ep(e, 'edge_color', queues[k].colors['edge_loop_color'])
+            g.set_ep(e, "edge_color", queues[k].colors["edge_loop_color"])
         else:
-            g.set_ep(e, 'edge_color', queues[k].colors['edge_color'])
+            g.set_ep(e, "edge_color", queues[k].colors["edge_color"])
 
     for v in g.nodes():
-        g.set_vp(v, 'vertex_pen_width', 1)
-        g.set_vp(v, 'vertex_size', 8)
+        g.set_vp(v, "vertex_pen_width", 1)
+        g.set_vp(v, "vertex_size", 8)
         e = (v, v)
         if g.is_edge(e):
-            g.set_vp(v, 'vertex_color', queues[g.edge_index[e]]._current_color(2))
-            g.set_vp(v, 'vertex_fill_color', queues[g.edge_index[e]]._current_color())
+            g.set_vp(v, "vertex_color", queues[g.edge_index[e]]._current_color(2))
+            g.set_vp(v, "vertex_fill_color", queues[g.edge_index[e]]._current_color())
         else:
-            g.set_vp(v, 'vertex_color', g_colors['vertex_color'])
-            g.set_vp(v, 'vertex_fill_color', g_colors['vertex_fill_color'])
+            g.set_vp(v, "vertex_color", g_colors["vertex_color"])
+            g.set_vp(v, "vertex_fill_color", g_colors["vertex_fill_color"])
 
     return g, queues
 
@@ -160,12 +160,12 @@ def _set_queues(g, q_cls, q_arg, has_cap):
     queues = [0 for k in range(g.number_of_edges())]
 
     for e in g.edges():
-        eType = g.ep(e, 'edge_type')
+        eType = g.ep(e, "edge_type")
         qedge = (e[0], e[1], g.edge_index[e], eType)
 
-        if has_cap and 'num_servers' not in q_arg[eType]:
-            cap = g.vp(e[1], 'cap') if g.vp(e[1], 'cap') is not None else 0
-            q_arg[eType]['num_servers'] = max(cap, 1)
+        if has_cap and "num_servers" not in q_arg[eType]:
+            cap = g.vp(e[1], "cap") if g.vp(e[1], "cap") is not None else 0
+            q_arg[eType]["num_servers"] = max(cap, 1)
 
         queues[qedge[2]] = q_cls[eType](edge=qedge, **q_arg[eType])
 
diff --git a/queueing_tool/graph/graph_wrapper.py b/queueing_tool/graph/graph_wrapper.py
index 656e13d..dffb832 100644
--- a/queueing_tool/graph/graph_wrapper.py
+++ b/queueing_tool/graph/graph_wrapper.py
@@ -5,7 +5,7 @@
     import matplotlib.pyplot as plt
     from matplotlib.collections import LineCollection
 
-    plt.style.use('ggplot')
+    plt.style.use("ggplot")
     HAS_MATPLOTLIB = True
 
 except ImportError:
@@ -25,7 +25,8 @@ def _matrix2dict(matrix, etype=False):
 
 
 def _dict2dict(adj_dict):
-    """Takes a dictionary based representation of an adjacency list
+    """
+    Takes a dictionary based representation of an adjacency list
     and returns a dict of dicts based representation.
     """
     item = adj_dict.popitem()
@@ -40,14 +41,14 @@ def _dict2dict(adj_dict):
 
 
 def _adjacency_adjust(adjacency, adjust, is_directed):
-    """Takes an adjacency list and returns a (possibly) modified
+    """
+    Takes an adjacency list and returns a (possibly) modified
     adjacency list.
     """
-
     for v, adj in adjacency.items():
         for properties in adj.values():
-            if properties.get('edge_type') is None:
-                properties['edge_type'] = 1
+            if properties.get("edge_type") is None:
+                properties["edge_type"] = 1
 
     if is_directed:
         if adjust == 2:
@@ -60,18 +61,19 @@ def _adjacency_adjust(adjacency, adjust, is_directed):
             for k, adj in adjacency.items():
                 for v in adj.keys():
                     if v in null_nodes:
-                        adj[v]['edge_type'] = 0
+                        adj[v]["edge_type"] = 0
 
         else:
             for k, adj in adjacency.items():
                 if len(adj) == 0:
-                    adj[k] = {'edge_type': 0}
+                    adj[k] = {"edge_type": 0}
 
     return adjacency
 
 
 def adjacency2graph(adjacency, edge_type=None, adjust=1, **kwargs):
-    """Takes an adjacency list, dict, or matrix and returns a graph.
+    """
+    Takes an adjacency list, dict, or matrix and returns a graph.
 
     The purpose of this function is take an adjacency list (or matrix)
     and return a :class:`.QueueNetworkDiGraph` that can be used with a
@@ -125,7 +127,7 @@ def adjacency2graph(adjacency, edge_type=None, adjust=1, **kwargs):
     >>> ans = qt.graph2dict(g)
     >>> sorted(ans.items())     # doctest: +NORMALIZE_WHITESPACE
     [(0, {1: {'edge_type': 1}}),
-     (1, {2: {'edge_type': 2}, 3: {'edge_type': 4}}), 
+     (1, {2: {'edge_type': 2}, 3: {'edge_type': 4}}),
      (2, {2: {'edge_type': 0}}),
      (3, {0: {'edge_type': 1}})]
 
@@ -151,28 +153,29 @@ def adjacency2graph(adjacency, edge_type=None, adjust=1, **kwargs):
      (1, {2: {'edge_type': 0}, 3: {'edge_type': 4}}),
      (2, {}),
      (3, {0: {'edge_type': 1}})]
-    """
 
+    """
     if isinstance(adjacency, np.ndarray):
         adjacency = _matrix2dict(adjacency)
     elif isinstance(adjacency, dict):
         adjacency = _dict2dict(adjacency)
     else:
-        msg = ("If the adjacency parameter is supplied it must be a "
-               "dict, or a numpy.ndarray.")
+        msg = (
+            "If the adjacency parameter is supplied it must be a "
+            "dict, or a numpy.ndarray."
+        )
         raise TypeError(msg)
 
     if edge_type is None:
         edge_type = {}
-    else:
-        if isinstance(edge_type, np.ndarray):
-            edge_type = _matrix2dict(edge_type, etype=True)
-        elif isinstance(edge_type, dict):
-            edge_type = _dict2dict(edge_type)
+    elif isinstance(edge_type, np.ndarray):
+        edge_type = _matrix2dict(edge_type, etype=True)
+    elif isinstance(edge_type, dict):
+        edge_type = _dict2dict(edge_type)
 
     for u, ty in edge_type.items():
         for v, et in ty.items():
-            adjacency[u][v]['edge_type'] = et
+            adjacency[u][v]["edge_type"] = et
 
     g = nx.from_dict_of_dicts(adjacency, create_using=nx.DiGraph())
     adjacency = nx.to_dict_of_dicts(g)
@@ -181,18 +184,20 @@ def adjacency2graph(adjacency, edge_type=None, adjust=1, **kwargs):
     return nx.from_dict_of_dicts(adjacency, create_using=nx.DiGraph())
 
 
-SAVEFIG_KWARGS = set([
-    'dpi',
-    'facecolorw',
-    'edgecolorw',
-    'orientationportrait',
-    'papertype',
-    'format',
-    'transparent',
-    'bbox_inches',
-    'pad_inches',
-    'frameon'
-])
+SAVEFIG_KWARGS = set(
+    [
+        "dpi",
+        "facecolorw",
+        "edgecolorw",
+        "orientationportrait",
+        "papertype",
+        "format",
+        "transparent",
+        "bbox_inches",
+        "pad_inches",
+        "frameon",
+    ]
+)
 
 
 class QueueNetworkDiGraph(nx.DiGraph):
@@ -228,7 +233,9 @@ class QueueNetworkDiGraph(nx.DiGraph):
     -----
     Not suitable for stand alone use; only use with a
     :class:`.QueueNetwork`.
+
     """
+
     def __init__(self, data=None, **kwargs):
         if isinstance(data, dict):
             data = adjacency2graph(data, **kwargs)
@@ -238,7 +245,7 @@ def __init__(self, data=None, **kwargs):
 
         self.edge_index = {e: k for k, e in enumerate(edges)}
 
-        pos = nx.get_node_attributes(self, name='pos')
+        pos = nx.get_node_attributes(self, name="pos")
         if len(pos) == self.number_of_nodes():
             self.pos = np.array([pos[v] for v in self.nodes()])
         else:
@@ -296,27 +303,28 @@ def edge_properties(self):
         return props
 
     def new_vertex_property(self, name):
-        values = {v: None for v in self.nodes()}
+        values = dict.fromkeys(self.nodes())
         nx.set_node_attributes(self, name=name, values=values)
-        if name == 'vertex_color':
+        if name == "vertex_color":
             self.vertex_color = [0 for v in range(self.number_of_nodes())]
-        if name == 'vertex_fill_color':
+        if name == "vertex_fill_color":
             self.vertex_fill_color = [0 for v in range(self.number_of_nodes())]
 
     def new_edge_property(self, name):
-        values = {e: None for e in self.edges()}
+        values = dict.fromkeys(self.edges())
         nx.set_edge_attributes(self, name=name, values=values)
-        if name == 'edge_color':
+        if name == "edge_color":
             self.edge_color = np.zeros((self.number_of_edges(), 4))
 
     def set_pos(self, pos=None):
         if pos is None:
             pos = nx.spring_layout(self)
-        nx.set_node_attributes(self, name='pos', values=pos)
+        nx.set_node_attributes(self, name="pos", values=pos)
         self.pos = np.array([pos[v] for v in self.nodes()])
 
     def get_edge_type(self, edge_type):
-        """Returns all edges with the specified edge type.
+        """
+        Returns all edges with the specified edge type.
 
         Parameters
         ----------
@@ -346,15 +354,17 @@ def get_edge_type(self, edge_type):
         >>> ans.sort()
         >>> ans
         [(0, 1), (2, 0)]
+
         """
         edges = []
         for e in self.edges():
-            if self.adj[e[0]][e[1]].get('edge_type') == edge_type:
+            if self.adj[e[0]][e[1]].get("edge_type") == edge_type:
                 edges.append(e)
         return edges
 
     def draw_graph(self, line_kwargs=None, scatter_kwargs=None, **kwargs):
-        """Draws the graph.
+        """
+        Draws the graph.
 
         Uses matplotlib, specifically
         :class:`~matplotlib.collections.LineCollection` and
@@ -389,17 +399,18 @@ def draw_graph(self, line_kwargs=None, scatter_kwargs=None, **kwargs):
         -----
         If the ``fname`` keyword is passed, then the figure is saved
         locally.
+
         """
         if not HAS_MATPLOTLIB:
             raise ImportError("Matplotlib is required to draw the graph.")
 
-        fig = plt.figure(figsize=kwargs.get('figsize', (7, 7)))
+        fig = plt.figure(figsize=kwargs.get("figsize", (7, 7)))
         ax = fig.gca()
 
         mpl_kwargs = {
-            'line_kwargs': line_kwargs,
-            'scatter_kwargs': scatter_kwargs,
-            'pos': kwargs.get('pos')
+            "line_kwargs": line_kwargs,
+            "scatter_kwargs": scatter_kwargs,
+            "pos": kwargs.get("pos"),
         }
 
         line_kwargs, scatter_kwargs = self.lines_scatter_args(**mpl_kwargs)
@@ -408,24 +419,25 @@ def draw_graph(self, line_kwargs=None, scatter_kwargs=None, **kwargs):
         ax.add_collection(edge_collection)
         ax.scatter(**scatter_kwargs)
 
-        if hasattr(ax, 'set_facecolor'):
-            ax.set_facecolor(kwargs.get('bgcolor', [1, 1, 1, 1]))
+        if hasattr(ax, "set_facecolor"):
+            ax.set_facecolor(kwargs.get("bgcolor", [1, 1, 1, 1]))
         else:
-            ax.set_axis_bgcolor(kwargs.get('bgcolor', [1, 1, 1, 1]))
+            ax.set_axis_bgcolor(kwargs.get("bgcolor", [1, 1, 1, 1]))
 
         ax.get_xaxis().set_visible(False)
         ax.get_yaxis().set_visible(False)
 
-        if 'fname' in kwargs:
+        if "fname" in kwargs:
             # savefig needs a positional argument for some reason
             new_kwargs = {k: v for k, v in kwargs.items() if k in SAVEFIG_KWARGS}
-            fig.savefig(kwargs['fname'], **new_kwargs)
+            fig.savefig(kwargs["fname"], **new_kwargs)
         else:
             plt.ion()
             plt.show()
 
     def lines_scatter_args(self, line_kwargs=None, scatter_kwargs=None, pos=None):
-        """Returns the arguments used when plotting.
+        """
+        Returns the arguments used when plotting.
 
         Takes any keyword arguments for
         :class:`~matplotlib.collections.LineCollection` and
@@ -454,6 +466,7 @@ def lines_scatter_args(self, line_kwargs=None, scatter_kwargs=None, pos=None):
         -----
         If a specific keyword argument is not passed then the defaults
         are used.
+
         """
         if pos is not None:
             self.set_pos(pos)
@@ -466,36 +479,36 @@ def lines_scatter_args(self, line_kwargs=None, scatter_kwargs=None, pos=None):
             edge_pos[ei] = (self.pos[e[0]], self.pos[e[1]])
 
         line_collecton_kwargs = {
-            'segments': edge_pos,
-            'colors': self.edge_color,
-            'linewidths': (1,),
-            'antialiaseds': (1,),
-            'linestyle': 'solid',
-            'transOffset': None,
-            'cmap': plt.cm.ocean_r,
-            'pickradius': 5,
-            'zorder': 0,
-            'facecolors': None,
-            'norm': None,
-            'offsets': None,
-            'hatch': None,
+            "segments": edge_pos,
+            "colors": self.edge_color,
+            "linewidths": (1,),
+            "antialiaseds": (1,),
+            "linestyle": "solid",
+            "transOffset": None,
+            "cmap": plt.cm.ocean_r,
+            "pickradius": 5,
+            "zorder": 0,
+            "facecolors": None,
+            "norm": None,
+            "offsets": None,
+            "hatch": None,
         }
         scatter_kwargs_ = {
-            'x': self.pos[:, 0],
-            'y': self.pos[:, 1],
-            's': 50,
-            'c': self.vertex_fill_color,
-            'alpha': None,
-            'norm': None,
-            'vmin': None,
-            'vmax': None,
-            'marker': 'o',
-            'zorder': 2,
-            'linewidths': 1,
-            'edgecolors': self.vertex_color,
-            'facecolors': None,
-            'antialiaseds': None,
-            'hatch': None,
+            "x": self.pos[:, 0],
+            "y": self.pos[:, 1],
+            "s": 50,
+            "c": self.vertex_fill_color,
+            "alpha": None,
+            "norm": None,
+            "vmin": None,
+            "vmax": None,
+            "marker": "o",
+            "zorder": 2,
+            "linewidths": 1,
+            "edgecolors": self.vertex_color,
+            "facecolors": None,
+            "antialiaseds": None,
+            "hatch": None,
         }
 
         line_kwargs = {} if line_kwargs is None else line_kwargs
diff --git a/queueing_tool/network/__init__.py b/queueing_tool/network/__init__.py
index 7f726f7..92c5093 100644
--- a/queueing_tool/network/__init__.py
+++ b/queueing_tool/network/__init__.py
@@ -23,13 +23,6 @@
 """
 
 from queueing_tool.network.priority_queue import PriorityQueue
-from queueing_tool.network.queue_network import (
-    QueueingToolError,
-    QueueNetwork
-)
+from queueing_tool.network.queue_network import QueueingToolError, QueueNetwork
 
-__all__ = [
-    'PriorityQueue',
-    'QueueingToolError',
-    'QueueNetwork'
-]
+__all__ = ["PriorityQueue", "QueueNetwork", "QueueingToolError"]
diff --git a/queueing_tool/network/queue_network.py b/queueing_tool/network/queue_network.py
index 3547f21..80d5956 100644
--- a/queueing_tool/network/queue_network.py
+++ b/queueing_tool/network/queue_network.py
@@ -1,7 +1,7 @@
+import array
 import collections
-import numbers
 import copy
-import array
+import numbers
 
 import numpy as np
 from numpy.random import uniform
@@ -11,24 +11,20 @@
     from matplotlib.animation import FuncAnimation
     from matplotlib.collections import LineCollection
 
-    plt.style.use('ggplot')
+    plt.style.use("ggplot")
     HAS_MATPLOTLIB = True
 
 except ImportError:
     HAS_MATPLOTLIB = False
 
 from queueing_tool.graph import _prepare_graph
-from queueing_tool.queues import (
-    NullQueue,
-    QueueServer,
-    LossQueue
-)
 from queueing_tool.network.priority_queue import PriorityQueue
+from queueing_tool.queues import LossQueue, NullQueue, QueueServer
 
 
 class QueueingToolError(Exception):
     """Base class for exceptions in Queueing-tool."""
-    pass
+
 
 
 EPS = np.float64(1e-7)
@@ -38,12 +34,13 @@ class QueueingToolError(Exception):
     [0, 0.5, 1, 1],
     [0.133, 0.545, 0.133, 1],
     [0.282, 0.239, 0.545, 1],
-    [1, 0.135, 0, 1]
+    [1, 0.135, 0, 1],
 ]
 
 
-class QueueNetwork(object):
-    """A class that simulates a network of queues.
+class QueueNetwork:
+    """
+    A class that simulates a network of queues.
 
     Takes a networkx :any:`DiGraph<networkx.DiGraph>` and places queues
     on each edge of the graph. The simulations are event based, and
@@ -274,18 +271,19 @@ class QueueNetwork(object):
 
     .. figure:: my_network1.png
        :align: center
+
     """
 
     default_colors = {
-        'vertex_fill_color': [0.95, 0.95, 0.95, 1.0],
-        'vertex_color': [0.0, 0.5, 1.0, 1.0],
-        'vertex_highlight': [0.5, 0.5, 0.5, 1.0],
-        'edge_departure': [0, 0, 0, 1],
-        'vertex_active': [0.1, 1.0, 0.5, 1.0],
-        'vertex_inactive': [0.95, 0.95, 0.95, 1.0],
-        'edge_active': [0.1, 0.1, 0.1, 1.0],
-        'edge_inactive': [0.8, 0.8, 0.8, 0.3],
-        'bgcolor': [1, 1, 1, 1]
+        "vertex_fill_color": [0.95, 0.95, 0.95, 1.0],
+        "vertex_color": [0.0, 0.5, 1.0, 1.0],
+        "vertex_highlight": [0.5, 0.5, 0.5, 1.0],
+        "edge_departure": [0, 0, 0, 1],
+        "vertex_active": [0.1, 1.0, 0.5, 1.0],
+        "vertex_inactive": [0.95, 0.95, 0.95, 1.0],
+        "edge_active": [0.1, 0.1, 0.1, 1.0],
+        "edge_inactive": [0.8, 0.8, 0.8, 0.3],
+        "bgcolor": [1, 1, 1, 1],
     }
 
     default_classes = {
@@ -293,20 +291,30 @@ class QueueNetwork(object):
         1: QueueServer,
         2: LossQueue,
         3: LossQueue,
-        4: LossQueue
+        4: LossQueue,
     }
 
     default_q_colors = {
-        k: {'edge_loop_color': [0, 0, 0, 0],
-            'edge_color': [0.8, 0.8, 0.8, 1.0],
-            'vertex_fill_color': [0.95, 0.95, 0.95, 1.0],
-            'vertex_color': v_pens[k]}
+        k: {
+            "edge_loop_color": [0, 0, 0, 0],
+            "edge_color": [0.8, 0.8, 0.8, 1.0],
+            "vertex_fill_color": [0.95, 0.95, 0.95, 1.0],
+            "vertex_color": v_pens[k],
+        }
         for k in range(5)
     }
 
-    def __init__(self, g, q_classes=None, q_args=None, seed=None, colors=None,
-                 max_agents=1000, blocking='BAS', adjust_graph=True):
-
+    def __init__(
+        self,
+        g,
+        q_classes=None,
+        q_args=None,
+        seed=None,
+        colors=None,
+        max_agents=1000,
+        blocking="BAS",
+        adjust_graph=True,
+    ):
         if not isinstance(blocking, str):
             raise TypeError("blocking must be a string")
 
@@ -317,7 +325,7 @@ def __init__(self, g, q_classes=None, q_args=None, seed=None, colors=None,
         self._initialized = False
         self._prev_edge = None
         self._fancy_heap = PriorityQueue()
-        self._blocking = True if blocking.lower() != 'rs' else False
+        self._blocking = True if blocking.lower() != "rs" else False
 
         if colors is None:
             colors = {}
@@ -341,8 +349,10 @@ def __init__(self, g, q_classes=None, q_args=None, seed=None, colors=None,
                 q_args[k] = {}
 
         for key, args in q_args.items():
-            if 'colors' not in args:
-                args['colors'] = self.default_q_colors.get(key, self.default_q_colors[1])
+            if "colors" not in args:
+                args["colors"] = self.default_q_colors.get(
+                    key, self.default_q_colors[1]
+                )
 
         if isinstance(seed, numbers.Integral):
             np.random.seed(seed)
@@ -361,7 +371,7 @@ def __init__(self, g, q_classes=None, q_args=None, seed=None, colors=None,
 
             for v in g.nodes():
                 vod = g.out_degree(v)
-                probs = array.array('d', [1. / vod for i in range(vod)])
+                probs = array.array("d", [1.0 / vod for i in range(vod)])
                 self.out_edges[v] = [g.edge_index[e] for e in sorted(g.out_edges(v))]
                 self.in_edges[v] = [g.edge_index[e] for e in sorted(g.in_edges(v))]
                 self._route_probs[v] = probs
@@ -370,18 +380,18 @@ def __init__(self, g, q_classes=None, q_args=None, seed=None, colors=None,
             self.g = g
 
     def __repr__(self):
-        the_string = 'QueueNetwork. # nodes: {0}, edges: {1}, agents: {2}'
+        the_string = "QueueNetwork. # nodes: {0}, edges: {1}, agents: {2}"
         return the_string.format(self.nV, self.nE, np.sum(self.num_agents))
 
     @property
     def blocking(self):
-        return 'BAS' if self._blocking else 'RS'
+        return "BAS" if self._blocking else "RS"
 
     @blocking.setter
     def blocking(self, tmp):
         if not isinstance(tmp, str):
             raise TypeError("blocking must be a string")
-        self._blocking = True if tmp.lower() != 'rs' else False
+        self._blocking = True if tmp.lower() != "rs" else False
 
     @property
     def num_vertices(self):
@@ -405,12 +415,14 @@ def time(self):
             e = self._fancy_heap.array_edges[0]
             t = self.edge2queue[e]._time
         else:
-            t = np.infty
+            t = np.inf
         return t
 
-    def animate(self, out=None, t=None, line_kwargs=None,
-                scatter_kwargs=None, **kwargs):
-        """Animates the network as it's simulating.
+    def animate(
+        self, out=None, t=None, line_kwargs=None, scatter_kwargs=None, **kwargs
+    ):
+        """
+        Animates the network as it's simulating.
 
         The animations can be saved to disk or viewed in interactive
         mode. Closing the window ends the animation if viewed in
@@ -497,11 +509,10 @@ def animate(self, out=None, t=None, line_kwargs=None,
         ...     'vertex_size': 15
         ... }
         >>> net.animate(**kwargs) # doctest: +SKIP
-        """
 
+        """
         if not self._initialized:
-            msg = ("Network has not been initialized. "
-                   "Call '.initialize()' first.")
+            msg = "Network has not been initialized. Call '.initialize()' first."
             raise QueueingToolError(msg)
 
         if not HAS_MATPLOTLIB:
@@ -509,15 +520,15 @@ def animate(self, out=None, t=None, line_kwargs=None,
             raise ImportError(msg)
 
         self._update_all_colors()
-        kwargs.setdefault('bgcolor', self.colors['bgcolor'])
+        kwargs.setdefault("bgcolor", self.colors["bgcolor"])
 
-        fig = plt.figure(figsize=kwargs.get('figsize', (7, 7)))
+        fig = plt.figure(figsize=kwargs.get("figsize", (7, 7)))
         ax = fig.gca()
 
         mpl_kwargs = {
-            'line_kwargs': line_kwargs,
-            'scatter_kwargs': scatter_kwargs,
-            'pos': kwargs.get('pos')
+            "line_kwargs": line_kwargs,
+            "scatter_kwargs": scatter_kwargs,
+            "pos": kwargs.get("pos"),
         }
 
         line_args, scat_args = self.g.lines_scatter_args(**mpl_kwargs)
@@ -526,7 +537,7 @@ def animate(self, out=None, t=None, line_kwargs=None,
         lines = ax.add_collection(lines)
         scatt = ax.scatter(**scat_args)
 
-        t = np.infty if t is None else t
+        t = np.inf if t is None else t
         now = self._t
 
         def update(frame_number):
@@ -534,51 +545,51 @@ def update(frame_number):
                 if self._t > now + t:
                     return False
             self._simulate_next_event(slow=True)
-            lines.set_color(line_args['colors'])
-            scatt.set_edgecolors(scat_args['edgecolors'])
-            scatt.set_facecolor(scat_args['c'])
+            lines.set_color(line_args["colors"])
+            scatt.set_edgecolors(scat_args["edgecolors"])
+            scatt.set_facecolor(scat_args["c"])
 
-        if hasattr(ax, 'set_facecolor'):
-            ax.set_facecolor(kwargs['bgcolor'])
+        if hasattr(ax, "set_facecolor"):
+            ax.set_facecolor(kwargs["bgcolor"])
         else:
-            ax.set_axis_bgcolor(kwargs['bgcolor'])
+            ax.set_axis_bgcolor(kwargs["bgcolor"])
 
         ax.get_xaxis().set_visible(False)
         ax.get_yaxis().set_visible(False)
 
         animation_args = {
-            'fargs': None,
-            'event_source': None,
-            'init_func': None,
-            'frames': None,
-            'blit': False,
-            'interval': 10,
-            'repeat': None,
-            'func': update,
-            'repeat_delay': None,
-            'fig': fig,
-            'save_count': None,
+            "fargs": None,
+            "event_source": None,
+            "init_func": None,
+            "frames": None,
+            "blit": False,
+            "interval": 10,
+            "repeat": None,
+            "func": update,
+            "repeat_delay": None,
+            "fig": fig,
+            "save_count": None,
         }
         for key, value in kwargs.items():
             if key in animation_args:
                 animation_args[key] = value
 
         animation = FuncAnimation(**animation_args)
-        if 'filename' not in kwargs:
+        if "filename" not in kwargs:
             plt.ioff()
             plt.show()
         else:
             save_args = {
-                'filename': None,
-                'writer': None,
-                'fps': None,
-                'dpi': None,
-                'codec': None,
-                'bitrate': None,
-                'extra_args': None,
-                'metadata': None,
-                'extra_anim': None,
-                'savefig_kwargs': None
+                "filename": None,
+                "writer": None,
+                "fps": None,
+                "dpi": None,
+                "codec": None,
+                "bitrate": None,
+                "extra_args": None,
+                "metadata": None,
+                "extra_anim": None,
+                "savefig_kwargs": None,
             }
             for key, value in kwargs.items():
                 if key in save_args:
@@ -587,7 +598,8 @@ def update(frame_number):
             animation.save(**save_args)
 
     def clear(self):
-        """Resets the queue to its initial state.
+        """
+        Resets the queue to its initial state.
 
         The attributes ``t``, ``num_events``, ``num_agents`` are set to
         zero, :meth:`.reset_colors` is called, and the
@@ -598,6 +610,7 @@ def clear(self):
         -----
         ``QueueNetwork`` must be re-initialized before any simulations
         can run.
+
         """
         self._t = 0
         self.num_events = 0
@@ -610,7 +623,8 @@ def clear(self):
             q.clear()
 
     def clear_data(self, queues=None, edge=None, edge_type=None):
-        """Clears data from all queues.
+        """
+        Clears data from all queues.
 
         If none of the parameters are given then every queue's data is
         cleared.
@@ -633,6 +647,7 @@ def clear_data(self, queues=None, edge=None, edge_type=None):
         edge_type : int or an iterable of int (optional)
             A integer, or a collection of integers identifying which
             edge types will have their data cleared.
+
         """
         queues = _get_queues(self.g, queues, edge, edge_type)
 
@@ -659,14 +674,14 @@ def copy(self):
         net._route_probs = copy.deepcopy(self._route_probs)
 
         if net._initialized:
-            keys = [q._key() for q in net.edge2queue if q._time < np.infty]
+            keys = [q._key() for q in net.edge2queue if q._time < np.inf]
             net._fancy_heap = PriorityQueue(keys, net.nE)
 
         return net
 
-    def draw(self, update_colors=True, line_kwargs=None,
-             scatter_kwargs=None, **kwargs):
-        """Draws the network. The coloring of the network corresponds
+    def draw(self, update_colors=True, line_kwargs=None, scatter_kwargs=None, **kwargs):
+        """
+        Draws the network. The coloring of the network corresponds
         to the number of agents at each queue.
 
         Parameters
@@ -727,6 +742,7 @@ def draw(self, update_colors=True, line_kwargs=None,
         following.
 
         >>> net.draw(line_kwargs={'linestyle': 'dashed'}) # doctest: +SKIP
+
         """
         if not HAS_MATPLOTLIB:
             raise ImportError("matplotlib is necessary to draw the network.")
@@ -734,14 +750,18 @@ def draw(self, update_colors=True, line_kwargs=None,
         if update_colors:
             self._update_all_colors()
 
-        if 'bgcolor' not in kwargs:
-            kwargs['bgcolor'] = self.colors['bgcolor']
+        if "bgcolor" not in kwargs:
+            kwargs["bgcolor"] = self.colors["bgcolor"]
 
-        self.g.draw_graph(line_kwargs=line_kwargs,
-                          scatter_kwargs=scatter_kwargs, **kwargs)
+        self.g.draw_graph(
+            line_kwargs=line_kwargs, scatter_kwargs=scatter_kwargs, **kwargs
+        )
 
-    def get_agent_data(self, queues=None, edge=None, edge_type=None, return_header=False):
-        """Gets data from queues and organizes it by agent.
+    def get_agent_data(
+        self, queues=None, edge=None, edge_type=None, return_header=False
+    ):
+        """
+        Gets data from queues and organizes it by agent.
 
         If none of the parameters are given then data from every
         :class:`.QueueServer` is retrieved.
@@ -788,6 +808,7 @@ def get_agent_data(self, queues=None, edge=None, edge_type=None, return_header=F
         headers : str (optional)
             A comma seperated string of the column headers. Returns
             ``'arrival,service,departure,num_queued,num_total,q_id'``
+
         """
         queues = _get_queues(self.g, queues, edge, edge_type)
 
@@ -803,25 +824,28 @@ def get_agent_data(self, queues=None, edge=None, edge_type=None, return_header=F
                     data[agent_id] = datum
 
         dType = [
-            ('a', float),
-            ('s', float),
-            ('d', float),
-            ('q', float),
-            ('n', float),
-            ('id', float)
+            ("a", float),
+            ("s", float),
+            ("d", float),
+            ("q", float),
+            ("n", float),
+            ("id", float),
         ]
         for agent_id, dat in data.items():
             datum = np.array([tuple(d) for d in dat.tolist()], dtype=dType)
-            datum = np.sort(datum, order='a')
+            datum = np.sort(datum, order="a")
             data[agent_id] = np.array([tuple(d) for d in datum])
 
         if return_header:
-            return data, 'arrival,service,departure,num_queued,num_total,q_id'
+            return data, "arrival,service,departure,num_queued,num_total,q_id"
 
         return data
 
-    def get_queue_data(self, queues=None, edge=None, edge_type=None, return_header=False):
-        """Gets data from all the queues.
+    def get_queue_data(
+        self, queues=None, edge=None, edge_type=None, return_header=False
+    ):
+        """
+        Gets data from all the queues.
 
         If none of the parameters are given then data from every
         :class:`.QueueServer` is retrieved.
@@ -888,6 +912,7 @@ def get_queue_data(self, queues=None, edge=None, edge_type=None, return_header=F
         You can specify the edge indices as well:
 
         >>> data = net.get_queue_data(queues=(20, 14, 0, 4))
+
         """
         queues = _get_queues(self.g, queues, edge, edge_type)
 
@@ -899,12 +924,13 @@ def get_queue_data(self, queues=None, edge=None, edge_type=None, return_header=F
                 data = np.vstack((data, dat))
 
         if return_header:
-            return data, 'arrival,service,departure,num_queued,num_total,q_id'
+            return data, "arrival,service,departure,num_queued,num_total,q_id"
 
         return data
 
     def initialize(self, nActive=1, queues=None, edges=None, edge_type=None):
-        """Prepares the ``QueueNetwork`` for simulation.
+        """
+        Prepares the ``QueueNetwork`` for simulation.
 
         Each :class:`.QueueServer` in the network starts inactive,
         which means they do not accept arrivals from outside the
@@ -952,6 +978,7 @@ def initialize(self, nActive=1, queues=None, edges=None, edge_type=None):
         :class:`NullQueues<.NullQueue>` cannot be activated, and are
         sifted out if they are specified. More specifically, every edge
         with edge type 0 is sifted out.
+
         """
         if queues is None and edges is None and edge_type is None:
             if nActive >= 1 and isinstance(nActive, numbers.Integral):
@@ -962,8 +989,7 @@ def initialize(self, nActive=1, queues=None, edges=None, edge_type=None):
                 msg = "If queues is None, then nActive must be an integer."
                 raise TypeError(msg)
             else:
-                msg = ("If queues is None, then nActive must be a "
-                       "positive int.")
+                msg = "If queues is None, then nActive must be a positive int."
                 raise ValueError(msg)
         else:
             queues = _get_queues(self.g, queues, edges, edge_type)
@@ -974,18 +1000,19 @@ def initialize(self, nActive=1, queues=None, edges=None, edge_type=None):
             raise QueueingToolError("There were no queues to initialize.")
 
         if len(queues) > self.max_agents:
-            queues = queues[:self.max_agents]
+            queues = queues[: self.max_agents]
 
         for ei in queues:
             self.edge2queue[ei].set_active()
             self.num_agents[ei] = self.edge2queue[ei]._num_total
 
-        keys = [q._key() for q in self.edge2queue if q._time < np.infty]
+        keys = [q._key() for q in self.edge2queue if q._time < np.inf]
         self._fancy_heap = PriorityQueue(keys, self.nE)
         self._initialized = True
 
     def next_event_description(self):
-        """Returns whether the next event is an arrival or a departure
+        """
+        Returns whether the next event is an arrival or a departure
         and the queue the event is accuring at.
 
         Returns
@@ -997,9 +1024,10 @@ def next_event_description(self):
         edge : int or ``None``
             The edge index of the edge that this event will occur at.
             If there are no events then ``None`` is returned.
+
         """
         if self._fancy_heap.size == 0:
-            event_type = 'Nothing'
+            event_type = "Nothing"
             edge_index = None
         else:
             s = [q._key() for q in self.edge2queue]
@@ -1007,19 +1035,20 @@ def next_event_description(self):
             e = s[0][1]
             q = self.edge2queue[e]
 
-            event_type = 'Arrival' if q.next_event_description() == 1 else 'Departure'
+            event_type = "Arrival" if q.next_event_description() == 1 else "Departure"
             edge_index = q.edge[2]
         return event_type, edge_index
 
     def reset_colors(self):
         """Resets all edge and vertex colors to their default values."""
         for k, e in enumerate(self.g.edges()):
-            self.g.set_ep(e, 'edge_color', self.edge2queue[k].colors['edge_color'])
+            self.g.set_ep(e, "edge_color", self.edge2queue[k].colors["edge_color"])
         for v in self.g.nodes():
-            self.g.set_vp(v, 'vertex_fill_color', self.colors['vertex_fill_color'])
+            self.g.set_vp(v, "vertex_fill_color", self.colors["vertex_fill_color"])
 
     def set_transitions(self, mat):
-        """Change the routing transitions probabilities for the
+        """
+        Change the routing transitions probabilities for the
         network.
 
         Parameters
@@ -1099,6 +1128,7 @@ def set_transitions(self, mat):
             probabilities.
         :func:`.generate_transition_matrix` : Generate a random routing
             matrix.
+
         """
         if isinstance(mat, dict):
             for key, value in mat.items():
@@ -1107,10 +1137,10 @@ def set_transitions(self, mat):
                 if key not in self.g.nodes:
                     msg = "One of the keys don't correspond to a vertex."
                     raise ValueError(msg)
-                elif len(self.out_edges[key]) > 0 and not np.isclose(sum(probs), 1):
+                if len(self.out_edges[key]) > 0 and not np.isclose(sum(probs), 1):
                     msg = "Sum of transition probabilities at a vertex was not 1."
                     raise ValueError(msg)
-                elif (np.array(probs) < 0).any():
+                if (np.array(probs) < 0).any():
                     msg = "Some transition probabilities were negative."
                     raise ValueError(msg)
 
@@ -1120,13 +1150,12 @@ def set_transitions(self, mat):
         elif isinstance(mat, np.ndarray):
             non_terminal = np.array([self.g.out_degree(v) > 0 for v in self.g.nodes()])
             if mat.shape != (self.nV, self.nV):
-                msg = ("Matrix is the wrong shape, should "
-                       "be {0} x {1}.").format(self.nV, self.nV)
+                msg = (f"Matrix is the wrong shape, should be {self.nV} x {self.nV}.")
                 raise ValueError(msg)
-            elif not np.allclose(np.sum(mat[non_terminal, :], axis=1), 1):
+            if not np.allclose(np.sum(mat[non_terminal, :], axis=1), 1):
                 msg = "Sum of transition probabilities at a vertex was not 1."
                 raise ValueError(msg)
-            elif (mat < 0).any():
+            if (mat < 0).any():
                 raise ValueError("Some transition probabilities were negative.")
 
             for k in range(self.nV):
@@ -1136,7 +1165,8 @@ def set_transitions(self, mat):
             raise TypeError("mat must be a numpy array or a dict.")
 
     def show_active(self, **kwargs):
-        """Draws the network, highlighting active queues.
+        """
+        Draws the network, highlighting active queues.
 
         The colored vertices represent vertices that have at least one
         queue on an in-edge that is active. Dark edges represent
@@ -1156,10 +1186,11 @@ def show_active(self, **kwargs):
         defined by the class attribute ``colors``. The relevant keys
         are ``vertex_active``, ``vertex_inactive``, ``edge_active``,
         and ``edge_inactive``.
+
         """
         g = self.g
         for v in g.nodes():
-            self.g.set_vp(v, 'vertex_color', [0, 0, 0, 0.9])
+            self.g.set_vp(v, "vertex_color", [0, 0, 0, 0.9])
             is_active = False
             my_iter = g.in_edges(v) if g.is_directed() else g.out_edges(v)
             for e in my_iter:
@@ -1168,22 +1199,23 @@ def show_active(self, **kwargs):
                     is_active = True
                     break
             if is_active:
-                self.g.set_vp(v, 'vertex_fill_color', self.colors['vertex_active'])
+                self.g.set_vp(v, "vertex_fill_color", self.colors["vertex_active"])
             else:
-                self.g.set_vp(v, 'vertex_fill_color', self.colors['vertex_inactive'])
+                self.g.set_vp(v, "vertex_fill_color", self.colors["vertex_inactive"])
 
         for e in g.edges():
             ei = g.edge_index[e]
             if self.edge2queue[ei]._active:
-                self.g.set_ep(e, 'edge_color', self.colors['edge_active'])
+                self.g.set_ep(e, "edge_color", self.colors["edge_active"])
             else:
-                self.g.set_ep(e, 'edge_color', self.colors['edge_inactive'])
+                self.g.set_ep(e, "edge_color", self.colors["edge_inactive"])
 
         self.draw(update_colors=False, **kwargs)
         self._update_all_colors()
 
     def show_type(self, edge_type, **kwargs):
-        """Draws the network, highlighting queues of a certain type.
+        """
+        Draws the network, highlighting queues of a certain type.
 
         The colored vertices represent self loops of type ``edge_type``.
         Dark edges represent queues of type ``edge_type``.
@@ -1216,28 +1248,32 @@ def show_type(self, edge_type, **kwargs):
 
         .. figure:: edge_type_2-1.png
            :align: center
+
         """
         for v in self.g.nodes():
             e = (v, v)
-            if self.g.is_edge(e) and self.g.ep(e, 'edge_type') == edge_type:
+            if self.g.is_edge(e) and self.g.ep(e, "edge_type") == edge_type:
                 ei = self.g.edge_index[e]
-                self.g.set_vp(v, 'vertex_fill_color', self.colors['vertex_highlight'])
-                self.g.set_vp(v, 'vertex_color', self.edge2queue[ei].colors['vertex_color'])
+                self.g.set_vp(v, "vertex_fill_color", self.colors["vertex_highlight"])
+                self.g.set_vp(
+                    v, "vertex_color", self.edge2queue[ei].colors["vertex_color"]
+                )
             else:
-                self.g.set_vp(v, 'vertex_fill_color', self.colors['vertex_inactive'])
-                self.g.set_vp(v, 'vertex_color', [0, 0, 0, 0.9])
+                self.g.set_vp(v, "vertex_fill_color", self.colors["vertex_inactive"])
+                self.g.set_vp(v, "vertex_color", [0, 0, 0, 0.9])
 
         for e in self.g.edges():
-            if self.g.ep(e, 'edge_type') == edge_type:
-                self.g.set_ep(e, 'edge_color', self.colors['edge_active'])
+            if self.g.ep(e, "edge_type") == edge_type:
+                self.g.set_ep(e, "edge_color", self.colors["edge_active"])
             else:
-                self.g.set_ep(e, 'edge_color', self.colors['edge_inactive'])
+                self.g.set_ep(e, "edge_color", self.colors["edge_inactive"])
 
         self.draw(update_colors=False, **kwargs)
         self._update_all_colors()
 
     def simulate(self, n=1, t=None):
-        """Simulates the network forward.
+        """
+        Simulates the network forward.
 
         Simulates either a specific number of events or for a specified
         amount of simulation time.
@@ -1286,10 +1322,10 @@ def simulate(self, n=1, t=None):
         >>> t1 = net.current_time
         >>> t1 - t0 # doctest: +ELLIPSIS
         75...
+
         """
         if not self._initialized:
-            msg = ("Network has not been initialized. "
-                   "Call '.initialize()' first.")
+            msg = "Network has not been initialized. Call '.initialize()' first."
             raise QueueingToolError(msg)
         if t is None:
             for dummy in range(n):
@@ -1301,7 +1337,7 @@ def simulate(self, n=1, t=None):
 
     def _simulate_next_event(self, slow=True):
         if self._fancy_heap.size == 0:
-            self._t = np.infty
+            self._t = np.inf
             return
 
         q1k = self._fancy_heap.pop()
@@ -1339,8 +1375,11 @@ def _simulate_next_event(self, slow=True):
                     self._update_graph_colors(qedge=q1.edge)
                     self._prev_edge = q1.edge
 
-                if q2._active and self.max_agents < np.infty and \
-                        np.sum(self.num_agents) > self.max_agents - 1:
+                if (
+                    q2._active
+                    and self.max_agents < np.inf
+                    and np.sum(self.num_agents) > self.max_agents - 1
+                ):
                     q2._active = False
 
                 q2.next_event()
@@ -1356,15 +1395,17 @@ def _simulate_next_event(self, slow=True):
             if new_q2k[0] != q2k[0]:
                 self._fancy_heap.push(*new_q2k)
 
-                if new_q1k[0] < np.infty and new_q1k != new_q2k:
-                    self._fancy_heap.push(*new_q1k)
-            else:
-                if new_q1k[0] < np.infty:
+                if new_q1k[0] < np.inf and new_q1k != new_q2k:
                     self._fancy_heap.push(*new_q1k)
+            elif new_q1k[0] < np.inf:
+                self._fancy_heap.push(*new_q1k)
 
         elif event == 1:  # This is an arrival
-            if q1._active and self.max_agents < np.infty and \
-                    np.sum(self.num_agents) > self.max_agents - 1:
+            if (
+                q1._active
+                and self.max_agents < np.inf
+                and np.sum(self.num_agents) > self.max_agents - 1
+            ):
                 q1._active = False
 
             q1.next_event()
@@ -1375,11 +1416,12 @@ def _simulate_next_event(self, slow=True):
                 self._prev_edge = q1.edge
 
             new_q1k = q1._key()
-            if new_q1k[0] < np.infty:
+            if new_q1k[0] < np.inf:
                 self._fancy_heap.push(*new_q1k)
 
     def start_collecting_data(self, queues=None, edge=None, edge_type=None):
-        """Tells the queues to collect data on agents' arrival, service
+        """
+        Tells the queues to collect data on agents' arrival, service
         start, and departure times.
 
         If none of the parameters are given then every
@@ -1403,6 +1445,7 @@ def start_collecting_data(self, queues=None, edge=None, edge_type=None):
         edge_type : int or an iterable of int (optional)
             A integer, or a collection of integers identifying which
             edge types will be set active.
+
         """
         queues = _get_queues(self.g, queues, edge, edge_type)
 
@@ -1410,7 +1453,8 @@ def start_collecting_data(self, queues=None, edge=None, edge_type=None):
             self.edge2queue[k].collect_data = True
 
     def stop_collecting_data(self, queues=None, edge=None, edge_type=None):
-        """Tells the queues to stop collecting data on agents.
+        """
+        Tells the queues to stop collecting data on agents.
 
         If none of the parameters are given then every
         :class:`.QueueServer` will stop collecting data.
@@ -1433,6 +1477,7 @@ def stop_collecting_data(self, queues=None, edge=None, edge_type=None):
         edge_type : int or an iterable of int (optional)
             A integer, or a collection of integers identifying which
             edge types will stop collecting data.
+
         """
         queues = _get_queues(self.g, queues, edge, edge_type)
 
@@ -1440,7 +1485,8 @@ def stop_collecting_data(self, queues=None, edge=None, edge_type=None):
             self.edge2queue[k].collect_data = False
 
     def transitions(self, return_matrix=True):
-        """Returns the routing probabilities for each vertex in the
+        """
+        Returns the routing probabilities for each vertex in the
         graph.
 
         Parameters
@@ -1517,6 +1563,7 @@ def transitions(self, return_matrix=True):
         ``0.438``, when at vertex ``6`` they will transition back to
         vertex ``2`` with probability ``0.673`` and route vertex ``4``
         probability ``0.326``, etc.
+
         """
         if return_matrix:
             mat = np.zeros((self.nV, self.nV))
@@ -1537,13 +1584,13 @@ def _update_all_colors(self):
             e = q.edge[:2]
             v = q.edge[1]
             if q.edge[0] == q.edge[1]:
-                self.g.set_ep(e, 'edge_color', q._current_color(1))
-                self.g.set_vp(v, 'vertex_color', q._current_color(2))
+                self.g.set_ep(e, "edge_color", q._current_color(1))
+                self.g.set_vp(v, "vertex_color", q._current_color(2))
                 if q.edge[3] != 0:
-                    self.g.set_vp(v, 'vertex_fill_color', q._current_color())
+                    self.g.set_vp(v, "vertex_fill_color", q._current_color())
                 do[v] = False
             else:
-                self.g.set_ep(e, 'edge_color', q._current_color())
+                self.g.set_ep(e, "edge_color", q._current_color())
                 if do[v]:
                     self._update_vertex_color(v)
                     do[v] = False
@@ -1563,14 +1610,14 @@ def _update_vertex_color(self, v):
                 nSy += self.edge2queue[ei].num_system
                 cap += self.edge2queue[ei].num_servers
 
-            div = (2 * cap) + 2.
-            tmp = 1. - min(nSy / div, 1.)
+            div = (2 * cap) + 2.0
+            tmp = 1.0 - min(nSy / div, 1.0)
 
-            color = [i * tmp for i in self.colors['vertex_fill_color']]
+            color = [i * tmp for i in self.colors["vertex_fill_color"]]
             color[3] = 1.0
-            self.g.set_vp(v, 'vertex_fill_color', color)
+            self.g.set_vp(v, "vertex_fill_color", color)
             if not ee_is_edge:
-                self.g.set_vp(v, 'vertex_color', self.colors['vertex_color'])
+                self.g.set_vp(v, "vertex_color", self.colors["vertex_color"])
 
     def _update_graph_colors(self, qedge):
         e = qedge[:2]
@@ -1581,24 +1628,24 @@ def _update_graph_colors(self, qedge):
             q = self.edge2queue[self._prev_edge[2]]
 
             if pe[0] == pe[1]:
-                self.g.set_ep(pe, 'edge_color', q._current_color(1))
-                self.g.set_vp(pv, 'vertex_color', q._current_color(2))
+                self.g.set_ep(pe, "edge_color", q._current_color(1))
+                self.g.set_vp(pv, "vertex_color", q._current_color(2))
                 if q.edge[3] != 0:
-                    self.g.set_vp(v, 'vertex_fill_color', q._current_color())
+                    self.g.set_vp(v, "vertex_fill_color", q._current_color())
 
             else:
-                self.g.set_ep(pe, 'edge_color', q._current_color())
+                self.g.set_ep(pe, "edge_color", q._current_color())
                 self._update_vertex_color(pv)
 
         q = self.edge2queue[qedge[2]]
         if qedge[0] == qedge[1]:
-            self.g.set_ep(e, 'edge_color', q._current_color(1))
-            self.g.set_vp(v, 'vertex_color', q._current_color(2))
+            self.g.set_ep(e, "edge_color", q._current_color(1))
+            self.g.set_vp(v, "vertex_color", q._current_color(2))
             if q.edge[3] != 0:
-                self.g.set_vp(v, 'vertex_fill_color', q._current_color())
+                self.g.set_vp(v, "vertex_fill_color", q._current_color())
 
         else:
-            self.g.set_ep(e, 'edge_color', q._current_color())
+            self.g.set_ep(e, "edge_color", q._current_color())
             self._update_vertex_color(v)
 
 
@@ -1625,7 +1672,7 @@ def _get_queues(g, queues, edge, edge_type):
                 edge_type = set([edge_type])
             tmp = []
             for e in g.edges():
-                if g.ep(e, 'edge_type') in edge_type:
+                if g.ep(e, "edge_type") in edge_type:
                     tmp.append(g.edge_index[e])
 
             queues = np.array(tmp, int)
diff --git a/queueing_tool/queues/__init__.py b/queueing_tool/queues/__init__.py
index 41a2963..5e6bd21 100644
--- a/queueing_tool/queues/__init__.py
+++ b/queueing_tool/queues/__init__.py
@@ -14,32 +14,29 @@
     ResourceQueue
 """
 
-from queueing_tool.queues.queue_servers import (
-    QueueServer,
-    LossQueue,
-    NullQueue,
-    poisson_random_measure
-)
-from queueing_tool.queues.agents import (
-    Agent,
-    GreedyAgent
-)
+from queueing_tool.queues.agents import Agent, GreedyAgent
 from queueing_tool.queues.queue_extentions import (
+    InfoAgent,
+    InfoQueue,
     ResourceAgent,
     ResourceQueue,
-    InfoAgent,
-    InfoQueue
+)
+from queueing_tool.queues.queue_servers import (
+    LossQueue,
+    NullQueue,
+    QueueServer,
+    poisson_random_measure,
 )
 
 __all__ = [
-    'InfoQueue',
-    'LossQueue',
-    'NullQueue',
-    'QueueServer',
-    'ResourceQueue',
-    'poisson_random_measure',
-    'Agent',
-    'GreedyAgent',
-    'InfoAgent',
-    'ResourceAgent'
+    "Agent",
+    "GreedyAgent",
+    "InfoAgent",
+    "InfoQueue",
+    "LossQueue",
+    "NullQueue",
+    "QueueServer",
+    "ResourceAgent",
+    "ResourceQueue",
+    "poisson_random_measure",
 ]
diff --git a/queueing_tool/queues/agents.py b/queueing_tool/queues/agents.py
index 74435bc..1204cb1 100644
--- a/queueing_tool/queues/agents.py
+++ b/queueing_tool/queues/agents.py
@@ -1,11 +1,12 @@
-from numpy import infty
+from numpy import inf
 from numpy.random import uniform
 
-from queueing_tool.queues.choice import _choice, _argmin
+from queueing_tool.queues.choice import _argmin, _choice
 
 
-class Agent(object):
-    """The base class for an agent.
+class Agent:
+    """
+    The base class for an agent.
 
     ``Agents`` are the objects that move throughout the network.
     ``Agents`` are instantiated by a queue, and once serviced the
@@ -34,14 +35,16 @@ class Agent(object):
     blocked : int
         Specifies how many times an agent has been blocked by a finite
         capacity queue.
+
     """
+
     def __init__(self, agent_id=(0, 0), **kwargs):
         self.agent_id = agent_id
         self.blocked = 0
         self._time = 0  # The agents arrival or departure time
 
     def __repr__(self):
-        return "Agent; agent_id:{0}. time: {1}".format(self.agent_id, round(self._time, 3))
+        return f"Agent; agent_id:{self.agent_id}. time: {round(self._time, 3)}"
 
     def __lt__(self, b):
         return self._time < b._time
@@ -59,13 +62,15 @@ def __ge__(self, b):
         return self._time >= b._time
 
     def add_loss(self, *args, **kwargs):
-        """Adds one to the number of times the agent has been blocked
+        """
+        Adds one to the number of times the agent has been blocked
         from entering a queue.
         """
         self.blocked += 1
 
     def desired_destination(self, network, edge):
-        """Returns the agents next destination given their current
+        """
+        Returns the agents next destination given their current
         location on the network.
 
         An ``Agent`` chooses one of the out edges at random. The
@@ -95,6 +100,7 @@ def desired_destination(self, network, edge):
         :meth:`.transitions` : :class:`QueueNetwork's<.QueueNetwork>`
             method that returns the transition probabilities for each
             edge in the graph.
+
         """
         n = len(network.out_edges[edge[1]])
         if n <= 1:
@@ -109,7 +115,8 @@ def desired_destination(self, network, edge):
         return network.out_edges[edge[1]][k]
 
     def queue_action(self, queue, *args, **kwargs):
-        """A method that acts on the queue the Agent is at. This method
+        """
+        A method that acts on the queue the Agent is at. This method
         is called when the Agent arrives at the queue (where
         ``args[0] == 0``), when service starts for the Agent (where
         ``args[0] == 1``), and when the Agent departs from the queue
@@ -117,11 +124,11 @@ def queue_action(self, queue, *args, **kwargs):
         to the queue, but is here if the Agent class is extended and
         this method is overwritten.
         """
-        pass
 
 
 class GreedyAgent(Agent):
-    """An agent that chooses the queue with the shortest line as their
+    """
+    An agent that chooses the queue with the shortest line as their
     next destination.
 
     Notes
@@ -129,7 +136,9 @@ class GreedyAgent(Agent):
     If there are any ties, the ``GreedyAgent`` chooses the first queue
     with the shortest line (where the ordering is given by
     :class:`QueueNetwork's<.QueueNetwork>` ``out_edges`` attribute).
+
     """
+
     def __init__(self, agent_id=(0, 0)):
         Agent.__init__(self, agent_id)
 
@@ -138,7 +147,8 @@ def __repr__(self):
         return msg.format(self.agent_id, round(self._time, 3))
 
     def desired_destination(self, network, edge):
-        """Returns the agents next destination given their current
+        """
+        Returns the agents next destination given their current
         location on the network.
 
         ``GreedyAgents`` choose their next destination with-in the
@@ -161,20 +171,23 @@ def desired_destination(self, network, edge):
         out : int
             Returns an the edge index corresponding to the agents next
             edge to visit in the network.
+
         """
         adjacent_edges = network.out_edges[edge[1]]
         d = _argmin([network.edge2queue[d].number_queued() for d in adjacent_edges])
         return adjacent_edges[d]
 
 
-class InftyAgent(object):
-    """An special agent that only operates within the
+class InftyAgent:
+    """
+    An special agent that only operates within the
     :class:`.QueueServer` class.
 
     This agent never interacts with the :class:`.QueueNetwork`.
     """
+
     def __init__(self):
-        self._time = infty
+        self._time = inf
 
     def __repr__(self):
         return "InftyAgent"
diff --git a/queueing_tool/queues/queue_extentions.py b/queueing_tool/queues/queue_extentions.py
index 2a5150f..6f9314e 100644
--- a/queueing_tool/queues/queue_extentions.py
+++ b/queueing_tool/queues/queue_extentions.py
@@ -1,14 +1,15 @@
-from heapq import heappush, heappop
+from heapq import heappop, heappush
 
-from numpy import logical_or, infty
 import numpy as np
+from numpy import logical_or
 
 from queueing_tool.queues.agents import Agent
 from queueing_tool.queues.queue_servers import LossQueue
 
 
 class ResourceAgent(Agent):
-    """An agent designed to interact with the :class:`.ResourceQueue`
+    """
+    An agent designed to interact with the :class:`.ResourceQueue`
     class.
 
     When an ``ResourceAgent`` departs from a :class:`.ResourceQueue`,
@@ -19,16 +20,18 @@ class ResourceAgent(Agent):
     a resource to that queue by increasing the number of
     servers there by one; the ``ResourceAgent`` is then deleted.
     """
+
     def __init__(self, agent_id=(0, 0)):
         super(ResourceAgent, self).__init__(agent_id)
         self._has_resource = False
         self._had_resource = False
 
     def __repr__(self):
-        return "ResourceAgent; agent_id:{0}. Time: {1}".format(self.agent_id, round(self._time, 3))
+        return f"ResourceAgent; agent_id:{self.agent_id}. Time: {round(self._time, 3)}"
 
     def queue_action(self, queue, *args, **kwargs):
-        """Function that specifies the interaction with a
+        """
+        Function that specifies the interaction with a
         :class:`.ResourceQueue` upon departure.
 
         When departuring from a :class:`.ResourceQueue` (or a
@@ -42,20 +45,21 @@ def queue_action(self, queue, *args, **kwargs):
         queue : :class:`.QueueServer`
             The instance of the queue that the ``ResourceAgent`` will
             interact with.
+
         """
         if isinstance(queue, ResourceQueue):
             if self._has_resource:
                 self._has_resource = False
                 self._had_resource = True
-            else:
-                if queue.num_servers > 0:
-                    queue.set_num_servers(queue.num_servers - 1)
-                    self._has_resource = True
-                    self._had_resource = False
+            elif queue.num_servers > 0:
+                queue.set_num_servers(queue.num_servers - 1)
+                self._has_resource = True
+                self._had_resource = False
 
 
 class ResourceQueue(LossQueue):
-    """An queue designed to interact with the :class:`.ResourceAgent`
+    """
+    An queue designed to interact with the :class:`.ResourceAgent`
     class.
 
     If a :class:`.ResourceAgent` does not have a resource already it
@@ -76,13 +80,14 @@ class ResourceQueue(LossQueue):
         the number of servers was at ``max_servers``.
     kwargs
         Any arguments to pass to :class:`.LossQueue`.
+
     """
 
     _default_colors = {
-        'edge_loop_color': [0.7, 0.7, 0.7, 0.50],
-        'edge_color': [0.7, 0.7, 0.7, 0.50],
-        'vertex_fill_color': [1.0, 1.0, 1.0, 1.0],
-        'vertex_pen_color': [0.0, 0.235, 0.718, 1.0]
+        "edge_loop_color": [0.7, 0.7, 0.7, 0.50],
+        "edge_color": [0.7, 0.7, 0.7, 0.50],
+        "vertex_fill_color": [1.0, 1.0, 1.0, 1.0],
+        "vertex_pen_color": [0.0, 0.235, 0.718, 1.0],
     }
 
     def __init__(self, num_servers=10, AgentFactory=ResourceAgent, qbuffer=0, **kwargs):
@@ -90,17 +95,25 @@ def __init__(self, num_servers=10, AgentFactory=ResourceAgent, qbuffer=0, **kwar
             num_servers=num_servers,
             AgentFactory=AgentFactory,
             qbuffer=qbuffer,
-            **kwargs
+            **kwargs,
         )
 
         self.max_servers = 2 * num_servers
         self.over_max = 0
 
     def __repr__(self):
-        my_str = ("ResourceQueue:{0}. Servers: {1}, max servers: {2}, "
-                  "arrivals: {3}, departures: {4}, next time: {5}")
-        arg = (self.edge[2], self.num_servers, self.max_servers,
-               self.num_arrivals, self.num_departures, round(self._time, 3))
+        my_str = (
+            "ResourceQueue:{0}. Servers: {1}, max servers: {2}, "
+            "arrivals: {3}, departures: {4}, next time: {5}"
+        )
+        arg = (
+            self.edge[2],
+            self.num_servers,
+            self.max_servers,
+            self.num_arrivals,
+            self.num_departures,
+            round(self._time, 3),
+        )
         return my_str.format(*arg)
 
     def set_num_servers(self, n):
@@ -109,7 +122,8 @@ def set_num_servers(self, n):
             self.over_max += 1
 
     def next_event(self):
-        """Simulates the queue forward one event.
+        """
+        Simulates the queue forward one event.
 
         This method behaves identically to a :class:`.LossQueue` if the
         arriving/departing agent is anything other than a
@@ -133,7 +147,7 @@ def next_event(self):
         if isinstance(self._arrivals[0], ResourceAgent):
             if self._departures[0]._time < self._arrivals[0]._time:
                 return super(ResourceQueue, self).next_event()
-            elif self._arrivals[0]._time < infty:
+            if self._arrivals[0]._time < np.inf:
                 if self._arrivals[0]._has_resource:
                     arrival = heappop(self._arrivals)
                     self._current_t = arrival._time
@@ -143,9 +157,13 @@ def next_event(self):
                     if self.collect_data:
                         t = arrival._time
                         if arrival.agent_id not in self.data:
-                            self.data[arrival.agent_id] = [[t, t, t, len(self.queue), self.num_system]]
+                            self.data[arrival.agent_id] = [
+                                [t, t, t, len(self.queue), self.num_system]
+                            ]
                         else:
-                            self.data[arrival.agent_id].append([t, t, t, len(self.queue), self.num_system])
+                            self.data[arrival.agent_id].append(
+                                [t, t, t, len(self.queue), self.num_system]
+                            )
 
                     if self._arrivals[0]._time < self._departures[0]._time:
                         self._time = self._arrivals[0]._time
@@ -164,9 +182,13 @@ def next_event(self):
 
                     if self.collect_data:
                         if arrival.agent_id not in self.data:
-                            self.data[arrival.agent_id] = [[arrival._time, 0, 0, len(self.queue), self.num_system]]
+                            self.data[arrival.agent_id] = [
+                                [arrival._time, 0, 0, len(self.queue), self.num_system]
+                            ]
                         else:
-                            self.data[arrival.agent_id].append([arrival._time, 0, 0, len(self.queue), self.num_system])
+                            self.data[arrival.agent_id].append(
+                                [arrival._time, 0, 0, len(self.queue), self.num_system]
+                            )
 
                     if self._arrivals[0]._time < self._departures[0]._time:
                         self._time = self._arrivals[0]._time
@@ -179,25 +201,25 @@ def _current_color(self, which=0):
         if which == 1:
             nSy = self.num_servers
             cap = self.max_servers
-            div = 5. if cap <= 1 else (3. * cap)
+            div = 5.0 if cap <= 1 else (3.0 * cap)
             tmp = 0.9 - min(nSy / div, 0.9)
 
-            color = [i * tmp / 0.9 for i in self.colors['edge_loop_color']]
+            color = [i * tmp / 0.9 for i in self.colors["edge_loop_color"]]
             color[3] = 0.0
 
         elif which == 2:
-            color = self.colors['vertex_pen_color']
+            color = self.colors["vertex_pen_color"]
         else:
             nSy = self.num_servers
             cap = self.max_servers
-            div = 5. if cap <= 1 else (3. * cap)
+            div = 5.0 if cap <= 1 else (3.0 * cap)
             tmp = 0.9 - min(nSy / div, 0.9)
 
             if self.edge[0] == self.edge[1]:
-                color = [i * tmp / 0.9 for i in self.colors['vertex_fill_color']]
+                color = [i * tmp / 0.9 for i in self.colors["vertex_fill_color"]]
                 color[3] = 1.0
             else:
-                color = [i * tmp / 0.9 for i in self.colors['edge_color']]
+                color = [i * tmp / 0.9 for i in self.colors["edge_color"]]
                 color[3] = 0.5
 
         return color
@@ -228,7 +250,9 @@ class InfoAgent(Agent):
         The number of edges in the network.
     **kwargs :
         Any arguments to pass to :class:`.Agent`.
+
     """
+
     def __init__(self, agent_id=(0, 0), net_size=1, **kwargs):
         super(InfoAgent, self).__init__(agent_id, **kwargs)
 
@@ -236,9 +260,11 @@ def __init__(self, agent_id=(0, 0), net_size=1, **kwargs):
         self.net_data = np.ones((net_size, 3)) * -1
 
     def __repr__(self):
-        return "InfoAgent; agent_id:{0}. Time: {1}".format(self.agent_id, round(self._time, 3))
+        return f"InfoAgent; agent_id:{self.agent_id}. Time: {round(self._time, 3)}"
 
-    def add_loss(self, qedge, *args, **kwargs):  # qedge[2] is the edge_index of the queue
+    def add_loss(
+        self, qedge, *args, **kwargs
+    ):  # qedge[2] is the edge_index of the queue
         self.stats[qedge[2], 2] += 1
 
     def get_beliefs(self):
@@ -247,20 +273,29 @@ def get_beliefs(self):
     def queue_action(self, queue, *args, **kwargs):
         if isinstance(queue, InfoQueue):
             # update information
-            a = logical_or(self.net_data[:, 0] < queue.net_data[:, 0], self.net_data[:, 0] == -1)
+            a = logical_or(
+                self.net_data[:, 0] < queue.net_data[:, 0], self.net_data[:, 0] == -1
+            )
             self.net_data[a, :] = queue.net_data[a, :]
 
             # stamp this information
-            n = queue.edge[2]    # This is the edge_index of the queue
+            n = queue.edge[2]  # This is the edge_index of the queue
             if self.agent_id in queue.data:
-                tmp = queue.data[self.agent_id][-1][1] - queue.data[self.agent_id][-1][0]
+                tmp = (
+                    queue.data[self.agent_id][-1][1] - queue.data[self.agent_id][-1][0]
+                )
                 self.stats[n, 0] = self.stats[n, 0] + tmp
                 self.stats[n, 1] += 1 if tmp > 0 else 0
-            self.net_data[n, :] = queue._current_t, queue.num_servers, queue.num_system / queue.num_servers
+            self.net_data[n, :] = (
+                queue._current_t,
+                queue.num_servers,
+                queue.num_system / queue.num_servers,
+            )
 
 
 class InfoQueue(LossQueue):
-    """A queue that stores information about the network.
+    """
+    A queue that stores information about the network.
 
     This queue gets information about the state of the network (number
     of :class:`Agent's<.Agent>` at other queues) from arriving
@@ -283,17 +318,29 @@ class InfoQueue(LossQueue):
         serviced.
     **kwargs :
         Extra parameters to pass to :class:`.LossQueue`.
+
     """
-    def __init__(self, net_size=1, AgentFactory=InfoAgent, qbuffer=np.infty, **kwargs):
-        super(InfoQueue, self).__init__(AgentFactory=AgentFactory, qbuffer=qbuffer, **kwargs)
+
+    def __init__(self, net_size=1, AgentFactory=InfoAgent, qbuffer=np.inf, **kwargs):
+        super(InfoQueue, self).__init__(
+            AgentFactory=AgentFactory, qbuffer=qbuffer, **kwargs
+        )
 
         self.networking(net_size)
 
     def __repr__(self):
-        my_str = ("InfoQueue:{0}. Servers: {1}, queued: {2}, "
-                  "arrivals: {3}, departures: {4}, next time: {5}")
-        arg = (self.edge[2], self.num_servers, len(self.queue),
-               self.num_arrivals, self.num_departures, round(self._time, 3))
+        my_str = (
+            "InfoQueue:{0}. Servers: {1}, queued: {2}, "
+            "arrivals: {3}, departures: {4}, next time: {5}"
+        )
+        arg = (
+            self.edge[2],
+            self.num_servers,
+            len(self.queue),
+            self.num_arrivals,
+            self.num_departures,
+            round(self._time, 3),
+        )
         return my_str.format(*arg)
 
     def networking(self, network_size):
@@ -308,23 +355,24 @@ def _add_arrival(self, agent=None):
         if agent is not None:
             self._num_total += 1
             heappush(self._arrivals, agent)
-        else:
-            if self._current_t >= self._next_ct:
-                self._next_ct = self.arrival_f(self._current_t)
+        elif self._current_t >= self._next_ct:
+            self._next_ct = self.arrival_f(self._current_t)
 
-                if self._next_ct >= self.deactive_t:
-                    self.active = False
-                    return
+            if self._next_ct >= self.deactive_t:
+                self.active = False
+                return
 
-                self._num_total += 1
-                new_agent = self.AgentFactory((self.edge[2], self._oArrivals), len(self.net_data))
-                new_agent._time = self._next_ct
-                heappush(self._arrivals, new_agent)
+            self._num_total += 1
+            new_agent = self.AgentFactory(
+                (self.edge[2], self._oArrivals), len(self.net_data)
+            )
+            new_agent._time = self._next_ct
+            heappush(self._arrivals, new_agent)
 
-                self._oArrivals += 1
+            self._oArrivals += 1
 
-                if self._oArrivals >= self.active_cap:
-                    self._active = False
+            if self._oArrivals >= self.active_cap:
+                self._active = False
 
         if self._arrivals[0]._time < self._departures[0]._time:
             self._time = self._arrivals[0]._time
diff --git a/queueing_tool/queues/queue_servers.py b/queueing_tool/queues/queue_servers.py
index 6a196cc..f57c9bf 100644
--- a/queueing_tool/queues/queue_servers.py
+++ b/queueing_tool/queues/queue_servers.py
@@ -1,17 +1,17 @@
 import collections
 import copy
 import numbers
-from heapq import heappush, heappop
+from heapq import heappop, heappush
 
-from numpy.random import uniform, exponential
-from numpy import infty
 import numpy as np
+from numpy.random import exponential, uniform
 
 from queueing_tool.queues.agents import Agent, InftyAgent
 
 
 def poisson_random_measure(t, rate, rate_max):
-    """A function that returns the arrival time of the next arrival for
+    """
+    A function that returns the arrival time of the next arrival for
     a Poisson random measure.
 
     Parameters
@@ -37,13 +37,13 @@ def poisson_random_measure(t, rate, rate_max):
 
     .. math::
 
-       \int_{t}^{t+s} dx \, r(x)
+       \\int_{t}^{t+s} dx \\, r(x)
 
     where :math:`r(t)` is the supplied ``rate`` function. This function
     can only simulate processes that have bounded intensity functions.
     See chapter 6 of [3]_ for more on the mathematics behind Poisson
     random measures; the book's publisher, Springer, has that chapter
-    available online for free at (`pdf`_\).
+    available online for free at (`pdf`_\\).
 
     A Poisson random measure is sometimes called a non-homogeneous
     Poisson process. A Poisson process is a special type of Poisson
@@ -54,7 +54,7 @@ def poisson_random_measure(t, rate, rate_max):
     Examples
     --------
     Suppose you wanted to model the arrival process as a Poisson
-    random measure with rate function :math:`r(t) = 2 + \sin( 2\pi t)`.
+    random measure with rate function :math:`r(t) = 2 + \\sin( 2\\pi t)`.
     Then you could do so as follows:
 
     >>> import queueing_tool as qt
@@ -70,6 +70,7 @@ def poisson_random_measure(t, rate, rate_max):
     .. [3] Cinlar, Erhan. *Probability and stochastics*. Graduate Texts in\
            Mathematics. Vol. 261. Springer, New York, 2011.\
            :doi:`10.1007/978-0-387-87859-1`
+
     """
     scale = 1.0 / rate_max
     t = t + exponential(scale)
@@ -78,8 +79,9 @@ def poisson_random_measure(t, rate, rate_max):
     return t
 
 
-class QueueServer(object):
-    """The base queue-server class.
+class QueueServer:
+    """
+    The base queue-server class.
 
     Built to work with the :class:`.QueueNetwork` class, but can stand
     alone as a multi-server queue. It supports a capped pool of
@@ -198,7 +200,7 @@ class QueueServer(object):
     The following code constructs an :math:`\\text{M}_t/\\text{GI}/5`
     :class:`.QueueServer` with mean utilization rate
     :math:`\\rho = 0.8`. The arrivals are modeled as a Poisson random
-    measure with rate function :math:`r(t) = 2 + 16 \sin^2(\pi t / 8)`
+    measure with rate function :math:`r(t) = 2 + 16 \\sin^2(\\pi t / 8)`
     and a service distribution that is gamma with shape and scale
     parameters 4 and 0.1 respectively. To create such a queue run:
 
@@ -223,8 +225,8 @@ class QueueServer(object):
     Notes
     -----
     This is a generic multi-server queue implimentation (see [4]_).
-    In `Kendall's notation`_\, this is a
-    :math:`\\text{GI}_t/\\text{GI}_t/c/\infty/N/\\text{FIFO}` queue
+    In `Kendall's notation`_\\, this is a
+    :math:`\\text{GI}_t/\\text{GI}_t/c/\\infty/N/\\text{FIFO}` queue
     class, where :math:`c` is set by ``num_servers`` and :math:`N` is set
     by ``active_cap``. See chapter 1 of [3]_ (pdfs from `the author`_
     and `the publisher`_) for a good introduction to the theory behind
@@ -266,40 +268,52 @@ class QueueServer(object):
     .. _the publisher: http://assets.cambridge.org/97811070/27503/excerpt/\
                         9781107027503_excerpt.pdf
     .. _9781107027503: http://www.cambridge.org/us/9781107027503
+
     """
 
     _default_colors = {
-        'edge_loop_color': [0, 0, 0, 0],
-        'edge_color': [0.9, 0.9, 0.9, 0.5],
-        'vertex_fill_color': [1.0, 1.0, 1.0, 1.0],
-        'vertex_color': [0.0, 0.5, 1.0, 1.0]
+        "edge_loop_color": [0, 0, 0, 0],
+        "edge_color": [0.9, 0.9, 0.9, 0.5],
+        "vertex_fill_color": [1.0, 1.0, 1.0, 1.0],
+        "vertex_color": [0.0, 0.5, 1.0, 1.0],
     }
 
-    def __init__(self, num_servers=1, arrival_f=None,
-                 service_f=None, edge=(0, 0, 0, 1),
-                 AgentFactory=Agent, collect_data=False, active_cap=infty,
-                 deactive_t=infty, colors=None, seed=None,
-                 coloring_sensitivity=2, **kwargs):
-
-        if not isinstance(num_servers, numbers.Integral) and num_servers is not infty:
+    def __init__(
+        self,
+        num_servers=1,
+        arrival_f=None,
+        service_f=None,
+        edge=(0, 0, 0, 1),
+        AgentFactory=Agent,
+        collect_data=False,
+        active_cap=np.inf,
+        deactive_t=np.inf,
+        colors=None,
+        seed=None,
+        coloring_sensitivity=2,
+        **kwargs,
+    ):
+        if not isinstance(num_servers, numbers.Integral) and num_servers is not np.inf:
             msg = "num_servers must be an integer or infinity."
             raise TypeError(msg)
-        elif num_servers <= 0:
+        if num_servers <= 0:
             msg = "num_servers must be a positive integer or infinity."
             raise ValueError(msg)
 
         self.edge = edge
-        self.num_servers = kwargs.get('nServers', num_servers)
+        self.num_servers = kwargs.get("nServers", num_servers)
         self.num_departures = 0
         self.num_system = 0
-        self.data = {}   # times; agent_id : [arrival, service start, departure]
+        self.data = {}  # times; agent_id : [arrival, service start, departure]
         self.queue = collections.deque()
 
         if arrival_f is None:
+
             def arrival_f(t):
                 return t + exponential(1.0)
 
         if service_f is None:
+
             def service_f(t):
                 return t + exponential(0.9)
 
@@ -311,15 +325,17 @@ def service_f(t):
         self.deactive_t = deactive_t
 
         inftyAgent = InftyAgent()
-        self._arrivals = [inftyAgent]    # A list of arriving agents.
-        self._departures = [inftyAgent]    # A list of departing agents.
+        self._arrivals = [inftyAgent]  # A list of arriving agents.
+        self._departures = [inftyAgent]  # A list of departing agents.
         self._num_arrivals = 0
         self._oArrivals = 0
-        self._num_total = 0       # The number of agents scheduled to arrive + num_system
+        self._num_total = 0  # The number of agents scheduled to arrive + num_system
         self._active = False
-        self._current_t = 0       # The time of the last event.
-        self._time = infty   # The time of the next event.
-        self._next_ct = 0       # The next time an arrival from outside the network can arrive.
+        self._current_t = 0  # The time of the last event.
+        self._time = np.inf  # The time of the next event.
+        self._next_ct = (
+            0  # The next time an arrival from outside the network can arrive.
+        )
         self.coloring_sensitivity = coloring_sensitivity
 
         if isinstance(seed, numbers.Integral):
@@ -349,50 +365,60 @@ def num_arrivals(self):
         return [self._num_arrivals, self._oArrivals]
 
     def __repr__(self):
-        my_str = ("QueueServer:{0}. Servers: {1}, queued: {2}, arrivals: {3}, "
-                  "departures: {4}, next time: {5}")
-        arg = (self.edge[2], self.num_servers, len(self.queue), self.num_arrivals,
-               self.num_departures, round(self._time, 3))
+        my_str = (
+            "QueueServer:{0}. Servers: {1}, queued: {2}, arrivals: {3}, "
+            "departures: {4}, next time: {5}"
+        )
+        arg = (
+            self.edge[2],
+            self.num_servers,
+            len(self.queue),
+            self.num_arrivals,
+            self.num_departures,
+            round(self._time, 3),
+        )
         return my_str.format(*arg)
 
     def _add_arrival(self, agent=None):
         if agent is not None:
             self._num_total += 1
             heappush(self._arrivals, agent)
-        else:
-            if self._current_t >= self._next_ct:
-                self._next_ct = self.arrival_f(self._current_t)
+        elif self._current_t >= self._next_ct:
+            self._next_ct = self.arrival_f(self._current_t)
 
-                if self._next_ct >= self.deactive_t:
-                    self._active = False
-                    return
+            if self._next_ct >= self.deactive_t:
+                self._active = False
+                return
 
-                self._num_total += 1
-                new_agent = self.AgentFactory((self.edge[2], self._oArrivals))
-                new_agent._time = self._next_ct
-                heappush(self._arrivals, new_agent)
+            self._num_total += 1
+            new_agent = self.AgentFactory((self.edge[2], self._oArrivals))
+            new_agent._time = self._next_ct
+            heappush(self._arrivals, new_agent)
 
-                self._oArrivals += 1
+            self._oArrivals += 1
 
-                if self._oArrivals >= self.active_cap:
-                    self._active = False
+            if self._oArrivals >= self.active_cap:
+                self._active = False
 
         if self._arrivals[0]._time < self._departures[0]._time:
             self._time = self._arrivals[0]._time
 
     def at_capacity(self):
-        """Returns whether the queue is at capacity or not.
+        """
+        Returns whether the queue is at capacity or not.
 
         Returns
         -------
         bool
             Always returns ``False``, since the ``QueueServer`` class
             has infinite capacity.
+
         """
         return False
 
     def clear(self):
-        """Clears out the queue. Removes all arrivals, departures, and
+        """
+        Clears out the queue. Removes all arrivals, departures, and
         queued agents from the :class:`.QueueServer`, resets
         :attr:`.num_arrivals`, :attr:`.num_departures`,
         :attr:`.num_system`, and the clock to zero. It also clears any
@@ -405,7 +431,7 @@ def clear(self):
         self.num_system = 0
         self._num_total = 0
         self._current_t = 0
-        self._time = infty
+        self._time = np.inf
         self._next_ct = 0
         self._active = False
         self.queue.clear()
@@ -418,7 +444,8 @@ def copy(self):
         return copy.deepcopy(self)
 
     def _current_color(self, which=0):
-        """Returns a color for the queue.
+        """
+        Returns a color for the queue.
 
         Parameters
         ----------
@@ -446,28 +473,30 @@ def _current_color(self, which=0):
               correspond to darker edge colors. Uses
               ``colors['vertex_fill_color']`` if the queue sits on a
               loop, and ``colors['edge_color']`` otherwise.
+
         """
         if which == 1:
-            color = self.colors['edge_loop_color']
+            color = self.colors["edge_loop_color"]
 
         elif which == 2:
-            color = self.colors['vertex_color']
+            color = self.colors["vertex_color"]
 
         else:
-            div = self.coloring_sensitivity * self.num_servers + 1.
-            tmp = 1. - min(self.num_system / div, 1)
+            div = self.coloring_sensitivity * self.num_servers + 1.0
+            tmp = 1.0 - min(self.num_system / div, 1)
 
             if self.edge[0] == self.edge[1]:
-                color = [i * tmp for i in self.colors['vertex_fill_color']]
+                color = [i * tmp for i in self.colors["vertex_fill_color"]]
                 color[3] = 1.0
             else:
-                color = [i * tmp for i in self.colors['edge_color']]
-                color[3] = 1 / 2.
+                color = [i * tmp for i in self.colors["edge_color"]]
+                color[3] = 1 / 2.0
 
         return color
 
     def delay_service(self, t=None):
-        """Adds an extra service time to the next departing
+        """
+        Adds an extra service time to the next departing
         :class:`Agent's<.Agent>` service time.
 
         Parameters
@@ -476,6 +505,7 @@ def delay_service(self, t=None):
             Specifies the departing time for the agent scheduled
             to depart next. If ``t`` is not given, then an additional
             service time is added to the next departing agent.
+
         """
         if len(self._departures) > 1:
             agent = heappop(self._departures)
@@ -498,7 +528,8 @@ def _service_from_queue(self):
         heappush(self._departures, agent)
 
     def fetch_data(self, return_header=False):
-        """Fetches data from the queue.
+        """
+        Fetches data from the queue.
 
         Parameters
         ----------
@@ -522,8 +553,8 @@ def fetch_data(self, return_header=False):
         headers : str (optional)
             A comma seperated string of the column headers. Returns
             ``'arrival,service,departure,num_queued,num_total,q_id'``
-        """
 
+        """
         qdata = []
         for d in self.data.values():
             qdata.extend(d)
@@ -534,19 +565,19 @@ def fetch_data(self, return_header=False):
             dat[:, 5] = self.edge[2]
 
             dType = [
-                ('a', float),
-                ('s', float),
-                ('d', float),
-                ('q', float),
-                ('n', float),
-                ('id', float)
+                ("a", float),
+                ("s", float),
+                ("d", float),
+                ("q", float),
+                ("n", float),
+                ("id", float),
             ]
             dat = np.array([tuple(d) for d in dat], dtype=dType)
-            dat = np.sort(dat, order='a')
+            dat = np.sort(dat, order="a")
             dat = np.array([tuple(d) for d in dat])
 
         if return_header:
-            return dat, 'arrival,service,departure,num_queued,num_total,q_id'
+            return dat, "arrival,service,departure,num_queued,num_total,q_id"
 
         return dat
 
@@ -554,17 +585,20 @@ def _key(self):
         return self._time, self.edge[2]
 
     def number_queued(self):
-        """Returns the number of agents waiting in line to be served.
+        """
+        Returns the number of agents waiting in line to be served.
 
         Returns
         -------
         out : int
             The number of agents waiting in line to be served.
+
         """
         return len(self.queue)
 
     def next_event(self):
-        """Simulates the queue forward one event.
+        """
+        Simulates the queue forward one event.
 
         Use :meth:`.simulate` instead.
 
@@ -577,6 +611,7 @@ def next_event(self):
         See Also
         --------
         :meth:`.simulate` : Simulates the queue forward.
+
         """
         if self._departures[0]._time < self._arrivals[0]._time:
             new_depart = heappop(self._departures)
@@ -590,7 +625,7 @@ def next_event(self):
 
             num_queued = len(self.queue)
             # This is the number of agents currently being serviced by
-            # the QueueServer. This number need not be equal to 
+            # the QueueServer. This number need not be equal to
             # num_servers - 1 (although it typically is), since there
             # can be a change to the number of servers after the queue was
             # created.
@@ -602,7 +637,7 @@ def next_event(self):
             self._update_time()
             return new_depart
 
-        elif self._arrivals[0]._time < infty:
+        if self._arrivals[0]._time < np.inf:
             arrival = heappop(self._arrivals)
             self._current_t = arrival._time
 
@@ -615,11 +650,13 @@ def next_event(self):
             if self.collect_data:
                 b = 0 if self.num_system <= self.num_servers else 1
                 if arrival.agent_id not in self.data:
-                    self.data[arrival.agent_id] = \
-                        [[arrival._time, 0, 0, len(self.queue) + b, self.num_system]]
+                    self.data[arrival.agent_id] = [
+                        [arrival._time, 0, 0, len(self.queue) + b, self.num_system]
+                    ]
                 else:
-                    self.data[arrival.agent_id]\
-                        .append([arrival._time, 0, 0, len(self.queue) + b, self.num_system])
+                    self.data[arrival.agent_id].append(
+                        [arrival._time, 0, 0, len(self.queue) + b, self.num_system]
+                    )
 
             arrival.queue_action(self, 0)
 
@@ -636,7 +673,8 @@ def next_event(self):
             self._update_time()
 
     def next_event_description(self):
-        """Returns an integer representing whether the next event is
+        """
+        Returns an integer representing whether the next event is
         an arrival, a departure, or nothing.
 
         Returns
@@ -646,16 +684,17 @@ def next_event_description(self):
             arrival or a departure: ``1`` corresponds to an arrival,
             ``2`` corresponds to a departure, and ``0`` corresponds to
             nothing scheduled to occur.
+
         """
         if self._departures[0]._time < self._arrivals[0]._time:
             return 2
-        elif self._arrivals[0]._time < infty:
+        if self._arrivals[0]._time < np.inf:
             return 1
-        else:
-            return 0
+        return 0
 
     def set_active(self):
-        """Changes the ``active`` attribute to True. Agents may now
+        """
+        Changes the ``active`` attribute to True. Agents may now
         arrive from outside the network.
         """
         if not self._active:
@@ -667,7 +706,8 @@ def set_inactive(self):
         self._active = False
 
     def set_num_servers(self, n):
-        """Change the number of servers in the queue to ``n``.
+        """
+        Change the number of servers in the queue to ``n``.
 
         Parameters
         ----------
@@ -682,26 +722,27 @@ def set_num_servers(self, n):
             error is raised.
         ValueError
             If ``n`` is not positive.
+
         """
-        if not isinstance(n, numbers.Integral) and n is not infty:
+        if not isinstance(n, numbers.Integral) and n is not np.inf:
             the_str = "n must be an integer or infinity.\n{0}"
             raise TypeError(the_str.format(str(self)))
-        elif n <= 0:
+        if n <= 0:
             the_str = "n must be a positive integer or infinity.\n{0}"
             raise ValueError(the_str.format(str(self)))
-        else:
-            agents_to_queue = max(min(n - self.num_servers, len(self.queue)), 0)
+        agents_to_queue = max(min(n - self.num_servers, len(self.queue)), 0)
 
-            for _ in range(agents_to_queue):
-                self._service_from_queue()
+        for _ in range(agents_to_queue):
+            self._service_from_queue()
 
-            if agents_to_queue > 0:
-                self._update_time()
+        if agents_to_queue > 0:
+            self._update_time()
 
-            self.num_servers = n
+        self.num_servers = n
 
     def simulate(self, n=1, t=None, nA=None, nD=None):
-        """This method simulates the queue forward for a specified
+        """
+        This method simulates the queue forward for a specified
         amount of simulation time, or for a specific number of
         events.
 
@@ -759,21 +800,22 @@ def simulate(self, n=1, t=None, nA=None, nD=None):
         >>> q.simulate(nA=1000)
         >>> q.num_departures - nD0, q.num_arrivals[1] - nA0,
         (987, 1000)
+
         """
         if t is None and nD is None and nA is None:
             for dummy in range(n):
                 self.next_event()
         elif t is not None:
             then = self._current_t + t
-            while self._current_t < then and self._time < infty:
+            while self._current_t < then and self._time < np.inf:
                 self.next_event()
         elif nD is not None:
             num_departures = self.num_departures + nD
-            while self.num_departures < num_departures and self._time < infty:
+            while self.num_departures < num_departures and self._time < np.inf:
                 self.next_event()
         elif nA is not None:
             num_arrivals = self._oArrivals + nA
-            while self._oArrivals < num_arrivals and self._time < infty:
+            while self._oArrivals < num_arrivals and self._time < np.inf:
                 self.next_event()
 
     def _update_time(self):
@@ -784,7 +826,8 @@ def _update_time(self):
 
 
 class LossQueue(QueueServer):
-    """A finite capacity queue.
+    """
+    A finite capacity queue.
 
     If an agent arrives to a queue that is at capacity, then the agent
     gets blocked. If this agent is arriving from inside the network
@@ -813,13 +856,14 @@ class LossQueue(QueueServer):
     :math:`\\text{GI}_t/\\text{GI}_t/c/c+b/N/\\text{FIFO}` queue, where
     :math:`b` is the ``qbuffer``. If the default parameters are used
     then the instance is an :math:`\\text{M}/\\text{M}/1/1` queue.
+
     """
 
     _default_colors = {
-        'edge_loop_color': [0, 0, 0, 0],
-        'edge_color': [0.7, 0.7, 0.7, 0.5],
-        'vertex_fill_color': [1.0, 1.0, 1.0, 1.0],
-        'vertex_color': [0.133, 0.545, 0.133, 1.0]
+        "edge_loop_color": [0, 0, 0, 0],
+        "edge_color": [0.7, 0.7, 0.7, 0.5],
+        "vertex_fill_color": [1.0, 1.0, 1.0, 1.0],
+        "vertex_color": [0.133, 0.545, 0.133, 1.0],
     }
 
     def __init__(self, qbuffer=0, **kwargs):
@@ -829,14 +873,23 @@ def __init__(self, qbuffer=0, **kwargs):
         self.buffer = qbuffer
 
     def __repr__(self):
-        tmp = ("LossQueue:{0}. Servers: {1}, queued: {2}, arrivals: {3}, "
-               "departures: {4}, next time: {5}")
-        arg = (self.edge[2], self.num_servers, len(self.queue), self.num_arrivals,
-               self.num_departures, round(self._time, 3))
+        tmp = (
+            "LossQueue:{0}. Servers: {1}, queued: {2}, arrivals: {3}, "
+            "departures: {4}, next time: {5}"
+        )
+        arg = (
+            self.edge[2],
+            self.num_servers,
+            len(self.queue),
+            self.num_arrivals,
+            self.num_departures,
+            round(self._time, 3),
+        )
         return tmp.format(*arg)
 
     def at_capacity(self):
-        """Returns whether the queue is at capacity or not.
+        """
+        Returns whether the queue is at capacity or not.
 
         Returns
         -------
@@ -844,6 +897,7 @@ def at_capacity(self):
             Returns whether the number of agents in the system -- the
             number of agents being serviced plus those waiting to be
             serviced -- is equal to ``num_servers + buffer``.
+
         """
         return self.num_system >= self.num_servers + self.buffer
 
@@ -852,10 +906,9 @@ def clear(self):
         self.num_blocked = 0
 
     def next_event(self):
-
         if self._departures[0]._time < self._arrivals[0]._time:
             return super(LossQueue, self).next_event()
-        elif self._arrivals[0]._time < infty:
+        if self._arrivals[0]._time < np.inf:
             if self.num_system < self.num_servers + self.buffer:
                 super(LossQueue, self).next_event()
             else:
@@ -872,9 +925,13 @@ def next_event(self):
 
                 if self.collect_data:
                     if arrival.agent_id in self.data:
-                        self.data[arrival.agent_id].append([arrival._time, 0, 0, len(self.queue), self.num_system])
+                        self.data[arrival.agent_id].append(
+                            [arrival._time, 0, 0, len(self.queue), self.num_system]
+                        )
                     else:
-                        self.data[arrival.agent_id] = [[arrival._time, 0, 0, len(self.queue), self.num_system]]
+                        self.data[arrival.agent_id] = [
+                            [arrival._time, 0, 0, len(self.queue), self.num_system]
+                        ]
 
                 if self._arrivals[0]._time < self._departures[0]._time:
                     self._time = self._arrivals[0]._time
@@ -883,7 +940,8 @@ def next_event(self):
 
 
 class NullQueue(QueueServer):
-    """A terminal queue.
+    """
+    A terminal queue.
 
     A queue that is used by the :class:`.QueueNetwork` class to
     represent agents leaving the network. Since the ``NullQueue``
@@ -899,21 +957,21 @@ class NullQueue(QueueServer):
     """
 
     _default_colors = {
-        'edge_loop_color': [0, 0, 0, 0],
-        'edge_color': [0.7, 0.7, 0.7, 0.5],
-        'vertex_fill_color': [1.0, 1.0, 1.0, 1.0],
-        'vertex_color': [0.5, 0.5, 0.5, 0.5]
+        "edge_loop_color": [0, 0, 0, 0],
+        "edge_color": [0.7, 0.7, 0.7, 0.5],
+        "vertex_fill_color": [1.0, 1.0, 1.0, 1.0],
+        "vertex_color": [0.5, 0.5, 0.5, 0.5],
     }
 
     def __init__(self, *args, **kwargs):
-        if 'edge' not in kwargs:
-            kwargs['edge'] = (0, 0, 0, 0)
+        if "edge" not in kwargs:
+            kwargs["edge"] = (0, 0, 0, 0)
 
         super(NullQueue, self).__init__(**kwargs)
         self.num_servers = 0
 
     def __repr__(self):
-        return "NullQueue:{0}.".format(self.edge[2])
+        return f"NullQueue:{self.edge[2]}."
 
     def initialize(self, *args, **kwargs):
         pass
@@ -942,14 +1000,13 @@ def next_event(self):
 
     def _current_color(self, which=0):
         if which == 1:
-            color = self.colors['edge_loop_color']
+            color = self.colors["edge_loop_color"]
         elif which == 2:
-            color = self.colors['vertex_color']
+            color = self.colors["vertex_color"]
+        elif self.edge[0] == self.edge[1]:
+            color = self.colors["vertex_fill_color"]
         else:
-            if self.edge[0] == self.edge[1]:
-                color = self.colors['vertex_fill_color']
-            else:
-                color = self.colors['edge_color']
+            color = self.colors["edge_color"]
         return color
 
     def clear(self):
diff --git a/queueing_tool/union_find.py b/queueing_tool/union_find.py
index 15a3a83..0787269 100644
--- a/queueing_tool/union_find.py
+++ b/queueing_tool/union_find.py
@@ -1,5 +1,6 @@
-class UnionFind(object):
-    """The union-find data structure with union by rank and path compression.
+class UnionFind:
+    """
+    The union-find data structure with union by rank and path compression.
 
     The UnionFind data structure is a collection of objects that supports
     the union and find operations (described below). Each object in the
@@ -17,20 +18,21 @@ class UnionFind(object):
     ----------
     nClusters : int
         The number of clusters contained in the data-structure.
+
     """
+
     def __init__(self, S):
-        self._leader   = dict((s, s) for s in S)
-        self._size     = dict((s, 1) for s in S)
-        self._rank     = dict((s, 0) for s in S)
+        self._leader = dict((s, s) for s in S)
+        self._size = dict((s, 1) for s in S)
+        self._rank = dict((s, 0) for s in S)
         self.nClusters = len(self._leader)
 
-
     def __repr__(self):
-        return "UnionFind: contains {0} clusters.".format(self.nClusters)
-
+        return f"UnionFind: contains {self.nClusters} clusters."
 
     def size(self, s):
-        """Returns the number of elements in the set that ``s`` belongs to.
+        """
+        Returns the number of elements in the set that ``s`` belongs to.
 
         Parameters
         ----------
@@ -41,13 +43,14 @@ def size(self, s):
         -------
         out : int
             The number of elements in the set that ``s`` belongs to.
+
         """
         leader = self.find(s)
         return self._size[leader]
 
-
     def find(self, s):
-        """Locates the leader of the set to which the element ``s`` belongs.
+        """
+        Locates the leader of the set to which the element ``s`` belongs.
 
         Parameters
         ----------
@@ -58,8 +61,9 @@ def find(self, s):
         -------
         object
             The leader of the set that contains ``s``.
+
         """
-        pSet   = [s]
+        pSet = [s]
         parent = self._leader[s]
 
         while parent != self._leader[parent]:
@@ -72,18 +76,19 @@ def find(self, s):
 
         return parent
 
-
     def union(self, a, b):
-        """Merges the set that contains ``a`` with the set that contains ``b``.
+        """
+        Merges the set that contains ``a`` with the set that contains ``b``.
 
         Parameters
         ----------
         a, b : objects
             Two objects whose sets are to be merged.
+
         """
         s1, s2 = self.find(a), self.find(b)
         if s1 != s2:
-            r1, r2  = self._rank[s1], self._rank[s2]
+            r1, r2 = self._rank[s1], self._rank[s2]
             if r2 > r1:
                 r1, r2 = r2, r1
                 s1, s2 = s2, s1
@@ -91,5 +96,5 @@ def union(self, a, b):
                 self._rank[s1] += 1
 
             self._leader[s2] = s1
-            self._size[s1]  += self._size[s2]
-            self.nClusters  -= 1
+            self._size[s1] += self._size[s2]
+            self.nClusters -= 1
diff --git a/ruff.toml b/ruff.toml
new file mode 100644
index 0000000..d80b034
--- /dev/null
+++ b/ruff.toml
@@ -0,0 +1,78 @@
+# Exclude a variety of commonly ignored directories.
+exclude = [
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".git-rewrite",
+    ".hg",
+    ".ipynb_checkpoints",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pyenv",
+    ".pytest_cache",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    ".vscode",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "site-packages",
+    "venv",
+]
+
+# Same as Black.
+line-length = 88
+indent-width = 4
+
+# Assume Python 3.9
+target-version = "py39"
+
+[lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`)  codes by default.
+# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
+# McCabe complexity (`C901`) by default.
+select = ["ALL"]
+ignore = ["D211", "D212", "COM812", "D203"]
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+# Allow unused variables when underscore-prefixed.
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+[format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Enable auto-formatting of code examples in docstrings. Markdown,
+# reStructuredText code/literal blocks and doctests are all supported.
+#
+# This is currently disabled by default, but it is planned for this
+# to be opt-out in the future.
+docstring-code-format = false
+
+# Set the line length limit used when formatting code snippets in
+# docstrings.
+#
+# This only has an effect when the `docstring-code-format` setting is
+# enabled.
+docstring-code-line-length = "dynamic"
+
diff --git a/setup.py b/setup.py
index 61f0d7f..7403ccd 100644
--- a/setup.py
+++ b/setup.py
@@ -1,15 +1,13 @@
-
-from setuptools import setup, Extension
-
+from setuptools import Extension, setup
 
 ext_modules = [
     Extension(
-        name='queueing_tool.network.priority_queue',
-        sources=['queueing_tool/network/priority_queue.c'],
+        name="queueing_tool.network.priority_queue",
+        sources=["queueing_tool/network/priority_queue.c"],
     ),
     Extension(
-        name='queueing_tool.queues.choice',
-        sources=['queueing_tool/queues/choice.c'],
+        name="queueing_tool.queues.choice",
+        sources=["queueing_tool/queues/choice.c"],
     ),
 ]
 
diff --git a/tests/test_graph_generation.py b/tests/test_graph_generation.py
index 6c91c75..3c0e4a1 100644
--- a/tests/test_graph_generation.py
+++ b/tests/test_graph_generation.py
@@ -1,4 +1,3 @@
-import unittest
 
 import networkx as nx
 import numpy as np
@@ -18,20 +17,20 @@ def generate_adjacency(a=3, b=25, c=6, n=12):
 @pytest.fixture(name="expected_response0")
 def fixture_expected_response0():
     return {
-        0: {1: {'edge_type': 5}},
-        1: {2: {'edge_type': 9}, 3: {'edge_type': 14}},
-        2: {0: {'edge_type': 1}},
-        3: {3: {'edge_type': 0}}
+        0: {1: {"edge_type": 5}},
+        1: {2: {"edge_type": 9}, 3: {"edge_type": 14}},
+        2: {0: {"edge_type": 1}},
+        3: {3: {"edge_type": 0}},
     }
 
 
 @pytest.fixture(name="expected_response1")
 def fixture_expected_response1():
     return {
-        0: {1: {'edge_type': 5}},
-        1: {2: {'edge_type': 9}, 3: {'edge_type': 0}},
-        2: {0: {'edge_type': 1}},
-        3: {}
+        0: {1: {"edge_type": 5}},
+        1: {2: {"edge_type": 9}, 3: {"edge_type": 0}},
+        2: {0: {"edge_type": 1}},
+        3: {},
     }
 
 
@@ -53,7 +52,7 @@ def test_add_edge_lengths():
         for key in g2.edge_properties():
             edge_props.add(key)
 
-        assert 'edge_length' in edge_props
+        assert "edge_length" in edge_props
 
     @staticmethod
     def test_generate_transition():
@@ -69,12 +68,8 @@ def test_generate_transition():
 
     @staticmethod
     def test_adjacency2graph_matrix_adjacency(expected_response1):
-
         # Test adjacency argument using ndarray work
-        adj = np.array([[0, 1, 0, 0],
-                        [0, 0, 1, 1],
-                        [1, 0, 0, 0],
-                        [0, 0, 0, 0]])
+        adj = np.array([[0, 1, 0, 0], [0, 0, 1, 1], [1, 0, 0, 0], [0, 0, 0, 0]])
         ety = {0: {1: 5}, 1: {2: 9, 3: 14}}
 
         g = qt.adjacency2graph(adj, edge_type=ety, adjust=2)
@@ -86,10 +81,7 @@ def test_adjacency2graph_matrix_adjacency(expected_response1):
     def test_adjacency2graph_matrix_etype(expected_response0):
         # Test adjacency argument using ndarrays work
         adj = {0: {1: {}}, 1: {2: {}, 3: {}}, 2: {0: {}}, 3: {}}
-        ety = np.array([[0, 5, 0, 0],
-                        [0, 0, 9, 14],
-                        [0, 0, 0, 0],
-                        [0, 0, 0, 0]])
+        ety = np.array([[0, 5, 0, 0], [0, 0, 9, 14], [0, 0, 0, 0], [0, 0, 0, 0]])
 
         g = qt.adjacency2graph(adj, edge_type=ety, adjust=1)
         ans = qt.graph2dict(g)
@@ -102,7 +94,6 @@ def test_adjacency2graph_errors():
 
     @staticmethod
     def test_set_types_random():
-
         nV = 1200
         nT = np.random.randint(5, 10)
         g = nx.random_geometric_graph(nV, 0.1).to_directed()
@@ -115,7 +106,7 @@ def test_set_types_random():
         g = qt.set_types_random(g, proportions=pType)
 
         non_loops = [e for e in g.edges() if e[0] != e[1]]
-        mat = [[g.ep(e, 'edge_type') == k for e in non_loops] for k in eType]
+        mat = [[g.ep(e, "edge_type") == k for e in non_loops] for k in eType]
         props = (np.array(mat).sum(1) + 0.0) / len(non_loops)
         ps = np.array([pType[k] for k in eType])
 
diff --git a/tests/test_network.py b/tests/test_network.py
index 412ff00..102e0db 100644
--- a/tests/test_network.py
+++ b/tests/test_network.py
@@ -1,8 +1,9 @@
 import os
-import unittest.mock as mock
+from unittest import mock
 
 try:
     import matplotlib.pyplot as plt
+
     HAS_MATPLOTLIB = True
 except ImportError:
     HAS_MATPLOTLIB = False
@@ -13,8 +14,8 @@
 
 import queueing_tool as qt
 
+CI_TEST = os.environ.get("CI_TEST", False)
 
-CI_TEST = os.environ.get('CI_TEST', False)
 
 @pytest.fixture(scope="module", name="queue_network")
 def fixture_queue_network():
@@ -73,42 +74,44 @@ def test_QueueNetwork_add_arrival(qn):
 
         trans = qn.transitions(False)
 
-        np.testing.assert_allclose(trans[1][2], p0, atol=10**(-2))
-        np.testing.assert_allclose(trans[1][3], p1, atol=10**(-2))
+        np.testing.assert_allclose(trans[1][2], p0, atol=10 ** (-2))
+        np.testing.assert_allclose(trans[1][3], p1, atol=10 ** (-2))
 
     @staticmethod
     @pytest.mark.filterwarnings("ignore:Matplotlib is currently using agg")
-    @pytest.mark.filterwarnings("ignore:Animation was deleted without rendering anything")
+    @pytest.mark.filterwarnings(
+        "ignore:Animation was deleted without rendering anything"
+    )
     def test_QueueNetwork_animate(qn):
         if HAS_MATPLOTLIB:
-            plt.switch_backend('Agg')
+            plt.switch_backend("Agg")
             qn.animate(frames=5)
         else:
-            with mock.patch('queueing_tool.network.queue_network.plt.show'):
+            with mock.patch("queueing_tool.network.queue_network.plt.show"):
                 qn.animate(frames=5)
 
     @staticmethod
     def test_QueueNetwork_blocking(qn):
         g = nx.random_geometric_graph(100, 0.2).to_directed()
-        g = qt.set_types_random(g, proportions={k: 1.0 / 6 for k in range(1, 7)})
+        g = qt.set_types_random(g, proportions=dict.fromkeys(range(1, 7), 1.0 / 6))
         q_cls = {
             1: qt.LossQueue,
             2: qt.QueueServer,
             3: qt.InfoQueue,
             4: qt.ResourceQueue,
             5: qt.ResourceQueue,
-            6: qt.QueueServer
+            6: qt.QueueServer,
         }
 
         q_arg = {
-            3: {'net_size': g.number_of_edges()},
-            4: {'num_servers': 500},
-            6: {'AgentFactory': qt.GreedyAgent}
+            3: {"net_size": g.number_of_edges()},
+            4: {"num_servers": 500},
+            6: {"AgentFactory": qt.GreedyAgent},
         }
 
         qn = qt.QueueNetwork(g, q_classes=q_cls, q_args=q_arg, seed=17)
-        qn.blocking = 'RS'
-        assert qn.blocking == 'RS'
+        qn.blocking = "RS"
+        assert qn.blocking == "RS"
         assert qn._blocking == False
 
         qn.clear()
@@ -116,7 +119,7 @@ def test_QueueNetwork_blocking(qn):
 
     @staticmethod
     def test_QueueNetwork_blocking_setter_error(qn):
-        qn.blocking = 'RS'
+        qn.blocking = "RS"
         with pytest.raises(TypeError):
             qn.blocking = 2
 
@@ -140,20 +143,19 @@ def test_QueueNetwork_closedness(qn):
     @staticmethod
     def test_QueueNetwork_copy(qn):
         g = nx.random_geometric_graph(100, 0.2).to_directed()
-        g = qt.set_types_random(g, proportions={k: 0.2 for k in range(1, 6)})
+        g = qt.set_types_random(g, proportions=dict.fromkeys(range(1, 6), 0.2))
         q_cls = {
             1: qt.LossQueue,
             2: qt.QueueServer,
             3: qt.InfoQueue,
             4: qt.ResourceQueue,
-            5: qt.ResourceQueue
+            5: qt.ResourceQueue,
         }
 
-        q_arg = {3: {'net_size': g.number_of_edges()},
-                 4: {'num_servers': 500}}
+        q_arg = {3: {"net_size": g.number_of_edges()}, 4: {"num_servers": 500}}
 
         qn = qt.QueueNetwork(g, q_classes=q_cls, q_args=q_arg, seed=17)
-        qn.max_agents = np.infty
+        qn.max_agents = np.inf
         qn.initialize(queues=range(g.number_of_edges()))
 
         qn.simulate(n=50000)
@@ -174,21 +176,23 @@ def test_QueueNetwork_copy(qn):
 
     @staticmethod
     def test_QueueNetwork_drawing(qn):
-        with mock.patch('queueing_tool.network.queue_network.HAS_MATPLOTLIB', True):
-            scatter_kwargs = {'c': 'b'}
-            kwargs = {'bgcolor': 'green'}
+        with mock.patch("queueing_tool.network.queue_network.HAS_MATPLOTLIB", True):
+            scatter_kwargs = {"c": "b"}
+            kwargs = {"bgcolor": "green"}
             qn.draw(scatter_kwargs=scatter_kwargs, **kwargs)
-            qn.g.draw_graph.assert_called_with(scatter_kwargs=scatter_kwargs,
-                                                    line_kwargs=None, **kwargs)
+            qn.g.draw_graph.assert_called_with(
+                scatter_kwargs=scatter_kwargs, line_kwargs=None, **kwargs
+            )
 
             qn.draw(scatter_kwargs=scatter_kwargs)
-            bgcolor = qn.colors['bgcolor']
-            qn.g.draw_graph.assert_called_with(scatter_kwargs=scatter_kwargs,
-                                                    line_kwargs=None, bgcolor=bgcolor)
+            bgcolor = qn.colors["bgcolor"]
+            qn.g.draw_graph.assert_called_with(
+                scatter_kwargs=scatter_kwargs, line_kwargs=None, bgcolor=bgcolor
+            )
 
     @staticmethod
     def test_QueueNetwork_drawing_importerror(qn):
-        with mock.patch('queueing_tool.network.queue_network.HAS_MATPLOTLIB', False):
+        with mock.patch("queueing_tool.network.queue_network.HAS_MATPLOTLIB", False):
             with pytest.raises(ImportError):
                 qn.draw()
 
@@ -199,7 +203,7 @@ def test_QueueNetwork_drawing_animation_error(qn):
             qn.animate()
 
         qn.initialize()
-        with mock.patch('queueing_tool.network.queue_network.HAS_MATPLOTLIB', False):
+        with mock.patch("queueing_tool.network.queue_network.HAS_MATPLOTLIB", False):
             with pytest.raises(ImportError):
                 qn.animate()
 
@@ -270,26 +274,15 @@ def ser_id(t):
             return t
 
         adj = {
-            0: {1: {'edge_type': 1}},
-            1: {
-                2: {'edge_type': 2},
-                3: {'edge_type': 2},
-                4: {'edge_type': 2}
-            }
+            0: {1: {"edge_type": 1}},
+            1: {2: {"edge_type": 2}, 3: {"edge_type": 2}, 4: {"edge_type": 2}},
         }
         g = qt.adjacency2graph(adj)
 
         qcl = {1: qt.QueueServer, 2: qt.QueueServer}
         arg = {
-            1: {
-                'arrival_f': arr,
-                'service_f': ser_id,
-                'AgentFactory': qt.GreedyAgent
-            },
-            2: {
-                'service_f': ser,
-                'num_servers': nSe
-            }
+            1: {"arrival_f": arr, "service_f": ser_id, "AgentFactory": qt.GreedyAgent},
+            2: {"service_f": ser, "num_servers": nSe},
         }
 
         qn = qt.QueueNetwork(g, q_classes=qcl, q_args=arg)
@@ -304,9 +297,11 @@ def ser_id(t):
 
         while c < num_events:
             qn.simulate(n=1)
-            if qn.next_event_description() == ('Departure', e01):
+            if qn.next_event_description() == ("Departure", e01):
                 d0 = qn.edge2queue[e01]._departures[0].desired_destination(qn, edg)
-                a1 = np.argmin([qn.edge2queue[e].number_queued() for e in qn.out_edges[1]])
+                a1 = np.argmin(
+                    [qn.edge2queue[e].number_queued() for e in qn.out_edges[1]]
+                )
                 d1 = qn.out_edges[1][a1]
                 ans[c] = d0 == d1
                 c += 1
@@ -324,7 +319,7 @@ def test_QueueNetwork_initialize_Error(qn):
 
         _get_queues_mock = mock.Mock()
         _get_queues_mock.return_value = []
-        mock_location = 'queueing_tool.network.queue_network._get_queues'
+        mock_location = "queueing_tool.network.queue_network._get_queues"
 
         with mock.patch(mock_location, _get_queues_mock):
             with pytest.raises(qt.QueueingToolError):
@@ -428,7 +423,7 @@ def test_QueueNetwork_max_agents(qn):
     @staticmethod
     def test_QueueNetwork_properties(qn):
         qn.clear()
-        assert qn.time == np.infty
+        assert qn.time == np.inf
         assert qn.num_edges == qn.nE
         assert qn.num_vertices == qn.nV
         assert qn.num_nodes == qn.nV
@@ -513,33 +508,24 @@ def test_QueueNetwork_simulate_slow(qn):
             if (edge[0] != edge[1]) == loop:
                 qn._simulate_next_event(slow=True)
                 break
-            else:
-                qn._simulate_next_event(slow=False)
+            qn._simulate_next_event(slow=False)
 
     @staticmethod
     def test_QueueNetwork_show_type(qn):
-        with mock.patch('queueing_tool.network.queue_network.HAS_MATPLOTLIB', True):
-            args = {'c': 'b', 'bgcolor': 'green'}
+        with mock.patch("queueing_tool.network.queue_network.HAS_MATPLOTLIB", True):
+            args = {"c": "b", "bgcolor": "green"}
             qn.show_type(edge_type=2, **args)
             qn.g.draw_graph.assert_called_with(
-                scatter_kwargs=None,
-                line_kwargs=None,
-                **args
+                scatter_kwargs=None, line_kwargs=None, **args
             )
 
     @staticmethod
     def test_QueueNetwork_show_active(qn):
-        with mock.patch('queueing_tool.network.queue_network.HAS_MATPLOTLIB', True):
-            args = {
-                'fname': 'types.png',
-                'figsize': (3, 3),
-                'bgcolor': 'green'
-            }
+        with mock.patch("queueing_tool.network.queue_network.HAS_MATPLOTLIB", True):
+            args = {"fname": "types.png", "figsize": (3, 3), "bgcolor": "green"}
             qn.show_active(**args)
             qn.g.draw_graph.assert_called_with(
-                scatter_kwargs=None,
-                line_kwargs=None,
-                **args
+                scatter_kwargs=None, line_kwargs=None, **args
             )
 
     @staticmethod
@@ -551,7 +537,7 @@ def test_QueueNetwork_sorting(qn):
             queue_times.sort()
             tmp = queue_times[0]
             qn.simulate(n=1)
-            ans[k] = (tmp == qn._qkey[0])
+            ans[k] = tmp == qn._qkey[0]
 
         qn.simulate(n=10000)
 
@@ -560,7 +546,7 @@ def test_QueueNetwork_sorting(qn):
             queue_times.sort()
             tmp = queue_times[0]
             qn.simulate(n=1)
-            ans[k] = (tmp == qn._qkey[0])
+            ans[k] = tmp == qn._qkey[0]
 
         assert ans.all()
 
diff --git a/tests/test_qndigraph.py b/tests/test_qndigraph.py
index b707655..7dd0589 100644
--- a/tests/test_qndigraph.py
+++ b/tests/test_qndigraph.py
@@ -1,5 +1,5 @@
 import os
-import unittest.mock as mock
+from unittest import mock
 
 import matplotlib
 import matplotlib.image
@@ -7,11 +7,9 @@
 import numpy as np
 import pytest
 
-
 import queueing_tool as qt
 
-
-CI_TEST = os.environ.get('CI_TEST', False)
+CI_TEST = os.environ.get("CI_TEST", False)
 
 a_mock = mock.Mock()
 a_mock.pyplot = mock.Mock()
@@ -19,12 +17,13 @@
 a_mock.collections = mock.Mock()
 
 matplotlib_mock = {
-    'matplotlib': a_mock,
-    'matplotlib.pyplot': a_mock.pyplot,
-    'matplotlib.animation': a_mock.animation,
-    'matplotlib.collections': a_mock.collections,
+    "matplotlib": a_mock,
+    "matplotlib.pyplot": a_mock.pyplot,
+    "matplotlib.animation": a_mock.animation,
+    "matplotlib.collections": a_mock.collections,
 }
 
+
 @pytest.fixture(scope="module", name="graph")
 def fixture_queue_network_graph():
     return qt.QueueNetworkDiGraph(nx.krackhardt_kite_graph())
@@ -33,54 +32,50 @@ def fixture_queue_network_graph():
 @pytest.fixture(scope="module", autouse=True)
 def fixture_set_seed():
     np.random.seed(10)
-    return
 
 
 class TestQueueNetworkDiGraph:
     @staticmethod
     def testlines_scatter_args(graph):
-        with mock.patch.dict('sys.modules', matplotlib_mock):
+        with mock.patch.dict("sys.modules", matplotlib_mock):
             ax = mock.Mock()
             ax.transData = mock.Mock()
-            line_args = {'linewidths': 77, 'vmax': 107}
-            scat_args = {'vmax': 107}
-            kwargs = {'pos': {v: (910, 10) for v in graph.nodes()}}
+            line_args = {"linewidths": 77, "vmax": 107}
+            scat_args = {"vmax": 107}
+            kwargs = {"pos": dict.fromkeys(graph.nodes(), (910, 10))}
 
             a, b = graph.lines_scatter_args(line_args, scat_args, **kwargs)
 
-            assert a['linewidths'] == 77
-            assert b['vmax'] == 107
-            assert 'beefy' not in a and 'beefy' not in b
+            assert a["linewidths"] == 77
+            assert b["vmax"] == 107
+            assert "beefy" not in a and "beefy" not in b
 
     @pytest.mark.xfail
     def test_draw_graph(graph):
         pos = np.random.uniform(size=(graph.number_of_nodes(), 2))
-        kwargs = {
-            'fname': 'test1.png',
-            'pos': pos
-        }
-        graph.draw_graph(scatter_kwargs={'s': 100}, **kwargs)
+        kwargs = {"fname": "test1.png", "pos": pos}
+        graph.draw_graph(scatter_kwargs={"s": 100}, **kwargs)
 
-        version = 1 if matplotlib.__version__.startswith('1') else 2
-        filename = 'test-mpl-{version}.x.png'.format(version=version)
+        version = 1 if matplotlib.__version__.startswith("1") else 2
+        filename = f"test-mpl-{version}.x.png"
 
-        img0 = matplotlib.image.imread('tests/img/{filename}'.format(filename=filename))
-        img1 = matplotlib.image.imread('test1.png')
+        img0 = matplotlib.image.imread(f"tests/img/{filename}")
+        img1 = matplotlib.image.imread("test1.png")
 
-        if os.path.exists('test1.png'):
-            os.remove('test1.png')
+        if os.path.exists("test1.png"):
+            os.remove("test1.png")
 
         pixel_diff = (img0 != img1).flatten()
         num_pixels = pixel_diff.shape[0] + 0.0
         assert pixel_diff.sum() / num_pixels < 0.0001
 
-        with mock.patch('queueing_tool.graph.graph_wrapper.HAS_MATPLOTLIB', False):
+        with mock.patch("queueing_tool.graph.graph_wrapper.HAS_MATPLOTLIB", False):
             with pytest.raises(ImportError):
                 graph.draw_graph()
 
-        kwargs = {'pos': 1}
+        kwargs = {"pos": 1}
         graph.set_pos = mock.MagicMock()
-        with mock.patch.dict('sys.modules', matplotlib_mock):
+        with mock.patch.dict("sys.modules", matplotlib_mock):
             graph.draw_graph(**kwargs)
 
         graph.set_pos.assert_called_once_with(1)
diff --git a/tests/test_queue_server.py b/tests/test_queue_server.py
index 174e0e2..d8abefb 100644
--- a/tests/test_queue_server.py
+++ b/tests/test_queue_server.py
@@ -18,7 +18,7 @@ def departure_rate():
 
 
 class TestQueueServers:
-    @staticmethod    
+    @staticmethod
     def test_QueueServer_init_errors():
         with pytest.raises(TypeError):
             qt.QueueServer(num_servers=3.0)
@@ -35,7 +35,7 @@ def test_QueueServer_set_num_servers():
         q.set_num_servers(2 * nSe)
 
         Se2 = q.num_servers
-        q.set_num_servers(np.infty)
+        q.set_num_servers(np.inf)
 
         assert Se1 == nSe
         assert Se2 == 2 * nSe
@@ -90,7 +90,6 @@ def r(t):
 
     @staticmethod
     def test_QueueServer_accounting(arrival_rate, departure_rate):
-
         nSe = np.random.randint(1, 10)
         mu = arrival_rate / (departure_rate * nSe)
 
@@ -226,10 +225,8 @@ def ser(t):
     @staticmethod
     def test_NullQueue_data_collection():
         adj = {
-            0: {1: {'edge_type': 1}},
-            1: {2: {'edge_type': 2},
-                3: {'edge_type': 2},
-                4: {'edge_type': 2}}
+            0: {1: {"edge_type": 1}},
+            1: {2: {"edge_type": 2}, 3: {"edge_type": 2}, 4: {"edge_type": 2}},
         }
         g = qt.adjacency2graph(adj)
 
@@ -251,7 +248,7 @@ def test_NullQueue_data_collection():
     def test_ResourceQueue_network():
         g = nx.random_geometric_graph(100, 0.2).to_directed()
         q_cls = {1: qt.ResourceQueue, 2: qt.ResourceQueue}
-        q_arg = {1: {'num_servers': 50}, 2: {'num_servers': 500}}
+        q_arg = {1: {"num_servers": 50}, 2: {"num_servers": 500}}
 
         qn = qt.QueueNetwork(g, q_classes=q_cls, q_args=q_arg)
         qn.max_agents = 400000
@@ -266,9 +263,10 @@ def test_ResourceQueue_network():
     def test_ResourceQueue_network_data_collection():
         g = qt.generate_random_graph(100)
         q_cls = {1: qt.ResourceQueue, 2: qt.ResourceQueue}
-        q_arg = {1: {'num_servers': 500},
-                 2: {'num_servers': 500,
-                     'AgentFactory': qt.Agent}}
+        q_arg = {
+            1: {"num_servers": 500},
+            2: {"num_servers": 500, "AgentFactory": qt.Agent},
+        }
 
         qn = qt.QueueNetwork(g, q_classes=q_cls, q_args=q_arg)
         qn.max_agents = 40000
@@ -283,27 +281,26 @@ def test_ResourceQueue_network_data_collection():
     def test_ResourceQueue_network_current_color():
         q = qt.ResourceQueue(num_servers=50)
         ans = q._current_color(0)
-        col = q.colors['vertex_fill_color']
-        col = [i * (0.9 - 1. / 6) / 0.9 for i in col]
+        col = q.colors["vertex_fill_color"]
+        col = [i * (0.9 - 1.0 / 6) / 0.9 for i in col]
         col[3] = 1.0
         assert ans == col
 
         ans = q._current_color(1)
-        col = q.colors['edge_loop_color']
-        col = [i * (0.9 - 1. / 6) / 0.9 for i in col]
+        col = q.colors["edge_loop_color"]
+        col = [i * (0.9 - 1.0 / 6) / 0.9 for i in col]
         col[3] = 0
         assert ans == col
 
         ans = q._current_color(2)
-        col = q.colors['vertex_pen_color']
+        col = q.colors["vertex_pen_color"]
         assert ans == col
 
     @staticmethod
     def test_InfoQueue_network():
-
         g = nx.random_geometric_graph(100, 0.2).to_directed()
         q_cls = {1: qt.InfoQueue}
-        q_arg = {1: {'net_size': g.number_of_edges()}}
+        q_arg = {1: {"net_size": g.number_of_edges()}}
 
         qn = qt.QueueNetwork(g, q_classes=q_cls, q_args=q_arg, seed=17)
         qn.max_agents = 40000
@@ -314,7 +311,6 @@ def test_InfoQueue_network():
 
     @staticmethod
     def test_Agent_compare():
-
         a0 = qt.Agent()
         a1 = qt.Agent()
         assert a0 == a1
@@ -367,7 +363,7 @@ def ser(t):
         queue.collect_data = True
         queue.set_active()
 
-        # We want 7 agents in the system. After the 
+        # We want 7 agents in the system. After the
         queue.simulate(nA=6)
         queue.set_inactive()
         queue.next_event()
diff --git a/tests/test_statistical_properties.py b/tests/test_statistical_properties.py
index 80aac84..cdb2b01 100644
--- a/tests/test_statistical_properties.py
+++ b/tests/test_statistical_properties.py
@@ -7,13 +7,13 @@
 
 import queueing_tool as qt
 
-
-CI_TEST = os.environ.get('CI_TEST', "false") == "true"
+CI_TEST = os.environ.get("CI_TEST", "false") == "true"
 
 
 def empirical_cdf0(x, z, n):
     return np.sum(z <= x) / n
 
+
 empirical_cdf = np.vectorize(empirical_cdf0, excluded={1, 2})
 
 
@@ -35,7 +35,6 @@ class TestQueueServers:
     @staticmethod
     @pytest.mark.skipif(CI_TEST, reason="Test takes long.")
     def test_Markovian_QueueServer(arrival_rate, departure_rate):
-
         nSe = np.random.randint(1, 10)
         mu = arrival_rate / (departure_rate * nSe)
 
@@ -49,7 +48,7 @@ def ser(t):
         n = 50000
 
         q.set_active()
-        q.simulate(n=20000)    # Burn in period
+        q.simulate(n=20000)  # Burn in period
         q.collect_data = True
         q.simulate(nD=n + 1)
         dat = q.fetch_data()
@@ -64,23 +63,22 @@ def ser(t):
         nbin = n // 6 - 1  # np.floor( np.exp( lam * upb) - 1 )
         bins = np.zeros(nbin + 2)
         bins[1:-1] = upb - np.log(np.arange(nbin, 0, -1)) / lamh
-        bins[-1] = np.infty
+        bins[-1] = np.inf
 
         N = np.histogram(dep, bins=bins)[0]
         pp = 1 - np.exp(-lamh * bins)
         pr = n * (pp[1:] - pp[:-1])
-        Q = np.sum((N - pr)**2 / pr)
+        Q = np.sum((N - pr) ** 2 / pr)
         p1 = 1 - chi2_cdf(Q, nbin - 1)
 
         x, y = dep[1:], dep[:-1]
         cc = np.corrcoef(x, y)[0, 1]
-        np.testing.assert_allclose(cc, 0, atol=10**(-1))
+        np.testing.assert_allclose(cc, 0, atol=10 ** (-1))
         assert p1 > 0.05
 
     @staticmethod
     @pytest.mark.skipif(CI_TEST, reason="Test takes long.")
     def test_QueueServer_Littleslaw(arrival_rate, departure_rate):
-
         nSe = np.random.randint(1, 10)
         mu = arrival_rate / (departure_rate * nSe)
 
@@ -94,7 +92,7 @@ def ser(t):
         n = 500000
 
         q.set_active()
-        q.simulate(n=n)    # Burn in period
+        q.simulate(n=n)  # Burn in period
         q.collect_data = True
         q.simulate(n=n + 1)
         data = q.fetch_data()
@@ -104,12 +102,11 @@ def ser(t):
         wait = data[ind, 1] - data[ind, 0]
         ans = np.mean(wait) * arrival_rate - np.mean(data[:, 3]) * departure_rate
 
-        np.testing.assert_allclose(ans, 0, atol=10**(-1))
+        np.testing.assert_allclose(ans, 0, atol=10 ** (-1))
 
     @staticmethod
     @pytest.mark.skipif(CI_TEST, reason="Test takes long.")
     def test_LossQueue_blocking(arrival_rate, departure_rate):
-
         nSe = np.random.randint(1, 10)
         mu = arrival_rate / (departure_rate * nSe)
         k = np.random.randint(5, 15)
@@ -123,7 +120,7 @@ def ser(t):
 
         q2 = qt.LossQueue(num_servers=nSe, arrival_f=arr, service_f=ser)
         q2.set_active()
-        q2.simulate(n=100000)    # Burn in period
+        q2.simulate(n=100000)  # Burn in period
 
         nA0 = q2.num_arrivals[1]
         nB0 = q2.num_blocked
@@ -137,9 +134,9 @@ def ser(t):
         f = np.array([math.factorial(j) for j in range(nSe + 1)])
 
         pois_pmf = np.exp(-a) * a**nSe / math.factorial(nSe)
-        pois_cdf = np.sum(np.exp(-a) * a**np.arange(nSe + 1) / f)
+        pois_cdf = np.sum(np.exp(-a) * a ** np.arange(nSe + 1) / f)
         p_block = (nB1 - nB0 + 0.0) / (nA1 - nA0)
-        np.testing.assert_allclose(pois_pmf / pois_cdf, p_block, atol=10**(-2))
+        np.testing.assert_allclose(pois_pmf / pois_cdf, p_block, atol=10 ** (-2))
 
 
 class TestRandomMeasure:
@@ -154,7 +151,7 @@ def test_poisson_random_measure():
         # This test should fail some percentage of the time
 
         def rate(t):
-            return 0.5 + 4 * np.sin(np.pi * t / 12)**2
+            return 0.5 + 4 * np.sin(np.pi * t / 12) ** 2
 
         arr_f = functools.partial(qt.poisson_random_measure, rate=rate, rate_max=4.5)
 
@@ -169,12 +166,16 @@ def rate(t):
                 if t > 12:
                     break
 
-        mu1 = 5 * np.sum(rate(np.linspace(3, 8, 200))) / 200  # or 2*(5 + (sqrt(3) + 2) * 3/pi) + 2.5
-        mu2 = 4 * np.sum(rate(np.linspace(8, 12, 200))) / 200  # or 2*(4 - 3*sqrt(3)/pi) + 2
+        mu1 = (
+            5 * np.sum(rate(np.linspace(3, 8, 200))) / 200
+        )  # or 2*(5 + (sqrt(3) + 2) * 3/pi) + 2.5
+        mu2 = (
+            4 * np.sum(rate(np.linspace(8, 12, 200))) / 200
+        )  # or 2*(4 - 3*sqrt(3)/pi) + 2
         mus = [mu1, mu2]
 
-        rv1 = np.sum(np.logical_and(3 < arrival_times, arrival_times < 8), axis=1)
-        rv2 = np.sum(np.logical_and(8 < arrival_times, arrival_times < 12), axis=1)
+        rv1 = np.sum(np.logical_and(arrival_times > 3, arrival_times < 8), axis=1)
+        rv2 = np.sum(np.logical_and(arrival_times > 8, arrival_times < 12), axis=1)
         rvs = [rv1, rv2]
         df = [max(rv1) + 2, max(rv2) + 2]
 
@@ -182,11 +183,11 @@ def rate(t):
 
         for i, sample in enumerate(rvs):
             for k in range(df[i] - 1):
-                pi_hat = nSamp * np.exp(-mus[i]) * mus[i]**k / math.factorial(k)
-                Q[k, i] = (np.sum(sample == k) - pi_hat)**2 / pi_hat
+                pi_hat = nSamp * np.exp(-mus[i]) * mus[i] ** k / math.factorial(k)
+                Q[k, i] = (np.sum(sample == k) - pi_hat) ** 2 / pi_hat
 
             ans = np.array([math.factorial(j) for j in range(k + 1)])
-            pois_cdf = np.sum(np.exp(-mus[i]) * mus[i]**np.arange(k + 1) / ans)
+            pois_cdf = np.sum(np.exp(-mus[i]) * mus[i] ** np.arange(k + 1) / ans)
             Q[k + 1, i] = nSamp * (1 - pois_cdf)
 
         Qs = np.sum(Q, axis=0)