Spaces:
Running
Running
File size: 11,171 Bytes
c2ea5ed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 |
#!/usr/bin/env python3
"""
Causal Graph Analysis
This module implements the core causal graph and analysis logic for the multi-agent system.
It handles perturbation propagation and effect calculation.
"""
from collections import defaultdict
import random
import json
import copy
import numpy as np
import os
from typing import Dict, Set, List, Tuple, Any, Optional, Union
class CausalGraph:
"""
Represents the causal graph of the multi-agent system derived from the knowledge graph.
Handles perturbation propagation and effect calculation.
"""
def __init__(self, knowledge_graph: Dict):
self.kg = knowledge_graph
self.entity_ids = [entity["id"] for entity in self.kg["entities"]]
self.relation_ids = [relation["id"] for relation in self.kg["relations"]]
# Extract outcomes and build dependency structure
self.relation_outcomes = {}
self.relation_dependencies = defaultdict(set)
self._build_dependency_graph()
def _build_dependency_graph(self):
"""Build the perturbation dependency graph based on the knowledge graph structure"""
for relation in self.kg["relations"]:
rel_id = relation["id"]
# Get perturbation outcome if available (now supports values between 0 and 1)
# Check for both 'purturbation' (current misspelling) and 'perturbation' (correct spelling)
y = relation.get("purturbation", relation.get("perturbation", relation.get("defense_success_rate", None)))
if y is not None:
# Store the perturbation value (can be any float between 0 and 1)
self.relation_outcomes[rel_id] = float(y)
# Process explicit dependencies
deps = relation.get("dependencies", {})
for dep_rel in deps.get("relations", []):
self.relation_dependencies[dep_rel].add(rel_id)
for dep_ent in deps.get("entities", []):
self.relation_dependencies[dep_ent].add(rel_id)
# Self-dependency: a relation can affect its own outcome
self.relation_dependencies[rel_id].add(rel_id)
# Add source and target entity dependencies automatically
source = relation.get("source", None)
target = relation.get("target", None)
if source:
self.relation_dependencies[source].add(rel_id)
if target:
self.relation_dependencies[target].add(rel_id)
def propagate_effects(self, perturbations: Dict[str, float]) -> Dict[str, float]:
"""
Propagate perturbation effects through the dependency graph.
Args:
perturbations: Dictionary mapping relation/entity IDs to their perturbation values (0-1)
Returns:
Dictionary mapping affected relation IDs to their outcome values
"""
affected_relations = set()
# Find all relations affected by the perturbation
for p in perturbations:
if p in self.relation_dependencies:
affected_relations.update(self.relation_dependencies[p])
# Calculate outcomes for affected relations
outcomes = {}
for rel_id in affected_relations:
if rel_id in self.relation_outcomes:
# If the relation itself is perturbed, use the perturbation value directly
if rel_id in perturbations:
outcomes[rel_id] = perturbations[rel_id]
else:
# Otherwise use the stored outcome value
outcomes[rel_id] = self.relation_outcomes[rel_id]
return outcomes
def calculate_outcome(self, perturbations: Optional[Dict[str, float]] = None) -> float:
"""
Calculate the final outcome score given a set of perturbations.
Args:
perturbations: Dictionary mapping relation/entity IDs to their perturbation values (0-1)
Returns:
Aggregate outcome score
"""
if perturbations is None:
perturbations = {}
affected_outcomes = self.propagate_effects(perturbations)
if not affected_outcomes:
return 0.0
# Aggregate outcomes (simple average for now)
outcome_value = sum(affected_outcomes.values()) / len(affected_outcomes)
return outcome_value
class CausalAnalyzer:
"""
Performs causal effect analysis on the multi-agent knowledge graph system.
Calculates Average Causal Effects (ACE) and Shapley values.
"""
def __init__(self, causal_graph: CausalGraph, n_shapley_samples: int = 200):
self.causal_graph = causal_graph
self.n_shapley_samples = n_shapley_samples
self.base_outcome = self.causal_graph.calculate_outcome({})
def set_perturbation_score(self, relation_id: str, score: float) -> None:
"""
Set the perturbation score for a specific relation ID.
This allows explicitly setting scores from external sources (like database queries).
Args:
relation_id: The ID of the relation to set the score for
score: The perturbation score value (typically between 0 and 1)
"""
# Update the relation_outcomes in the causal graph
self.causal_graph.relation_outcomes[relation_id] = float(score)
def calculate_ace(self) -> Dict[str, float]:
"""
Calculate Average Causal Effect (ACE) for each entity and relation.
Returns:
Dictionary mapping IDs to their ACE scores
"""
ace_scores = {}
# Calculate ACE for relations
for rel_id in self.causal_graph.relation_ids:
if rel_id in self.causal_graph.relation_outcomes:
# Use the actual perturbation value from the outcomes
perturbed_outcome = self.causal_graph.calculate_outcome({rel_id: self.causal_graph.relation_outcomes[rel_id]})
ace_scores[rel_id] = perturbed_outcome - self.base_outcome
else:
# Default to maximum perturbation (1.0) if no value is available
perturbed_outcome = self.causal_graph.calculate_outcome({rel_id: 1.0})
ace_scores[rel_id] = perturbed_outcome - self.base_outcome
# Calculate ACE for entities
for entity_id in self.causal_graph.entity_ids:
# Default to maximum perturbation (1.0) for entities
perturbed_outcome = self.causal_graph.calculate_outcome({entity_id: 1.0})
ace_scores[entity_id] = perturbed_outcome - self.base_outcome
return ace_scores
def calculate_shapley_values(self) -> Dict[str, float]:
"""
Calculate Shapley values to fairly attribute causal effects.
Uses sampling for approximation with larger graphs.
Returns:
Dictionary mapping IDs to their Shapley values
"""
# Combine entities and relations as "players" in the Shapley calculation
all_ids = self.causal_graph.entity_ids + self.causal_graph.relation_ids
shapley_values = {id_: 0.0 for id_ in all_ids}
# Generate random permutations for Shapley approximation
for _ in range(self.n_shapley_samples):
perm = random.sample(all_ids, len(all_ids))
current_set = {} # Empty dictionary instead of empty set
current_outcome = self.base_outcome
for id_ in perm:
# Determine perturbation value to use
if id_ in self.causal_graph.relation_outcomes:
pert_value = self.causal_graph.relation_outcomes[id_]
else:
pert_value = 1.0 # Default to maximum perturbation
# Add current ID to the coalition with its perturbation value
new_set = current_set.copy()
new_set[id_] = pert_value
new_outcome = self.causal_graph.calculate_outcome(new_set)
# Calculate marginal contribution
marginal = new_outcome - current_outcome
shapley_values[id_] += marginal
# Update for next iteration
current_outcome = new_outcome
current_set = new_set
# Normalize the values
for id_ in shapley_values:
shapley_values[id_] /= self.n_shapley_samples
return shapley_values
def analyze(self) -> Tuple[Dict[str, float], Dict[str, float]]:
"""
Perform complete causal analysis.
Returns:
Tuple of (ACE scores, Shapley values)
"""
ace_scores = self.calculate_ace()
shapley_values = self.calculate_shapley_values()
return ace_scores, shapley_values
def enrich_knowledge_graph(kg: Dict, ace_scores: Dict[str, float],
shapley_values: Dict[str, float]) -> Dict:
"""
Enrich the knowledge graph with causal attribution scores.
Args:
kg: Original knowledge graph
ace_scores: Dictionary of ACE scores
shapley_values: Dictionary of Shapley values
Returns:
Enriched knowledge graph
"""
enriched_kg = copy.deepcopy(kg)
# Add scores to entities
for entity in enriched_kg["entities"]:
entity_id = entity["id"]
entity["causal_attribution"] = {
"ACE": ace_scores.get(entity_id, 0),
"Shapley": shapley_values.get(entity_id, 0)
}
# Add scores to relations
for relation in enriched_kg["relations"]:
relation_id = relation["id"]
relation["causal_attribution"] = {
"ACE": ace_scores.get(relation_id, 0),
"Shapley": shapley_values.get(relation_id, 0)
}
return enriched_kg
def generate_summary_report(ace_scores: Dict[str, float],
shapley_values: Dict[str, float],
kg: Dict) -> List[Dict]:
"""
Generate a summary report of causal attributions.
Args:
ace_scores: Dictionary of ACE scores
shapley_values: Dictionary of Shapley values
kg: Knowledge graph
Returns:
List of attribution data for each entity/relation
"""
entity_ids = [entity["id"] for entity in kg["entities"]]
report = []
for id_ in ace_scores:
if id_ in entity_ids:
type_ = "entity"
else:
type_ = "relation"
report.append({
"id": id_,
"ACE": ace_scores.get(id_, 0),
"Shapley": shapley_values.get(id_, 0),
"type": type_
})
# Sort by Shapley value to highlight most important factors
report.sort(key=lambda x: abs(x["Shapley"]), reverse=True)
return report |