Spaces:

shield137
/

shockmap-api

Sleeping

shockmap-api / backend /app /services /graph_service.py

Dhruva P Gowda

Deploy ShockMap MVP: High-fidelity UI overhaul, animated graph particles, live GDELT ingestion, and production scaffolding

f59a149 6 months ago

Raw History Blame Contribute Delete

25.2 kB

"""
Engine 2 — Supply Chain Knowledge Graph with PageRank + Louvain Community Detection.
Maintains a directed graph of drugs, APIs, rare earths, and production provinces.
Provides:
  - Personalized PageRank for shock propagation (replaces naive BFS)
  - Louvain community detection for supply cluster identification
  - HHI concentration scoring
  - Combined risk formula: R_i = PR_i × (1 - S_i) × e^(-B_i/τ) × C_community
Implements:
  - GNN-based shock propagation concepts from AAAI 2025
  - Causal ML supply chain intervention from arXiv 2024
"""

import math
import logging
from typing import List, Dict, Optional, Any, Literal, Tuple

import networkx as nx

logger = logging.getLogger("shockmap.graph")

# ─── Sector Lead-Time Constants (τ in days) ──────────────────────────────────
SECTOR_TAU = {
    "pharma": 30,
    "rare_earth": 60,
    "semiconductor": 45,
    "ev_battery": 50,
    "solar": 40,
    "defence": 90,
}


class GraphService:
    """
    Engine 2: Maintains a directed graph representation of the multi-sector
    supply chain with PageRank propagation and Louvain community detection.
    """

    def __init__(self, data_loader: Any):
        """
        Initializes the GraphService and builds the DiGraph from seed data.
        Args:
            data_loader: The data provider for drugs, APIs, rare earths, and dependencies.
        """
        self.data_loader = data_loader
        self.graph = nx.DiGraph()
        self._undirected = None  # For Louvain (needs undirected graph)
        self._communities = {}   # node_id → community_id
        self._community_labels = {}  # community_id → descriptive label
        self._community_detection_enabled = True
        self._community_detection_method = "louvain"
        self._pagerank_cache = {}
        self._build_graph()
        self._detect_communities()

    # ─── Graph Construction ──────────────────────────────────────────────

    def _build_graph(self):
        """Constructs the NetworkX graph from data_loader seed data."""
        # Add Drug nodes
        for drug in self.data_loader.get_drugs():
            self.graph.add_node(
                drug.id,
                type="drug",
                name=drug.name,
                sector="pharma",
                nlem_tier=drug.nlem_tier,
                patient_population_estimate=drug.patient_population_estimate,
                has_substitute=drug.has_substitute,
                therapeutic_class=drug.therapeutic_class,
                current_risk=drug.current_risk or 0.0,
                # Engine 2 enrichments
                substitutability=self._drug_substitutability(drug),
                buffer_days=self._estimate_buffer_days(drug),
                strategic_priority=1.0 if drug.nlem_tier.value == "TIER_1" else 0.5,
            )

        # Add API nodes and collect unique provinces
        provinces = set()
        for api in self.data_loader.get_apis():
            self.graph.add_node(
                api.id,
                type="api",
                name=api.name,
                sector="pharma",
                china_share=api.china_share,
                primary_provinces=api.primary_provinces,
                monthly_import_value_usd_millions=api.monthly_import_value_usd_millions,
                substitutability=self._api_substitutability(api),
                buffer_days=max(5, int(30 * (1 - api.china_share))),
            )
            for p in api.primary_provinces:
                provinces.add(p)

        # Add Province nodes
        for p_name in provinces:
            self.graph.add_node(
                p_name,
                type="province",
                name=p_name,
                country="China",
                sector="multi",
            )

        # Add Edges with trade_volume_share
        for edge in self.data_loader.get_dependencies():
            self.graph.add_edge(
                edge.source,
                edge.target,
                weight=edge.weight,
                edge_type=edge.edge_type,
                trade_volume_share=edge.weight,  # normalized 0-1
            )

        # Add Rare Earth nodes (from rare_earths.json if available)
        self._add_rare_earth_nodes()

        logger.info(
            f"Graph built: {self.graph.number_of_nodes()} nodes, "
            f"{self.graph.number_of_edges()} edges"
        )

    def _add_rare_earth_nodes(self):
        """Adds rare earth mineral nodes and their province edges."""
        try:
            import json
            import os
            seed_dir = self.data_loader.seed_dir
            re_path = os.path.join(seed_dir, "rare_earths.json")
            if not os.path.exists(re_path):
                return
            with open(re_path, "r", encoding="utf-8") as f:
                rare_earths = json.load(f)

            for re in rare_earths:
                node_id = re["id"]
                self.graph.add_node(
                    node_id,
                    type="input",
                    name=re["name"],
                    sector="rare_earth",
                    china_share=re.get("china_share", 0.9),
                    primary_provinces=re.get("primary_provinces", []),
                    monthly_import_value_usd_millions=re.get("monthly_import_value_usd_millions", 0),
                    end_use_sectors=re.get("end_use_sectors", []),
                    substitutability=max(0.0, 1.0 - re.get("china_share", 0.9)),
                    buffer_days=re.get("india_stockpile_days", 15),
                    strategic_priority=1.0 if re.get("criticality", 0) > 0.8 else 0.7,
                    criticality=re.get("criticality", 0.5),
                )

                # Create edges to provinces
                for prov in re.get("primary_provinces", []):
                    if not self.graph.has_node(prov):
                        self.graph.add_node(
                            prov, type="province", name=prov,
                            country="China", sector="multi"
                        )
                    share = re.get("china_share", 0.9) / max(1, len(re.get("primary_provinces", [prov])))
                    self.graph.add_edge(
                        node_id, prov,
                        weight=share,
                        edge_type="input_from_province",
                        trade_volume_share=share,
                    )
        except Exception as e:
            logger.warning(f"Failed to load rare earth nodes: {e}")

    # ─── Louvain Community Detection ─────────────────────────────────────

    def _detect_communities(self):
        """
        Runs Louvain community detection on the undirected version of the graph.
        Identifies supply chain clusters — groups of nodes that are tightly
        interconnected and likely to co-propagate shocks.
        """
        try:
            import community as community_louvain

            self._undirected = self.graph.to_undirected()
            self._communities = community_louvain.best_partition(
                self._undirected, random_state=42
            )

            # Build community summaries
            community_members: Dict[int, List[str]] = {}
            for node_id, comm_id in self._communities.items():
                community_members.setdefault(comm_id, []).append(node_id)

            # Label each community by its dominant province or sector
            for comm_id, members in community_members.items():
                provinces = [m for m in members if self.graph.nodes[m].get("type") == "province"]
                if provinces:
                    self._community_labels[comm_id] = f"{', '.join(provinces[:2])} Supply Cluster"
                else:
                    sectors = [self.graph.nodes[m].get("sector", "unknown") for m in members[:3]]
                    self._community_labels[comm_id] = f"{'|'.join(set(sectors))} Cluster {comm_id}"

            logger.info(
                f"Louvain detected {len(community_members)} communities "
                f"across {len(self._communities)} nodes"
            )

        except ImportError:
            logger.warning("python-louvain not installed. Community detection disabled.")
            self._community_detection_enabled = False
            self._community_detection_method = "disabled_missing_dependency"
            self._communities = {n: 0 for n in self.graph.nodes()}
        except Exception as e:
            logger.warning(f"Community detection failed: {e}")
            self._community_detection_enabled = False
            self._community_detection_method = "disabled_runtime_error"
            self._communities = {n: 0 for n in self.graph.nodes()}

    def get_communities(self) -> Dict[str, Any]:
        """Returns community detection results for the frontend."""
        community_data: Dict[int, Dict[str, Any]] = {}

        for node_id, comm_id in self._communities.items():
            if comm_id not in community_data:
                community_data[comm_id] = {
                    "id": comm_id,
                    "label": self._community_labels.get(comm_id, f"Cluster {comm_id}"),
                    "members": [],
                    "provinces": [],
                    "drugs": [],
                    "inputs": [],
                    "size": 0,
                }
            entry = community_data[comm_id]
            node = self.graph.nodes[node_id]
            node_type = node.get("type", "unknown")
            entry["members"].append(node_id)
            entry["size"] += 1
            if node_type == "province":
                entry["provinces"].append(node_id)
            elif node_type == "drug":
                entry["drugs"].append(node_id)
            elif node_type in ("api", "input"):
                entry["inputs"].append(node_id)

        return {
            "community_detection_enabled": self._community_detection_enabled,
            "community_detection_method": self._community_detection_method,
            "total_communities": len(community_data),
            "communities": list(community_data.values()),
        }

    def get_node_community(self, node_id: str) -> Optional[int]:
        """Returns the community ID for a given node."""
        return self._communities.get(node_id)

    def get_community_amplifier(self, community_id: int) -> float:
        """
        Returns the community amplifier C_community.
        Larger, more interconnected communities amplify shocks more.
        """
        members = [n for n, c in self._communities.items() if c == community_id]
        if not members:
            return 1.0
        # Amplifier based on community density (more edges = more amplification)
        subgraph = self._undirected.subgraph(members) if self._undirected else self.graph.subgraph(members)
        density = nx.density(subgraph) if len(members) > 1 else 0
        return 1.0 + (density * 0.5)  # Range: 1.0 to 1.5

    # ─── Personalized PageRank ───────────────────────────────────────────

    def compute_pagerank(
        self,
        shocked_node: Optional[str] = None,
        alpha: float = 0.85,
    ) -> Dict[str, float]:
        """
        Computes Personalized PageRank from a shock origin node.
        IMPORTANT: Since graph edges point drug→api→province, we run PageRank
        on the REVERSED graph when propagating from a province. This makes
        the random walk flow province→api→drug, correctly modeling how supply
        shocks propagate upstream to find affected end products.
        Args:
            shocked_node: The node where the shock originates (e.g., "Hebei")
            alpha: Damping factor (default 0.85)
        """
        cache_key = shocked_node or "__global__"
        if cache_key in self._pagerank_cache:
            return self._pagerank_cache[cache_key]

        personalization = None
        if shocked_node and self.graph.has_node(shocked_node):
            personalization = {shocked_node: 1.0}

        # Use reversed graph so PageRank flows from province → api → drug
        # This correctly models shock propagation through the supply chain
        target_graph = self.graph.reverse(copy=False)

        try:
            scores = nx.pagerank(
                target_graph,
                alpha=alpha,
                weight="trade_volume_share",
                personalization=personalization,
                max_iter=200,
            )
            self._pagerank_cache[cache_key] = scores
            return scores
        except nx.PowerIterationFailedConvergence:
            logger.warning("PageRank failed to converge, using uniform distribution")
            n = self.graph.number_of_nodes()
            return {node: 1.0 / n for node in self.graph.nodes()}

    def compute_province_risk(self, province_id: str, sector: str = "pharma") -> float:
        """
        Calculates the aggregated risk score for a Chinese province.
        Risk is based on:
          - Concentration (HHI) of the products produced in this province.
          - Total monthly import value dependent on this province.
          - Strategic priority of the dependent products.
        """
        if not self.graph.has_node(province_id):
            return 0.0

        dependent_products = []
        # Find APIs/Inputs that are LOCATED_IN or supplied by this province
        # In our graph, edges are drug -> api -> province, so predecessors(province) are the APIs
        for product_id in self.graph.predecessors(province_id):
            node_data = self.graph.nodes.get(product_id, {})
            if sector == "both" or node_data.get("sector") == sector:
                dependent_products.append(product_id)

        if not dependent_products:
            return 0.0

        total_risk = 0.0
        for product_id in dependent_products:
            node_data = self.graph.nodes.get(product_id, {})
            import_value = node_data.get("monthly_import_value_usd_millions", 1.0)
            china_share = node_data.get("china_share", 0.5)
            # High import value + High China share = High Dependency Risk
            product_risk = (import_value * china_share)
            total_risk += product_risk

        # Normalize to 0-100 scale (heuristic normalization)
        normalized_risk = min(100.0, math.log1p(total_risk) * 15)
        return round(normalized_risk, 2)

    # ─── Combined Risk Formula ───────────────────────────────────────────

    def compute_combined_risk(
        self,
        shocked_node: str,
        sector: str = "pharma",
    ) -> Dict[str, Dict[str, Any]]:
        """
        Computes the full risk score using the combined formula:
        R_i = PR_i(shock) × (1 - S_i) × e^(-B_i/τ) × C_community
        Where:
          PR_i  = Personalized PageRank from shock origin (on reversed graph)
          S_i   = Substitutability (0-1)
          B_i   = Buffer stock in days
          τ     = Sector lead time (pharma=30d, rare_earth=60d)
          C     = Community amplifier
        Returns a dict of {node_id: {score, components}} for all affected nodes.
        """
        tau = SECTOR_TAU.get(sector, 30)
        pr_scores = self.compute_pagerank(shocked_node)

        # Find max PR score for normalization (excluding the shocked node itself)
        pr_values = [v for k, v in pr_scores.items() if k != shocked_node and v > 0]
        pr_max = max(pr_values) if pr_values else 1.0

        results = {}
        for node_id, pr_score in pr_scores.items():
            if node_id == shocked_node:
                continue  # Skip the shock origin itself

            node_data = self.graph.nodes.get(node_id, {})
            node_type = node_data.get("type", "")

            # Only score drugs, APIs, and inputs — not provinces
            if node_type not in ("drug", "api", "input"):
                continue

            substitutability = node_data.get("substitutability", 0.3)
            buffer_days = max(1, node_data.get("buffer_days", 15))
            community_id = self._communities.get(node_id, 0)
            community_amp = self.get_community_amplifier(community_id)

            # Normalize PR score relative to max (so top node → ~1.0)
            pr_normalized = pr_score / pr_max if pr_max > 0 else 0

            # The combined formula
            non_substitutability = 1.0 - substitutability
            buffer_decay = math.exp(-buffer_days / tau)
            risk_raw = pr_normalized * non_substitutability * buffer_decay * community_amp

            # Scale to 0-100
            risk_score_normalized = min(100.0, risk_raw * 100)

            if risk_score_normalized > 0.5:  # Include anything meaningful
                results[node_id] = {
                    "risk_score": round(risk_score_normalized, 2),
                    "components": {
                        "pagerank": round(pr_score, 6),
                        "pagerank_normalized": round(pr_normalized, 4),
                        "non_substitutability": round(non_substitutability, 3),
                        "buffer_decay": round(buffer_decay, 3),
                        "community_amplifier": round(community_amp, 3),
                        "community_id": community_id,
                        "community_label": self._community_labels.get(community_id, ""),
                        "buffer_days": buffer_days,
                        "substitutability": round(substitutability, 3),
                    },
                    "type": node_type,
                    "name": node_data.get("name", node_id),
                    "sector": node_data.get("sector", sector),
                }

        # Sort by risk score descending
        results = dict(
            sorted(results.items(), key=lambda x: x[1]["risk_score"], reverse=True)
        )

        return results

    # ─── Existing Methods (Preserved) ────────────────────────────────────

    def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
        """Returns node data in a flat dictionary format."""
        if not self.graph.has_node(node_id):
            return None

        node_data = self.graph.nodes[node_id]
        return {
            "id": node_id,
            "type": node_data.get("type"),
            "name": node_data.get("name"),
            "attributes": {k: v for k, v in node_data.items() if k not in ["type", "name"]},
        }

    def get_neighbors(
        self,
        node_id: str,
        direction: Literal["in", "out", "both"] = "out",
    ) -> List[str]:
        """Returns the IDs of neighboring nodes based on direction."""
        if not self.graph.has_node(node_id):
            return []

        if direction == "out":
            return list(self.graph.successors(node_id))
        elif direction == "in":
            return list(self.graph.predecessors(node_id))
        else:
            return list(nx.all_neighbors(self.graph, node_id))

    def get_dependency_chain(self, drug_id: str) -> List[List[str]]:
        """
        Finds all simple paths of length 1-3 starting from a drug.
        Typically captures drug -> api -> province paths.
        """
        if not self.graph.has_node(drug_id):
            return []

        all_paths = []
        # Find paths up to length 3 (Drug -> API -> Province)
        for target in self.graph.nodes():
            if target == drug_id:
                continue
            # Use short path limits to capture the direct supply chain
            paths = list(nx.all_simple_paths(self.graph, source=drug_id, target=target, cutoff=3))
            all_paths.extend(paths)

        return all_paths

    def compute_concentration_hhi(self, drug_id: str) -> float:
        """
        Calculates the Herfindahl-Hirschman Index (HHI) for a drug's geographic dependency.
        Ranges from 0 to 10000. High values indicate extreme concentration.
        """
        if not self.graph.has_node(drug_id):
            return 0.0

        # 1. Map drug to provinces with effective weights
        province_weights: Dict[str, float] = {}

        # Get APIs
        for api_id in self.graph.successors(drug_id):
            edge_da = self.graph[drug_id][api_id]
            w_da = edge_da.get("weight", 1.0)

            # Get Provinces for this API
            for prov_id in self.graph.successors(api_id):
                edge_ap = self.graph[api_id][prov_id]
                w_ap = edge_ap.get("weight", 1.0)

                # Cumulative dependency share on this province
                share = w_da * w_ap
                province_weights[prov_id] = province_weights.get(prov_id, 0.0) + share

        if not province_weights:
            return 0.0

        # 2. Calculate HHI
        total_weight = sum(province_weights.values())
        if total_weight == 0:
            return 0.0

        hhi = 0.0
        for weight in province_weights.values():
            percentage = (weight / total_weight) * 100
            hhi += percentage ** 2

        return hhi

    def get_drugs_dependent_on_province(self, province_id: str) -> List[str]:
        """
        Returns a list of drug IDs that are dependent on the given province
        by traversing the graph in reverse.
        """
        if not self.graph.has_node(province_id):
            return []

        dependent_drugs = set()
        # Traverse backwards: Province <- API <- Drug
        for api_id in self.graph.predecessors(province_id):
            for drug_id in self.graph.predecessors(api_id):
                node_data = self.graph.nodes[drug_id]
                if node_data.get("type") == "drug":
                    dependent_drugs.add(drug_id)

        # Also check rare earth inputs
        for input_id in self.graph.predecessors(province_id):
            node_data = self.graph.nodes.get(input_id, {})
            if node_data.get("type") == "input":
                dependent_drugs.add(input_id)

        return list(dependent_drugs)

    def nodes_by_type(self, node_type: str) -> List[Dict[str, Any]]:
        """Returns all nodes of a specific type (e.g., 'api', 'province')."""
        return [
            self.get_node(n)
            for n, d in self.graph.nodes(data=True)
            if d.get("type") == node_type
        ]

    def to_serializable_dict(self) -> Dict[str, Any]:
        """Converts the graph structure to a JSON-serializable dictionary."""
        nodes = []
        for n, d in self.graph.nodes(data=True):
            node_entry = {
                "id": n,
                "type": d.get("type"),
                "name": d.get("name"),
                "attributes": {k: v for k, v in d.items() if k not in ["type", "name"]},
            }
            # Add community info
            if n in self._communities:
                node_entry["community_id"] = self._communities[n]
            nodes.append(node_entry)

        edges = []
        for u, v, d in self.graph.edges(data=True):
            edges.append({
                "source": u,
                "target": v,
                "weight": d.get("weight"),
                "edge_type": d.get("edge_type"),
            })

        return {"nodes": nodes, "edges": edges}

    def apply_risk_scores(self, risk_map: Dict[str, float]) -> None:
        """
        Updates the 'current_risk' attribute on drug nodes based on a map of {drug_id: risk_score}.
        """
        for drug_id, risk in risk_map.items():
            if self.graph.has_node(drug_id):
                self.graph.nodes[drug_id]["current_risk"] = risk

    # ─── Private Helpers ─────────────────────────────────────────────────

    @staticmethod
    def _drug_substitutability(drug: Any) -> float:
        """Estimates substitutability score for a drug (0=no alternatives, 1=many)."""
        if drug.has_substitute:
            return 0.6
        if drug.nlem_tier.value == "TIER_1":
            return 0.1  # Critical, hard to replace
        return 0.3

    @staticmethod
    def _api_substitutability(api: Any) -> float:
        """Estimates substitutability for an API based on China share."""
        # Higher China share = fewer alternatives = lower substitutability
        return max(0.0, 1.0 - api.china_share)

    @staticmethod
    def _estimate_buffer_days(drug: Any) -> int:
        """Estimates buffer stock in days based on drug tier and population."""
        if drug.nlem_tier.value == "TIER_1":
            return 10  # High-demand, low buffer
        elif drug.nlem_tier.value == "TIER_2":
            return 20
        return 30

    def invalidate_cache(self):
        """Clears the PageRank cache (call after graph modifications)."""
        self._pagerank_cache.clear()

	"""
	Engine 2 — Supply Chain Knowledge Graph with PageRank + Louvain Community Detection.

	Maintains a directed graph of drugs, APIs, rare earths, and production provinces.
	Provides:
	- Personalized PageRank for shock propagation (replaces naive BFS)
	- Louvain community detection for supply cluster identification
	- HHI concentration scoring
	- Combined risk formula: R_i = PR_i × (1 - S_i) × e^(-B_i/τ) × C_community

	Implements:
	- GNN-based shock propagation concepts from AAAI 2025
	- Causal ML supply chain intervention from arXiv 2024
	"""

	import math
	import logging
	from typing import List, Dict, Optional, Any, Literal, Tuple

	import networkx as nx

	logger = logging.getLogger("shockmap.graph")

	# ─── Sector Lead-Time Constants (τ in days) ──────────────────────────────────
	SECTOR_TAU = {
	"pharma": 30,
	"rare_earth": 60,
	"semiconductor": 45,
	"ev_battery": 50,
	"solar": 40,
	"defence": 90,
	}


	class GraphService:
	"""
	Engine 2: Maintains a directed graph representation of the multi-sector
	supply chain with PageRank propagation and Louvain community detection.
	"""

	def __init__(self, data_loader: Any):
	"""
	Initializes the GraphService and builds the DiGraph from seed data.

	Args:
	data_loader: The data provider for drugs, APIs, rare earths, and dependencies.
	"""
	self.data_loader = data_loader
	self.graph = nx.DiGraph()
	self._undirected = None # For Louvain (needs undirected graph)
	self._communities = {} # node_id → community_id
	self._community_labels = {} # community_id → descriptive label
	self._community_detection_enabled = True
	self._community_detection_method = "louvain"
	self._pagerank_cache = {}
	self._build_graph()
	self._detect_communities()

	# ─── Graph Construction ──────────────────────────────────────────────

	def _build_graph(self):
	"""Constructs the NetworkX graph from data_loader seed data."""
	# Add Drug nodes
	for drug in self.data_loader.get_drugs():
	self.graph.add_node(
	drug.id,
	type="drug",
	name=drug.name,
	sector="pharma",
	nlem_tier=drug.nlem_tier,
	patient_population_estimate=drug.patient_population_estimate,
	has_substitute=drug.has_substitute,
	therapeutic_class=drug.therapeutic_class,
	current_risk=drug.current_risk or 0.0,
	# Engine 2 enrichments
	substitutability=self._drug_substitutability(drug),
	buffer_days=self._estimate_buffer_days(drug),
	strategic_priority=1.0 if drug.nlem_tier.value == "TIER_1" else 0.5,
	)

	# Add API nodes and collect unique provinces
	provinces = set()
	for api in self.data_loader.get_apis():
	self.graph.add_node(
	api.id,
	type="api",
	name=api.name,
	sector="pharma",
	china_share=api.china_share,
	primary_provinces=api.primary_provinces,
	monthly_import_value_usd_millions=api.monthly_import_value_usd_millions,
	substitutability=self._api_substitutability(api),
	buffer_days=max(5, int(30 * (1 - api.china_share))),
	)
	for p in api.primary_provinces:
	provinces.add(p)

	# Add Province nodes
	for p_name in provinces:
	self.graph.add_node(
	p_name,
	type="province",
	name=p_name,
	country="China",
	sector="multi",
	)

	# Add Edges with trade_volume_share
	for edge in self.data_loader.get_dependencies():
	self.graph.add_edge(
	edge.source,
	edge.target,
	weight=edge.weight,
	edge_type=edge.edge_type,
	trade_volume_share=edge.weight, # normalized 0-1
	)

	# Add Rare Earth nodes (from rare_earths.json if available)
	self._add_rare_earth_nodes()

	logger.info(
	f"Graph built: {self.graph.number_of_nodes()} nodes, "
	f"{self.graph.number_of_edges()} edges"
	)

	def _add_rare_earth_nodes(self):
	"""Adds rare earth mineral nodes and their province edges."""
	try:
	import json
	import os
	seed_dir = self.data_loader.seed_dir
	re_path = os.path.join(seed_dir, "rare_earths.json")
	if not os.path.exists(re_path):
	return
	with open(re_path, "r", encoding="utf-8") as f:
	rare_earths = json.load(f)

	for re in rare_earths:
	node_id = re["id"]
	self.graph.add_node(
	node_id,
	type="input",
	name=re["name"],
	sector="rare_earth",
	china_share=re.get("china_share", 0.9),
	primary_provinces=re.get("primary_provinces", []),
	monthly_import_value_usd_millions=re.get("monthly_import_value_usd_millions", 0),
	end_use_sectors=re.get("end_use_sectors", []),
	substitutability=max(0.0, 1.0 - re.get("china_share", 0.9)),
	buffer_days=re.get("india_stockpile_days", 15),
	strategic_priority=1.0 if re.get("criticality", 0) > 0.8 else 0.7,
	criticality=re.get("criticality", 0.5),
	)

	# Create edges to provinces
	for prov in re.get("primary_provinces", []):
	if not self.graph.has_node(prov):
	self.graph.add_node(
	prov, type="province", name=prov,
	country="China", sector="multi"
	)
	share = re.get("china_share", 0.9) / max(1, len(re.get("primary_provinces", [prov])))
	self.graph.add_edge(
	node_id, prov,
	weight=share,
	edge_type="input_from_province",
	trade_volume_share=share,
	)
	except Exception as e:
	logger.warning(f"Failed to load rare earth nodes: {e}")

	# ─── Louvain Community Detection ─────────────────────────────────────

	def _detect_communities(self):
	"""
	Runs Louvain community detection on the undirected version of the graph.
	Identifies supply chain clusters — groups of nodes that are tightly
	interconnected and likely to co-propagate shocks.
	"""
	try:
	import community as community_louvain

	self._undirected = self.graph.to_undirected()
	self._communities = community_louvain.best_partition(
	self._undirected, random_state=42
	)

	# Build community summaries
	community_members: Dict[int, List[str]] = {}
	for node_id, comm_id in self._communities.items():
	community_members.setdefault(comm_id, []).append(node_id)

	# Label each community by its dominant province or sector
	for comm_id, members in community_members.items():
	provinces = [m for m in members if self.graph.nodes[m].get("type") == "province"]
	if provinces:
	self._community_labels[comm_id] = f"{', '.join(provinces[:2])} Supply Cluster"
	else:
	sectors = [self.graph.nodes[m].get("sector", "unknown") for m in members[:3]]
	self._community_labels[comm_id] = f"{'\|'.join(set(sectors))} Cluster {comm_id}"

	logger.info(
	f"Louvain detected {len(community_members)} communities "
	f"across {len(self._communities)} nodes"
	)

	except ImportError:
	logger.warning("python-louvain not installed. Community detection disabled.")
	self._community_detection_enabled = False
	self._community_detection_method = "disabled_missing_dependency"
	self._communities = {n: 0 for n in self.graph.nodes()}
	except Exception as e:
	logger.warning(f"Community detection failed: {e}")
	self._community_detection_enabled = False
	self._community_detection_method = "disabled_runtime_error"
	self._communities = {n: 0 for n in self.graph.nodes()}

	def get_communities(self) -> Dict[str, Any]:
	"""Returns community detection results for the frontend."""
	community_data: Dict[int, Dict[str, Any]] = {}

	for node_id, comm_id in self._communities.items():
	if comm_id not in community_data:
	community_data[comm_id] = {
	"id": comm_id,
	"label": self._community_labels.get(comm_id, f"Cluster {comm_id}"),
	"members": [],
	"provinces": [],
	"drugs": [],
	"inputs": [],
	"size": 0,
	}
	entry = community_data[comm_id]
	node = self.graph.nodes[node_id]
	node_type = node.get("type", "unknown")
	entry["members"].append(node_id)
	entry["size"] += 1
	if node_type == "province":
	entry["provinces"].append(node_id)
	elif node_type == "drug":
	entry["drugs"].append(node_id)
	elif node_type in ("api", "input"):
	entry["inputs"].append(node_id)

	return {
	"community_detection_enabled": self._community_detection_enabled,
	"community_detection_method": self._community_detection_method,
	"total_communities": len(community_data),
	"communities": list(community_data.values()),
	}

	def get_node_community(self, node_id: str) -> Optional[int]:
	"""Returns the community ID for a given node."""
	return self._communities.get(node_id)

	def get_community_amplifier(self, community_id: int) -> float:
	"""
	Returns the community amplifier C_community.
	Larger, more interconnected communities amplify shocks more.
	"""
	members = [n for n, c in self._communities.items() if c == community_id]
	if not members:
	return 1.0
	# Amplifier based on community density (more edges = more amplification)
	subgraph = self._undirected.subgraph(members) if self._undirected else self.graph.subgraph(members)
	density = nx.density(subgraph) if len(members) > 1 else 0
	return 1.0 + (density * 0.5) # Range: 1.0 to 1.5

	# ─── Personalized PageRank ───────────────────────────────────────────

	def compute_pagerank(
	self,
	shocked_node: Optional[str] = None,
	alpha: float = 0.85,
	) -> Dict[str, float]:
	"""
	Computes Personalized PageRank from a shock origin node.

	IMPORTANT: Since graph edges point drug→api→province, we run PageRank
	on the REVERSED graph when propagating from a province. This makes
	the random walk flow province→api→drug, correctly modeling how supply
	shocks propagate upstream to find affected end products.

	Args:
	shocked_node: The node where the shock originates (e.g., "Hebei")
	alpha: Damping factor (default 0.85)
	"""
	cache_key = shocked_node or "__global__"
	if cache_key in self._pagerank_cache:
	return self._pagerank_cache[cache_key]

	personalization = None
	if shocked_node and self.graph.has_node(shocked_node):
	personalization = {shocked_node: 1.0}

	# Use reversed graph so PageRank flows from province → api → drug
	# This correctly models shock propagation through the supply chain
	target_graph = self.graph.reverse(copy=False)

	try:
	scores = nx.pagerank(
	target_graph,
	alpha=alpha,
	weight="trade_volume_share",
	personalization=personalization,
	max_iter=200,
	)
	self._pagerank_cache[cache_key] = scores
	return scores
	except nx.PowerIterationFailedConvergence:
	logger.warning("PageRank failed to converge, using uniform distribution")
	n = self.graph.number_of_nodes()
	return {node: 1.0 / n for node in self.graph.nodes()}

	def compute_province_risk(self, province_id: str, sector: str = "pharma") -> float:
	"""
	Calculates the aggregated risk score for a Chinese province.
	Risk is based on:
	- Concentration (HHI) of the products produced in this province.
	- Total monthly import value dependent on this province.
	- Strategic priority of the dependent products.
	"""
	if not self.graph.has_node(province_id):
	return 0.0

	dependent_products = []
	# Find APIs/Inputs that are LOCATED_IN or supplied by this province
	# In our graph, edges are drug -> api -> province, so predecessors(province) are the APIs
	for product_id in self.graph.predecessors(province_id):
	node_data = self.graph.nodes.get(product_id, {})
	if sector == "both" or node_data.get("sector") == sector:
	dependent_products.append(product_id)

	if not dependent_products:
	return 0.0

	total_risk = 0.0
	for product_id in dependent_products:
	node_data = self.graph.nodes.get(product_id, {})
	import_value = node_data.get("monthly_import_value_usd_millions", 1.0)
	china_share = node_data.get("china_share", 0.5)
	# High import value + High China share = High Dependency Risk
	product_risk = (import_value * china_share)
	total_risk += product_risk

	# Normalize to 0-100 scale (heuristic normalization)
	normalized_risk = min(100.0, math.log1p(total_risk) * 15)
	return round(normalized_risk, 2)

	# ─── Combined Risk Formula ───────────────────────────────────────────

	def compute_combined_risk(
	self,
	shocked_node: str,
	sector: str = "pharma",
	) -> Dict[str, Dict[str, Any]]:
	"""
	Computes the full risk score using the combined formula:

	R_i = PR_i(shock) × (1 - S_i) × e^(-B_i/τ) × C_community

	Where:
	PR_i = Personalized PageRank from shock origin (on reversed graph)
	S_i = Substitutability (0-1)
	B_i = Buffer stock in days
	τ = Sector lead time (pharma=30d, rare_earth=60d)
	C = Community amplifier

	Returns a dict of {node_id: {score, components}} for all affected nodes.
	"""
	tau = SECTOR_TAU.get(sector, 30)
	pr_scores = self.compute_pagerank(shocked_node)

	# Find max PR score for normalization (excluding the shocked node itself)
	pr_values = [v for k, v in pr_scores.items() if k != shocked_node and v > 0]
	pr_max = max(pr_values) if pr_values else 1.0

	results = {}
	for node_id, pr_score in pr_scores.items():
	if node_id == shocked_node:
	continue # Skip the shock origin itself

	node_data = self.graph.nodes.get(node_id, {})
	node_type = node_data.get("type", "")

	# Only score drugs, APIs, and inputs — not provinces
	if node_type not in ("drug", "api", "input"):
	continue

	substitutability = node_data.get("substitutability", 0.3)
	buffer_days = max(1, node_data.get("buffer_days", 15))
	community_id = self._communities.get(node_id, 0)
	community_amp = self.get_community_amplifier(community_id)

	# Normalize PR score relative to max (so top node → ~1.0)
	pr_normalized = pr_score / pr_max if pr_max > 0 else 0

	# The combined formula
	non_substitutability = 1.0 - substitutability
	buffer_decay = math.exp(-buffer_days / tau)
	risk_raw = pr_normalized * non_substitutability * buffer_decay * community_amp

	# Scale to 0-100
	risk_score_normalized = min(100.0, risk_raw * 100)

	if risk_score_normalized > 0.5: # Include anything meaningful
	results[node_id] = {
	"risk_score": round(risk_score_normalized, 2),
	"components": {
	"pagerank": round(pr_score, 6),
	"pagerank_normalized": round(pr_normalized, 4),
	"non_substitutability": round(non_substitutability, 3),
	"buffer_decay": round(buffer_decay, 3),
	"community_amplifier": round(community_amp, 3),
	"community_id": community_id,
	"community_label": self._community_labels.get(community_id, ""),
	"buffer_days": buffer_days,
	"substitutability": round(substitutability, 3),
	},
	"type": node_type,
	"name": node_data.get("name", node_id),
	"sector": node_data.get("sector", sector),
	}

	# Sort by risk score descending
	results = dict(
	sorted(results.items(), key=lambda x: x[1]["risk_score"], reverse=True)
	)

	return results

	# ─── Existing Methods (Preserved) ────────────────────────────────────

	def get_node(self, node_id: str) -> Optional[Dict[str, Any]]:
	"""Returns node data in a flat dictionary format."""
	if not self.graph.has_node(node_id):
	return None

	node_data = self.graph.nodes[node_id]
	return {
	"id": node_id,
	"type": node_data.get("type"),
	"name": node_data.get("name"),
	"attributes": {k: v for k, v in node_data.items() if k not in ["type", "name"]},
	}

	def get_neighbors(
	self,
	node_id: str,
	direction: Literal["in", "out", "both"] = "out",
	) -> List[str]:
	"""Returns the IDs of neighboring nodes based on direction."""
	if not self.graph.has_node(node_id):
	return []

	if direction == "out":
	return list(self.graph.successors(node_id))
	elif direction == "in":
	return list(self.graph.predecessors(node_id))
	else:
	return list(nx.all_neighbors(self.graph, node_id))

	def get_dependency_chain(self, drug_id: str) -> List[List[str]]:
	"""
	Finds all simple paths of length 1-3 starting from a drug.
	Typically captures drug -> api -> province paths.
	"""
	if not self.graph.has_node(drug_id):
	return []

	all_paths = []
	# Find paths up to length 3 (Drug -> API -> Province)
	for target in self.graph.nodes():
	if target == drug_id:
	continue
	# Use short path limits to capture the direct supply chain
	paths = list(nx.all_simple_paths(self.graph, source=drug_id, target=target, cutoff=3))
	all_paths.extend(paths)

	return all_paths

	def compute_concentration_hhi(self, drug_id: str) -> float:
	"""
	Calculates the Herfindahl-Hirschman Index (HHI) for a drug's geographic dependency.
	Ranges from 0 to 10000. High values indicate extreme concentration.
	"""
	if not self.graph.has_node(drug_id):
	return 0.0

	# 1. Map drug to provinces with effective weights
	province_weights: Dict[str, float] = {}

	# Get APIs
	for api_id in self.graph.successors(drug_id):
	edge_da = self.graph[drug_id][api_id]
	w_da = edge_da.get("weight", 1.0)

	# Get Provinces for this API
	for prov_id in self.graph.successors(api_id):
	edge_ap = self.graph[api_id][prov_id]
	w_ap = edge_ap.get("weight", 1.0)

	# Cumulative dependency share on this province
	share = w_da * w_ap
	province_weights[prov_id] = province_weights.get(prov_id, 0.0) + share

	if not province_weights:
	return 0.0

	# 2. Calculate HHI
	total_weight = sum(province_weights.values())
	if total_weight == 0:
	return 0.0

	hhi = 0.0
	for weight in province_weights.values():
	percentage = (weight / total_weight) * 100
	hhi += percentage ** 2

	return hhi

	def get_drugs_dependent_on_province(self, province_id: str) -> List[str]:
	"""
	Returns a list of drug IDs that are dependent on the given province
	by traversing the graph in reverse.
	"""
	if not self.graph.has_node(province_id):
	return []

	dependent_drugs = set()
	# Traverse backwards: Province <- API <- Drug
	for api_id in self.graph.predecessors(province_id):
	for drug_id in self.graph.predecessors(api_id):
	node_data = self.graph.nodes[drug_id]
	if node_data.get("type") == "drug":
	dependent_drugs.add(drug_id)

	# Also check rare earth inputs
	for input_id in self.graph.predecessors(province_id):
	node_data = self.graph.nodes.get(input_id, {})
	if node_data.get("type") == "input":
	dependent_drugs.add(input_id)

	return list(dependent_drugs)

	def nodes_by_type(self, node_type: str) -> List[Dict[str, Any]]:
	"""Returns all nodes of a specific type (e.g., 'api', 'province')."""
	return [
	self.get_node(n)
	for n, d in self.graph.nodes(data=True)
	if d.get("type") == node_type
	]

	def to_serializable_dict(self) -> Dict[str, Any]:
	"""Converts the graph structure to a JSON-serializable dictionary."""
	nodes = []
	for n, d in self.graph.nodes(data=True):
	node_entry = {
	"id": n,
	"type": d.get("type"),
	"name": d.get("name"),
	"attributes": {k: v for k, v in d.items() if k not in ["type", "name"]},
	}
	# Add community info
	if n in self._communities:
	node_entry["community_id"] = self._communities[n]
	nodes.append(node_entry)

	edges = []
	for u, v, d in self.graph.edges(data=True):
	edges.append({
	"source": u,
	"target": v,
	"weight": d.get("weight"),
	"edge_type": d.get("edge_type"),
	})

	return {"nodes": nodes, "edges": edges}

	def apply_risk_scores(self, risk_map: Dict[str, float]) -> None:
	"""
	Updates the 'current_risk' attribute on drug nodes based on a map of {drug_id: risk_score}.
	"""
	for drug_id, risk in risk_map.items():
	if self.graph.has_node(drug_id):
	self.graph.nodes[drug_id]["current_risk"] = risk

	# ─── Private Helpers ─────────────────────────────────────────────────

	@staticmethod
	def _drug_substitutability(drug: Any) -> float:
	"""Estimates substitutability score for a drug (0=no alternatives, 1=many)."""
	if drug.has_substitute:
	return 0.6
	if drug.nlem_tier.value == "TIER_1":
	return 0.1 # Critical, hard to replace
	return 0.3

	@staticmethod
	def _api_substitutability(api: Any) -> float:
	"""Estimates substitutability for an API based on China share."""
	# Higher China share = fewer alternatives = lower substitutability
	return max(0.0, 1.0 - api.china_share)

	@staticmethod
	def _estimate_buffer_days(drug: Any) -> int:
	"""Estimates buffer stock in days based on drug tier and population."""
	if drug.nlem_tier.value == "TIER_1":
	return 10 # High-demand, low buffer
	elif drug.nlem_tier.value == "TIER_2":
	return 20
	return 30

	def invalidate_cache(self):
	"""Clears the PageRank cache (call after graph modifications)."""
	self._pagerank_cache.clear()