update graph builder and covid analysis.py

graph_builder.py

@@ -1,50 +1,97 @@
 """
-Graph builder module for converting GDELT data to graph formats
+Graph builder module for converting GDELT data to graph formats.
+
+This module provides classes for building graphs from GDELT data in various formats,
+including NetworkX, Neo4j, and st-link-analysis. It supports batch processing,
+input validation, custom properties, and logging for robust and efficient graph construction.
 """
 import pandas as pd
 import networkx as nx
+import logging
 import json
+from neo4j import GraphDatabase
+from typing import Dict, List, Optional, Set, Tuple
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 class GraphBuilder:
-    """Base class for building graph from GDELT data"""
-    def process_entities(self, row):
-        """Process entities from a row and return nodes and relationships"""
+    """
+    Base class for building graphs from GDELT data.
+
+    Attributes:
+        ENTITY_MAPPINGS (dict): Mapping of GDELT fields to node types and relationships.
+        logger (Logger): Logger instance for tracking progress and errors.
+    """
+    ENTITY_MAPPINGS = {
+        "V2EnhancedPersons": ("Person", "MENTIONED_IN"),
+        "V2EnhancedOrganizations": ("Organization", "MENTIONED_IN"),
+        "V2EnhancedLocations": ("Location", "LOCATED_IN"),
+        "V2EnhancedThemes": ("Theme", "CATEGORIZED_AS"),
+        "V2.1AllNames": ("Name", "MENTIONED_IN"),
+        "V2.1Counts": ("Count", "MENTIONED_IN"),
+        "V2.1Amounts": ("Amount", "MENTIONED_IN"),
+    }
+
+    def __init__(self):
+        self.logger = logger
+
+    def validate_input(self, df):
+        """
+        Validate input DataFrame for required columns.
+
+        Args:
+            df (pd.DataFrame): Input DataFrame containing GDELT data.
+
+        Raises:
+            ValueError: If any required column is missing.
+        """
+        required_columns = ["GKGRECORDID", "DATE", "SourceCommonName", "DocumentIdentifier"]
+        for col in required_columns:
+            if col not in df.columns:
+                raise ValueError(f"Missing required column: {col}")
+
+    def process_entities(self, row, custom_node_props=None, custom_edge_props=None):
+        """
+        Process entities from a row and return nodes and relationships.
+
+        Args:
+            row (pd.Series): A row of GDELT data.
+            custom_node_props (dict, optional): Custom properties for nodes.
+            custom_edge_props (dict, optional): Custom properties for edges.
+
+        Returns:
+            Tuple[List[Dict], List[Dict]]: Lists of nodes and relationships.
+        """
         nodes = []
         relationships = []
         event_id = row["GKGRECORDID"]
         event_date = row["DATE"]
         event_source = row["SourceCommonName"]
         event_document_id = row["DocumentIdentifier"]
-        # event_image = row["V2.1SharingImage"] if pd.notna(row["V2.1SharingImage"]) else ""
         event_quotations = row["V2.1Quotations"] if pd.notna(row["V2.1Quotations"]) else ""
         event_tone = float(row["tone"]) if pd.notna(row["tone"]) else 0.0
 
-        # Add event node
+        # Add event node with custom properties
+        event_props = {
+            "date": event_date,
+            "source": event_source,
+            "document": event_document_id,
+            "quotations": event_quotations,
+            "tone": event_tone
+        }
+        if custom_node_props:
+            event_props.update(custom_node_props)
+
         nodes.append({
             "id": event_id,
             "type": "event",
-            "properties": {
-                "date": event_date,
-                "source": event_source,
-                "document": event_document_id,
-                # "image": event_image,
-                "quotations": event_quotations,
-                "tone": event_tone
-            }
+            "properties": event_props
         })
 
         # Process each entity type
-        entity_mappings = {
-            "V2EnhancedPersons": ("Person", "MENTIONED_IN"),
-            "V2EnhancedOrganizations": ("Organization", "MENTIONED_IN"),
-            "V2EnhancedLocations": ("Location", "LOCATED_IN"),
-            "V2EnhancedThemes": ("Theme", "CATEGORIZED_AS"),
-            "V2.1AllNames": ("Name", "MENTIONED_IN"),
-            "V2.1Counts": ("Count", "MENTIONED_IN"),
-            "V2.1Amounts": ("Amount", "MENTIONED_IN"),
-        }
-
-        for field, (label, relationship) in entity_mappings.items():
+        for field, (label, relationship) in self.ENTITY_MAPPINGS.items():
             if pd.notna(row[field]):
                 entities = [e.strip() for e in row[field].split(';') if e.strip()]
                 for entity in entities:
@@ -53,168 +100,230 @@ class GraphBuilder:
                         "type": label.lower(),
                         "properties": {"name": entity}
                     })
+                    edge_props = {"created_at": event_date}
+                    if custom_edge_props:
+                        edge_props.update(custom_edge_props)
                     relationships.append({
                         "from": entity,
                         "to": event_id,
                         "type": relationship,
-                        "properties": {"created_at": event_date}
+                        "properties": edge_props
                     })
 
         return nodes, relationships
 
+    def validate_graph(self, G):
+        """
+        Validate the graph for consistency.
+
+        Args:
+            G (nx.Graph): The graph to validate.
+
+        Raises:
+            ValueError: If the graph is invalid.
+        """
+        if not isinstance(G, nx.Graph):
+            raise ValueError("Input is not a valid NetworkX graph.")
+        if len(G.nodes) == 0:
+            raise ValueError("Graph has no nodes.")
+        if len(G.edges) == 0:
+            raise ValueError("Graph has no edges.")
+        self.logger.info("Graph validation passed.")
+
 class NetworkXBuilder(GraphBuilder):
-    """Builder for NetworkX graphs"""
-    def build_graph(self, df):
-        G = nx.Graph()
+    """
+    Builder for NetworkX graphs.
 
-        for _, row in df.iterrows():
-            nodes, relationships = self.process_entities(row)
+    Attributes:
+        directed (bool): Whether to create a directed graph.
+    """
+    def __init__(self, directed=False):
+        super().__init__()
+        self.directed = directed
 
-            # Add nodes
-            for node in nodes:
-                G.add_node(node["id"],
-                          type=node["type"],
-                          **node["properties"])
+    def build_graph(self, df, custom_node_props=None, custom_edge_props=None):
+        """
+        Build a NetworkX graph from the DataFrame.
 
-            # Add relationships
-            for rel in relationships:
-                G.add_edge(rel["from"],
-                          rel["to"],
-                          relationship=rel["type"],
-                          **rel["properties"])
+        Args:
+            df (pd.DataFrame): Input DataFrame containing GDELT data.
+            custom_node_props (dict, optional): Custom properties for nodes.
+            custom_edge_props (dict, optional): Custom properties for edges.
 
-        return G
+        Returns:
+            nx.Graph: The constructed graph.
+        """
+        self.validate_input(df)
+        G = nx.DiGraph() if self.directed else nx.Graph()
+
+        try:
+            for _, row in df.iterrows():
+                nodes, relationships = self.process_entities(row, custom_node_props, custom_edge_props)
+
+                # Add nodes
+                for node in nodes:
+                    G.add_node(node["id"], type=node["type"], **node["properties"])
+
+                # Add relationships
+                for rel in relationships:
+                    G.add_edge(rel["from"], rel["to"], relationship=rel["type"], **rel["properties"])
+
+            self.validate_graph(G)
+            self.logger.info("NetworkX graph built successfully.")
+            return G
+        except Exception as e:
+            self.logger.error(f"Error building NetworkX graph: {str(e)}")
+            raise
 
 class Neo4jBuilder(GraphBuilder):
+    """
+    Builder for Neo4j graphs.
+
+    Attributes:
+        driver (neo4j.Driver): Neo4j driver instance.
+        logger (Logger): Logger instance for tracking progress and errors.
+    """
     def __init__(self, uri, user, password):
+        super().__init__()
         self.driver = GraphDatabase.driver(uri, auth=(user, password))
-        self.logger = logging.getLogger(__name__)
 
     def close(self):
+        """Close the Neo4j driver."""
         self.driver.close()
 
-    def build_graph(self, df):
+    def build_graph(self, df, batch_size=1000, custom_node_props=None, custom_edge_props=None):
+        """
+        Build a Neo4j graph from the DataFrame with batch processing.
+
+        Args:
+            df (pd.DataFrame): Input DataFrame containing GDELT data.
+            batch_size (int): Number of rows to process in each batch.
+            custom_node_props (dict, optional): Custom properties for nodes.
+            custom_edge_props (dict, optional): Custom properties for edges.
+        """
+        self.validate_input(df)
+
         with self.driver.session() as session:
+            batch = []
             for _, row in df.iterrows():
-                nodes, relationships = self.process_entities(row)
-
-                # Create nodes and relationships in Neo4j
-                try:
-                    session.execute_write(self._create_graph_elements,
-                                        nodes, relationships)
-                except Exception as e:
-                    self.logger.error(f"Error processing row {row['GKGRECORDID']}: {str(e)}")
-
-    def _create_graph_elements(self, tx, nodes, relationships):
-        # Create nodes
-        for node in nodes:
-            query = f"""
-            MERGE (n:{node['type']} {{id: $id}})
-            SET n += $properties
-            """
-            tx.run(query, id=node["id"], properties=node["properties"])
-
-        # Create relationships
-        for rel in relationships:
-            query = f"""
-            MATCH (a {{id: $from_id}})
-            MATCH (b {{id: $to_id}})
-            MERGE (a)-[r:{rel['type']}]->(b)
-            SET r += $properties
-            """
-            tx.run(query,
-                  from_id=rel["from"],
-                  to_id=rel["to"],
-                  properties=rel["properties"])
-
-class StreamlitGraphBuilder:
-    """Adapted graph builder for Streamlit visualization"""
-    def __init__(self):
-        self.G = nx.Graph()
-        
-    def process_row(self, row):
-        """Process a single row of data"""
-        event_id = row["GKGRECORDID"]
-        event_props = {
-            "type": "event",  # already in lowercase
-            "date": row["DATE"],
-            "source": row["SourceCommonName"],
-            "document": row["DocumentIdentifier"],
-            "tone": row["tone"],
-            # Store display name in its original format if needed.
-            "name": row["SourceCommonName"]
-        }
-        
-        self.G.add_node(event_id, **event_props)
-        
-        # Use lowercase node types for consistency in lookups.
-        entity_types = {
-            "V2EnhancedPersons": ("person", "MENTIONED_IN"),
-            "V2EnhancedOrganizations": ("organization", "MENTIONED_IN"),
-            "V2EnhancedLocations": ("location", "LOCATED_IN"),
-            "V2EnhancedThemes": ("theme", "CATEGORIZED_AS"),
-            "V2.1AllNames": ("name", "MENTIONED_IN"),
-            "V2.1Counts": ("count", "MENTIONED_IN"),
-            "V2.1Amounts": ("amount", "MENTIONED_IN"),
-        }
-        
-        for col, (node_type, rel_type) in entity_types.items():
-            if pd.notna(row[col]):
-                # The actual display value (which may be in Parent Case) is preserved in the "name" attribute.
-                entities = [e.strip() for e in row[col].split(';') if e.strip()]
-                for entity in entities:
-                    self.G.add_node(entity, type=node_type, name=entity)
-                    self.G.add_edge(entity, event_id, 
-                                    relationship=rel_type, 
-                                    date=row["DATE"])
+                nodes, relationships = self.process_entities(row, custom_node_props, custom_edge_props)
+                batch.append((nodes, relationships))
+                if len(batch) >= batch_size:
+                    session.execute_write(self._create_graph_elements_batch, batch)
+                    batch = []
+            if batch:
+                session.execute_write(self._create_graph_elements_batch, batch)
+
+            self.logger.info("Neo4j graph built successfully.")
+
+    def _create_graph_elements_batch(self, tx, batch):
+        """
+        Create nodes and relationships in Neo4j in batches.
+
+        Args:
+            tx (neo4j.Transaction): Neo4j transaction.
+            batch (List[Tuple[List[Dict], List[Dict]]]): Batch of nodes and relationships.
+        """
+        for nodes, relationships in batch:
+            # Create nodes
+            for node in nodes:
+                query = f"""
+                MERGE (n:{node['type']} {{id: $id}})
+                SET n += $properties
+                """
+                tx.run(query, id=node["id"], properties=node["properties"])
+
+            # Create relationships
+            for rel in relationships:
+                query = f"""
+                MATCH (a {{id: $from_id}})
+                MATCH (b {{id: $to_id}})
+                MERGE (a)-[r:{rel['type']}]->(b)
+                SET r += $properties
+                """
+                tx.run(query, from_id=rel["from"], to_id=rel["to"], properties=rel["properties"])
 
 class StLinkBuilder(GraphBuilder):
-    """Builder for st-link-analysis compatible graphs"""
-    def build_graph(self, df):
-        """Build graph in st-link-analysis format"""
+    """
+    Builder for st-link-analysis compatible graphs.
+
+    Attributes:
+        logger (Logger): Logger instance for tracking progress and errors.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def build_graph(self, df, custom_node_props=None, custom_edge_props=None):
+        """
+        Build graph in st-link-analysis format.
+
+        Args:
+            df (pd.DataFrame): Input DataFrame containing GDELT data.
+            custom_node_props (dict, optional): Custom properties for nodes.
+            custom_edge_props (dict, optional): Custom properties for edges.
+
+        Returns:
+            Dict: Graph data in st-link-analysis format.
+        """
+        self.validate_input(df)
         all_nodes = []
         all_edges = []
         edge_counter = 0
-
-        # Track nodes we've already added to avoid duplicates
         added_nodes = set()
-        
-        for _, row in df.iterrows():
-            nodes, relationships = self.process_entities(row)
-            
-            # Process nodes
-            for node in nodes:
-                if node["id"] not in added_nodes:
-                    stlink_node = {
+
+        try:
+            for _, row in df.iterrows():
+                nodes, relationships = self.process_entities(row, custom_node_props, custom_edge_props)
+
+                # Process nodes
+                for node in nodes:
+                    if node["id"] not in added_nodes:
+                        stlink_node = {
+                            "data": {
+                                "id": str(node["id"]),
+                                "label": node["type"].upper(),
+                                **node["properties"]
+                            }
+                        }
+                        all_nodes.append(stlink_node)
+                        added_nodes.add(node["id"])
+
+                # Process relationships/edges
+                for rel in relationships:
+                    edge_counter += 1
+                    stlink_edge = {
                         "data": {
-                            "id": str(node["id"]),
-                            "label": node["type"].upper(),
-                            **node["properties"]
+                            "id": f"e{edge_counter}",
+                            "source": str(rel["from"]),
+                            "target": str(rel["to"]),
+                            "label": rel["type"],
+                            **rel["properties"]
                         }
                     }
-                    all_nodes.append(stlink_node)
-                    added_nodes.add(node["id"])
-            
-            # Process relationships/edges
-            for rel in relationships:
-                edge_counter += 1
-                stlink_edge = {
-                    "data": {
-                        "id": f"e{edge_counter}",
-                        "source": str(rel["from"]),
-                        "target": str(rel["to"]),
-                        "label": rel["type"],
-                        **rel["properties"]
-                    }
-                }
-                all_edges.append(stlink_edge)
-        
-        return {
-            "nodes": all_nodes,
-            "edges": all_edges
-        }
+                    all_edges.append(stlink_edge)
+
+            self.logger.info("st-link-analysis graph built successfully.")
+            return {
+                "nodes": all_nodes,
+                "edges": all_edges
+            }
+        except Exception as e:
+            self.logger.error(f"Error building st-link-analysis graph: {str(e)}")
+            raise
 
     def write_json(self, graph_data, filename):
-        """Write graph to JSON file"""
-        with open(filename, 'w') as f:
-            json.dump(graph_data, f, indent=2)
+        """
+        Write graph to JSON file with streaming.
+
+        Args:
+            graph_data (Dict): Graph data in st-link-analysis format.
+            filename (str): Output file name.
+        """
+        try:
+            with open(filename, 'w') as f:
+                json.dump(graph_data, f, indent=2)
+            self.logger.info(f"Graph data written to {filename} successfully.")
+        except Exception as e:
+            self.logger.error(f"Error writing JSON file: {str(e)}")
+            raise
+        

pages/2_🌐_COVID_Network_Analysis.py

@@ -13,7 +13,7 @@ from typing import Dict, List, Set, Tuple, Optional
 from pathlib import Path
 
 from data_access import get_gdelt_data, filter_dataframe, GDELT_CATEGORIES
-from graph_builder import StreamlitGraphBuilder
+from graph_builder import NetworkXBuilder  # Updated to use NetworkXBuilder
 from graph_config import NODE_TYPES
 
 # Type aliases for clarity
@@ -252,12 +252,10 @@ def main():
             themes_filter=themes_filter
         )
         
-        # Build graph
+        # Build graph using NetworkXBuilder
         with st.spinner("Building knowledge graph..."):
-            builder = StreamlitGraphBuilder()
-            for _, row in df.iterrows():
-                builder.process_row(row)
-            G = builder.G
+            builder = NetworkXBuilder()  # Use NetworkXBuilder
+            G = builder.build_graph(df)  # Build graph from DataFrame
         
         if G.number_of_nodes() == 0:
             st.warning("No data found matching the specified criteria.")