275 lines
10 KiB
Python
275 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import random
|
|
from dataclasses import dataclass
|
|
from itertools import combinations
|
|
from typing import Any
|
|
|
|
from .models import BehaviorCase, BehaviorSkeleton, Condition, HiddenRule, HiddenSystem
|
|
|
|
|
|
LENGTH_PRESETS = {"short": 4, "medium": 6, "long": 8}
|
|
ANSWER_OPTIONS = ("PRESENT", "ABSENT")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _State:
|
|
present: bool = False
|
|
invalidated: bool = False
|
|
flags: frozenset[str] = frozenset()
|
|
|
|
|
|
def resolve_length(value: str | int) -> int:
|
|
"""Resolve the number of distinct events in the held-out query."""
|
|
if isinstance(value, int):
|
|
result = value
|
|
elif value in LENGTH_PRESETS:
|
|
result = LENGTH_PRESETS[value]
|
|
else:
|
|
try:
|
|
result = int(value)
|
|
except ValueError as exc:
|
|
choices = ", ".join(LENGTH_PRESETS)
|
|
raise ValueError(f"length must be an integer or one of: {choices}") from exc
|
|
if not 4 <= result <= 10:
|
|
raise ValueError("length must be between 4 and 10")
|
|
return result
|
|
|
|
|
|
def _condition_holds(condition: Condition, attributes: dict[str, Any], state: _State) -> bool:
|
|
if condition.source == "attribute":
|
|
actual = attributes.get(condition.key)
|
|
elif condition.source == "state":
|
|
actual = condition.key in state.flags
|
|
else:
|
|
raise ValueError(f"unknown condition source: {condition.source}")
|
|
|
|
operations = {
|
|
"EQ": lambda left, right: left == right,
|
|
"LE": lambda left, right: left <= right,
|
|
"GT": lambda left, right: left > right,
|
|
"IS_SET": lambda left, right: left is bool(right),
|
|
}
|
|
try:
|
|
return operations[condition.operator](actual, condition.value)
|
|
except KeyError as exc:
|
|
raise ValueError(f"unknown condition operator: {condition.operator}") from exc
|
|
|
|
|
|
def _apply_action(state: _State, rule: HiddenRule) -> _State:
|
|
if rule.action == "SET_FLAG":
|
|
return _State(state.present, state.invalidated, state.flags | {rule.target})
|
|
if rule.action == "ADD_RECORD":
|
|
if state.invalidated:
|
|
return state
|
|
return _State(True, False, state.flags)
|
|
if rule.action == "INVALIDATE_RECORD":
|
|
return _State(False, True, state.flags)
|
|
raise ValueError(f"unknown action: {rule.action}")
|
|
|
|
|
|
def exact_present_probability(
|
|
system: HiddenSystem, attributes: dict[str, Any], events: tuple[str, ...]
|
|
) -> float:
|
|
"""Enumerate all stochastic branches while respecting event and state order."""
|
|
initial = _State(present=system.initial_record == "PRESENT")
|
|
distribution: dict[_State, float] = {initial: 1.0}
|
|
for event in events:
|
|
for rule in (item for item in system.rules if item.on_event == event):
|
|
next_distribution: dict[_State, float] = {}
|
|
for state, mass in distribution.items():
|
|
if not all(_condition_holds(item, attributes, state) for item in rule.conditions):
|
|
next_distribution[state] = next_distribution.get(state, 0.0) + mass
|
|
continue
|
|
inactive_mass = mass * (1.0 - rule.probability)
|
|
active_state = _apply_action(state, rule)
|
|
next_distribution[state] = next_distribution.get(state, 0.0) + inactive_mass
|
|
next_distribution[active_state] = (
|
|
next_distribution.get(active_state, 0.0) + mass * rule.probability
|
|
)
|
|
distribution = next_distribution
|
|
return sum(mass for state, mass in distribution.items() if state.present)
|
|
|
|
|
|
def _simulate_once(
|
|
system: HiddenSystem,
|
|
attributes: dict[str, Any],
|
|
events: tuple[str, ...],
|
|
rng: random.Random,
|
|
) -> bool:
|
|
state = _State(present=system.initial_record == "PRESENT")
|
|
for event in events:
|
|
for rule in (item for item in system.rules if item.on_event == event):
|
|
if all(_condition_holds(item, attributes, state) for item in rule.conditions):
|
|
if rng.random() < rule.probability:
|
|
state = _apply_action(state, rule)
|
|
return state.present
|
|
|
|
|
|
def _observe_case(
|
|
case_id: str,
|
|
system: HiddenSystem,
|
|
attributes: dict[str, Any],
|
|
events: tuple[str, ...],
|
|
trials: int,
|
|
rng: random.Random,
|
|
) -> BehaviorCase:
|
|
present = sum(_simulate_once(system, attributes, events, rng) for _ in range(trials))
|
|
return BehaviorCase(
|
|
id=case_id,
|
|
attributes=attributes,
|
|
events=events,
|
|
trials=trials,
|
|
observed_present=present,
|
|
observed_absent=trials - present,
|
|
)
|
|
|
|
|
|
def _build_system(index: int, complexity: int, rng: random.Random) -> HiddenSystem:
|
|
"""Build a finite stochastic program with a gated rule and veto conflicts."""
|
|
present_favored = index % 2 == 0
|
|
add_probability = rng.choice((0.85, 0.9, 0.95)) if present_favored else rng.choice((0.65, 0.75))
|
|
veto_probability = rng.choice((0.15, 0.25)) if present_favored else rng.choice((0.7, 0.8, 0.9))
|
|
marker_probability = rng.choice((0.6, 0.75, 0.9))
|
|
gated_action = "ADD_RECORD" if present_favored else "INVALIDATE_RECORD"
|
|
gated_probability = rng.choice((0.6, 0.75, 0.85))
|
|
|
|
rules = [
|
|
HiddenRule("R0", "E0", (), add_probability, "ADD_RECORD", "record"),
|
|
HiddenRule(
|
|
"R1",
|
|
"E1",
|
|
(Condition("attribute", "load", "LE", 4),),
|
|
veto_probability,
|
|
"INVALIDATE_RECORD",
|
|
"record",
|
|
),
|
|
HiddenRule("R2", "E2", (), marker_probability, "SET_FLAG", "marked"),
|
|
HiddenRule(
|
|
"R3",
|
|
"E3",
|
|
(Condition("state", "marked", "IS_SET", True),),
|
|
gated_probability,
|
|
gated_action,
|
|
"record",
|
|
),
|
|
]
|
|
for number in range(4, complexity):
|
|
action = "ADD_RECORD" if (number + index) % 3 else "INVALIDATE_RECORD"
|
|
probability = rng.choice((0.2, 0.35, 0.5))
|
|
condition = Condition("attribute", "class", "EQ", "priority")
|
|
rules.append(
|
|
HiddenRule(f"R{number}", f"E{number}", (condition,), probability, action, "record")
|
|
)
|
|
return HiddenSystem(
|
|
initial_record="ABSENT",
|
|
conflict_policy="PERSISTENT_INVALIDATION_VETO",
|
|
rules=tuple(rules),
|
|
)
|
|
|
|
|
|
def _support_event_sets(complexity: int) -> list[tuple[str, ...]]:
|
|
events = [f"E{number}" for number in range(complexity)]
|
|
cases: list[tuple[str, ...]] = [(event,) for event in events]
|
|
cases.extend((("E2", "E3"), ("E3", "E2"), ("E0", "E1")))
|
|
cases.extend(tuple(pair) for pair in combinations(events, 2) if tuple(pair) not in cases)
|
|
return cases
|
|
|
|
|
|
def _generate_one(index: int, complexity: int, rng: random.Random) -> BehaviorSkeleton:
|
|
system = _build_system(index, complexity, rng)
|
|
low = {"load": 3, "class": "priority"}
|
|
high = {"load": 8, "class": "standard"}
|
|
query_attributes = dict(low)
|
|
query_events = tuple(f"E{number}" for number in range(complexity))
|
|
probability = exact_present_probability(system, query_attributes, query_events)
|
|
if abs(probability - 0.5) < 0.08:
|
|
raise RuntimeError("generated query has an insufficient probability margin")
|
|
gold = "PRESENT" if probability > 0.5 else "ABSENT"
|
|
|
|
trials = 48 if complexity <= 4 else 64
|
|
support: list[BehaviorCase] = []
|
|
for case_number, events in enumerate(_support_event_sets(complexity)):
|
|
attributes = low if case_number % 3 else high
|
|
support.append(
|
|
_observe_case(
|
|
f"history-{case_number:03d}", system, dict(attributes), events, trials, rng
|
|
)
|
|
)
|
|
# Keeping events fixed while changing load makes the veto guard observable.
|
|
next_id = len(support)
|
|
support.append(
|
|
_observe_case(f"history-{next_id:03d}", system, dict(low), ("E0", "E1"), trials, rng)
|
|
)
|
|
support.append(
|
|
_observe_case(
|
|
f"history-{next_id + 1:03d}", system, dict(high), ("E0", "E1"), trials, rng
|
|
)
|
|
)
|
|
support.append(
|
|
_observe_case(
|
|
f"history-{next_id + 2:03d}",
|
|
system,
|
|
dict(low),
|
|
("E0", "E2", "E3"),
|
|
trials,
|
|
rng,
|
|
)
|
|
)
|
|
|
|
query = BehaviorCase(
|
|
id="query",
|
|
attributes=query_attributes,
|
|
events=query_events,
|
|
trials=1,
|
|
observed_present=None,
|
|
observed_absent=None,
|
|
)
|
|
trace = tuple(
|
|
[f"query events: {' -> '.join(query_events)}"]
|
|
+ [
|
|
f"{rule.id}: on {rule.on_event}, p={rule.probability}, {rule.action}({rule.target})"
|
|
for rule in system.rules
|
|
]
|
|
+ [f"exact P(PRESENT)={probability:.8f}", f"gold={gold}"]
|
|
)
|
|
return BehaviorSkeleton(
|
|
id=f"dynamic-{index:04d}",
|
|
family="stochastic_behavioral_system",
|
|
complexity=complexity,
|
|
event_vocabulary=tuple(f"E{number}" for number in range(complexity)),
|
|
attribute_vocabulary={"load": "integer", "class": "category"},
|
|
observation_protocol=(
|
|
"Each trial begins with an absent record and no internal flags.",
|
|
"Events occur from left to right; the same hidden system is used in every trial.",
|
|
"PRESENT and ABSENT counts are final observations over independent repeated trials.",
|
|
"For the query, choose which final outcome is more likely; an exact probability is not required.",
|
|
),
|
|
hidden_system=system,
|
|
support_cases=tuple(support),
|
|
query_case=query,
|
|
answer_options=ANSWER_OPTIONS,
|
|
gold_answer=gold,
|
|
gold_present_probability=probability,
|
|
latent_trace=trace,
|
|
)
|
|
|
|
|
|
def generate_skeletons(count: int, length: str | int, seed: int = 0) -> list[BehaviorSkeleton]:
|
|
if count < 1:
|
|
raise ValueError("count must be positive")
|
|
complexity = resolve_length(length)
|
|
rng = random.Random(seed)
|
|
generated: list[BehaviorSkeleton] = []
|
|
for index in range(count):
|
|
for _attempt in range(100):
|
|
try:
|
|
generated.append(_generate_one(index, complexity, rng))
|
|
break
|
|
except RuntimeError as exc:
|
|
if "insufficient probability margin" not in str(exc):
|
|
raise
|
|
else:
|
|
raise RuntimeError("could not generate a query with a stable probability margin")
|
|
return generated
|