second round of edits from günther and andreas

app.py

@@ -6,7 +6,9 @@ from PIL import Image
 import numpy as np
 import gradio as gr
 import os
-from tirex import load_model, ForecastModel
+from plotly.subplots import make_subplots
+
+# from tirex import load_model, ForecastModel
 
 # ----------------------------
 # Helper functions (logic mostly unchanged)
@@ -15,21 +17,23 @@ from tirex import load_model, ForecastModel
 torch.manual_seed(42)
 
 def model_forecast(input_data, forecast_length=256, file_name=None):
-    if os.path.basename(file_name) == "merged_ett2_loop.csv":
-        _forecast_tensor = torch.load("data/merged_ett2_loop_forecast_256.pt")
+    if os.path.basename(file_name) == "loop.csv":
+        _forecast_tensor = torch.load("data/loop_forecast_512.pt")
+        return _forecast_tensor[:,:forecast_length,:]
+    elif os.path.basename(file_name) == "ett2.csv":
+        _forecast_tensor = torch.load("data/ett2_forecast_512.pt")
         return _forecast_tensor[:,:forecast_length,:]
     elif os.path.basename(file_name) == "air_passangers.csv":
-        _forecast_tensor = torch.load("data/air_passengers_forecast_256.pt")
+        _forecast_tensor = torch.load("data/air_passengers_forecast_512.pt")
         return _forecast_tensor[:,:forecast_length,:]
     else:
-        model: ForecastModel = load_model("NX-AI/TiRex",device='cuda')
-        forecast = model.forecast(context=input_data, prediction_length=forecast_length)
-        return forecast[0]
+        # model: ForecastModel = load_model("NX-AI/TiRex",device='cuda')
+        # forecast = model.forecast(context=input_data, prediction_length=forecast_length)
+        # return forecast[0]
+        pass
 
     
 
-
-
 def plot_forecast_plotly(timeseries, quantile_predictions, timeseries_name):
     """
     - timeseries: 1D list/array of historical values.
@@ -66,7 +70,7 @@ def plot_forecast_plotly(timeseries, quantile_predictions, timeseries_name):
         y=lower_q,
         mode="lines",
         line=dict(color="rgba(0, 0, 0, 0)", width=0),
-        name=f"{timeseries_name} – Lower Bound",
+        name=f"{timeseries_name} – 10% Quantile",
         hovertemplate="Lower: %{y:.2f}<extra></extra>"
     ))
 
@@ -78,7 +82,7 @@ def plot_forecast_plotly(timeseries, quantile_predictions, timeseries_name):
         line=dict(color="rgba(0, 0, 0, 0)", width=0),
         fill="tonexty",
         fillcolor="rgba(128, 128, 128, 0.3)",
-        name=f"{timeseries_name} – Upper Bound",
+        name=f"{timeseries_name} – 90% Quantile",
         hovertemplate="Upper: %{y:.2f}<extra></extra>"
     ))
 
@@ -139,24 +143,6 @@ def load_table(file_path):
         raise ValueError("Unsupported format. Use CSV, XLS, XLSX, or PARQUET.")
 
 
-# def extract_names_and_update(file, preset_filename):
-#     try:
-#         if file is not None:
-#             df = load_table(file.name)
-#         else:
-#             if not preset_filename:
-#                 return gr.update(choices=[], value=[]), []
-#             df = load_table(preset_filename)
-
-#         if df.shape[1] > 0 and df.iloc[:, 0].dtype == object and not df.iloc[:, 0].str.isnumeric().all():
-#             names = df.iloc[:, 0].tolist()
-#         else:
-#             names = [f"Series {i}" for i in range(len(df))]
-#         return gr.update(choices=names, value=names), names
-#     except Exception:
-#         return gr.update(choices=[], value=[]), []
-
-
 def extract_names_and_update(file, preset_filename):
     try:
         # Determine which file to use and get default forecast length
@@ -206,7 +192,7 @@ def get_default_forecast_length(file_path):
         return 64
     
     filename = os.path.basename(file_path)
-    if filename == "merged_ett2_loop.csv":
+    if filename == "loop.csv" or filename == "ett2.csv":
         return 256
     elif filename == "air_passangers.csv":
         return 48
@@ -216,17 +202,19 @@ def get_default_forecast_length(file_path):
 
 def display_filtered_forecast(file, preset_filename, selected_names, forecast_length):
     try:
-        # If no file uploaded and no valid preset chosen, return early
+        # 1) If no file or preset selected, show an error
         if file is None and (preset_filename is None or preset_filename == "-- No preset selected --"):
             return None, "No file selected."
 
-        # Load data
+        # 2) Load DataFrame and remember which filename to pass to model_forecast
         if file is not None:
             df = load_table(file.name)
+            file_name = file.name
         else:
             df = load_table(preset_filename)
+            file_name = preset_filename
 
-        # Determine names vs numeric data
+        # 3) Determine whether first column is names or numeric
         if (
             df.shape[1] > 0
             and df.iloc[:, 0].dtype == object
@@ -238,63 +226,128 @@ def display_filtered_forecast(file, preset_filename, selected_names, forecast_le
             all_names = [f"Series {i}" for i in range(len(df))]
             data_only = df.astype(float)
 
-        # Build a boolean mask for selected series
+        # 4) Build mask from selected_names
         mask = [name in selected_names for name in all_names]
         if not any(mask):
             return None, "No timeseries chosen to plot."
 
-        # Extract the filtered historical data (numpy array of shape (n_selected, seq_len))
-        filtered_data = data_only.iloc[mask, :].values
+        filtered_data = data_only.iloc[mask, :].values   # shape = (n_selected, seq_len)
         filtered_names = [all_names[i] for i, m in enumerate(mask) if m]
+        n_selected = filtered_data.shape[0]
+
+        # 5) First call model_forecast on all series, then select only the masked rows
+        full_data = data_only.values  # shape = (n_all, seq_len)
+        full_out = model_forecast(full_data, forecast_length=forecast_length, file_name=file_name)
+
+        # Now pick only the rows we actually filtered
+        out = full_out[mask, :, :]    # shape = (n_selected, pred_len, n_q)
+        inp = torch.tensor(filtered_data)
+
+        # 6) Create one subplot per selected series, with vertical spacing
+        fig = make_subplots(
+            rows=n_selected,
+            cols=1,
+            shared_xaxes=False,
+            vertical_spacing=0.3,   # more space between subplots
+            subplot_titles=filtered_names
+        )
 
-        # Get the full forecast tensor (n_series, pred_len, n_q)
-        file_path = file.name if file is not None else preset_filename
-        out = model_forecast(filtered_data, forecast_length=forecast_length, file_name=file_path)
-        inp = torch.tensor(filtered_data)  # shape = (n_selected, seq_len)
-
-        # If only one series is selected, we can just call plot_forecast_plotly directly:
-        if inp.shape[0] == 1:
-            ts = inp[0].numpy().tolist()
-            qp = out[0].numpy()
-            fig = plot_forecast_plotly(ts, qp, filtered_names[0])
-            return fig, ""
-
-        # If multiple series are selected, build a master figure by concatenating traces
-        master_fig = go.Figure()
-        for idx in range(inp.shape[0]):
+        for idx in range(n_selected):
             ts = inp[idx].numpy().tolist()
             qp = out[idx].numpy()
             series_name = filtered_names[idx]
 
-            # Get a “per‐series” figure
-            small_fig = plot_forecast_plotly(ts, qp, series_name)
-            # Append each trace from small_fig into master_fig
-            for trace in small_fig.data:
-                master_fig.add_trace(trace)
+            # a) plot historical data (blue line)
+            x_hist = list(range(len(ts)))
+            fig.add_trace(
+                go.Scatter(
+                    x=x_hist,
+                    y=ts,
+                    mode="lines",
+                    name=f"{series_name} – Given Data",
+                    line=dict(color="blue", width=2),
+                    showlegend=False
+                ),
+                row=idx + 1, col=1
+            )
+
+            # b) compute forecast indices
+            pred_len = qp.shape[0]
+            x_pred = list(range(len(ts) - 1, len(ts) - 1 + pred_len))
+
+            lower_q = qp[:, 0]
+            upper_q = qp[:, -1]
+            n_q = qp.shape[1]
+            median_idx = n_q // 2
+            median_q = qp[:, median_idx]
+
+            # c) lower‐bound (invisible)
+            fig.add_trace(
+                go.Scatter(
+                    x=x_pred,
+                    y=lower_q,
+                    mode="lines",
+                    line=dict(color="rgba(0,0,0,0)", width=0),
+                    name=f"{series_name} – 10% Quantile",
+                    hovertemplate="10% Quantile: %{y:.2f}<extra></extra>",
+                    showlegend=False
+                ),
+                row=idx + 1, col=1
+            )
+
+            # d) upper‐bound (shaded area)
+            fig.add_trace(
+                go.Scatter(
+                    x=x_pred,
+                    y=upper_q,
+                    mode="lines",
+                    line=dict(color="rgba(0,0,0,0)", width=0),
+                    fill="tonexty",
+                    fillcolor="rgba(128,128,128,0.3)",
+                    name=f"{series_name} – 90% Quantile",
+                    hovertemplate="90% Quantile: %{y:.2f}<extra></extra>",
+                    showlegend=False
+                ),
+                row=idx + 1, col=1
+            )
+
+            # e) median forecast (orange line)
+            fig.add_trace(
+                go.Scatter(
+                    x=x_pred,
+                    y=median_q,
+                    mode="lines",
+                    name=f"{series_name} – Median Forecast",
+                    line=dict(color="orange", width=2),
+                    hovertemplate="Median: %{y:.2f}<extra></extra>",
+                    showlegend=False
+                ),
+                row=idx + 1, col=1
+            )
+
+            # f) label axes for each subplot
+            fig.update_xaxes(title_text="Time", row=idx + 1, col=1)
+            fig.update_yaxes(title_text="Value", row=idx + 1, col=1)
 
-        # Finally, update the layout of master_fig (you can tweak the title if you want)
-        master_fig.update_layout(
+        # 7) Global layout tweaks
+        fig.update_layout(
             template="plotly_dark",
+            height=300 * n_selected,  # 300px per subplot
             title=dict(
                 text="Forecasts for Selected Timeseries",
                 x=0.5,
-                font=dict(size=16, family="Arial", color="white")
+                font=dict(size=20, family="Arial", color="white")
             ),
-            xaxis=dict(
-                rangeslider=dict(visible=True),  # <-- put rangeslider here
-                fixedrange=False
-            ),
-            xaxis_title="Time",
-            yaxis_title="Value",
             hovermode="x unified",
-            # height=plot_height,
-            # width=plot_width   # ← add these here
-            autosize=True,
+            margin=dict(t=120, b=40, l=60, r=40),
+            showlegend=False
         )
-        return master_fig, ""
+
+        return fig, ""
 
     except Exception as e:
-        return None, f"Error: {e}. Use CSV, XLS, XLSX, or PARQUET."
+        return None, f"Error: {str(e)}"
+
 
 
 # ----------------------------
@@ -313,7 +366,7 @@ with gr.Blocks(fill_width=True,theme=gr.themes.Ocean()) as demo:
                 label="Upload CSV / XLSX / PARQUET",
                 file_types=[".csv", ".xls", ".xlsx", ".parquet"]
             )
-            preset_choices = ["-- No preset selected --", "data/merged_ett2_loop.csv", "data/air_passangers.csv"]
+            preset_choices = ["-- No preset selected --", "data/loop.csv", "data/air_passangers.csv", 'data/ett2.csv']
 
             preset_dropdown = gr.Dropdown(
                 label="Or choose a preset:",
@@ -347,7 +400,6 @@ with gr.Blocks(fill_width=True,theme=gr.themes.Ocean()) as demo:
                 gr.Image("static/nxai_logo.png", width=150, show_label=False, container=False)
                 gr.Image("static/tirex.jpeg", width=150, show_label=False, container=False)
 
-        # Right column: interactive plot + instructions
         with gr.Column(scale=5):
             gr.Markdown("## Forecast Plot")
             plot_output = gr.Plot()
@@ -418,10 +470,10 @@ with gr.Blocks(fill_width=True,theme=gr.themes.Ocean()) as demo:
 
     # Plot button
     plot_button.click(
-        fn=display_filtered_forecast,
-        inputs=[file_input, preset_dropdown, filter_checkbox, forecast_length_slider],  # <-- add slider here
-        outputs=[plot_output, errbox]
-    )
+    fn=display_filtered_forecast,
+    inputs=[file_input, preset_dropdown, filter_checkbox, forecast_length_slider],
+    outputs=[plot_output, errbox]
+)
     demo.launch()