REFACTOR: refactor code + add doc

app.py

@@ -1,241 +1,121 @@
-from concrete import fhe
-import folium
-import geopandas
-from shapely.geometry import Point
 import streamlit as st
-from streamlit_folium import st_folium
-
+from shapely.geometry import Point
+import geopandas
 import config 
+from utils import generate_path, set_up_server, set_up_client, display_encrypted, add_marker,\
+                  display_map, init_session, add_to_server_side, add_to_client_side, display_client_side,\
+                  display_server_side, restart_session, compute_shortest_path, decrypt_shortest_path
 from network import get_frames
-from utils import generate_path, set_up_server, set_up_client, display_encrypted
-
 
-st.set_page_config(layout="wide")
 ways = geopandas.read_file(config.roads_filepath)
 nodes, _, rel = get_frames(ways)
 
 server = set_up_server()
 client = set_up_client(server.client_specs.serialize())
 
-m = nodes.explore(
-    color="red",  
-    marker_kwds=dict(radius=5, fill=True, name='node_id'),  
-    tooltip="node_id",  
-    tooltip_kwds=dict(labels=False), 
-    zoom_control=False,
-)
-
-
-c1, c2, c3 = st.columns([1, 3, 1])
+c1, c2, c3 = init_session()
 
 with c1:
-    st.write("**Client-side**")
+    display_client_side()
 
 with c3:
-    st.write("**Server-side**")
-
+    display_server_side()
 
+# keys generation view
 if 'evaluation_key' not in st.session_state:
     with c1:
         if st.button('Generate keys'):
             with st.spinner('Generating keys'):
                 client.keys.load_if_exists_generate_and_save_otherwise(config.keys_filepath)
-                # client.keys.save(config.keys_filepath)
                 st.session_state['evaluation_key'] = client.evaluation_keys.serialize()
+                add_to_client_side("Encryption/decryption keys and evaluation keys are generated.")
+                add_to_client_side("The evaluation key is sent to the server.")
+                add_to_server_side(f"Evaluation key: {display_encrypted(st.session_state['evaluation_key'])}")
             st.rerun()
 else:
-    with c1:
-        st.write("Encryption/decryption keys and evaluation keys are generated.")
-        st.write("The evaluation key is sent to the server.")
-    with c3:
-        st.write(f"Evaluation key: {display_encrypted(st.session_state['evaluation_key'])}")
-
-            
+    # Origin selection view
     if 'origin' not in st.session_state :
         with c1:
             st.write("Select the origin on the map")
-
         with c2:
-            st_data_origin = st_folium(m, width=725, key="origin", returned_objects=["last_object_clicked"])
-
-        with c3:
-            st.write("")
+            st_data_origin = display_map(nodes, returned_objects=["last_object_clicked"])
             
         if st_data_origin["last_object_clicked"]:
             origin = Point(st_data_origin["last_object_clicked"]["lng"], st_data_origin["last_object_clicked"]["lat"])
+            add_marker(origin, 'origin')
             origin_node = nodes[nodes['geometry'] == origin]['node_id'].values[0]
             st.session_state['origin'] = origin
             st.session_state['origin_node'] = origin_node
-            st.write(f"Selected node is node number: {origin_node}")  
+            add_to_client_side(f"Selected origin is node number: {origin_node}")  
             st.rerun()
 
+    # Destination selection view
     if 'origin' in st.session_state and 'destination' not in st.session_state:
-        origin = st.session_state['origin']
-        folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)    
         origin_node = st.session_state['origin_node']
         with c1:
-            st.write(f"Selected origin is node number: {origin_node}")   
             st.write("Select the destination on the map")         
         with c2:
-            st_data_destination = st_folium(m, width=725, key="destination", returned_objects=["last_object_clicked"])
-        with c3:
-            st.write("")
-            st.write("")
+            st_data_destination = display_map(nodes, returned_objects=["last_object_clicked"])
 
         if st_data_destination["last_object_clicked"]:
             destination = Point(st_data_destination["last_object_clicked"]["lng"], st_data_destination["last_object_clicked"]["lat"])
+            add_marker(destination, 'destination')
+            destination_node = nodes[nodes['geometry'] == destination]['node_id'].values[0]
             st.session_state['destination'] = destination
-            st.session_state['destination_node'] = nodes[nodes['geometry'] == destination]['node_id'].values[0]
-            st.write(st_data_destination["last_object_clicked"])
+            st.session_state['destination_node'] = destination_node
+            add_to_client_side(f"Selected destination is node number: {destination_node}")
             st.rerun()
 
-    if 'origin' in st.session_state and 'destination' in st.session_state and 'encrypted_origin' not in st.session_state :
-        origin = st.session_state['origin']
-        folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)
-        destination = st.session_state['destination']
-        folium.Marker([destination.y, destination.x], popup="Destination", tooltip="Destination").add_to(m)
-        # breakpoint()
+    # Origin/Destination encryption view
+    if 'destination' in st.session_state and 'encrypted_origin' not in st.session_state :
         origin_node = st.session_state['origin_node']
         destination_node = st.session_state['destination_node']
         with c1:
-            st.write(f"Selected origin is node number: {origin_node}")   
-            st.write(f"Selected destination is node number: {destination_node}")   
             if st.button('Encrypt and send inputs'):
                 with st.spinner('Encrypting inputs'):
                     client.keys.load(config.keys_filepath)
                     origin, destination = client.encrypt(origin_node, destination_node)
                     st.session_state['encrypted_origin'] = origin.serialize()
                     st.session_state['encrypted_destination'] = destination.serialize()
+                    add_to_server_side(f"Received origin: {display_encrypted(st.session_state['encrypted_origin'])}")
+                    add_to_server_side(f"Received destination: {display_encrypted(st.session_state['encrypted_destination'])}")
                 st.rerun()
         with c2:
-            st_data_final = st_folium(m, width=725, key="destination", returned_objects=[])
-        with c3:
-            st.write("")
-            st.write("")
+            st_data_final = display_map(nodes)
 
+    # Shortest path computation view
     if 'encrypted_origin' in st.session_state and 'encrypted_shortest_path' not in st.session_state:
-        origin = st.session_state['origin']
-        folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)
-        destination = st.session_state['destination']
-        folium.Marker([destination.y, destination.x], popup="Destination", tooltip="Destination").add_to(m)
-        
         origin_node = st.session_state['origin_node']
         destination_node = st.session_state['destination_node']
-        with c1:
-            st.write(f"Selected origin is node number: {origin_node}")   
-            st.write(f"Selected destination is node number: {destination_node}")
         with c2:
-            origin = st.session_state['origin']
-            folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)  
-            destination = st.session_state['destination']
-            folium.Marker([destination.y, destination.x], popup="Destination", tooltip="Destination").add_to(m)
-            st_data_final = st_folium(m, width=725, key="destination", returned_objects=[])
+            st_data_final = display_map(nodes)
         with c3:
-            st.write("")    
-            st.write("")    
-            st.write(f"Received origin: {display_encrypted(st.session_state['encrypted_origin'])}")
-            st.write(f"Received destination: {display_encrypted(st.session_state['encrypted_destination'])}")
-            if st.button('Compute shortest path'):
+            if st.button('Compute and send shortest path'):
                 with st.spinner('Computing'):
-                    deserialized_origin = fhe.Value.deserialize(st.session_state['encrypted_origin'])
-                    deserialized_destination = fhe.Value.deserialize(st.session_state['encrypted_destination'])
-                    deserialized_evaluation_keys = fhe.EvaluationKeys.deserialize(st.session_state['evaluation_key'])
-                    client.keys.load_if_exists_generate_and_save_otherwise(config.keys_filepath)
-                    origin = st.session_state['origin_node']
-                    destination = st.session_state['destination_node']
-                    path = [origin, ]
-                    encrypted_path = [st.session_state['encrypted_origin'], ]
-                    o, d = deserialized_origin, deserialized_destination
-                    for _ in range(nodes.shape[0] - 1):
-                        # Careful: breaking early could lead to information leak
-                        if origin == destination:
-                            break
-                        o = server.run(o, d, evaluation_keys=client.evaluation_keys)
-                        origin = client.decrypt(o)
-                        encrypted_path.append(o.serialize())
-                        path.append(origin)
-                        
+                    encrypted_path = compute_shortest_path(nodes.shape[0], client, server)
                     st.session_state['encrypted_shortest_path'] = encrypted_path
-                    st.session_state['decrypted_shortest_path'] = path
+                    add_to_client_side("Received path is:")
+                    add_to_client_side(f"{display_encrypted(st.session_state['encrypted_shortest_path'][0])}")
+                    add_to_client_side("...")
+                    add_to_client_side(f"{display_encrypted(st.session_state['encrypted_shortest_path'][-1])}")
                 st.rerun()
-
-    if 'encrypted_shortest_path' in st.session_state and 'decrypted_result' not in st.session_state :  
-        origin = st.session_state['origin']
-        folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)
-        destination = st.session_state['destination']
-        folium.Marker([destination.y, destination.x], popup="Destination", tooltip="Destination").add_to(m)
-        
-        origin_node = st.session_state['origin_node']
-        destination_node = st.session_state['destination_node']
+                
+    # Result decryption view
+    if 'encrypted_shortest_path' in st.session_state and 'decrypted_result' not in st.session_state :
         with c1:
-            st.write(f"Selected origin is node number: {origin_node}")   
-            st.write(f"Selected destination is node number: {destination_node}") 
-            st.write("Received path is:")
-            st.write(f"{display_encrypted(st.session_state['encrypted_shortest_path'][0])}")
-            st.write("...")
-            st.write(f"{display_encrypted(st.session_state['encrypted_shortest_path'][-1])}")
-            
             if st.button('Decrypt and show shortest path'):
                 with st.spinner('Computing'):
-                    client.keys.load_if_exists_generate_and_save_otherwise(config.keys_filepath)
-                    path = []
-                    for enc_value in st.session_state['encrypted_shortest_path']:
-                        deserialized_result = fhe.Value.deserialize(enc_value)
-                        path.append(client.decrypt(deserialized_result))
-                    st.session_state['decrypted_result'] = path
+                    decrypt_shortest_path(client)
                 st.rerun()
         with c2:
+            st_data_final = display_map(nodes)
             
-            origin = st.session_state['origin']
-            folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)  
-            
-            destination = st.session_state['destination']
-            folium.Marker([destination.y, destination.x], popup="Destination", tooltip="Destination").add_to(m)
-            st_data_final = st_folium(m, width=725, key="destination", returned_objects=[])
-        with c3:
-            st.write("")    
-            st.write("")    
-            st.write(f"Received origin: {display_encrypted(st.session_state['encrypted_origin'])}")
-            st.write(f"Received destination: {display_encrypted(st.session_state['encrypted_destination'])}")
-            # st.write(f"Next node is {display_encrypted(st.session_state['encrypted_shortest_path'])} and is sent to client")
+    # Display result view
     if 'decrypted_result' in st.session_state :
-        origin = st.session_state['origin']
-        folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)
-        destination = st.session_state['destination']
-        folium.Marker([destination.y, destination.x], popup="Destination", tooltip="Destination").add_to(m)
-        
-        origin_node = st.session_state['origin_node']
-        destination_node = st.session_state['destination_node']
         with c1:
-            st.write(f"Selected origin is node number: {origin_node}")   
-            st.write(f"Selected destination is node number: {destination_node}") 
-            st.write("Received path is:")
-            st.write(f"{display_encrypted(st.session_state['encrypted_shortest_path'][0])}")
-            st.write("...")
-            st.write(f"{display_encrypted(st.session_state['encrypted_shortest_path'][-1])}")
             st.write(f"Decrypted result is: {st.session_state['decrypted_result']}")
-            if st.button('Restart'):
-                for key in st.session_state.keys():
-                    del st.session_state[key]
-                st.rerun() 
+            restart_session()
         with c2:
-            origin = st.session_state['origin']
-            folium.Marker([origin.y, origin.x], popup="Origin", tooltip="Origin").add_to(m)  
-            destination = st.session_state['destination']
-            folium.Marker([destination.y, destination.x], popup="Destination", tooltip="Destination").add_to(m)
             shortest_path_list = st.session_state['decrypted_result']
             final_result = generate_path(shortest_path_list, rel)
-            final_result.explore(
-                m=m,  
-                color="green",
-                style_kwds = {"weight":5},
-                tooltip="name",
-                popup=["name"],
-                name="Quadratic-Paris", 
-            )
-            st_data_final = st_folium(m, width=725, key="destination", returned_objects=[])
-        with c3:
-            st.write("")    
-            st.write("")    
-            st.write(f"Received origin: {display_encrypted(st.session_state['encrypted_origin'])}")
-            st.write(f"Received destination: {display_encrypted(st.session_state['encrypted_destination'])}")
+            display_map(nodes, path=final_result)

network.py

@@ -4,6 +4,17 @@ import numpy
 
 
 def get_frames(ways: geopandas.GeoDataFrame):
+    """Extract nodes, edges, and relations from a GeoDataFrame of ways.
+
+    Args:
+        ways (geopandas.GeoDataFrame): A GeoDataFrame containing information about ways.
+
+    Returns:
+        nodes (geopandas.GeoDataFrame): A GeoDataFrame containing node geometries.
+        edges (geopandas.GeoDataFrame): A GeoDataFrame containing ways information.
+        rel (geopandas.GeoDataFrame): A GeoDataFrame representing relations between nodes and edges.
+
+    """
     edges = ways.explode(index_parts=False)[["id", "name", "geometry"]]
     edges.index.name = "way_id"
     edges = edges.reset_index()
@@ -15,6 +26,14 @@ def get_frames(ways: geopandas.GeoDataFrame):
 
 
 def get(ways: geopandas.GeoDataFrame) -> networkx.Graph:
+    """ Convert a GeoDataFrame of ways into a networkx Graph.
+
+    Args:
+        ways (geopandas.GeoDataFrame): A GeoDataFrame containing information about ways.
+
+    Returns:
+        networkx.Graph: A networkx Graph representing the geographical network.
+    """
     nodes, edges, rel = get_frames(ways)
     graph = networkx.Graph()
     for node in rel['node_id'].unique():
@@ -22,7 +41,6 @@ def get(ways: geopandas.GeoDataFrame) -> networkx.Graph:
     for idx in rel['way_id'].unique():
         node_id1, node_id2 = rel.loc[idx]['node_id']
         way = edges.loc[idx]
-        way_length = rel.loc[idx]['geometry'].iloc[0].length
         graph.add_edge(
             node_id1, node_id2,
             weight=way.geometry.length,
@@ -33,6 +51,14 @@ def get(ways: geopandas.GeoDataFrame) -> networkx.Graph:
 
 
 def weighted_adjacency_matrix(graph: networkx.Graph) -> numpy.ndarray:
+    """Generate the weighted adjacency matrix from a networkx Graph.
+
+    Args:
+        graph (networkx.Graph):  A networkx Graph representing the network.
+
+    Returns:
+        numpy.ndarray: A NumPy array representing the weighted adjacency matrix.
+    """
     weights = numpy.full((graph.number_of_nodes(),)*2, numpy.inf)
     numpy.fill_diagonal(weights, 0., wrap=False)
     for i in graph:

utils.py

@@ -1,47 +1,220 @@
+# Contains a set of useful functions
 import pandas as pd
 from concrete import fhe
-from config import circuit_filepath
-
+from config import circuit_filepath, keys_filepath
+import streamlit as st
+import folium
+from streamlit_folium import st_folium
 
 def set_up_server():
-    # Setting up a server
+    """Load a server instance from a specified circuit file
+
+    Raises:
+        OSError: If there is an issue loading the FHE server.
+
+    Returns:
+       concrete.fhe.compilation.server.Server: A server instance loaded from the circuit file.
+    """
     try:
         server = fhe.Server.load(circuit_filepath)
-    except Exception:
-        raise Exception(f"Something went wrong with the circuit. Make sure that the circuit exists in {circuit_filepath}. If not run python generate_circuit.py.") 
+    except OSError as e:
+        raise OSError(f"Something went wrong with the circuit. Make sure that the circuit \
+            exists in {circuit_filepath}.If not run python generate_circuit.py.") from e
 
     return server
 
 def set_up_client(serialized_client_specs):
-    # Setting up client
+    """Generate a client instance from a specified circuit file
+
+    Args:
+        serialized_client_specs (bytes): A serialized client specs
+
+    Returns:
+        concrete.fhe.compilation.client.Client: A client instance created from the client specs
+    """
+    
     client_specs = fhe.ClientSpecs.deserialize(serialized_client_specs)
     client = fhe.Client(client_specs)
 
     return client
 
 def display_encrypted(encrypted_object):
+    """Display a truncated representation of an encrypted object as a hexadecimal string
+
+    Args:
+        encrypted_object (bytes): A serialized encrypted object to display
+
+    Returns:
+        str: A truncated hexadecimal representation of the encrypted object
+    """
     encoded_text = encrypted_object.hex()
     res = '...' + encoded_text[-10:]
     return res
 
-def shortest_path(circuit, N_NODES, origin, destination):
-    path = [origin, ]
-    o, d = circuit.encrypt(origin, destination)
-    for _ in range(N_NODES - 1):
-        # Careful: breaking early could lead to information leak
+def compute_shortest_path(nodes_nb, client, server):
+    """Calculate the shortest path between two nodes
+
+    Args:
+        nodes_nb (int): The number of nodes in the network
+        client (concrete.fhe.compilation.client.Client): A client instance
+        server (concrete.fhe.compilation.server.Server): A server instance
+
+    Returns:
+        List[bytes]: A list of encrypted values representing the path
+    """
+    deserialized_origin = fhe.Value.deserialize(st.session_state['encrypted_origin'])
+    deserialized_destination = fhe.Value.deserialize(st.session_state['encrypted_destination'])
+    deserialized_evaluation_keys = fhe.EvaluationKeys.deserialize(st.session_state['evaluation_key'])
+    client.keys.load_if_exists_generate_and_save_otherwise(keys_filepath)
+    origin = st.session_state['origin_node']
+    destination = st.session_state['destination_node']
+    encrypted_path = [st.session_state['encrypted_origin'], ]
+    o, d = deserialized_origin, deserialized_destination
+    for _ in range(nodes_nb):
         if origin == destination:
             break
-        o = circuit.run(o, d)
-        origin = circuit.decrypt(o)
-        path.append(origin)
-    return path
+        o = server.run(o, d, evaluation_keys=deserialized_evaluation_keys)
+        origin = client.decrypt(o)
+        encrypted_path.append(o.serialize())
+    return encrypted_path
 
-def generate_path(shortest_path_list, rel):
+    
+def generate_path(shortest_path, roads):
+    """Generate a path from a list of nodes using road data.
+
+    Args:
+        shortest_path (List[int]): A list of nodes representing the shortest path.
+        roads (geopandas.DataFrame): A DataFrame containing the ways between the nodes.
+
+    Returns:
+       pandas.DataFrame: A DataFrame representing the road segments that form the shortest path.
+    """
     pairs_list = []
-    for i in range(len(shortest_path_list) - 1):
-        current_element = shortest_path_list[i]
-        next_element = shortest_path_list[i + 1]
-        result = rel.groupby('way_id').filter(lambda x: set([current_element, next_element]).issubset(x['node_id']))
+    for i in range(len(shortest_path) - 1):
+        current_element = shortest_path[i]
+        next_element = shortest_path[i + 1]
+        result = roads.groupby('way_id').filter(lambda x: set([current_element, next_element]).issubset(x['node_id']))
         pairs_list.append(result)
     final_result = pd.concat(pairs_list)
     return final_result
+
+
+
+
+def decrypt_shortest_path(client):
+    """Decrypt and store the shortest path in the session state
+
+    Args:
+        client (concrete.fhe.compilation.client.Client): A client instance
+    """
+    client.keys.load_if_exists_generate_and_save_otherwise(keys_filepath)
+    path = []
+    for enc_value in st.session_state['encrypted_shortest_path']:
+        deserialized_result = fhe.Value.deserialize(enc_value)
+        path.append(client.decrypt(deserialized_result))
+    st.session_state['decrypted_result'] = path
+
+def init_session():
+    """Initialize the Streamlit session and layout configuration.
+
+    Returns:
+        Streamlit.columns: A tuple of Streamlit columns for layout customization.
+    """
+    st.set_page_config(layout="wide")
+    
+    if 'markers' not in st.session_state:
+        st.session_state['markers'] = []
+    if 'server_side' not in st.session_state:
+            st.session_state['server_side'] = []    
+    if 'client_side' not in st.session_state:
+            st.session_state['client_side'] = []
+
+    c1, c2, c3 = st.columns([1, 3, 1])
+
+    return c1, c2, c3
+
+def add_marker(coordinates, name):
+    """Add a marker with coordinates and a name to the Streamlit session.
+
+    Args:
+        coordinates (Point): The coordinates of the marker 
+        name (str): The name or label for the marker
+    """
+    data = {'coordinates': coordinates, 'name': name}
+    st.session_state['markers'].append(data)
+
+
+def display_map(nodes, returned_objects=None, path=None):
+    """Display the map with nodes and optional markers and paths.
+
+    Args:
+        nodes (geopandas.DataFrame): A dataframe containing the nodes to display
+        returned_objects (List[str], optional): Objects to be returned when interacting with the map. Defaults to None.
+        path (pandas.DataFrame, optional): A DataFrame representing the road segments that form the shortest path. Defaults to None.
+
+    Returns:
+        Streamlit.FoliumMap: An interactive map displaying nodes, markers, and paths
+    """
+    m = nodes.explore(
+        color="red",  
+        marker_kwds=dict(radius=5, fill=True, name='node_id'),  
+        tooltip="node_id",  
+        tooltip_kwds=dict(labels=False), 
+        zoom_control=False,
+    )
+
+    if 'markers' in st.session_state:
+        for mrk in st.session_state['markers']:
+            folium.Marker([mrk['coordinates'].y, mrk['coordinates'].x], popup=mrk['name'], tooltip=mrk['name']).add_to(m)
+
+    if path is not None:
+        path.explore(
+                m=m,  
+                color="green",
+                style_kwds = {"weight":5},
+                tooltip="name",
+                popup=["name"],
+                name="Quadratic-Paris", 
+            )
+        
+    return st_folium(m, width=725, key="origin", returned_objects=returned_objects)
+
+
+def add_to_server_side(message):
+    """Add a message to the server side of the view
+
+    Args:
+        message (str): The message to be added to the server side
+    """
+    st.session_state['server_side'].append(message)
+
+def add_to_client_side(message):
+    """Add a message to the client side of the view
+
+    Args:
+        message (str): The message to be added to the client side
+    """
+    st.session_state['client_side'].append(message)
+
+def display_server_side():
+    """Display the messages stored in the server-side view.
+    """
+    st.write("**Server-side**")
+    for message in st.session_state['server_side']:
+        st.write(message)
+
+def display_client_side():
+    """Display the messages stored in the client-side view.
+    """
+    st.write("**Client-side**")
+    for message in st.session_state['client_side']:
+        st.write(message)
+
+def restart_session():
+    """Clear the session state to restart
+    """
+    if st.button('Restart'):
+        for key in st.session_state.items():
+            if key != 'evaluation_key':
+                del st.session_state[key]
+        st.rerun()