staplebridge/training/trajectory.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from staplebridge.chemistry.state import StapleState


@dataclass
class TransitionStep:
    state: StapleState
    next_state: StapleState
    candidates: list[StapleState]
    chosen_idx: int
    t: int


@dataclass
class WeightedTrajectory:
    steps: list[TransitionStep]
    terminal_state: StapleState
    terminal_energy: float
    weight: float = 0.0
    context: dict[str, Any] = field(default_factory=dict)
    # Hydrocarbon Exact-SB uses these fields to keep plan-conditioned positive
    # paths separate from off-plan/unfinished failures.  Defaults preserve all
    # legacy callers.
    is_positive: bool = True
    failure_reason: str | None = None