SMC inference (#1230)

* WIP SMC

Signed-off-by: Guillaume Baudart <[email protected]>

* Bug fix: HMM example

Signed-off-by: Guillaume Baudart <[email protected]>

* Patch performance SMC

Signed-off-by: Guillaume Baudart <[email protected]>

* Add hmm_nl_priors

Signed-off-by: Guillaume Baudart <[email protected]>

* feat: replay an execution (#1211)

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

* refactor: performance optimizations (#1228)

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

* WIP: Parallel loop SMC

Signed-off-by: Guillaume Baudart <[email protected]>

---------

Signed-off-by: Guillaume Baudart <[email protected]>
Signed-off-by: Louis Mandel <[email protected]>
Co-authored-by: Louis Mandel <[email protected]>

Author: gbdrt

examples/ppdl/hmm.pdl

@@ -0,0 +1,41 @@
+defs:
+  step:
+    function:
+      pre_x: number
+      y: number
+    return:
+      defs:
+        x:
+          lang: python
+          code: |
+            from mu_ppl import Gaussian
+            result = Gaussian(pre_x, 1).sample()
+        score: 
+          lang: python
+          code: |
+            from mu_ppl import Gaussian
+            result = Gaussian(x, 1).log_prob(y)
+      lastOf:
+        - factor: ${score}
+        - data: ${x}
+  obs:
+    array: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 25, 26, 27, 28, 29, 30]
+
+  pre_x: 0
+
+for:
+  y: ${obs}
+repeat:
+  defs:
+    pre_x:
+      call: ${step}
+      args:
+        pre_x: ${pre_x}
+        y: ${y}
+  data: ${pre_x}
+join:
+  as: array  
+  
+
+
+

examples/ppdl/hmm_nl_priors.pdl

@@ -0,0 +1,48 @@
+defs:
+  step:
+    function:
+      pre_x: number
+      y: number
+    return:
+      defs:
+        x:
+          model: ollama_chat/granite3.3:2b
+          parameters:
+            temperature: 1
+          input: We are modeling a random walk along a line. Generate a random number that is Gaussian distributed around ${pre_x}. We do not now the parameters of the random walk but we suspect that the new value should be greater than ${pre_x}. DO NOT GENERATE A PYTHON CODE, JUST ANSWER WITH THE NUMBER
+          parser: json
+          spec: number
+          fallback:
+            lang: python
+            code: |
+              from mu_ppl import Gaussian
+              result = Gaussian(pre_x, 1).sample()
+        score: 
+          lang: python
+          code: |
+            from mu_ppl import Gaussian
+            result = Gaussian(x, 1).log_prob(y)
+      lastOf:
+        - factor: ${score}
+        - data: ${x}
+  obs:
+    array: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+
+  pre_x: 0
+
+for:
+  y: ${obs}
+repeat:
+  defs:
+    pre_x:
+      call: ${step}
+      args:
+        pre_x: ${pre_x}
+        y: ${y}
+  data: ${pre_x}
+join:
+  as: array  
+  
+
+
+

src/pdl/pdl_infer.py

@@ -8,7 +8,9 @@
 from .pdl import InterpreterConfig, exec_program
 from .pdl_ast import get_default_model_parameters
 from .pdl_parser import parse_file
+from .pdl_smc import infer_smc
 from .pdl_utils import validate_scope
+from matplotlib import pyplot as plt
 
 
 def main():
@@ -33,7 +35,11 @@ def main():
         default=5,
     )
     parser.add_argument(
-        "-v", "--viz", help="Display the distribution of results", default=False
+        "-v",
+        "--viz",
+        help="Display the distribution of results",
+        default=False,
+        action="store_true",
     )
     parser.add_argument(
         "--version",
@@ -67,12 +73,22 @@ def main():
         yield_result=False, yield_background=False, batch=1, cwd=Path(args.pdl).parent
     )
     program, loc = parse_file(args.pdl)
-    with ImportanceSampling(num_particles=args.num_particles):
-        dist = infer(
-            lambda: exec_program(program, config, initial_scope, loc, "result")
-        )
+    # with ImportanceSampling(num_particles=args.num_particles):
+    # dist = infer(
+    #     lambda: exec_program(program, config, initial_scope, loc, "result")
+    # )
+
+    def model(replay):
+        config["replay"] = replay
+        result = exec_program(program, config, initial_scope, loc, "all")
+        state = result["replay"]
+        return result["result"], state
+
+    dist = infer_smc(args.num_particles, model)
+
     if args.viz:
         viz(dist)
+        plt.show()
     print(dist.sample())
     return 0
 

src/pdl/pdl_interpreter.py

@@ -144,6 +144,7 @@
 )
 from .pdl_schema_utils import get_json_schema  # noqa: E402
 from .pdl_schema_validator import type_check_args, type_check_spec  # noqa: E402
+from .pdl_smc import Resample
 from .pdl_utils import (  # noqa: E402
     GeneratorWrapper,
     apply_defaults,
@@ -314,7 +315,7 @@ def process_prog(
         loc,
     )
 
-    stdlib_scope = scope | PdlDict({"stdlib": stdlib_dict})
+    stdlib_scope = scope  # | PdlDict({"stdlib": stdlib_dict})
 
     result, document, final_scope, trace = process_block(
         state, stdlib_scope, block=prog.root, loc=loc
@@ -567,6 +568,8 @@ def process_advance_block_retry(  # noqa: C901
             break
         except KeyboardInterrupt as exc:
             raise exc from exc
+        except Resample as exc:
+            raise exc from exc
         except Exception as exc:
             do_retry = block.retry and trial_idx + 1 < trial_total
             if block.fallback is None and not do_retry:
@@ -610,7 +613,8 @@ def process_advance_block_retry(  # noqa: C901
                     trace=trace,
                 )
                 result = lazy_apply(checker, result)
-    factor(score)
+    if score != 0:
+        factor(score)
     return result, background, new_scope, trace
 
 
@@ -1147,6 +1151,9 @@ def loop_body(iidx, items):
             factor(weight)
             result = PdlConst(None)
             background = DependentContext([])
+            assert block.pdl__id is not None
+            state.replay[block.pdl__id] = None
+            raise Resample(state.replay)
         case EmptyBlock():
             result = PdlConst("")
             background = DependentContext([])

src/pdl/pdl_smc.py

@@ -0,0 +1,115 @@
+from typing import TypeVar, ParamSpec, Callable, Any
+from mu_ppl.distributions import Categorical
+from mu_ppl import ImportanceSampling
+from tqdm import tqdm
+from copy import deepcopy
+from concurrent.futures import ThreadPoolExecutor
+import asyncio
+
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+class Resample(Exception):
+    def __init__(self, state):
+        self.state = state
+
+
+def resample(particles: list[Any], scores: list[float]) -> list[Any]:
+    d = Categorical(list(zip(particles, scores)))
+    return [
+        d.sample() for _ in range(len(particles))
+    ]  # resample a new set of particles
+
+
+def _process_particle(state, model, num_particles):
+    """Process a single particle and return (result, state, score)"""
+    with ImportanceSampling(0) as sampler:
+        try:
+            result, new_state = model(state)
+            return result, new_state, sampler.score
+        except Resample as exn:
+            return None, exn.state, sampler.score
+
+
+def infer_smc(num_particles: int, model) -> Categorical[Any]:
+    """Sequential version"""
+    particles = [{} for _ in range(num_particles)]  # initialise the particles
+    results: list[Any] = []
+    scores: list[float] = []
+    while len(results) < num_particles:
+        states = []
+        scores = []
+        results = []
+        for state in particles:
+            result, state, score = _process_particle(state, model, num_particles)
+            if result is not None:
+                results.append(result)
+            states.append(state)
+            scores.append(score)
+        particles = resample(states, scores)
+    return Categorical(list(zip(results, scores)))
+
+
+# Warning: Parallel version conflict with the context managers for inference. Need fix!
+
+# def infer_smc(num_particles:int, model) -> Categorical[Any]:
+#     """Parallelized version using ThreadPoolExecutor"""
+#     particles = [{} for _ in range(num_particles)]  # initialise the particles
+#     results: list[Any] = []
+#     scores: list[float] = []
+#     while len(results) < num_particles:
+#         states = []
+#         scores = []
+#         results = []
+#         with ThreadPoolExecutor() as executor:
+#             future_to_particle = {
+#                 executor.submit(_process_particle, state, model, num_particles): state
+#                 for state in particles
+#             }
+#             for future in future_to_particle:
+#                 result, state, score = future.result()
+#                 if result is not None:
+#                     results.append(result)  # execute all the particles
+#                 states.append(state)
+#                 scores.append(score)
+#         particles = resample(states, scores)
+#     return Categorical(list(zip(results, scores)))
+
+
+# async def _process_particle_async(state, model, num_particles):
+#     with ImportanceSampling(num_particles) as sampler:
+#         try:
+#             loop = asyncio.get_event_loop()
+#             result, new_state = await loop.run_in_executor(None, lambda: model(state))
+#             return result, new_state, sampler.score
+#         except Resample as exn:
+#             return None, exn.state, sampler.score
+
+
+# async def infer_smc_async(num_particles: int, model) -> Categorical[Any]:
+#     """Parallelized version using Async"""
+#     particles = [{} for _ in range(num_particles)]  # initialise the particles
+#     results: list[Any] = []
+#     scores: list[float] = []
+#     while len(results) < num_particles:
+#         states = []
+#         scores = []
+#         results = []
+#         tasks = [
+#             _process_particle_async(state, model, num_particles)
+#             for state in particles
+#         ]
+#         particle_results = await asyncio.gather(*tasks)
+#         for result, state, score in particle_results:
+#             if result is not None:
+#                 results.append(result)
+#             states.append(state)
+#             scores.append(score)
+#         particles = resample(states, scores)
+#     return Categorical(list(zip(results, scores)))
+
+
+# def infer_smc_async(num_particles: int, model) -> Categorical[Any]:
+#     return asyncio.run(infer_smc_async(num_particles, model))