Replace ballistic by mode

src/simulated_bifurcation/core/ising.py

@@ -18,7 +18,7 @@
 
 """
 
-from typing import Optional, TypeVar, Union
+from typing import Literal, Optional, TypeVar, Union
 
 import torch
 from numpy import ndarray
@@ -241,7 +241,7 @@ def minimize(
         self,
         agents: int = 128,
         max_steps: int = 10000,
-        ballistic: bool = False,
+        mode: Literal["ballistic", "discrete"] = "ballistic",
         heated: bool = False,
         verbose: bool = True,
         *,
@@ -263,7 +263,7 @@ def minimize(
         max_steps : int, default=10_000
             Number of iterations after which the algorithm is stopped
             regardless of whether convergence has been achieved.
-        ballistic : bool, default=False
+        mode : "ballistic" | "discrete", optional, default = "ballistic"
             Whether to use the ballistic or the discrete SB algorithm.
             See Notes for further information about the variants of the SB
             algorithm.
@@ -390,12 +390,13 @@ def minimize(
         https://doi.org/10.1038/s42005-022-00929-9
 
         """
-        engine = SimulatedBifurcationEngine.get_engine(ballistic, heated)
+        engine = SimulatedBifurcationEngine.get_engine(mode)
         optimizer = SimulatedBifurcationOptimizer(
             agents,
             max_steps,
             timeout,
             engine,
+            heated,
             verbose,
             sampling_period,
             convergence_threshold,

src/simulated_bifurcation/core/quadratic_polynomial.py

@@ -21,7 +21,7 @@
 """
 
 import re
-from typing import Optional, Tuple, Union
+from typing import Literal, Optional, Tuple, Union
 
 import numpy as np
 import torch
@@ -313,7 +313,7 @@ def optimize(
         agents: int = 128,
         max_steps: int = 10000,
         best_only: bool = True,
-        ballistic: bool = False,
+        mode: Literal["ballistic", "discrete"] = "ballistic",
         heated: bool = False,
         minimize: bool = True,
         verbose: bool = True,
@@ -379,9 +379,10 @@ def optimize(
         timeout : float | None, default=None
             Time in seconds after which the simulation is stopped.
             None means no timeout.
-        ballistic : bool, optional
-            if True, the ballistic SB will be used, else it will be the
-            discrete SB (default is True)
+        mode : "ballistic" | "discrete", optional, default = "ballistic"
+            Whether to use the ballistic or the discrete SB algorithm.
+            See Notes for further information about the variants of the SB
+            algorithm.
         heated : bool, optional
             if True, the heated SB will be used, else it will be the non-heated
             SB (default is True)
@@ -407,7 +408,7 @@ def optimize(
         ising_equivalent.minimize(
             agents,
             max_steps,
-            ballistic,
+            mode,
             heated,
             verbose,
             use_window=use_window,
@@ -430,7 +431,7 @@ def minimize(
         agents: int = 128,
         max_steps: int = 10000,
         best_only: bool = True,
-        ballistic: bool = False,
+        mode: Literal["ballistic", "discrete"] = "ballistic",
         heated: bool = False,
         verbose: bool = True,
         *,
@@ -495,9 +496,10 @@ def minimize(
         timeout : float | None, default=None
             Time in seconds after which the simulation is stopped.
             None means no timeout.
-        ballistic : bool, optional
-            if True, the ballistic SB will be used, else it will be the
-            discrete SB (default is True)
+        mode : "ballistic" | "discrete", optional, default = "ballistic"
+            Whether to use the ballistic or the discrete SB algorithm.
+            See Notes for further information about the variants of the SB
+            algorithm.
         heated : bool, optional
             if True, the heated SB will be used, else it will be the non-heated
             SB (default is True)
@@ -518,7 +520,7 @@ def minimize(
             agents,
             max_steps,
             best_only,
-            ballistic,
+            mode,
             heated,
             True,
             verbose,
@@ -534,7 +536,7 @@ def maximize(
         agents: int = 128,
         max_steps: int = 10000,
         best_only: bool = True,
-        ballistic: bool = False,
+        mode: Literal["ballistic", "discrete"] = "ballistic",
         heated: bool = False,
         verbose: bool = True,
         *,
@@ -598,9 +600,10 @@ def maximize(
         timeout : float | None, default=None
             Time in seconds after which the simulation is stopped.
             None means no timeout.
-        ballistic : bool, optional
-            if True, the ballistic SB will be used, else it will be the
-            discrete SB (default is True)
+        mode : "ballistic" | "discrete", optional, default = "ballistic"
+            Whether to use the ballistic or the discrete SB algorithm.
+            See Notes for further information about the variants of the SB
+            algorithm.
         heated : bool, optional
             if True, the heated SB will be used, else it will be the non-heated
             SB (default is True)
@@ -621,7 +624,7 @@ def maximize(
             agents,
             max_steps,
             best_only,
-            ballistic,
+            mode,
             heated,
             False,
             verbose,

src/simulated_bifurcation/models/abc_model.py

@@ -1,5 +1,5 @@
 from abc import ABC
-from typing import Optional, Tuple
+from typing import Literal, Optional, Tuple
 
 import torch
 
@@ -26,7 +26,7 @@ def optimize(
         agents: int = 128,
         max_steps: int = 10000,
         best_only: bool = True,
-        ballistic: bool = False,
+        mode: Literal["ballistic", "discrete"] = "ballistic",
         heated: bool = False,
         minimize: bool = True,
         verbose: bool = True,
@@ -41,7 +41,7 @@ def optimize(
             agents,
             max_steps,
             best_only,
-            ballistic,
+            mode,
             heated,
             minimize,
             verbose,
@@ -56,7 +56,7 @@ def minimize(
         agents: int = 128,
         max_steps: int = 10000,
         best_only: bool = True,
-        ballistic: bool = False,
+        mode: Literal["ballistic", "discrete"] = "ballistic",
         heated: bool = False,
         verbose: bool = True,
         *,
@@ -69,7 +69,7 @@ def minimize(
             agents,
             max_steps,
             best_only,
-            ballistic,
+            mode,
             heated,
             True,
             verbose,
@@ -84,7 +84,7 @@ def maximize(
         agents: int = 128,
         max_steps: int = 10000,
         best_only: bool = True,
-        ballistic: bool = False,
+        mode: Literal["ballistic", "discrete"] = "ballistic",
         heated: bool = False,
         verbose: bool = True,
         *,
@@ -97,7 +97,7 @@ def maximize(
             agents,
             max_steps,
             best_only,
-            ballistic,
+            mode,
             heated,
             False,
             verbose,

src/simulated_bifurcation/optimizer/simulated_bifurcation_engine.py

@@ -1,34 +1,31 @@
 from enum import Enum
+from typing import Callable, Literal
 
 import torch
 
 
 class SimulatedBifurcationEngine(Enum):
     """
-    Enum class that gathers the 4 variants of the Simulated Bifurcation
-    algorithm:
-
-    1. Ballistic SB (bSB)
-    2. Discrete SB (dSB)
-    3. Heated ballistic SB (HbSB)
-    4. Heated discrete SB (HdSB)
+    Variants of the Simulated Bifurcation algorithm.
     """
 
-    bSB = torch.nn.Identity(), False
-    dSB = torch.sign, False
-    HbSB = torch.nn.Identity(), True
-    HdSB = torch.sign, True
+    bSB = torch.nn.Identity()
+    dSB = torch.sign
+
+    def __init__(
+        self, activation_function: Callable[[torch.Tensor], torch.Tensor]
+    ) -> None:
+        self.__activation_function = activation_function
 
-    def __init__(self, activation_function, heated: bool) -> None:
-        self.activation_function = activation_function
-        self.heated = heated
+    @property
+    def activation_function(self) -> Callable[[torch.Tensor], torch.Tensor]:
+        return self.__activation_function
 
     @staticmethod
-    def get_engine(ballistic: bool, heated: bool):
-        if ballistic:
-            if heated:
-                return SimulatedBifurcationEngine.HbSB
+    def get_engine(engine_name: Literal["ballistic", "discrete"]):
+        if engine_name == "ballistic":
             return SimulatedBifurcationEngine.bSB
-        if heated:
-            return SimulatedBifurcationEngine.HdSB
-        return SimulatedBifurcationEngine.dSB
+        elif engine_name == "discrete":
+            return SimulatedBifurcationEngine.dSB
+        else:
+            raise ValueError(f"Unknwown Simulated Bifurcation engine: {engine_name}.")

src/simulated_bifurcation/optimizer/simulated_bifurcation_optimizer.py

@@ -69,6 +69,7 @@ def __init__(
         max_steps: Optional[int],
         timeout: Optional[float],
         engine: SimulatedBifurcationEngine,
+        heated: bool,
         verbose: bool,
         sampling_period: int,
         convergence_threshold: int,
@@ -78,7 +79,7 @@ def __init__(
         self.window = None
         self.symplectic_integrator = None
         self.heat_coefficient = ENVIRONMENT.heat_coefficient
-        self.heated = engine.heated
+        self.heated = heated
         self.verbose = verbose
         self.start_time = None
         self.simulation_time = None

src/simulated_bifurcation/simulated_bifurcation.py

@@ -28,7 +28,7 @@
 
 """
 
-from typing import Optional, Tuple, Union
+from typing import Literal, Optional, Tuple, Union
 
 import torch
 
@@ -164,7 +164,7 @@ def optimize(
     agents: int = 128,
     max_steps: int = 10_000,
     best_only: bool = True,
-    ballistic: bool = False,
+    mode: Literal["ballistic", "discrete"] = "ballistic",
     heated: bool = False,
     minimize: bool = True,
     verbose: bool = True,
@@ -228,7 +228,7 @@ def optimize(
         polynomial at this vector. Otherwise, returns all the vectors
         found by the SB algorithm and the values of polynomial at these
         points.
-    ballistic : bool, default=False, keyword-only
+    mode : "ballistic" | "discrete",  default="ballistic, keyword-only
         Whether to use the ballistic or the discrete SB algorithm.
         See Notes for further information about the variants of the SB
         algorithm.
@@ -422,7 +422,7 @@ def optimize(
         agents=agents,
         max_steps=max_steps,
         best_only=best_only,
-        ballistic=ballistic,
+        mode=mode,
         heated=heated,
         minimize=minimize,
         verbose=verbose,
@@ -442,7 +442,7 @@ def minimize(
     agents: int = 128,
     max_steps: int = 10_000,
     best_only: bool = True,
-    ballistic: bool = False,
+    mode: Literal["ballistic", "discrete"] = "ballistic",
     heated: bool = False,
     verbose: bool = True,
     use_window: bool = True,
@@ -505,7 +505,7 @@ def minimize(
         polynomial at this vector. Otherwise, returns all the vectors
         found by the SB algorithm and the values of polynomial at these
         points.
-    ballistic : bool, default=False, keyword-only
+    mode : "ballistic" | "discrete",  default="ballistic, keyword-only
         Whether to use the ballistic or the discrete SB algorithm.
         See Notes for further information about the variants of the SB
         algorithm.
@@ -686,7 +686,7 @@ def minimize(
         agents=agents,
         max_steps=max_steps,
         best_only=best_only,
-        ballistic=ballistic,
+        mode=mode,
         heated=heated,
         minimize=True,
         verbose=verbose,
@@ -705,7 +705,7 @@ def maximize(
     agents: int = 128,
     max_steps: int = 10_000,
     best_only: bool = True,
-    ballistic: bool = False,
+    mode: Literal["ballistic", "discrete"] = "ballistic",
     heated: bool = False,
     verbose: bool = True,
     use_window: bool = True,
@@ -768,7 +768,7 @@ def maximize(
         polynomial at this vector. Otherwise, returns all the vectors
         found by the SB algorithm and the values of polynomial at these
         points.
-    ballistic : bool, default=False, keyword-only
+    mode : "ballistic" | "discrete",  default="ballistic, keyword-only
         Whether to use the ballistic or the discrete SB algorithm.
         See Notes for further information about the variants of the SB
         algorithm.
@@ -949,7 +949,7 @@ def maximize(
         agents=agents,
         max_steps=max_steps,
         best_only=best_only,
-        ballistic=ballistic,
+        mode=mode,
         heated=heated,
         minimize=False,
         verbose=verbose,

tests/models/test_ising_model.py

@@ -12,7 +12,7 @@ def test_ising():
         agents=10,
         verbose=False,
         best_only=True,
-        ballistic=True,
+        mode="ballistic",
     )
     assert torch.equal(
         torch.tensor([-1.0, 1.0, -1.0], dtype=torch.float32), spin_vector

tests/models/test_knapsack.py

@@ -16,7 +16,7 @@ def test_knapsack():
         "status": "not optimized",
     }
 
-    model.minimize(ballistic=True, verbose=False, agents=1000)
+    model.minimize(mode="ballistic", verbose=False, agents=1000)
     assert model.summary["items"] == [1, 2, 3, 4]
     assert model.summary["total_cost"] == 15
     assert model.summary["total_weight"] == 8