Update logging of y-values and parameters in Spectra class

Author: icarosadero

pyduino/optimization/nelder_mead.py

@@ -1,17 +1,23 @@
 import numpy as np
-from typing import Callable, List
+import warnings
+from typing import Union, List, Iterable
 from . import Optimizer, linmap
 
 def sigmoid(x):
     return 1 / (1 + np.exp(-x))
 def logit(x, inf=100):
     return np.where(x==0, -inf, np.where(x==1, inf, np.log(x) - np.log(1-x)))
+def softmax(x):
+    assert x.ndim == 2, "Softmax only implemented for 2D arrays"
+    return np.exp(x) / np.sum(np.exp(x), axis=1, keepdims=True)
 
-class NelderMead(Optimizer):
+class NelderMeadBounded(Optimizer):
     def __init__(self, population_size: int, ranges: List[float], rng_seed: int = 0, hypercube_radius = 100):
         """
         This Nelder-Mead algorithm assumes that the optimization function 'f' is to be minimized and time independent.
 
+        The space is assumed to be a hypercube, and the algorithm will map the hypercube to the unit hypercube, bounded by a sigmoid function.
+
         ranges list of pairs:
             Maxima and minima that each parameter can cover.
             Example: [(0, 10), (0, 10)] for two parameters ranging from 0 to 10.
@@ -59,6 +65,126 @@ def ask_oracle(self, X: np.ndarray) -> np.ndarray:
     def init_oracle(self):
         return self.ask_oracle(self.view_g())
 
+    def step(self):
+        """
+        This function performs a single step of the Nelder-Mead algorithm.
+        """
+        
+        # Sort the population by the value of the oracle
+        y = self.ask_oracle(self.view_g())
+        idx = np.argsort(y)
+        self.population = self.population[idx]
+
+        # Compute the centroid of the population
+        centroid = self.population[:-1].mean(axis=0)
+
+        # Reflect the worst point through the centroid
+        reflected = centroid + (centroid - self.population[-1])
+
+        # Evaluate the reflected point
+        
+        y_reflected = self.ask_oracle(self.view(reflected.reshape(1,-1)))
+
+        # If the reflected point is better than the second worst, but not better than the best, then expand
+        
+        if y_reflected < y[-2] and y_reflected > y[0]:
+            expanded = centroid + (reflected - centroid)
+            y_expanded = self.ask_oracle(self.view(expanded.reshape(1,-1)))
+            if y_expanded < y_reflected:
+                self.population[-1] = expanded
+            else:
+                self.population[-1] = reflected
+        # If the reflected point is better than the best, then expand
+        elif y_reflected < y[0]:
+            expanded = centroid + 2 * (reflected - centroid)
+            y_expanded = self.ask_oracle(self.view(expanded.reshape(1,-1)))
+            if y_expanded < y_reflected:
+                self.population[-1] = expanded
+            else:
+                self.population[-1] = reflected
+        # If the reflected point is worse than the second worst, then contract
+        elif y_reflected > y[-2]:
+            contracted = centroid + 0.5 * (self.population[-1] - centroid)
+            y_contracted = self.ask_oracle(self.view(contracted.reshape(1,-1)))
+            if y_contracted < y[-1]:
+                self.population[-1] = contracted
+            else:
+                for i in range(1, len(self.population)):
+                    self.population[i] = 0.5 * (self.population[i] + self.population[0])
+        # If the reflected point is worse than the worst, then shrink
+        elif y_reflected > y[-1]:
+            for i in range(1, len(self.population)):
+                self.population[i] = 0.5 * (self.population[i] + self.population[0])
+
+        self.y = y.copy()
+        return self.view_g()
+
+class NelderMeadConstant(Optimizer):
+    def __init__(self, population_size: int, ranges: Union[int, Iterable], rng_seed: int = 0, energy=10):
+        """
+        This Nelder-Mead algorithm assumes that the optimization function 'f' is to be minimized and time independent.
+
+        The parameters are constrained so that their sum always add up to `energy` though a softmax function.
+
+        Parameters:
+        - population_size (int): The size of the population.
+        - ranges (int or list): The number of parameters or a list of ranges for each parameter.
+        - rng_seed (int): The seed for the random number generator.
+        - energy (float): The energy parameter used in the optimization.
+
+        Note: The ranges parameter is merely a placeholder. The ranges are set to (0, energy) for all parameters.
+        """
+        self.population_size = population_size
+        if isinstance(ranges, Iterable):
+            self.ranges = np.array([(0, energy) for _ in range(len(ranges))])
+            warnings.warn("While using Nelder-MeadConstant, the ranges are set to (0, energy) for all parameters. The parameter `ranges` is merely a placeholder.")
+        elif isinstance(ranges, int):
+            self.ranges = np.array([(0, energy) for _ in range(ranges)])
+        self.rng_seed = np.random.default_rng(rng_seed)
+
+        # Initialize the population (random position and initial momentum)
+        self.population = self.rng_seed.random((self.population_size, len(self.ranges)))
+
+        # Derived attributes
+        self.energy = energy
+        self.e_summation = 1
+        # Initialize y as vector of nans
+        self.y = np.full(self.population_size, np.nan)
+    
+    def forward(self, x):
+        self.e_summation = np.sum(np.exp(x), axis=1, keepdims=True)
+        return self.energy * softmax(x)
+    
+    def backward(self, x):
+        """
+        Softmax is not injective. This is a pseudo-inverse.
+        """
+        return np.log(self.e_summation * x/self.energy)
+
+    def view(self, x):
+        """
+        Maps the input from the domain to the codomain.
+        """
+        return self.forward(x)
+    
+    def view_g(self):
+        """
+        Maps the input from the domain to the codomain.
+        """
+        return self.forward(self.population)
+
+    def inverse_view(self, x):
+        """
+        Maps the input from the codomain to the domain.
+        """
+        return self.backward(x)
+    
+    def ask_oracle(self, X: np.ndarray) -> np.ndarray:
+        return super().ask_oracle()
+
+    def init_oracle(self):
+        return self.ask_oracle(self.view_g())
+
     def step(self):
         """
         This function performs a single step of the Nelder-Mead algorithm.

pyduino/spectra.py

@@ -1,4 +1,4 @@
-from pyduino.optimization.nelder_mead import NelderMead
+from pyduino.optimization.nelder_mead import NelderMeadBounded
 from pyduino.pyduino2 import ReactorManager, chunks, PATHS
 import numpy as np
 import json
@@ -75,7 +75,7 @@ def parse_dados(X,param):
     """
     return np.array(list(map(seval,map(lambda x: x[1].get(param,0),sorted(X.items(),key=lambda x: x[0])))))
 
-class Spectra(RangeParser,ReactorManager,NelderMead):
+class Spectra(RangeParser,ReactorManager,NelderMeadBounded):
     def __init__(self,ranges,density_param,maximize=True,log_name=None,reset_density=False,**kwargs):
         """
         Args:
@@ -99,7 +99,7 @@ def __init__(self,ranges,density_param,maximize=True,log_name=None,reset_density
         self.irradiance = PATHS.SYSTEM_PARAMETERS['irradiance']
 
         ReactorManager.__init__(self)
-        NelderMead.__init__(
+        NelderMeadBounded.__init__(
             self,
             population_size=len(self.parameters),
             ranges=self.ranges_as_list(),

tests/test_optim.py

@@ -8,10 +8,24 @@
 dotenv.load_dotenv(dotenv_path=".env.local")
 
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-from pyduino.optimization.nelder_mead import NelderMead
+from pyduino.optimization.nelder_mead import NelderMeadBounded, NelderMeadConstant, softmax
 
-class TestOptim:
-    optimizer = NelderMead(population_size=10, ranges=[(-10, 10)] * 4)
+def test_softmax():
+    # Test case 1: 2D array
+    x = np.array([[1, 2, 3], [4, 5, 6]])
+    expected_output = np.array([[0.09003057, 0.24472847, 0.66524096], [0.09003057, 0.24472847, 0.66524096]])
+    assert np.allclose(softmax(x), expected_output)
+
+    # Test case 2: 1D array (should raise an assertion error)
+    x = np.array([1, 2, 3])
+    try:
+        softmax(x)
+        assert False, "Expected an assertion error"
+    except AssertionError:
+        pass
+
+class TestOptimNelderMeadBounded:
+    optimizer = NelderMeadBounded(population_size=10, ranges=[(-10, 10)] * 4)
 
     def test_logic(self):
         x = (2*np.random.random((10, 4))-1)*100
@@ -58,4 +72,57 @@ def ask(self, X):
         assert self.optimizer.view(self.optimizer.population).shape == (10, 4)
         assert self.optimizer.view_g().shape == (10, 4)
         assert self.optimizer.inverse_view(self.optimizer.view(self.optimizer.population)).shape == (10, 4)
-        assert np.isclose(self.optimizer.y.min(), 0, atol=1e-4), f"Oracle: {self.optimizer.y.min()}"
+        assert np.isclose(self.optimizer.y.min(), 0, atol=1e-4), f"Oracle: {self.optimizer.y.min()}"
+
+class TestOptimNelderMeadConstant:
+    energy = 10
+    optimizer = NelderMeadConstant(population_size=10, ranges=[(0, 10)] * 4, energy=energy)
+
+    def test_logic(self):
+        x = softmax((2*np.random.random((10, 4))-1)*100)
+        assert np.allclose(x.sum(axis=1), 1), f"Sum of x: {x.sum(axis=1)}"
+
+        print("X")
+        print(x)
+        assert np.allclose(self.energy, self.optimizer.forward(x).sum(axis=1)), f"Sum of forward: {self.optimizer.forward(x).sum(axis=1)}"
+        print("Forward -> Backward")
+        print(self.optimizer.backward(self.optimizer.forward(x)))
+
+        assert np.allclose(x, self.optimizer.backward(self.optimizer.forward(x)))
+        assert np.allclose(x, self.optimizer.inverse_view(self.optimizer.view(x)))
+
+        assert np.all(self.optimizer.backward(np.zeros((10,4))) < 1)
+
+    def test_nelder_mead_parabola(self):
+        class Oracle:
+            def ask(self, X):
+                return (X**2).sum(axis=1)
+        oracle = Oracle()
+        self.optimizer.ask_oracle = oracle.ask
+
+        # Set initial population to oracle
+        self.optimizer.init_oracle()
+
+        for i in range(100):
+            self.optimizer.step()
+        assert self.optimizer.view(self.optimizer.population).shape == (10, 4)
+        assert self.optimizer.view_g().shape == (10, 4)
+        assert self.optimizer.inverse_view(self.optimizer.view(self.optimizer.population)).shape == (10, 4)
+        assert np.allclose(self.energy, self.optimizer.view_g().sum(axis=1), atol=1e-4), f"Energy wasn't conserved: {self.optimizer.view_g().sum(axis=1)}"
+
+    def test_nelder_mead_rosenbrock(self):
+        class Oracle:
+            def ask(self, X):
+                return ((1 - X[:, :-1])**2).sum(axis=1) + 100 * ((X[:, 1:] - X[:, :-1]**2)**2).sum(axis=1)
+        oracle = Oracle()
+        self.optimizer.ask_oracle = oracle.ask
+
+        # Set initial population to oracle
+        self.optimizer.init_oracle()
+
+        for i in range(1000):
+            self.optimizer.step()
+        assert self.optimizer.view(self.optimizer.population).shape == (10, 4)
+        assert self.optimizer.view_g().shape == (10, 4)
+        assert self.optimizer.inverse_view(self.optimizer.view(self.optimizer.population)).shape == (10, 4)
+        assert np.allclose(self.energy, self.optimizer.view_g().sum(axis=1), atol=1e-4), f"Energy wasn't conserved: {self.optimizer.view_g().sum(axis=1)}"