feat: probabilistic inference returns traces

Signed-off-by: Louis Mandel <lmandel@us.ibm.com>

Author: mandel

src/pdl/pdl_distributions.py

@@ -0,0 +1,70 @@
+# Adapted from mu-ppl: https://github.com/gbdrt/mu-ppl/blob/main/mu_ppl/distributions.py
+
+from typing import Any, Generic, TypeVar
+import numpy as np
+import numpy.random as rand
+from scipy.special import logsumexp
+import seaborn as sns
+
+T = TypeVar("T")
+
+class Categorical(Generic[T]):
+    """
+    Categorical distribution, i.e., finite support distribution where values can be of arbitrary type.
+    """
+
+    def __init__(self, tuples: list[tuple[T, float, list[Any]]]):
+        """
+        Args:
+            tuples: List of tuples (value, score, metadata), where the score is in log scale.
+        """
+        self.values, self.logits, self.metadata = zip(*tuples)
+        lse = logsumexp(self.logits)
+        self.probs = np.exp(self.logits - lse)  # type: ignore
+
+    def shrink(self) -> "Categorical[T]":
+        """
+        Create an equivalent distribution without duplicated values.
+        """
+        res: dict[T, tuple[float, list]] = {}
+        for v, w, m in zip(self.values, self.probs, self.metadata):
+            if v in res:
+                w_v, m_v = res[v]
+                res[v] = (w_v + w, m_v + m)
+            else:
+                res[v] = (w, m)
+        return Categorical([(v, w, m) for v, (w, m) in res.items()])
+
+
+    def sample(self) -> T:
+        u = rand.rand()
+        i = np.searchsorted(np.cumsum(self.probs), u)
+        return self.values[i]
+
+
+    def sort(self) -> "Categorical[T]":
+        d = self.shrink()
+        sorted_indices = np.argsort(d.logits)[::-1]
+        d.values = [d.values[i] for i in sorted_indices]
+        d.logits = np.array(d.logits)[sorted_indices]
+        d.probs = np.array(d.probs)[sorted_indices]
+        d.metadata = [d.metadata[i] for i in sorted_indices]
+        return d
+
+
+def viz(dist: Categorical[float], **kwargs):
+    """
+    Visualize a distribution
+    """
+    dist = dist.shrink()
+    if len(dist.values) < 100:
+        sns.barplot(x=dist.values, y=dist.probs, errorbar=None, **kwargs)
+    else:
+        sns.histplot(
+            x=dist.values,
+            weights=dist.probs,
+            bins=50,
+            kde=True,
+            stat="probability",
+            **kwargs,
+        )

src/pdl/pdl_infer.py

@@ -5,16 +5,16 @@
 import yaml
 from matplotlib import pyplot as plt
 from mu_ppl import viz
-from mu_ppl.distributions import Categorical
 
 from ._version import version
 from .pdl import InterpreterConfig
 from .pdl_ast import PdlLocationType, Program, ScopeType, get_default_model_parameters
+from .pdl_distributions import Categorical
 from .pdl_inference import (
     infer_importance_sampling,
     infer_importance_sampling_parallel,
-    infer_rejection,
-    infer_rejection_parallel,
+    infer_rejection_sampling,
+    infer_rejection_sampling_parallel,
     infer_smc,
     infer_smc_parallel,
 )
@@ -42,7 +42,7 @@ def exec_program(  # pylint: disable=too-many-arguments, too-many-positional-arg
     ppdl_config: Optional[PpdlConfig] = None,
     scope: Optional[ScopeType | dict[str, Any]] = None,
     loc: Optional[PdlLocationType] = None,
-    output: Literal["result", "all"] = "result",
+    # output: Literal["result", "all"] = "result",
 ) -> Categorical[Any]:
     ppdl_config = ppdl_config or PpdlConfig()
 
@@ -56,46 +56,42 @@ def exec_program(  # pylint: disable=too-many-arguments, too-many-positional-arg
     config["batch"] = 1
     config["event_loop"] = _LOOP
 
+    dist: Categorical[Any]
     match algo:
         case "is":
             dist = infer_importance_sampling(
-                prog, config, scope, loc, output, num_particles=num_particles
+                prog, config, scope, loc, num_particles=num_particles
             )
         case "parallel-is":
             dist = infer_importance_sampling_parallel(
                 prog,
                 config,
                 scope,
                 loc,
-                output,
                 num_particles=num_particles,
                 max_workers=max_workers,
             )
         case "smc":
-            dist = infer_smc(
-                prog, config, scope, loc, output, num_particles=num_particles
-            )
+            dist = infer_smc(prog, config, scope, loc, num_particles=num_particles)
         case "parallel-smc":
             dist = infer_smc_parallel(
                 prog,
                 config,
                 scope,
                 loc,
-                output,
                 num_particles=num_particles,
                 max_workers=max_workers,
             )
         case "rejection":
-            dist = infer_rejection(
-                prog, config, scope, loc, output, num_samples=num_particles
+            dist = infer_rejection_sampling(
+                prog, config, scope, loc, num_samples=num_particles
             )
         case "parallel-rejection":
-            dist = infer_rejection_parallel(
+            dist = infer_rejection_sampling_parallel(
                 prog,
                 config,
                 scope,
                 loc,
-                output,
                 num_samples=num_particles,
                 max_workers=max_workers,
             )
@@ -110,10 +106,10 @@ def exec_dict(  # pylint: disable=too-many-arguments, too-many-positional-argume
     ppdl_config: Optional[PpdlConfig] = None,
     scope: Optional[ScopeType | dict[str, Any]] = None,
     loc: Optional[PdlLocationType] = None,
-    output: Literal["result", "all"] = "result",
+    # output: Literal["result", "all"] = "result",
 ) -> Any:
     program = parse_dict(prog)
-    result = exec_program(program, config, ppdl_config, scope, loc, output)
+    result = exec_program(program, config, ppdl_config, scope, loc)
     return result
 
 
@@ -122,10 +118,10 @@ def exec_str(
     config: Optional[InterpreterConfig] = None,
     ppdl_config: Optional[PpdlConfig] = None,
     scope: Optional[ScopeType | dict[str, Any]] = None,
-    output: Literal["result", "all"] = "result",
+    # output: Literal["result", "all"] = "result",
 ) -> Any:
     program, loc = parse_str(prog)
-    result = exec_program(program, config, ppdl_config, scope, loc, output)
+    result = exec_program(program, config, ppdl_config, scope, loc)
     return result
 
 
@@ -134,14 +130,14 @@ def exec_file(
     config: Optional[InterpreterConfig] = None,
     ppdl_config: Optional[PpdlConfig] = None,
     scope: Optional[ScopeType | dict[str, Any]] = None,
-    output: Literal["result", "all"] = "result",
+    # output: Literal["result", "all"] = "result",
 ) -> Any:
     program, loc = parse_file(prog)
     if config is None:
         config = InterpreterConfig()
     if config.get("cwd") is None:
         config["cwd"] = Path(prog).parent
-    result = exec_program(program, config, ppdl_config, scope, loc, output)
+    result = exec_program(program, config, ppdl_config, scope, loc)
     return result