Test functions package

examples/viz_optimizers.py

@@ -1,35 +1,16 @@
-import math
+import matplotlib.pyplot as plt
 import numpy as np
-import torch_optimizer as optim
 import torch
 from hyperopt import fmin, tpe, hp
-import matplotlib.pyplot as plt
-
 
-plt.style.use('seaborn-white')
-
-
-def rosenbrock(tensor):
-    # https://en.wikipedia.org/wiki/Test_functions_for_optimization
-    x, y = tensor
-    return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2
+import torch_optimizer as optim
+from test_functions import Rastrigin, Rosenbrock
 
+plt.style.use("seaborn-white")
 
-def rastrigin(tensor, lib=torch):
-    # https://en.wikipedia.org/wiki/Test_functions_for_optimization
-    x, y = tensor
-    A = 10
-    f = (
-        A * 2
-        + (x ** 2 - A * lib.cos(x * math.pi * 2))
-        + (y ** 2 - A * lib.cos(y * math.pi * 2))
-    )
-    return f
 
-
-def execute_steps(
-    func, initial_state, optimizer_class, optimizer_config, num_iter=500
-):
+def execute_steps(func, initial_state, optimizer_class, optimizer_config, num_iter=500):
+    """ Execute one steps of the optimizer """
     x = torch.Tensor(initial_state).requires_grad_(True)
     optimizer = optimizer_class([x], **optimizer_config)
     steps = []
@@ -45,116 +26,77 @@ def execute_steps(
     return steps
 
 
-def objective_rastrigin(params):
-    lr = params['lr']
-    optimizer_class = params['optimizer_class']
-    initial_state = (-2.0, 3.5)
-    minimum = (0, 0)
-    optimizer_config = dict(lr=lr)
-    num_iter = 100
-    steps = execute_steps(
-        rastrigin, initial_state, optimizer_class, optimizer_config, num_iter
-    )
-    return (steps[0][-1] - minimum[0]) ** 2 + (steps[1][-1] - minimum[1]) ** 2
-
-
-def objective_rosenbrok(params):
-    lr = params['lr']
-    optimizer_class = params['optimizer_class']
-    minimum = (1.0, 1.0)
-    initial_state = (-2.0, 2.0)
+def objective(params, func):
+    """ Execute objective """
+    lr = params["lr"]
+    optimizer_class = params["optimizer_class"]
+    minimum = func.minimum
+    initial_state = func.initial_state
     optimizer_config = dict(lr=lr)
     num_iter = 100
     steps = execute_steps(
-        rosenbrock, initial_state, optimizer_class, optimizer_config, num_iter
+        func, initial_state, optimizer_class, optimizer_config, num_iter
     )
     return (steps[0][-1] - minimum[0]) ** 2 + (steps[1][-1] - minimum[1]) ** 2
 
 
-def plot_rastrigin(grad_iter, optimizer_name, lr):
-    x = np.linspace(-4.5, 4.5, 250)
-    y = np.linspace(-4.5, 4.5, 250)
-    minimum = (0, 0)
+def plot(func, grad_iter, optimizer_name, lr):
+    """ Plot result of a simulation """
+    x = torch.linspace(func.x_domain[0], func.x_domain[1], func.num_pt)
+    y = torch.linspace(func.y_domain[0], func.y_domain[1], func.num_pt)
+    minimum = func.minimum
 
-    X, Y = np.meshgrid(x, y)
-    Z = rastrigin([X, Y], lib=np)
+    X, Y = torch.meshgrid(x, y)
+    Z = func([X, Y])
 
     iter_x, iter_y = grad_iter[0, :], grad_iter[1, :]
 
     fig = plt.figure(figsize=(8, 8))
 
     ax = fig.add_subplot(1, 1, 1)
-    ax.contour(X, Y, Z, 20, cmap='jet')
-    ax.plot(iter_x, iter_y, color='r', marker='x')
-    ax.set_title(
-        f'Rastrigin func: {optimizer_name} with '
-        f'{len(iter_x)} iterations, lr={lr:.6}'
-    )
-    plt.plot(*minimum, 'gD')
-    plt.plot(iter_x[-1], iter_y[-1], 'rD')
-    plt.savefig(f'docs/rastrigin_{optimizer_name}.png')
-
-
-def plot_rosenbrok(grad_iter, optimizer_name, lr):
-    x = np.linspace(-2, 2, 250)
-    y = np.linspace(-1, 3, 250)
-    minimum = (1.0, 1.0)
-
-    X, Y = np.meshgrid(x, y)
-    Z = rosenbrock([X, Y])
-
-    iter_x, iter_y = grad_iter[0, :], grad_iter[1, :]
-
-    fig = plt.figure(figsize=(8, 8))
-
-    ax = fig.add_subplot(1, 1, 1)
-    ax.contour(X, Y, Z, 90, cmap='jet')
-    ax.plot(iter_x, iter_y, color='r', marker='x')
+    ax.contour(X, Y, Z, func.levels, cmap="jet")
+    ax.plot(iter_x, iter_y, color="r", marker="x")
 
+    func_name = func.__name__()
     ax.set_title(
-        f'Rosenbrock func: {optimizer_name} with {len(iter_x)} '
-        f'iterations, lr={lr:.6}'
+        f"{func_name} func: {optimizer_name} with {len(iter_x)} "
+        f"iterations, lr={lr:.6}"
     )
-    plt.plot(*minimum, 'gD')
-    plt.plot(iter_x[-1], iter_y[-1], 'rD')
-    plt.savefig(f'docs/rosenbrock_{optimizer_name}.png')
+    plt.plot(*minimum, "gD")
+    plt.plot(iter_x[-1], iter_y[-1], "rD")
+    plt.savefig(f"docs/{func_name}_{optimizer_name}.png")
 
 
-def execute_experiments(
-    optimizers, objective, func, plot_func, initial_state, seed=1
-):
+def execute_experiments(optimizers, func, seed=1):
+    """ Execute simulation on a list of optimizers using a test function """
     seed = seed
     for item in optimizers:
         optimizer_class, lr_low, lr_hi = item
         space = {
-            'optimizer_class': hp.choice('optimizer_class', [optimizer_class]),
-            'lr': hp.loguniform('lr', lr_low, lr_hi),
+            "optimizer_class": hp.choice("optimizer_class", [optimizer_class]),
+            "lr": hp.loguniform("lr", lr_low, lr_hi),
         }
         best = fmin(
-            fn=objective,
+            fn=lambda x: objective(x, func),
             space=space,
             algo=tpe.suggest,
             max_evals=200,
             rstate=np.random.RandomState(seed),
         )
-        print(best['lr'], optimizer_class)
+        print(best["lr"], optimizer_class)
 
         steps = execute_steps(
-            func,
-            initial_state,
-            optimizer_class,
-            {'lr': best['lr']},
-            num_iter=500,
+            func, func.initial_state, optimizer_class, {"lr": best["lr"]}, num_iter=500,
         )
-        plot_func(steps, optimizer_class.__name__, best['lr'])
+        plot(func, steps, optimizer_class.__name__, best["lr"])
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     # python examples/viz_optimizers.py
 
     # Each optimizer has tweaked search space to produce better plots and
     # help to converge on better lr faster.
-    optimizers = [
+    list_optimizers = [
         # Adam based
         (optim.AdaBound, -8, 0.3),
         (optim.AdaMod, -8, 0.2),
@@ -168,18 +110,9 @@ def execute_experiments(
         (optim.SGDW, -8, -1.5),
         (optim.PID, -8, -1.0),
     ]
-    execute_experiments(
-        optimizers,
-        objective_rastrigin,
-        rastrigin,
-        plot_rastrigin,
-        (-2.0, 3.5),
-    )
 
-    execute_experiments(
-        optimizers,
-        objective_rosenbrok,
-        rosenbrock,
-        plot_rosenbrok,
-        (-2.0, 2.0),
-    )
+    for test_func in [Rastrigin(), Rosenbrock()]:
+        print(f"Test function {test_func.__name__()}")
+        execute_experiments(
+            list_optimizers, test_func,
+        )

test_functions/__init__.py

@@ -0,0 +1,2 @@
+from .rosenbrock import Rosenbrock
+from .rastrigin import Rastrigin

test_functions/rastrigin.py

@@ -0,0 +1,47 @@
+import torch
+import math
+
+from .test_function import TestFunction
+
+
+class Rastrigin(TestFunction):
+    r"""Rastrigin test function.
+
+    Example:
+    >>> import test_functions
+    >>> rastrigin = test_functions.Rastrigin()
+    >>> x = torch.linspace(
+            rastrigin.x_domain[0],
+            rastrigin.x_domain[1],
+            rastrigin.num_pt
+        )
+    >>> y = torch.linspace(
+            rastrigin.y_domain[0],
+            rastrigin.y_domain[1],
+            rastrigin.num_pt
+        )
+    >>> Y, X = torch.meshgrid(x, y)
+    >>> Z = rastrigin([X, Y])
+
+    __ https://en.wikipedia.org/wiki/Test_functions_for_optimization
+    __ https://en.wikipedia.org/wiki/Rastrigin_function
+    """
+
+    def __init__(self):
+        super(Rastrigin, self).__init__(
+            x_domain=(-4.5, 4.5),
+            y_domain=(-4.5, 4.5),
+            minimum=(0, 0),
+            initial_state=(-2.0, 3.5),
+            levels=20,
+        )
+
+    def __call__(self, tensor, lib=torch):
+        x, y = tensor
+        A = 10
+        f = (
+            A * 2
+            + (x ** 2 - A * lib.cos(x * math.pi * 2))
+            + (y ** 2 - A * lib.cos(y * math.pi * 2))
+        )
+        return f

test_functions/rosenbrock.py

@@ -0,0 +1,40 @@
+import torch
+
+from .test_function import TestFunction
+
+
+class Rosenbrock(TestFunction):
+    r"""Rosenbrock test function.
+
+    Example:
+    >>> import test_functions
+    >>> rosenbrock = test_functions.Rosenbrock()
+    >>> x = torch.linspace(
+            rosenbrock.x_domain[0],
+            rosenbrock.x_domain[1],
+            rosenbrock.num_pt
+        )
+    >>> y = torch.linspace(
+            rosenbrock.y_domain[0],
+            rosenbrock.y_domain[1],
+            rosenbrock.num_pt
+        )
+    >>> Y, X = torch.meshgrid(x, y)
+    >>> Z = rosenbrock([X, Y])
+
+    __ https://en.wikipedia.org/wiki/Test_functions_for_optimization
+    __ https://en.wikipedia.org/wiki/Rosenbrock_function
+    """
+
+    def __init__(self):
+        super(Rosenbrock, self).__init__(
+            x_domain=(-2.0, 2.0),
+            y_domain=(-1.0, 3.0),
+            minimum=(1.0, 1.0),
+            initial_state=(-2.0, 2.0),
+            levels=90,
+        )
+
+    def __call__(self, tensor, lib=torch):
+        x, y = tensor
+        return (1 - x) ** 2 + 100 * (y - x ** 2) ** 2

test_functions/test_function.py

@@ -0,0 +1,58 @@
+class TestFunction:
+    r"""Super class to implement a test function.
+
+    Arguments:
+        x_domain: iterable of int for the x search domain
+        y_domain: iterable of int for the y search domain
+        minimum: global minimum (x0, y0) of the
+            function for the given domain
+        initial_state: starting point
+        num_pt: number of point in the search domain
+        levels: determines the number and positions of the contour lines / regions.
+    """
+
+    def __init__(
+        self,
+        x_domain: iter,
+        y_domain: iter,
+        minimum: iter,
+        initial_state: iter,
+        num_pt: int = 250,
+        levels: int = 50,
+    ):
+        self._x_domain = x_domain
+        self._y_domain = y_domain
+        self._minimum = minimum
+        self._initial_state = initial_state
+        self._num_pt = num_pt
+        self._levels = levels
+
+    def __call__(self, tensor, lib):
+        raise NotImplementedError
+
+    def __name__(self):
+        return self.__class__.__name__
+
+    @property
+    def x_domain(self):
+        return self._x_domain
+
+    @property
+    def y_domain(self):
+        return self._y_domain
+
+    @property
+    def num_pt(self):
+        return self._num_pt
+
+    @property
+    def minimum(self):
+        return self._minimum
+
+    @property
+    def initial_state(self):
+        return self._initial_state
+
+    @property
+    def levels(self):
+        return self._levels