|
22 | 22 | from pm4py.util.constants import CASE_CONCEPT_NAME, PARAMETER_CONSTANT_CASEID_KEY, PARAMETER_CONSTANT_ACTIVITY_KEY |
23 | 23 | from enum import Enum |
24 | 24 | from pm4py.util import exec_utils, constants, xes_constants |
| 25 | +from pm4py.statistics.variants.polars import get as variants_get |
25 | 26 | from typing import Optional, Dict, Any, Union, List |
26 | 27 | import polars as pl |
27 | 28 |
|
@@ -92,6 +93,44 @@ def apply( |
92 | 93 | return ret |
93 | 94 |
|
94 | 95 |
|
| 96 | +def filter_variants_top_k( |
| 97 | + log: pl.LazyFrame, |
| 98 | + k: int, |
| 99 | + parameters: Optional[Dict[Union[str, Parameters], Any]] = None, |
| 100 | +) -> pl.LazyFrame: |
| 101 | + """ |
| 102 | + Keeps the top-k variants of the log |
| 103 | +
|
| 104 | + Parameters |
| 105 | + ------------- |
| 106 | + log |
| 107 | + Event log LazyFrame |
| 108 | + k |
| 109 | + Number of variants that should be kept |
| 110 | + parameters |
| 111 | + Parameters |
| 112 | +
|
| 113 | + Returns |
| 114 | + ------------- |
| 115 | + filtered_log |
| 116 | + Filtered log |
| 117 | + """ |
| 118 | + if parameters is None: |
| 119 | + parameters = {} |
| 120 | + |
| 121 | + variants = variants_get.get_variants_count(log, parameters=parameters) |
| 122 | + variant_count = [] |
| 123 | + for variant, count in variants.items(): |
| 124 | + variant_count.append([variant, count]) |
| 125 | + variant_count = sorted( |
| 126 | + variant_count, key=lambda x: (x[1], x[0]), reverse=True |
| 127 | + ) |
| 128 | + variant_count = variant_count[: min(k, len(variant_count))] |
| 129 | + variants_to_filter = [x[0] for x in variant_count] |
| 130 | + |
| 131 | + return apply(log, variants_to_filter, parameters=parameters) |
| 132 | + |
| 133 | + |
95 | 134 | def apply_auto_filter( |
96 | 135 | df: pl.LazyFrame, |
97 | 136 | parameters: Optional[Dict[Union[str, Parameters], Any]] = None, |
|
0 commit comments