🎨 add typehints and always return figure for plotting fcts

double check type hints
- also FloatArray is not very informative about the layout of the array (2D, 3D, 4D?)

src/move/analysis/metrics.py

@@ -83,7 +83,9 @@ def norm(x: np.ma.MaskedArray, axis: int = 1) -> np.ma.MaskedArray:
     return np.sqrt(np.sum(x**2, axis=axis))
 
 
-def get_2nd_order_polynomial(x_array, y_array, n_points=100):
+def get_2nd_order_polynomial(
+    x_array: FloatArray, y_array: FloatArray, n_points=100
+) -> tuple[FloatArray, FloatArray, tuple[float, float, float]]:
     """
     Given a set of x an y values, find the 2nd oder polynomial fitting best the data.
 

src/move/tasks/identify_associations.py

@@ -885,5 +885,5 @@ def identify_associations(config: MOVEConfig) -> None:
         association_df = pd.read_csv(
             output_path / f"results_sig_assoc_{task_type}.tsv", sep="\t"
         )
-        plot_feature_association_graph(association_df, output_path)
-        plot_feature_association_graph(association_df, output_path, layout="spring")
+        _ = plot_feature_association_graph(association_df, output_path)
+        _ = plot_feature_association_graph(association_df, output_path, layout="spring")

src/move/visualization/dataset_distributions.py

@@ -85,7 +85,7 @@ def plot_reconstruction_diff(
 
 def plot_feature_association_graph(
     association_df, output_path, layout="circular", style: str = DEFAULT_PLOT_STYLE
-):
+) -> matplotlib.figure.Figure:
     """
     This function plots a graph where each node corresponds to a feature and the edges
     represent the associations between features. Edge width represents the probability of
@@ -183,6 +183,7 @@ def plot_feature_association_graph(
         fig.savefig(
             output_path / f"Feature_association_graph_{layout}.png", format="png"
         )
+    return fig
 
 
 def plot_feature_mean_median(
@@ -282,8 +283,15 @@ def plot_cumulative_distributions(
 
 
 def plot_correlations(
-    x, y, x_pol, y_pol, a2, a1, a, k, style: str = DEFAULT_PLOT_STYLE
-):
+        x: FloatArray,
+        y: FloatArray,
+        x_pol: FloatArray,
+        y_pol: FloatArray,
+        a2: float,
+        a1: float,
+        a: float,
+        k: int,
+        style: str = DEFAULT_PLOT_STYLE) -> matplotlib.figure.Figure:
     """
     Plot y vs x and the corresponding polynomial fit.
     """

src/move/visualization/latent_space.py

@@ -108,17 +108,17 @@ def plot_latent_space_with_con(
 
 
 def plot_3D_latent_and_displacement(
-    mu_baseline,
-    mu_perturbed,
-    feature_values,
-    feature_name,
-    show_baseline=True,
-    show_perturbed=True,
-    show_arrows=True,
+    mu_baseline: FloatArray,
+    mu_perturbed: FloatArray,
+    feature_values: FloatArray,
+    feature_name: str,
+    show_baseline: bool = True,
+    show_perturbed: bool = True,
+    show_arrows: bool = True,
     step: int = 1,
     altitude: int = 30,
     azimuth: int = 45,
-):
+) -> matplotlib.figure.Figure:
     """
     Plot the movement of the samples in the 3D latent space after perturbing one
     input variable.

src/move/visualization/vae_visualization.py

@@ -25,7 +25,7 @@ def plot_vae(
     output_sample: Optional[torch.Tensor] = None,
     mu: Optional[torch.Tensor] = None,
     logvar: Optional[torch.Tensor] = None,
-):
+) -> matplotlib.figure.Figure:
     """
     This function is aimed to visualize MOVE's architecture.
 
@@ -265,3 +265,4 @@ def plot_vae(
     plt.colorbar(sm_edge, label="Edge value", shrink=0.2)
     plt.tight_layout()
     fig.savefig(savepath / f"{title}.png", format="png", dpi=200)
+    return fig