Floordiv, modulus, consistent ordering (#15)

* floordiv

* sort variables in correlations for consistent results

* add Mod, i.e. the % operator

Author: tommyod

src/probabilit/modeling.py

@@ -3,7 +3,9 @@
 --------
 
 Probabilit lets the user perform Monte-Carlo sampling using a high-level
-modeling language, which creates a computational graph.
+modeling language.
+The modeling language creates a lazy computational graph.
+When a node is sampled, all ancestor nodes are sampled in turn and the samples are propagated down in the graph, from parent nodes to child nodes.
 
 For instance, to compute the shipping cost of a box where we are uncertain
 about the measurements:
@@ -20,21 +22,21 @@
 20.00139737515...
 
 Distributions are built on top of scipy, so "norm" refers to the name of the
-normal distribution as given in `scipy.stats`, and the arguments to the distribution
-must also match those given by `scipy.stats.norm`.
+normal distribution as given in `scipy.stats`, and the arguments to the
+distribution must also match those given by `scipy.stats.norm`.
 
-Here is another example showing composite distributions, where the argument
-to one distribution is another distribution:
+Here is another example demonstrating composite distributions, where an
+argument to one distribution is another distribution:
 
 >>> eggs_per_nest = Distribution("poisson", mu=3)
 >>> survivial_prob = Distribution("beta", a=10, b=15)
 >>> survived = Distribution("binom", n=eggs_per_nest, p=survivial_prob)
 >>> survived.sample(9, random_state=rng)
 array([0., 1., 2., 0., 3., 1., 1., 0., 2.])
 
-To understand and examine the modeling language, we can do some computations
-with constants. The computational graph carries out arithmetic operations lazily
-once a model is sampled. Mixing numbers with nodes is allowed, but at least one
+To understand and examine the modeling language, we can perform  computations
+using constants. The computational graph carries out arithmetic operations
+when the model is sampled. Mixing numbers with nodes is allowed, but at least one
 expression or term must be a probabilit class instance:
 
 >>> a = Constant(1)
@@ -62,7 +64,7 @@
 Multiply(Distribution("expon", scale=1), Constant(5))
 Add(Add(Power(Distribution("norm", loc=5, scale=1), Distribution("expon", scale=1)), Multiply(Distribution("norm", loc=5, scale=1), Distribution("expon", scale=1))), Multiply(Distribution("expon", scale=1), Constant(5)))
 
-Sampling the expression is simple enough:
+Sampling the expression is simple:
 
 >>> expression.sample(5, random_state=rng)
 array([ 2.70764145, 36.58578812,  7.07064239,  1.84433247,  3.90951632])
@@ -73,10 +75,10 @@
 >>> a.samples_
 array([4.51589278, 4.37788659, 5.25960812, 5.80609507, 4.33770499])
 
-To sample using e.g. Latin Hypercube, do the following:
+To sample using e.g. the Latin Hypercube algorithm, do the following:
 
 >>> from scipy.stats.qmc import LatinHypercube
->>> d = expression.get_number_of_distribution_nodes()
+>>> d = expression.num_distribution_nodes()
 >>> hypercube = LatinHypercube(d=d, rng=rng)
 >>> hypercube_samples = hypercube.random(5) # Draw 5 samples
 >>> expression.sample_from_cube(hypercube_samples)
@@ -91,6 +93,19 @@
 >>> expression.sample(5, random_state=rng)
 array([ 4.70542018, 14.43250192,  6.74494838, -0.14020459, -3.27334554])
 
+Nodes are hashable and can be used in sets, so __hash__ and __eq__ must both
+be defined. We cannot use `==` for modeling since equality in that context has
+another meaning. Use the Equal node instead. This is only relevant in cases
+when equality is part of a model. For real-valued distribution (e.g. Normal)
+equality does not make sense since the probability that two floats are equal
+is zero.
+
+>>> dice1 = Distribution("uniform", loc=1, scale=6) // 1
+>>> dice2 = Distribution("uniform", loc=1, scale=6) // 1
+>>> equal_result = Equal(dice1, dice2)
+>>> float(equal_result.sample(999, random_state=42).mean())
+0.166...
+
 Functions
 ---------
 
@@ -328,7 +343,7 @@ def nodes(self):
             yield (node := queue.pop())
             queue.extend(node.get_parents())
 
-    def get_number_of_distribution_nodes(self):
+    def num_distribution_nodes(self):
         return sum(1 for node in set(self.nodes()) if isinstance(node, Distribution))
 
     def sample(self, size=None, random_state=None):
@@ -337,7 +352,7 @@ def sample(self, size=None, random_state=None):
         random_state = check_random_state(random_state)
 
         # Draw a cube of random variables in [0, 1]
-        cube = random_state.random((size, self.get_number_of_distribution_nodes()))
+        cube = random_state.random((size, self.num_distribution_nodes()))
 
         return self.sample_from_cube(cube)
 
@@ -346,7 +361,7 @@ def sample_from_cube(self, cube):
         distributions. The cube must have shape (dimensionality, num_samples)."""
         assert nx.is_directed_acyclic_graph(self.to_graph())
         size, n_dim = cube.shape
-        assert n_dim == self.get_number_of_distribution_nodes()
+        assert n_dim == self.num_distribution_nodes()
 
         # Prepare columns of quantiles, one column for each Distribution
         columns = iter(list(cube.T))
@@ -377,6 +392,7 @@ def sample_from_cube(self, cube):
             node.samples_ = node._sample(q=next(columns))
 
         # Go through all ancestor nodes and create a list [(var, corr), ...]
+        # that contains all correlations we must induce
         correlations = []
         for node in set(self.nodes()):
             if hasattr(node, "_correlations"):
@@ -397,6 +413,8 @@ def sample_from_cube(self, cube):
 
         # Map all variables to integers
         all_variables = list(functools.reduce(set.union, variable_sets, set()))
+        # Ensure consistent ordering for reproducible results
+        all_variables = sorted(all_variables, key=lambda n: n._id)
         var_to_int = {v: i for (i, v) in enumerate(all_variables)}
         correlations = [
             (tuple(var_to_int[var] for var in variables), corrmat)
@@ -508,12 +526,24 @@ def __mul__(self, other):
     def __rmul__(self, other):
         return Multiply(self, other)
 
+    def __floordiv__(self, other):
+        return FloorDivide(self, other)
+
+    def __rfloordiv__(self, other):
+        return FloorDivide(other, self)
+
     def __truediv__(self, other):
         return Divide(self, other)
 
     def __rtruediv__(self, other):
         return Divide(other, self)
 
+    def __mod__(self, other):
+        return Mod(self, other)
+
+    def __rmod__(self, other):
+        return Mod(other, self)
+
     def __sub__(self, other):
         return Subtract(self, other)
 
@@ -691,6 +721,14 @@ def get_parents(self):
         yield from self.parents
 
 
+class FloorDivide(BinaryTransform):
+    op = np.floor_divide
+
+
+class Mod(BinaryTransform):
+    op = np.mod
+
+
 class Divide(BinaryTransform):
     op = operator.truediv