Repository: xgi-org/xgi (language: python)

Issue:
Hypergraph() can take sparse arrays as input - document this explicitly
Right now we just say that, for example, the incidence matrix can be an scipy or numpy array.


Discussion / hints:

Fix patch:
diff --git a/xgi/convert/bipartite_edges.py b/xgi/convert/bipartite_edges.py
index 09b644dc8..dfd01199c 100644
--- a/xgi/convert/bipartite_edges.py
+++ b/xgi/convert/bipartite_edges.py
@@ -1,3 +1,5 @@
+"""Methods for converting to and from bipartite edgelists."""
+
 from ..generators import empty_hypergraph
 
 __all__ = ["from_bipartite_edgelist", "to_bipartite_edgelist"]
diff --git a/xgi/convert/bipartite_graph.py b/xgi/convert/bipartite_graph.py
index 7fa6a037f..40f3e0ea9 100644
--- a/xgi/convert/bipartite_graph.py
+++ b/xgi/convert/bipartite_graph.py
@@ -1,3 +1,5 @@
+"""Methods for converting to and from bipartite graphs."""
+
 import networkx as nx
 from networkx import bipartite
 
diff --git a/xgi/convert/graph.py b/xgi/convert/graph.py
index 51c7ab53d..94499a66e 100644
--- a/xgi/convert/graph.py
+++ b/xgi/convert/graph.py
@@ -1,3 +1,5 @@
+"""Method for projecting a hypergraph to a graph."""
+
 import networkx as nx
 
 from ..linalg import adjacency_matrix
diff --git a/xgi/convert/higher_order_network.py b/xgi/convert/higher_order_network.py
index e5dd66a01..81574b70f 100644
--- a/xgi/convert/higher_order_network.py
+++ b/xgi/convert/higher_order_network.py
@@ -1,3 +1,5 @@
+"""Methods for converting to higher-order network objects."""
+
 from copy import deepcopy
 
 import pandas as pd
diff --git a/xgi/convert/hyperedges.py b/xgi/convert/hyperedges.py
index 0429a750b..e954c024d 100644
--- a/xgi/convert/hyperedges.py
+++ b/xgi/convert/hyperedges.py
@@ -1,3 +1,5 @@
+"""Methods for converting to and from hyperedge lists."""
+
 from ..core import SimplicialComplex
 from ..generators import empty_hypergraph
 
diff --git a/xgi/convert/hypergraph_dict.py b/xgi/convert/hypergraph_dict.py
index b962c4e57..f47b1f8f9 100644
--- a/xgi/convert/hypergraph_dict.py
+++ b/xgi/convert/hypergraph_dict.py
@@ -1,3 +1,5 @@
+"""Method for converting from a standardized dictionary."""
+
 from ..exception import XGIError
 from ..generators import empty_hypergraph
 
diff --git a/xgi/convert/incidence.py b/xgi/convert/incidence.py
index d01ad0adb..c9728a49c 100644
--- a/xgi/convert/incidence.py
+++ b/xgi/convert/incidence.py
@@ -1,3 +1,5 @@
+"""Methods for converting to and from an incidence matrix."""
+
 from scipy.sparse import coo_array
 
 from ..exception import XGIError
diff --git a/xgi/convert/line_graph.py b/xgi/convert/line_graph.py
index cf492bb7b..4ccdf3b04 100644
--- a/xgi/convert/line_graph.py
+++ b/xgi/convert/line_graph.py
@@ -1,3 +1,5 @@
+"""Method for converting to a line graph."""
+
 from itertools import combinations
 
 import networkx as nx
diff --git a/xgi/convert/pandas.py b/xgi/convert/pandas.py
index d905ee898..d91a4e78b 100644
--- a/xgi/convert/pandas.py
+++ b/xgi/convert/pandas.py
@@ -1,3 +1,5 @@
+"""Methods for converting to and from a Pandas dataframe."""
+
 from collections import defaultdict
 
 import pandas as pd
diff --git a/xgi/convert/simplex.py b/xgi/convert/simplex.py
index 2c71d4f04..765e30ce3 100644
--- a/xgi/convert/simplex.py
+++ b/xgi/convert/simplex.py
@@ -1,3 +1,5 @@
+"""Methods for simplicial complexes."""
+
 from ..core import Hypergraph, SimplicialComplex
 from ..exception import XGIError
 from ..generators import empty_simplicial_complex
diff --git a/xgi/core/hypergraph.py b/xgi/core/hypergraph.py
index a7f77966e..3f77768b6 100644
--- a/xgi/core/hypergraph.py
+++ b/xgi/core/hypergraph.py
@@ -34,7 +34,7 @@ class Hypergraph:
         * hyperedge list
         * hyperedge dictionary
         * 2-column Pandas dataframe (bipartite edges)
-        * Scipy/Numpy incidence matrix
+        * Incidence matrix: numpy ndarray or scipy.sparse array
         * Hypergraph object
         * SimplicialComplex object
 
diff --git a/xgi/core/simplicialcomplex.py b/xgi/core/simplicialcomplex.py
index 1b59a7d36..17e65198b 100644
--- a/xgi/core/simplicialcomplex.py
+++ b/xgi/core/simplicialcomplex.py
@@ -44,7 +44,7 @@ class SimplicialComplex(Hypergraph):
         * simplex list
         * simplex dictionary
         * 2-column Pandas dataframe (bipartite edges)
-        * Scipy/Numpy incidence matrix
+        * Incidence matrix: numpy ndarray or scipy.sparse array
         * SimplicialComplex object
         * Hypergraph object
 
diff --git a/xgi/drawing/draw.py b/xgi/drawing/draw.py
index 238de0514..6194a06ce 100644
--- a/xgi/drawing/draw.py
+++ b/xgi/drawing/draw.py
@@ -1,24 +1,27 @@
 """Draw hypergraphs and simplicial complexes with matplotlib."""
 
-from collections.abc import Iterable
 from inspect import signature
 from itertools import combinations
 
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib import cm
-from matplotlib.colors import LinearSegmentedColormap, ListedColormap
 from matplotlib.patches import FancyArrow
 from mpl_toolkits.mplot3d.art3d import Line3DCollection, Poly3DCollection
 from networkx import spring_layout
-from numpy import ndarray
 from scipy.spatial import ConvexHull
 
 from .. import convert
-from ..algorithms import max_edge_order
+from ..algorithms import max_edge_order, unique_edge_sizes
 from ..core import DiHypergraph, Hypergraph, SimplicialComplex
 from ..exception import XGIError
-from ..stats import IDStat
+from .draw_utils import (
+    _CCW_sort,
+    _color_arg_to_dict,
+    _draw_init,
+    _scalar_arg_to_dict,
+    _update_lims,
+)
 from .layout import _augmented_projection, barycenter_spring_layout
 
 __all__ = [
@@ -112,8 +115,10 @@ def draw(
         If True, draw ids on the hyperedges. If a dict, must contain (edge_id: label)
         pairs.  By default, False.
     aspect : {"auto", "equal"} or float, optional
-        Set the aspect ratio of the axes scaling, i.e. y/x-scale. `aspect` is passed 
-        directly to matplotlib's `ax.set_aspect()`. Default is `equal`. See full 
+        Set the aspect ratio of the axes scaling, i.e. y/x-scale. `aspect` is passed
+        directly to matplotlib's `ax.set_aspect()`. Default is `equal`. See full
+        Set the aspect ratio of the axes scaling, i.e. y/x-scale. `aspect` is passed
+        directly to matplotlib's `ax.set_aspect()`. Default is `equal`. See full
         description at
         https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.set_aspect.html
     **kwargs : optional args
@@ -163,15 +168,7 @@ def draw(
     if edge_fc is None:
         edge_fc = H.edges.size
 
-    if pos is None:
-        pos = barycenter_spring_layout(H)
-
-    if ax is None:
-        ax = plt.gca()
-
-    ax.get_xaxis().set_ticks([])
-    ax.get_yaxis().set_ticks([])
-    ax.axis("off")
+    ax, pos = _draw_init(H, ax, pos)
 
     if not max_order:
         max_order = max_edge_order(H)
@@ -311,15 +308,7 @@ def draw_nodes(
 
     settings.update(kwargs)
 
-    if pos is None:
-        pos = barycenter_spring_layout(H)
-
-    if ax is None:
-        ax = plt.gca()
-
-    ax.get_xaxis().set_ticks([])
-    ax.get_yaxis().set_ticks([])
-    ax.axis("off")
+    ax, pos = _draw_init(H, ax, pos)
 
     # Note Iterable covers lists, tuples, ranges, generators, np.ndarrays, etc
     node_fc = _color_arg_to_dict(node_fc, H.nodes, settings["node_fc_cmap"])
@@ -436,11 +425,7 @@ def draw_hyperedges(
 
     """
 
-    if pos is None:
-        pos = barycenter_spring_layout(H)
-
-    if ax is None:
-        ax = plt.gca()
+    ax, pos = _draw_init(H, ax, pos)
 
     if max_order is None:
         max_order = max_edge_order(H)
@@ -448,10 +433,6 @@ def draw_hyperedges(
     if edge_fc is None:
         edge_fc = H.edges.size
 
-    ax.get_xaxis().set_ticks([])
-    ax.get_yaxis().set_ticks([])
-    ax.axis("off")
-
     if settings is None:
         settings = {
             "min_dyad_lw": 2.0,
@@ -604,19 +585,11 @@ def draw_simplices(
     if not max_order:
         max_order = max_edge_order(H_)
 
-    if pos is None:
-        pos = barycenter_spring_layout(H_)
-
-    if ax is None:
-        ax = plt.gca()
+    ax, pos = _draw_init(H_, ax, pos)
 
     if edge_fc is None:
         edge_fc = H_.edges.size
 
-    ax.get_xaxis().set_ticks([])
-    ax.get_yaxis().set_ticks([])
-    ax.axis("off")
-
     if settings is None:
         settings = {
             "min_dyad_lw": 2.0,
@@ -699,215 +672,6 @@ def draw_simplices(
     return ax
 
 
-def _scalar_arg_to_dict(scalar_arg, ids, min_val, max_val):
-    """Map different types of arguments for drawing style to a dict with scalar values.
-
-    Parameters
-    ----------
-    scalar_arg : int, float, dict, iterable, or NodeStat/EdgeStat
-        Attributes for drawing parameter.
-    ids : NodeView or EdgeView
-        This is the node or edge IDs that attributes get mapped to.
-    min_val : int or float
-        The minimum value of the drawing parameter.
-    max_val : int or float
-        The maximum value of the drawing parameter.
-
-    Returns
-    -------
-    dict
-        An ID: scalar dictionary.
-
-    Raises
-    ------
-    TypeError
-        If a int, float, list, dict, or NodeStat/EdgeStat is not passed.
-    """
-    if isinstance(scalar_arg, str):
-        raise TypeError(
-            "Argument must be int, float, dict, iterable, "
-            f"or NodeStat/EdgeStat. Received {type(scalar_arg)}"
-        )
-
-    # Single argument
-    if isinstance(scalar_arg, (int, float)):
-        return {id: scalar_arg for id in ids}
-
-    # IDStat
-    if isinstance(scalar_arg, IDStat):
-        vals = np.interp(
-            scalar_arg.asnumpy(),
-            [scalar_arg.min(), scalar_arg.max()],
-            [min_val, max_val],
-        )
-        return dict(zip(ids, vals))
-
-    # Iterables of floats or ints
-    if isinstance(scalar_arg, Iterable):
-        if isinstance(scalar_arg, dict):
-            try:
-                return {id: float(scalar_arg[id]) for id in scalar_arg if id in ids}
-            except ValueError as e:
-                raise TypeError(
-                    "The input dict must have values that can be cast to floats."
-                )
-
-        elif isinstance(scalar_arg, (list, ndarray)):
-            try:
-                return {id: float(scalar_arg[idx]) for idx, id in enumerate(ids)}
-            except ValueError as e:
-                raise TypeError(
-                    "The input list or array must have values that can be cast to floats."
-                )
-        else:
-            raise TypeError(
-                "Argument must be an dict, list, or numpy array of floats or ints."
-            )
-
-    raise TypeError(
-        "Argument must be int, float, dict, iterable, "
-        f"or NodeStat/EdgeStat. Received {type(scalar_arg)}"
-    )
-
-
-def _color_arg_to_dict(color_arg, ids, cmap):
-    """Map different types of arguments for drawing style to a dict with color values.
-
-    Parameters
-    ----------
-    color_arg : Several formats are accepted:
-
-        Single color values
-
-        * str
-        * 3- or 4-tuple
-
-        Iterable of colors (each color specified as above)
-
-        * numpy array
-        * list
-        * dict {id: color} pairs
-
-        Iterable of numerical values (floats or ints)
-
-        * list
-        * dict
-        * numpy array
-
-        Stats
-
-        * NodeStat
-        * EdgeStat
-
-        Attributes for drawing parameter.
-    ids : NodeView or EdgeView
-        This is the node or edge IDs that attributes get mapped to.
-    cmap : ListedColormap or LinearSegmentedColormap
-        colormap to use for NodeStat/EdgeStat.
-
-    Returns
-    -------
-    dict
-        An ID: color dictionary.
-
-    Raises
-    ------
-    TypeError
-        If a string, dict, iterable, or NodeStat/EdgeStat is not passed.
-
-    Notes
-    -----
-    For the iterable of values, we do not accept tuples,
-    because there is the potential for ambiguity.
-    """
-
-    # single argument. Must be a string or a tuple of floats
-    if isinstance(color_arg, str) or (
-        isinstance(color_arg, tuple) and isinstance(color_arg[0], float)
-    ):
-        return {id: color_arg for id in ids}
-
-    # Iterables of colors. The values of these iterables must strings or tuples. As of now,
-    # there is not a check to verify that the tuples contain floats.
-    if isinstance(color_arg, Iterable):
-        if isinstance(color_arg, dict) and isinstance(
-            next(iter(color_arg.values())), (str, tuple, ndarray)
-        ):
-            return {id: color_arg[id] for id in color_arg if id in ids}
-        if isinstance(color_arg, (list, ndarray)) and isinstance(
-            color_arg[0], (str, tuple, ndarray)
-        ):
-            return {id: color_arg[idx] for idx, id in enumerate(ids)}
-
-    # Stats or iterable of values
-    if isinstance(color_arg, (Iterable, IDStat)):
-        # set max and min of interpolation based on color map
-        if isinstance(cmap, ListedColormap):
-            minval = 0
-            maxval = cmap.N
-        elif isinstance(cmap, LinearSegmentedColormap):
-            minval = 0.1
-            maxval = 0.9
-        else:
-            raise XGIError("Invalid colormap!")
-
-        # handle the case of IDStat vs iterables
-        if isinstance(color_arg, IDStat):
-            vals = np.interp(
-                color_arg.asnumpy(),
-                [color_arg.min(), color_arg.max()],
-                [minval, maxval],
-            )
-            return {
-                id: np.array(cmap(vals[i])).reshape(1, -1) for i, id in enumerate(ids)
-            }
-
-        elif isinstance(color_arg, Iterable):
-            if isinstance(color_arg, dict) and isinstance(
-                next(iter(color_arg.values())), (int, float)
-            ):
-                v = list(color_arg.values())
-                vals = np.interp(v, [np.min(v), np.max(v)], [minval, maxval])
-                # because we have ids, we can't just assume that the keys of arg correspond to
-                # the ids.
-                return {
-                    id: np.array(cmap(v)).reshape(1, -1)
-                    for v, id in zip(vals, color_arg.keys())
-                    if id in ids
-                }
-
-            if isinstance(color_arg, (list, ndarray)) and isinstance(
-                color_arg[0], (int, float)
-            ):
-                vals = np.interp(
-                    color_arg, [np.min(color_arg), np.max(color_arg)], [minval, maxval]
-                )
-                return {
-                    id: np.array(cmap(vals[i])).reshape(1, -1)
-                    for i, id in enumerate(ids)
-                }
-            else:
-                raise TypeError(
-                    "Argument must be an dict, list, or numpy array of floats."
-                )
-
-    raise TypeError(
-        "Argument must be str, dict, iterable, or "
-        f"NodeStat/EdgeStat. Received {type(color_arg)}"
-    )
-
-
-def _CCW_sort(p):
-    """
-    Sort the input 2D points counterclockwise.
-    """
-    p = np.array(p)
-    mean = np.mean(p, axis=0)
-    d = p - mean
-    s = np.arctan2(d[:, 0], d[:, 1])
-    return p[np.argsort(s), :]
-
-
 def draw_node_labels(
     H,
     pos,
@@ -1141,24 +905,6 @@ def draw_hyperedge_labels(
     return text_items
 
 
-def _update_lims(pos, ax):
-    """Update Axis limits based on node positions"""
-
-    # compute axis limits
-    pos_arr = np.asarray([[x, y] for n, (x, y) in pos.items()])
-
-    maxx, maxy = np.max(pos_arr, axis=0)
-    minx, miny = np.min(pos_arr, axis=0)
-    w = maxx - minx
-    h = maxy - miny
-
-    # update view after drawing
-    padx, pady = 0.05 * w, 0.05 * h
-    corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
-    ax.update_datalim(corners)
-    ax.autoscale_view()
-
-
 def _draw_hull(node_pos, ax, edges_ec, facecolor, alpha, zorder, radius):
     """Draw a convex hull encompassing the nodes in node_pos
 
@@ -1283,8 +1029,8 @@ def draw_hypergraph_hull(
     radius : float, optional
         Radius of the convex hull in the vicinity of the nodes, by default 0.05.
     aspect : {"auto", "equal"} or float, optional
-        Set the aspect ratio of the axes scaling, i.e. y/x-scale. `aspect` is passed 
-        directly to matplotlib's `ax.set_aspect()`. Default is `equal`. See full 
+        Set the aspect ratio of the axes scaling, i.e. y/x-scale. `aspect` is passed
+        directly to matplotlib's `ax.set_aspect()`. Default is `equal`. See full
         description at
         https://matplotlib.org/stable/api/_as_gen/matplotlib.axes.Axes.set_aspect.html
     **kwargs : optional args
@@ -1337,15 +1083,7 @@ def draw_hypergraph_hull(
 
     settings.update(kwargs)
 
-    if pos is None:
-        pos = barycenter_spring_layout(H)
-
-    if ax is None:
-        ax = plt.gca()
-
-    ax.get_xaxis().set_ticks([])
-    ax.get_yaxis().set_ticks([])
-    ax.axis("off")
+    ax, pos = _draw_init(H, ax, pos)
 
     if not max_order:
         max_order = max_edge_order(H)
@@ -1419,19 +1157,23 @@ def draw_multilayer(
     H,
     pos=None,
     ax=None,
-    node_fc="tab:blue",
+    dyad_color="black",
+    dyad_lw=0.5,
+    edge_fc=None,
+    node_fc="white",
     node_ec="black",
     node_lw=0.5,
     node_size=5,
+    plane_color="grey",
     max_order=None,
-    palette="jet",
     conn_lines=True,
     conn_lines_style="dotted",
     width=5,
     height=5,
     h_angle=10,
-    v_angle=0,
+    v_angle=20,
     sep=1,
+    **kwargs,
 ):
     """Draw a hypergraph or simplicial complex visualized in 3D
     showing hyperedges/simplices of different orders on superimposed layers.
@@ -1441,21 +1183,57 @@ def draw_multilayer(
     H : Hypergraph or SimplicialComplex.
         Higher-order network to plot.
     pos : dict or None, optional
-        The positions of the nodes in the multilayer network. If None, a default layout will be computed using xgi.barycenter_spring_layout(). Default is None.
+        The positions of the nodes in the multilayer network.
+        If None, a default layout will be computed using
+        xgi.barycenter_spring_layout(). Default is None.
     ax : matplotlib Axes3DSubplot or None, optional
-        The subplot to draw the visualization on. If None, a new subplot will be created. Default is None.
-    node_fc : color or sequence of colors, optional
-        The face color(s) of the nodes. Default is "tab:blue".
-    node_ec : color or sequence of colors, optional
-        The edge color(s) of the nodes. Default is "black".
-    node_lw : float or sequence of floats, optional
-        The linewidth(s) of the node edges. Default is 0.5.
-    node_size : scalar or array-like, optional
-        The size(s) of the nodes. Default is 5.
-    max_order : int or None, optional
-        The maximum order of hyperedges/simplices to consider for coloring. If None edges up to the maximal order are drawn. Default is None.
-    palette : str, optional
-        The name of the matplotlib color palette to use. Default is 'jet'.
+        The subplot to draw the visualization on.
+        If None, a new subplot will be created. Default is None.
+    dyad_color : str, dict, iterable, or EdgeStat, optional
+        Color of the dyadic links.  If str, use the same color for all edges. If a dict,
+        must contain (edge_id: color_str) pairs.  If iterable, assume the colors are
+        specified in the same order as the edges are found in H.edges. If EdgeStat, use
+        a colormap (specified with dyad_color_cmap) associated to it. By default,
+        "black".
+    dyad_lw : int, float, dict, iterable, or EdgeStat, optional
+        Line width of edges of order 1 (dyadic links).  If int or float, use the same
+        width for all edges.  If a dict, must contain (edge_id: width) pairs.  If
+        iterable, assume the widths are specified in the same order as the edges are
+        found in H.edges. If EdgeStat, use a monotonic linear interpolation defined
+        between min_dyad_lw and max_dyad_lw. By default, 0.5.
+    edge_fc : str, dict, iterable, or EdgeStat, optional
+        Color of the hyperedges.  If str, use the same color for all nodes.  If a dict,
+        must contain (edge_id: color_str) pairs.  If other iterable, assume the colors
+        are specified in the same order as the hyperedges are found in H.edges. If
+        EdgeStat, use the colormap specified with edge_fc_cmap. If None (default), use
+        the H.edges.size.
+    node_fc : str, dict, iterable, or NodeStat, optional
+        Color of the nodes.  If str, use the same color for all nodes.  If a dict, must
+        contain (node_id: color_str) pairs.  If other iterable, assume the colors are
+        specified in the same order as the nodes are found in H.nodes. If NodeStat, use
+        the colormap specified with node_fc_cmap. By default, "white".
+    node_ec : str, dict, iterable, or NodeStat, optional
+        Color of node borders.  If str, use the same color for all nodes.  If a dict,
+        must contain (node_id: color_str) pairs.  If other iterable, assume the colors
+        are specified in the same order as the nodes are found in H.nodes. If NodeStat,
+        use the colormap specified with node_ec_cmap. By default, "black".
+    node_lw : int, float, dict, iterable, or NodeStat, optional
+        Line width of the node borders in pixels.  If int or float, use the same width
+        for all node borders.  If a dict, must contain (node_id: width) pairs.  If
+        iterable, assume the widths are specified in the same order as the nodes are
+        found in H.nodes. If NodeStat, use a monotonic linear interpolation defined
+        between min_node_lw and max_node_lw. By default, 1.
+    node_size : int, float, dict, iterable, or NodeStat, optional
+        Radius of the nodes in pixels.  If int or float, use the same radius for all
+        nodes.  If a dict, must contain (node_id: radius) pairs.  If iterable, assume
+        the radiuses are specified in the same order as the nodes are found in
+        H.nodes. If NodeStat, use a monotonic linear interpolation defined between
+        min_node_size and max_node_size. By default, 5.
+    plane_color : color (str or tuple) or iterable (dict, list, or numpy array), optional
+        Color of each plane. If a dict, must contain (edge size: color) pairs.
+        By default, "grey".
+    max_order : int, optional
+        Maximum of hyperedges to plot. If None (default), plots all orders.
     conn_lines : bool, optional
         Whether to draw connections between layers. Default is True.
     conn_lines_style : str, optional
@@ -1470,91 +1248,163 @@ def draw_multilayer(
         The rotation angle around the vertical axis in degrees. Default is 0.
     sep : float, optional
         The separation between layers. Default is 1.
+    **kwargs : optional args
+        Alternate default values. Values that can be overwritten are the following:
+        * min_node_size
+        * max_node_size
+        * min_node_lw
+        * max_node_lw
+        * min_dyad_lw
+        * max_dyad_lw
+        * node_fc_cmap
+        * node_ec_cmap
+        * dyad_color_cmap
+        * edge_fc_cmap
 
     Returns
     -------
     ax : matplotlib Axes3DSubplot
         The subplot with the multilayer network visualization.
     """
+    settings = {
+        "min_node_size": 10.0,
+        "max_node_size": 30.0,
+        "min_dyad_lw": 2.0,
+        "max_dyad_lw": 10.0,
+        "min_node_lw": 1.0,
+        "max_node_lw": 5.0,
+        "node_fc_cmap": cm.Reds,
+        "node_ec_cmap": cm.Greys,
+        "edge_fc_cmap": cm.Blues,
+        "dyad_color_cmap": cm.Greys,
+        "plane_color_cmap": cm.Greys,
+    }
+
+    settings.update(kwargs)
+
+    if edge_fc is None:
+        edge_fc = H.edges.size
 
     if pos is None:
         pos = barycenter_spring_layout(H)
 
     if ax is None:
-        fig, ax = plt.subplots(
+        _, ax = plt.subplots(
             1, 1, figsize=(width, height), dpi=600, subplot_kw={"projection": "3d"}
         )
 
+    s = unique_edge_sizes(H)
     if max_order is None:
-        max_order = max_edge_order(H)
-
-    cmap = cm.get_cmap(palette, max_order)
+        max_order = max(s) - 1
+    else:
+        max_order = min(max_order, max(s) - 1)
+    min_order = min(s) - 1
 
     xs, ys = zip(*pos.values())
 
-    for order in range(1, max_order + 1):
-        zs = [order * sep] * len(xs)
-
-        # draw lines connecting points on the different planes
-        if conn_lines and order > 1:
-            thru_nodes = H.nodes
-            lines3d_between = [
-                (list(pos[i]) + [order * sep - sep], list(pos[i]) + [order * sep])
-                for i in thru_nodes
-            ]
-            between_lines = Line3DCollection(
-                lines3d_between,
-                zorder=order,
-                color=".5",
-                alpha=0.4,
-                linestyle=conn_lines_style,
-                linewidth=1,
-            )
-            ax.add_collection3d(between_lines)
+    dyad_color = _color_arg_to_dict(dyad_color, H.edges, settings["dyad_color_cmap"])
+    dyad_lw = _scalar_arg_to_dict(
+        dyad_lw, H.edges, settings["min_dyad_lw"], settings["max_dyad_lw"]
+    )
+
+    edge_fc = _color_arg_to_dict(edge_fc, H.edges, settings["edge_fc_cmap"])
+
+    node_fc = _color_arg_to_dict(node_fc, H.nodes, settings["node_fc_cmap"])
+    node_ec = _color_arg_to_dict(node_ec, H.nodes, settings["node_ec_cmap"])
+    node_lw = _scalar_arg_to_dict(
+        node_lw,
+        H.nodes,
+        settings["min_node_lw"],
+        settings["max_node_lw"],
+    )
+    node_size = _scalar_arg_to_dict(
+        node_size, H.nodes, settings["min_node_size"], settings["max_node_size"]
+    )
+
+    plane_color = _color_arg_to_dict(
+        plane_color,
+        [i for i in range(min_order, max_order + 1)],
+        settings["plane_color_cmap"],
+    )
+
+    for id, he in H.edges.members(dtype=dict).items():
+        d = len(he) - 1
+        zs = d * sep
 
-        # draw the edges/simplices of given order
-        edges = H.edges.filterby("order", order).members()
         # dyads
-        if order == 1:
-            lines3d = [
-                (list(pos[i]) + [order * sep], list(pos[j]) + [order * sep])
-                for i, j in edges
-            ]
-            line_collection = Line3DCollection(
-                lines3d,
-                zorder=order - 1,
-                color=cmap(order - 1),
-                alpha=1,
-                linewidth=0.7,
+        if d > max_order:
+            continue
+
+        if d == 1:
+            he = list(he)
+            x1 = [pos[he[0]][0], pos[he[0]][1], zs]
+            x2 = [pos[he[1]][0], pos[he[1]][1], zs]
+            l = Line3DCollection(
+                [(x1, x2)],
+                color=dyad_color[id],
+                linewidth=dyad_lw[id],
             )
-            ax.add_collection3d(line_collection)
+            ax.add_collection3d(l)
         # higher-orders
         else:
             poly = []
-            for e in edges:
-                vertices = np.array([[xs[i - 1], ys[i - 1], zs[i - 1]] for i in e])
-                vertices = _CCW_sort(vertices)
-                poly.append(vertices)
+            vertices = np.array([[pos[i][0], pos[i][1], zs] for i in he])
+            vertices = _CCW_sort(vertices)
+            poly.append(vertices)
             poly = Poly3DCollection(
                 poly,
-                zorder=order - 1,
-                color=cmap(order - 1),
+                zorder=d - 1,
+                color=edge_fc[id],
                 alpha=0.5,
+                edgecolor=None,
             )
             ax.add_collection3d(poly)
+
+    # now draw by order
+    # draw lines connecting points on the different planes
+    if conn_lines:
+        lines3d_between = [
+            (list(pos[i]) + [min_order * sep], list(pos[i]) + [max_order * sep])
+            for i in H.nodes
+        ]
+        between_lines = Line3DCollection(
+            lines3d_between,
+            zorder=d,
+            color=".5",
+            alpha=0.4,
+            linestyle=conn_lines_style,
+            linewidth=1,
+        )
+        ax.add_collection3d(between_lines)
+
+    (x, y, s, c, ec, lw,) = zip(
+        *[
+            (
+                pos[i][0],
+                pos[i][1],
+                node_size[i] ** 2,
+                node_fc[i],
+                node_ec[i],
+                node_lw[i],
+            )
+            for i in H.nodes
+        ]
+    )
+    for d in range(min_order, max_order + 1):
         # draw nodes
+        z = [sep * d] * H.num_nodes
         ax.scatter(
-            xs,
-            ys,
-            zs,
-            c=node_fc,
-            edgecolors=node_ec,
-            linewidths=node_lw,
-            marker="o",
+            x,
+            y,
+            z,
+            s=s,
+            c=c,
+            edgecolors=ec,
+            linewidths=lw,
+            zorder=max_order + 1,
             alpha=1,
-            zorder=order + 1,
-            s=node_size,
         )
+
         # draw surfaces corresponding to the different orders
         xdiff = np.max(xs) - np.min(xs)
         ydiff = np.max(ys) - np.min(ys)
@@ -1563,14 +1413,14 @@ def draw_multilayer(
         xmin = np.min(xs) - xdiff * 0.1 * (width / height)
         xmax = np.max(xs) + xdiff * 0.1 * (width / height)
         xx, yy = np.meshgrid([xmin, xmax], [ymin, ymax])
-        zz = np.zeros(xx.shape) + order * sep
+        zz = np.zeros(xx.shape) + d * sep
         ax.plot_surface(
             xx,
             yy,
             zz,
-            color=cmap(order - 1),
+            color=plane_color[d],
             alpha=0.1,
-            zorder=order,
+            zorder=d,
         )
 
     ax.view_init(h_angle, v_angle)
diff --git a/xgi/drawing/draw_utils.py b/xgi/drawing/draw_utils.py
new file mode 100644
index 000000000..dbc051580
--- /dev/null
+++ b/xgi/drawing/draw_utils.py
@@ -0,0 +1,272 @@
+"""Draw hypergraphs and simplicial complexes with matplotlib."""
+
+from collections.abc import Iterable
+
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.colors import LinearSegmentedColormap, ListedColormap
+from numpy import ndarray
+
+from ..exception import XGIError
+from ..stats import IDStat
+from .layout import barycenter_spring_layout
+
+
+def _draw_init(H, ax, pos):
+    """Initializes the axis and node positions
+
+    Parameters
+    ----------
+    H : Hypergraph
+        The hypergraph to draw
+    ax : Matplotlib axis
+        The axis on which to plot
+    pos : dict of lists
+        The output of the layout functions.
+        If None, uses the barycenter spring layout.
+
+    Returns
+    -------
+    ax, pos
+        ax: axis updated with the XGI format
+        pos: dict of lists with the node positions
+    """
+
+    if pos is None:
+        pos = barycenter_spring_layout(H)
+
+    if ax is None:
+        ax = plt.gca()
+
+    ax.get_xaxis().set_ticks([])
+    ax.get_yaxis().set_ticks([])
+    ax.axis("off")
+
+    return ax, pos
+
+
+def _update_lims(pos, ax):
+    """Update Axis limits based on node positions"""
+
+    # compute axis limits
+    pos_arr = np.asarray([[x, y] for n, (x, y) in pos.items()])
+
+    maxx, maxy = np.max(pos_arr, axis=0)
+    minx, miny = np.min(pos_arr, axis=0)
+    w = maxx - minx
+    h = maxy - miny
+
+    # update view after drawing
+    padx, pady = 0.05 * w, 0.05 * h
+    corners = (minx - padx, miny - pady), (maxx + padx, maxy + pady)
+    ax.update_datalim(corners)
+    ax.autoscale_view()
+
+
+def _scalar_arg_to_dict(scalar_arg, ids, min_val, max_val):
+    """Map different types of arguments for drawing style to a dict with scalar values.
+
+    Parameters
+    ----------
+    scalar_arg : int, float, dict, iterable, or NodeStat/EdgeStat
+        Attributes for drawing parameter.
+    ids : NodeView or EdgeView
+        This is the node or edge IDs that attributes get mapped to.
+    min_val : int or float
+        The minimum value of the drawing parameter.
+    max_val : int or float
+        The maximum value of the drawing parameter.
+
+    Returns
+    -------
+    dict
+        An ID: scalar dictionary.
+
+    Raises
+    ------
+    TypeError
+        If a int, float, list, dict, or NodeStat/EdgeStat is not passed.
+    """
+    if isinstance(scalar_arg, str):
+        raise TypeError(
+            "Argument must be int, float, dict, iterable, "
+            f"or NodeStat/EdgeStat. Received {type(scalar_arg)}"
+        )
+
+    # Single argument
+    if isinstance(scalar_arg, (int, float)):
+        return {id: scalar_arg for id in ids}
+
+    # IDStat
+    if isinstance(scalar_arg, IDStat):
+        vals = np.interp(
+            scalar_arg.asnumpy(),
+            [scalar_arg.min(), scalar_arg.max()],
+            [min_val, max_val],
+        )
+        return dict(zip(ids, vals))
+
+    # Iterables of floats or ints
+    if isinstance(scalar_arg, Iterable):
+        if isinstance(scalar_arg, dict):
+            try:
+                return {id: float(scalar_arg[id]) for id in scalar_arg if id in ids}
+            except ValueError as e:
+                raise TypeError(
+                    "The input dict must have values that can be cast to floats."
+                )
+
+        elif isinstance(scalar_arg, (list, ndarray)):
+            try:
+                return {id: float(scalar_arg[idx]) for idx, id in enumerate(ids)}
+            except ValueError as e:
+                raise TypeError(
+                    "The input list or array must have values that can be cast to floats."
+                )
+        else:
+            raise TypeError(
+                "Argument must be an dict, list, or numpy array of floats or ints."
+            )
+
+    raise TypeError(
+        "Argument must be int, float, dict, iterable, "
+        f"or NodeStat/EdgeStat. Received {type(scalar_arg)}"
+    )
+
+
+def _color_arg_to_dict(color_arg, ids, cmap):
+    """Map different types of arguments for drawing style to a dict with color values.
+
+    Parameters
+    ----------
+    color_arg : Several formats are accepted:
+
+        Single color values
+
+        * str
+        * 3- or 4-tuple
+
+        Iterable of colors (each color specified as above)
+
+        * numpy array
+        * list
+        * dict {id: color} pairs
+
+        Iterable of numerical values (floats or ints)
+
+        * list
+        * dict
+        * numpy array
+
+        Stats
+
+        * NodeStat
+        * EdgeStat
+
+        Attributes for drawing parameter.
+    ids : NodeView or EdgeView
+        This is the node or edge IDs that attributes get mapped to.
+    cmap : ListedColormap or LinearSegmentedColormap
+        colormap to use for NodeStat/EdgeStat.
+
+    Returns
+    -------
+    dict
+        An ID: color dictionary.
+
+    Raises
+    ------
+    TypeError
+        If a string, dict, iterable, or NodeStat/EdgeStat is not passed.
+
+    Notes
+    -----
+    For the iterable of values, we do not accept tuples,
+    because there is the potential for ambiguity.
+    """
+
+    # single argument. Must be a string or a tuple of floats
+    if isinstance(color_arg, str) or (
+        isinstance(color_arg, tuple) and isinstance(color_arg[0], float)
+    ):
+        return {id: color_arg for id in ids}
+
+    # Iterables of colors. The values of these iterables must strings or tuples. As of now,
+    # there is not a check to verify that the tuples contain floats.
+    if isinstance(color_arg, Iterable):
+        if isinstance(color_arg, dict) and isinstance(
+            next(iter(color_arg.values())), (str, tuple, ndarray)
+        ):
+            return {id: color_arg[id] for id in color_arg if id in ids}
+        if isinstance(color_arg, (list, ndarray)) and isinstance(
+            color_arg[0], (str, tuple, ndarray)
+        ):
+            return {id: color_arg[idx] for idx, id in enumerate(ids)}
+
+    # Stats or iterable of values
+    if isinstance(color_arg, (Iterable, IDStat)):
+        # set max and min of interpolation based on color map
+        if isinstance(cmap, ListedColormap):
+            minval = 0
+            maxval = cmap.N
+        elif isinstance(cmap, LinearSegmentedColormap):
+            minval = 0.1
+            maxval = 0.9
+        else:
+            raise XGIError("Invalid colormap!")
+
+        # handle the case of IDStat vs iterables
+        if isinstance(color_arg, IDStat):
+            vals = np.interp(
+                color_arg.asnumpy(),
+                [color_arg.min(), color_arg.max()],
+                [minval, maxval],
+            )
+            return {
+                id: np.array(cmap(vals[i])).reshape(1, -1) for i, id in enumerate(ids)
+            }
+
+        elif isinstance(color_arg, Iterable):
+            if isinstance(color_arg, dict) and isinstance(
+                next(iter(color_arg.values())), (int, float)
+            ):
+                v = list(color_arg.values())
+                vals = np.interp(v, [np.min(v), np.max(v)], [minval, maxval])
+                # because we have ids, we can't just assume that the keys of arg correspond to
+                # the ids.
+                return {
+                    id: np.array(cmap(v)).reshape(1, -1)
+                    for v, id in zip(vals, color_arg.keys())
+                    if id in ids
+                }
+
+            if isinstance(color_arg, (list, ndarray)) and isinstance(
+                color_arg[0], (int, float)
+            ):
+                vals = np.interp(
+                    color_arg, [np.min(color_arg), np.max(color_arg)], [minval, maxval]
+                )
+                return {
+                    id: np.array(cmap(vals[i])).reshape(1, -1)
+                    for i, id in enumerate(ids)
+                }
+            else:
+                raise TypeError(
+                    "Argument must be an dict, list, or numpy array of floats."
+                )
+
+    raise TypeError(
+        "Argument must be str, dict, iterable, or "
+        f"NodeStat/EdgeStat. Received {type(color_arg)}"
+    )
+
+
+def _CCW_sort(p):
+    """
+    Sort the input 2D points counterclockwise.
+    """
+    p = np.array(p)
+    mean = np.mean(p, axis=0)
+    d = p - mean
+    s = np.arctan2(d[:, 0], d[:, 1])
+    return p[np.argsort(s), :]
diff --git a/xgi/drawing/layout.py b/xgi/drawing/layout.py
index d6790800c..9d3033dd3 100644
--- a/xgi/drawing/layout.py
+++ b/xgi/drawing/layout.py
@@ -22,7 +22,7 @@
 ]
 
 
-def random_layout(H, center=None, dim=2, seed=None):
+def random_layout(H, center=None, seed=None):
     """Position nodes uniformly at random in the unit square.
 
     For every node, a position is generated by choosing each of dim coordinates
@@ -36,8 +36,6 @@ def random_layout(H, center=None, dim=2, seed=None):
     center : array-like, optional
         Coordinate pair around which to center the layout.
         If None (default), does not center the positions.
-    dim : int, optional
-        Dimension of layout, by default 2.
     seed : int, optional
         Set the random state for deterministic node layouts.
         If int, `seed` is the seed used by the random number generator,
@@ -76,15 +74,15 @@ def random_layout(H, center=None, dim=2, seed=None):
     if seed is not None:
         np.random.seed(seed)
 
-    H, center = nx.drawing.layout._process_params(H, center, dim)
-    pos = np.random.rand(len(H), dim) + center
+    H, center = nx.drawing.layout._process_params(H, center, 2)
+    pos = np.random.rand(len(H), 2) + center
     pos = pos.astype(np.float32)
     pos = dict(zip(H, pos))
 
     return pos
 
 
-def pairwise_spring_layout(H, seed=None):
+def pairwise_spring_layout(H, seed=None, k=None, **kwargs):
     """
     Position the nodes using Fruchterman-Reingold force-directed
     algorithm using the graph projection of the hypergraph
@@ -99,6 +97,13 @@ def pairwise_spring_layout(H, seed=None):
         If int, `seed` is the seed used by the random number generator,
         If None (default), random numbers are sampled from the
         numpy random number generator without initialization.
+    k : float
+        The spring constant of the links. When k=None (default),
+        k = 1/sqrt(N). For more information, see the documentation
+        for the NetworkX spring_layout() function.
+    kwargs :
+        Optional arguments for the NetworkX spring_layout() function.
+        See https://networkx.org/documentation/stable/reference/generated/networkx.drawing.layout.spring_layout.html
 
     Returns
     -------
@@ -131,7 +136,7 @@ def pairwise_spring_layout(H, seed=None):
     if isinstance(H, SimplicialComplex):
         H = convert.from_max_simplices(H)
     G = convert.convert_to_graph(H)
-    pos = nx.spring_layout(G, seed=seed)
+    pos = nx.spring_layout(G, seed=seed, k=k, **kwargs)
     return pos
 
 
@@ -190,7 +195,9 @@ def _augmented_projection(H, weighted=False):
     return G
 
 
-def barycenter_spring_layout(H, return_phantom_graph=False, seed=None):
+def barycenter_spring_layout(
+    H, return_phantom_graph=False, seed=None, k=None, **kwargs
+):
     """
     Position the nodes using Fruchterman-Reingold force-directed
     algorithm using an augmented version of the the graph projection
@@ -211,6 +218,14 @@ def barycenter_spring_layout(H, return_phantom_graph=False, seed=None):
         If int, `seed` is the seed used by the random number generator,
         If None (default), random numbers are sampled from the
         numpy random number generator without initialization.
+    k : float
+        The spring constant of the links. When k=None (default),
+        k = 1/sqrt(N). For more information, see the documentation
+        for the NetworkX spring_layout() function.
+    kwargs :
+        Optional arguments for the NetworkX spring_layout() function.
+        See https://networkx.org/documentation/stable/reference/generated/networkx.drawing.layout.spring_layout.html
+
 
     Returns
     -------
@@ -241,7 +256,7 @@ def barycenter_spring_layout(H, return_phantom_graph=False, seed=None):
 
     # Creating a dictionary for the position of the nodes with the standard spring
     # layout
-    pos_with_phantom_nodes = nx.spring_layout(G, seed=seed)
+    pos_with_phantom_nodes = nx.spring_layout(G, seed=seed, k=k, **kwargs)
 
     # Retaining only the positions of the real nodes
     pos = {k: pos_with_phantom_nodes[k] for k in list(H.nodes)}
@@ -252,7 +267,9 @@ def barycenter_spring_layout(H, return_phantom_graph=False, seed=None):
         return pos
 
 
-def weighted_barycenter_spring_layout(H, return_phantom_graph=False, seed=None):
+def weighted_barycenter_spring_layout(
+    H, return_phantom_graph=False, seed=None, k=None, **kwargs
+):
     """Position the nodes using Fruchterman-Reingold force-directed algorithm.
 
     This uses an augmented version of the the graph projection of the hypergraph (or
@@ -275,6 +292,14 @@ def weighted_barycenter_spring_layout(H, return_phantom_graph=False, seed=None):
         If int, `seed` is the seed used by the random number generator,
         If None (default), random numbers are sampled from the
         numpy random number generator without initialization.
+    k : float
+        The spring constant of the links. When k=None (default),
+        k = 1/sqrt(N). For more information, see the documentation
+        for the NetworkX spring_layout() function.
+    kwargs :
+        Optional arguments for the NetworkX spring_layout() function.
+        See https://networkx.org/documentation/stable/reference/generated/networkx.drawing.layout.spring_layout.html
+
 
     Returns
     -------
@@ -305,7 +330,9 @@ def weighted_barycenter_spring_layout(H, return_phantom_graph=False, seed=None):
     G = _augmented_projection(H, weighted=True)
 
     # Creating a dictionary for node position with the standard spring layout
-    pos_with_phantom_nodes = nx.spring_layout(G, weight="weight", seed=seed)
+    pos_with_phantom_nodes = nx.spring_layout(
+        G, weight="weight", seed=seed, k=k, **kwargs
+    )
 
     # Retaining only the positions of the real nodes
     pos = {k: pos_with_phantom_nodes[k] for k in list(H.nodes)}
@@ -452,7 +479,7 @@ def spiral_layout(H, center=None, resolution=0.35, equidistant=False):
     return pos
 
 
-def barycenter_kamada_kawai_layout(H, return_phantom_graph=False):
+def barycenter_kamada_kawai_layout(H, return_phantom_graph=False, **kwargs):
     """Position nodes using Kamada-Kawai path-length cost-function
     using an augmented version of the the graph projection
     of the hypergraph (or simplicial complex), where phantom nodes
@@ -467,6 +494,9 @@ def barycenter_kamada_kawai_layout(H, return_phantom_graph=False):
     return_phantom_graph: bool (default=False)
         If True the function returns also the augmented version of the
         the graph projection of the hypergraph (or simplicial complex).
+    kwargs :
+        Optional arguments for the NetworkX spring_layout() function.
+        See https://networkx.org/documentation/stable/reference/generated/networkx.drawing.layout.kamada_kawai_layout.html
 
     Returns
     -------
@@ -479,7 +509,7 @@ def barycenter_kamada_kawai_layout(H, return_phantom_graph=False):
     G = _augmented_projection(H)
 
     # Creating a dictionary for the position of the nodes with the standard spring layout
-    pos_with_phantom_nodes = nx.kamada_kawai_layout(G)
+    pos_with_phantom_nodes = nx.kamada_kawai_layout(G, **kwargs)
 
     # Retaining only the positions of the real nodes
     pos = {k: pos_with_phantom_nodes[k] for k in list(H.nodes)}


Test patch:
diff --git a/tests/drawing/test_draw.py b/tests/drawing/test_draw.py
index aa8ec3a1e..e09f69413 100644
--- a/tests/drawing/test_draw.py
+++ b/tests/drawing/test_draw.py
@@ -1,149 +1,10 @@
 import matplotlib.pyplot as plt
-import numpy as np
 import pytest
-from matplotlib import cm
 
 import xgi
-from xgi.drawing.draw import _CCW_sort, _color_arg_to_dict, _scalar_arg_to_dict
 from xgi.exception import XGIError
 
 
-def test_CCW_sort():
-    coords = [[0.919, 0.145], [0.037, 0.537], [0.402, 0.56]]
-    sorted_coords = _CCW_sort(coords)
-    assert np.all(
-        sorted_coords == np.array([[0.037, 0.537], [0.402, 0.56], [0.919, 0.145]])
-    )
-
-    coords = [[0.037, 0.537], [0.402, 0.56], [0.791, 0.91], [0.0, 0.868]]
-    sorted_coords = _CCW_sort(coords)
-    assert np.all(
-        sorted_coords
-        == np.array([[0.037, 0.537], [0.0, 0.868], [0.791, 0.91], [0.402, 0.56]])
-    )
-
-
-def test_scalar_arg_to_dict(edgelist4):
-    ids = [1, 2, 3]
-    min_val = 1
-    max_val = 5
-
-    arg = 1
-    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-    assert d == {1: 1, 2: 1, 3: 1}
-
-    arg = 0.3
-    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-    assert d == {1: 0.3, 2: 0.3, 3: 0.3}
-
-    arg = [0.2, 3, 4]
-    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-    assert d == {1: 0.2, 2: 3, 3: 4}
-
-    arg = np.array([0.2, 3, 4])
-    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-    assert d == {1: 0.2, 2: 3, 3: 4}
-
-    arg = {1: 0.2, 2: 3, 3: 4}
-    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-    assert d == {1: 0.2, 2: 3, 3: 4}
-
-    H = xgi.Hypergraph(edgelist4)
-    arg = H.nodes.degree
-    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-    assert d == {1: 1.0, 2: 3.0, 3: 5.0}
-
-    with pytest.raises(TypeError):
-        arg = "2"
-        d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-
-    with pytest.raises(TypeError):
-        arg = (1, 2, 3)
-        d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
-
-
-def test_color_arg_to_dict(edgelist4):
-    ids = [1, 2, 3]
-
-    # single values
-    arg1 = "black"
-    arg2 = (0.1, 0.2, 0.3)
-    arg3 = (0.1, 0.2, 0.3, 0.5)
-
-    # test iterables of colors
-    arg4 = [(0.1, 0.2, 0.3), (0.1, 0.2, 0.4), (0.1, 0.2, 0.5)]
-    arg5 = ["blue", "black", "red"]
-    arg6 = np.array(["blue", "black", "red"])
-    arg7 = {0: (0.1, 0.2, 0.3), 1: (0.1, 0.2, 0.4), 2: (0.1, 0.2, 0.5)}
-    arg8 = {0: "blue", 1: "black", 2: "red"}
-
-    # test iterables of values
-    arg9 = [0, 0.1, 0.2]
-    arg10 = {1: 0, 2: 0.1, 3: 0.2}
-    arg11 = np.array([0, 0.1, 0.2])
-
-    # test single values
-    d = _color_arg_to_dict(arg1, ids, None)
-    assert d == {1: "black", 2: "black", 3: "black"}
-
-    d = _color_arg_to_dict(arg2, ids, None)
-    assert d == {1: (0.1, 0.2, 0.3), 2: (0.1, 0.2, 0.3), 3: (0.1, 0.2, 0.3)}
-
-    d = _color_arg_to_dict(arg3, ids, None)
-    for i in d:
-        assert np.allclose(d[i], np.array([0.1, 0.2, 0.3, 0.5]))
-
-    # Test iterables of colors
-    d = _color_arg_to_dict(arg4, ids, None)
-    assert d == {1: (0.1, 0.2, 0.3), 2: (0.1, 0.2, 0.4), 3: (0.1, 0.2, 0.5)}
-
-    d = _color_arg_to_dict(arg5, ids, None)
-    assert d == {1: "blue", 2: "black", 3: "red"}
-
-    d = _color_arg_to_dict(arg6, ids, None)
-    assert d == {1: "blue", 2: "black", 3: "red"}
-
-    d = _color_arg_to_dict(arg7, ids, None)
-    assert d == {1: (0.1, 0.2, 0.4), 2: (0.1, 0.2, 0.5)}
-
-    d = _color_arg_to_dict(arg8, ids, None)
-    assert d == {1: "black", 2: "red"}
-
-    # Test iterables of values
-    cdict = {
-        1: np.array([[0.89173395, 0.93510188, 0.97539408, 1.0]]),
-        2: np.array([[0.41708574, 0.68063053, 0.83823145, 1.0]]),
-        3: np.array([[0.03137255, 0.28973472, 0.57031911, 1.0]]),
-    }
-    d = _color_arg_to_dict(arg9, ids, cm.Blues)
-    for i in d:
-        assert np.allclose(d[i], cdict[i])
-
-    d = _color_arg_to_dict(arg10, ids, cm.Blues)
-    for i in d:
-        assert np.allclose(d[i], cdict[i])
-
-    d = _color_arg_to_dict(arg11, ids, cm.Blues)
-    for i in d:
-        assert np.allclose(d[i], cdict[i])
-
-    H = xgi.Hypergraph(edgelist4)
-    arg = H.nodes.degree
-    d = _color_arg_to_dict(arg, ids, cm.Reds)
-    assert np.allclose(d[1], np.array([[0.99692426, 0.89619377, 0.84890427, 1.0]]))
-    assert np.allclose(d[2], np.array([[0.98357555, 0.41279508, 0.28835063, 1.0]]))
-    assert np.allclose(d[3], np.array([[0.59461745, 0.0461361, 0.07558631, 1.0]]))
-
-    # Test bad calls
-    with pytest.raises(TypeError):
-        arg = 0.3
-        d = _color_arg_to_dict(arg, ids, None)
-
-    with pytest.raises(TypeError):
-        arg = 1
-        d = _color_arg_to_dict(arg, ids, None)
-
-
 def test_draw(edgelist8):
     H = xgi.Hypergraph(edgelist8)
 
@@ -244,18 +105,48 @@ def test_correct_number_of_collections_draw_multilayer(edgelist8):
     # hypergraph
     H = xgi.Hypergraph(edgelist8)
     ax1 = xgi.draw_multilayer(H)
-    assert xgi.max_edge_order(H) * 4 - 1 == len(ax1.collections)
+    sizes = xgi.unique_edge_sizes(H)
+    num_planes = max(sizes) - min(sizes) + 1
+    num_node_collections = max(sizes) - min(sizes) + 1
+    num_edge_collections = H.num_edges
+    num_thru_lines_collections = 1
+
+    assert (
+        num_planes
+        + num_node_collections
+        + num_edge_collections
+        + num_thru_lines_collections
+        == len(ax1.collections)
+    )
     plt.close()
 
     # max_order parameter
     max_order = 2
     ax2 = xgi.draw_multilayer(H, max_order=max_order)
-    assert max_order * 4 - 1 == len(ax2.collections)
+    sizes = [2, 3]
+    num_planes = max(sizes) - min(sizes) + 1
+    num_node_collections = max(sizes) - min(sizes) + 1
+    num_edge_collections = len(H.edges.filterby("size", [2, 3], "between"))
+    num_thru_lines_collections = 1
+
+    assert (
+        num_planes
+        + num_node_collections
+        + num_edge_collections
+        + num_thru_lines_collections
+        == len(ax2.collections)
+    )
     plt.close()
 
     # conn_lines parameter
     ax3 = xgi.draw_multilayer(H, conn_lines=False)
-    assert xgi.max_edge_order(H) * 3 == len(ax3.collections)
+    sizes = xgi.unique_edge_sizes(H)
+    num_planes = max(sizes) - min(sizes) + 1
+    num_node_collections = max(sizes) - min(sizes) + 1
+    num_edge_collections = H.num_edges
+    assert num_planes + num_node_collections + num_edge_collections == len(
+        ax3.collections
+    )
     plt.close()
 
     # custom parameters
@@ -274,7 +165,19 @@ def test_correct_number_of_collections_draw_multilayer(edgelist8):
         v_angle=15,
         sep=2,
     )
-    assert xgi.max_edge_order(H) * 4 - 1 == len(ax4.collections)
+    sizes = xgi.unique_edge_sizes(H)
+    num_planes = max(sizes) - min(sizes) + 1
+    num_node_collections = max(sizes) - min(sizes) + 1
+    num_edge_collections = H.num_edges
+    num_thru_lines_collections = 1
+
+    assert (
+        num_planes
+        + num_node_collections
+        + num_edge_collections
+        + num_thru_lines_collections
+        == len(ax4.collections)
+    )
     plt.close()
 
 
diff --git a/tests/drawing/test_draw_utils.py b/tests/drawing/test_draw_utils.py
new file mode 100644
index 000000000..9625069c7
--- /dev/null
+++ b/tests/drawing/test_draw_utils.py
@@ -0,0 +1,142 @@
+import numpy as np
+import pytest
+from matplotlib import cm
+
+import xgi
+from xgi.drawing.draw import _CCW_sort, _color_arg_to_dict, _scalar_arg_to_dict
+
+
+def test_CCW_sort():
+    coords = [[0.919, 0.145], [0.037, 0.537], [0.402, 0.56]]
+    sorted_coords = _CCW_sort(coords)
+    assert np.all(
+        sorted_coords == np.array([[0.037, 0.537], [0.402, 0.56], [0.919, 0.145]])
+    )
+
+    coords = [[0.037, 0.537], [0.402, 0.56], [0.791, 0.91], [0.0, 0.868]]
+    sorted_coords = _CCW_sort(coords)
+    assert np.all(
+        sorted_coords
+        == np.array([[0.037, 0.537], [0.0, 0.868], [0.791, 0.91], [0.402, 0.56]])
+    )
+
+
+def test_scalar_arg_to_dict(edgelist4):
+    ids = [1, 2, 3]
+    min_val = 1
+    max_val = 5
+
+    arg = 1
+    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+    assert d == {1: 1, 2: 1, 3: 1}
+
+    arg = 0.3
+    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+    assert d == {1: 0.3, 2: 0.3, 3: 0.3}
+
+    arg = [0.2, 3, 4]
+    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+    assert d == {1: 0.2, 2: 3, 3: 4}
+
+    arg = np.array([0.2, 3, 4])
+    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+    assert d == {1: 0.2, 2: 3, 3: 4}
+
+    arg = {1: 0.2, 2: 3, 3: 4}
+    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+    assert d == {1: 0.2, 2: 3, 3: 4}
+
+    H = xgi.Hypergraph(edgelist4)
+    arg = H.nodes.degree
+    d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+    assert d == {1: 1.0, 2: 3.0, 3: 5.0}
+
+    with pytest.raises(TypeError):
+        arg = "2"
+        d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+
+    with pytest.raises(TypeError):
+        arg = (1, 2, 3)
+        d = _scalar_arg_to_dict(arg, ids, min_val, max_val)
+
+
+def test_color_arg_to_dict(edgelist4):
+    ids = [1, 2, 3]
+
+    # single values
+    arg1 = "black"
+    arg2 = (0.1, 0.2, 0.3)
+    arg3 = (0.1, 0.2, 0.3, 0.5)
+
+    # test iterables of colors
+    arg4 = [(0.1, 0.2, 0.3), (0.1, 0.2, 0.4), (0.1, 0.2, 0.5)]
+    arg5 = ["blue", "black", "red"]
+    arg6 = np.array(["blue", "black", "red"])
+    arg7 = {0: (0.1, 0.2, 0.3), 1: (0.1, 0.2, 0.4), 2: (0.1, 0.2, 0.5)}
+    arg8 = {0: "blue", 1: "black", 2: "red"}
+
+    # test iterables of values
+    arg9 = [0, 0.1, 0.2]
+    arg10 = {1: 0, 2: 0.1, 3: 0.2}
+    arg11 = np.array([0, 0.1, 0.2])
+
+    # test single values
+    d = _color_arg_to_dict(arg1, ids, None)
+    assert d == {1: "black", 2: "black", 3: "black"}
+
+    d = _color_arg_to_dict(arg2, ids, None)
+    assert d == {1: (0.1, 0.2, 0.3), 2: (0.1, 0.2, 0.3), 3: (0.1, 0.2, 0.3)}
+
+    d = _color_arg_to_dict(arg3, ids, None)
+    for i in d:
+        assert np.allclose(d[i], np.array([0.1, 0.2, 0.3, 0.5]))
+
+    # Test iterables of colors
+    d = _color_arg_to_dict(arg4, ids, None)
+    assert d == {1: (0.1, 0.2, 0.3), 2: (0.1, 0.2, 0.4), 3: (0.1, 0.2, 0.5)}
+
+    d = _color_arg_to_dict(arg5, ids, None)
+    assert d == {1: "blue", 2: "black", 3: "red"}
+
+    d = _color_arg_to_dict(arg6, ids, None)
+    assert d == {1: "blue", 2: "black", 3: "red"}
+
+    d = _color_arg_to_dict(arg7, ids, None)
+    assert d == {1: (0.1, 0.2, 0.4), 2: (0.1, 0.2, 0.5)}
+
+    d = _color_arg_to_dict(arg8, ids, None)
+    assert d == {1: "black", 2: "red"}
+
+    # Test iterables of values
+    cdict = {
+        1: np.array([[0.89173395, 0.93510188, 0.97539408, 1.0]]),
+        2: np.array([[0.41708574, 0.68063053, 0.83823145, 1.0]]),
+        3: np.array([[0.03137255, 0.28973472, 0.57031911, 1.0]]),
+    }
+    d = _color_arg_to_dict(arg9, ids, cm.Blues)
+    for i in d:
+        assert np.allclose(d[i], cdict[i])
+
+    d = _color_arg_to_dict(arg10, ids, cm.Blues)
+    for i in d:
+        assert np.allclose(d[i], cdict[i])
+
+    d = _color_arg_to_dict(arg11, ids, cm.Blues)
+    for i in d:
+        assert np.allclose(d[i], cdict[i])
+
+    H = xgi.Hypergraph(edgelist4)
+    arg = H.nodes.degree
+    d = _color_arg_to_dict(arg, ids, cm.Reds)
+    assert np.allclose(d[1], np.array([[0.99692426, 0.89619377, 0.84890427, 1.0]]))
+    assert np.allclose(d[2], np.array([[0.98357555, 0.41279508, 0.28835063, 1.0]]))
+    assert np.allclose(d[3], np.array([[0.59461745, 0.0461361, 0.07558631, 1.0]]))
+
+    # Test bad calls
+    with pytest.raises(TypeError):
+        arg = 0.3
+        d = _color_arg_to_dict(arg, ids, None)
+
+    with pytest.raises(TypeError):
+        arg = 1
+        d = _color_arg_to_dict(arg, ids, None)