Spaces:
Running
Running
| """Learned row-conditional transition head for N4 (Equations 10 and 19).""" | |
| from __future__ import annotations | |
| import math | |
| from bisect import bisect_right | |
| from collections import defaultdict | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| class TransitionHead(nn.Module): | |
| """Shares the belief backbone's query/candidate embeddings; has separate parameters.""" | |
| def __init__(self, hidden_dim: int = 128) -> None: | |
| super().__init__() | |
| self.mlp = nn.Sequential( | |
| nn.Linear(3 * hidden_dim + 4, hidden_dim), | |
| nn.GELU(), | |
| nn.Linear(hidden_dim, 1), | |
| ) | |
| def forward(self, context: torch.Tensor, candidates: torch.Tensor, | |
| horizon_s: torch.Tensor, candidate_mask: torch.Tensor) -> torch.Tensor: | |
| batch, states, hidden = candidates.shape | |
| if context.shape != (batch, hidden) or horizon_s.shape != (batch,): | |
| raise ValueError("transition context or horizon has incorrect shape") | |
| if torch.any(horizon_s < 0): | |
| raise ValueError("transition horizon must be nonnegative") | |
| phase = horizon_s.float() * (2 * math.pi / 86400.0) | |
| time = torch.stack((torch.log1p(horizon_s.float()), | |
| torch.sin(phase), torch.cos(phase), | |
| horizon_s.float() / 86400.0), dim=-1) | |
| features = torch.cat(( | |
| context[:, None, None, :].expand(-1, states, states, -1), | |
| candidates[:, :, None, :].expand(-1, -1, states, -1), | |
| candidates[:, None, :, :].expand(-1, states, -1, -1), | |
| time[:, None, None, :].expand(-1, states, states, -1), | |
| ), dim=-1) | |
| logits = self.mlp(features).squeeze(-1) | |
| logits = logits.masked_fill(~candidate_mask[:, None, :].bool(), -1e4) | |
| probabilities = torch.softmax(logits, dim=-1) | |
| return probabilities * candidate_mask[:, :, None].float() | |
| def transition_nll(kernel: torch.Tensor, source_index: torch.Tensor, | |
| destination_index: torch.Tensor) -> torch.Tensor: | |
| rows = kernel[torch.arange(len(kernel), device=kernel.device), source_index.long()] | |
| selected = rows.gather(1, destination_index.long()[:, None]).squeeze(1) | |
| return -selected.clamp_min(1e-9).log().mean() | |
| def chronological_pairs(*, instance_ids: np.ndarray, world_ids: np.ndarray, | |
| times_s: np.ndarray, states: np.ndarray, | |
| max_horizon_s: float) -> list[tuple[int, int, float, int, int]]: | |
| """Create train-only H<=ta, s(ta), delta, s(ta+delta) labels.""" | |
| groups: dict[tuple[str, int], list[int]] = defaultdict(list) | |
| for index, (instance, world) in enumerate(zip(instance_ids, world_ids, strict=True)): | |
| groups[(str(instance), int(world))].append(index) | |
| pairs = [] | |
| for indices in groups.values(): | |
| ordered = sorted(indices, key=lambda i: float(times_s[i])) | |
| for source, destination in zip(ordered, ordered[1:], strict=False): | |
| horizon = float(times_s[destination] - times_s[source]) | |
| if 0 < horizon <= max_horizon_s: | |
| pairs.append((source, destination, horizon, | |
| int(states[source]), int(states[destination]))) | |
| return pairs | |
| def event_horizon_pairs(*, instance_ids: np.ndarray, world_ids: np.ndarray, | |
| query_times_s: np.ndarray, query_states: np.ndarray, | |
| events: list[dict], horizons_s: tuple[float, ...] | |
| ) -> list[tuple[int, float, float, int, int]]: | |
| """Sample short stationary and event-crossing tuples from one split's causal snapshots.""" | |
| if not horizons_s or any(h <= 0 for h in horizons_s): | |
| raise ValueError("horizons must be positive") | |
| groups: dict[tuple[str, int], list[int]] = defaultdict(list) | |
| event_groups: dict[tuple[str, int], list[dict]] = defaultdict(list) | |
| for index, (instance, world) in enumerate(zip(instance_ids, world_ids, strict=True)): | |
| groups[(str(instance), int(world))].append(index) | |
| for event in events: | |
| event_groups[(str(event["instance_uuid"]), int(event.get("world_id", 0)))].append(event) | |
| pairs: list[tuple[int, float, float, int, int]] = [] | |
| for (instance, world), indices in groups.items(): | |
| ordered = sorted(indices, key=lambda i: float(query_times_s[i])) | |
| times = [float(query_times_s[i]) for i in ordered] | |
| relevant = sorted(event_groups.get((instance, world), []), | |
| key=lambda event: float(event["event_time_s"])) | |
| event_times = [float(event["event_time_s"]) for event in relevant] | |
| for position, source in enumerate(ordered): | |
| horizon = float(horizons_s[position % len(horizons_s)]) | |
| anchor = times[position] | |
| if anchor + horizon > times[-1]: | |
| continue | |
| if bisect_right(event_times, anchor + horizon) == bisect_right(event_times, anchor): | |
| state = int(query_states[source]) | |
| pairs.append((source, anchor, horizon, state, state)) | |
| for event_index, event in enumerate(relevant): | |
| event_time = float(event["event_time_s"]) | |
| for horizon in horizons_s: | |
| anchor = event_time - horizon / 2.0 | |
| future = anchor + horizon | |
| if anchor < times[0] or future > times[-1]: | |
| continue | |
| if (event_index > 0 and event_times[event_index - 1] >= anchor) or ( | |
| event_index + 1 < len(event_times) and event_times[event_index + 1] <= future | |
| ): | |
| continue | |
| source_position = bisect_right(times, anchor) - 1 | |
| if source_position < 0: | |
| continue | |
| pairs.append(( | |
| ordered[source_position], anchor, float(horizon), | |
| int(event["source_state_id"]), int(event["destination_state_id"]), | |
| )) | |
| return pairs | |
| class NeuralTransition: | |
| """Inference adapter that never receives transition labels or mobility tags.""" | |
| dynamic = True | |
| def __init__(self, belief_model: nn.Module, head: TransitionHead, | |
| packed_batch: dict[str, torch.Tensor]) -> None: | |
| self.model = belief_model.eval() | |
| self.head = head.eval() | |
| self.batch = {key: value.clone() for key, value in packed_batch.items()} | |
| if self.batch["candidate_state_ids"].shape[0] != 1: | |
| raise ValueError("N4 transition adapter expects one active episode") | |
| self.elapsed_s = 0.0 | |
| def add_candidate(self, state_id: int) -> None: | |
| present = self.batch["candidate_state_ids"][0].tolist() | |
| if state_id in present: | |
| return | |
| self.batch["candidate_state_ids"] = torch.cat(( | |
| self.batch["candidate_state_ids"], | |
| torch.tensor([[state_id]], dtype=self.batch["candidate_state_ids"].dtype), | |
| ), dim=1) | |
| self.batch["candidate_mask"] = torch.cat(( | |
| self.batch["candidate_mask"], | |
| torch.ones((1, 1), dtype=self.batch["candidate_mask"].dtype), | |
| ), dim=1) | |
| def matrix(self, states: list[int], elapsed_s: float) -> np.ndarray: | |
| if elapsed_s < 0: | |
| raise ValueError("time cannot go backwards") | |
| if elapsed_s == 0: | |
| return np.eye(len(states), dtype=float) | |
| device = next(self.model.parameters()).device | |
| batch = {key: value.to(device) for key, value in self.batch.items()} | |
| batch["query_time_days"] = batch["query_time_days"] + self.elapsed_s / 86400.0 | |
| batch["elapsed_since_last_positive_days"] = ( | |
| batch["elapsed_since_last_positive_days"] + self.elapsed_s / 86400.0 | |
| ) | |
| day_phase = 2 * math.pi * batch["query_time_days"] | |
| batch["query_time_of_day_sin_cos"] = torch.stack( | |
| (day_phase.sin(), day_phase.cos()), dim=-1 | |
| ) | |
| with torch.inference_mode(): | |
| context, candidates = self.model.backbone(batch) | |
| kernel = self.head( | |
| context, candidates, torch.tensor([elapsed_s], device=device), | |
| batch["candidate_mask"], | |
| )[0] | |
| candidate_ids = batch["candidate_state_ids"][0].tolist() | |
| index = [candidate_ids.index(state) for state in states] | |
| return kernel[index][:, index].cpu().numpy() | |
| def advance_clock(self, elapsed_s: float) -> None: | |
| if elapsed_s < 0: | |
| raise ValueError("time cannot go backwards") | |
| self.elapsed_s += elapsed_s | |