""" analysis.trajectory_extractor ============================= Convert raw WorldSmithAI simulation history into structured trajectory data. This module is intentionally placed above the simulation engine. It consumes world history, metric logs, and event logs without mutating runtime agents, resources, behaviors, policies, schedulers, or worlds. The extractor is designed to be tolerant of several history formats because the core engine may evolve independently. It accepts either: - a world-like object with ``history`` - a raw sequence of snapshot mappings - a single snapshot mapping The output is a ``TrajectoryData`` dataclass consumed by later analysis modules. Example ------- extractor = TrajectoryExtractor() trajectory = extractor.extract(world) print(trajectory.step_count) print(trajectory.agent_type_counts) print(trajectory.population_timeseries) Architecture ------------ world.history ↓ TrajectoryExtractor.extract(...) ↓ TrajectoryData ↓ simulation_analyzer / report_generator / scenario comparator """ from __future__ import annotations import copy import logging import math from collections import defaultdict from collections.abc import Mapping, Sequence from dataclasses import dataclass, field from typing import Any, Protocol logger = logging.getLogger(__name__) class WorldHistoryProtocol(Protocol): """Structural protocol for world-like objects with history.""" history: Any step_count: int @dataclass(frozen=True) class TrajectoryExtractionConfig: """Configuration for trajectory extraction. Attributes: count_alive_only: If true, population counts exclude agents with ``alive=False``. include_snapshot_events: If true, event-like records inside each snapshot are added to the normalized event log. include_snapshot_metrics: If true, metric-like records inside each snapshot are added to the normalized metric history. preserve_unknown_snapshot_fields: If true, selected unrecognized top-level snapshot keys are preserved in metadata for debugging. max_preserved_unknown_values: Maximum number of unknown snapshot values preserved in metadata. """ count_alive_only: bool = True include_snapshot_events: bool = True include_snapshot_metrics: bool = True preserve_unknown_snapshot_fields: bool = False max_preserved_unknown_values: int = 20 @dataclass(frozen=True) class TrajectoryData: """Structured simulation trajectory data. Attributes: population_timeseries: Mapping from step to counts by agent type. resource_timeseries: Mapping from step to total resource amount by resource type. event_log: Normalized event records. Each record should include a ``step`` key when the step is known. metric_history: Mapping from metric name to a list of per-step metric records. agent_type_counts: Final observed agent counts by type. step_count: Number of simulated steps represented by the trajectory. This is the highest observed step plus one when possible. first_step: First observed step, or ``None`` when there is no history. last_step: Last observed step, or ``None`` when there is no history. snapshot_count: Number of snapshots consumed. metadata: Additional extraction diagnostics. """ population_timeseries: Mapping[int, Mapping[str, int]] = field(default_factory=dict) resource_timeseries: Mapping[int, Mapping[str, float]] = field(default_factory=dict) event_log: tuple[Mapping[str, Any], ...] = field(default_factory=tuple) metric_history: Mapping[str, tuple[Mapping[str, Any], ...]] = field(default_factory=dict) agent_type_counts: Mapping[str, int] = field(default_factory=dict) step_count: int = 0 first_step: int | None = None last_step: int | None = None snapshot_count: int = 0 metadata: Mapping[str, Any] = field(default_factory=dict) @property def is_empty(self) -> bool: """Return whether this trajectory contains no snapshots.""" return self.snapshot_count == 0 @property def final_population(self) -> Mapping[str, int]: """Return the final population count by agent type.""" if self.last_step is None: return {} return dict(self.population_timeseries.get(self.last_step, {})) @property def final_resources(self) -> Mapping[str, float]: """Return the final resource totals by resource type.""" if self.last_step is None: return {} return dict(self.resource_timeseries.get(self.last_step, {})) def population_delta(self) -> dict[str, int]: """Return final-minus-initial population delta by agent type.""" if self.first_step is None or self.last_step is None: return {} initial = self.population_timeseries.get(self.first_step, {}) final = self.population_timeseries.get(self.last_step, {}) keys = set(initial.keys()) | set(final.keys()) return { key: int(final.get(key, 0)) - int(initial.get(key, 0)) for key in sorted(keys) } def resource_delta(self) -> dict[str, float]: """Return final-minus-initial resource amount delta by resource type.""" if self.first_step is None or self.last_step is None: return {} initial = self.resource_timeseries.get(self.first_step, {}) final = self.resource_timeseries.get(self.last_step, {}) keys = set(initial.keys()) | set(final.keys()) return { key: float(final.get(key, 0.0)) - float(initial.get(key, 0.0)) for key in sorted(keys) } def to_dict(self) -> dict[str, Any]: """Return a JSON-friendly representation of the trajectory.""" return { "population_timeseries": { str(step): dict(counts) for step, counts in sorted(self.population_timeseries.items()) }, "resource_timeseries": { str(step): dict(amounts) for step, amounts in sorted(self.resource_timeseries.items()) }, "event_log": [dict(event) for event in self.event_log], "metric_history": { metric_name: [dict(record) for record in records] for metric_name, records in self.metric_history.items() }, "agent_type_counts": dict(self.agent_type_counts), "step_count": self.step_count, "first_step": self.first_step, "last_step": self.last_step, "snapshot_count": self.snapshot_count, "population_delta": self.population_delta(), "resource_delta": self.resource_delta(), "metadata": copy.deepcopy(dict(self.metadata)), } @dataclass class TrajectoryExtractor: """Extract structured trajectory data from completed simulations.""" config: TrajectoryExtractionConfig = field(default_factory=TrajectoryExtractionConfig) def extract( self, world_or_history: Any, *, metric_history: Mapping[str, Any] | Sequence[Any] | None = None, event_log: Sequence[Any] | Mapping[str, Any] | None = None, ) -> TrajectoryData: """Extract trajectory data from a world-like object or raw history. Args: world_or_history: Either a world-like object with ``history`` or a raw history object. metric_history: Optional external metric history to merge with snapshot metrics. event_log: Optional external event log to merge with snapshot events. Returns: Normalized ``TrajectoryData``. """ world = self._as_world_like(world_or_history) raw_history = self._history_from_input(world_or_history) snapshots = self._normalize_history(raw_history) if not snapshots and world is not None: current_snapshot = self._snapshot_from_world(world) if current_snapshot: snapshots = (current_snapshot,) population_timeseries: dict[int, dict[str, int]] = {} resource_timeseries: dict[int, dict[str, float]] = {} normalized_event_log: list[dict[str, Any]] = [] normalized_metric_history: dict[str, list[dict[str, Any]]] = defaultdict(list) unknown_snapshot_metadata: list[dict[str, Any]] = [] for index, raw_snapshot in enumerate(snapshots): snapshot = self._mapping_from_any(raw_snapshot) step = self._step_from_snapshot(snapshot, fallback=index) population_counts = self._population_counts_from_snapshot(snapshot) resource_amounts = self._resource_amounts_from_snapshot(snapshot) population_timeseries[step] = population_counts resource_timeseries[step] = resource_amounts if self.config.include_snapshot_events: self._append_snapshot_events( snapshot=snapshot, step=step, output=normalized_event_log, ) if self.config.include_snapshot_metrics: self._append_snapshot_metrics( snapshot=snapshot, step=step, output=normalized_metric_history, ) if self.config.preserve_unknown_snapshot_fields: preserved = self._preserve_unknown_snapshot_fields(snapshot, step=step) if preserved: unknown_snapshot_metadata.append(preserved) self._append_external_events( event_log=event_log, world=world, output=normalized_event_log, ) self._append_external_metrics( metric_history=metric_history, world=world, output=normalized_metric_history, ) observed_steps = sorted(population_timeseries.keys() | resource_timeseries.keys()) first_step = observed_steps[0] if observed_steps else None last_step = observed_steps[-1] if observed_steps else None if last_step is None: step_count = int(getattr(world, "step_count", 0)) if world is not None else 0 else: step_count = max(last_step + 1, int(getattr(world, "step_count", 0)) if world is not None else 0) final_counts = ( population_timeseries.get(last_step, {}) if last_step is not None else {} ) metadata: dict[str, Any] = { "source": "world" if world is not None else "history", "history_length": len(snapshots), "count_alive_only": self.config.count_alive_only, } if unknown_snapshot_metadata: metadata["unknown_snapshot_fields"] = unknown_snapshot_metadata[ : self.config.max_preserved_unknown_values ] return TrajectoryData( population_timeseries={ step: dict(counts) for step, counts in sorted(population_timeseries.items()) }, resource_timeseries={ step: dict(amounts) for step, amounts in sorted(resource_timeseries.items()) }, event_log=tuple(normalized_event_log), metric_history={ metric_name: tuple(records) for metric_name, records in sorted(normalized_metric_history.items()) }, agent_type_counts=dict(final_counts), step_count=step_count, first_step=first_step, last_step=last_step, snapshot_count=len(snapshots), metadata=metadata, ) def extract_from_world(self, world: WorldHistoryProtocol) -> TrajectoryData: """Extract trajectory data directly from a world-like object.""" return self.extract(world) def extract_from_history(self, history: Sequence[Any] | Mapping[str, Any]) -> TrajectoryData: """Extract trajectory data directly from a raw history object.""" return self.extract(history) @staticmethod def _as_world_like(value: Any) -> Any | None: """Return value when it appears to be a world-like object.""" if hasattr(value, "history"): return value return None @staticmethod def _history_from_input(value: Any) -> Any: """Return raw history from a world-like object or raw history input.""" if hasattr(value, "history"): return getattr(value, "history") return value def _normalize_history(self, raw_history: Any) -> tuple[Any, ...]: """Normalize raw history into a tuple of snapshots.""" if raw_history is None: return () if isinstance(raw_history, Mapping): if "history" in raw_history: return self._normalize_history(raw_history["history"]) if "snapshots" in raw_history: return self._normalize_history(raw_history["snapshots"]) return (raw_history,) if isinstance(raw_history, Sequence) and not isinstance(raw_history, (str, bytes)): return tuple(raw_history) logger.warning( "TrajectoryExtractor received unsupported history type: %s", raw_history.__class__.__name__, ) return () def _snapshot_from_world(self, world: Any) -> dict[str, Any]: """Build a single best-effort snapshot from a world-like object.""" snapshot_method = getattr(world, "snapshot", None) if callable(snapshot_method): try: snapshot = snapshot_method() if isinstance(snapshot, Mapping): return dict(snapshot) except Exception: logger.debug("world.snapshot() failed during trajectory extraction", exc_info=True) to_dict = getattr(world, "to_dict", None) if callable(to_dict): try: snapshot = to_dict() if isinstance(snapshot, Mapping): return dict(snapshot) except Exception: logger.debug("world.to_dict() failed during trajectory extraction", exc_info=True) snapshot: dict[str, Any] = { "step": getattr(world, "step_count", 0), "agents": getattr(world, "agents", {}), "resources": getattr(world, "resources", {}), } metrics = getattr(world, "metrics", None) if metrics is not None: snapshot["metrics"] = metrics return snapshot def _population_counts_from_snapshot(self, snapshot: Mapping[str, Any]) -> dict[str, int]: """Extract population counts by agent type from a snapshot.""" if "agent_type_counts" in snapshot and isinstance(snapshot["agent_type_counts"], Mapping): return { str(agent_type): int(count) for agent_type, count in snapshot["agent_type_counts"].items() if self._is_numeric(count) } if "population" in snapshot and isinstance(snapshot["population"], Mapping): return { str(agent_type): int(count) for agent_type, count in snapshot["population"].items() if self._is_numeric(count) } agents = self._collection_values( self._first_present( snapshot, ("agents", "agent_states", "agent_snapshots"), default=(), ) ) counts: dict[str, int] = defaultdict(int) for agent in agents: agent_type = self._field(agent, "type", "agent_type", "kind", default="unknown") alive = self._field(agent, "alive", default=True) if self.config.count_alive_only and alive is False: continue counts[str(agent_type)] += 1 return dict(sorted(counts.items())) def _resource_amounts_from_snapshot(self, snapshot: Mapping[str, Any]) -> dict[str, float]: """Extract resource total amounts by resource type from a snapshot.""" for key in ("resource_timeseries", "resource_type_amounts", "resource_amounts"): value = snapshot.get(key) if isinstance(value, Mapping): return { str(resource_type): float(amount) for resource_type, amount in value.items() if self._is_numeric(amount) } resources = self._collection_values( self._first_present( snapshot, ("resources", "resource_states", "resource_snapshots"), default=(), ) ) amounts: dict[str, float] = defaultdict(float) for resource in resources: resource_type = self._field(resource, "type", "resource_type", "kind", default="unknown") amount = self._field(resource, "amount", "quantity", "value", default=0.0) if self._is_numeric(amount): amounts[str(resource_type)] += float(amount) return dict(sorted(amounts.items())) def _append_snapshot_events( self, *, snapshot: Mapping[str, Any], step: int, output: list[dict[str, Any]], ) -> None: """Append normalized event records from a snapshot.""" raw_events = self._first_present( snapshot, ("event_log", "events_processed", "events_fired", "events"), default=(), ) for record in self._collection_values(raw_events): normalized = self._normalize_event_record(record, default_step=step) if normalized: output.append(normalized) def _append_snapshot_metrics( self, *, snapshot: Mapping[str, Any], step: int, output: dict[str, list[dict[str, Any]]], ) -> None: """Append metric records from a snapshot.""" raw_metrics = self._first_present( snapshot, ("metrics", "metric_values", "metric_history"), default=None, ) if raw_metrics is None: return self._append_metric_records(raw_metrics, default_step=step, output=output) def _append_external_events( self, *, event_log: Sequence[Any] | Mapping[str, Any] | None, world: Any | None, output: list[dict[str, Any]], ) -> None: """Append externally supplied or world-level event log records.""" raw_event_log = event_log if raw_event_log is None and world is not None: raw_event_log = self._first_present_from_object( world, ("event_log", "events_processed", "events_fired"), default=None, ) if raw_event_log is None: return for record in self._collection_values(raw_event_log): normalized = self._normalize_event_record(record, default_step=None) if normalized: output.append(normalized) def _append_external_metrics( self, *, metric_history: Mapping[str, Any] | Sequence[Any] | None, world: Any | None, output: dict[str, list[dict[str, Any]]], ) -> None: """Append externally supplied or world-level metric history.""" raw_metric_history = metric_history if raw_metric_history is None and world is not None: raw_metric_history = self._first_present_from_object( world, ("metric_history", "metrics_history", "metrics"), default=None, ) if raw_metric_history is None: return self._append_metric_records(raw_metric_history, default_step=None, output=output) def _append_metric_records( self, raw_metrics: Any, *, default_step: int | None, output: dict[str, list[dict[str, Any]]], ) -> None: """Normalize and append metric records.""" if isinstance(raw_metrics, Mapping): for metric_name, metric_value in raw_metrics.items(): if isinstance(metric_value, Sequence) and not isinstance(metric_value, (str, bytes, Mapping)): for item in metric_value: output[str(metric_name)].append( self._normalize_metric_record( item, metric_name=str(metric_name), default_step=default_step, ) ) else: output[str(metric_name)].append( self._normalize_metric_record( metric_value, metric_name=str(metric_name), default_step=default_step, ) ) return if isinstance(raw_metrics, Sequence) and not isinstance(raw_metrics, (str, bytes)): for item in raw_metrics: item_mapping = self._mapping_from_any(item) metric_name = str( self._first_present( item_mapping, ("name", "metric", "metric_name"), default="unknown_metric", ) ) output[metric_name].append( self._normalize_metric_record( item_mapping, metric_name=metric_name, default_step=default_step, ) ) def _normalize_event_record( self, record: Any, *, default_step: int | None, ) -> dict[str, Any]: """Normalize one event record.""" mapping = self._mapping_from_any(record) if not mapping: return {} if default_step is not None: mapping.setdefault("step", default_step) event_name = self._first_present( mapping, ("name", "event", "event_name", "type"), default=None, ) if event_name is not None: mapping.setdefault("name", str(event_name)) return self._json_safe_mapping(mapping) def _normalize_metric_record( self, record: Any, *, metric_name: str, default_step: int | None, ) -> dict[str, Any]: """Normalize one metric record.""" if isinstance(record, Mapping): output = self._json_safe_mapping(record) else: output = {"value": self._json_safe(record)} output.setdefault("name", metric_name) if default_step is not None: output.setdefault("step", default_step) return output def _preserve_unknown_snapshot_fields( self, snapshot: Mapping[str, Any], *, step: int, ) -> dict[str, Any]: """Preserve selected unknown fields from a snapshot for diagnostics.""" known_keys = { "step", "step_count", "time", "tick", "agents", "agent_states", "agent_snapshots", "agent_type_counts", "population", "resources", "resource_states", "resource_snapshots", "resource_timeseries", "resource_type_amounts", "resource_amounts", "events", "event_log", "events_processed", "events_fired", "metrics", "metric_values", "metric_history", } unknown: dict[str, Any] = {} for key, value in snapshot.items(): if key in known_keys: continue unknown[str(key)] = self._json_safe(value) if not unknown: return {} return { "step": step, "fields": unknown, } @staticmethod def _step_from_snapshot(snapshot: Mapping[str, Any], *, fallback: int) -> int: """Return step value from snapshot with a fallback index.""" for key in ("step", "step_count", "time", "tick"): if key in snapshot and TrajectoryExtractor._is_numeric(snapshot[key]): return int(snapshot[key]) return int(fallback) @staticmethod def _collection_values(value: Any) -> tuple[Any, ...]: """Return collection values from mapping or sequence.""" if value is None: return () if isinstance(value, Mapping): return tuple(value.values()) if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): return tuple(value) return (value,) @staticmethod def _mapping_from_any(value: Any) -> dict[str, Any]: """Return a mapping representation of an arbitrary value.""" if value is None: return {} if isinstance(value, Mapping): return dict(value) for method_name in ("to_dict", "snapshot", "model_dump"): method = getattr(value, method_name, None) if callable(method): try: result = method() if isinstance(result, Mapping): return dict(result) except Exception: logger.debug( "Could not convert %s using %s()", value.__class__.__name__, method_name, exc_info=True, ) output: dict[str, Any] = {} for attribute in ("id", "name", "type", "amount", "position", "alive", "state", "metadata"): if hasattr(value, attribute): output[attribute] = getattr(value, attribute) return output @staticmethod def _field(value: Any, *names: str, default: Any = None) -> Any: """Read the first present field from mapping or object.""" if isinstance(value, Mapping): for name in names: if name in value: return value[name] return default for name in names: if hasattr(value, name): return getattr(value, name) return default @staticmethod def _first_present(mapping: Mapping[str, Any], keys: Sequence[str], *, default: Any) -> Any: """Return first present mapping value.""" for key in keys: if key in mapping and mapping[key] is not None: return mapping[key] return default @staticmethod def _first_present_from_object(value: Any, keys: Sequence[str], *, default: Any) -> Any: """Return first present object attribute value.""" for key in keys: if hasattr(value, key): candidate = getattr(value, key) if candidate is not None: return candidate return default @staticmethod def _is_numeric(value: Any) -> bool: """Return whether value is a finite int or float but not bool.""" if isinstance(value, bool): return False if not isinstance(value, (int, float)): return False return math.isfinite(float(value)) @classmethod def _json_safe_mapping(cls, mapping: Mapping[str, Any]) -> dict[str, Any]: """Return JSON-safe copy of a mapping.""" return { str(key): cls._json_safe(value) for key, value in mapping.items() } @classmethod def _json_safe(cls, value: Any) -> Any: """Return a JSON-friendly representation.""" if value is None or isinstance(value, (str, bool)): return value if isinstance(value, int) and not isinstance(value, bool): return value if isinstance(value, float): return value if math.isfinite(value) else None if isinstance(value, Mapping): return { str(key): cls._json_safe(item) for key, item in value.items() } if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): return [cls._json_safe(item) for item in value] for method_name in ("to_dict", "snapshot", "model_dump"): method = getattr(value, method_name, None) if callable(method): try: return cls._json_safe(method()) except Exception: logger.debug( "Could not JSON-normalize %s using %s()", value.__class__.__name__, method_name, exc_info=True, ) if hasattr(value, "tolist") and callable(value.tolist): try: return cls._json_safe(value.tolist()) except Exception: logger.debug("Could not JSON-normalize tolist() value", exc_info=True) return repr(value) def extract_trajectory( world_or_history: Any, *, metric_history: Mapping[str, Any] | Sequence[Any] | None = None, event_log: Sequence[Any] | Mapping[str, Any] | None = None, config: TrajectoryExtractionConfig | None = None, ) -> TrajectoryData: """Convenience function for extracting trajectory data.""" extractor = TrajectoryExtractor(config=config or TrajectoryExtractionConfig()) return extractor.extract( world_or_history, metric_history=metric_history, event_log=event_log, ) __all__ = [ "TrajectoryData", "TrajectoryExtractionConfig", "TrajectoryExtractor", "WorldHistoryProtocol", "extract_trajectory", ]