Repository: jnothman/UpSetPlot

Issue:
Specify maximum number of intersections to display
Hi,

Apologies if this has already been implemented, but after going through the [API documentation](https://upsetplot.readthedocs.io/en/stable/api.html), I'm unable to find an option that limits the number of displayed intersections, which forces me to manually truncate my dataframe I'm trying to plot from. 

If you can point me to an argument that lets me do this that would be great. If this is not already implemented, it certainly would be useful!

Thanks!

Comments:

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 9624e07..27421f4 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -6,6 +6,7 @@ What's new in version 0.9
   introduced in version 0.7. (:issue:`248`)
 - Ability to disable totals plot with `totals_plot_elements=0`. (:issue:`246`)
 - Ability to set totals y axis label (:issue:`243`)
+- Added ``max_subset_rank`` to get only n most populous subsets.
 
 What's new in version 0.8
 -------------------------
diff --git a/upsetplot/plotting.py b/upsetplot/plotting.py
index 0451215..a4d85cf 100644
--- a/upsetplot/plotting.py
+++ b/upsetplot/plotting.py
@@ -28,6 +28,7 @@ def _process_data(
     sum_over,
     min_subset_size=None,
     max_subset_size=None,
+    max_subset_rank=None,
     min_degree=None,
     max_degree=None,
     reverse=False,
@@ -41,6 +42,7 @@ def _process_data(
         sum_over=sum_over,
         min_subset_size=min_subset_size,
         max_subset_size=max_subset_size,
+        max_subset_rank=max_subset_rank,
         min_degree=min_degree,
         max_degree=max_degree,
         include_empty_subsets=include_empty_subsets,
@@ -200,6 +202,11 @@ class UpSet:
         a size greater than this threshold will be omitted from plotting.
 
         .. versionadded:: 0.5
+    max_subset_rank : int, optional
+        Limit to the top N ranked subsets in descending order of size.
+        All tied subsets are included.
+
+        .. versionadded:: 0.9
     min_degree : int, optional
         Minimum degree of a subset to be shown in the plot.
 
@@ -270,6 +277,7 @@ class UpSet:
         sum_over=None,
         min_subset_size=None,
         max_subset_size=None,
+        max_subset_rank=None,
         min_degree=None,
         max_degree=None,
         facecolor="auto",
@@ -324,6 +332,7 @@ class UpSet:
             sum_over=sum_over,
             min_subset_size=min_subset_size,
             max_subset_size=max_subset_size,
+            max_subset_rank=max_subset_rank,
             min_degree=min_degree,
             max_degree=max_degree,
             reverse=not self._horizontal,
@@ -345,6 +354,7 @@ class UpSet:
         absent=None,
         min_subset_size=None,
         max_subset_size=None,
+        max_subset_rank=None,
         min_degree=None,
         max_degree=None,
         facecolor=None,
@@ -371,6 +381,11 @@ class UpSet:
             Minimum size of a subset to be styled.
         max_subset_size : int, optional
             Maximum size of a subset to be styled.
+        max_subset_rank : int, optional
+            Limit to the top N ranked subsets in descending order of size.
+            All tied subsets are included.
+
+            .. versionadded:: 0.9
         min_degree : int, optional
             Minimum degree of a subset to be styled.
         max_degree : int, optional
@@ -405,6 +420,7 @@ class UpSet:
             absent=absent,
             min_subset_size=min_subset_size,
             max_subset_size=max_subset_size,
+            max_subset_rank=max_subset_rank,
             min_degree=min_degree,
             max_degree=max_degree,
         )
diff --git a/upsetplot/reformat.py b/upsetplot/reformat.py
index aaca407..821a359 100644
--- a/upsetplot/reformat.py
+++ b/upsetplot/reformat.py
@@ -94,7 +94,14 @@ def _scalar_to_list(val):
 
 
 def _get_subset_mask(
-    agg, min_subset_size, max_subset_size, min_degree, max_degree, present, absent
+    agg,
+    min_subset_size,
+    max_subset_size,
+    max_subset_rank,
+    min_degree,
+    max_degree,
+    present,
+    absent,
 ):
     """Get a mask over subsets based on size, degree or category presence"""
     subset_mask = True
@@ -102,6 +109,10 @@ def _get_subset_mask(
         subset_mask = np.logical_and(subset_mask, agg >= min_subset_size)
     if max_subset_size is not None:
         subset_mask = np.logical_and(subset_mask, agg <= max_subset_size)
+    if max_subset_rank is not None:
+        subset_mask = np.logical_and(
+            subset_mask, agg.rank(method="min", ascending=False) <= max_subset_rank
+        )
     if (min_degree is not None and min_degree >= 0) or max_degree is not None:
         degree = agg.index.to_frame().sum(axis=1)
         if min_degree is not None:
@@ -121,12 +132,21 @@ def _get_subset_mask(
 
 
 def _filter_subsets(
-    df, agg, min_subset_size, max_subset_size, min_degree, max_degree, present, absent
+    df,
+    agg,
+    min_subset_size,
+    max_subset_size,
+    max_subset_rank,
+    min_degree,
+    max_degree,
+    present,
+    absent,
 ):
     subset_mask = _get_subset_mask(
         agg,
         min_subset_size=min_subset_size,
         max_subset_size=max_subset_size,
+        max_subset_rank=max_subset_rank,
         min_degree=min_degree,
         max_degree=max_degree,
         present=present,
@@ -189,6 +209,7 @@ def query(
     absent=None,
     min_subset_size=None,
     max_subset_size=None,
+    max_subset_rank=None,
     min_degree=None,
     max_degree=None,
     sort_by="degree",
@@ -221,6 +242,11 @@ def query(
         Size may be a sum of values, see `subset_size`.
     max_subset_size : int, optional
         Maximum size of a subset to be reported.
+    max_subset_rank : int, optional
+        Limit to the top N ranked subsets in descending order of size.
+        All tied subsets are included.
+
+        .. versionadded:: 0.9
     min_degree : int, optional
         Minimum degree of a subset to be reported.
     max_degree : int, optional
@@ -348,6 +374,7 @@ def query(
         agg,
         min_subset_size=min_subset_size,
         max_subset_size=max_subset_size,
+        max_subset_rank=max_subset_rank,
         min_degree=min_degree,
         max_degree=max_degree,
         present=present,


Test patch:
diff --git a/upsetplot/tests/test_upsetplot.py b/upsetplot/tests/test_upsetplot.py
index 673f893..80dfa76 100644
--- a/upsetplot/tests/test_upsetplot.py
+++ b/upsetplot/tests/test_upsetplot.py
@@ -753,6 +753,14 @@ def test_index_must_be_bool(x):
                 (True, True, True): 990,
             },
         ),
+        (
+            {"max_subset_rank": 3},
+            {
+                (True, False, False): 884,
+                (True, True, False): 1547,
+                (True, True, True): 990,
+            },
+        ),
         (
             {"min_subset_size": 800, "max_subset_size": 990},
             {
@@ -822,6 +830,29 @@ def test_filter_subsets(filter_params, expected, sort_by):
     assert upset_full.total == pytest.approx(upset_filtered.total)
 
 
+def test_filter_subsets_max_subset_rank_tie():
+    data = generate_samples(seed=0, n_samples=5, n_categories=3)
+    tested_non_tie = False
+    tested_tie = True
+    full = UpSet(data, subset_size="count").intersections
+    prev = None
+    for max_rank in range(1, 5):
+        cur = UpSet(data, subset_size="count", max_subset_rank=max_rank).intersections
+        if prev is not None:
+            if cur.shape[0] > prev.shape[0]:
+                # check we add rows only when they are new
+                assert cur.min() < prev.min()
+                tested_non_tie = True
+            elif cur.shape[0] != full.shape[0]:
+                assert (cur == cur.min()).sum() > 1
+                tested_tie = True
+
+        prev = cur
+    assert tested_non_tie
+    assert tested_tie
+    assert cur.shape[0] == full.shape[0]
+
+
 @pytest.mark.parametrize(
     "x",
     [