🧪 Add smoke tests to CI pipeline

.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"

.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

.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

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",
 ]

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,17 +40,17 @@ 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:
         return np.array([self.fitness_fn(ind) for ind in population])
 
     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)

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):