gausse/src/gausse/optim/search_space.py

"""Пространство поиска: геном переменной длины (число ступеней эволюционирует).

Геном кодирует только ИНДЕКСЫ в базу реальных компонентов (не сами
параметры) + непрерывные величины (расстояния, напряжение, геометрия
снаряда). Это гарантирует, что что бы ни нашёл поиск, оно собрано из
реально продающихся деталей, а не из выдуманных чисел.
"""

import copy
import math
import random
from dataclasses import dataclass, field

from gausse.components.database import ComponentDatabase
from gausse.sim.coilgun import CoilgunConfig
from gausse.sim.stage import ProjectileConfig, StageConfig

TUBE_WALL_CLEARANCE_M = 0.001


@dataclass(frozen=True)
class SearchBounds:
    min_stages: int = 1
    max_stages: int = 4
    turns_per_layer_min: int = 5
    turns_per_layer_max: int = 100
    layers_min: int = 1
    layers_max: int = 8
    sensor_to_coil_distance_m_min: float = 0.005
    sensor_to_coil_distance_m_max: float = 0.05
    inter_stage_gap_m_min: float = 0.01
    inter_stage_gap_m_max: float = 0.10
    tube_od_m_min: float = 0.009
    tube_od_m_max: float = 0.02
    projectile_diameter_m_min: float = 0.004
    projectile_diameter_m_max: float = 0.008
    projectile_length_m_min: float = 0.01
    projectile_length_m_max: float = 0.03
    charge_voltage_fraction_min: float = 0.5
    charge_voltage_fraction_max: float = 1.0
    initial_launch_velocity_mps: float = 3.0
    initial_approach_margin_m: float = 0.02


@dataclass
class StageGene:
    wire_idx: int
    capacitor_idx: int
    switch_idx: int
    sensor_idx: int
    turns_per_layer: int
    layers: int
    sensor_to_coil_distance_m: float
    charge_voltage_fraction: float


@dataclass
class ProjectileGene:
    material_idx: int
    diameter_m: float
    length_m: float


@dataclass
class Genome:
    tube_od_m: float
    stages: list[StageGene]
    inter_stage_gaps_m: list[float]
    projectile: ProjectileGene


def _clip(value: float, lo: float, hi: float) -> float:
    return max(lo, min(hi, value))


def sample_stage_gene(db: ComponentDatabase, bounds: SearchBounds, rng: random.Random) -> StageGene:
    return StageGene(
        wire_idx=rng.randrange(len(db.wires)),
        capacitor_idx=rng.randrange(len(db.capacitors)),
        switch_idx=rng.randrange(len(db.switches)),
        sensor_idx=rng.randrange(len(db.sensors)),
        turns_per_layer=rng.randint(bounds.turns_per_layer_min, bounds.turns_per_layer_max),
        layers=rng.randint(bounds.layers_min, bounds.layers_max),
        sensor_to_coil_distance_m=rng.uniform(
            bounds.sensor_to_coil_distance_m_min, bounds.sensor_to_coil_distance_m_max
        ),
        charge_voltage_fraction=rng.uniform(
            bounds.charge_voltage_fraction_min, bounds.charge_voltage_fraction_max
        ),
    )


def sample_genome(db: ComponentDatabase, bounds: SearchBounds, rng: random.Random) -> Genome:
    tube_od_m = rng.uniform(bounds.tube_od_m_min, bounds.tube_od_m_max)
    max_diameter = min(bounds.projectile_diameter_m_max, tube_od_m - TUBE_WALL_CLEARANCE_M)
    diameter_m = rng.uniform(bounds.projectile_diameter_m_min, max(max_diameter, bounds.projectile_diameter_m_min))

    n_stages = rng.randint(bounds.min_stages, bounds.max_stages)
    stages = [sample_stage_gene(db, bounds, rng) for _ in range(n_stages)]
    gaps = [
        rng.uniform(bounds.inter_stage_gap_m_min, bounds.inter_stage_gap_m_max)
        for _ in range(n_stages - 1)
    ]
    projectile = ProjectileGene(
        material_idx=rng.randrange(len(db.projectile_materials)),
        diameter_m=diameter_m,
        length_m=rng.uniform(bounds.projectile_length_m_min, bounds.projectile_length_m_max),
    )
    return Genome(tube_od_m=tube_od_m, stages=stages, inter_stage_gaps_m=gaps, projectile=projectile)


def repair(genome: Genome, db: ComponentDatabase, bounds: SearchBounds) -> Genome:
    """Приводит геном в границы после мутации/скрещивания (клэмп, не отбраковка)."""
    genome.tube_od_m = _clip(genome.tube_od_m, bounds.tube_od_m_min, bounds.tube_od_m_max)
    max_diameter = max(genome.tube_od_m - TUBE_WALL_CLEARANCE_M, bounds.projectile_diameter_m_min)
    genome.projectile.diameter_m = _clip(
        genome.projectile.diameter_m, bounds.projectile_diameter_m_min, max_diameter
    )
    genome.projectile.length_m = _clip(
        genome.projectile.length_m, bounds.projectile_length_m_min, bounds.projectile_length_m_max
    )
    genome.projectile.material_idx = genome.projectile.material_idx % len(db.projectile_materials)

    for stage in genome.stages:
        stage.wire_idx %= len(db.wires)
        stage.capacitor_idx %= len(db.capacitors)
        stage.switch_idx %= len(db.switches)
        stage.sensor_idx %= len(db.sensors)
        stage.turns_per_layer = int(_clip(stage.turns_per_layer, bounds.turns_per_layer_min, bounds.turns_per_layer_max))
        stage.layers = int(_clip(stage.layers, bounds.layers_min, bounds.layers_max))
        stage.sensor_to_coil_distance_m = _clip(
            stage.sensor_to_coil_distance_m,
            bounds.sensor_to_coil_distance_m_min,
            bounds.sensor_to_coil_distance_m_max,
        )
        stage.charge_voltage_fraction = _clip(
            stage.charge_voltage_fraction,
            bounds.charge_voltage_fraction_min,
            bounds.charge_voltage_fraction_max,
        )
    genome.inter_stage_gaps_m = [
        _clip(g, bounds.inter_stage_gap_m_min, bounds.inter_stage_gap_m_max)
        for g in genome.inter_stage_gaps_m[: len(genome.stages) - 1]
    ]
    while len(genome.inter_stage_gaps_m) < len(genome.stages) - 1:
        genome.inter_stage_gaps_m.append(
            (bounds.inter_stage_gap_m_min + bounds.inter_stage_gap_m_max) / 2
        )
    return genome


def decode(genome: Genome, db: ComponentDatabase, bounds: SearchBounds) -> tuple[CoilgunConfig, float, float]:
    """Возвращает (CoilgunConfig, initial_x_m, initial_v_mps)."""
    material = db.projectile_materials[genome.projectile.material_idx % len(db.projectile_materials)]
    projectile = ProjectileConfig(
        material=material,
        diameter_m=genome.projectile.diameter_m,
        length_m=genome.projectile.length_m,
    )

    stage_configs = []
    for gene in genome.stages:
        wire = db.wires[gene.wire_idx % len(db.wires)]
        capacitor = db.capacitors[gene.capacitor_idx % len(db.capacitors)]
        switch = db.switches[gene.switch_idx % len(db.switches)]
        sensor = db.sensors[gene.sensor_idx % len(db.sensors)]
        max_voltage = min(capacitor.voltage_v, switch.max_voltage_v)
        charge_voltage_v = gene.charge_voltage_fraction * max_voltage
        stage_configs.append(
            StageConfig(
                wire=wire,
                capacitor=capacitor,
                switch=switch,
                sensor=sensor,
                tube_od_m=genome.tube_od_m,
                turns_per_layer=gene.turns_per_layer,
                layers=gene.layers,
                sensor_to_coil_distance_m=gene.sensor_to_coil_distance_m,
                charge_voltage_v=charge_voltage_v,
            )
        )

    config = CoilgunConfig(
        stages=stage_configs,
        inter_stage_gaps_m=list(genome.inter_stage_gaps_m),
        projectile=projectile,
        initial_x_m=-(stage_configs[0].sensor_to_coil_distance_m + bounds.initial_approach_margin_m),
        initial_v_mps=bounds.initial_launch_velocity_mps,
    )
    return config, config.initial_x_m, config.initial_v_mps


def mutate(genome: Genome, db: ComponentDatabase, bounds: SearchBounds, rng: random.Random, rate: float = 0.2) -> Genome:
    child = copy.deepcopy(genome)

    if rng.random() < rate:
        child.tube_od_m = rng.uniform(bounds.tube_od_m_min, bounds.tube_od_m_max)

    for stage in child.stages:
        if rng.random() < rate:
            stage.wire_idx = rng.randrange(len(db.wires))
        if rng.random() < rate:
            stage.capacitor_idx = rng.randrange(len(db.capacitors))
        if rng.random() < rate:
            stage.switch_idx = rng.randrange(len(db.switches))
        if rng.random() < rate:
            stage.sensor_idx = rng.randrange(len(db.sensors))
        if rng.random() < rate:
            stage.turns_per_layer += rng.randint(-10, 10)
        if rng.random() < rate:
            stage.layers += rng.randint(-1, 1)
        if rng.random() < rate:
            stage.sensor_to_coil_distance_m *= rng.uniform(0.7, 1.3)
        if rng.random() < rate:
            stage.charge_voltage_fraction *= rng.uniform(0.9, 1.1)

    for i in range(len(child.inter_stage_gaps_m)):
        if rng.random() < rate:
            child.inter_stage_gaps_m[i] *= rng.uniform(0.7, 1.3)

    if rng.random() < rate:
        child.projectile.material_idx = rng.randrange(len(db.projectile_materials))
    if rng.random() < rate:
        child.projectile.diameter_m *= rng.uniform(0.8, 1.2)
    if rng.random() < rate:
        child.projectile.length_m *= rng.uniform(0.8, 1.2)

    # структурные операторы: число ступеней тоже эволюционирует
    structural_roll = rng.random()
    if structural_roll < rate / 3 and len(child.stages) < bounds.max_stages:
        new_stage = sample_stage_gene(db, bounds, rng)
        child.stages.append(new_stage)
        child.inter_stage_gaps_m.append(
            rng.uniform(bounds.inter_stage_gap_m_min, bounds.inter_stage_gap_m_max)
        )
    elif structural_roll < 2 * rate / 3 and len(child.stages) > bounds.min_stages:
        idx = rng.randrange(len(child.stages))
        child.stages.pop(idx)
        if child.inter_stage_gaps_m:
            child.inter_stage_gaps_m.pop(min(idx, len(child.inter_stage_gaps_m) - 1))
    elif structural_roll < rate and len(child.stages) < bounds.max_stages:
        idx = rng.randrange(len(child.stages))
        child.stages.insert(idx, copy.deepcopy(child.stages[idx]))
        child.inter_stage_gaps_m.append(
            rng.uniform(bounds.inter_stage_gap_m_min, bounds.inter_stage_gap_m_max)
        )

    return repair(child, db, bounds)


def crossover(parent_a: Genome, parent_b: Genome, rng: random.Random) -> Genome:
    """Обмен целыми ступенями между родителями — ступень физически цельная единица."""
    n = rng.choice([len(parent_a.stages), len(parent_b.stages)])
    stages = []
    for i in range(n):
        source = parent_a if rng.random() < 0.5 else parent_b
        if i < len(source.stages):
            stages.append(copy.deepcopy(source.stages[i]))
        else:
            fallback = parent_a if source is parent_b else parent_b
            stages.append(copy.deepcopy(fallback.stages[i % len(fallback.stages)]))

    gaps = []
    for i in range(n - 1):
        source = parent_a if rng.random() < 0.5 else parent_b
        if i < len(source.inter_stage_gaps_m):
            gaps.append(source.inter_stage_gaps_m[i])
        else:
            gaps.append(parent_a.inter_stage_gaps_m[i % max(len(parent_a.inter_stage_gaps_m), 1)])

    projectile_source = parent_a if rng.random() < 0.5 else parent_b
    tube_source = parent_a if rng.random() < 0.5 else parent_b

    return Genome(
        tube_od_m=tube_source.tube_od_m,
        stages=stages,
        inter_stage_gaps_m=gaps,
        projectile=copy.deepcopy(projectile_source.projectile),
    )


def genome_to_dict(genome: Genome) -> dict:
    return {
        "tube_od_m": genome.tube_od_m,
        "stages": [vars(s) for s in genome.stages],
        "inter_stage_gaps_m": genome.inter_stage_gaps_m,
        "projectile": vars(genome.projectile),
    }


def genome_from_dict(data: dict) -> Genome:
    return Genome(
        tube_od_m=data["tube_od_m"],
        stages=[StageGene(**s) for s in data["stages"]],
        inter_stage_gaps_m=list(data["inter_stage_gaps_m"]),
        projectile=ProjectileGene(**data["projectile"]),
    )