move function to pymc3.math

Author: junpenglao

pymc3/math.py

@@ -14,13 +14,51 @@
 
 import sys
 import theano.tensor as tt
+
 # pylint: disable=unused-import
 import theano
 from theano.tensor import (
-    constant, flatten, zeros_like, ones_like, stack, concatenate, sum, prod,
-    lt, gt, le, ge, eq, neq, switch, clip, where, and_, or_, abs_, exp, log,
-    cos, sin, tan, cosh, sinh, tanh, sqr, sqrt, erf, erfc, erfinv, erfcinv, dot,
-    maximum, minimum, sgn, ceil, floor)
+    constant,
+    flatten,
+    zeros_like,
+    ones_like,
+    stack,
+    concatenate,
+    sum,
+    prod,
+    lt,
+    gt,
+    le,
+    ge,
+    eq,
+    neq,
+    switch,
+    clip,
+    where,
+    and_,
+    or_,
+    abs_,
+    exp,
+    log,
+    cos,
+    sin,
+    tan,
+    cosh,
+    sinh,
+    tanh,
+    sqr,
+    sqrt,
+    erf,
+    erfc,
+    erfinv,
+    erfcinv,
+    dot,
+    maximum,
+    minimum,
+    sgn,
+    ceil,
+    floor,
+)
 from theano.tensor.nlinalg import det, matrix_inverse, extract_diag, matrix_dot, trace
 import theano.tensor.slinalg
 import theano.sparse
@@ -62,7 +100,7 @@ def cartesian(*arrays):
             1D arrays where earlier arrays loop more slowly than later ones
     """
     N = len(arrays)
-    return np.stack(np.meshgrid(*arrays, indexing='ij'), -1).reshape(-1, N)
+    return np.stack(np.meshgrid(*arrays, indexing="ij"), -1).reshape(-1, N)
 
 
 def kron_matrix_op(krons, m, op):
@@ -81,6 +119,7 @@ def kron_matrix_op(krons, m, op):
     -------
     numpy array
     """
+
     def flat_matrix_op(flat_mat, mat):
         Nmat = mat.shape[1]
         flat_shape = flat_mat.shape
@@ -93,7 +132,7 @@ def kron_vector_op(v):
     if m.ndim == 1:
         m = m[:, None]  # Treat 1D array as Nx1 matrix
     if m.ndim != 2:  # Has not been tested otherwise
-        raise ValueError('m must have ndim <= 2, not {}'.format(m.ndim))
+        raise ValueError("m must have ndim <= 2, not {}".format(m.ndim))
     res = kron_vector_op(m)
     res_shape = res.shape
     return tt.reshape(res, (res_shape[1], res_shape[0])).T
@@ -104,6 +143,7 @@ def kron_vector_op(v):
 kron_solve_lower = partial(kron_matrix_op, op=tt.slinalg.solve_lower_triangular)
 kron_solve_upper = partial(kron_matrix_op, op=tt.slinalg.solve_upper_triangular)
 
+
 def flat_outer(a, b):
     return tt.outer(a, b).ravel()
 
@@ -124,7 +164,7 @@ def tround(*args, **kwargs):
     Temporary function to silence round warning in Theano. Please remove
     when the warning disappears.
     """
-    kwargs['mode'] = 'half_to_even'
+    kwargs["mode"] = "half_to_even"
     return tt.round(*args, **kwargs)
 
 
@@ -136,19 +176,28 @@ def logsumexp(x, axis=None):
 
 def logaddexp(a, b):
     diff = b - a
-    return tt.switch(diff > 0,
-                     b + tt.log1p(tt.exp(-diff)),
-                     a + tt.log1p(tt.exp(diff)))
+    return tt.switch(diff > 0, b + tt.log1p(tt.exp(-diff)), a + tt.log1p(tt.exp(diff)))
 
 
 def logdiffexp(a, b):
     """log(exp(a) - exp(b))"""
     return a + log1mexp(a - b)
 
 
+def logdiffexp_numpy(a, b):
+    """log(exp(a) - exp(b))"""
+    return a + log1mexp_numpy(a - b)
+
+
 def invlogit(x, eps=sys.float_info.epsilon):
     """The inverse of the logit function, 1 / (1 + exp(-x))."""
-    return (1. - 2. * eps) / (1. + tt.exp(-x)) + eps
+    return (1.0 - 2.0 * eps) / (1.0 + tt.exp(-x)) + eps
+
+
+def logbern(log_p):
+    if np.isnan(log_p):
+        raise FloatingPointError("log_p can't be nan.")
+    return np.log(nr.uniform()) < log_p
 
 
 def logit(p):
@@ -171,10 +220,16 @@ def log1mexp(x):
     For details, see
     https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
     """
-    return tt.switch(
-        tt.lt(x, 0.683),
-        tt.log(-tt.expm1(-x)),
-        tt.log1p(-tt.exp(-x)))
+    return tt.switch(tt.lt(x, 0.683), tt.log(-tt.expm1(-x)), tt.log1p(-tt.exp(-x)))
+
+
+def log1mexp_numpy(x):
+    """Return log(1 - exp(-x)).
+    This function is numerically more stable than the naive approach.
+    For details, see
+    https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf
+    """
+    return np.where(x < 0.683, np.log(-np.expm1(-x)), np.log1p(-np.exp(-x)))
 
 
 def flatten_list(tensors):
@@ -191,6 +246,7 @@ class LogDet(Op):
     Once PR #3959 (https://github.com/Theano/Theano/pull/3959/) by harpone is merged,
     this must be removed.
     """
+
     def make_node(self, x):
         x = theano.tensor.as_tensor_variable(x)
         o = theano.tensor.scalar(dtype=x.dtype)
@@ -204,7 +260,7 @@ def perform(self, node, inputs, outputs, params=None):
             log_det = np.sum(np.log(np.abs(s)))
             z[0] = np.asarray(log_det, dtype=x.dtype)
         except Exception:
-            print('Failed to compute logdet of {}.'.format(x))
+            print("Failed to compute logdet of {}.".format(x))
             raise
 
     def grad(self, inputs, g_outputs):
@@ -215,19 +271,20 @@ def grad(self, inputs, g_outputs):
     def __str__(self):
         return "LogDet"
 
+
 logdet = LogDet()
 
 
 def probit(p):
-    return -sqrt(2.) * erfcinv(2. * p)
+    return -sqrt(2.0) * erfcinv(2.0 * p)
 
 
 def invprobit(x):
-    return .5 * erfc(-x / sqrt(2.))
+    return 0.5 * erfc(-x / sqrt(2.0))
 
 
 def expand_packed_triangular(n, packed, lower=True, diagonal_only=False):
-    R"""Convert a packed triangular matrix into a two dimensional array.
+    r"""Convert a packed triangular matrix into a two dimensional array.
 
     Triangular matrices can be stored with better space efficiancy by
     storing the non-zero values in a one-dimensional array. We number
@@ -250,9 +307,9 @@ def expand_packed_triangular(n, packed, lower=True, diagonal_only=False):
         If true, return only the diagonal of the matrix.
     """
     if packed.ndim != 1:
-        raise ValueError('Packed triagular is not one dimensional.')
+        raise ValueError("Packed triagular is not one dimensional.")
     if not isinstance(n, int):
-        raise TypeError('n must be an integer')
+        raise TypeError("n must be an integer")
 
     if diagonal_only and lower:
         diag_idxs = np.arange(1, n + 1).cumsum() - 1
@@ -274,12 +331,13 @@ class BatchedDiag(tt.Op):
     """
     Fast BatchedDiag allocation
     """
+
     __props__ = ()
 
     def make_node(self, diag):
         diag = tt.as_tensor_variable(diag)
         if diag.type.ndim != 2:
-            raise TypeError('data argument must be a matrix', diag.type)
+            raise TypeError("data argument must be a matrix", diag.type)
 
         return tt.Apply(self, [diag], [tt.tensor3(dtype=diag.dtype)])
 
@@ -301,7 +359,7 @@ def grad(self, inputs, gout):
         return [gz[..., idx, idx]]
 
     def infer_shape(self, nodes, shapes):
-        return [(shapes[0][0], ) + (shapes[0][1],) * 2]
+        return [(shapes[0][0],) + (shapes[0][1],) * 2]
 
 
 def batched_diag(C):
@@ -315,54 +373,60 @@ def batched_diag(C):
         idx = tt.arange(dim)
         return C[..., idx, idx]
     else:
-        raise ValueError('Input should be 2 or 3 dimensional')
+        raise ValueError("Input should be 2 or 3 dimensional")
 
 
 class BlockDiagonalMatrix(Op):
-    __props__ = ('sparse', 'format')
+    __props__ = ("sparse", "format")
 
-    def __init__(self, sparse=False, format='csr'):
-        if format not in ('csr', 'csc'):
-            raise ValueError("format must be one of: 'csr', 'csc', got {}".format(format))
+    def __init__(self, sparse=False, format="csr"):
+        if format not in ("csr", "csc"):
+            raise ValueError(
+                "format must be one of: 'csr', 'csc', got {}".format(format)
+            )
         self.sparse = sparse
         self.format = format
 
     def make_node(self, *matrices):
         if not matrices:
-            raise ValueError('no matrices to allocate')
+            raise ValueError("no matrices to allocate")
         matrices = list(map(tt.as_tensor, matrices))
         if any(mat.type.ndim != 2 for mat in matrices):
-            raise TypeError('all data arguments must be matrices')
+            raise TypeError("all data arguments must be matrices")
         if self.sparse:
-            out_type = theano.sparse.matrix(self.format, dtype=largest_common_dtype(matrices))
+            out_type = theano.sparse.matrix(
+                self.format, dtype=largest_common_dtype(matrices)
+            )
         else:
             out_type = theano.tensor.matrix(dtype=largest_common_dtype(matrices))
         return tt.Apply(self, matrices, [out_type])
 
     def perform(self, node, inputs, output_storage, params=None):
         dtype = largest_common_dtype(inputs)
         if self.sparse:
-            output_storage[0][0] = sp.sparse.block_diag(
-                inputs, self.format, dtype
-            )
+            output_storage[0][0] = sp.sparse.block_diag(inputs, self.format, dtype)
         else:
             output_storage[0][0] = scipy_block_diag(*inputs).astype(dtype)
 
     def grad(self, inputs, gout):
         shapes = tt.stack([i.shape for i in inputs])
         index_end = shapes.cumsum(0)
         index_begin = index_end - shapes
-        slices = [ix_(tt.arange(index_begin[i, 0], index_end[i, 0]),
-                      tt.arange(index_begin[i, 1], index_end[i, 1])
-                      ) for i in range(len(inputs))]
+        slices = [
+            ix_(
+                tt.arange(index_begin[i, 0], index_end[i, 0]),
+                tt.arange(index_begin[i, 1], index_end[i, 1]),
+            )
+            for i in range(len(inputs))
+        ]
         return [gout[0][slc] for slc in slices]
 
     def infer_shape(self, nodes, shapes):
         first, second = zip(*shapes)
         return [(tt.add(*first), tt.add(*second))]
 
 
-def block_diagonal(matrices, sparse=False, format='csr'):
+def block_diagonal(matrices, sparse=False, format="csr"):
     r"""See scipy.sparse.block_diag or
     scipy.linalg.block_diag for reference