oh shit im scared, but its alive

main_hypot/best_model_and_plots.py

@@ -223,6 +223,7 @@ def plot_density_scatter(
     rolling_window: int = DEFAULT_ROLLING_WINDOW,
     savgol_window: int = DEFAULT_SAVGOL_WINDOW,
     savgol_poly: int = DEFAULT_SAVGOL_POLY,
+    return_fig: bool = False,
 ) -> None:
     fig, ax = plt.subplots(figsize=(8, 8))
     alpha_values = compute_density_alpha(
@@ -246,6 +247,7 @@ def plot_density_scatter(
         linewidths=0,
     )
 
+    trend_data = None
     if with_trend:
         tx, ty = compute_trend(
             df,
@@ -260,6 +262,7 @@ def plot_density_scatter(
         if len(tx):
             ax.plot(tx, ty, color=trend_color, linewidth=trend_linewidth, label=f"{trend_method} тренд")
             ax.legend()
+            trend_data = (tx, ty)
 
     ax.set_xlim(0, x_max)
     ax.set_ylim(y_min, y_max)
@@ -272,6 +275,8 @@ def plot_density_scatter(
     out_path.parent.mkdir(parents=True, exist_ok=True)
     fig.tight_layout()
     fig.savefig(out_path, dpi=150)
+    if return_fig:
+        return fig, ax, trend_data
     plt.close(fig)
     print(f"Saved {out_path}")
 
@@ -426,7 +431,7 @@ def plot_clean_trend_scatter(
         q_low=q_low,
         q_high=q_high,
     )
-    plot_density_scatter(
+    fig_ax = plot_density_scatter(
         cleaned,
         y_col=y_col,
         title=f"Облако без выбросов + тренд {y_col} vs {x_col}",
@@ -450,9 +455,11 @@ def plot_clean_trend_scatter(
         rolling_window=rolling_window,
         savgol_window=savgol_window,
         savgol_poly=savgol_poly,
+        return_fig=return_components,
     )
     if return_components:
-        return fig, ax, cleaned
+        fig, ax, trend_data = fig_ax
+        return fig, ax, cleaned, trend_data
 
 
 def generate_scatter_set(