Added all brands

app.py

@@ -4,12 +4,54 @@ from pyvis.network import Network
 import pickle
 import math
 
-# Load the graph from a pickle file
+# Dictionary to map brands to their respective pickle files
+BRAND_GRAPHS = {
+    'drumeo': 'drumeo_pop_items_labels.pkl',
+    'pianote': 'pianote_pop_items_labels.pkl',
+    'singeo': 'singeo_pop_items_labels.pkl',
+    'guitareo': 'guitareo_pop_items_labels.pkl'
+}
+
 @st.cache_resource
-def load_graph():
-    with open('drumeo_multi_student_graph_labels.pkl', 'rb') as f:
+def load_graph(brand):
+    """
+    Load the graph for the selected brand.
+    Parameters:
+        brand (str): The brand name corresponding to the graph to load.
+    Returns:
+        nx.DiGraph: The loaded graph.
+    """
+    with open(BRAND_GRAPHS[brand], 'rb') as f:
         return pickle.load(f)
 
+def filter_graph(graph, node_threshold=10, edge_threshold=5):
+    """
+    Filters the graph to include only popular nodes and edges.
+
+    Parameters:
+        graph (nx.DiGraph): The original graph.
+        node_threshold (int): Minimum degree for a node to be included.
+        edge_threshold (int): Minimum weight for an edge to be included.
+
+    Returns:
+        nx.DiGraph: A filtered graph with popular nodes and edges.
+    """
+    # Identify popular nodes based on their degree
+    popular_nodes = [
+        node for node in graph.nodes
+        if graph.degree(node) >= node_threshold
+    ]
+
+    # Create a subgraph with only popular nodes
+    filtered_graph = graph.subgraph(popular_nodes).copy()
+
+    # Remove edges that don't meet the weight threshold
+    for u, v, data in list(filtered_graph.edges(data=True)):
+        if data.get("weight", 0) < edge_threshold:
+            filtered_graph.remove_edge(u, v)
+
+    return filtered_graph
+
 def dynamic_visualize_graph(graph, start_node, layers=3, top_k=5):
     net = Network(notebook=False, width="100%", height="600px", directed=True)
     net.set_options("""
@@ -62,18 +104,58 @@ def dynamic_visualize_graph(graph, start_node, layers=3, top_k=5):
     html_content = net.generate_html()
     st.components.v1.html(html_content, height=600, scrolling=False)
 
+def display_node_info(graph, node_id):
+    """
+    Display all attributes of a node and its edges in the graph.
+
+    Parameters:
+        graph (nx.DiGraph): The graph containing the node.
+        node_id (int or str): The ID of the node to inspect.
+
+    Returns:
+        None
+    """
+    if node_id not in graph:
+        print(f"Node {node_id} does not exist in the graph.")
+        return
+
+    # Display node attributes
+    print(f"Attributes of node {node_id}:")
+    for attr, value in graph.nodes[node_id].items():
+        print(f"  {attr}: {value}")
+
+    # Display incoming edges
+    print(f"\nIncoming edges to node {node_id}:")
+    for u, v, data in graph.in_edges(node_id, data=True):
+        print(f"  From {u} to {v} with attributes: {data}")
+
+    # Display outgoing edges
+    print(f"\nOutgoing edges from node {node_id}:")
+    for u, v, data in graph.out_edges(node_id, data=True):
+        print(f"  From {u} to {v} with attributes: {data}")
+
 
 # Streamlit interface
-st.title("Popular Path Expander")
+st.title("Interactive Graph Expansion with Tooltips")
+
+# Brand Selection
+selected_brand = st.selectbox("Select a brand:", options=list(BRAND_GRAPHS.keys()))
 
-# Load the graph
-G = load_graph()
+# Load the graph for the selected brand
+G = load_graph(selected_brand)
 
+# Filter the graph for popular nodes and edges
+node_degree_threshold = 1  # Minimum degree for nodes
+edge_weight_threshold = 1   # Minimum weight for edges
+G_filtered = filter_graph(G, node_threshold=node_degree_threshold, edge_threshold=edge_weight_threshold)
+#print('spot check degree', G.degree(389062))
+#print('spot check degree', G_filtered.degree(389062))
+#display_node_info(G_filtered, 389062)
 import random
 # Randomly sample a starting node
 # Initialize session state for the starting node
-if "start_node" not in st.session_state:
-    st.session_state.start_node = random.choice(list(G.nodes))  # Randomly select a node once
+#if "start_node" not in st.session_state:
+st.session_state.start_node = random.choice(list(G_filtered.nodes))  # Randomly select a node once
 
 # Input: Starting node
 start_node = st.number_input(
@@ -86,7 +168,7 @@ top_k = st.slider("Branching factor (per node):", 1, 6, value=3)
 
 # Trigger the visualization
 if st.button("Expand Graph"):
-    if start_node in G:
-        dynamic_visualize_graph(G, start_node, layers=layers, top_k=top_k)
+    if start_node in G_filtered:
+        dynamic_visualize_graph(G_filtered, start_node, layers=layers, top_k=top_k)
     else:
         st.error("The starting node is not in the graph!")