Files
PSO/tests/test_deep_pso_methods.py
T
jung-geun 813433000a feat: modernize PSO and add convergence research
Migrate the package and examples to the tensor-native PyTorch implementation, add benchmark evidence, and add the guarded post-training convergence protocol with TensorBoard progress monitoring and hash-verified recovery.

Constraint: Preserve one-shot official-test sealing and auditable research artifacts

Rejected: Commit local .omc runs and downloaded datasets | multi-gigabyte runtime state is machine-local

Confidence: high

Scope-risk: broad

Not-tested: Production CUDA run on pieroot-server
2026-09-07 22:03:25 +09:00

236 lines
8.8 KiB
Python

"""
Focused Offline Unit Tests for MNIST Deep PSO Methods (Protocol MNIST-PSO-RAW-V5 1.0.0)
Tests:
1. Split balance and nested subset inclusion invariant (I_2k subset of I_10k subset of I_50k).
2. Latent transform exactness (z_0 = 0 -> theta_0), antithetic symmetry (even swarm pairing), and subspace dimensions.
3. Stage transition pbest re-evaluation, gbest rebuild, and exact query/sample accounting.
4. Validation-only pilot and elite ensemble selection with greedy disagreement.
5. CLI argument validation.
"""
import sys
from pathlib import Path
# Insert repo test/ path so deep_pso_methods can be imported
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "test"))
import math
import numpy as np
import pytest
import torch
import torch.nn as nn
from deep_pso_methods import (
CompactCNN,
LatentTransform,
build_nested_stratified_subsets,
evaluate_latent_batch,
evaluate_probabilistic_metrics,
make_compact_cnn,
run_latent_pso,
select_diverse_candidates,
validate_cli_args,
build_parser,
)
def test_split_balance_and_nesting():
"""Verify stratified split balance and nested index inclusion: I_2k subset of I_10k subset of I_50k."""
N = 50000
num_classes = 10
samples_per_class = N // num_classes
y_search = torch.cat([torch.full((samples_per_class,), c, dtype=torch.long) for c in range(num_classes)])
split_seed = 20260902
nested_subsets = build_nested_stratified_subsets(
y_search=y_search,
subset_sizes=[2000, 10000, 50000],
subset_seed=split_seed,
)
idx_2k = nested_subsets[2000]
idx_10k = nested_subsets[10000]
idx_50k = nested_subsets[50000]
assert len(idx_2k) == 2000
assert len(idx_10k) == 10000
assert len(idx_50k) == 50000
set_2k = set(idx_2k.numpy().tolist())
set_10k = set(idx_10k.numpy().tolist())
set_50k = set(idx_50k.numpy().tolist())
# Strict nesting: I_2k subset of I_10k subset of I_50k
assert set_2k.issubset(set_10k), "I_2k must be a strict subset of I_10k"
assert set_10k.issubset(set_50k), "I_10k must be a strict subset of I_50k"
# Exact stratification (equal counts per class)
y_2k = y_search[idx_2k].numpy()
y_10k = y_search[idx_10k].numpy()
counts_2k = np.bincount(y_2k, minlength=10)
counts_10k = np.bincount(y_10k, minlength=10)
for c in range(num_classes):
assert counts_2k[c] == 200, f"Class {c} in 2k subset must have 200 samples; got {counts_2k[c]}"
assert counts_10k[c] == 1000, f"Class {c} in 10k subset must have 1000 samples; got {counts_10k[c]}"
def test_latent_transform_exactness_and_antithetic_symmetry():
"""Verify transform exactness (z_0 = 0 -> theta_0), even swarm pairing, zero centroid, and subspace shapes."""
device = torch.device("cpu")
base_model = make_compact_cnn(seed=41).to(device)
base_vec = torch.cat([p.detach().view(-1) for p in base_model.parameters()])
dims = [290, 1024, 4096, "full"]
swarm_size = 30 # Even swarm size
for d in dims:
transform = LatentTransform(base_model, latent_dim=d, device=device)
Z = transform.init_swarm(swarm_size=swarm_size, seed=91)
# 1. Particle 0 is exact zero (base model)
decoded_p0 = transform.decode(Z[0:1]).squeeze(0)
assert torch.allclose(decoded_p0, base_vec, atol=1e-6), f"Particle 0 for dim={d} must match exact base vector"
# 2. For even N=30: particle N-1 (index 29) is also exact zero
decoded_plast = transform.decode(Z[29:30]).squeeze(0)
assert torch.allclose(decoded_plast, base_vec, atol=1e-6), f"Particle N-1 for dim={d} must match exact zero"
# 3. Antithetic pairs (particles 1..28 in 14 exact pairs)
z1 = Z[1:2]
z2 = Z[2:3]
assert torch.allclose(z1 + z2, torch.zeros_like(z1), atol=1e-6), "Antithetic pair latent sum must be zero"
delta1 = transform.decode(z1).squeeze(0) - transform.base_vec
delta2 = transform.decode(z2).squeeze(0) - transform.base_vec
assert torch.allclose(delta1 + delta2, torch.zeros_like(delta1), atol=1e-5), "Antithetic pair delta sum must be zero"
# 4. Latent swarm centroid is strictly zero
centroid_z = Z.mean(dim=0)
assert torch.allclose(centroid_z, torch.zeros_like(centroid_z), atol=1e-6), "Overall swarm centroid must be zero"
def test_transition_reevaluation_and_exact_accounting():
"""Verify objective size transition re-evaluates all pbests, rebuilds gbest, and asserts exact query/sample counts."""
device = torch.device("cpu")
base_model = make_compact_cnn(seed=41).to(device)
N_samples = 100
x_synth = torch.randn(N_samples, 1, 28, 28)
y_synth = torch.randint(0, 10, (N_samples,))
nested_subsets = {
20: torch.arange(20, dtype=torch.long),
50: torch.arange(50, dtype=torch.long),
}
transform = LatentTransform(base_model, latent_dim=290, device=device)
# Run 2-stage PSO: stage 0 (20 samples, 2 epochs), stage 1 (50 samples, 2 epochs), swarm_size = 10
# Stage 0: 2 * 10 = 20 queries, 20 * 20 = 400 sample evals. No transition re-eval.
# Stage 1 transition: 1 * 10 = 10 queries, 10 * 50 = 500 sample evals. Transition count = 10.
# Stage 1: 2 * 10 = 20 queries, 20 * 50 = 1000 sample evals.
# Total queries = 20 + 10 + 20 = 50.
# Total sample evals = 400 + 500 + 1000 = 1900.
res = run_latent_pso(
transform=transform,
base_model=base_model,
x_search=x_synth,
y_search=y_synth,
nested_subsets=nested_subsets,
schedule_str="20:2,50:2",
epochs=4,
swarm_size=10,
seed=42,
device=device,
)
assert res["transition_reevaluation_counts"] == 10, f"Expected 10 transition re-evaluations; got {res['transition_reevaluation_counts']}"
assert res["total_queries"] == 50, f"Expected 50 total queries; got {res['total_queries']}"
assert res["total_sample_evaluations"] == 1900, f"Expected 1900 sample evaluations; got {res['total_sample_evaluations']}"
assert len(res["stage_histories"]) == 4
def test_validation_only_elite_and_ensemble_selection():
"""Verify validation metrics correctly rank candidates and metric routines compute expected values."""
N = 100
C = 10
y_val = torch.randint(0, C, (N,))
# Candidate 1: Perfect predictions
probs_perfect = torch.zeros((N, C), dtype=torch.float32)
probs_perfect[torch.arange(N), y_val] = 1.0
# Candidate 2: Random noise
probs_random = torch.full((N, C), 1.0 / C, dtype=torch.float32)
m1 = evaluate_probabilistic_metrics(probs_perfect, y_val)
m2 = evaluate_probabilistic_metrics(probs_random, y_val)
assert m1["accuracy"] == 100.0
assert m1["nll"] < m2["nll"]
assert m1["brier"] < m2["brier"]
assert m1["ece"] <= 0.01
candidates = [
{"id": "cand2", "val_loss": m2["nll"], "val_acc": m2["accuracy"]},
{"id": "cand1", "val_loss": m1["nll"], "val_acc": m1["accuracy"]},
]
candidates.sort(key=lambda c: (c["val_loss"], -c["val_acc"]))
assert candidates[0]["id"] == "cand1"
diverse_candidates = [
{
"seed": 7,
"particle_idx": particle_idx,
"val_loss": 0.2 + 0.01 * particle_idx,
"val_acc": 90.0 - 0.25 * particle_idx,
"val_probs": torch.roll(probs_perfect, shifts=particle_idx, dims=1),
"latent_z": torch.full((4,), float(particle_idx)),
}
for particle_idx in range(3)
]
selected = select_diverse_candidates(
diverse_candidates, max_size=3, accuracy_window=2.0
)
assert len(selected) == 3
assert len({
(candidate["seed"], candidate["particle_idx"])
for candidate in selected
}) == 3
def test_cli_argument_validation():
"""Verify CLI argument validation rejects invalid parameters and accepts valid settings."""
parser = build_parser()
# Valid args
valid_args = parser.parse_args([
"--pilot-epochs", "160",
"--confirmation-epochs", "600",
"--confirmation-schedule", "2000:420,10000:135,50000:45",
"--seeds", "101", "102", "103",
"--dimensions", "290", "1024", "4096", "full"
])
validate_cli_args(valid_args)
# Invalid schedule sum mismatch
invalid_schedule = parser.parse_args([
"--confirmation-epochs", "600",
"--confirmation-schedule", "2000:400,10000:100,50000:50" # Sums to 550 != 600
])
with pytest.raises(ValueError, match="Schedule epoch sum"):
validate_cli_args(invalid_schedule)
# Invalid negative seed
invalid_seed = parser.parse_args(["--seeds", "-1"])
with pytest.raises(ValueError, match="Seeds must be non-negative"):
validate_cli_args(invalid_seed)
# Invalid duplicate seed
duplicate_seed = parser.parse_args(["--seeds", "101", "101"])
with pytest.raises(ValueError, match="Confirmation seeds must be unique"):
validate_cli_args(duplicate_seed)