Enable state probing for OptimizationSolver on CPU and Loihi backend

src/lava/lib/optimization/solvers/generic/builder.py

@@ -178,6 +178,9 @@ def constructor(
             self.variable_assignment = Var(
                 shape=(problem.variables.num_variables,)
             )
+            self.best_variable_assignment = Var(
+                shape=(problem.variables.num_variables,)
+            )
             self.optimality = Var(shape=(1,))
             self.optimum = Var(shape=(2,))
             self.feasibility = Var(shape=(1,))
@@ -242,7 +245,12 @@ def constructor(self, proc):
             if hasattr(proc, "cost_coefficients"):
                 proc.vars.optimum.alias(self.solution_reader.min_cost)
                 proc.vars.optimality.alias(proc.finders[0].cost)
-            proc.vars.variable_assignment.alias(self.solution_reader.solution)
+            proc.vars.variable_assignment.alias(
+                proc.finders[0].variables_assignment
+            )
+            proc.vars.best_variable_assignment.alias(
+                self.solution_reader.solution
+            )
             proc.vars.solution_step.alias(self.solution_reader.solution_step)
 
             # Connect processes

src/lava/lib/optimization/solvers/generic/monitoring_processes/solution_readout/models.py

@@ -42,26 +42,44 @@ def run_spk(self):
             return
         raw_cost, min_cost_id = self.cost_in.recv()
         if raw_cost != 0:
-            timestep = self.timestep_in.recv()[0]
-            # The following casts cost as a signed 24-bit value (8 = 32 - 24)
-            cost = (np.array([raw_cost]).astype(np.int32) << 8) >> 8
-            raw_solution = self.read_solution.recv()
-            raw_solution &= 0x1F  # AND with 0x1F (=0b11111) retains 5 LSBs
-            # The binary solution was attained 2 steps ago. Shift down by 4.
-            self.solution[:] = raw_solution.astype(np.int8) >> 4
+            timestep, raw_solution = self._receive_data()
+            cost = self._decode_cost(raw_cost)
             self.solution_step = abs(timestep)
+            self.solution[:] = self._decode_solution(raw_solution)
             self.min_cost[:] = np.asarray([cost[0], min_cost_id])
             if cost[0] < 0:
-                print(
-                    f"Host: better solution found by network {min_cost_id} at "
-                    f"step {abs(timestep)-2} "
-                    f"with cost {cost[0]}: {self.solution}"
-                )
+                self._printout_new_solution(cost, min_cost_id, timestep)
+            self._printout_if_converged()
+            self._stop_if_requested(timestep, min_cost_id)
 
-            if (
+    def _receive_data(self):
+        timestep = self.timestep_in.recv()[0]
+        raw_solution = self.read_solution.recv()
+        return timestep, raw_solution
+
+    def _decode_cost(self, raw_cost):
+        # The following casts cost as a signed 24-bit value (8 = 32 - 24)
+        return (np.array([raw_cost]).astype(np.int32) << 8) >> 8
+
+    def _decode_solution(self, raw_solution):
+        raw_solution &= 0x1F  # AND with 0x1F (=0b11111) retains 5 LSBs
+        # The binary solution was attained 2 steps ago. Shift down by 4.
+        return raw_solution.astype(np.int8) >> 4
+
+    def _printout_new_solution(self, cost, min_cost_id, timestep):
+        print(
+            f"Host: better solution found by network {min_cost_id} at "
+            f"step {abs(timestep) - 2} "
+            f"with cost {cost[0]}: {self.solution}"
+        )
+
+    def _printout_if_converged(self):
+        if (
                 self.min_cost[0] is not None
                 and self.min_cost[0] <= self.target_cost
-            ):
-                print(f"Host: network reached target cost {self.target_cost}.")
-            if timestep > 0 or timestep == -1:
-                self.stop = True
+        ):
+            print(f"Host: network reached target cost {self.target_cost}.")
+
+    def _stop_if_requested(self, timestep, min_cost_id):
+        if (timestep > 0 or timestep == -1) and min_cost_id != -1:
+            self.stop = True

src/lava/lib/optimization/solvers/generic/solver.py

@@ -137,6 +137,7 @@ class SolverConfig:
     backend: BACKENDS = CPU
     hyperparameters: HP_TYPE = None
     probe_cost: bool = False
+    probe_state: bool = False
     probe_time: bool = False
     probe_energy: bool = False
     log_level: int = 40
@@ -165,6 +166,7 @@ class SolverReport:
     best_state: np.ndarray = None
     best_timestep: int = None
     cost_timeseries: np.ndarray = None
+    state_timeseries: np.ndarray = None
     solver_config: SolverConfig = None
     profiler: Profiler = None
 
@@ -217,6 +219,7 @@ def __init__(self, problem: OptimizationProblem):
         self.solver_model = None
         self._profiler = None
         self._cost_tracker = None
+        self._state_tracker = None
 
     def solve(self, config: SolverConfig = SolverConfig()) -> SolverReport:
         """
@@ -235,7 +238,7 @@ def solve(self, config: SolverConfig = SolverConfig()) -> SolverReport:
         run_condition, run_cfg = self._prepare_solver(config)
         self.solver_process.run(condition=run_condition, run_cfg=run_cfg)
         best_state, best_cost, best_timestep = self._get_results(config)
-        cost_timeseries = self._get_cost_tracking()
+        cost_timeseries, state_timeseries = self._get_probing(config)
         self.solver_process.stop()
         return SolverReport(
             best_cost=best_cost,
@@ -244,25 +247,42 @@ def solve(self, config: SolverConfig = SolverConfig()) -> SolverReport:
             solver_config=config,
             profiler=self._profiler,
             cost_timeseries=cost_timeseries,
+            state_timeseries=state_timeseries,
         )
 
     def _prepare_solver(self, config: SolverConfig):
         self._create_solver_process(config=config)
         hps = config.hyperparameters
         num_in_ports = len(hps) if isinstance(hps, list) else 1
-        if config.probe_cost:
-            if config.backend in NEUROCORES:
-                from lava.utils.loihi2_state_probes import StateProbe
+        probes = []
+        if config.backend in NEUROCORES:
+            from lava.utils.loihi2_state_probes import StateProbe
+            if config.probe_cost:
                 self._cost_tracker = StateProbe(self.solver_process.optimality)
-            if config.backend in CPUS:
+                probes.append(self._cost_tracker)
+            if config.probe_state:
+                self._state_tracker = StateProbe(
+                    self.solver_process.variable_assignment
+                )
+                probes.append(self._state_tracker)
+        elif config.backend in CPUS:
+            if config.probe_cost:
                 self._cost_tracker = Monitor()
                 self._cost_tracker.probe(
                     target=self.solver_process.optimality,
                     num_steps=config.timeout,
                 )
+                probes.append(self._cost_tracker)
+            if config.probe_state:
+                self._state_tracker = Monitor()
+                self._state_tracker.probe(
+                    target=self.solver_process.variable_assignment,
+                    num_steps=config.timeout,
+                )
+                probes.append(self._state_tracker)
         run_cfg = self._get_run_config(
             backend=config.backend,
-            probes=[self._cost_tracker] if self._cost_tracker else None,
+            probes=probes,
             num_in_ports=num_in_ports,
         )
         run_condition = RunSteps(num_steps=config.timeout)
@@ -308,15 +328,37 @@ def _get_requirements_and_protocol(
         """
         return [CPU] if backend in CPUS else [Loihi2NeuroCore], LoihiProtocol
 
-    def _get_cost_tracking(self):
-        if self._cost_tracker is None:
+    def _get_probing(
+        self, config: SolverConfig()
+    ) -> ty.Tuple[np.ndarray, np.ndarray]:
+        """
+        Return the cost and state timeseries if probed.
+
+        Parameters
+        ----------
+        config: SolverConfig
+            Solver configuraiton used. Refers to SolverConfig documentation.
+        """
+        cost_timeseries = self._get_probed_data(
+            tracker=self._cost_tracker, var_name="optimality"
+        )
+        state_timeseries = self._get_probed_data(
+            tracker=self._state_tracker, var_name="variable_assignment"
+        )
+        if state_timeseries is not None:
+            state_timeseries &= 0x1F
+            state_timeseries = state_timeseries.astype(np.int8) >> 4
+        return cost_timeseries, state_timeseries
+
+    def _get_probed_data(self, tracker, var_name):
+        if tracker is None:
             return None
-        if isinstance(self._cost_tracker, Monitor):
-            return self._cost_tracker.get_data()[self.solver_process.name][
-                self.solver_process.optimality.name
-            ].T.astype(np.int32)
+        if isinstance(tracker, Monitor):
+            return tracker.get_data()[self.solver_process.name][
+                getattr(self.solver_process, var_name).name
+            ].astype(np.int32)
         else:
-            return self._cost_tracker.time_series
+            return tracker.time_series
 
     def _get_run_config(
         self, backend: BACKENDS, probes=None, num_in_ports: int = None
@@ -380,10 +422,12 @@ def _get_results(self, config: SolverConfig):
 
     def _get_best_state(self, config: SolverConfig, idx: int):
         if isinstance(config.hyperparameters, list):
+            idx = int(idx)
             raw_solution = np.asarray(
-                self.solver_process.finders[int(idx)].variables_assignment.get()
+                self.solver_process.finders[idx].variables_assignment.get()
             ).astype(np.int32)
             raw_solution &= 0x3F
             return raw_solution.astype(np.int8) >> 5
         else:
-            return self.solver_process.variable_assignment.aliased_var.get()
+            best_assignment = self.solver_process.best_variable_assignment
+            return best_assignment.aliased_var.get()

tests/lava/lib/optimization/solvers/generic/test_solver_cpu_backend.py

@@ -34,7 +34,6 @@ def test_create_obj(self):
         self.assertIsInstance(self.solver, OptimizationSolver)
 
     def test_solution_has_expected_shape(self):
-        print("test_solution_has_expected_shape")
         report = self.solver.solve(config=SolverConfig(timeout=3000))
         self.assertEqual(report.best_state.shape, self.solution.shape)
 
@@ -47,7 +46,6 @@ def test_solve_method_nebm(self):
             hyperparameters={"neuron_model": "nebm"},
         )
         report = self.solver.solve(config=config)
-        print(report)
         self.assertTrue((report.best_state == self.solution).all())
         self.assertEqual(report.best_cost, self.problem.evaluate_cost(
             report.best_state))
@@ -61,7 +59,6 @@ def test_solve_method_scif(self):
                 hyperparameters={"neuron_model": "scif", "noise_precision": 5},
             )
         )
-        print(report)
         self.assertTrue((report.best_state == self.solution).all())
         self.assertEqual(report.best_cost, self.problem.evaluate_cost(
             report.best_state))
@@ -85,16 +82,14 @@ def test_subprocesses_connections(self):
         pm = self.solver.solver_process.model_class(self.solver.solver_process)
         solution_finder = pm.finder_0
         solution_reader = pm.solution_reader
+        best_assignment = self.solver.solver_process.best_variable_assignment
         self.assertIs(
             solution_finder.cost_out.out_connections[0].process,
             solution_reader,
         )
+        self.assertIs(best_assignment.aliased_var, solution_reader.solution)
         self.assertIs(
-            self.solver.solver_process.variable_assignment.aliased_var,
-            solution_reader.solution,
-        )
-        self.assertIs(
-            self.solver.solver_process.variable_assignment.aliased_var.process,
+            best_assignment.aliased_var.process,
             solution_reader,
         )
 
@@ -109,7 +104,6 @@ def test_qubo_cost_defines_weights(self):
 
     def test_cost_tracking(self):
         np.random.seed(77)
-        print("test_cost_tracking")
         config = SolverConfig(
             timeout=50,
             target_cost=-20,
@@ -119,7 +113,22 @@ def test_cost_tracking(self):
         report = self.solver.solve(config=config)
         self.assertIsInstance(report.cost_timeseries, np.ndarray)
         self.assertEqual(report.best_cost,
-                         report.cost_timeseries[0][report.best_timestep])
+                         report.cost_timeseries.T[0][report.best_timestep])
+
+    def test_state_tracking(self):
+        np.random.seed(77)
+        config = SolverConfig(
+            timeout=50,
+            target_cost=-20,
+            backend="CPU",
+            probe_state=True
+        )
+        report = self.solver.solve(config=config)
+        states = report.state_timeseries
+        self.assertIsInstance(states, np.ndarray)
+        self.assertTrue(
+            np.all(report.best_state == states[report.best_timestep])
+        )
 
 
 def solve_workload(problem, reference_solution, noise_precision=5,