added option to min max scatter

mutagenesis_visualization/main/bar_graphs/differential.py

@@ -1,7 +1,7 @@
 """
 This module contains the class that plots the differential bar plot.
 """
-from typing import Union, Dict, Any, Literal, TYPE_CHECKING
+from typing import Union, Dict, Any, Literal, TYPE_CHECKING, Optional
 from pathlib import Path
 from pandas import DataFrame
 import matplotlib.pyplot as plt
@@ -28,6 +28,8 @@ def __call__(
         metric: METRIC = 'rmse',
         plot_type: str = 'bar',
         show_cartoon: bool = False,
+        min_score: Optional[float] = None,
+        max_score: Optional[float] = None,
         replicate: int = -1,
         output_file: Union[None, str, Path] = None,
         **kwargs: Any,
@@ -52,6 +54,16 @@ def __call__(
             structure. The user must have added the secondary structure
             to the object.
 
+        min_score : float, default None
+            Change values below a minimum score to be that score.
+            i.e., setting min_score = -1 will change any value smaller
+            than -1 to -1.
+
+        max_score : float, default None
+            Change values below a maximum score to be that score.
+            i.e., setting max_score = 1 will change any value greater
+            than 1 to 1.
+
         replicate : int, default -1
             Set the replicate to plot. By default, the mean is plotted.
             First replicate start with index 0.
@@ -69,8 +81,8 @@ def __call__(
         self.graph_parameters()
 
         self.df_output: DataFrame = process_rmse_residue(
-            self.dataframes.df_notstopcodons[replicate],
-            screen_object.dataframes.df_notstopcodons[replicate], metric
+            self.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate],
+            screen_object.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate], metric
         )
         # make cartoon
         if show_cartoon:

mutagenesis_visualization/main/bar_graphs/enrichment_bar.py

@@ -1,7 +1,7 @@
 """
 This module contains the class that plots the mean position bar plot.
 """
-from typing import Union, Dict, Any
+from typing import Union, Dict, Any, Optional
 from pathlib import Path
 import matplotlib.pyplot as plt
 import numpy as np
@@ -21,6 +21,8 @@ def __call__(
         self,
         mode: str = "mean",
         show_cartoon: bool = False,
+        min_score: Optional[float] = None,
+        max_score: Optional[float] = None,
         replicate: int = -1,
         output_file: Union[None, str, Path] = None,
         **kwargs: Any,
@@ -37,6 +39,16 @@ def __call__(
             show the alanine substitution profile. Can be used for each
             amino acid. Use the one-letter code and upper case.
 
+        min_score : float, default None
+            Change values below a minimum score to be that score.
+            i.e., setting min_score = -1 will change any value smaller
+            than -1 to -1.
+
+        max_score : float, default None
+            Change values below a maximum score to be that score.
+            i.e., setting max_score = 1 will change any value greater
+            than 1 to 1.
+
         show_cartoon : boolean, default False
             If true, the plot will display a cartoon with the secondary
             structure. The user must have added the secondary structure
@@ -64,12 +76,12 @@ def __call__(
 
         # Select grouping
         if mode.lower() == 'mean':
-            self.df_output = self.dataframes.df_notstopcodons[replicate].groupby(
+            self.df_output = self.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate].groupby(
                 'Position', as_index=False
             ).mean()
         else:
-            self.df_output = self.dataframes.df_notstopcodons[replicate].loc[
-                self.dataframes.df_notstopcodons[replicate]['Aminoacid'] == mode].copy()
+            self.df_output = self.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate].loc[
+                self.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate]['Aminoacid'] == mode].copy()
 
         self.df_output['Color'] = self.df_output.apply(
             color_data, axis=1, args=(

mutagenesis_visualization/main/bar_graphs/secondary.py

@@ -1,7 +1,7 @@
 """
 This module contains the box plot class.
 """
-from typing import Union, Dict, Any
+from typing import Union, Dict, Any, Optional
 from pathlib import Path
 import matplotlib.pyplot as plt
 from pandas.core.frame import DataFrame
@@ -27,6 +27,8 @@ class Secondary(Pyplot):
     """
     def __call__(
         self,
+        min_score: Optional[float] = None,
+        max_score: Optional[float] = None,
         replicate: int = -1,
         output_file: Union[None, str, Path] = None,
         **kwargs: Any,
@@ -37,6 +39,16 @@ def __call__(
 
         Parameters
         -----------
+        min_score : float, default None
+            Change values below a minimum score to be that score.
+            i.e., setting min_score = -1 will change any value smaller
+            than -1 to -1.
+
+        max_score : float, default None
+            Change values below a maximum score to be that score.
+            i.e., setting max_score = 1 will change any value greater
+            than 1 to 1.
+
         replicate : int, default -1
             Set the replicate to plot. By default, the mean is plotted.
             First replicate start with index 0.
@@ -54,7 +66,7 @@ def __call__(
 
         # Get data
         df_output: DataFrame = _calculate_secondary(
-            self.dataframes.df_notstopcodons[replicate], self.secondary_dup
+            self.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate], self.secondary_dup
         )
 
         # Color

mutagenesis_visualization/main/classes/screen.py

@@ -232,7 +232,6 @@ def __init__(
             aminoacids=self.aminoacids
         )
 
-        # scatter
         self.scatter: Scatter = Scatter(dataframes=self.dataframes, aminoacids=self.aminoacids)
 
         self.correlation: Correlation = Correlation(

mutagenesis_visualization/main/process_data/calculate_enrichment.py

@@ -145,7 +145,6 @@ def calculate_enrichment(
     log10_counts, output_lib_np = get_log_enrichment(
         pre_lib_df, post_lib_df, input_stopcodon, output_stopcodon, min_counts, stopcodon, infinite
     )
-    print(output_lib_np)
     # Group by amino acid
     log10_counts_grouped: DataFrame = group_by_aa(DataFrame(log10_counts), aminoacids)
     log10_counts_mad = np.ravel(np.array(log10_counts_grouped))

mutagenesis_visualization/main/scatter/scatter.py

@@ -1,7 +1,7 @@
 """
 This module contains the class that plots scatters.
 """
-from typing import Union, Dict, Any, TYPE_CHECKING, Literal
+from typing import Union, Dict, Any, TYPE_CHECKING, Literal, Optional
 from pathlib import Path
 import numpy as np
 from pandas.core.frame import DataFrame
@@ -25,6 +25,8 @@ def __call__(
         self,
         screen_object: Union['Screen', Any],
         mode: Literal["mean", "pointmutant"] = 'pointmutant',
+        min_score: Optional[float] = None,
+        max_score: Optional[float] = None,
         replicate: int = -1,
         replicate_second_object: int = -1,
         output_file: Union[None, str, Path] = None,
@@ -41,6 +43,16 @@ def __call__(
         mode : str, default 'pointmutant'.
             Alternative set to "mean" for the mean of each position.
 
+        min_score : float, default None
+            Change values below a minimum score to be that score.
+            i.e., setting min_score = -1 will change any value smaller
+            than -1 to -1.
+
+        max_score : float, default None
+            Change values below a maximum score to be that score.
+            i.e., setting max_score = 1 will change any value greater
+            than 1 to 1.
+
         replicate : int, default -1
             Set the replicate to plot. By default, the mean is plotted.
             First replicate start with index 0.
@@ -65,13 +77,13 @@ def __call__(
         # Chose mode:
         if mode.lower() == 'pointmutant':
             df_output: DataFrame = process_by_pointmutant(
-                self.dataframes.df_notstopcodons[replicate],
-                screen_object.dataframes.df_notstopcodons[replicate_second_object]
+                self.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate],
+                screen_object.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate_second_object]
             )
         else:
             df_output = process_mean_residue(
-                self.dataframes.df_notstopcodons[replicate],
-                screen_object.dataframes.df_notstopcodons[replicate_second_object]
+                self.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate],
+                screen_object.dataframes.df_notstopcodons_limit_score(min_score, max_score)[replicate_second_object]
             )
 
         # create figure
@@ -149,4 +161,4 @@ def _tune_plot(self, temp_kwargs: Dict[str, Any]) -> None:
         plt.draw()
 
         # Legend
-        plt.legend(loc='upper left', handlelength=0, handletextpad=0, frameon=False, fontsize=10)
+        plt.legend(loc='upper left', handlelength=0, handletextpad=0, frameon=False, fontsize=temp_kwargs['legend_fontsize'])

mutagenesis_visualization/main/scatter/scatter_replicates.py

@@ -2,7 +2,7 @@
 This module will have the wrapper function to plot scatters for all
 replicates within a Screen object.
 """
-from typing import Any, Union, Literal
+from typing import Any, Union, Literal, Optional
 from pathlib import Path
 from itertools import combinations
 from copy import deepcopy
@@ -18,6 +18,8 @@ def __call__(
         self,
         wt_counts_only: bool = False,
         mode: Literal["mean", "pointmutant"] = 'pointmutant',
+        min_score: Optional[float] = None,
+        max_score: Optional[float] = None,
         output_file: Union[None, str, Path] = None,
         **kwargs: Any
     ) -> None:
@@ -33,6 +35,16 @@ def __call__(
         mode : str, default 'pointmutant'.
             Alternative set to "mean" for the mean of each position.
 
+        min_score : float, default None
+            Change values below a minimum score to be that score.
+            i.e., setting min_score = -1 will change any value smaller
+            than -1 to -1.
+
+        max_score : float, default None
+            Change values below a maximum score to be that score.
+            i.e., setting max_score = 1 will change any value greater
+            than 1 to 1.
+
         output_file : str, default None
             If you want to export the generated graph, add the path and name
             of the file. Example: 'path/filename.png' or 'path/filename.svg'.
@@ -62,6 +74,8 @@ def __call__(
             scatter_obj_1(
                 scatter_obj_2,
                 mode,
+                min_score,
+                max_score,
                 combination[0],
                 combination[1],
                 output_file_edited,

mutagenesis_visualization/main/utils/kwargs.py

@@ -82,6 +82,7 @@ def generate_default_kwargs() -> Dict[str, Any]:
 
     close: boolean, default False
         Whether to execute plt.close() or not on a matplotlib object.
+
     random_state : int, default 554
         Random state used for PCA function.
 
@@ -96,6 +97,7 @@ def generate_default_kwargs() -> Dict[str, Any]:
 
     figsize_y : int
 
+    legend_fontsize: int, default 10
 
     Returns
     --------
@@ -137,6 +139,7 @@ def generate_default_kwargs() -> Dict[str, Any]:
         'title_fontsize': 12,
         'x_label_fontsize': 10,
         'y_label_fontsize': 10,
+        'legend_fontsize': 10,
     }
 
     return default_kwargs

mutagenesis_visualization/main/utils/replicates_screen_input.py

@@ -5,6 +5,7 @@
 
 from typing import List, Union
 from dataclasses import dataclass
+from copy import deepcopy
 import numpy as np
 from numpy import typing as npt
 from pandas import DataFrame
@@ -47,6 +48,18 @@ def df_wildtype_scores(self) -> List[DataFrame]:
     def df_notstopcodons_mean(self) -> DataFrame:
         return [df.groupby(['Position'], as_index=False).mean() for df in self.df_notstopcodons]
 
+    def df_notstopcodons_limit_score(self, min_score: float, max_score: float) -> List[DataFrame]:
+        """
+        Make values above a threshold to be the threshold.
+        using df.where was not respecting the NaN values.
+        """
+        df_list: List[DataFrame] = deepcopy(self.df_notstopcodons)
+        for df in df_list:
+            df.loc[df["Score"] < min_score, "Score"] = min_score
+            df.loc[df["Score"] > max_score, "Score"] = max_score
+            df.loc[df["Score_NaN"] < min_score, "Score_NaN"] = min_score
+            df.loc[df["Score_NaN"] > max_score, "Score_NaN"] = max_score
+        return df_list
 
 def handle_input_datasets(
     datasets: Union[npt.NDArray, DataFrame, List[Union[npt.NDArray, DataFrame]]]