added colorcoded points

app.py

@@ -156,20 +156,8 @@ class ArxivSearch:
             raise ValueError(f"Unsupported embedding type: {self.embedding}")
         
         results_scores = [i[1] for i in self.last_results]
-        # Normalize scores to [0, 1] for color mapping
-        if results_scores:
-            min_score = min(results_scores)
-            max_score = max(results_scores)
-            if max_score > min_score:
-            color_scale = [(score - min_score) / (max_score - min_score) for score in results_scores]
-            else:
-            color_scale = [0.5 for _ in results_scores]
-        else:
-            color_scale = []
 
-        # Use a "hot" colorscale (yellow/red for high, black for low)
-        cmap = matplotlib.cm.get_cmap('hot')
-        results_colors = [matplotlib.colors.rgb2hex(cmap(val)) for val in color_scale]
+        traces = []
 
         trace = go.Scatter3d(
             x=reduced_data[:, 0],
@@ -183,39 +171,60 @@ class ArxivSearch:
             text=[f"<br>: {self.arxiv_ids[i] if self.arxiv_ids[i] else self.documents[i].split()[:10]}" for i in range(len(self.documents))],
             hoverinfo='text'
         )
+
+        traces.append(trace)
+
         layout = go.Layout(
             margin=dict(l=0, r=0, b=0, t=0),
             scene=dict(
-                xaxis_title='PCA 1',
-                yaxis_title='PCA 2',
-                zaxis_title='PCA 3',
-                xaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
-                yaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
-                zaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
+            xaxis_title='PCA 1',
+            yaxis_title='PCA 2',
+            zaxis_title='PCA 3',
+            xaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
+            yaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
+            zaxis=dict(backgroundcolor='black', color='white', gridcolor='gray', zerolinecolor='gray'),
             ),
             paper_bgcolor='black',   # Outside the plotting area
             plot_bgcolor='black',    # Plotting area
             font=dict(color='white'), # Axis and legend text
             legend=dict(
-                bgcolor='rgba(0,0,0,0)',   # Transparent legend background
-                bordercolor='rgba(0,0,0,0)', # No border
-                x=0.01,  # Place legend inside plot area (adjust as needed)
-                y=0.99,
-                xanchor='left',
-                yanchor='top'
+            bgcolor='rgba(0,0,0,0)',   # Transparent legend background
+            bordercolor='rgba(0,0,0,0)', # No border
+            x=0.01,  # Place legend inside plot area (adjust as needed)
+            y=0.99,
+            xanchor='left',
+            yanchor='top'
             )
         )
+
         if len(reduced_results_points) > 0:
+            custom_colorscale = [
+                [0.0, "#00ffea"],  # Start color (e.g., bright cyan)
+                [1.0, "#ffea00"],  # End color (e.g., bright yellow)
+            ]
+
             results_trace = go.Scatter3d(
                 x=reduced_results_points[:, 0],
                 y=reduced_results_points[:, 1],
                 z=reduced_results_points[:, 2],
                 mode='markers',
-                marker=dict(size=3.5, color='orange', opacity=0.75),
+                marker=dict(size=4.25, 
+                            color=results_scores, 
+                            colorscale=custom_colorscale, 
+                            opacity=0.99, 
+                            colorbar=dict(
+                                    title="Score",
+                                    bgcolor='rgba(0,0,0,0)',  # <-- Transparent background for colorbar
+                                    bordercolor='rgba(0,0,0,0)' # No border
+                            )
+                        ),
                 name='Results',
                 text=[f"<br>{self.documents[i][:100]}" for i in results_indices],
                 hoverinfo='text'
             )
+
+            traces.append(results_trace)
+
             if not self.embedding == "tfidf" and self.query_encoding is not None and self.query_encoding.shape[0] > 0:
                 query_trace = go.Scatter3d(
                     x=query_point[:, 0],
@@ -227,11 +236,10 @@ class ArxivSearch:
                     text=[f"<br>Query: {self.query}"],
                     hoverinfo='text'
                 )
-                fig = go.Figure(data=[trace, results_trace, query_trace], layout=layout)
-            else:
-                fig = go.Figure(data=[trace, results_trace], layout=layout)
-        else:
-            fig = go.Figure(data=[trace], layout=layout)
+                traces.append(query_trace)
+
+        fig = go.Figure(data=traces, layout=layout)
+
         return fig
     
     def keyword_match_ranking(self, query, top_n=10):