Update scripts/app.py

scripts/app.py

@@ -139,93 +139,51 @@ def calculate_metrics_pro(portfolio_values, freq=252, rf=0.0):
 # XAI: Feature Importance Function
 # =========================================
 def calculate_feature_importance(model, obs):
-    """
-    Calculates feature importance using Integrated Gradients on the RL agent's policy network.
-    """
-    # Convert observation to torch tensor and enable gradient tracking
     obs_tensor = torch.as_tensor(obs, dtype=torch.float32, device=model.device)
+    if obs_tensor.dim() == 1: obs_tensor = obs_tensor.unsqueeze(0)
     obs_tensor.requires_grad_()
-
-    # Get the policy network (actor)
+    
     actor = model.policy.actor
-
-    # Define a baseline (e.g., a zero observation)
     baseline = torch.zeros_like(obs_tensor)
-
-    # Number of steps for integral approximation
     steps = 50
-
-    # Generate scaled inputs along the path from baseline to input
     scaled_inputs = [baseline + (float(i) / steps) * (obs_tensor - baseline) for i in range(steps + 1)]
-
+    
     grads = []
     for scaled_input in scaled_inputs:
-        # Forward pass to get action distribution parameters (mean)
         action_mean = actor(scaled_input)
-
-        # We need a scalar output to calculate gradients against.
-        # Here we sum, representing overall sensitivity of the action vector.
         target_output = action_mean.sum()
-
-        # Calculate gradients of the target output with respect to the input features
         grad = torch.autograd.grad(outputs=target_output, inputs=scaled_input)[0]
         grads.append(grad)
 
-    # Average the gradients using the trapezoidal rule approximation
-    avg_grads = (grads[:-1] + grads[1:]) / 2.0
-    avg_grads = torch.stack(avg_grads).mean(dim=0)
+    # --- Stack gradients first, then perform arithmetic ---
+    stacked_grads = torch.stack(grads)
+    avg_grads = (stacked_grads[:-1] + stacked_grads[1:]) / 2.0
+    avg_grads = avg_grads.mean(dim=0)
+    # -----------------------------------------------------------
 
-    # Calculate Integrated Gradients: (input - baseline) * average_gradients
     integrated_grads = (obs_tensor - baseline) * avg_grads
-
-    # Detach, move to cpu, and convert to numpy array
     importance_scores = integrated_grads.detach().cpu().numpy().flatten()
-
-    # Feature Names mapping
-    num_assets = len(ASSETS)
-    num_macro = len(MACRO_COLS)
-
-    # Create feature names based on the observation structure
+    
     feature_names = []
     for i in range(WINDOW_SIZE):
-        for asset in ASSETS:
-            feature_names.append(f"{asset}_t-{WINDOW_SIZE-1-i}")
+        for asset in ASSETS: feature_names.append(f"{asset}_t-{WINDOW_SIZE-1-i}")
     for i in range(WINDOW_SIZE):
-        for macro in MACRO_COLS:
-            feature_names.append(f"{macro}_t-{WINDOW_SIZE-1-i}")
+        for macro in MACRO_COLS: feature_names.append(f"{macro}_t-{WINDOW_SIZE-1-i}")
 
-    # Combine into a dictionary and sort by absolute importance
     feature_importance_dict = dict(zip(feature_names, importance_scores))
-
-    # Aggregate importance by feature type (sum of absolute values across time steps)
     aggregated_importance = {}
     for base_feature in ASSETS + MACRO_COLS:
         total_imp = sum(abs(val) for key, val in feature_importance_dict.items() if key.startswith(base_feature))
         aggregated_importance[base_feature] = total_imp
 
-    # Sort and take top N for display
     top_features = dict(sorted(aggregated_importance.items(), key=lambda item: item[1], reverse=True)[:8])
 
-    # Create a Plotly bar chart
-    fig = px.bar(
-        x=list(top_features.values()),
-        y=list(top_features.keys()),
-        orientation='h',
-        title="Top Influential Features (XAI)",
-        labels={'x': 'Relative Importance Score', 'y': 'Feature'},
-        color=list(top_features.values()),
-        color_continuous_scale=px.colors.sequential.Viridis
-    )
-    fig.update_layout(
-        template="plotly_dark",
-        paper_bgcolor='rgba(0,0,0,0)',
-        plot_bgcolor='rgba(0,0,0,0)',
-        yaxis={'categoryorder':'total ascending'},
-        coloraxis_showscale=False,
-        margin=dict(l=10, r=10, t=40, b=10),
-        height=300 # Keep it compact
-    )
-
+    fig = px.bar(x=list(top_features.values()), y=list(top_features.keys()), orientation='h',
+        title="Top Influential Features (XAI)", labels={'x': 'Importance', 'y': 'Feature'},
+        color=list(top_features.values()), color_continuous_scale=px.colors.sequential.Viridis)
+    fig.update_layout(template="plotly_dark", paper_bgcolor='rgba(0,0,0,0)', plot_bgcolor='rgba(0,0,0,0)',
+        yaxis={'categoryorder':'total ascending'}, coloraxis_showscale=False, margin=dict(l=10, r=10, t=40, b=10), height=300,
+        hoverlabel=dict(bgcolor="white", font_size=14, font_family="Roboto", font_color="black"))
     return fig
 
 # =========================================
@@ -411,83 +369,54 @@ def prepare_observation(data_window):
     return obs.flatten().astype(np.float32), obs.astype(np.float32), data_window
 
 def predict_and_analyze():
-    """Main function for Forecast Tab."""
-    status_msg = "Starting process..."
-    loading_html = """<div style="color: #9ca3af;">🔄 Fetching data & running prediction...</div>"""
-    # Update to yield an empty plot for the XAI chart initially
-    yield status_msg, None, go.Figure(), loading_html
-
+    yield "Starting...", None, go.Figure(), "Loading..."
     try:
-        data_window = get_latest_data_window(WINDOW_SIZE)
-        # Get flattened obs for prediction and raw obs for XAI
+        data_window = get_prediction_data(WINDOW_SIZE)
         flat_obs, raw_obs, df_window_for_analyst = prepare_observation(data_window)
-
-        if not os.path.exists(MODEL_PATH): raise FileNotFoundError(f"Model not found: {MODEL_PATH}")
+        
+        if not os.path.exists(MODEL_PATH): raise FileNotFoundError("Model not found.")
         model = SAC.load(MODEL_PATH)
+        
+        # --- Pass the FLATTENED observation to XAI function ---
+        # The XAI function logic expects an input that matches the model's input layer.
+        yield "XAI Calc...", None, go.Figure(), "Calculating XAI..."
+        xai_plot = calculate_feature_importance(model, flat_obs)
 
-        # --- XAI: Calculate Feature Importance ---
-        status_msg = "Calculating feature importance..."
-        yield status_msg, None, go.Figure(), loading_html
-        xai_plot = calculate_feature_importance(model, raw_obs)
-
-        # --- Prediction ---
         action, _ = model.predict(flat_obs, deterministic=True)
-        exp_action = np.exp(np.asarray(action).flatten())
-        weights = exp_action / np.sum(exp_action)
-        allocations_dict = {asset: weights[i] for i, asset in enumerate(ASSETS)}
-        allocations_dict['Cash'] = weights[-1]
-        alloc_df = pd.DataFrame(list(allocations_dict.items()), columns=['Asset', 'Proposed Allocation'])
-        alloc_df['Proposed Allocation'] = alloc_df['Proposed Allocation'].apply(lambda x: f"{x:.2%}")
-
-        status_msg = "Prediction done. Running AI Risk Analysis..."
-        analysing_html = """<div style="color: #9ca3af;">🤖 Running Qwen-2.5-3B Risk Analysis...</div>"""
-        # Yield XAI plot along with other outputs
-        yield status_msg, alloc_df, xai_plot, analysing_html
-
-        allocations_for_llm = {k: float(v) for k, v in allocations_dict.items()}
-        analysis_result = analyze_agent_decision(df_window_for_analyst, allocations_for_llm)
-        status_msg = "Analysis complete!"
-
-        if isinstance(analysis_result, dict):
-            strat = analysis_result.get('strategy_summary', 'N/A')
-            risk = analysis_result.get('risk_level', 'N/A').upper()
-            just = analysis_result.get('justification', 'N/A')
-            conf = analysis_result.get('confidence_score', 'N/A')
-            if 'HIGH' in risk:
-                risk_css = "color: #ef4444; font-weight: bold;"
-                status_bg = "#7f1d1d"
-                status_border = "#ef4444"
-                status_icon = "⛔"
-                status_text = "TRADE BLOCKED: High Risk Detected"
-            else:
-                risk_css = "color: #10b981; font-weight: bold;"
-                status_bg = "#064e3b"
-                status_border = "#10b981"
-                status_icon = "🚀"
-                status_text = "TRADE APPROVED"
-
-            report_html = f"""
-            <div style="background-color: #1f2937; padding: 20px; border-radius: 12px 12px 0 0; border: 1px solid #374151; border-bottom: none;">
-                <h3 style="margin-top: 0; color: #e5e7eb;">🤖 AI Risk Analyst Report</h3>
-                <div style="margin-bottom: 15px;"><strong style="color: #9ca3af;">Strategy:</strong><br><span style="color: #d1d5db;">{strat}</span></div>
-                <div style="margin-bottom: 15px;"><strong style="color: #9ca3af;">Risk Level:</strong><span style="margin-left: 8px; {risk_css}">{risk}</span></div>
-                <div style="margin-bottom: 15px;"><strong style="color: #9ca3af;">Justification:</strong><br><span style="color: #d1d5db;">{just}</span></div>
-                <div><strong style="color: #9ca3af;">Confidence:</strong> <span style="color: #d1d5db;">{conf}/10</span></div>
-            </div>
-            <div style="background-color: {status_bg}; color: white; padding: 15px; border-radius: 0 0 12px 12px; border: 2px solid {status_border}; text-align: center; font-size: 1.2em; font-weight: bold; display: flex; align-items: center; justify-content: center;">
-                <span style="margin-right: 10px; font-size: 1.4em;">{status_icon}</span>{status_text}
-            </div>"""
+        exp_act = np.exp(np.asarray(action).flatten())
+        weights = exp_act / np.sum(exp_act)
+        
+        allocs = {ASSETS[i]: weights[i] for i in range(len(ASSETS))}
+        allocs['Cash'] = weights[-1]
+        alloc_df = pd.DataFrame(list(allocs.items()), columns=['Asset', 'Alloc'])
+        alloc_df['Alloc'] = alloc_df['Alloc'].apply(lambda x: f"{x:.2%}")
+
+        yield "AI Analysis...", alloc_df, xai_plot, "Running AI..."
+        llm_allocs = {k: float(v) for k, v in allocs.items()}
+        res = analyze_agent_decision(df_window_for_analyst, llm_allocs)
+        
+        if isinstance(res, dict):
+            strat, risk, just, conf = res.get('strategy_summary','N/A'), res.get('risk_level','N/A').upper(), res.get('justification','N/A'), res.get('confidence_score','N/A')
+            border_col = "#ef4444" if 'HIGH' in risk else "#10b981"
+            bg_col = "#7f1d1d" if 'HIGH' in risk else "#064e3b"
+            icon = "⛔" if 'HIGH' in risk else "🚀"
+            status = "TRADE BLOCKED" if 'HIGH' in risk else "TRADE APPROVED"
+            
+            html = f"""<div style="background-color: #1f2937; padding: 20px; border-radius: 12px; border: 1px solid #374151;">
+                <h3 style="margin-top: 0; color: #e5e7eb;">🤖 AI Report</h3>
+                <p><strong>Strategy:</strong> <span style="color:#d1d5db">{strat}</span></p>
+                <p><strong>Risk:</strong> <span style="color:{border_col}; font-weight:bold">{risk}</span></p>
+                <p><strong>Reason:</strong> <span style="color:#d1d5db">{just}</span></p>
+                <p><strong>Conf:</strong> <span style="color:#d1d5db">{conf}/10</span></p></div>
+                <div style="background-color:{bg_col}; color:white; padding:15px; margin-top:10px; border-radius:12px; text-align:center; font-weight:bold;">{icon} {status}</div>"""
         else:
-            report_html = f"""<div style="padding: 20px; background-color: #7f1d1d; color: #fca5a5; border-radius: 12px;"><h3>❌ Analysis Failed to Parse</h3><p>{str(analysis_result)}</p></div>"""
-        # Final yield with all outputs including XAI plot
-        yield status_msg, alloc_df, xai_plot, report_html
+            html = f"<div style='color:red'>{str(res)}</div>"
+        
+        yield "Done", alloc_df, xai_plot, html
     except Exception as e:
         import traceback
         traceback.print_exc()
-        status_msg = f"Error: {str(e)}"
-        error_html = f"""<div style="padding: 20px; background-color: #7f1d1d; color: #fca5a5; border-radius: 12px;"><h3>❌ Process Error</h3><p>{str(e)}</p></div>"""
-        # Final yield in case of error
-        yield status_msg, None, go.Figure(), error_html
+        yield f"Error: {str(e)}", None, go.Figure(), f"Error: {str(e)}"
 
 
 # =========================================