evolution-strategies/functions.py

import numpy as np

from dataclasses import dataclass
from typing import Callable, Tuple


@dataclass
class Function:
    xlim: Tuple[float, float]
    ylim: Tuple[float, float]
    minimum: Tuple[float, float]
    eval: Callable

    def grid(self, x_dim=256, y_dim=256):
        x = np.linspace(self.xlim[0], self.xlim[1], x_dim)
        y = np.linspace(self.ylim[0], self.ylim[1], y_dim)
        X, Y = np.meshgrid(x, y)
        return x, y, self.eval(X, Y)


class Rastrigin(Function):
    def __init__(self, A: int = 10):
        super().__init__(
                xlim=(-6.12, 6.12), 
                ylim=(-6.12, 6.12), 
                minimum=(0, 0),
                eval=lambda x, y: self.A * 2 + \
                        (x**2 - self.A * np.cos(2 * np.pi * x)) + \
                        (y**2 - self.A * np.cos(2 * np.pi * y)))
        self.A = A


class Schaffer2(Function):
    def __init__(self):
        super().__init__(
                xlim=(-50, 50),
                ylim=(-50, 50),
                minimum=(0, 0),
                eval=lambda x, y: 0.5 + \
                        (np.sin(x**2 - y**2)**2 - 0.5) / \
                        ((1 + 0.001 * (x**2 + y**2))**2))

class Sphere(Function):
    def __init__(self):
        super().__init__(
                xlim=(-5, 5),
                ylim=(-5, 5),
                minimum=(0, 0),
                eval=lambda x, y: x**2 + y**2)