diff --git a/.claw/claw.yaml b/.claw/claw.yaml new file mode 100644 index 0000000..d5ce500 --- /dev/null +++ b/.claw/claw.yaml @@ -0,0 +1,47 @@ +# The executable name of the LLM CLI tool that exists in your PATH. +# Change this to "gemini", "ollama", or any other tool you use. +receiver_type: Generic +llm_command: "codex" + +# (Optional) The argument pattern for passing the prompt to the LLM. +# The "{{prompt}}" placeholder will be replaced with the final rendered prompt. +# The default is just "{{prompt}}". +# +# Example for gemini-cli: +prompt_arg_template: "{{prompt}}" + +# Context Management 2.0 Configuration +# These settings control how claw processes files passed via --context parameter + +# Maximum file size in KB that can be included as context (default: 1024 = 1 MB) +max_file_size_kb: 3072 + +# Maximum number of files per directory when scanning (default: 50) +max_files_per_directory: 50 + +# How to handle errors during context processing (default: flexible) +# Options: +# strict: Fail immediately on any error +# flexible: Collect all errors and prompt user for approval before proceeding +# ignore: Log warnings but continue processing valid files +error_handling_mode: flexible + +# Directories to exclude when scanning for context files +excluded_directories: + - ".git" + - "node_modules" + - "target" + - ".venv" + - "__pycache__" + +# File extensions to exclude when scanning for context files +excluded_extensions: + - "exe" + - "bin" + - "so" + - "dylib" + - "dll" + - "o" + - "a" + - "lock" + - "pdf" diff --git a/.claw/goals/pr-notes/prompt.yaml b/.claw/goals/pr-notes/prompt.yaml new file mode 100644 index 0000000..a48a77c --- /dev/null +++ b/.claw/goals/pr-notes/prompt.yaml @@ -0,0 +1,49 @@ +name: "Pull Request Notes" +description: "Create pull request notes based on the changes made in the repo" + +context_scripts: + branch_diff: "git diff master \":(exclude)*.lock\"" + +prompt: | + Based on the following git diff, write PR notes following this exact format. Do NOT run any commands yourself - only analyze the provided diff. + + {{ Context.branch_diff }} + + CRITICAL: Output raw markdown text only. Do not render or format the markdown. I need the literal markdown characters (*, #, `, etc.) visible in plain text format. + + Generate PR notes using this exact structure: + + ## Required Format + + **Title (first line)** + - Maximum 56 characters including emoji + - Start with an appropriate emoji followed by a space + - Use sentence case, no period at the end + - Example: šŸš€ Add user authentication system + + **# What does this PR do?** (heading with single # character) + - Write 2-3 complete sentences in paragraph form + - Summarize the overall accomplishment and impact + - No bullet points in this section + + **# Details** (heading with single # character) + - Use markdown bullet points starting with asterisk and space: * + - Each bullet point must be a complete thought + - Maximum 160 characters per bullet point + - List specific changes, additions, or modifications + + **# Highlights** (heading with single # character, optional section) + - Only include if there are important code changes worth showcasing + - Use proper markdown code blocks with triple backticks and language identifiers + - Format: ```language on first line, code content, closing ``` on last line + - Add brief context before each code snippet if needed + + ## Output Requirements + - Output ONLY the raw markdown text + - Start with the emoji title on the first line + - Use literal # characters for headings + - Use literal * characters for bullet points + - Use literal ``` characters for code blocks + - Include blank lines between sections + - Do not add any meta-commentary or explanations around the markdown + - The output should be ready to copy-paste directly into a PR description diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..09e9d17 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,31 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + +jobs: + tests: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.11"] + + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v4 + + - name: Set up Python + run: uv python install ${{ matrix.python-version }} + + - name: Install dependencies + run: uv sync --dev + + - name: Run tests + env: + MPLBACKEND: Agg + run: uv run pytest \ No newline at end of file diff --git a/examples/two_peaks.py b/examples/two_peaks.py deleted file mode 100644 index d9dfd65..0000000 --- a/examples/two_peaks.py +++ /dev/null @@ -1,68 +0,0 @@ -"""Two-peak optimization example with visualization.""" - -import matplotlib.pyplot as plt -import numpy as np -from numpy.typing import NDArray - -from devol import DiffusionConfig, DiffusionEvolution - - -def two_peaks_function(x: NDArray) -> float: - """Two Gaussian peaks at (1,1) and (-1,-1).""" - peak1 = np.exp(-np.sum((x - np.array([1.0, 1.0])) ** 2) / 0.1) - peak2 = np.exp(-np.sum((x - np.array([-1.0, -1.0])) ** 2) / 0.1) - return (peak1 + peak2) / 2 - - -def run_two_peaks() -> None: - config = DiffusionConfig( - population_size=512, - num_steps=50, - param_dim=2, - sigma_m=1.0, - seed=42, - ) - - algo = DiffusionEvolution(config, two_peaks_function) - final_population = algo.run() - - fitness_values = np.array([two_peaks_function(ind) for ind in final_population]) - top_indices = np.argsort(fitness_values)[-20:] - top_solutions = final_population[top_indices] - - print("Top 20 solutions:") - for i, (sol, fit) in enumerate(zip(top_solutions, fitness_values[top_indices])): - print(f"{i + 1:2d}. x={sol[0]:6.3f}, y={sol[1]:6.3f}, fitness={fit:.6f}") - - peak1 = np.array([1.0, 1.0]) - peak2 = np.array([-1.0, -1.0]) - near_peak1 = np.sum(np.linalg.norm(top_solutions - peak1, axis=1) < 0.5) - near_peak2 = np.sum(np.linalg.norm(top_solutions - peak2, axis=1) < 0.5) - - print("\nDiversity analysis:") - print(f" Solutions near peak (1,1): {near_peak1}") - print(f" Solutions near peak (-1,-1): {near_peak2}") - - x = np.linspace(-2, 2, 100) - y = np.linspace(-2, 2, 100) - X, Y = np.meshgrid(x, y) - Z = np.zeros_like(X) - - for i in range(X.shape[0]): - for j in range(X.shape[1]): - Z[i, j] = two_peaks_function(np.array([X[i, j], Y[i, j]])) - - plt.figure(figsize=(10, 8)) - plt.contourf(X, Y, Z, levels=20, cmap="viridis", alpha=0.6) - plt.colorbar(label="Fitness") - plt.scatter(final_population[:, 0], final_population[:, 1], c="red", s=10, alpha=0.5) - plt.scatter(top_solutions[:, 0], top_solutions[:, 1], c="white", s=50, edgecolors="black") - plt.xlabel("x") - plt.ylabel("y") - plt.title("Two Peaks: Final Population Distribution") - plt.savefig("two_peaks_result.png", dpi=150, bbox_inches="tight") - print("\nVisualization saved to two_peaks_result.png") - - -if __name__ == "__main__": - run_two_peaks() diff --git a/pyproject.toml b/pyproject.toml index 8cbe48e..0ea401f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,6 +10,7 @@ dependencies = [ "pydantic>=2.0.0", "pydantic-settings>=2.0.0", "pydantic-yaml>=1.6.0", + "pytest>=8.4.2", "torch>=2.9.0", "torchvision>=0.24.0", ] diff --git a/src/devol/algorithm.py b/src/devol/algorithm.py index e953702..df53bab 100644 --- a/src/devol/algorithm.py +++ b/src/devol/algorithm.py @@ -8,7 +8,7 @@ from devol.config import DiffusionConfig from devol.distance import create_distance_computer from devol.evolution import compute_epsilon_hat, estimate_x0, evolution_step -from devol.fitness import create_fitness_mapper +from devol.fitness import create_fitness_mapper, create_fitness_normalizer from devol.schedules import create_alpha_schedule, create_sigma_schedule @@ -18,9 +18,7 @@ def __init__(self, config: DiffusionConfig, fitness_fn: Callable[[NDArray], floa self.fitness_fn = fitness_fn self.rng = np.random.default_rng(config.seed) - self.alpha = create_alpha_schedule( - config.schedule.type.value, config.num_steps, config.schedule.epsilon - ) + self.alpha = create_alpha_schedule(config.schedule.type.value, config.num_steps, config.schedule.epsilon) self.sigma = create_sigma_schedule(self.alpha, config.sigma_m) self.distance_computer = create_distance_computer( @@ -30,8 +28,9 @@ def __init__(self, config: DiffusionConfig, fitness_fn: Callable[[NDArray], floa config.seed, ) + self.fitness_normalizer = create_fitness_normalizer(config.fitness.normalize) self.fitness_mapper = create_fitness_mapper( - config.fitness.mapping.value, + config.fitness.mapping, config.fitness.temperature, ) @@ -41,9 +40,7 @@ def __init__(self, config: DiffusionConfig, fitness_fn: Callable[[NDArray], floa # Make it a docstring # Explain how the noising op is shifting the original pdf to a ~N(0, 1) def initialize_population(self) -> NDArray: # TODO: maybe make it of type Population - self.population = self.rng.standard_normal( - (self.config.population_size, self.config.param_dim) - ) + self.population = self.rng.standard_normal((self.config.population_size, self.config.param_dim)) return self.population def evaluate_fitness(self, population: NDArray) -> NDArray: @@ -51,7 +48,9 @@ def evaluate_fitness(self, population: NDArray) -> NDArray: def step(self, timestamp: int, population: NDArray) -> NDArray: fitness = self.evaluate_fitness(population) - fitness_weights = self.fitness_mapper(fitness) + normalized_fitness = self.fitness_normalizer(fitness) + + fitness_weights = self.fitness_mapper(normalized_fitness) alpha_t = self.alpha[timestamp] alpha_t_minus_1 = self.alpha[timestamp - 1] @@ -63,14 +62,14 @@ def step(self, timestamp: int, population: NDArray) -> NDArray: x_t = population[i] x_hat_0 = estimate_x0(x_t, population, fitness_weights, alpha_t, self.distance_computer) epsilon_hat = compute_epsilon_hat(x_t, x_hat_0, alpha_t) - new_population[i] = evolution_step( - x_t, x_hat_0, epsilon_hat, alpha_t, alpha_t_minus_1, sigma_t, self.rng - ) + new_population[i] = evolution_step(x_t, x_hat_0, epsilon_hat, alpha_t, alpha_t_minus_1, sigma_t, self.rng) return new_population - def run(self) -> NDArray: - population = self.initialize_population() + def run(self, initial_population: NDArray | None) -> NDArray: + population = initial_population + if population is None: + population = self.initialize_population() for timestamp in range(self.config.num_steps, 0, -1): population = self.step(timestamp, population) diff --git a/src/devol/config.py b/src/devol/config.py index ee3960c..9c418dc 100644 --- a/src/devol/config.py +++ b/src/devol/config.py @@ -1,35 +1,47 @@ """Configuration models for Diffusion Evolution.""" -from enum import Enum +from enum import StrEnum from pydantic import BaseModel, Field, field_validator -class ScheduleType(str, Enum): +class ScheduleType(StrEnum): LINEAR = "linear" COSINE = "cosine" DDPM = "ddpm" -class FitnessMapping(str, Enum): +class FitnessMapping(StrEnum): + DIRECT = "direct" + IDENTITY = "identity" + ENERGY = "energy" EXPONENTIAL = "exponential" RANK = "rank" -class DistanceType(str, Enum): +class DistanceType(StrEnum): EUCLIDEAN = "euclidean" LATENT = "latent" COSINE = "cosine" +class NormalType(StrEnum): + MAX_SCALE = "max_scale" + MIN_MAX = "min_max" + Z_SCORE = "z_score" + SUM_TO_ONE = "sum_to_one" + IDENTITY = "identity" + + class ScheduleConfig(BaseModel, frozen=True): type: ScheduleType = ScheduleType.COSINE epsilon: float = Field(default=1e-4, gt=0, lt=1) class FitnessConfig(BaseModel, frozen=True): - mapping: FitnessMapping = FitnessMapping.EXPONENTIAL + mapping: FitnessMapping = FitnessMapping.DIRECT temperature: float = Field(default=1.0, gt=0) + normalize: NormalType = NormalType.MIN_MAX class DistanceConfig(BaseModel, frozen=True): diff --git a/src/devol/fitness.py b/src/devol/fitness.py index 2e75a2f..981db3e 100644 --- a/src/devol/fitness.py +++ b/src/devol/fitness.py @@ -5,6 +5,8 @@ import numpy as np from numpy.typing import NDArray +from devol.config import FitnessMapping, NormalType + class FitnessMapper(Protocol): def __call__(self, fitness: NDArray) -> NDArray: @@ -12,6 +14,54 @@ def __call__(self, fitness: NDArray) -> NDArray: ... +class FitnessNormalizer(Protocol): + def __call__(self, fitness: NDArray) -> NDArray: + """Normalize fitness values to be within acceptable range or gaussian.""" + ... + + +class Identity: + """Identity fitness mapping function.""" + + def __init__(self, l2_factor=0.0): + self.l2_factor = l2_factor + + def l2(self, x): + return np.linalg.norm(x, axis=-1) ** 2 + + def forward(self, x): + return x + + def __call__(self, fitness: NDArray) -> NDArray: + return self.forward(fitness) * np.exp(-1.0 * self.l2(fitness) * self.l2_factor) + + +class DirectMapper: + def __call__(self, fitness: NDArray) -> NDArray: + return fitness + + +class Energy(Identity): + """Fitness mapping function that treats the fitness as energy. + + Args: + temperature: float, the temperature of the system. + + Returns: + p: torch.Tensor, the probability of the fitness. Compute by exp(-x / temperature). + """ + + def __init__(self, temperature=1.0, l2_factor=0.0): + super().__init__(l2_factor=l2_factor) + self.temperature = temperature + + def forward(self, x): + power = -x / self.temperature + power = power - power.max() + 5 # avoid overflow + p = np.exp(power) + return p + + class ExponentialMapper: def __init__(self, temperature: float = 1.0): self.temperature = temperature @@ -29,15 +79,88 @@ def __call__(self, fitness: NDArray) -> NDArray: def create_fitness_mapper( - mapping_type: str, + mapping_type: FitnessMapping, temperature: float = 1.0, ) -> FitnessMapper: mapper: FitnessMapper - if mapping_type == "exponential": + if mapping_type == FitnessMapping.EXPONENTIAL: mapper = ExponentialMapper(temperature) - elif mapping_type == "rank": + elif mapping_type == FitnessMapping.RANK: mapper = RankMapper() + elif mapping_type == FitnessMapping.DIRECT: + mapper = DirectMapper() + elif mapping_type == FitnessMapping.ENERGY: + return Energy(temperature=temperature) + elif mapping_type == FitnessMapping.IDENTITY: + return Identity() else: raise ValueError(f"Unknown fitness mapping: {mapping_type}") return mapper + + +class MaxScaleNormalizer: + def __init__(self, epsilon: float = 1e-12): + self.epsilon = epsilon + + def __call__(self, fitness: NDArray) -> NDArray: + max_abs = np.max(np.abs(fitness)) + if max_abs < self.epsilon: + return np.zeros_like(fitness) + return fitness / max_abs + + +class MinMaxNormalizer: + def __init__(self, epsilon: float = 1e-12): + self.epsilon = epsilon + + def __call__(self, fitness: NDArray) -> NDArray: + min_val = np.min(fitness) + max_val = np.max(fitness) + span = max_val - min_val + if span < self.epsilon: + return np.zeros_like(fitness) + return (fitness - min_val) / span + + +class ZScoreNormalizer: + def __init__(self, epsilon: float = 1e-12): + self.epsilon = epsilon + + def __call__(self, fitness: NDArray) -> NDArray: + mean = np.mean(fitness) + std = np.std(fitness) + if std < self.epsilon: + return np.zeros_like(fitness) + return (fitness - mean) / std + + +class SumToOneNormalizer: + def __init__(self, epsilon: float = 1e-12): + self.epsilon = epsilon + + def __call__(self, fitness: NDArray) -> NDArray: + total = np.sum(np.abs(fitness)) + if total < self.epsilon: + return np.zeros_like(fitness) + return fitness / total + + +class IdentityNormalizer: + def __call__(self, fitness: NDArray) -> NDArray: + return fitness + + +def create_fitness_normalizer(normalize_type: NormalType = NormalType.MAX_SCALE) -> FitnessNormalizer: + if normalize_type == NormalType.MAX_SCALE: + return MaxScaleNormalizer() + if normalize_type == NormalType.MIN_MAX: + return MinMaxNormalizer() + if normalize_type == NormalType.Z_SCORE: + return ZScoreNormalizer() + if normalize_type == NormalType.SUM_TO_ONE: + return SumToOneNormalizer() + if normalize_type == NormalType.IDENTITY: + return IdentityNormalizer() + + raise ValueError(f"Unknown normalizer type: {normalize_type}") diff --git a/test_day2.py b/test_day2.py deleted file mode 100644 index 8bdb034..0000000 --- a/test_day2.py +++ /dev/null @@ -1,43 +0,0 @@ -"""Test Day 2 implementations.""" - -import numpy as np - -from devol.distance import create_distance_computer -from devol.fitness import create_fitness_mapper - - -def test_fitness_pipeline(): - fitness = np.array([-1.0, 0.5, 2.0, 1.5]) - print(f"āœ“ Raw fitness: {fitness}") - - mapper = create_fitness_mapper("exponential", temperature=1.0) - weights = mapper(fitness) - print(f"āœ“ Exponential weights (sum={weights.sum():.4f}): {weights}") - - rank_mapper = create_fitness_mapper("rank") - ranks = rank_mapper(fitness) - print(f"āœ“ Rank weights: {ranks}") - - -def test_distance_pipeline(): - population = np.random.randn(10, 5) - x_i = population[0] - - euclidean = create_distance_computer("euclidean", param_dim=5) - distances = euclidean.compute_distances(x_i, population) - print(f"āœ“ Euclidean distances (first 3): {distances[:3]}") - - latent = create_distance_computer("latent", param_dim=5, latent_dim=2, seed=42) - distances_latent = latent.compute_distances(x_i, population) - print(f"āœ“ Latent distances (first 3): {distances_latent[:3]}") - - cosine = create_distance_computer("cosine", param_dim=5) - distances_cosine = cosine.compute_distances(x_i, population) - print(f"āœ“ Cosine distances (first 3): {distances_cosine[:3]}") - - -if __name__ == "__main__": - test_fitness_pipeline() - print() - test_distance_pipeline() - print("\nāœ“ Day 2 complete!") diff --git a/test_day3.py b/test_day3.py deleted file mode 100644 index 105c932..0000000 --- a/test_day3.py +++ /dev/null @@ -1,92 +0,0 @@ -"""Test Day 3 implementations.""" - -import numpy as np - -from devol import DiffusionConfig, DiffusionEvolution - - -def simple_sphere(x: np.ndarray) -> float: - """Simple sphere function: maximize -(x^2 + y^2).""" - return -np.sum(x**2) - - -def two_peaks(x: np.ndarray) -> float: - """Two Gaussian peaks at (1,1) and (-1,-1).""" - peak1 = np.exp(-np.sum((x - np.array([1.0, 1.0])) ** 2) / 0.1) - peak2 = np.exp(-np.sum((x - np.array([-1.0, -1.0])) ** 2) / 0.1) - return (peak1 + peak2) / 2 - - -def test_sphere_optimization(): - print("Testing sphere function optimization...") - config = DiffusionConfig( - population_size=128, - num_steps=25, - param_dim=2, - sigma_m=0.5, - seed=43, - ) - - algo = DiffusionEvolution(config, simple_sphere) - final_population = algo.run() - - best_individual, best_fitness = algo.get_best_individual() - print(f"āœ“ Best individual: {best_individual}") - print(f"āœ“ Best fitness: {best_fitness:.6f}") - print(f"āœ“ Distance from origin: {np.linalg.norm(best_individual):.6f}") - - assert np.linalg.norm(best_individual) < 0.5, "Should converge near origin" - - -def test_two_peaks(): - print("\nTesting two-peak function...") - config = DiffusionConfig( - population_size=256, - num_steps=50, - param_dim=2, - sigma_m=1.0, - seed=43, - ) - - algo = DiffusionEvolution(config, two_peaks) - final_population = algo.run() - - fitness = np.array([two_peaks(ind) for ind in final_population]) - best_indices = np.argsort(fitness)[-10:] - top_solutions = final_population[best_indices] - - print("āœ“ Top 10 solutions found:") - for i, sol in enumerate(top_solutions[-5:]): - print(f" {i + 1}. {sol} (fitness: {fitness[best_indices[-5 + i]]:.6f})") - - peak1_count = np.sum(np.linalg.norm(top_solutions - [1, 1], axis=1) < 0.5) - peak2_count = np.sum(np.linalg.norm(top_solutions - [-1, -1], axis=1) < 0.5) - print(f"āœ“ Solutions near peak (1,1): {peak1_count}") - print(f"āœ“ Solutions near peak (-1,-1): {peak2_count}") - - -def test_latent_space(): - print("\nTesting latent space distance...") - config = DiffusionConfig( - population_size=128, - num_steps=25, - param_dim=10, - distance={"type": "latent", "latent_dim": 3}, - sigma_m=0.8, - seed=42, - ) - - algo = DiffusionEvolution(config, simple_sphere) - final_population = algo.run() - best_individual, best_fitness = algo.get_best_individual() - - print("āœ“ Latent space evolution successful") - print(f"āœ“ Best fitness in 10D: {best_fitness:.6f}") - print(f"āœ“ Norm: {np.linalg.norm(best_individual):.6f}") - - -if __name__ == "__main__": - test_sphere_optimization() - test_two_peaks() - test_latent_space() - print("\nāœ“ Day 3 complete!") diff --git a/test_setup.py b/test_setup.py deleted file mode 100644 index d3cbe0c..0000000 --- a/test_setup.py +++ /dev/null @@ -1,27 +0,0 @@ -"""Quick test to verify setup.""" - -from devol import DiffusionConfig, DistanceType, FitnessMapping, ScheduleType -from devol.schedules import create_alpha_schedule, create_sigma_schedule - - -def test_config(): - config = DiffusionConfig( - param_dim=10, - population_size=128, - num_steps=25, - distance={"type": DistanceType.LATENT, "latent_dim": 5}, - ) - print(f"āœ“ Config created: {config.population_size} individuals, {config.num_steps} steps") - - -def test_schedules(): - alpha = create_alpha_schedule("cosine", 50, 1e-4) - sigma = create_sigma_schedule(alpha, 1.0) - print(f"āœ“ Schedules created: alpha range [{alpha.min():.4f}, {alpha.max():.4f}]") - print(f" sigma range [{sigma.min():.4f}, {sigma.max():.4f}]") - - -if __name__ == "__main__": - test_config() - test_schedules() - print("\nāœ“ Day 1 setup complete!") diff --git a/tests/ci/__init__.py b/tests/ci/__init__.py new file mode 100644 index 0000000..135bd01 --- /dev/null +++ b/tests/ci/__init__.py @@ -0,0 +1 @@ +# CI utilities and smoke tests. diff --git a/tests/ci/n_peaks.py b/tests/ci/n_peaks.py new file mode 100644 index 0000000..a63b252 --- /dev/null +++ b/tests/ci/n_peaks.py @@ -0,0 +1,262 @@ +"""N-peak optimization smoke test. + +Run with: `python tests/ci/n_peaks.py 3 --plot prime_peaks_3.png` +""" + +from __future__ import annotations + +import argparse +import math +from pathlib import Path +from typing import Sequence + +import numpy as np +from numpy.typing import NDArray + +try: # Matplotlib is optional, only needed when plotting. + import matplotlib.pyplot as plt +except ImportError: # pragma: no cover - optional dependency + plt = None + +from devol import DiffusionConfig, DiffusionEvolution +from devol.config import FitnessConfig, FitnessMapping, NormalType + + +def create_peak_positions( + num_peaks: int, + *, + bounds: tuple[float, float] = (-1.0, 1.0), + seed: int | None = 123, +) -> NDArray: + """Place peaks on a lattice inside the provided bounds. + + Peaks are spread evenly across the square area, using interior lattice points + (endpoints are omitted so peaks stay visible on plots). When fewer peaks than + lattice slots are needed, the order is shuffled deterministically with the + provided seed before truncation. + """ + if num_peaks < 1: + raise ValueError("num_peaks must be > 0") + + low, high = bounds + + grid_size = math.ceil(math.sqrt(num_peaks)) + lattice_coords = np.linspace(low, high, grid_size + 2)[1:-1] + grid_x, grid_y = np.meshgrid(lattice_coords, lattice_coords) + lattice_points = np.stack([grid_x.ravel(), grid_y.ravel()], axis=1) + + rng = np.random.default_rng(seed) + rng.shuffle(lattice_points) + + return lattice_points[:num_peaks] + + +def make_multi_peak_function(peaks: NDArray, width: float = 0.02): + """Return a callable fitness function for the provided peak coordinates.""" + + def _fitness(x: NDArray) -> float: + diffs = x - peaks + dist_sq = np.sum(diffs * diffs, axis=1) + contributions = np.exp(-dist_sq / width) + return float(np.mean(contributions)) + + return _fitness + + +def verify_convergence(population: NDArray, peaks: NDArray, tolerance: float) -> list[bool]: + """Check that each peak has at least one individual within the tolerance.""" + flags: list[bool] = [] + for peak in peaks: + dists = np.linalg.norm(population - peak, axis=1) + flags.append(np.min(dists) <= tolerance) + return flags + + +def average_min_peak_distance(population: NDArray, peaks: NDArray) -> float: + """Average distance from each individual to its closest peak.""" + distances = np.linalg.norm(population[:, None, :] - peaks[None, :, :], axis=2) + nearest = np.min(distances, axis=1) + return float(np.mean(nearest)) + + +def render_population( + population: NDArray, + peaks: NDArray, + fitness_fn, + out_path: Path | None, + bounds: tuple[float, float] = (-1.2, 1.2), +) -> None: + """Save a contour plot of the multi-peak landscape if matplotlib is ready.""" + if out_path is None: + return + + if plt is None: + print("Matplotlib not available; skipping visualization.") + return + + x = np.linspace(bounds[0], bounds[1], 120) + y = np.linspace(bounds[0], bounds[1], 120) + X, Y = np.meshgrid(x, y) + Z = np.zeros_like(X) + + for i in range(X.shape[0]): + for j in range(X.shape[1]): + Z[i, j] = fitness_fn(np.array([X[i, j], Y[i, j]])) + + plt.figure(figsize=(8, 7)) + plt.contourf(X, Y, Z, levels=30, cmap="viridis", alpha=0.6) + plt.colorbar(label="Fitness") + plt.scatter(population[:, 0], population[:, 1], c="red", s=10, alpha=0.4, label="Population") + plt.scatter(peaks[:, 0], peaks[:, 1], c="white", s=80, edgecolors="black", label="Peaks") + plt.xlabel("x") + plt.ylabel("y") + plt.title(f"Multi-peak population snapshot ({len(peaks)} peaks)") + plt.legend() + plt.savefig(out_path, dpi=150, bbox_inches="tight") + plt.close() + print(f"Visualization saved to {out_path}") + + +def run_multi_peak( + num_peaks: int, + population_size: int = 512, + num_steps: int = 50, + convergence_radius: float = 0.25, + plot_path: Path | None = None, + peak_seed: int | None = 123, + fitness_mapping: FitnessMapping = FitnessMapping.DIRECT, + normalization: NormalType = NormalType.SUM_TO_ONE, + min_improvement_ratio: float = 2.0, + temperature: float = 0.25, +) -> None: + """Run diffusion evolution and assert convergence for each target peak.""" + peaks = create_peak_positions(num_peaks, seed=peak_seed) + fitness_fn = make_multi_peak_function(peaks) + + config = DiffusionConfig( + population_size=population_size, + num_steps=num_steps, + param_dim=2, + sigma_m=1.0, + seed=333, + fitness=FitnessConfig( + mapping=fitness_mapping, + normalize=normalization, + temperature=temperature, + ), + ) + + algo = DiffusionEvolution(config, fitness_fn) + initial_population = algo.initialize_population() + initial_avg_distance = average_min_peak_distance(initial_population, peaks) + + final_population = algo.run(initial_population) + final_avg_distance = average_min_peak_distance(final_population, peaks) + improvement_ratio = math.inf if final_avg_distance == 0 else initial_avg_distance / final_avg_distance + + flags = verify_convergence(final_population, peaks, tolerance=convergence_radius) + + for idx, success in enumerate(flags, start=1): + status = "PASS" if success else "FAIL" + peak_coords = peaks[idx - 1] + print(f"{status} Peak {idx}: ({peak_coords[0]:+.3f}, {peak_coords[1]:+.3f})") + + print( + "Average nearest-peak distance: " + f"start {initial_avg_distance:.3f} -> end {final_avg_distance:.3f}" + ) + print(f"Improvement ratio (start/end): {improvement_ratio:.2f}x") + + render_population(final_population, peaks, fitness_fn, plot_path) + + missing_peaks = [str(i + 1) for i, ok in enumerate(flags) if not ok] + if missing_peaks: + print(f"Note: peaks lacking neighbors within radius {convergence_radius}: {', '.join(missing_peaks)}") + + if improvement_ratio < min_improvement_ratio: + raise RuntimeError( + f"Insufficient improvement: {improvement_ratio:.2f}x (<{min_improvement_ratio:.2f}x target) " + f"[start {initial_avg_distance:.3f}, end {final_avg_distance:.3f}]" + ) + + print( + f"SUCCESS: Achieved {improvement_ratio:.2f}x improvement on average nearest-peak distance " + f"(target: {min_improvement_ratio:.2f}x)." + ) + + +def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Smoke test for diffusion evolution on N peaks.") + parser.add_argument("num_peaks", type=int, help="Number of target peaks (>=1)") + parser.add_argument("-o", "--plot", type=Path, default=None, help="Optional path to save a contour visualization.") + parser.add_argument("-p", "--population", type=int, default=512, help="Population size (default: 512)") + parser.add_argument("-s", "--steps", type=int, default=50, help="Number of denoising steps (default: 50)") + parser.add_argument( + "-r", + "--radius", + type=float, + default=0.25, + help="Distance threshold to count a peak as converged (default: 0.25)", + ) + parser.add_argument( + "-t", + "--temperature", + type=float, + default=0.25, + help="Temperature to use in fitness mapping (default: 0.25)", + ) + parser.add_argument( + "-k", + "--peak-seed", + type=int, + default=123, + help="Seed used to shuffle lattice peak positions (default: 123)", + ) + parser.add_argument( + "-m", + "--mapping", + type=FitnessMapping, + choices=list(FitnessMapping), + default=FitnessMapping.DIRECT, + help="Fitness mapping strategy (default: direct)", + ) + parser.add_argument( + "-n", + "--normalize", + type=NormalType, + choices=list(NormalType), + default=NormalType.SUM_TO_ONE, + help="Fitness normalization strategy (default: sum_to_one)", + ) + parser.add_argument( + "-i", + "--improvement", + type=float, + default=2.0, + help="Required improvement ratio of average nearest-peak distance (start/end). Default: 2.0x", + ) + return parser.parse_args(argv) + + +def main(argv: Sequence[str] | None = None) -> None: + args = parse_args(argv) + try: + run_multi_peak( + num_peaks=args.num_peaks, + population_size=args.population, + num_steps=args.steps, + convergence_radius=args.radius, + plot_path=args.plot, + peak_seed=args.peak_seed, + fitness_mapping=args.mapping, + normalization=args.normalize, + min_improvement_ratio=args.improvement, + temperature=args.temperature, + ) + except RuntimeError as exc: # Ensure CI failure on missed peaks. + print(exc) + raise SystemExit(1) from exc + + +if __name__ == "__main__": + main() diff --git a/tests/test_ci_n_peaks.py b/tests/test_ci_n_peaks.py new file mode 100644 index 0000000..5a9bba1 --- /dev/null +++ b/tests/test_ci_n_peaks.py @@ -0,0 +1,19 @@ +import os + +import pytest + +from tests.ci.n_peaks import FitnessMapping, NormalType, run_multi_peak + + +@pytest.mark.parametrize("num_peaks", [2, 3, 5, 7, 11]) +def test_multi_peak_convergence(num_peaks: int) -> None: + os.environ.setdefault("DEVOL_DEBUG_FITNESS", "0") + run_multi_peak( + num_peaks=num_peaks, + population_size=512, + num_steps=50, + convergence_radius=0.1, + peak_seed=123, + fitness_mapping=FitnessMapping.DIRECT, + normalization=NormalType.MIN_MAX, + ) diff --git a/uv.lock b/uv.lock index dc05de8..24d0569 100644 --- a/uv.lock +++ b/uv.lock @@ -135,6 +135,7 @@ dependencies = [ { name = "pydantic" }, { name = "pydantic-settings" }, { name = "pydantic-yaml" }, + { name = "pytest" }, { name = "torch" }, { name = "torchvision" }, ] @@ -167,6 +168,7 @@ requires-dist = [ { name = "pydantic", specifier = ">=2.0.0" }, { name = "pydantic-settings", specifier = ">=2.0.0" }, { name = "pydantic-yaml", specifier = ">=1.6.0" }, + { name = "pytest", specifier = ">=8.4.2" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, { name = "rich", marker = "extra == 'benchmark'", specifier = ">=13.0.0" }, { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.1.0" }, @@ -1221,6 +1223,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6b/fa/3234f913fe9a6525a7b97c6dad1f51e72b917e6872e051a5e2ffd8b16fbb/ruamel.yaml.clib-0.2.14-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:70eda7703b8126f5e52fcf276e6c0f40b0d314674f896fc58c47b0aef2b9ae83", size = 137970, upload-time = "2025-09-22T19:51:09.472Z" }, { url = "https://files.pythonhosted.org/packages/ef/ec/4edbf17ac2c87fa0845dd366ef8d5852b96eb58fcd65fc1ecf5fe27b4641/ruamel.yaml.clib-0.2.14-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:a0cb71ccc6ef9ce36eecb6272c81afdc2f565950cdcec33ae8e6cd8f7fc86f27", size = 739639, upload-time = "2025-09-22T19:51:10.566Z" }, { url = "https://files.pythonhosted.org/packages/15/18/b0e1fafe59051de9e79cdd431863b03593ecfa8341c110affad7c8121efc/ruamel.yaml.clib-0.2.14-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:e7cb9ad1d525d40f7d87b6df7c0ff916a66bc52cb61b66ac1b2a16d0c1b07640", size = 764456, upload-time = "2025-09-22T19:51:11.736Z" }, + { url = "https://files.pythonhosted.org/packages/e7/cd/150fdb96b8fab27fe08d8a59fe67554568727981806e6bc2677a16081ec7/ruamel_yaml_clib-0.2.14-cp314-cp314-win32.whl", hash = "sha256:9b4104bf43ca0cd4e6f738cb86326a3b2f6eef00f417bd1e7efb7bdffe74c539", size = 102394, upload-time = "2025-11-14T21:57:36.703Z" }, + { url = "https://files.pythonhosted.org/packages/bd/e6/a3fa40084558c7e1dc9546385f22a93949c890a8b2e445b2ba43935f51da/ruamel_yaml_clib-0.2.14-cp314-cp314-win_amd64.whl", hash = "sha256:13997d7d354a9890ea1ec5937a219817464e5cc344805b37671562a401ca3008", size = 122673, upload-time = "2025-11-14T21:57:38.177Z" }, ] [[package]]