matrix_patches

  1"""
  2[matrix_patches](https://github.com/isblab/af_pipeline/tree/main/af_pipeline/tools/matrix_patches.py)
  3==============================
  4
  5Class to get interacting patches from a binary matrix.
  6"""
  7
  8import numpy as np
  9import pandas as pd
 10from collections import defaultdict
 11from af_pipeline.utils.misc_utils import (
 12    get_key_from_res_range,
 13    get_res_range_from_key,
 14)
 15from af_pipeline.constants.af_constants import MiscStrEnum
 16
 17class MatrixPatches:
 18    """Class to get interacting patches from a binary matrix"""
 19
 20    matrix: np.ndarray
 21    """ Binary matrix where rows and columns represent different objects
 22    (e.g., chains in a protein complex). """
 23
 24    row_obj: str
 25    """ Identifier for the rows in the matrix. """
 26
 27    col_obj: str
 28    """ Identifier for the columns in the matrix. """
 29
 30    def __init__(
 31        self,
 32        matrix: np.ndarray,
 33        row_obj: str = "row_obj",
 34        col_obj: str = "col_obj",
 35    ):
 36        self.matrix = matrix
 37        self.row_obj = row_obj
 38        self.col_obj = col_obj
 39
 40    def get_patches_from_matrix(self):
 41        """Get all interacting patches from a binary matrix
 42
 43        Arguments:
 44
 45        - **matrix (np.ndarray)**:<br />
 46            Binary matrix where rows and columns represent different objects
 47            (e.g., chains in a protein complex).
 48
 49        - **row_obj (str)**:<br />
 50            Identifier for the rows in the matrix
 51
 52        - **col_obj (str)**:<br />
 53            Identifier for the columns in the matrix
 54
 55        Returns:
 56
 57        - **patches (dict)**:<br />
 58            Dictionary of interacting patches
 59
 60        Example:
 61
 62            >>> matrix = np.array([
 63            ... [0, 0, 0, 1],
 64            ... [0, 1, 1, 1],
 65            ... [0, 0, 1, 1],
 66            ... [0, 1, 0, 0],
 67            ... [0, 1, 0, 1]
 68            ... ])
 69            >>> matrix_patches = MatrixPatches(
 70            ... matrix, row_obj="A", col_obj="B"
 71            ... )
 72            >>> matrix_patches.get_patches_from_matrix()
 73                       A          B
 74            0  {0, 1, 2}        {3}
 75            1        {1}  {1, 2, 3}
 76            2     {1, 2}     {2, 3}
 77            3     {3, 4}        {1}
 78            4        {4}        {3}
 79        """
 80
 81        assert np.isin(self.matrix, [0, 1]).all() and np.any(self.matrix), (
 82            f"Matrix must be binary and non-empty, got {np.unique(self.matrix)}"
 83        )
 84
 85        row_sets = self.get_one_sets_from_matrix(self.matrix, axis=0)
 86        col_sets = self.get_one_sets_from_matrix(self.matrix, axis=1)
 87
 88        split_row_sets = self.extend_one_sets_by_subsets(row_sets)
 89        split_col_sets = self.extend_one_sets_by_subsets(col_sets)
 90
 91        df_row = self.one_sets_to_df(
 92            split_row_sets, [self.row_obj, self.col_obj]
 93        )
 94        df_col = self.one_sets_to_df(
 95            split_col_sets, [self.col_obj, self.row_obj]
 96        )
 97
 98        df_row = self.aggregate_df_rows(df_row, self.col_obj, self.row_obj)
 99        df_col = self.aggregate_df_rows(df_col, self.row_obj, self.col_obj)
100
101        combined_df = self.combine_dfs(
102            df_row, df_col, self.row_obj, self.col_obj
103        )
104
105        for col in [self.row_obj, self.col_obj]:
106            combined_df[col] = combined_df[col].apply(
107                get_res_range_from_key, return_type="set"
108            )
109
110        combined_df = self.remove_subset_rows(
111            combined_df, self.row_obj, self.col_obj
112        )
113
114        return combined_df
115
116    @staticmethod
117    def get_one_sets_from_matrix(matrix: np.ndarray, axis: int = 0):
118        """Get the indices of 1s in a binary matrix rowwise or columnwise.
119
120        Arguments:
121
122        - **matrix (np.ndarray)**:<br />
123            Binary matrix where rows and columns represent different objects
124            (e.g., chains in a protein complex).
125
126        - **axis (int, optional)**:<br />
127            0 for rowwise, 1 for columnwise.
128
129        Returns:
130
131        - **one_sets (dict)**:<br />
132            `{k:v}` where `v` is a set of indices of 1s for key `k`.
133
134        Example:
135
136            >>> matrix = np.array([
137            ... [1, 0, 1],
138            ... [0, 1, 0],
139            ... [1, 1, 0]
140            ... ])
141            >>> MatrixPatches.get_one_sets_from_matrix(matrix=matrix, axis=0)
142            {0: {np.int64(0), np.int64(2)}, 1: {np.int64(1)}, 2: {np.int64(0), np.int64(1)}}
143            >>> MatrixPatches.get_one_sets_from_matrix(matrix=matrix, axis=1)
144            {0: {np.int64(0), np.int64(2)}, 1: {np.int64(1), np.int64(2)}, 2: {np.int64(0)}}
145        """
146
147        assert np.isin(matrix, [0, 1]).all() and np.any(matrix), (
148            f"Matrix must be binary and non-empty, got {np.unique(matrix)}"
149        )
150
151        one_sets = {}
152
153        if axis == 0:  # row_sets
154            for i in range(matrix.shape[0]):
155                one_sets[i] = set(np.where(matrix[i] == 1)[0])
156
157        elif axis == 1:  # col_sets
158            for j in range(matrix.shape[1]):
159                one_sets[j] = set(np.where(matrix[:, j] == 1)[0])
160
161        return one_sets
162
163    @staticmethod
164    def extend_one_sets_by_subsets(one_sets: dict) -> dict:
165        """Add the subsets of the sets in list_of_sets to the one_sets.
166
167        Arguments:
168
169        - **one_sets (dict)**:<br />
170            `{k:v}` where `v` is a set of indices of 1s for key `k`.
171
172        Returns:
173
174        - **new_one_sets (dict)**:<br />
175            `{k:v}` where `v` is a list of sets of indices of 1s for key `k`
176            each set is a subset of the original set and is present in
177            the values of `one_sets`.
178
179        Example:
180
181            >>> one_sets = {
182            ... 0: {0, 1, 2, 3, 5, 6},
183            ... 1: {1},
184            ... 2: {0, 1}
185            ... }
186            >>> MatrixPatches.extend_one_sets_by_subsets(one_sets)
187            {0: [{0, 1, 2, 3}, {5, 6}, {1}, {0, 1}], 1: [{1}], 2: [{1}, {0, 1}]}
188        """
189
190        split_sets = MatrixPatches.split_one_sets(one_sets)
191
192        new_one_sets = defaultdict(list)
193        list_of_sets = []  # unique sets from split_sets
194        list_of_sets = [
195            set(x)
196            for xs in split_sets.values()
197            for x in xs
198            if set(x) not in list_of_sets
199        ]
200
201        for set1 in list_of_sets:
202            for idx, one_set in one_sets.items():
203                if set1.issubset(one_set):
204                    (
205                        new_one_sets[idx].append(set1)
206                        if set1 not in new_one_sets[idx]
207                        else None
208                    )
209
210        return dict(new_one_sets)
211
212    @staticmethod
213    def split_one_sets(one_sets: dict) -> dict:
214        """Split the sets in `one_sets` into sub-sets such that
215        each subset only contains consecutive indices.
216
217        Arguments:
218
219        - **one_sets (dict)**:<br />
220            `{k:v}` where `v` is a set of indices of 1s for key `k`.
221
222        Returns:
223
224        - **new_one_sets (dict)**:<br />
225            dictionary of lists of lists where each list contains the
226            indices of 1s.
227
228        Example:
229
230            >>> one_sets = {0: {0, 1, 2, 3, 5, 6}, 1: {1}, 2: {0, 1}}
231            >>> MatrixPatches.split_one_sets(one_sets)
232            {0: [[0, 1, 2, 3], [5, 6]], 1: [[1]], 2: [[0, 1]]}
233        """
234
235        new_one_sets = {}
236
237        for i, one_set in one_sets.items():
238
239            if not isinstance(one_set, set):
240                raise TypeError("one_set must be a set")
241
242            sub_sets = MatrixPatches.split_one_set(one_set)
243            new_one_sets[i] = sub_sets
244
245        return new_one_sets
246
247    @staticmethod
248    def split_one_set(one_set: set | list) -> list:
249        """Split a set of indices into sub-sets such that
250        each subset only contains consecutive indices.
251
252        Arguments:
253
254        - **one_set (set | list)**:<br />
255            Set of indices of 1s.
256
257        Returns:
258
259        - **sub_sets (list)**:<br />
260            List of lists where each list contains the indices of 1s.
261
262        Example:
263
264            >>> one_set = {0, 1, 2, 3, 5, 6} \n
265            >>> MatrixPatches.split_one_set(one_set)
266            [[0, 1, 2, 3], [5, 6]]
267        """
268
269        assert isinstance(
270            one_set, set | list
271        ), "one_set must be a set or a list"
272
273        sub_sets = []
274
275        if isinstance(one_set, list):
276            # need to remove duplicates if any
277            one_set = set(one_set)
278
279        one_set = sorted(list(one_set))
280
281        for idx, one_pos in enumerate(one_set):
282
283            curr_pos = one_pos
284            prev_pos = one_set[idx - 1] if idx > 0 else None
285
286            if idx == 0:
287                # If it's the first position, create a new sub-set
288                sub_sets.append([curr_pos])
289
290            elif curr_pos - prev_pos == 1:
291                # If the current position is consecutive to the previous one
292                # add it to the last sub-set
293                sub_sets[-1].append(one_pos)
294
295            else:
296                # If the current position is not consecutive to the previous one
297                # create a new sub-set
298                sub_sets.append([curr_pos])
299
300        return sub_sets
301
302    @staticmethod
303    def one_sets_to_df(
304        one_sets: dict,
305        columns: list
306    ):
307        """Convert a dictionary to a pandas DataFrame.
308
309        Arguments:
310
311        - **one_sets (dict)**:<br />
312            Dictionary to convert.
313
314        - **columns (list)**:<br />
315            Column names.
316
317        Returns:
318
319        - df (pd.DataFrame)**:<br />
320            DataFrame with the dictionary keys as first column and values
321            as second column in columns.
322
323        Example:
324
325            >>> one_sets = {
326            ... 1: [{1, 2}, {5}],
327            ... 2: [{4, 5}, {6}]
328            ... }
329            >>> columns = ["A", "B"]
330            >>> MatrixPatches.one_sets_to_df(one_sets, columns)
331               A       B
332            0  1  {1, 2}
333            1  1     {5}
334            2  2  {4, 5}
335            3  2     {6}
336        """
337
338        if all([isinstance(val, list) for val in one_sets.values()]):
339
340            df_rows = []
341
342            for k, v in one_sets.items():
343                for val in v:
344                    df_rows.append([str(k), val])
345
346            df = pd.DataFrame(df_rows, columns=columns)
347
348        else:
349            raise ValueError("All values in the dictionary must be lists.")
350
351        return df
352
353    @staticmethod
354    def aggregate_df_rows(
355        df: pd.DataFrame,
356        groupby_col: str,
357        agg_col: str
358    ):
359        """Group a DataFrame by a column and aggregate another column.
360
361        Arguments:
362
363        - **df (pd.DataFrame)**:<br />
364            DataFrame with groupby_col and agg_col.
365
366        - **groupby_col (str)**:<br />
367            Column to group by (each value is a set).
368
369        - **agg_col (str)**:<br />
370            Column to aggregate (each value is a string).
371
372        Returns:
373
374        - **df_group (pd.DataFrame)**:<br />
375            Grouped DataFrame with both columns as a set.
376
377        Example:
378
379            >>> df = pd.DataFrame({
380            ... "A": ["1", "1", "1", "2", "3", "4"],
381            ... "B": [{1}, {1,2}, {5}, {4,5}, {1,2}, {1,2}]
382            ... })
383            >>> MatrixPatches.aggregate_df_rows(df, "B", "A")
384                    B          A
385            0     {1}        {1}
386            1  {1, 2}  {1, 3, 4}
387            2  {4, 5}        {2}
388            3     {5}        {1}
389        """
390
391        df_group = (
392            df.groupby(df[groupby_col].map(tuple))[agg_col]
393            .apply(",".join)
394            .reset_index()
395        )
396        df_group[agg_col] = df_group[agg_col].astype(object)
397        for idx, row in df_group.iterrows():
398            one_set = row[agg_col].split(",")
399            one_set = [int(x) for x in one_set]
400            one_set = sorted(one_set)
401            df_group.at[idx, agg_col] = set(one_set)
402
403        df_group[groupby_col] = df_group[groupby_col].apply(set)
404
405        return df_group
406
407    @staticmethod
408    def combine_dfs(
409        df1: pd.DataFrame,
410        df2: pd.DataFrame,
411        colname_1: str,
412        colname_2: str
413    ):
414        """Combine two DataFrames with columns colname_1 and colname_2
415        into a new DataFrame with interacting residues ranges without duplicates.
416
417        Arguments:
418
419        - **df1 (pd.DataFrame)**:<br />
420            DataFrame 1.
421
422        - **df2 (pd.DataFrame)**:<br />
423            DataFrame 2.
424
425        - **colname_1 (str)**:<br />
426            Column name 1.
427
428        - **colname_2 (str)**:<br />
429            Column name 2.
430
431        Returns:
432
433        - **new_df (pd.DataFrame)**:<br />
434            Combined DataFrame of interacting residues ranges without
435            duplicates.
436
437        Example:
438
439            >>> df1 = pd.DataFrame({
440            ... "A":[{1, 3, 4}, {1}, {1, 2}, {0, 1, 2, 4}],
441            ... "B":[{1}, {1, 2, 3}, {2, 3}, {3}]
442            ... })
443            >>> df2 = pd.DataFrame({
444            ... "A": [{0, 1, 2}, {1}, {1, 2}, {3, 4}, {4}],
445            ... "B": [{3}, {1, 2, 3}, {2, 3}, {1}, {1, 3}]
446            ... })
447
448            >>> MatrixPatches.combine_dfs(df1, df2, "A", "B")
449                 A    B
450            0  0-2    3
451            1    1  1-3
452            2  1-2  2-3
453            3  3-4    1
454            4    4    1
455            5    4    3
456            6    1    1
457        """
458
459        combined_df = pd.concat([df2, df1], axis=0)
460        combined_df.reset_index(drop=True, inplace=True)
461
462        df_rows = []
463
464        for _, row in combined_df.iterrows():
465            if isinstance(row[colname_1], set) and isinstance(
466                row[colname_2], set
467            ):
468                ranges1 = get_key_from_res_range(row[colname_1], as_list=True)
469                ranges2 = get_key_from_res_range(row[colname_2], as_list=True)
470                assert isinstance(ranges1, list) and isinstance(ranges2, list)
471                for res_range1 in ranges1:
472                    for res_range2 in ranges2:
473                        df_rows.append([res_range1, res_range2])
474
475        new_df = pd.DataFrame(df_rows, columns=[colname_1, colname_2])
476        new_df.drop_duplicates(inplace=True, keep=MiscStrEnum.FIRST)
477        new_df.reset_index(drop=True, inplace=True)
478
479        return new_df
480
481    @staticmethod
482    def remove_subset_rows(
483        df: pd.DataFrame,
484        colname_1: str,
485        colname_2: str
486    ):
487        """Remove rows that are subsets of other rows.
488        (from chatgpt)
489
490        Arguments:
491
492        - **df (pd.DataFrame)**:<br />
493            DataFrame with columns `colname_1` and `colname_2`.
494
495        - **colname_1 (str)**:<br />
496            column name 1.
497
498        - **colname_2 (str)**:<br />
499            column name 2.
500
501        Returns:
502
503        - **filtered_df (pd.DataFrame)**:<br />
504            DataFrame with subset rows removed.
505
506        Example:
507
508        >>> df = pd.DataFrame({
509        ... "A": [{0, 1, 2}, {1}, {1, 2}, {3, 4}, {4}, {4}, {1}],
510        ... "B": [{3}, {1, 2, 3}, {2, 3}, {1}, {1}, {3}, {1}]
511        ... })
512        >>> MatrixPatches.remove_subset_rows(df, "A", "B")
513                   A          B
514        0  {0, 1, 2}        {3}
515        1        {1}  {1, 2, 3}
516        2     {1, 2}     {2, 3}
517        3     {3, 4}        {1}
518        4        {4}        {3}
519        """
520
521        rows_to_keep = []
522
523        for i, row in df.iterrows():
524
525            if not any(
526                MatrixPatches.is_subset(
527                    row, df.iloc[j], colname_1, colname_2
528                )
529                for j in range(len(df))
530                if i != j
531            ):
532                rows_to_keep.append(i)
533
534        filtered_df = df.loc[rows_to_keep].reset_index(drop=True)
535
536        return filtered_df
537
538    @staticmethod
539    def is_subset(
540        row: pd.Series,
541        other_row: pd.Series,
542        colname_1: str,
543        colname_2: str,
544    ):
545        """Check if row is a subset of other_row for two specified columns.
546
547        Arguments:
548
549        - **row (pd.Series)**:<br />
550            Row to check if it is a subset of other_row.
551
552        - **other_row (pd.Series)**:<br />
553            Row to check against.
554
555        - **colname_1 (str)**:<br />
556            Column name 1.
557
558        - **colname_2 (str)**:<br />
559            Column name 2.
560
561        Returns:
562
563        - **(bool)**:<br />
564            `True` if row is a subset of `other_row`, `False` otherwise.
565
566        Example:
567
568            >>> row = pd.Series({"A": {0, 1, 2}, "B": {3}})
569            >>> other_row = pd.Series({"A": {0, 1, 2, 3}, "B": {3}})
570            >>> MatrixPatches.is_subset(row, other_row, "A", "B")
571            True
572            >>> other_row = pd.Series({"A": {1, 2}, "B": {3}})
573            >>> MatrixPatches.is_subset(row, other_row, "A", "B")
574            False
575        """
576
577        return (
578            row[colname_1].issubset(other_row[colname_1])
579            and row[colname_2].issubset(other_row[colname_2])
580        )
581
582
583if __name__ == "__main__":
584
585    import doctest
586    doctest.testmod()