Enable state probing for OptimizationSolver on CPU and Loihi backend (#…

…217)

* Enable cost tracking with Loihi backend (#212)

* Refactor solution readout by method extraction.

Signed-off-by: GaboFGuerra <[email protected]>

* Enable probing of variable assignment on OptimizationSolver.

Signed-off-by: GaboFGuerra <[email protected]>

* Bugfix in SolutionReadout pymodel

* Fix linting issues

* Add test for state tracking on OptimizationSolver

* Improve OptimizationSolver unittests

* Bugfix in OptimizationSolver unittest

* Bugfix in OptimizationSolver unittest

* Update OptimizationSolver to rtack live cost

* Bugfix in OptimizationSolver and linting issues

* Update poetry.lock

* Add decoding step in state tracking

* Refactor to consolidate cost/state decoding

* Fix linting issues

---------

Signed-off-by: GaboFGuerra <[email protected]>
Co-authored-by: GaboFGuerra <[email protected]>

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,46 @@ 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
+
+    @staticmethod
+    def decode_cost(raw_cost) -> np.ndarray:
+        # The following casts cost as a signed 24-bit value (8 = 32 - 24)
+        return (np.array([raw_cost]).astype(np.int32) << 8) >> 8
+
+    @staticmethod
+    def decode_solution(raw_solution) -> np.ndarray:
+        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

@@ -1,37 +1,13 @@
 # Copyright (C) 2021 Intel Corporation
 # SPDX-License-Identifier: BSD-3-Clause
 # See: https://spdx.org/licenses/
-import numpy as np
 import typing as ty
-
 from dataclasses import dataclass
-from lava.lib.optimization.problems.problems import OptimizationProblem
-from lava.lib.optimization.solvers.generic.builder import SolverProcessBuilder
-from lava.lib.optimization.solvers.generic.hierarchical_processes import (
-    NEBMAbstract,
-    NEBMSimulatedAnnealingAbstract,
-)
-
-from lava.lib.optimization.solvers.generic.scif.models import (
-    PyModelQuboScifFixed,
-)
-from lava.lib.optimization.solvers.generic.nebm.models import NEBMPyModel
-from lava.lib.optimization.solvers.generic.scif.process import QuboScif
-from lava.lib.optimization.solvers.generic.nebm.process import NEBM
-from lava.lib.optimization.solvers.generic.cost_integrator.process import (
-    CostIntegrator,
-)
-from lava.lib.optimization.solvers.generic.nebm.process import (
-    NEBMSimulatedAnnealing,
-)
-from lava.lib.optimization.solvers.generic.sub_process_models import (
-    NEBMAbstractModel,
-    NEBMSimulatedAnnealingAbstractModel,
-)
 
+import numpy as np
 from lava.magma.core.resources import (
-    AbstractComputeResource,
     CPU,
+    AbstractComputeResource,
     Loihi2NeuroCore,
     NeuroCore,
 )
@@ -44,17 +20,43 @@
 from lava.proc.monitor.process import Monitor
 from lava.utils.profiler import Profiler
 
+from lava.lib.optimization.problems.problems import OptimizationProblem
+from lava.lib.optimization.solvers.generic.builder import SolverProcessBuilder
+from lava.lib.optimization.solvers.generic.cost_integrator.process import (
+    CostIntegrator,
+)
+from lava.lib.optimization.solvers.generic.hierarchical_processes import (
+    NEBMAbstract,
+    NEBMSimulatedAnnealingAbstract,
+)
+from lava.lib.optimization.solvers.generic.monitoring_processes.\
+    solution_readout.models import SolutionReadoutPyModel
+from lava.lib.optimization.solvers.generic.nebm.models import NEBMPyModel
+from lava.lib.optimization.solvers.generic.nebm.process import (
+    NEBM,
+    NEBMSimulatedAnnealing,
+)
+from lava.lib.optimization.solvers.generic.scif.models import (
+    PyModelQuboScifFixed,
+)
+from lava.lib.optimization.solvers.generic.scif.process import QuboScif
+from lava.lib.optimization.solvers.generic.sub_process_models import (
+    NEBMAbstractModel,
+    NEBMSimulatedAnnealingAbstractModel,
+)
+
 try:
-    from lava.lib.optimization.solvers.generic.read_gate.ncmodels import (
-        ReadGateCModel,
-    )
     from lava.proc.dense.ncmodels import NcModelDense
+
+    from lava.lib.optimization.solvers.generic.cost_integrator.ncmodels import (
+        CostIntegratorNcModel,
+    )
     from lava.lib.optimization.solvers.generic.nebm.ncmodels import (
         NEBMNcModel,
         NEBMSimulatedAnnealingNcModel,
     )
-    from lava.lib.optimization.solvers.generic.cost_integrator.ncmodels import (
-        CostIntegratorNcModel,
+    from lava.lib.optimization.solvers.generic.read_gate.ncmodels import (
+        ReadGateCModel,
     )
 except ImportError:
 
@@ -137,6 +139,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 +168,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 +221,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 +240,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 +249,43 @@ 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,25 +331,48 @@ 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 = SolutionReadoutPyModel.decode_solution(
+                state_timeseries
+            )
+        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
     ):
-        from lava.lib.optimization.solvers.generic.read_gate.process import (
-            ReadGate,
-        )
         from lava.lib.optimization.solvers.generic.read_gate.models import (
             get_read_gate_model_class,
         )
+        from lava.lib.optimization.solvers.generic.read_gate.process import (
+            ReadGate,
+        )
 
         if backend in CPUS:
             ReadGatePyModel = get_read_gate_model_class(num_in_ports)
@@ -373,17 +419,19 @@ def _prepare_profiler(self, config: SolverConfig, run_cfg) -> None:
 
     def _get_results(self, config: SolverConfig):
         best_cost, idx = self.solver_process.optimum.get()
-        best_cost = (np.asarray([best_cost]).astype(np.int32) << 8) >> 8
+        best_cost = SolutionReadoutPyModel.decode_cost(best_cost)
         best_state = self._get_best_state(config, idx)
         best_timestep = self.solver_process.solution_step.aliased_var.get() - 2
         return best_state, int(best_cost), int(best_timestep)
 
     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()