exploit_orthogonal.py

from __future__ import annotations

from app.core.schema import Candidate, Session
from app.samplers.base import clamp_vector, make_rng


class ExploitOrthogonalSampler:
    """Structured sampler that mixes exploit, refine, and orthogonal probes."""

    name = "exploit_orthogonal"

    def propose(self, session: Session, seed: int) -> list[Candidate]:
        """Generate a batch with deliberate directional roles."""

        rng = make_rng(seed + 17)
        base = session.current_z
        dimensions = max(1, len(base))
        candidates = []
        for index in range(session.config.candidate_count):
            direction, role = self._pattern_for_index(index, dimensions)
            jitter = [rng.uniform(-0.03, 0.03) for _ in direction]
            z = clamp_vector(
                [current + delta + noise for current, delta, noise in zip(base, direction, jitter, strict=False)],
                session.config.trust_radius,
            )
            candidates.append(
                Candidate(
                    round_id="",
                    candidate_index=index,
                    z=z,
                    sampler_role=role,
                    predicted_score=sum(z) + 0.05,
                    predicted_uncertainty=0.12 + (index * 0.03),
                    seed=seed,
                    generation_params={"image_size": session.config.image_size},
                )
            )
        return candidates

    @staticmethod
    def _pattern_for_index(index: int, dimensions: int) -> tuple[list[float], str]:
        """Build a structured direction pattern for any steering dimensionality."""

        patterns = [
            ("exploit", ((0, 0.18), (1, 0.04), (2, -0.02))),
            ("refine", ((0, 0.16), (1, -0.08), (2, 0.06))),
            ("orthogonal", ((0, -0.02), (1, 0.18), (2, 0.08))),
            ("mirror", ((0, -0.14), (1, -0.04), (2, 0.12))),
        ]
        role, anchors = patterns[index % len(patterns)]
        direction = [0.0 for _ in range(dimensions)]
        for axis, value in anchors:
            direction[axis % dimensions] += value
        if dimensions > 3:
            extra_axis = (index + 3) % dimensions
            direction[extra_axis] += 0.05 if index % 2 == 0 else -0.05
        return direction, role