Update app.py

app.py

@@ -157,7 +157,7 @@ def create_plot(data, forecast_data, time_col, target_col):
     )
     return fig
 
-def full_forecast_pipeline(file_obj, time_col, target_col, forecast_horizon, finetune_steps, freq, start_date, end_date, start_time, end_time, resample_freq, merge_data) -> Tuple[str, object, str, str]:
+def full_forecast_pipeline(file_obj, time_col, target_col, finetune_steps, freq, start_date, end_date, start_time, end_time, resample_freq, merge_data, forecast_start_date, forecast_end_date) -> Tuple[str, object, str, str]:
     """
     Full pipeline: loads the data, calls the forecast function, and then processes the data.
     """
@@ -172,13 +172,6 @@ def full_forecast_pipeline(file_obj, time_col, target_col, forecast_horizon, fin
         # Sort the DataFrame by the time column
         data = data.sort_values(by=time_col)
 
-        # Get min and max dates from the data
-        min_date = data[time_col].min().strftime('%Y-%m-%d')
-        max_date = data[time_col].max().strftime('%Y-%m-%d')
-
-        # Fill missing values with 0
-        data = data.fillna(0)
-
         # Apply date range selection
         if start_date and end_date:
             start_datetime = pd.to_datetime(start_date)
@@ -186,12 +179,40 @@ def full_forecast_pipeline(file_obj, time_col, target_col, forecast_horizon, fin
             data = data[(data[time_col] >= start_datetime) & (data[time_col] <= end_datetime)]
             logger.info(f"Data filtered from {start_datetime} to {end_datetime}. Shape: {data.shape}")
 
-        data = data.set_index(time_col)
-
         # Resample the data
         data = data.resample(resample_freq).mean()
         data.reset_index(inplace=True)
 
+        # Convert forecast start and end dates to datetime
+        if forecast_start_date and forecast_end_date:
+            forecast_start_datetime = pd.to_datetime(forecast_start_date)
+            forecast_end_datetime = pd.to_datetime(forecast_end_date)
+
+            # Calculate the time difference
+            time_difference = forecast_end_datetime - forecast_start_datetime
+
+            # Calculate forecast horizon based on frequency
+            if freq == 'D':
+                forecast_horizon = time_difference.days
+            elif freq == 'W':
+                forecast_horizon = time_difference.days / 7
+            elif freq == 'M':
+                forecast_horizon = time_difference.days / 30.44  # Average days in a month
+            elif freq == 'Y':
+                forecast_horizon = time_difference.days / 365.25  # Average days in a year
+            elif 'min' in freq:
+                minutes = int(freq.replace('min', ''))
+                forecast_horizon = time_difference.total_seconds() / (minutes * 60)
+            elif 'H' in freq:
+                hours = int(freq.replace('H', ''))
+                forecast_horizon = time_difference.total_seconds() / (hours * 3600)
+            else:
+                raise ValueError("Unsupported frequency. Please select a valid frequency.")
+
+            forecast_horizon = int(forecast_horizon)
+        else:
+            raise ValueError("Forecast start and end dates must be provided.")
+
         forecast_result = forecast_nixtla(data, forecast_horizon, finetune_steps, freq, time_col, target_col)
         processed_data = process_forecast_data(forecast_result, time_col)
         processed_data = apply_zero_patterns(data.copy(), processed_data, time_col, target_col)
@@ -201,8 +222,13 @@ def full_forecast_pipeline(file_obj, time_col, target_col, forecast_horizon, fin
         else:
             merged_data = processed_data
 
-        plot = create_plot(data, processed_data, time_col, target_col)
-        csv_data = processed_data.to_csv(index=False)
+        # Filter forecast data based on forecast start and end dates
+        merged_data[time_col] = pd.to_datetime(merged_data[time_col])  # Ensure time_col is datetime
+        merged_data = merged_data[(merged_data[time_col] >= forecast_start_datetime) & (merged_data[time_col] <= forecast_end_datetime)]
+        logger.info(f"Forecast data filtered from {forecast_start_datetime} to {forecast_end_datetime}. Shape: {merged_data.shape}")
+
+        plot = create_plot(data, merged_data, time_col, target_col)
+        csv_data = merged_data.to_csv(index=False)
 
         # Create a temporary file and write the CSV data to it
         with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".csv") as tmpfile:
@@ -277,7 +303,7 @@ def create_interface():
         btn = gr.Button("Generate Forecast")
         btn.click(
             fn=full_forecast_pipeline,
-            inputs=[file_input, time_col_input, target_col_input, forecast_horizon_input, finetune_steps_input, freq_dropdown, start_date_input, end_date_input, start_time_input, end_time_input, resample_freq_dropdown, gr.Checkbox(label="Merge Data", value=False)],
+            inputs=[file_input, time_col_input, target_col_input, forecast_horizon_input, finetune_steps_input, freq_dropdown, start_date_input, end_date_input, start_time_input, end_time_input, resample_freq_dropdown, gr.Checkbox(label="Merge Data", value=False), gr.Textbox(label="Forecast Start Date", placeholder="YYYY-MM-DD", value="2023-01-01"), gr.Textbox(label="Forecast End Date", placeholder="YYYY-MM-DD", value="2023-12-31")],
             outputs=[output_csv, output_plot, download_button, error_output]
         )
     return iface