diff --git a/metasyn/distribution/categorical.py b/metasyn/distribution/categorical.py index bc1141bb..086579de 100644 --- a/metasyn/distribution/categorical.py +++ b/metasyn/distribution/categorical.py @@ -68,6 +68,25 @@ def _param_schema(cls): "probs": {"type": "array", "items": {"type": "number"}}, } + @property + def _params_formatted(self) -> str: + n_categories = len(self.labels) + top_bottom_counts = (4, 3) + if n_categories <= sum(top_bottom_counts): + return super()._params_formatted + sorted_indices = np.argsort(-self.probs) + top_indices = sorted_indices[:top_bottom_counts[0]] + bottom_indices = sorted_indices[-top_bottom_counts[1]:] + top_labels = ", ".join(repr(label) for label in self.labels[top_indices].tolist()) + top_probs = ", ".join(f"{prob:.3f}" for prob in self.probs[top_indices]) + bottom_labels = ", ".join(repr(label) for label in self.labels[bottom_indices].tolist()) + bottom_probs = ", ".join(f"{prob:.3f}" for prob in self.probs[bottom_indices]) + return ( + f"\t- labels: [{top_labels}, ..., {bottom_labels}]\n" + f"\t- probs: [{top_probs}, ..., {bottom_probs}]\n" + f"\tShowing {top_bottom_counts[0]+top_bottom_counts[1]} of {n_categories} categories" + ) + def draw(self): return np.random.choice(self.labels, p=self.probs)