ChokePointAI / src /chokepoint /graph /simulator.py
KateShkuryna
feat: NetworkX graph engine, shock propagation, GDELT ingestion
9580dda
Raw History Blame Contribute Delete
2.32 kB
from __future__ import annotations
from collections import deque
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
import networkx as nx
if TYPE_CHECKING:
Graph = nx.DiGraph[str, dict[str, Any], dict[str, Any]]
@dataclass
class NodeState:
impact: float
eta: float
hops: int
path: list[str] = field(default_factory=list)
def substitution_factor(graph: Graph, node_id: str, state: dict[str, NodeState]) -> float:
"""Reduce impact when a node still has healthy substitute nodes."""
substitutes: list[str] = graph.nodes[node_id].get("substitutes", [])
if not substitutes:
return 1.0
unaffected = sum(1 for substitute in substitutes if substitute not in state)
return 1.0 - 0.5 * (unaffected / len(substitutes))
def propagate(
graph: Graph,
epicentres: dict[str, float],
*,
max_hops: int = 4,
decay: float = 0.75,
min_impact: float = 0.02,
) -> dict[str, NodeState]:
"""Propagate impact with a deterministic, cycle-safe weighted BFS."""
state: dict[str, NodeState] = {
node_id: NodeState(impact=impact, eta=0.0, hops=0, path=[node_id])
for node_id, impact in epicentres.items()
}
frontier: deque[str] = deque(epicentres)
while frontier:
source = frontier.popleft()
source_state = state[source]
if source_state.hops >= max_hops:
continue
for target in graph.successors(source):
edge = graph.edges[source, target]
node = graph.nodes[target]
contribution = source_state.impact * edge["weight"] * decay ** (source_state.hops + 1)
contribution *= 1.0 - node.get("resilience", 0.0)
contribution *= substitution_factor(graph, target, state)
if contribution < min_impact:
continue
candidate = NodeState(
impact=min(1.0, contribution),
eta=source_state.eta + edge.get("lead_time_days", 0),
hops=source_state.hops + 1,
path=[*source_state.path, target],
)
previous = state.get(target)
if previous is None or candidate.impact > previous.impact:
state[target] = candidate
frontier.append(target)
return state