Update AALogo().get_conservation

Author: breimanntools

aaanalysis/_utils/check_data.py

@@ -234,16 +234,19 @@ def check_superset_subset(subset=None, superset=None, name_subset=None, name_sup
 
 # df checking functions
 def check_df(name="df", df=None, accept_none=False, accept_nan=True, check_all_positive=False,
-             cols_requiered=None, cols_forbidden=None, cols_nan_check=None, str_add=None):
+             check_series=False, cols_requiered=None, cols_forbidden=None, cols_nan_check=None,
+             str_add=None):
     """Check if the provided DataFrame meets various criteria such as NaN values, required/forbidden columns, etc."""
     # Check DataFrame and values
     if df is None:
         if not accept_none:
             raise ValueError(f"'{name}' should not be None")
         else:
             return None
-    if not isinstance(df, pd.DataFrame):
-        str_error = add_str(str_error=f"'{name}' ({type(df)}) should be DataFrame",
+    _check_dtype = pd.DataFrame if not check_series else pd.Series
+    _str_check_dtype = "DataFrame" if not check_series else "Series"
+    if not isinstance(df, _check_dtype):
+        str_error = add_str(str_error=f"'{name}' ({type(df)}) should be {_str_check_dtype}",
                             str_add=str_add)
         raise ValueError(str_error)
     if not accept_nan and df.isna().any().any():
@@ -278,11 +281,12 @@ def check_df(name="df", df=None, accept_none=False, accept_nan=True, check_all_p
             str_error = add_str(str_error=f"NaN values are not allowed in '{cols_nan_check}'.",
                                 str_add=str_add)
             raise ValueError(str_error)
-    columns = list(df)
-    cols_duplicated = [x for x in columns if columns.count(x) > 1]
-    if len(cols_duplicated) > 0:
-        str_error = add_str(str_error=f"The following columns are duplicated '{cols_duplicated}'.", str_add=str_add)
-        raise ValueError(str_error)
+    if _check_dtype == pd.DataFrame:
+        columns = list(df)
+        cols_duplicated = [x for x in columns if columns.count(x) > 1]
+        if len(cols_duplicated) > 0:
+            str_error = add_str(str_error=f"The following columns are duplicated '{cols_duplicated}'.", str_add=str_add)
+            raise ValueError(str_error)
 
 
 def check_warning_consecutive_index(name="df", df=None):

aaanalysis/seq_analysis/_aalogo.py

@@ -9,6 +9,15 @@
 
 
 # I Helper function
+def check_df_logo_info(df_logo_info=None):
+    """Check if df_logo_info has correct format"""
+    ut.check_df(name="df_logo_info", df=df_logo_info,
+                check_series=True, accept_none=False, accept_nan=False)
+    # Additional check specific to logo info: index name
+    if df_logo_info.index.name != "pos":
+        raise ValueError("Index name must be 'pos'")
+
+
 def _adjust_tmd(df_parts=None, tmd_len=None, start_n=False, ):
     """Adjust TMD to have similar length for df logo"""
     if ut.COL_TMD in list(df_parts):
@@ -92,22 +101,19 @@ def get_df_logo_info(self,
         df_logo_info = df_logo.sum(axis=1)  # vals_sum_per_pos
         return df_logo_info
 
-    def get_conservation(self,
-                         df_logo=None,
+    @staticmethod
+    def get_conservation(df_logo_info=None,
                          value_type: Literal["min", "mean", "median", "max"] = "mean"):
         """Compute conservation scores from sequence logos, ranging from 0 (no conservation) to
         4.248 (completely conserved)."""
-
-        if self._logo_type != "information":
-            raise ValueError("Conservation can only be computed for 'logo_type'='information'")
-        vals_sum_per_pos = df_logo.sum(axis=1)
+        check_df_logo_info(df_logo_info=df_logo_info)
         # Compute the statistic for each scale
         if value_type == "min":
-            conservation = vals_sum_per_pos.min()
+            cons_val = df_logo_info.min()
         elif value_type == "mean":
-            conservation = vals_sum_per_pos.mean()
+            cons_val = df_logo_info.mean()
         elif value_type == "median":
-            conservation = vals_sum_per_pos.median()
+            cons_val = df_logo_info.median()
         else:
-            conservation = vals_sum_per_pos.max()
-        return conservation
+            cons_val = df_logo_info.max()
+        return cons_val