Invalidate cached properties in delay distribution classes upon parameter updates and add a new DoublePoissonDelay class for discrete delays. Update method signatures for clarity.

Author: yichao-liang

predicators/approaches/pp_param_learning_approach.py

@@ -1,6 +1,7 @@
 import logging
 from collections import defaultdict
 from typing import Dict, List, Optional, Sequence, Set, Tuple
+from pprint import pformat
 
 import numpy as np
 from gym.spaces import Box
@@ -142,9 +143,9 @@ def _learn_process_parameters(self,
         progress_bar = tqdm(desc="Optim. params.", unit="iter")
 
         # 2. Define objective and optimize
-        def objective(params):
+        def objective(params) -> float:
             """Objective function for scipy.optimize.minimize to minimize.
-
+ˍ
             It does some preparation and then calls the -ELBO function.
             """
             nonlocal best_elbo
@@ -191,6 +192,10 @@ def objective(params):
             },
             method="L-BFGS-B")  # terminate in 19464iter
         progress_bar.close()
+        # Display the learned processes
+        logging.debug("Learned processes:")
+        for proc in self._processes:
+            logging.debug(pformat(proc))
         logging.info(f"Best likelihood bound: {-result.fun}")
         breakpoint()
 

predicators/structs.py

@@ -2460,6 +2460,11 @@ def filter_predicates(self, kept: Collection[Predicate]) -> CausalProcess:
     def _set_parameters(self, parameters: Sequence[float]) -> None:
         self.strength = parameters[0]
         self.delay_distribution.set_parameters(parameters[1:])
+        # Invalidate cached properties
+        if '_str' in self.__dict__:
+            del self.__dict__['_str']
+        if '_hash' in self.__dict__:
+            del self.__dict__['_hash']
 
     def delay_probability(self, delay: int) -> float:
         return self.delay_distribution.probability(delay)

predicators/utils.py

@@ -4316,6 +4316,11 @@ def sample(self):
 
     def set_parameters(self, parameters):
         self.delay = parameters[0]
+        # Invalidate cached properties
+        if '_str' in self.__dict__:
+            del self.__dict__['_str']
+        if '_hash' in self.__dict__:
+            del self.__dict__['_hash']
 
     def probability(self, k: int) -> float:
         return 1.0 if k == self.delay else 0.0
@@ -4341,6 +4346,11 @@ def sample(self):
     def set_parameters(self, parameters):
         self.mean = parameters[0]
         self.std = parameters[1]
+        # Invalidate cached properties
+        if '_str' in self.__dict__:
+            del self.__dict__['_str']
+        if '_hash' in self.__dict__:
+            del self.__dict__['_hash']
 
     def probability(self, k: int) -> float:
         return norm.pdf(k, self.mean, self.std)
@@ -4353,14 +4363,14 @@ def _str(self) -> str:
 class CMPDelay(DelayDistribution):
     """Conway-Maxwell-Poisson (CMP) distribution for delays."""
 
-    def __init__(self, lam: float, nu: float, rng: np.random.Generator):
+    def __init__(self, lam: float, nu: float, rng: np.random.Generator) -> None:
         self.lam = lam
         self.nu = nu
         self.rng = rng
         self._max_k = 100
         self._update_cache()
 
-    def _update_cache(self):
+    def _update_cache(self) -> None:
         """Precompute and cache PMF and CDF."""
         # Calculate log factorial values once
         log_factorials = np.array([gammaln(n + 1) for n in range(self._max_k)])
@@ -4381,10 +4391,15 @@ def _update_cache(self):
         self._pmf = masses / np.sum(masses)
         self._cdf = np.cumsum(self._pmf)
 
-    def set_parameters(self, parameters):
+    def set_parameters(self, parameters: Sequence[float]) -> None:
         self.lam = parameters[0]
         self.nu = parameters[1]
         self._update_cache()
+        # Invalidate cached properties
+        if '_str' in self.__dict__:
+            del self.__dict__['_str']
+        if '_hash' in self.__dict__:
+            del self.__dict__['_hash']
 
     def probability(self, k: int) -> float:
         """Return the probability of delay k."""
@@ -4396,12 +4411,85 @@ def probability(self, k: int) -> float:
     def _str(self) -> str:
         return f"CMPDelay({self.lam}, {self.nu})"
 
-    def sample(self):
+    def sample(self) -> None:
         """Sample from the CMP distribution using cached CDF."""
         u = self.rng.random()
         idx = np.searchsorted(self._cdf, u)
         return idx
 
+class DoublePoissonDelay(DelayDistribution):
+    """
+    Double-Poisson distribution for discrete delays.
+
+    Parameters
+    ----------
+    mu : float
+        Mean-like parameter ( > 0 ).
+    phi : float
+        Dispersion parameter ( > 0 ).
+        • phi < 1  → over-dispersion
+        • phi = 1 → Poisson
+        • phi > 1 → under-dispersion
+    rng : np.random.Generator
+        Numpy random generator used by `sample()`.
+    max_k : int, optional
+        Cache PMF/CDF up to this value (default 100).
+        Increase if your delays can be very large.
+    """
+
+    def __init__(self,
+                 mu: float,
+                 phi: float,
+                 rng: np.random.Generator,
+                 max_k: int = 50):
+        self.mu = mu
+        self.phi = phi
+        self.rng = rng
+        self._max_k = max_k
+        self._update_cache()
+
+    # ------------------------------------------------------------------ #
+    # Internals
+    # ------------------------------------------------------------------ #
+    def _update_cache(self) -> None:
+        """Vectorised pre-computation of PMF and CDF."""
+        ks = np.arange(self._max_k)                     # 0, 1, …, max_k-1
+        log_fact = gammaln(ks + 1)                      # log(k!)
+        # Saddle-point normaliser  C(μ, φ) ≈ (2π φ μ)^(-1/2)
+        log_C = -0.5 * (np.log(2 * np.pi) + np.log(self.phi) + np.log(self.mu))
+
+        # log P(Y=k)
+        log_p = (log_C
+                 + self.phi * (ks * np.log(self.mu) - log_fact)
+                 - self.phi * self.mu)
+
+        # Stabilise → exponentiate → renormalise
+        log_p -= log_p.max()                            # avoid overflow
+        pmf = np.exp(log_p)
+        pmf /= pmf.sum()
+
+        self._pmf = pmf
+        self._cdf = np.cumsum(pmf)
+
+    def set_parameters(self, parameters: Sequence[float]) -> None:
+        """Update μ and φ, then rebuild the cache."""
+        self.mu, self.phi = parameters
+        self._update_cache()
+
+    def probability(self, k: int) -> float:
+        """Fast O(1) lookup of P(delay = k)."""
+        if 0 <= k < self._max_k:
+            return float(self._pmf[k])
+        return 0.0
+
+    def sample(self) -> int:
+        """Inverse-CDF sampling using cached CDF (O(log max_k))."""
+        u = self.rng.random()
+        return int(np.searchsorted(self._cdf, u))
+
+    @cached_property
+    def _str(self) -> str:
+        return f"DoublePoissonDelay({self.mu}, {self.phi})"
 
 @functools.lru_cache(maxsize=None)
 def get_git_commit_hash() -> str:

scripts/configs/mara_bench.yaml

@@ -14,7 +14,7 @@ APPROACHES:
     ARGS:
       - "load_data"
     FLAGS:
-      boil_use_constant_delay: True
+      # boil_use_constant_delay: True
       demonstrator: "oracle_process_planning"
       # This following is needed to generate a successful demo trajectory, or 
       # else the demo collection will fail with an option plan exhausted error.