Let Model rather than _Graph hold .states and .objective

dwave/optimization/_graph.pxd

@@ -36,40 +36,6 @@ cdef class _Graph:
 
     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 object 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`.
-    """
-
     # The number of times "lock()" has been called.
     cdef readonly Py_ssize_t _lock_count
 

dwave/optimization/_graph.pyi

@@ -32,11 +32,6 @@ _GraphSubclass = typing.TypeVar("_GraphSubclass", bound="_Graph")
 class _Graph:
     def __init__(self, *args, **kwargs) -> typing.NoReturn: ...
 
-    @property
-    def objective(self) -> ArraySymbol: ...
-    @property
-    def states(self) -> States: ...
-
     def add_constraint(self, value: ArraySymbol) -> ArraySymbol: ...
     def decision_state_size(self) -> int: ...
 

dwave/optimization/_graph.pyx

@@ -41,7 +41,7 @@ __all__ = []
 
 cdef class _Graph:
     def __cinit__(self):
-        self.states = States(self)
+        self._lock_count = 0
         self._data_sources = []
 
     def __init__(self, *args, **kwargs):
@@ -213,7 +213,10 @@ cdef class _Graph:
     def _header_data(self, *, only_decision, max_num_states=float('inf')):
         """The header data associated with the model (but not the states)."""
         num_nodes = self.num_decisions() if only_decision else self.num_nodes()
-        num_states = max(0, min(self.states.size(), max_num_states))
+        try:
+            num_states = max(0, min(self.states.size(), max_num_states))
+        except AttributeError:
+            num_states = 0
 
         decision_state_size = self.decision_state_size()
         state_size = decision_state_size if only_decision else self.state_size()
@@ -484,7 +487,6 @@ cdef class _Graph:
         if value.size() > 1:
             raise ValueError("the value of an array with more than one element is ambiguous")
         self._graph.set_objective(value.array_ptr)
-        self.objective = value
 
     cpdef Py_ssize_t num_constraints(self) noexcept:
         """Number of constraints in the model.
@@ -672,16 +674,10 @@ cdef class _Graph:
 
         self._lock_count -= 1
 
-        cdef States states = self.states  # for Cython access
-
         # if we're now unlocked, then reset the topological sort and the
         # non-decision states
         if self._lock_count < 1:
             self._graph.reset_topological_sort()
-            for i in range(states.size()):
-                # this might actually increase the size of the states in some
-                # cases, but that's fine
-                states._states[i].resize(self.num_decisions())
 
 
 cdef class Symbol:
@@ -791,6 +787,9 @@ cdef class Symbol:
             raise TypeError("the state of an intermediate variable cannot be accessed without "
                             "locking the model first. See model.lock().")
 
+        if not hasattr(self.model, "states"):
+            return False
+
         cdef States states = self.model.states  # for Cython access
 
         states.resolve()

dwave/optimization/model.py

@@ -25,8 +25,10 @@
 
 import contextlib
 import tempfile
+import typing
 
 from dwave.optimization._graph import ArraySymbol, _Graph, Symbol
+from dwave.optimization.states import States
 
 __all__ = ["Model"]
 
@@ -61,8 +63,43 @@ class Model(_Graph):
         >>> model = flow_shop_scheduling(processing_times=processing_times)
     """
 
+    objective: typing.Optional[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
+    """
+
+    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`.
+    """
+
     def __init__(self):
-        pass
+        self.objective = None
+        self.states = States(self)
 
     def binary(self, shape=None):
         r"""Create a binary symbol as a decision variable.
@@ -285,6 +322,11 @@ def lock(self):
         super().lock()
         return locked(self)
 
+    def minimize(self, value: ArraySymbol):
+        # inherit the docstring from _Graph
+        super().minimize(value)
+        self.objective = value
+
     def quadratic_model(self, x, quadratic, linear=None):
         """Create a quadratic model from an array and a quadratic model.
 
@@ -414,3 +456,13 @@ def to_networkx(self):
             G.add_edge(repr(symbol), "constraint(s)")
 
         return G
+
+    def unlock(self):
+        # inherit the docstring from _Graph
+        if not self.is_locked():
+            return
+
+        super().unlock()
+
+        if not self.is_locked():
+            self.states._reset_intermediate_states()

dwave/optimization/states.pyi

@@ -39,6 +39,7 @@ class States:
         file: typing.Union[typing.BinaryIO, collections.abc.ByteString, str],
         ): ...
 
+    def _reset_intermediate_states(self): ...
     def resize(self, n: int): ...
     def resolve(self): ...
     def size(self) -> int: ...

dwave/optimization/states.pyx

@@ -222,6 +222,12 @@ cdef class States:
             raise ReferenceError("accessing the states of a garbage collected model")
         return m
 
+    def _reset_intermediate_states(self):
+        """Reset any non-decision states."""
+        cdef Py_ssize_t num_decisions = self._model().num_decisions()
+        for i in range(self.size()):
+            self._states[i].resize(num_decisions)
+
     def resize(self, Py_ssize_t n):
         """Resize the number of states.