feat: add prob method to to the categorical distribution

Signed-off-by: Louis Mandel <[email protected]>

Author: mandel

src/pdl/pdl_distributions.py

@@ -51,6 +51,15 @@ def sort(self) -> "Categorical[T]":
         d.metadata = [d.metadata[i] for i in sorted_indices]
         return d
 
+    def prob(self, x: T) -> float:
+        dist = self.shrink()
+        try:
+            i = dist.values.index(x)
+            p = dist.probs[i]
+        except ValueError:
+            p = 0.0
+        return p
+
 
 def viz(dist: Categorical[float], **kwargs):
     """