Add AlgoTune tasks for OpenEvolve

Author: richardcsuwandi

examples/algotune/AlgoTuneTasks/__init__.py

@@ -0,0 +1,154 @@
+import hmac
+import logging
+import os
+from typing import Any
+
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+
+from AlgoTuneTasks.base import register_task, Task
+
+
+# Define standard key sizes and nonce size
+AES_KEY_SIZES = [16, 24, 32]  # AES-128, AES-192, AES-256
+GCM_NONCE_SIZE = 12  # Recommended nonce size for GCM
+GCM_TAG_SIZE = 16  # Standard tag size for GCM
+
+
+@register_task("aes_gcm_encryption")
+class AesGcmEncryption(Task):
+    """
+    AesGcmEncryption Task:
+
+    Encrypt plaintext using AES-GCM from the `cryptography` library.
+    The difficulty scales with the size of the plaintext.
+    """
+
+    DEFAULT_KEY_SIZE = 16  # Default to AES-128
+    DEFAULT_PLAINTEXT_MULTIPLIER = 1024  # Bytes of plaintext per unit of n
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def generate_problem(self, n: int, random_seed: int = 1) -> dict[str, Any]:
+        """
+        Generate inputs for the AES-GCM encryption task.
+
+        Args:
+            n (int): Scaling parameter, determines the plaintext size.
+            random_seed (int): Seed for reproducibility (used for key, nonce, plaintext).
+
+        Returns:
+            dict: A dictionary containing the problem parameters (key, nonce, plaintext, associated_data).
+        """
+        # Use n to scale the plaintext size
+        plaintext_size = max(1, n * self.DEFAULT_PLAINTEXT_MULTIPLIER)  # Ensure at least 1 byte
+
+        # Use fixed key size and nonce size for simplicity in scaling
+        key_size = self.DEFAULT_KEY_SIZE
+        nonce_size = GCM_NONCE_SIZE
+
+        logging.debug(
+            f"Generating AES-GCM problem with n={n} (plaintext_size={plaintext_size}), "
+            f"key_size={key_size}, nonce_size={nonce_size}, random_seed={random_seed}"
+        )
+
+        if key_size not in AES_KEY_SIZES:
+            raise ValueError(f"Unsupported key size: {key_size}. Must be one of {AES_KEY_SIZES}.")
+
+        # Use os.urandom for cryptographic randomness.
+        # While random_seed is an input, for crypto tasks, using truly random
+        # keys/nonces per problem instance is generally better practice,
+        # even if it makes the benchmark non-deterministic w.r.t random_seed.
+        # If strict determinism based on random_seed is required, we'd need
+        # to use a seeded PRNG like random.getrandbits, but that's less secure.
+        key = os.urandom(key_size)
+        nonce = os.urandom(nonce_size)
+        plaintext = os.urandom(plaintext_size)
+        # Generate some associated data for testing, can be empty
+        associated_data = os.urandom(32) if n % 2 == 0 else b""  # Example: include AAD for even n
+
+        return {
+            "key": key,
+            "nonce": nonce,
+            "plaintext": plaintext,
+            "associated_data": associated_data,
+        }
+
+    def solve(self, problem: dict[str, Any]) -> dict[str, bytes]:
+        """
+        Encrypt the plaintext using AES-GCM from the `cryptography` library.
+
+        Args:
+            problem (dict): The problem dictionary generated by `generate_problem`.
+
+        Returns:
+            dict: A dictionary containing 'ciphertext' and 'tag'.
+        """
+        key = problem["key"]
+        nonce = problem["nonce"]
+        plaintext = problem["plaintext"]
+        associated_data = problem["associated_data"]
+
+        try:
+            # Validate key size based on provided key length
+            if len(key) not in AES_KEY_SIZES:
+                raise ValueError(f"Invalid key size: {len(key)}. Must be one of {AES_KEY_SIZES}.")
+
+            aesgcm = AESGCM(key)
+            ciphertext = aesgcm.encrypt(nonce, plaintext, associated_data)
+
+            # GCM ciphertext includes the tag appended at the end. We need to split them.
+            # The tag length is fixed (usually 16 bytes / 128 bits).
+            if len(ciphertext) < GCM_TAG_SIZE:
+                raise ValueError("Encrypted output is shorter than the expected tag size.")
+
+            actual_ciphertext = ciphertext[:-GCM_TAG_SIZE]
+            tag = ciphertext[-GCM_TAG_SIZE:]
+
+            return {"ciphertext": actual_ciphertext, "tag": tag}
+
+        except Exception as e:
+            logging.error(f"Error during AES-GCM encryption in solve: {e}")
+            raise  # Re-raise exception
+
+    def is_solution(self, problem: dict[str, Any], solution: dict[str, bytes] | Any) -> bool:
+        """
+        Verify the provided solution by comparing its ciphertext and tag
+        against the result obtained from calling the task's own solve() method.
+
+        Args:
+            problem (dict): The problem dictionary.
+            solution (dict): The proposed solution dictionary with 'ciphertext' and 'tag'.
+
+        Returns:
+            bool: True if the solution matches the result from self.solve().
+        """
+        if not isinstance(solution, dict) or "ciphertext" not in solution or "tag" not in solution:
+            logging.error(
+                f"Invalid solution format. Expected dict with 'ciphertext' and 'tag'. Got: {type(solution)}"
+            )
+            return False
+
+        try:
+            # Get the correct result by calling the solve method
+            reference_result = self.solve(problem)
+            reference_ciphertext = reference_result["ciphertext"]
+            reference_tag = reference_result["tag"]
+        except Exception as e:
+            # If solve itself fails, we cannot verify the solution
+            logging.error(f"Failed to generate reference solution in is_solution: {e}")
+            return False
+
+        solution_ciphertext = solution["ciphertext"]
+        solution_tag = solution["tag"]
+
+        # Ensure types are bytes before comparison
+        if not isinstance(solution_ciphertext, bytes) or not isinstance(solution_tag, bytes):
+            logging.error("Solution 'ciphertext' or 'tag' is not bytes.")
+            return False
+
+        # Constant-time comparison for security
+        ciphertext_match = hmac.compare_digest(reference_ciphertext, solution_ciphertext)
+        tag_match = hmac.compare_digest(reference_tag, solution_tag)
+
+        return ciphertext_match and tag_match

examples/algotune/AlgoTuneTasks/aes_gcm_encryption/aes_gcm_encryption.py

@@ -0,0 +1,34 @@
+AesGcmEncryption Task:
+
+Task Description:
+Encrypt a given plaintext using AES (Advanced Encryption Standard) in GCM (Galois/Counter Mode) with a provided key, nonce (Initialization Vector - IV), and optional associated data (AAD). This task uses `cryptography.hazmat.primitives.ciphers.aead.AESGCM`. AES-GCM provides both confidentiality and data authenticity. The primary computational cost scales with the length of the plaintext.
+
+Input:
+A dictionary with keys:
+  - "key": A bytes object representing the AES key (e.g., 16, 24, or 32 bytes for AES-128, AES-192, AES-256).
+  - "nonce": A bytes object representing the nonce (IV). For GCM, 12 bytes is commonly recommended. Must be unique for each encryption with the same key.
+  - "plaintext": A bytes object representing the data to encrypt. The size of this data will scale with the problem size 'n'.
+  - "associated_data": A bytes object representing additional data to authenticate but not encrypt (optional, can be `None` or empty bytes `b''`).
+
+Example input:
+{
+    "key": b'\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b\x0c\x0d\x0e\x0f\x10', # 16 bytes key for AES-128
+    "nonce": b'\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0a\x0b', # 12 bytes nonce
+    "plaintext": b'data to encrypt' * 100, # Example scaled plaintext
+    "associated_data": b'metadata'
+}
+
+Output:
+A dictionary containing:
+  - "ciphertext": A bytes object representing the encrypted data.
+  - "tag": A bytes object representing the GCM authentication tag (typically 16 bytes).
+
+Example output:
+# The actual output depends on the exact inputs (key, nonce, plaintext, aad).
+# This is a conceptual placeholder.
+{
+    "ciphertext": b'\xencrypted...\data',
+    "tag": b'\xauthentication-tag' # 16 bytes
+}
+
+Category: cryptography

examples/algotune/AlgoTuneTasks/aes_gcm_encryption/description.txt

@@ -0,0 +1,186 @@
+import logging
+import random
+from typing import Any
+
+import numpy as np
+import scipy.ndimage
+
+from AlgoTuneTasks.base import register_task, Task
+
+
+@register_task("affine_transform_2d")
+class AffineTransform2D(Task):
+    def __init__(self, **kwargs):
+        """
+        Initialize the AffineTransform2D task.
+
+        Performs an affine transformation on a 2D image using scipy.ndimage.affine_transform.
+        The transformation is defined by a 2x3 matrix. Uses cubic spline interpolation
+        (order=3) and 'constant' mode for handling boundaries.
+        """
+        super().__init__(**kwargs)
+        self.order = 3
+        self.mode = "constant"  # Or 'nearest', 'reflect', 'mirror', 'wrap'
+
+    def generate_problem(self, n: int, random_seed: int = 1) -> dict[str, Any]:
+        """
+        Generates a 2D affine transformation problem.
+
+        Creates a random 2D array (image) of size n x n and a 2x3 affine
+        transformation matrix.
+
+        :param n: The side dimension of the square input image.
+        :param random_seed: Seed for reproducibility.
+        :return: A dictionary representing the problem with keys:
+                 "image": Input image as a numpy array (n x n).
+                 "matrix": The 2x3 affine transformation matrix (numpy array).
+        """
+        logging.debug(
+            f"Generating 2D Affine Transform problem with n={n}, random_seed={random_seed}"
+        )
+        random.seed(random_seed)
+        np.random.seed(random_seed)
+
+        # Generate random n x n image
+        image = np.random.rand(n, n) * 255.0  # Example: 0-255 range image
+
+        # Generate a random affine matrix (e.g., combining rotation, scale, shear, translation)
+        angle = np.random.uniform(-np.pi / 6, np.pi / 6)  # +/- 30 degrees
+        scale = np.random.uniform(0.8, 1.2, size=2)
+        shear = np.random.uniform(-0.2, 0.2)
+        translation = np.random.uniform(-n * 0.1, n * 0.1, size=2)  # Translate up to 10%
+
+        # Rotation matrix
+        cos_a, sin_a = np.cos(angle), np.sin(angle)
+        R = np.array([[cos_a, -sin_a], [sin_a, cos_a]])
+
+        # Scale matrix
+        S = np.array([[scale[0], 0], [0, scale[1]]])
+
+        # Shear matrix
+        H = np.array([[1, shear], [0, 1]])
+
+        # Combine transformations (excluding translation for the 2x2 part)
+        M_2x2 = R @ S @ H
+
+        # Create the 2x3 matrix [ M_2x2 | translation ]
+        matrix = np.hstack((M_2x2, translation[:, np.newaxis]))
+
+        problem = {"image": image, "matrix": matrix}
+        logging.debug(f"Generated 2D Affine Transform problem for image shape ({n},{n})")
+        return problem
+
+    def solve(self, problem: dict[str, Any]) -> dict[str, list[list[float]]]:
+        """
+        Solves the 2D affine transformation problem using scipy.ndimage.affine_transform.
+
+        :param problem: A dictionary representing the problem.
+        :return: A dictionary with key "transformed_image":
+                 "transformed_image": The transformed image as a list of lists.
+        """
+        image = problem["image"]
+        matrix = problem["matrix"]
+
+        # Perform affine transformation
+        try:
+            # output_shape can be specified, default is same as input
+            transformed_image = scipy.ndimage.affine_transform(
+                image, matrix, order=self.order, mode=self.mode
+            )
+        except Exception as e:
+            logging.error(f"scipy.ndimage.affine_transform failed: {e}")
+            # Return an empty list to indicate failure? Adjust based on benchmark policy.
+            return {"transformed_image": []}
+
+        solution = {"transformed_image": transformed_image}
+        return solution
+
+    def is_solution(self, problem: dict[str, Any], solution: dict[str, list[list[float]]]) -> bool:
+        """
+        Check if the provided affine transformation solution is valid.
+
+        Checks structure, dimensions, finite values, and numerical closeness to
+        the reference scipy.ndimage.affine_transform output.
+
+        :param problem: The problem definition dictionary.
+        :param solution: The proposed solution dictionary.
+        :return: True if the solution is valid, False otherwise.
+        """
+        if not all(k in problem for k in ["image", "matrix"]):
+            logging.error("Problem dictionary missing 'image' or 'matrix'.")
+            return False
+        image = problem["image"]
+        matrix = problem["matrix"]
+
+        if not isinstance(solution, dict) or "transformed_image" not in solution:
+            logging.error("Solution format invalid: missing 'transformed_image' key.")
+            return False
+
+        proposed_list = solution["transformed_image"]
+
+        # Handle potential failure case from solve()
+        if proposed_list == []:
+            logging.warning("Proposed solution is empty list (potential failure).")
+            # Check if reference solver also fails/produces empty-like result
+            try:
+                ref_output = scipy.ndimage.affine_transform(
+                    image, matrix, order=self.order, mode=self.mode
+                )
+                if ref_output.size == 0:  # Check if reference is also effectively empty
+                    logging.info(
+                        "Reference solver also produced empty result. Accepting empty solution."
+                    )
+                    return True
+                else:
+                    logging.error("Reference solver succeeded, but proposed solution was empty.")
+                    return False
+            except Exception:
+                logging.info("Reference solver also failed. Accepting empty solution.")
+                return True  # Both failed, likely invalid input
+
+        if not isinstance(proposed_list, list):
+            logging.error("'transformed_image' is not a list.")
+            return False
+
+        try:
+            proposed_array = np.asarray(proposed_list, dtype=float)
+        except ValueError:
+            logging.error("Could not convert 'transformed_image' list to numpy float array.")
+            return False
+
+        # Expected output shape is usually same as input for affine_transform unless specified
+        if proposed_array.shape != image.shape:
+            logging.error(f"Output shape {proposed_array.shape} != input shape {image.shape}.")
+            # This might be acceptable if output_shape was used, but base case expects same shape.
+            # Adjust if the task allows different output shapes.
+            return False  # Assuming same shape output for now
+
+        if not np.all(np.isfinite(proposed_array)):
+            logging.error("Proposed 'transformed_image' contains non-finite values.")
+            return False
+
+        # Re-compute reference solution
+        try:
+            ref_array = scipy.ndimage.affine_transform(
+                image, matrix, order=self.order, mode=self.mode
+            )
+        except Exception as e:
+            logging.error(f"Error computing reference solution: {e}")
+            return False  # Cannot verify if reference fails
+
+        # Compare results
+        rtol = 1e-5
+        atol = 1e-7  # Slightly tighter atol for image data often in 0-255 range
+        is_close = np.allclose(proposed_array, ref_array, rtol=rtol, atol=atol)
+
+        if not is_close:
+            abs_diff = np.abs(proposed_array - ref_array)
+            max_abs_err = np.max(abs_diff) if abs_diff.size > 0 else 0
+            logging.error(
+                f"Solution verification failed: Output mismatch. "
+                f"Max absolute error: {max_abs_err:.3f} (rtol={rtol}, atol={atol})"
+            )
+            return False
+
+        logging.debug("Solution verification successful.")
+        return True