Feature/coloring: Four colorings for Zephyr and Pegasus, and a two coloring for Chimera

dwave_networkx/generators/chimera.py

@@ -35,6 +35,7 @@
            'linear_to_chimera',
            'chimera_sublattice_mappings',
            'chimera_torus',
+           'chimera_two_color',
            ]
 
 
@@ -866,3 +867,30 @@ def chimera_torus(m, n=None, t=None, node_list=None, edge_list=None):
     G.graph['boundary_condition'] = 'torus'
 
     return G
+
+
+def chimera_two_color(q):
+    """
+    Node color assignment sufficient for four coloring of a Chimera graph. 
+
+    Parameters
+    ----------
+        q : tuple
+            Qubit label in standard coordinate format: i, j, u, k
+
+    Returns
+    -------
+        color : int
+            Colors 0, 1, 2 or 3
+
+    Examples
+    ========
+    A mapping of every qubit (default integer labels) in the Chimera[m, t]
+    graph to one of 2 colors
+    >>> m = 2
+    >>> G = dnx.chimera_graph(m)
+    >>> colors = {q: dnx.chimera_two_color(dnx.chimera_coordinates(m,t).linear_to_chimera(q)) for q in G.nodes()}    # doctest: +SKIP
+
+    """
+    i, j, u, _ = q
+    return (i ^ j ^ u) & 1

dwave_networkx/generators/pegasus.py

@@ -30,6 +30,7 @@
            'pegasus_coordinates',
            'pegasus_sublattice_mappings',
            'pegasus_torus',
+           'pegasus_four_color',
            ]
 
 def pegasus_graph(m, create_using=None, node_list=None, edge_list=None, data=True,
@@ -538,7 +539,7 @@ def fragmented_edges(pegasus_graph):
             else:
                 yield ((fw1, fw0, u0, k0&1), (fw1, fw0, u1, k1&1))
 
-
+    
 # Developer note: we could implement a function that creates the iter_*_to_* and
 # iter_*_to_*_pairs methods just-in-time, but there are a small enough number
 # that for now it makes sense to do them by hand.
@@ -1303,3 +1304,28 @@ def relabel(u,w,k,z):
     G.graph['boundary_condition'] = 'torus'
 
     return G
+
+def pegasus_four_color(q):
+    """
+    Node color assignment sufficient for four coloring of a pegasus graph. 
+
+    Parameters
+    ----------
+        q : tuple
+            Qubit label in standard coordinate format.
+
+    Returns
+    -------
+        color : int
+            Colors 0, 1, 2 or 3
+    Examples
+    ========
+    A mapping of every qubit (default integer labels) in the Pegasus[m]
+    graph to one of 4 colors
+    >>> m = 2
+    >>> G = dnx.pegasus_graph(m)
+    >>> colors = {q: dnx.pegasus_four_color(dnx.pegasus_coordinates(m).linear_to_zephyr(q)) for q in G.nodes()}    # doctest: +SKIP
+
+    """
+    u, w, k, z = q
+    return 2 * u + ((k ^ z) & 1)

dwave_networkx/generators/zephyr.py

@@ -30,7 +30,8 @@
 __all__ = ['zephyr_graph',
            'zephyr_coordinates',
            'zephyr_sublattice_mappings',
-           'zephyr_torus'
+           'zephyr_torus',
+           'zephyr_four_color',
            ]
 
 def zephyr_graph(m, t=4, create_using=None, node_list=None, edge_list=None,
@@ -797,3 +798,35 @@ def relabel(u, w, k, j, z):
     G.graph['boundary_condition'] = 'torus'
 
     return G
+
+
+def zephyr_four_color(q, scheme=0):
+    """
+    Node color assignment sufficient for four coloring of a Zephyr graph. 
+
+    Parameters
+    ----------
+        q : tuple
+            Qubit label in standard coordinate format: u, w, k, j, z
+        scheme : int
+            Two patterns not related by automorphism are supported 
+    Returns
+    -------
+        color : int
+            Colors 0, 1, 2 or 3
+    Examples
+    ========
+    A mapping of every qubit (default integer labels) in the Zephyr[m, t]
+    graph to one of 4 colors
+    >>> m = 2
+    >>> G = dnx.zephyr_graph(m)
+    >>> colors = {q: dnx.zephyr_four_color(dnx.zephyr_coordinates(m,t).linear_to_zephyr(q)) for q in G.nodes()}    # doctest: +SKIP
+
+    """
+    u, w, _, j, z = q
+
+    if scheme == 0:
+        return j + ((w + 2*(z+u) + j)&2)
+    else:
+        return (2*u + w + 2*z + j) & 3
+

tests/test_regular_coloring.py

@@ -0,0 +1,69 @@
+# Copyright 2024 D-Wave Systems Inc.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import unittest
+
+import dwave_networkx as dnx
+import numpy as np
+
+class TestRegularColoring(unittest.TestCase):
+
+    def test_valid(self):
+        test_cases = {'chimera': [(3,3), (4,4)],
+                      'pegasus': [(2,), (4,)],
+                      'zephyr': [(1,2), (3,4)]}
+
+        for topology_type, topology_shapes in test_cases.items():
+            if topology_type == 'zephyr':
+                graph = dnx.zephyr_graph
+                color = dnx.zephyr_four_color
+                num_colors = 4
+            elif topology_type == 'pegasus':
+                graph = dnx.pegasus_graph
+                color = dnx.pegasus_four_color
+                num_colors = 4
+            elif topology_type == 'chimera':
+                graph = dnx.chimera_graph
+                color = dnx.chimera_two_color
+                num_colors = 2
+            else:
+                raise ValueError('unrecognized topology')
+
+            for topology_shape in topology_shapes:
+                G = graph(*topology_shape, coordinates=True)
+                col_dict = {q: color(q) for q in G.nodes}
+                self.assertTrue(np.all(np.unique(list(col_dict.values())) ==
+                                       np.arange(num_colors)))
+                self.assertTrue(all(col_dict[q1] != col_dict[q2]
+                                    for q1, q2 in G.edges),
+                                f'{topology_type}[{topology_shape}]')
+
+    def test_non_default_schemes(self):
+        graph = dnx.zephyr_graph
+        color = dnx.zephyr_four_color
+        num_colors = 4
+        topology_type = 'zephyr'
+        topology_shapes =  [(2,2), (3,1)]
+        for topology_shape in topology_shapes:
+            G = graph(*topology_shape, coordinates=True)
+            col_dicts = {}
+            for scheme in [0,1]:
+                col_dict = {q: color(q, scheme=scheme) for q in G.nodes}
+                self.assertTrue(np.all(np.unique(list(col_dict.values())) ==
+                                       np.arange(num_colors)))
+                self.assertTrue(all(col_dict[q1] != col_dict[q2]
+                                    for q1, q2 in G.edges),
+                                f'{topology_type}[{topology_shape}]')
+                col_dicts[scheme] = col_dict
+            self.assertFalse(all(col_dicts[0][q] == col_dicts[1][q] for q in G.nodes))