Separate out Model class into cython _GraphManager and python Model

dwave/optimization/model.pxd → dwave/optimization/graph_manager.pxd

@@ -23,56 +23,20 @@ from dwave.optimization.libcpp.graph cimport ArrayNode as cppArrayNode, Node as
 from dwave.optimization.libcpp.graph cimport Graph as cppGraph
 from dwave.optimization.libcpp.state cimport State as cppState
 
-__all__ = ["Model"]
 
+cdef class _GraphManager:
+    cpdef bool _is_locked(self) noexcept
+    cpdef Py_ssize_t _num_nodes(self) noexcept
 
-cdef class Model:
-    cpdef bool is_locked(self) noexcept
-    cpdef Py_ssize_t num_decisions(self) noexcept
-    cpdef Py_ssize_t num_nodes(self) noexcept
-    cpdef Py_ssize_t num_constraints(self) noexcept
-
-    # Allow dynamic attributes on the Model class
+    # Allow dynamic attributes on the _GraphManager class
     cdef dict __dict__
 
-    # Make the Model class weak referenceable
+    # Make the _GraphManager class weak referenceable
     cdef object __weakref__
 
     cdef cppGraph _graph
 
-    cdef readonly object objective  # todo: cdef ArraySymbol?
-    """Objective to be minimized.
-
-    Examples:
-        This example prints the value of the objective of a model representing 
-        the simple polynomial, :math:`y = i^2 - 4i`, for a state with value 
-        :math:`i=2.0`. 
-
-        >>> from dwave.optimization import Model
-        ...
-        >>> model = Model()
-        >>> i = model.integer(lower_bound=-5, upper_bound=5)
-        >>> c = model.constant(4)
-        >>> y = i**2 - c*i
-        >>> model.minimize(y) 
-        >>> with model.lock():
-        ...     model.states.resize(1)
-        ...     i.set_state(0, 2.0)
-        ...     print(f"Objective = {model.objective.state(0)}")
-        Objective = -4.0
-    """
-
-    cdef readonly States states
-    """States of the model.
-
-    :ref:`States <intro_optimization_states>` represent assignments of values
-    to a symbol.
-
-    See also:
-        :ref:`States methods <optimization_models>` such as 
-        :meth:`~dwave.optimization.model.States.size` and 
-        :meth:`~dwave.optimization.model.States.resize`.
-    """
+    cdef readonly States _states
 
     # The number of times "lock()" has been called.
     cdef readonly Py_ssize_t _lock_count
@@ -82,6 +46,11 @@ cdef class Model:
     # memory for easier cleanup later if that becomes a concern.
     cdef object _data_sources
 
+    cpdef Py_ssize_t _num_constraints(self) noexcept
+    cpdef Py_ssize_t _num_decisions(self) noexcept
+
+    cdef readonly object _objective  # todo: cdef ArraySymbol?
+
 
 cdef class States:
     """The states/solutions of the model."""
@@ -91,7 +60,8 @@ cdef class States:
     cpdef resolve(self)
     cpdef Py_ssize_t size(self) except -1
 
-    cdef Model _model(self)
+    # TODO: rename this to _graph_manager?
+    cdef _GraphManager _model(self)
 
     # In order to not create a circular reference, we only hold a weakref
     # to the model from the states. This introduces some overhead, but it
@@ -114,20 +84,20 @@ cdef class States:
 
 cdef class Symbol:
     # Inheriting nodes must call this method from their __init__()
-    cdef void initialize_node(self, Model model, cppNode* node_ptr) noexcept
+    cdef void initialize_node(self, _GraphManager model, cppNode* node_ptr) noexcept
 
     cpdef uintptr_t id(self) noexcept
 
     # Exactly deref(self.expired_ptr)
     cpdef bool expired(self) noexcept
 
     @staticmethod
-    cdef Symbol from_ptr(Model model, cppNode* ptr)
+    cdef Symbol from_ptr(_GraphManager model, cppNode* ptr)
 
     # Hold on to a reference to the Model, both for access but also, importantly,
     # to ensure that the model doesn't get garbage collected unless all of
     # the observers have also been garbage collected.
-    cdef readonly Model model
+    cdef readonly _GraphManager model
 
     # Hold Node* pointer. This is redundant as most observers will also hold
     # a pointer to their observed node with the correct type. But the cost
@@ -145,7 +115,7 @@ cdef class Symbol:
 # also Symbols (probably a fair assumption)
 cdef class ArraySymbol(Symbol):
     # Inheriting symbols must call this method from their __init__()
-    cdef void initialize_arraynode(self, Model model, cppArrayNode* array_ptr) noexcept
+    cdef void initialize_arraynode(self, _GraphManager model, cppArrayNode* array_ptr) noexcept
 
     # Hold ArrayNode* pointer. Again this is redundant, because we're also holding
     # a pointer to Node* and we can theoretically dynamic cast each time.

dwave/optimization/graph_manager.pyi

@@ -0,0 +1,180 @@
+# 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 collections.abc
+import typing
+
+import numpy
+import numpy.typing
+
+from dwave.optimization.model import Model
+from dwave.optimization.symbols import (
+    Absolute,
+    Add,
+    AdvancedIndexing,
+    All,
+    Any,
+    BasicIndexing,
+    Constant,
+    Equal,
+    LessEqual,
+    Max,
+    Min,
+    Modulus,
+    Multiply,
+    NaryAdd,
+    NaryMultiply,
+    Negative,
+    PartialSum,
+    Permutation,
+    Prod,
+    Reshape,
+    Size,
+    Subtract,
+    Sum,
+)
+
+_ShapeLike: typing.TypeAlias = typing.Union[int, collections.abc.Sequence[int]]
+
+_GraphManagerT = typing.TypeVar('_GraphManagerT', bound='_GraphManager')
+
+
+class _GraphManager:
+    def __init__(self): ...
+
+    @property
+    def _states(self) -> "States": ...
+    def _constant(self, array_like: numpy.typing.ArrayLike) -> Constant: ...
+
+    @classmethod
+    def _from_file(
+        cls: typing.Type[_GraphManagerT],
+        file: typing.Union[typing.BinaryIO, collections.abc.ByteString, str],
+        *,
+        check_header: bool = True,
+    ) -> _GraphManagerT: ...
+
+    def _into_file(
+        self,
+        file: typing.Union[typing.BinaryIO, collections.abc.ByteString, str],
+        *,
+        max_num_states: int = 0,
+        only_decision: bool = False,
+    ): ...
+    def _is_locked(self) -> bool: ...
+    def _iter_symbols(self) -> collections.abc.Iterator["Symbol"]: ...
+    def _num_nodes(self) -> int: ...
+    def _num_symbols(self) -> int: ...
+    def _remove_unused_symbols(self) -> int: ...
+    def _state_size(self) -> int: ...
+
+    def _to_file(
+        self,
+        *,
+        max_num_states: int = 0,
+        only_decision: bool = False,
+    ) -> typing.BinaryIO: ...
+
+    # networkx might not be installed, so we just say we return an object.
+    def _to_networkx(self) -> object: ...
+    def _unlock(self): ...
+
+    @property
+    def _objective(self) -> "ArraySymbol": ...
+    def _add_constraint(self, value: "ArraySymbol") -> "ArraySymbol": ...
+    def _iter_constraints(self) -> collections.abc.Iterator["ArraySymbol"]: ...
+    def _iter_decisions(self) -> collections.abc.Iterator["Symbol"]: ...
+    def _minimize(self, value: "ArraySymbol"): ...
+    def _num_constraints(self) -> int: ...
+    def _num_decisions(self) -> int: ...
+
+
+class States:
+    def __init__(self, model: Model): ...
+    def __len__(self) -> int: ...
+    def clear(self): ...
+
+    def from_file(
+        self,
+        file: typing.Union[typing.BinaryIO, collections.abc.ByteString, str],
+        *,
+        replace: bool = True,
+        check_header: bool = True,
+    ) -> Model: ...
+    def from_future(self, future: object, result_hook: collections.abc.Callable): ...
+    def initialize(self): ...
+
+    def into_file(
+        self,
+        file: typing.Union[typing.BinaryIO, collections.abc.ByteString, str],
+    ): ...
+
+    def resize(self, n: int): ...
+    def resolve(self): ...
+    def size(self) -> int: ...
+    def to_file(self) -> typing.BinaryIO: ...
+
+
+class Symbol:
+    def __init__(self, *args, **kwargs) -> typing.NoReturn: ...
+    def equals(self, other: "Symbol") -> bool: ...
+    def expired(self) -> bool: ...
+    def has_state(self, index: int = 0) -> bool: ...
+    def id(self) -> int: ...
+    def iter_predecessors(self) -> collections.abc.Iterator["Symbol"]: ...
+    def iter_successors(self) -> collections.abc.Iterator["Symbol"]: ...
+    def maybe_equals(self, other: "Symbol") -> int: ...
+    def reset_state(self, index: int): ...
+    def shares_memory(self, other: "Symbol") -> bool: ...
+    def state_size(self) -> int: ...
+    def topological_index(self) -> int: ...
+
+
+class ArraySymbol(Symbol):
+    def __init__(self, *args, **kwargs) -> typing.NoReturn: ...
+    def __abs__(self) -> Absolute: ...
+    def __add__(self, rhs: "ArraySymbol") -> Add: ...
+    def __bool__(self) -> typing.NoReturn: ...
+    def __eq__(self, rhs: "ArraySymbol") -> Equal: ...  # type: ignore
+
+    def __getitem__(
+        self,
+        index: typing.Union[Symbol, int, slice, tuple],
+    ) -> typing.Union[AdvancedIndexing, BasicIndexing, Permutation]: ...
+
+    def __iadd__(self, rhs: "ArraySymbol") -> NaryAdd: ...
+    def __imul__(self, rhs: "ArraySymbol") -> NaryMultiply: ...
+    def __le__(self, rhs: "ArraySymbol") -> LessEqual: ...
+    def __mod__(self, rhs: "ArraySymbol") -> Modulus: ...
+    def __mul__(self, rhs: "ArraySymbol") -> Multiply: ...
+    def __neg__(self) -> Negative: ...
+    def __pow__(self, exponent: int) -> "ArraySymbol": ...
+    def __sub__(self, rhs: "ArraySymbol") -> Subtract: ...
+    def all(self) -> All: ...
+    def any(self) -> Any: ...
+    def max(self) -> Max: ...
+    def min(self) -> Min: ...
+    def ndim(self) -> int: ...
+    def prod(self) -> Prod: ...
+    def reshape(self, shape: _ShapeLike) -> Reshape: ...
+    def shape(self) -> tuple[int, ...]: ...
+    def size(self) -> typing.Union[int, Size]: ...
+    def sqrt(self) -> "ArraySymbol": ...
+    def state(self, index: int = 0, *, copy: bool = True) -> numpy.ndarray: ...
+    def state_size(self) -> int: ...
+    def strides(self) -> tuple[int, ...]: ...
+
+    def sum(
+        self, axis: typing.Optional[int] = None
+    ) -> typing.Union[Sum, PartialSum]: ...