qoro-divi 0.3.5__py3-none-any.whl → 0.4.0__py3-none-any.whl

divi/qprog/optimizers.py

@@ -6,6 +6,7 @@ import time
 from abc import ABC, abstractmethod
 from collections.abc import Callable
 from enum import Enum
+from typing import Any
 
 import numpy as np
 from pymoo.algorithms.soo.nonconvex.cmaes import CMAES
@@ -37,16 +38,25 @@ class Optimizer(ABC):
         self,
         cost_fn: Callable[[np.ndarray], float],
         initial_params: np.ndarray,
-        callback_fn: Callable | None = None,
+        callback_fn: Callable[[OptimizeResult], Any] | None = None,
         **kwargs,
     ) -> OptimizeResult:
-        """
-        Optimize the given cost function starting from initial parameters.
+        """Optimize the given cost function starting from initial parameters.
 
         Parameters:
             cost_fn: The cost function to minimize.
             initial_params: Initial parameters for the optimization.
-            **kwargs: Additional keyword arguments for the optimizer.
+            **kwargs: Additional keyword arguments for the optimizer:
+
+                - maxiter (int, optional): Maximum number of iterations.
+                  Defaults vary by optimizer (e.g., 5 for population-based optimizers,
+                  None for some scipy methods).
+                - rng (np.random.Generator, optional): Random number generator for
+                  stochastic optimizers (PymooOptimizer, MonteCarloOptimizer).
+                  Defaults to a new generator if not provided.
+                - jac (Callable, optional): Gradient/Jacobian function for
+                  gradient-based optimizers (only used by ScipyOptimizer with
+                  L_BFGS_B method). Defaults to None.
 
         Returns:
             Optimized parameters.
@@ -157,9 +167,9 @@ class PymooOptimizer(Optimizer):
         )
 
         while optimizer_obj.has_next():
-            X = pop.get("X")
+            evaluated_X = pop.get("X")
 
-            curr_losses = cost_fn(X)
+            curr_losses = cost_fn(evaluated_X)
             static = StaticProblem(problem, F=curr_losses)
             Evaluator().eval(static, pop)
 
@@ -168,7 +178,7 @@ class PymooOptimizer(Optimizer):
             pop = optimizer_obj.ask()
 
             if callback_fn:
-                callback_fn(OptimizeResult(x=pop.get("X"), fun=curr_losses))
+                callback_fn(OptimizeResult(x=evaluated_X, fun=curr_losses))
 
         result = optimizer_obj.result()
 
@@ -207,7 +217,7 @@ class ScipyOptimizer(Optimizer):
         self.method = method
 
     @property
-    def n_param_sets(self):
+    def n_param_sets(self) -> int:
         """
         Get the number of parameter sets used by this optimizer.
 
@@ -220,9 +230,9 @@ class ScipyOptimizer(Optimizer):
         self,
         cost_fn: Callable[[np.ndarray], float],
         initial_params: np.ndarray,
-        callback_fn: Callable | None = None,
+        callback_fn: Callable[[OptimizeResult], Any] | None = None,
         **kwargs,
-    ):
+    ) -> OptimizeResult:
         """
         Run the scipy optimization algorithm.
 
@@ -231,8 +241,8 @@ class ScipyOptimizer(Optimizer):
                 parameters and return a scalar cost value.
             initial_params (np.ndarray): Initial parameter values as a 1D or 2D array.
                 If 2D with shape (1, n_params), it will be squeezed to 1D.
-            callback_fn (Callable, optional): Function called after each iteration.
-                Defaults to None.
+            callback_fn (Callable, optional): Function called after each iteration
+                with an `OptimizeResult` object. Defaults to None.
             **kwargs: Additional keyword arguments:
                 - maxiter (int): Maximum number of iterations
                 - jac (Callable): Gradient function (only used for L-BFGS-B)
@@ -242,6 +252,24 @@ class ScipyOptimizer(Optimizer):
         """
         max_iterations = kwargs.pop("maxiter", None)
 
+        # If a callback is provided, we wrap the cost function and callback
+        # to ensure the data passed to the callback has a consistent shape.
+        if callback_fn:
+
+            def callback_wrapper(intermediate_result: OptimizeResult):
+                # Create a dictionary from the intermediate result to preserve all of its keys.
+                result_dict = dict(intermediate_result)
+
+                # Overwrite 'x' and 'fun' to ensure they have consistent dimensions.
+                result_dict["x"] = np.atleast_2d(intermediate_result.x)
+                result_dict["fun"] = np.atleast_1d(intermediate_result.fun)
+
+                # Create a new OptimizeResult and pass it to the user's callback.
+                return callback_fn(OptimizeResult(**result_dict))
+
+        else:
+            callback_wrapper = None
+
         if max_iterations is None or self.method == ScipyMethod.COBYLA:
             # COBYLA perceive maxiter as maxfev so we need
             # to use the callback fn for counting instead.
@@ -263,7 +291,7 @@ class ScipyOptimizer(Optimizer):
             jac=(
                 kwargs.pop("jac", None) if self.method == ScipyMethod.L_BFGS_B else None
             ),
-            callback=callback_fn,
+            callback=callback_wrapper,
             options={"maxiter": maxiter},
         )
 
@@ -277,45 +305,65 @@ class MonteCarloOptimizer(Optimizer):
     decreasing variance. This implements a simple but effective evolutionary strategy.
     """
 
-    def __init__(self, n_param_sets: int = 10, n_best_sets: int = 3):
+    def __init__(
+        self,
+        population_size: int = 10,
+        n_best_sets: int = 3,
+        keep_best_params: bool = False,
+    ):
         """
         Initialize a Monte Carlo optimizer.
 
         Args:
-            n_param_sets (int, optional): Total number of parameter sets to evaluate
-                per iteration. Defaults to 10.
+            population_size (int, optional): Size of the population for the algorithm.
+                Defaults to 10.
             n_best_sets (int, optional): Number of top-performing parameter sets to
                 use as seeds for the next generation. Defaults to 3.
+            keep_best_params (bool, optional): If True, includes the best parameter sets
+                directly in the new population. If False, generates all new parameters
+                by sampling around the best ones. Defaults to False.
 
         Raises:
-            ValueError: If n_best_sets is greater than n_param_sets.
+            ValueError: If n_best_sets is greater than population_size.
+            ValueError: If keep_best_params is True and n_best_sets equals population_size.
         """
         super().__init__()
 
-        if n_best_sets > n_param_sets:
-            raise ValueError("n_best_sets must be less than or equal to n_param_sets.")
+        if n_best_sets > population_size:
+            raise ValueError(
+                "n_best_sets must be less than or equal to population_size."
+            )
 
-        self._n_param_sets = n_param_sets
-        self._n_best_sets = n_best_sets
+        if keep_best_params and n_best_sets == population_size:
+            raise ValueError(
+                "If keep_best_params is True, n_best_sets must be less than population_size."
+            )
 
-        # Calculate how many times each of the best sets should be repeated
-        samples_per_best = self.n_param_sets // self.n_best_sets
-        remainder = self.n_param_sets % self.n_best_sets
-        self._repeat_counts = np.full(self.n_best_sets, samples_per_best)
-        self._repeat_counts[:remainder] += 1
+        self._population_size = population_size
+        self._n_best_sets = n_best_sets
+        self._keep_best_params = keep_best_params
 
     @property
-    def n_param_sets(self):
+    def population_size(self) -> int:
+        """
+        Get the size of the population.
+
+        Returns:
+            int: Size of the population.
         """
-        Get the number of parameter sets evaluated per iteration.
+        return self._population_size
+
+    @property
+    def n_param_sets(self) -> int:
+        """Number of parameter sets (population size), per the Optimizer interface.
 
         Returns:
-            int: Total number of parameter sets.
+            int: The population size.
         """
-        return self._n_param_sets
+        return self._population_size
 
     @property
-    def n_best_sets(self):
+    def n_best_sets(self) -> int:
         """
         Get the number of best parameter sets used for seeding the next generation.
 
@@ -324,6 +372,16 @@ class MonteCarloOptimizer(Optimizer):
         """
         return self._n_best_sets
 
+    @property
+    def keep_best_params(self) -> bool:
+        """
+        Get whether the best parameters are kept in the new population.
+
+        Returns:
+            bool: True if best parameters are included in new population, False otherwise.
+        """
+        return self._keep_best_params
+
     def _compute_new_parameters(
         self,
         params: np.ndarray,
@@ -338,34 +396,57 @@ class MonteCarloOptimizer(Optimizer):
         # 1. Select the best parameter sets from the current population
         best_params = params[best_indices]
 
-        # 2. Prepare the means for sampling by repeating each best parameter set
-        # according to its assigned count
-        new_means = np.repeat(best_params, self._repeat_counts, axis=0)
+        # 2. Determine how many new samples to generate and calculate repeat counts
+        if self._keep_best_params:
+            n_new_samples = self._population_size - self._n_best_sets
+            # Calculate repeat counts for new samples only
+            samples_per_best = n_new_samples // self._n_best_sets
+            remainder = n_new_samples % self._n_best_sets
+        else:
+            # Calculate repeat counts for the entire population
+            samples_per_best = self._population_size // self._n_best_sets
+            remainder = self._population_size % self._n_best_sets
+
+        repeat_counts = np.full(self._n_best_sets, samples_per_best)
+        repeat_counts[:remainder] += 1
+
+        # 3. Prepare the means for sampling by repeating each best parameter set
+        new_means = np.repeat(best_params, repeat_counts, axis=0)
 
-        # 3. Define the standard deviation (scale), which shrinks over iterations
+        # 4. Define the standard deviation (scale), which shrinks over iterations
         scale = 1.0 / (2.0 * (curr_iteration + 1.0))
 
-        # 4. Generate all new parameters in a single vectorized call
+        # 5. Generate new parameters by sampling around the best ones
         new_params = rng.normal(loc=new_means, scale=scale)
 
-        # Apply periodic boundary conditions
-        return new_params % (2 * np.pi)
+        # 6. Apply periodic boundary conditions
+        new_params = new_params % (2 * np.pi)
+
+        # 7. Conditionally combine with best params if keeping them
+        if self._keep_best_params:
+            return np.vstack([best_params, new_params])
+        else:
+            return new_params
 
     def optimize(
         self,
         cost_fn: Callable[[np.ndarray], float],
         initial_params: np.ndarray,
-        callback_fn: Callable[[OptimizeResult], float | np.ndarray] | None = None,
+        callback_fn: Callable[[OptimizeResult], Any] | None = None,
         **kwargs,
     ) -> OptimizeResult:
-        """
-        Perform Monte Carlo optimization on the cost function.
+        """Perform Monte Carlo optimization on the cost function.
 
         Parameters:
             cost_fn: The cost function to minimize.
             initial_params: Initial parameters for the optimization.
             callback_fn: Optional callback function to monitor progress.
-            **kwargs: Additional keyword arguments for the optimizer.
+            **kwargs: Additional keyword arguments:
+
+                - maxiter (int, optional): Maximum number of iterations. Defaults to 5.
+                - rng (np.random.Generator, optional): Random number generator for
+                  parameter sampling. Defaults to a new generator if not provided.
+
         Returns:
             Optimized parameters.
         """
@@ -373,30 +454,32 @@ class MonteCarloOptimizer(Optimizer):
         max_iterations = kwargs.pop("maxiter", 5)
 
         population = np.copy(initial_params)
+        evaluated_population = population
 
         for curr_iter in range(max_iterations):
             # Evaluate the entire population once
             losses = cost_fn(population)
+            evaluated_population = population
 
-            # Find the indices of the best-performing parameter sets (only once)
+            if callback_fn:
+                callback_fn(OptimizeResult(x=evaluated_population, fun=losses))
+
+            # Find the indices of the best-performing parameter sets
             best_indices = np.argpartition(losses, self.n_best_sets - 1)[
                 : self.n_best_sets
             ]
 
-            if callback_fn:
-                callback_fn(
-                    OptimizeResult(x=population[best_indices], fun=losses[best_indices])
-                )
-
             # Generate the next generation of parameters
             population = self._compute_new_parameters(
-                population, curr_iter, best_indices, rng
+                evaluated_population, curr_iter, best_indices, rng
             )
 
+        # Note: 'losses' here are from the last successfully evaluated population
         best_idx = np.argmin(losses)
+
         # Return the best results from the LAST EVALUATED population
         return OptimizeResult(
-            x=population[best_idx],
+            x=evaluated_population[best_idx],
             fun=losses[best_idx],
             nit=max_iterations,
         )