Dev: add mklsts [--random] [--modulo] [--cont] [--plot] flags

Author: niklases

pypef/main.py

@@ -111,6 +111,7 @@
 
 Usage:
     pypef mklsts --wt WT_FASTA --input CSV_FILE
+        [--random] [--modulo] [--cont] [--plot]
         [--drop THRESHOLD] [--sep CSV_COLUMN_SEPARATOR] [--mutation_sep MUTATION_SEPARATOR]
         [--numrnd NUMBER] [--ls_proportion LS_PROPORTION]
     pypef mkps --wt WT_FASTA [--input CSV_FILE]
@@ -303,6 +304,7 @@
 schema = Schema({
     Optional('--all'): bool,
     Optional('--conc'): bool,
+    Optional('--cont'): bool,
     Optional('--csvaa'): bool,
     Optional('--ddiverse'): bool,
     Optional('--drecomb'): bool,
@@ -319,6 +321,7 @@
     Optional('--ls'): Or(None, str),
     Optional('--ls_proportion'): Or(None, Use(float)),
     Optional('--model'): Or(None, str),
+    Optional('--modulo'): bool,
     Optional('--msa'): Or(None, str),
     Optional('--mutation_sep'): Or(None, str),
     Optional('--negative'): bool,
@@ -331,7 +334,9 @@
     Optional('--params'): Or(None, str),
     Optional('--pdb'): Or(None, str),    
     Optional('--pmult'): bool,
+    Optional('--plot'): bool,
     Optional('--ps'): Or(None, str),
+    Optional('--random'): bool,
     Optional('--qdiverse'): bool,
     Optional('--qarecomb'): bool,
     Optional('--qirecomb'): bool,

pypef/utils/split.py

@@ -17,22 +17,25 @@ def __init__(
             self, 
             df_or_csv_file: str | PathLike | pd.DataFrame, 
             n_cv: int | None = None,
-            mutation_column: str | None = None, 
-            separator: str | None = None
+            mutation_column: str | None = None,
+            mutation_separator: str | None = None,
+            csv_separator: str | None = None
     ):
-        if mutation_column is None:
-            mutation_column = 'mutant'
         self.mutation_column = mutation_column
-        if separator is None:
-            separator = ','
-        self.separator = separator
+        if csv_separator is None:
+            csv_separator = ','
+        if mutation_separator is None:
+            mutation_separator = '/'
+        self.mutation_separator = mutation_separator
+        self.csv_separator = csv_separator
         if n_cv is None:
             n_cv = 5
         self.n_cv = n_cv
         if type(df_or_csv_file) == pd.DataFrame:
             self.df = df_or_csv_file
         else:
-            self.df = pd.read_csv(self.csv_file, sep=self.separator)
+            self.df = pd.read_csv(self.df_or_csv_file, sep=self.csv_separator)
+        print(f'Dataframe size: {self.df.shape[0]}')
         self.random_splits_train_indices_combined, self.random_splits_test_indices_combined = None, None
         self.modulo_splits_train_indices_combined, self.modulo_splits_test_indices_combined = None, None
         self.cont_splits_train_indices_combined, self.cont_splits_test_indices_combined = None, None
@@ -43,9 +46,23 @@ def __init__(
 
     def order_by_pos(self):
         if self.mutation_column is None:
-            self.mutation_column = 'mutant'
-        variants = self.df[self.mutation_column].to_list()
-        self.df['variant_pos'] = [int(v[1:-1]) for v in variants]
+            variants = self.df.iloc[:, 0].to_list()
+        else:
+            variants = self.df[self.mutation_column].to_list()
+        single_mut_idxs = []
+        for i, variant in enumerate(variants):
+            if not self.mutation_separator in variant:
+                single_mut_idxs.append(i)
+        if single_mut_idxs:   
+            self.df = self.df.loc[single_mut_idxs, :]
+            if len(single_mut_idxs) != self.df.size:
+                print(f'Removed multimutated variants from dataframe... '
+                      f'new dataframe size: {self.df.shape[0]}')
+        if self.mutation_column is None:
+            variants = self.df.iloc[:, 0].to_list()
+        else:
+            variants = self.df[self.mutation_column].to_list()
+        self.df.loc[:, 'variant_pos'] = [int(v[1:-1]) for v in variants]
         self.df['substitutions'] = [v[-1] for v in variants]
         self.df.sort_values(['variant_pos', 'substitutions'], ascending=[True, True], inplace=True)
         self.min_pos, self.max_pos = self.df['variant_pos'].to_numpy()[0], self.df['variant_pos'].to_numpy()[-1]
@@ -144,6 +161,31 @@ def get_all_split_indices(self):
             [self.modulo_splits_train_indices_combined, self.modulo_splits_test_indices_combined],
             [self.cont_splits_train_indices_combined, self.cont_splits_test_indices_combined]
         ]
+    
+    def _get_df_split_data(self, combined_train_indices, combined_test_indices):
+        train_split_data, test_split_data = [], []
+        for train_split, test_split in zip(combined_train_indices, combined_test_indices):
+            train_split_data.append(self.df.iloc[train_split, :])
+            test_split_data.append(self.df.iloc[test_split, :])
+        return train_split_data, test_split_data
+    
+    def get_random_df_split_data(self):
+        return self._get_df_split_data(
+            self.random_splits_train_indices_combined, 
+            self.random_splits_test_indices_combined
+        )
+
+    def get_modulo_df_split_data(self):
+        return self._get_df_split_data(
+            self.modulo_splits_train_indices_combined, 
+            self.modulo_splits_test_indices_combined
+        )
+
+    def get_continuous_df_split_data(self):
+        return self._get_df_split_data(
+            self.cont_splits_train_indices_combined, 
+            self.cont_splits_test_indices_combined
+        )
 
     def plot_distributions(self):
         fig, axs = plt.subplots(

pypef/utils/utils_run.py

@@ -34,12 +34,19 @@
 
 
 def run_pypef_utils(arguments):
-    if arguments['mklsts'] or ['mklsts_rnd'] or ['mklsts_mod'] or ['mklsts_cont'] or ['mklsts_plot']:
+    if arguments['mklsts']:
         wt_sequence = get_wt_sequence(arguments['--wt'])
         t_drop = float(arguments['--drop'])
         ls_proportion = arguments['--ls_proportion']
         logger.info(f'Length of provided sequence: {len(wt_sequence)} amino acids.')
-        logger.info(f'Training set proportion (--ls_proportion): {ls_proportion}.')
+        if True in [
+            arguments['--random'], arguments['--modulo'], 
+            arguments['--cont'], arguments['--plot']
+        ]:
+            logger.info(f'Ignoring set proportion (--ls_proportion).')
+        else:
+            logger.info(f'Training set proportion (--ls_proportion): {ls_proportion}.')
+            
         df = drop_rows(arguments['--input'], amino_acids, t_drop, 
                        arguments['--sep'], arguments['--mutation_sep'])
         no_rnd = arguments['--numrnd']
@@ -49,18 +56,33 @@ def run_pypef_utils(arguments):
         if len(single_variants) == 0:
             logger.info('Found no single substitution variants for possible recombination!')
 
-        if arguments['mklsts']:
+        if True in [
+            arguments['--random'], arguments['--modulo'], 
+            arguments['--cont'], arguments['--plot']
+        ]:
+            ds = DatasetSplitter(df)
+            if arguments['--random']:
+                train_data, test_data = ds.get_random_df_split_data()
+                print(train_data)
+            elif arguments['--modulo']:
+                train_data, test_data = ds.get_modulo_df_split_data()
+                print(train_data)
+            elif arguments['--cont']:
+                train_data, test_data = ds.get_continuous_df_split_data()
+                print(train_data)
+            elif arguments['--plot']:
+                ds.print_shapes()
+                ds.plot_distributions() 
+        else:
             sub_ls, val_ls, sub_ts, val_ts = make_sub_ls_ts(
                 single_variants, single_values, 
                 higher_variants, higher_values, 
                 ls_proportion
             )
             logger.info('Tip: You can edit your LS and TS datasets just by '
                         'cutting/pasting between the LS and TS fasta datasets.')
-
             make_fasta_ls_ts('LS.fasl', wt_sequence, sub_ls, val_ls)
             make_fasta_ls_ts('TS.fasl', wt_sequence, sub_ts, val_ts)
-
             try:
                 no_rnd = int(no_rnd)
             except ValueError:
@@ -77,19 +99,6 @@ def run_pypef_utils(arguments):
                     make_fasta_ls_ts('LS_random_' + str(random_set_counter) + '.fasl', wt_sequence, sub_ls, val_ls)
                     make_fasta_ls_ts('TS_random_' + str(random_set_counter) + '.fasl', wt_sequence, sub_ts, val_ts)
                     random_set_counter += 1
-        else:
-            ds = DatasetSplitter(df)
-            if arguments['mklsts_rnd']:
-                pass # TODO
-
-            elif arguments['mklsts_mod']:
-                pass # TODO
-
-            elif arguments['mklsts_cont']:
-                pass # TODO   
-
-            elif arguments['mklsts_plot']:
-                pass # TODO  
 
     elif arguments['mkps']:
         wt_sequence = get_wt_sequence(arguments['--wt'])