diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 388d7d7..753c108 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -6,15 +6,51 @@ on:
   pull_request:
 
 jobs:
-  tests:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
+      - name: Set up Python
+        run: uv python install 3.13
+
+      - name: Install dev dependencies
+        run: uv sync --extra dev
+
+      - name: Ruff check
+        run: uv run ruff check
+
+      - name: Ruff format check
+        run: uv run ruff format --check
+
+  typecheck:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
+      - name: Set up Python
+        run: uv python install 3.13
+
+      - name: Install dev dependencies
+        run: uv sync --extra dev
+
+      - name: Mypy
+        run: uv run mypy src/devol
+
+  test:
     runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
-        python-version: ["3.11"]
-
+        python-version: ["3.11", "3.12", "3.13"]
     steps:
-      - name: Checkout repository
-        uses: actions/checkout@v4
+      - uses: actions/checkout@v4
 
       - name: Install uv
         uses: astral-sh/setup-uv@v4
@@ -22,10 +58,27 @@ jobs:
       - name: Set up Python
         run: uv python install ${{ matrix.python-version }}
 
-      - name: Install dependencies
+      - name: Install dev dependencies
         run: uv sync --extra dev
 
       - name: Run tests
         env:
           MPLBACKEND: Agg
-        run: uv run pytest
\ No newline at end of file
+        run: uv run pytest
+
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
+      - name: Set up Python
+        run: uv python install 3.13
+
+      - name: Build sdist and wheel
+        run: uv build
+
+      - name: Validate package metadata
+        run: uv run --with twine twine check dist/*
diff --git a/benchmark/display.py b/benchmark/display.py
index 1d1f9ce..d9fedf0 100644
--- a/benchmark/display.py
+++ b/benchmark/display.py
@@ -28,9 +28,7 @@ def display_results(results: list[BenchmarkMetrics]) -> None:
     display_parameter_sensitivity(results, console)
 
 
-def display_schedule_summary(
-    results: list[BenchmarkMetrics], console: Console
-) -> None:
+def display_schedule_summary(results: list[BenchmarkMetrics], console: Console) -> None:
     """Display summary statistics for each schedule type."""
     table = Table(title="Schedule Performance Summary")
 
@@ -92,9 +90,7 @@ def display_best_configs(results: list[BenchmarkMetrics], console: Console) -> N
     console.print(table)
 
 
-def display_parameter_sensitivity(
-    results: list[BenchmarkMetrics], console: Console
-) -> None:
+def display_parameter_sensitivity(results: list[BenchmarkMetrics], console: Console) -> None:
     """Display how performance varies with each parameter."""
     params = ["population_size", "num_steps", "param_dim", "sigma_m"]
 
diff --git a/benchmark/metrics.py b/benchmark/metrics.py
index 0e71928..5a7419a 100644
--- a/benchmark/metrics.py
+++ b/benchmark/metrics.py
@@ -64,9 +64,7 @@ def calculate_population_diversity(fitness: NDArray) -> float:
     return float(np.std(fitness))
 
 
-def evaluate_population_fitness(
-    population: NDArray, fitness_fn: FitnessFunction
-) -> NDArray:
+def evaluate_population_fitness(population: NDArray, fitness_fn: FitnessFunction) -> NDArray:
     """Evaluate fitness for entire population.
 
     Args:
diff --git a/benchmark/plots.py b/benchmark/plots.py
index 1a897ff..b9f49c1 100644
--- a/benchmark/plots.py
+++ b/benchmark/plots.py
@@ -65,12 +65,8 @@ def create_heatmaps(results: list[BenchmarkMetrics], output_dir: str) -> None:
             global_max = max(global_max, valid_data.max())
 
     # Second pass: plot with consistent scale
-    for ax, sched, (heatmap_data, pop_sizes, step_counts) in zip(
-        axes, schedules, all_heatmaps
-    ):
-        im = ax.imshow(
-            heatmap_data, cmap="RdYlGn_r", aspect="auto", vmin=global_min, vmax=global_max
-        )
+    for ax, sched, (heatmap_data, pop_sizes, step_counts) in zip(axes, schedules, all_heatmaps):
+        im = ax.imshow(heatmap_data, cmap="RdYlGn_r", aspect="auto", vmin=global_min, vmax=global_max)
         ax.set_title(f"{sched} Schedule")
         ax.set_xlabel("Population Size")
         ax.set_ylabel("Num Steps")
@@ -111,11 +107,7 @@ def create_line_plots(results: list[BenchmarkMetrics], output_dir: str) -> None:
             stds = []
 
             for val in param_vals:
-                distances = [
-                    r.distance_from_origin
-                    for r in sched_results
-                    if getattr(r, param_name) == val
-                ]
+                distances = [r.distance_from_origin for r in sched_results if getattr(r, param_name) == val]
                 means.append(np.mean(distances))
                 stds.append(np.std(distances))
 
@@ -139,9 +131,7 @@ def create_line_plots(results: list[BenchmarkMetrics], output_dir: str) -> None:
             ax.set_xscale("log")
 
         plt.tight_layout()
-        plt.savefig(
-            f"{output_dir}/lineplot_{param_name}.png", dpi=150, bbox_inches="tight"
-        )
+        plt.savefig(f"{output_dir}/lineplot_{param_name}.png", dpi=150, bbox_inches="tight")
         plt.close()
 
 
@@ -155,10 +145,7 @@ def create_boxplots(results: list[BenchmarkMetrics], output_dir: str) -> None:
     fig, ax = plt.subplots(figsize=(10, 6))
 
     schedules = sorted(set(r.schedule_type for r in results))
-    data = [
-        [r.distance_from_origin for r in results if r.schedule_type == sched]
-        for sched in schedules
-    ]
+    data = [[r.distance_from_origin for r in results if r.schedule_type == sched] for sched in schedules]
 
     bp = ax.boxplot(data, labels=schedules, patch_artist=True)
 
diff --git a/benchmark/runner.py b/benchmark/runner.py
index dd92ccc..7bae13c 100644
--- a/benchmark/runner.py
+++ b/benchmark/runner.py
@@ -90,9 +90,7 @@ def __init__(
             n_workers: Number of parallel workers (default: CPU count)
         """
         self.fitness_fn = fitness_fn
-        self.schedule_types = [
-            ScheduleType(s) if isinstance(s, str) else s for s in schedule_types
-        ]
+        self.schedule_types = [ScheduleType(s) if isinstance(s, str) else s for s in schedule_types]
         self.population_sizes = population_sizes
         self.num_steps_list = num_steps_list
         self.param_dims = param_dims
diff --git a/examples/cartpole.py b/examples/cartpole.py
index 956eb6e..e2168c2 100644
--- a/examples/cartpole.py
+++ b/examples/cartpole.py
@@ -92,7 +92,7 @@ def run_cartpole() -> None:
     if test_reward >= 475:
         print("✓ Task solved! (reward >= 475)")
     else:
-        print(f"Task not yet solved. Try more steps or larger population.")
+        print("Task not yet solved. Try more steps or larger population.")
 
 
 if __name__ == "__main__":
diff --git a/examples/cartpole_visual.py b/examples/cartpole_visual.py
index 68af336..d10965d 100644
--- a/examples/cartpole_visual.py
+++ b/examples/cartpole_visual.py
@@ -150,7 +150,7 @@ def run_cartpole(show_progress: bool = False, compare: bool = True) -> None:
         seed=42,
     )
 
-    print(f"\nRunning evolution:")
+    print("\nRunning evolution:")
     print(f"  Population size: {config.population_size}")
     print(f"  Evolution steps: {config.num_steps}")
     #    print(f"  Distance metric: latent space (dim={config.distance.latent_dim})")
@@ -170,7 +170,7 @@ def run_cartpole(show_progress: bool = False, compare: bool = True) -> None:
     final_population = algo.run()
     best_individual, best_fitness = algo.get_best_individual()
 
-    print(f"\n" + "=" * 60)
+    print("\n" + "=" * 60)
     print("EVOLUTION COMPLETE")
     print("=" * 60)
     print(f"Best fitness achieved: {best_fitness:.2f}")
@@ -182,7 +182,7 @@ def run_cartpole(show_progress: bool = False, compare: bool = True) -> None:
     if test_reward >= 475:
         print("✓ Task SOLVED! (reward >= 475)")
     else:
-        print(f"✗ Task not solved yet (need >= 475)")
+        print("✗ Task not solved yet (need >= 475)")
 
     print("\nDemonstrating best evolved controller:")
     demonstrate_controller(best_individual, num_demos=3)
diff --git a/examples/mnist/fitness.py b/examples/mnist/fitness.py
index aa33e29..e483db9 100644
--- a/examples/mnist/fitness.py
+++ b/examples/mnist/fitness.py
@@ -1,11 +1,10 @@
 """GPU-batched fitness evaluation for MNIST."""
 
+import numpy as np
 import torch
-import torch.nn as nn
+from numpy.typing import NDArray
 from torch.utils.data import DataLoader, Subset
 from torchvision import datasets, transforms
-import numpy as np
-from numpy.typing import NDArray
 
 from examples.mnist.serialization import deserialize_model
 
diff --git a/examples/mnist/serialization.py b/examples/mnist/serialization.py
index a8fb1b7..7b896cc 100644
--- a/examples/mnist/serialization.py
+++ b/examples/mnist/serialization.py
@@ -1,8 +1,8 @@
 """Parameter serialization and deserialization for LeNet models."""
 
+import numpy as np
 import torch
 import torch.nn as nn
-import numpy as np
 from numpy.typing import NDArray
 
 from examples.mnist.lenet5 import LeNet5, LeNetMini, create_lenet5, create_lenet_mini
diff --git a/examples/mnist/train.py b/examples/mnist/train.py
index f75d4c1..e002531 100644
--- a/examples/mnist/train.py
+++ b/examples/mnist/train.py
@@ -2,7 +2,6 @@
 
 import json
 from pathlib import Path
-from typing import Any
 
 import numpy as np
 import torch
@@ -12,8 +11,6 @@
 from devol.algorithm import DiffusionEvolution
 from examples.mnist.config import MNISTConfig
 from examples.mnist.fitness import MNISTFitnessEvaluator
-from examples.mnist.lenet5 import count_parameters, create_lenet5, create_lenet_mini
-from examples.mnist.serialization import create_random_individual, serialize_model
 
 
 class MNISTEvolution(DiffusionEvolution):
@@ -99,7 +96,7 @@ def run(self) -> NDArray:
                     break
 
                 if self.patience_counter >= self.mnist_config.early_stopping_patience:
-                    print(f"\n✗ Early stopping triggered (patience exhausted)")
+                    print("\n✗ Early stopping triggered (patience exhausted)")
                     break
 
         self.population = population
diff --git a/pyproject.toml b/pyproject.toml
index 1900610..aef8618 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -75,6 +75,9 @@ packages = ["src/devol", "examples"]
 [tool.ruff]
 line-length = 120
 target-version = "py311"
+# Demo scripts are linted separately; they use argparse globals, numpy meshgrid
+# conventions, and half-finished code that doesn't belong in the published package.
+exclude = ["examples"]
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "N", "W", "UP"]
@@ -84,3 +87,4 @@ strict = true
 warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
+files = ["src/devol"]
diff --git a/src/devol/algorithm.py b/src/devol/algorithm.py
index df53bab..3a16c6b 100644
--- a/src/devol/algorithm.py
+++ b/src/devol/algorithm.py
@@ -3,17 +3,16 @@
 from collections.abc import Callable
 
 import numpy as np
-from numpy.typing import NDArray
 
 from devol.config import DiffusionConfig
-from devol.distance import create_distance_computer
+from devol.distance import FloatArray, create_distance_computer
 from devol.evolution import compute_epsilon_hat, estimate_x0, evolution_step
 from devol.fitness import create_fitness_mapper, create_fitness_normalizer
 from devol.schedules import create_alpha_schedule, create_sigma_schedule
 
 
 class DiffusionEvolution:
-    def __init__(self, config: DiffusionConfig, fitness_fn: Callable[[NDArray], float]) -> None:
+    def __init__(self, config: DiffusionConfig, fitness_fn: Callable[[FloatArray], float]) -> None:
         self.config = config
         self.fitness_fn = fitness_fn
         self.rng = np.random.default_rng(config.seed)
@@ -34,19 +33,16 @@ def __init__(self, config: DiffusionConfig, fitness_fn: Callable[[NDArray], floa
             config.fitness.temperature,
         )
 
-        self.population: NDArray | None = None
+        self.population: FloatArray | None = None
 
-    # TODO: Is this init optimal? do we want to abstract it?
-    # 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
+    def initialize_population(self) -> FloatArray:
         self.population = self.rng.standard_normal((self.config.population_size, self.config.param_dim))
         return self.population
 
-    def evaluate_fitness(self, population: NDArray) -> NDArray:
+    def evaluate_fitness(self, population: FloatArray) -> FloatArray:
         return np.array([self.fitness_fn(ind) for ind in population])
 
-    def step(self, timestamp: int, population: NDArray) -> NDArray:
+    def step(self, timestamp: int, population: FloatArray) -> FloatArray:
         fitness = self.evaluate_fitness(population)
         normalized_fitness = self.fitness_normalizer(fitness)
 
@@ -66,7 +62,7 @@ def step(self, timestamp: int, population: NDArray) -> NDArray:
 
         return new_population
 
-    def run(self, initial_population: NDArray | None) -> NDArray:
+    def run(self, initial_population: FloatArray | None = None) -> FloatArray:
         population = initial_population
         if population is None:
             population = self.initialize_population()
@@ -77,9 +73,9 @@ def run(self, initial_population: NDArray | None) -> NDArray:
         self.population = population
         return population
 
-    def get_best_individual(self) -> tuple[NDArray, float]:
+    def get_best_individual(self) -> tuple[FloatArray, float]:
         if self.population is None:
             raise ValueError("Algorithm has not been run yet")
         fitness = self.evaluate_fitness(self.population)
-        best_idx = np.argmax(fitness)
-        return self.population[best_idx], fitness[best_idx]
+        best_idx = int(np.argmax(fitness))
+        return self.population[best_idx], float(fitness[best_idx])
diff --git a/src/devol/cli.py b/src/devol/cli.py
index 3a15e8b..ab3bfa8 100644
--- a/src/devol/cli.py
+++ b/src/devol/cli.py
@@ -2,14 +2,15 @@
 
 import argparse
 import sys
+from collections.abc import Callable
 from pathlib import Path
-from typing import Any
 
 import numpy as np
 from pydantic_yaml import parse_yaml_file_as
 
 from devol.algorithm import DiffusionEvolution
 from devol.config import DiffusionConfig
+from devol.distance import FloatArray
 
 
 def load_config(config_path: str) -> DiffusionConfig:
@@ -17,12 +18,12 @@ def load_config(config_path: str) -> DiffusionConfig:
     return parse_yaml_file_as(DiffusionConfig, Path(config_path))
 
 
-def sphere_function(x: np.ndarray) -> float:
+def sphere_function(x: FloatArray) -> float:
     """Sphere function: maximize -(x^2)."""
-    return -np.sum(x**2)
+    return float(-np.sum(x**2))
 
 
-def rosenbrock_function(x: np.ndarray) -> float:
+def rosenbrock_function(x: FloatArray) -> float:
     """Rosenbrock function (minimization converted to maximization)."""
     result = 0.0
     for i in range(len(x) - 1):
@@ -30,7 +31,7 @@ def rosenbrock_function(x: np.ndarray) -> float:
     return -result
 
 
-BUILTIN_FUNCTIONS: dict[str, Any] = {
+BUILTIN_FUNCTIONS: dict[str, Callable[[FloatArray], float]] = {
     "sphere": sphere_function,
     "rosenbrock": rosenbrock_function,
 }
diff --git a/src/devol/distance.py b/src/devol/distance.py
index b16afb6..eb8cbc5 100644
--- a/src/devol/distance.py
+++ b/src/devol/distance.py
@@ -5,17 +5,20 @@
 import numpy as np
 from numpy.typing import NDArray
 
+FloatArray = NDArray[np.float64]
+
 
 class DistanceComputer(Protocol):
-    def compute_distances(self, x_i: NDArray, population: NDArray) -> NDArray:
+    def compute_distances(self, x_i: FloatArray, population: FloatArray) -> FloatArray:
         """Compute distances from x_i to all individuals in population."""
         ...
 
 
 class EuclideanDistance:
-    def compute_distances(self, x_i: NDArray, population: NDArray) -> NDArray:
+    def compute_distances(self, x_i: FloatArray, population: FloatArray) -> FloatArray:
         diff = population - x_i
-        return np.sum(diff * diff, axis=1)
+        result: FloatArray = np.sum(diff * diff, axis=1)
+        return result
 
 
 class LatentDistance:
@@ -23,15 +26,16 @@ def __init__(self, param_dim: int, latent_dim: int, seed: int | None = None):
         rng = np.random.default_rng(seed)
         self.projection = rng.normal(0, 1 / np.sqrt(param_dim), (latent_dim, param_dim))
 
-    def compute_distances(self, x_i: NDArray, population: NDArray) -> NDArray:
+    def compute_distances(self, x_i: FloatArray, population: FloatArray) -> FloatArray:
         z_i = self.projection @ x_i
         z_pop = self.projection @ population.T
         diff = z_pop.T - z_i
-        return np.sum(diff * diff, axis=1)
+        result: FloatArray = np.sum(diff * diff, axis=1)
+        return result
 
 
 class CosineDistance:
-    def compute_distances(self, x_i: NDArray, population: NDArray) -> NDArray:
+    def compute_distances(self, x_i: FloatArray, population: FloatArray) -> FloatArray:
         norm_i = np.linalg.norm(x_i)
         norms_pop = np.linalg.norm(population, axis=1)
 
@@ -40,7 +44,8 @@ def compute_distances(self, x_i: NDArray, population: NDArray) -> NDArray:
 
         similarity = population @ x_i / (norms_pop * norm_i)
         similarity = np.clip(similarity, -1, 1)
-        return 1 - similarity
+        result: FloatArray = 1 - similarity
+        return result
 
 
 def create_distance_computer(
diff --git a/src/devol/evolution.py b/src/devol/evolution.py
index eec50f9..2a84f2f 100644
--- a/src/devol/evolution.py
+++ b/src/devol/evolution.py
@@ -1,26 +1,26 @@
 """Core evolution step logic."""
 
 import numpy as np
-from numpy.typing import NDArray
 
-from devol.distance import DistanceComputer
+from devol.distance import DistanceComputer, FloatArray
 
 
-def compute_gaussian_weights(distances_squared: NDArray, alpha_t: float, epsilon: float = 1e-10) -> NDArray:
+def compute_gaussian_weights(distances_squared: FloatArray, alpha_t: float, epsilon: float = 1e-10) -> FloatArray:
     variance = 1 - alpha_t
     if variance < epsilon:
         return np.ones_like(distances_squared)
     exp_term = -distances_squared / (2 * variance)
-    return np.exp(np.clip(exp_term, -50, 50))
+    result: FloatArray = np.exp(np.clip(exp_term, -50, 50))
+    return result
 
 
 def estimate_x0(
-    x_t: NDArray,
-    population: NDArray,
-    fitness_weights: NDArray,
+    x_t: FloatArray,
+    population: FloatArray,
+    fitness_weights: FloatArray,
     alpha_t: float,
     distance_computer: DistanceComputer,
-) -> NDArray:
+) -> FloatArray:
     """Estimate the clean data point x₀ from noisy observation xₜ using Bayesian inference.
 
     Implements Equations 8 and 9 from "Diffusion Models are Evolutionary Algorithms"
@@ -52,35 +52,37 @@ def estimate_x0(
     gaussian_weights = compute_gaussian_weights(distances_squared, alpha_t)
     combined_weights = fitness_weights * gaussian_weights
 
-    z = np.sum(combined_weights)
+    z = float(np.sum(combined_weights))
     if z < 1e-10:
         combined_weights = np.ones_like(combined_weights) / len(combined_weights)
         z = 1.0
 
     normalized_weights = combined_weights / z
-    x_hat_0 = np.sum(normalized_weights[:, np.newaxis] * population, axis=0)
+    x_hat_0: FloatArray = np.sum(normalized_weights[:, np.newaxis] * population, axis=0)
     return x_hat_0
 
 
-def compute_epsilon_hat(x_t: NDArray, x_hat_0: NDArray, alpha_t: float) -> NDArray:
+def compute_epsilon_hat(x_t: FloatArray, x_hat_0: FloatArray, alpha_t: float) -> FloatArray:
     """
     This is implementing equation 10 from section 3
     """
     numerator = x_t - np.sqrt(alpha_t) * x_hat_0
     denominator = np.sqrt(1 - alpha_t)
-    return numerator / denominator if denominator > 1e-10 else np.zeros_like(x_t)
+    result: FloatArray = numerator / denominator if denominator > 1e-10 else np.zeros_like(x_t)
+    return result
 
 
 def evolution_step(
-    x_t: NDArray,
-    x_hat_0: NDArray,
-    epsilon_hat: NDArray,
+    x_t: FloatArray,
+    x_hat_0: FloatArray,
+    epsilon_hat: FloatArray,
     alpha_t: float,
     alpha_t_minus_1: float,
     sigma_t: float,
     rng: np.random.Generator,
-) -> NDArray:
+) -> FloatArray:
     sqrt_alpha_prev = np.sqrt(alpha_t_minus_1)
     direction_term = np.sqrt(max(0, 1 - alpha_t_minus_1 - sigma_t**2))
     mutation = sigma_t * rng.standard_normal(x_t.shape)
-    return sqrt_alpha_prev * x_hat_0 + direction_term * epsilon_hat + mutation
+    result: FloatArray = sqrt_alpha_prev * x_hat_0 + direction_term * epsilon_hat + mutation
+    return result
diff --git a/src/devol/fitness.py b/src/devol/fitness.py
index 981db3e..d35201a 100644
--- a/src/devol/fitness.py
+++ b/src/devol/fitness.py
@@ -7,15 +7,17 @@
 
 from devol.config import FitnessMapping, NormalType
 
+FloatArray = NDArray[np.float64]
+
 
 class FitnessMapper(Protocol):
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         """Map fitness values to probability weights."""
         ...
 
 
 class FitnessNormalizer(Protocol):
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         """Normalize fitness values to be within acceptable range or gaussian."""
         ...
 
@@ -23,21 +25,23 @@ def __call__(self, fitness: NDArray) -> NDArray:
 class Identity:
     """Identity fitness mapping function."""
 
-    def __init__(self, l2_factor=0.0):
+    def __init__(self, l2_factor: float = 0.0) -> None:
         self.l2_factor = l2_factor
 
-    def l2(self, x):
-        return np.linalg.norm(x, axis=-1) ** 2
+    def l2(self, x: FloatArray) -> FloatArray:
+        result: FloatArray = np.linalg.norm(x, axis=-1) ** 2
+        return result
 
-    def forward(self, x):
+    def forward(self, x: FloatArray) -> FloatArray:
         return x
 
-    def __call__(self, fitness: NDArray) -> NDArray:
-        return self.forward(fitness) * np.exp(-1.0 * self.l2(fitness) * self.l2_factor)
+    def __call__(self, fitness: FloatArray) -> FloatArray:
+        result: FloatArray = self.forward(fitness) * np.exp(-1.0 * self.l2(fitness) * self.l2_factor)
+        return result
 
 
 class DirectMapper:
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         return fitness
 
 
@@ -51,14 +55,14 @@ class Energy(Identity):
         p: torch.Tensor, the probability of the fitness. Compute by exp(-x / temperature).
     """
 
-    def __init__(self, temperature=1.0, l2_factor=0.0):
+    def __init__(self, temperature: float = 1.0, l2_factor: float = 0.0) -> None:
         super().__init__(l2_factor=l2_factor)
         self.temperature = temperature
 
-    def forward(self, x):
+    def forward(self, x: FloatArray) -> FloatArray:
         power = -x / self.temperature
         power = power - power.max() + 5  # avoid overflow
-        p = np.exp(power)
+        p: FloatArray = np.exp(power)
         return p
 
 
@@ -66,16 +70,18 @@ class ExponentialMapper:
     def __init__(self, temperature: float = 1.0):
         self.temperature = temperature
 
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         scaled = fitness / self.temperature
         exp_fitness = np.exp(scaled - np.max(scaled))
-        return exp_fitness / np.sum(exp_fitness)
+        result: FloatArray = exp_fitness / np.sum(exp_fitness)
+        return result
 
 
 class RankMapper:
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         ranks = np.argsort(np.argsort(fitness)) + 1
-        return ranks / len(ranks)
+        result: FloatArray = ranks / len(ranks)
+        return result
 
 
 def create_fitness_mapper(
@@ -103,51 +109,55 @@ class MaxScaleNormalizer:
     def __init__(self, epsilon: float = 1e-12):
         self.epsilon = epsilon
 
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         max_abs = np.max(np.abs(fitness))
         if max_abs < self.epsilon:
             return np.zeros_like(fitness)
-        return fitness / max_abs
+        result: FloatArray = fitness / max_abs
+        return result
 
 
 class MinMaxNormalizer:
     def __init__(self, epsilon: float = 1e-12):
         self.epsilon = epsilon
 
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         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
+        result: FloatArray = (fitness - min_val) / span
+        return result
 
 
 class ZScoreNormalizer:
     def __init__(self, epsilon: float = 1e-12):
         self.epsilon = epsilon
 
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         mean = np.mean(fitness)
         std = np.std(fitness)
         if std < self.epsilon:
             return np.zeros_like(fitness)
-        return (fitness - mean) / std
+        result: FloatArray = (fitness - mean) / std
+        return result
 
 
 class SumToOneNormalizer:
     def __init__(self, epsilon: float = 1e-12):
         self.epsilon = epsilon
 
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         total = np.sum(np.abs(fitness))
         if total < self.epsilon:
             return np.zeros_like(fitness)
-        return fitness / total
+        result: FloatArray = fitness / total
+        return result
 
 
 class IdentityNormalizer:
-    def __call__(self, fitness: NDArray) -> NDArray:
+    def __call__(self, fitness: FloatArray) -> FloatArray:
         return fitness
 
 
diff --git a/src/devol/schedules.py b/src/devol/schedules.py
index 9a7d0d7..b014589 100644
--- a/src/devol/schedules.py
+++ b/src/devol/schedules.py
@@ -5,6 +5,8 @@
 import numpy as np
 from numpy.typing import NDArray
 
+FloatArray = NDArray[np.float64]
+
 
 class Schedule(Protocol):
     def __call__(self, t: int, total_steps: int) -> float:
@@ -27,7 +29,7 @@ class CosineSchedule:
     """
 
     def __call__(self, t: int, total_steps: int) -> float:
-        return 0.5 * np.cos(np.pi * t / total_steps) + 0.5
+        return float(0.5 * np.cos(np.pi * t / total_steps) + 0.5)
 
 
 class DDPMSchedule:
@@ -46,7 +48,7 @@ def __call__(self, t: int, total_steps: int) -> float:
         return float(np.exp(-beta_0 * t - gamma * t * t / total_steps))
 
 
-def create_alpha_schedule(schedule_type: str, total_steps: int, epsilon: float) -> NDArray:
+def create_alpha_schedule(schedule_type: str, total_steps: int, epsilon: float) -> FloatArray:
     schedule: Schedule
     if schedule_type == "linear":
         schedule = LinearSchedule()
@@ -60,7 +62,7 @@ def create_alpha_schedule(schedule_type: str, total_steps: int, epsilon: float)
     return np.array([schedule(t, total_steps) for t in range(total_steps + 1)])
 
 
-def create_sigma_schedule(alpha: NDArray, sigma_m: float) -> NDArray:
+def create_sigma_schedule(alpha: FloatArray, sigma_m: float) -> FloatArray:
     """
     This is following the paper's equation 17 from appendix A.2
     """
diff --git a/tests/ci/n_peaks.py b/tests/ci/n_peaks.py
index a63b252..5a269f4 100644
--- a/tests/ci/n_peaks.py
+++ b/tests/ci/n_peaks.py
@@ -7,8 +7,8 @@
 
 import argparse
 import math
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Sequence
 
 import numpy as np
 from numpy.typing import NDArray
@@ -96,8 +96,8 @@ def render_population(
 
     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)
+    X, Y = np.meshgrid(x, y)  # noqa: N806
+    Z = np.zeros_like(X)  # noqa: N806
 
     for i in range(X.shape[0]):
         for j in range(X.shape[1]):
@@ -161,10 +161,7 @@ def run_multi_peak(
         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"Average nearest-peak distance: 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)
diff --git a/tests/test_distance.py b/tests/test_distance.py
index 7bf45d6..041cf4c 100644
--- a/tests/test_distance.py
+++ b/tests/test_distance.py
@@ -1,7 +1,6 @@
 """Tests for distance computation strategies."""
 
 import numpy as np
-import pytest
 
 from devol.distance import CosineDistance, EuclideanDistance, LatentDistance
 
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index 826d5df..97bc6b9 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -1,7 +1,6 @@
 """Tests for schedule implementations."""
 
 import numpy as np
-import pytest
 
 from devol.schedules import CosineSchedule, DDPMSchedule, LinearSchedule, create_alpha_schedule