Add Zephyr symbol

Author: arcondello

dwave/optimization/include/dwave-optimization/nodes/quadratic_model.hpp

@@ -168,6 +168,16 @@ class LatticeNode : public ScalarOutputMixin<ArrayNode> {
     /// Return a reference to the lattice structure of the node.
     const LatticeType& lattice() const noexcept { return lattice_; }
 
+    template<class... Args>
+    static ssize_t lattice_num_edges(Args... args) {
+        return LatticeType(args...).num_edges();
+    }
+
+    template<class... Args>
+    static ssize_t lattice_num_nodes(Args... args) {
+        return LatticeType(args...).num_nodes();
+    }
+
     /// Get the linear bias associated with `u`. Returns `0` if `u` is out-of-bounds.
     double linear(int u) const noexcept;
 

dwave/optimization/libcpp/__init__.pxd

@@ -12,8 +12,15 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 
+from libcpp.functional cimport function
+
 # As of Cython 3.0.8 these are not in Cython's libcpp
 
+cdef extern from "<functional>" namespace "std" nogil:
+    # We just do the overloads we need
+    function[double(int)] bind_front(double(void*, int), void*)
+    function[double(int, int)] bind_front(double(void*, int, int), void*)
+
 cdef extern from "<span>" namespace "std" nogil:
     cdef cppclass span[T]:
         ctypedef size_t size_type

dwave/optimization/libcpp/nodes.pxd

@@ -262,6 +262,14 @@ cdef extern from "dwave-optimization/nodes/quadratic_model.hpp" namespace "dwave
     cdef cppclass QuadraticModelNode(ArrayNode):
         QuadraticModel* get_quadratic_model()
 
+    cdef cppclass ZephyrNode(ArrayNode):
+        @staticmethod
+        Py_ssize_t lattice_num_nodes(Py_ssize_t m)
+        @staticmethod
+        Py_ssize_t lattice_num_edges(Py_ssize_t m)
+        double linear(int v)
+        double quadratic(int u, int v)
+
 
 cdef extern from "dwave-optimization/nodes/testing.hpp" namespace "dwave::optimization" nogil:
     cdef cppclass ArrayValidationNode(Node):

dwave/optimization/src/nodes/quadratic_model.cpp

@@ -484,12 +484,15 @@ LatticeNode<T>::LatticeNode(ArrayNode* x_ptr, T lattice, std::function<double(in
     // Add the nodes (with their weights) to the adjacency
     adj_.resize(lattice_.num_nodes());
     for (ssize_t u = 0, N = adj_.size(); u < N; ++u) {
-        adj_[u].bias = linear(u);
+        const double bias = linear(u);
+        if (!std::isfinite(bias)) throw std::invalid_argument("biases must be finite");
+        adj_[u].bias = bias;
     }
 
     // Add the edges (with their weights) to the adjacency
     for (const auto& [u, v] : lattice_.edges()) {
         const double bias = quadratic(u, v);
+        if (!std::isfinite(bias)) throw std::invalid_argument("biases must be finite");
         adj_[u].neighbors.emplace_back(v, bias);
         adj_[v].neighbors.emplace_back(u, bias);
     }

dwave/optimization/symbols.pyx

@@ -34,7 +34,7 @@ from libcpp.unordered_map cimport unordered_map
 from libcpp.utility cimport move
 from libcpp.vector cimport vector
 
-from dwave.optimization.libcpp cimport get, holds_alternative, span
+from dwave.optimization.libcpp cimport bind_front, get, holds_alternative, span
 from dwave.optimization.libcpp.array cimport (
     Array as cppArray,
     SizeInfo as cppSizeInfo,
@@ -108,6 +108,7 @@ from dwave.optimization.libcpp.nodes cimport (
     SumNode as cppSumNode,
     WhereNode as cppWhereNode,
     XorNode as cppXorNode,
+    ZephyrNode as cppZephyrNode,
     )
 from dwave.optimization.model cimport ArraySymbol, _Graph, Symbol
 from dwave.optimization.states cimport States
@@ -172,6 +173,7 @@ __all__ = [
     "Sum",
     "Where",
     "Xor",
+    "Zephyr",
     ]
 
 # We would like to be able to do constructions like dynamic_cast[cppConstantNode*](...)
@@ -4003,3 +4005,110 @@ cdef class Xor(ArraySymbol):
     cdef cppXorNode* ptr
 
 _register(Xor, typeid(cppXorNode))
+
+
+cdef double _read_linear(void* obj, int v) noexcept nogil:
+    with gil:
+        try:
+            return (<object>obj)(v)
+        except Exception as ex:
+            (<object>obj).ex = ex
+        # If we got here that means an exception was raised. So we trigger C++'s
+        # error handling by returning a non-finite value.
+        return float('nan')
+
+cdef double _read_quadratic(void* obj, int u, int v) noexcept nogil:
+    # object must be a Python callable!
+    with gil:
+        try:
+            return (<object>obj)(u, v)
+        except Exception as ex:
+            (<object>obj).ex = ex
+        # If we got here that means an exception was raised. So we trigger C++'s
+        # error handling by returning a non-finite value.
+        return float('nan')
+
+# Currently we don't have a matching dwave.optimization.mathematical function
+# so we document it better here than we usually do.
+cdef class Zephyr(ArraySymbol):
+    """A quadratic model with biases arranged according to a Zephyr lattice.
+
+    Encode a quadratic model arranged according to a :term:`Zephyr` lattice.
+
+    Args:
+        x: A 1D array symbol giving the assignments to the variables.
+        m: Grid parameter for the Zephyr lattice.
+        linear:
+            A function that will be called for each node in the Zephyr lattice.
+            Must accept a single ``int`` giving the node index.
+            Must return a finite number giving the linear bias.
+            Will be called once for each node.
+        quadratic:
+            A function that will be called for each edge in the Zephyr lattice.
+            Must accept two ``int`` giving the indices of the nodes that share an edge.
+            Must return a finite number giving the quadratic bias.
+            Will be called once for each edge.
+
+    See also:
+        :func:`dwave_networkx.zephyr_graph()`
+
+    ..versionadded:: 0.6.2
+    """
+    def __init__(self, ArraySymbol x, Py_ssize_t m, *, object linear, object quadratic):
+        cdef _Graph model = x.model
+
+        # todo: we could accept dictionaries directly (probably adding uv + vu for quadratic)
+        if not callable(linear):
+            raise TypeError("linear must be a callable that accepts one int and returns a float")
+        if not callable(quadratic):
+            raise TypeError("quadratic must be a callable that accepts two ints and returns a float")
+
+        # We have to do a bit of convolution in order to do error handling
+        # We start by making some aliases of our input functions so that we
+        # can add an .ex attribute that will signal if an exception was raised
+        # in the C++ code.
+        def _linear(v):
+            return linear(v)
+        _linear.ex = None
+        def _quadratic(u, v):
+            return quadratic(u, v)
+        _quadratic.ex = None
+
+        try:
+            self.ptr = model._graph.emplace_node[cppZephyrNode](
+                x.array_ptr,
+                m,
+                bind_front(_read_linear, <void*>_linear),
+                bind_front(_read_quadratic, <void*>_quadratic),
+            )
+        except Exception as ex:
+            # Determine if the error was raised by the Python function or by C++
+            if _linear.ex:
+                raise _linear.ex from None
+            if _quadratic.ex:
+                raise _quadratic.ex from None
+            raise ex
+
+        self.initialize_arraynode(model, self.ptr)
+
+    @staticmethod
+    def lattice_num_edges(Py_ssize_t m):
+        """Return the number of edges in a Zephyr lattice with grid parameter ``m``."""
+        return cppZephyrNode.lattice_num_edges(m)
+
+    @staticmethod
+    def lattice_num_nodes(Py_ssize_t m):
+        """Return the number of nodes in a Zephyr lattice with grid parameter ``m``."""
+        return cppZephyrNode.lattice_num_nodes(m)
+
+    def linear(self, Py_ssize_t v):
+        """Return the linear bias associated with node ``v``, or ``0`` if the node
+        is not in the model."""
+        return self.ptr.linear(v)
+
+    def quadratic(self, Py_ssize_t u, Py_ssize_t v):
+        """Return the quadratic bias associated with edge ``u, v``, or ``0`` if
+        the edge is not in the model."""
+        return self.ptr.quadratic(u, v)
+
+    cdef cppZephyrNode* ptr

releasenotes/notes/zephyr-3ef838b269142a72.yaml

@@ -0,0 +1,3 @@
+---
+features:
+  - Add Python ``Zephyr`` symbol and C++ ``ZephyrNode``.

tests/requirements.txt

@@ -1 +1,2 @@
 dimod==0.12.17
+dwave-networkx==0.8.17

tests/test_symbols.py

@@ -2936,3 +2936,84 @@ def test_array_xor(self):
 
         np.testing.assert_array_equal(ab.state(0), [0, 1, 0, 1])
         np.testing.assert_array_equal(ac.state(0), [1, 0, 1, 0])
+
+
+class TestZephyr(unittest.TestCase):
+    def test(self):
+        from dwave.optimization.symbols import Zephyr
+
+        model = Model()
+
+        num_nodes = Zephyr.lattice_num_nodes(2)
+
+        x = model.binary(num_nodes)
+        z = Zephyr(x, 2, linear=lambda v: 5, quadratic=lambda u, v: 1)
+
+    def test_exceptions(self):
+        from dwave.optimization.symbols import Zephyr
+
+        model = Model()
+        x = model.binary(Zephyr.lattice_num_nodes(2))
+
+        def _raise(*args):
+            raise ValueError("boom")
+
+        with self.assertRaisesRegex(ValueError, "boom"):
+            Zephyr(x, 2, linear=_raise, quadratic=lambda u, v: 1)
+        with self.assertRaisesRegex(ValueError, "boom"):
+            Zephyr(x, 2, linear=lambda u: 1, quadratic=_raise)
+
+        def _infinite(*args):
+            return float('inf')
+        with self.assertRaisesRegex(ValueError, "biases must be finite"):
+            Zephyr(x, 2, linear=_infinite, quadratic=lambda u, v: 1)
+
+        # and, very importantly, there shouldn't be any side effects nodes
+        self.assertEqual(model.num_symbols(), 1)
+        self.assertEqual(len(list(x.iter_successors())), 0)
+
+    def test_graph(self):
+        try:
+            import dwave_networkx as dnx
+        except ImportError:
+            return self.skipTest("no dwave_networkx")
+
+        from dwave.optimization.symbols import Zephyr
+
+        model = Model()
+
+        for m in range(1, 4):
+            with self.subTest(m=m):
+                G = dnx.zephyr_graph(m)
+
+                for v in G.nodes:
+                    G.nodes[v]["count"] = 0
+                for u, v in G.edges:
+                    G.edges[u, v]["count"] = 0
+
+                x = model.binary(len(G.nodes))
+
+                def linear(v):
+                    if not G.has_node(v):
+                        raise ValueError("no node")
+                    G.nodes[v]["count"] += 1
+                    return v
+
+                def quadratic(u, v):
+                    if not G.has_edge(u, v):
+                        raise ValueError("no edge")
+                    G.edges[u, v]["count"] += 1
+                    return u*v
+
+                z = Zephyr(x, m, linear=linear, quadratic=quadratic)
+
+                # each node/edge should have been called exactly once
+                self.assertTrue(all(G.nodes[v]["count"] == 1 for v in G.nodes))
+                self.assertTrue(all(G.edges[u, v]["count"] == 1 for u, v in G.edges))
+
+                # and they got the bias we expected
+                for v in G.nodes:
+                    self.assertEqual(z.linear(v), v)
+                for u, v in G.edges:
+                    self.assertEqual(z.quadratic(u, v), u * v)
+                    self.assertEqual(z.quadratic(v, u), u * v)