Commit 0ebedb02 authored by Junpeng Lao's avatar Junpeng Lao

move function to pymc3.math

parent 196d1fad
......@@ -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_matrix_op(krons, m, op):
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 @@ kron_dot = partial(kron_matrix_op, op=tt.dot)
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,9 +176,7 @@ 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):
......@@ -146,9 +184,20 @@ def logdiffexp(a, 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 @@ class LogDet(Op):
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 @@ class LogDet(Op):
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 @@ class BatchedDiag(tt.Op):
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,26 +373,30 @@ 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])
......@@ -342,9 +404,7 @@ class BlockDiagonalMatrix(Op):
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)
......@@ -352,9 +412,13 @@ class BlockDiagonalMatrix(Op):
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):
......@@ -362,7 +426,7 @@ class BlockDiagonalMatrix(Op):
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
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment