pyRDDLGym-jax 2.2__py3-none-any.whl → 2.4__py3-none-any.whl

pyRDDLGym_jax/core/logic.py

@@ -29,15 +29,29 @@
 #
 # ***********************************************************************
 
-from typing import Callable, Dict, Union
+
+from abc import ABCMeta, abstractmethod
+import traceback
+from typing import Callable, Dict, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 import jax.random as random
 import jax.scipy as scipy 
 
+from pyRDDLGym.core.debug.exception import raise_warning
+
+# more robust approach - if user does not have this or broken try to continue
+try:
+    from tensorflow_probability.substrates import jax as tfp
+except Exception:
+    raise_warning('Failed to import tensorflow-probability: '
+                  'compilation of some probability distributions will fail.', 'red')
+    traceback.print_exc()
+    tfp = None
+
 
-def enumerate_literals(shape, axis, dtype=jnp.int32):
+def enumerate_literals(shape: Tuple[int, ...], axis: int, dtype: type=jnp.int32) -> jnp.ndarray:
     literals = jnp.arange(shape[axis], dtype=dtype)
     literals = literals[(...,) + (jnp.newaxis,) * (len(shape) - 1)]
     literals = jnp.moveaxis(literals, source=0, destination=axis)
@@ -52,30 +66,35 @@ def enumerate_literals(shape, axis, dtype=jnp.int32):
 #
 # ===========================================================================
 
-class Comparison:
+class Comparison(metaclass=ABCMeta):
     '''Base class for approximate comparison operations.'''
     
+    @abstractmethod
     def greater_equal(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def greater(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def equal(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def sgn(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def argmax(self, id, init_params):
-        raise NotImplementedError
+        pass
     
 
 class SigmoidComparison(Comparison):
     '''Comparison operations approximated using sigmoid functions.'''
     
-    def __init__(self, weight: float=10.0):
-        self.weight = weight
+    def __init__(self, weight: float=10.0) -> None:
+        self.weight = float(weight)
         
     # https://arxiv.org/abs/2110.05651
     def greater_equal(self, id, init_params):
@@ -127,21 +146,23 @@ class SigmoidComparison(Comparison):
 #
 # ===========================================================================
 
-class Rounding:
+class Rounding(metaclass=ABCMeta):
     '''Base class for approximate rounding operations.'''
     
+    @abstractmethod
     def floor(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def round(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class SoftRounding(Rounding):
     '''Rounding operations approximated using soft operations.'''
     
-    def __init__(self, weight: float=10.0):
-        self.weight = weight
+    def __init__(self, weight: float=10.0) -> None:
+        self.weight = float(weight)
         
     # https://www.tensorflow.org/probability/api_docs/python/tfp/substrates/jax/bijectors/Softfloor
     def floor(self, id, init_params):
@@ -177,11 +198,12 @@ class SoftRounding(Rounding):
 #
 # ===========================================================================
 
-class Complement:
+class Complement(metaclass=ABCMeta):
     '''Base class for approximate logical complement operations.'''
     
+    @abstractmethod
     def __call__(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class StandardComplement(Complement):
@@ -210,16 +232,18 @@ class StandardComplement(Complement):
 # https://www.sciencedirect.com/science/article/abs/pii/016501149190171L
 # ===========================================================================
 
-class TNorm:
+class TNorm(metaclass=ABCMeta):
     '''Base class for fuzzy differentiable t-norms.'''
     
+    @abstractmethod
     def norm(self, id, init_params):
         '''Elementwise t-norm of x and y.'''
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def norms(self, id, init_params):
         '''T-norm computed for tensor x along axis.'''
-        raise NotImplementedError
+        pass
     
 
 class ProductTNorm(TNorm):
@@ -291,7 +315,7 @@ class YagerTNorm(TNorm):
     '''Yager t-norm given by the expression 
     (x, y) -> max(1 - ((1 - x)^p + (1 - y)^p)^(1/p)).'''
     
-    def __init__(self, p=2.0):
+    def __init__(self, p: float=2.0) -> None:
         self.p = float(p)
     
     def norm(self, id, init_params):
@@ -327,23 +351,32 @@ class YagerTNorm(TNorm):
 #
 # ===========================================================================
 
-class RandomSampling:
+class RandomSampling(metaclass=ABCMeta):
     '''Describes how non-reparameterizable random variables are sampled.'''
     
+    @abstractmethod
     def discrete(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def poisson(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def binomial(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
+    def negative_binomial(self, id, init_params, logic):
+        pass
+    
+    @abstractmethod
     def geometric(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def bernoulli(self, id, init_params, logic):
-        raise NotImplementedError
+        pass
     
     def __str__(self) -> str:
         return 'RandomSampling'
@@ -386,8 +419,7 @@ class SoftRandomSampling(RandomSampling):
     def _poisson_gumbel_softmax(self, id, init_params, logic):        
         argmax_approx = logic.argmax(id, init_params)
         def _jax_wrapped_calc_poisson_gumbel_softmax(key, rate, params):
-            ks = jnp.arange(0, self.poisson_bins)
-            ks = ks[(jnp.newaxis,) * jnp.ndim(rate) + (...,)]
+            ks = jnp.arange(self.poisson_bins)[(jnp.newaxis,) * jnp.ndim(rate) + (...,)]
             rate = rate[..., jnp.newaxis]
             log_prob = ks * jnp.log(rate + logic.eps) - rate - scipy.special.gammaln(ks + 1)
             Gumbel01 = random.gumbel(key=key, shape=jnp.shape(log_prob), dtype=logic.REAL)
@@ -400,10 +432,7 @@ class SoftRandomSampling(RandomSampling):
         less_approx = logic.less(id, init_params)
         def _jax_wrapped_calc_poisson_exponential(key, rate, params):
             Exp1 = random.exponential(
-                key=key, 
-                shape=(self.poisson_bins,) + jnp.shape(rate), 
-                dtype=logic.REAL
-            )
+                key=key,  shape=(self.poisson_bins,) + jnp.shape(rate), dtype=logic.REAL)
             delta_t = Exp1 / rate[jnp.newaxis, ...]
             times = jnp.cumsum(delta_t, axis=0)
             indicator, params = less_approx(times, 1.0, params)
@@ -411,72 +440,98 @@ class SoftRandomSampling(RandomSampling):
             return sample, params
         return _jax_wrapped_calc_poisson_exponential
 
+    # normal approximation to Poisson: Poisson(rate) -> Normal(rate, rate)
+    def _poisson_normal_approx(self, logic):
+        def _jax_wrapped_calc_poisson_normal_approx(key, rate, params):
+            normal = random.normal(key=key, shape=jnp.shape(rate), dtype=logic.REAL)
+            sample = rate + jnp.sqrt(rate) * normal
+            return sample, params    
+        return _jax_wrapped_calc_poisson_normal_approx
+    
     def poisson(self, id, init_params, logic):
-        def _jax_wrapped_calc_poisson_exact(key, rate, params):
-            sample = random.poisson(key=key, lam=rate, dtype=logic.INT)
-            sample = jnp.asarray(sample, dtype=logic.REAL)
-            return sample, params      
-              
         if self.poisson_exp_method:
             _jax_wrapped_calc_poisson_diff = self._poisson_exponential(
                 id, init_params, logic)
         else:
             _jax_wrapped_calc_poisson_diff = self._poisson_gumbel_softmax(
                 id, init_params, logic)
+        _jax_wrapped_calc_poisson_normal = self._poisson_normal_approx(logic)
         
+        # for small rate use the Poisson process or gumbel-softmax reparameterization
+        # for large rate use the normal approximation
         def _jax_wrapped_calc_poisson_approx(key, rate, params):
-            
-            # determine if error of truncation at rate is acceptable
             if self.poisson_bins > 0:
                 cuml_prob = scipy.stats.poisson.cdf(self.poisson_bins, rate)
-                approx_cond = jax.lax.stop_gradient(
-                    jnp.min(cuml_prob) > self.poisson_min_cdf)
+                small_rate = jax.lax.stop_gradient(cuml_prob >= self.poisson_min_cdf)
+                small_sample, params = _jax_wrapped_calc_poisson_diff(key, rate, params)
+                large_sample, params = _jax_wrapped_calc_poisson_normal(key, rate, params)
+                sample = jnp.where(small_rate, small_sample, large_sample)
+                return sample, params
             else:
-                approx_cond = False
-            
-            # for acceptable truncation use the approximation, use exact otherwise
-            return jax.lax.cond(
-                approx_cond, 
-                _jax_wrapped_calc_poisson_diff, 
-                _jax_wrapped_calc_poisson_exact, 
-                key, rate, params
-            )
+                return _jax_wrapped_calc_poisson_normal(key, rate, params)
         return _jax_wrapped_calc_poisson_approx        
     
-    def binomial(self, id, init_params, logic):
-        def _jax_wrapped_calc_binomial_exact(key, trials, prob, params):            
-            trials = jnp.asarray(trials, dtype=logic.REAL)
-            prob = jnp.asarray(prob, dtype=logic.REAL)
-            sample = random.binomial(key=key, n=trials, p=prob, dtype=logic.REAL)
-            return sample, params     
+    # normal approximation to Binomial: Bin(n, p) -> Normal(np, np(1-p))
+    def _binomial_normal_approx(self, logic):
+        def _jax_wrapped_calc_binomial_normal_approx(key, trials, prob, params):
+            normal = random.normal(key=key, shape=jnp.shape(trials), dtype=logic.REAL)
+            mean = trials * prob
+            std = jnp.sqrt(trials * prob * (1.0 - prob))
+            sample = mean + std * normal
+            return sample, params    
+        return _jax_wrapped_calc_binomial_normal_approx
         
-        # Binomial(n, p) = sum_{i = 1 ... n} Bernoulli(p)
-        bernoulli_approx = self.bernoulli(id, init_params, logic)
-        def _jax_wrapped_calc_binomial_sum(key, trials, prob, params):
-            prob_full = jnp.broadcast_to(
-                prob[..., jnp.newaxis], shape=jnp.shape(prob) + (self.binomial_bins,))
-            sample_bern, params = bernoulli_approx(key, prob_full, params)
-            indices = jnp.arange(self.binomial_bins)[
-                (jnp.newaxis,) * jnp.ndim(prob) + (...,)]
-            mask = indices < trials[..., jnp.newaxis]
-            sample = jnp.sum(sample_bern * mask, axis=-1)
-            return sample, params
+    def _binomial_gumbel_softmax(self, id, init_params, logic):
+        argmax_approx = logic.argmax(id, init_params)
+        def _jax_wrapped_calc_binomial_gumbel_softmax(key, trials, prob, params):
+            ks = jnp.arange(self.binomial_bins)[(jnp.newaxis,) * jnp.ndim(trials) + (...,)]
+            trials = trials[..., jnp.newaxis]
+            prob = prob[..., jnp.newaxis]
+            in_support = ks <= trials
+            ks = jnp.minimum(ks, trials)
+            log_prob = ((scipy.special.gammaln(trials + 1) - 
+                         scipy.special.gammaln(ks + 1) - 
+                         scipy.special.gammaln(trials - ks + 1)) +
+                        ks * jnp.log(prob + logic.eps) + 
+                        (trials - ks) * jnp.log1p(-prob + logic.eps))
+            log_prob = jnp.where(in_support, log_prob, jnp.log(logic.eps))
+            Gumbel01 = random.gumbel(key=key, shape=jnp.shape(log_prob), dtype=logic.REAL)
+            sample = Gumbel01 + log_prob
+            return argmax_approx(sample, axis=-1, params=params)
+        return _jax_wrapped_calc_binomial_gumbel_softmax
+        
+    def binomial(self, id, init_params, logic):
+        _jax_wrapped_calc_binomial_normal = self._binomial_normal_approx(logic)  
+        _jax_wrapped_calc_binomial_gs = self._binomial_gumbel_softmax(id, init_params, logic)
         
-        # for trials not too large use the Bernoulli relaxation, use exact otherwise
+        # for small trials use the Bernoulli relaxation
+        # for large trials use the normal approximation
         def _jax_wrapped_calc_binomial_approx(key, trials, prob, params):
-            return jax.lax.cond(
-                jax.lax.stop_gradient(jnp.max(trials) < self.binomial_bins), 
-                _jax_wrapped_calc_binomial_sum, 
-                _jax_wrapped_calc_binomial_exact, 
-                key, trials, prob, params
-            )
+            small_trials = jax.lax.stop_gradient(trials < self.binomial_bins)
+            small_sample, params = _jax_wrapped_calc_binomial_gs(key, trials, prob, params)
+            large_sample, params = _jax_wrapped_calc_binomial_normal(key, trials, prob, params)
+            sample = jnp.where(small_trials, small_sample, large_sample)
+            return sample, params
         return _jax_wrapped_calc_binomial_approx
     
+    # https://en.wikipedia.org/wiki/Negative_binomial_distribution#Gamma%E2%80%93Poisson_mixture
+    def negative_binomial(self, id, init_params, logic):
+        poisson_approx = self.poisson(id, init_params, logic)
+        def _jax_wrapped_calc_negative_binomial_approx(key, trials, prob, params):
+            key, subkey = random.split(key)
+            trials = jnp.asarray(trials, dtype=logic.REAL)
+            Gamma = random.gamma(key=key, a=trials, dtype=logic.REAL)
+            scale = (1.0 - prob) / prob
+            poisson_rate = scale * Gamma
+            return poisson_approx(subkey, poisson_rate, params)
+        return _jax_wrapped_calc_negative_binomial_approx
+
     def geometric(self, id, init_params, logic):
         approx_floor = logic.floor(id, init_params)
         def _jax_wrapped_calc_geometric_approx(key, prob, params):
             U = random.uniform(key=key, shape=jnp.shape(prob), dtype=logic.REAL)
-            floor, params = approx_floor(jnp.log1p(-U) / jnp.log1p(-prob), params)
+            floor, params = approx_floor(
+                jnp.log1p(-U) / jnp.log1p(-prob + logic.eps), params)
             sample = floor + 1
             return sample, params
         return _jax_wrapped_calc_geometric_approx
@@ -532,6 +587,14 @@ class Determinization(RandomSampling):
     def binomial(self, id, init_params, logic):
         return self._jax_wrapped_calc_binomial_determinized
     
+    @staticmethod
+    def _jax_wrapped_calc_negative_binomial_determinized(key, trials, prob, params):
+        sample = trials * ((1.0 / prob) - 1.0)
+        return sample, params
+
+    def negative_binomial(self, id, init_params, logic):
+        return self._jax_wrapped_calc_negative_binomial_determinized
+    
     @staticmethod
     def _jax_wrapped_calc_geometric_determinized(key, prob, params):
         sample = 1.0 / prob
@@ -558,21 +621,23 @@ class Determinization(RandomSampling):
 #
 # ===========================================================================
 
-class ControlFlow:
+class ControlFlow(metaclass=ABCMeta):
     '''A base class for control flow, including if and switch statements.'''
     
+    @abstractmethod
     def if_then_else(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def switch(self, id, init_params):
-        raise NotImplementedError
+        pass
 
 
 class SoftControlFlow(ControlFlow):
     '''Soft control flow using a probabilistic interpretation.'''
     
     def __init__(self, weight: float=10.0) -> None:
-        self.weight = weight
+        self.weight = float(weight)
         
     @staticmethod
     def _jax_wrapped_calc_if_then_else_soft(c, a, b, params):
@@ -606,15 +671,15 @@ class SoftControlFlow(ControlFlow):
 # ===========================================================================
 
 
-class Logic:
+class Logic(metaclass=ABCMeta):
     '''A base class for representing logic computations in JAX.'''
     
     def __init__(self, use64bit: bool=False) -> None:
         self.set_use64bit(use64bit)
     
-    def summarize_hyperparameters(self) -> None:
-        print(f'model relaxation:\n'
-              f'    use_64_bit    ={self.use64bit}')
+    def summarize_hyperparameters(self) -> str:
+        return (f'model relaxation:\n'
+                f'    use_64_bit    ={self.use64bit}')
     
     def set_use64bit(self, use64bit: bool) -> None:
         '''Toggles whether or not the JAX system will use 64 bit precision.'''
@@ -712,123 +777,158 @@ class Logic:
                 'Discrete': self.discrete,
                 'Poisson': self.poisson,
                 'Geometric': self.geometric,
-                'Binomial': self.binomial
+                'Binomial': self.binomial,
+                'NegativeBinomial': self.negative_binomial
             }
         }
 
     # ===========================================================================
     # logical operators
     # ===========================================================================
-     
+    
+    @abstractmethod
     def logical_and(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def logical_not(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def logical_or(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def xor(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def implies(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def equiv(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def forall(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def exists(self, id, init_params):    
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # comparison operators
     # ===========================================================================
     
+    @abstractmethod
     def greater_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def greater(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def less_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def less(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def not_equal(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # special functions
     # ===========================================================================
      
+    @abstractmethod
     def sgn(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def floor(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def round(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def ceil(self, id, init_params):
-        raise NotImplementedError 
+        pass 
     
+    @abstractmethod
     def div(self, id, init_params):
-        raise NotImplementedError 
+        pass 
     
+    @abstractmethod
     def mod(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
+    @abstractmethod
     def sqrt(self, id, init_params):
-        raise NotImplementedError 
+        pass
     
     # ===========================================================================
     # indexing
     # ===========================================================================
-     
+    
+    @abstractmethod
     def argmax(self, id, init_params):   
-        raise NotImplementedError  
+        pass
     
+    @abstractmethod
     def argmin(self, id, init_params):   
-        raise NotImplementedError     
+        pass
     
     # ===========================================================================
     # control flow
     # ===========================================================================
      
+    @abstractmethod
     def control_if(self, id, init_params):
-        raise NotImplementedError
+        pass
         
+    @abstractmethod
     def control_switch(self, id, init_params):
-        raise NotImplementedError
+        pass
     
     # ===========================================================================
     # random variables
     # ===========================================================================
      
+    @abstractmethod
     def discrete(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def bernoulli(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def poisson(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def geometric(self, id, init_params):
-        raise NotImplementedError
+        pass
     
+    @abstractmethod
     def binomial(self, id, init_params):
-        raise NotImplementedError
+        pass
+
+    @abstractmethod
+    def negative_binomial(self, id, init_params):
+        pass
 
 
 class ExactLogic(Logic):
@@ -1005,6 +1105,17 @@ class ExactLogic(Logic):
             sample = jnp.asarray(sample, dtype=self.INT)
             return sample, params
         return _jax_wrapped_calc_binomial_exact
+    
+    # note: for some reason tfp defines it as number of successes before trials failures
+    # I will define it as the number of failures before trials successes
+    def negative_binomial(self, id, init_params):
+        def _jax_wrapped_calc_negative_binomial_exact(key, trials, prob, params):
+            trials = jnp.asarray(trials, dtype=self.REAL)
+            prob = jnp.asarray(prob, dtype=self.REAL)
+            dist = tfp.distributions.NegativeBinomial(total_count=trials, probs=1.0 - prob)
+            sample = jnp.asarray(dist.sample(seed=key), dtype=self.INT)
+            return sample, params
+        return _jax_wrapped_calc_negative_binomial_exact
 
     
 class FuzzyLogic(Logic):
@@ -1049,8 +1160,8 @@ class FuzzyLogic(Logic):
                 f'    underflow_tol={self.eps}\n'
                 f'    use_64_bit   ={self.use64bit}\n')
 
-    def summarize_hyperparameters(self) -> None:
-        print(self.__str__())
+    def summarize_hyperparameters(self) -> str:
+        return self.__str__()
         
     # ===========================================================================
     # logical operators
@@ -1234,6 +1345,9 @@ class FuzzyLogic(Logic):
     
     def binomial(self, id, init_params):
         return self.sampling.binomial(id, init_params, self)
+    
+    def negative_binomial(self, id, init_params):
+        return self.sampling.negative_binomial(id, init_params, self)
 
 
 # ===========================================================================