implement forward_val in transform as a numpy function (#2920)

* implement `forward_val` in transform as a numpy function

it is faster to do forward transformation in numpy, could be useful for sampling from prior.

* add test

* change transform useage in SMC

* rename test

Author: Junpeng Lao

pymc3/distributions/transforms.py

@@ -7,6 +7,7 @@
 from ..math import logit, invlogit
 from .distribution import draw_values
 import numpy as np
+from scipy.special import logit as nplogit
 
 __all__ = ['transform', 'stick_breaking', 'logodds', 'interval', 'log_exp_m1',
            'lowerbound', 'upperbound', 'log', 'sum_to_1', 't_stick_breaking']
@@ -97,7 +98,7 @@ def forward(self, x):
         return tt.log(x)
 
     def forward_val(self, x, point=None):
-        return self.forward(x)
+        return np.log(x)
 
     def jacobian_det(self, x):
         return x
@@ -119,7 +120,7 @@ def forward(self, x):
         return tt.log(1.-tt.exp(-x)) + x
 
     def forward_val(self, x, point=None):
-        return self.forward(x)
+        return np.log(1.-np.exp(-x)) + x
 
     def jacobian_det(self, x):
         return -tt.nnet.softplus(-x)
@@ -137,7 +138,7 @@ def forward(self, x):
         return logit(x)
 
     def forward_val(self, x, point=None):
-        return self.forward(x)
+        return nplogit(x)
 
 logodds = LogOdds()
 
@@ -166,7 +167,7 @@ def forward_val(self, x, point=None):
         # For an explanation see pull/2328#issuecomment-309303811
         a, b = draw_values([self.a-0., self.b-0.],
                             point=point)
-        return floatX(tt.log(x - a) - tt.log(b - x))
+        return floatX(np.log(x - a) - np.log(b - x))
 
     def jacobian_det(self, x):
         s = tt.nnet.softplus(-x)
@@ -198,7 +199,7 @@ def forward_val(self, x, point=None):
         # For an explanation see pull/2328#issuecomment-309303811
         a = draw_values([self.a-0.],
                         point=point)[0]
-        return floatX(tt.log(x - a))
+        return floatX(np.log(x - a))
 
     def jacobian_det(self, x):
         return x
@@ -229,7 +230,7 @@ def forward_val(self, x, point=None):
         # For an explanation see pull/2328#issuecomment-309303811
         b = draw_values([self.b-0.],
                         point=point)[0]
-        return floatX(tt.log(b - x))
+        return floatX(np.log(b - x))
 
     def jacobian_det(self, x):
         return x
@@ -249,7 +250,7 @@ def forward(self, x):
         return x[:-1]
 
     def forward_val(self, x, point=None):
-        return self.forward(x)
+        return x[:-1]
 
     def jacobian_det(self, x):
         return 0
@@ -284,8 +285,17 @@ def forward(self, x_):
         y = logit(z) - eq_share
         return floatX(y.T)
 
-    def forward_val(self, x, point=None):
-        return self.forward(x)
+    def forward_val(self, x_, point=None):
+        x = x_.T
+        # reverse cumsum
+        x0 = x[:-1]
+        s = np.cumsum(x0[::-1], 0)[::-1] + x[-1]
+        z = x0 / s
+        Km1 = x.shape[0] - 1
+        k = np.arange(Km1)[(slice(None),) + (None,) * (x.ndim - 1)]
+        eq_share = nplogit(1. / (Km1 + 1 - k).astype(str(x_.dtype)))
+        y = nplogit(z) - eq_share
+        return floatX(y.T)
 
     def backward(self, y_):
         y = y_.T
@@ -326,7 +336,7 @@ def forward(self, x):
         return tt.as_tensor_variable(x)
 
     def forward_val(self, x, point=None):
-        return self.forward(x)
+        return x
 
     def jacobian_det(self, x):
         return 0
@@ -346,8 +356,9 @@ def backward(self, x):
     def forward(self, y):
         return tt.advanced_set_subtensor1(y, tt.log(y[self.diag_idxs]), self.diag_idxs)
 
-    def forward_val(self, x, point=None):
-        return self.forward(x)
+    def forward_val(self, y, point=None):
+        y[self.diag_idxs] = np.log(y[self.diag_idxs])
+        return y
 
     def jacobian_det(self, y):
         return tt.sum(y[self.diag_idxs])

pymc3/step_methods/smc.py

@@ -186,9 +186,9 @@ def __init__(self, vars=None, out_vars=None, samples=1000, n_chains=100, n_steps
         init_rnd = {}
         for v in vars:
             if pm.util.is_transformed_name(v.name):
-                trans = v.distribution.transform_used.forward
-                rnd = trans(v.distribution.dist.random(size=self.n_chains, point=start))
-                init_rnd[v.name] = rnd.eval()
+                trans = v.distribution.transform_used.forward_val
+                init_rnd[v.name] = trans(v.distribution.dist.random(
+                    size=self.n_chains, point=start))
             else:
                 init_rnd[v.name] = v.random(size=self.n_chains, point=start)
 

pymc3/tests/test_transforms.py

@@ -13,17 +13,21 @@
 tol = 1e-7 if theano.config.floatX == 'flaot64' else 1e-6
 
 
-def check_transform_identity(transform, domain, constructor=tt.dscalar, test=0):
+def check_transform(transform, domain, constructor=tt.dscalar, test=0):
     x = constructor('x')
     x.tag.test_value = test
+    # test forward and forward_val
+    forward_f = theano.function([x], transform.forward(x))
+    # test transform identity
     identity_f = theano.function([x], transform.backward(transform.forward(x)))
 
     for val in domain.vals:
         close_to(val, identity_f(val), tol)
+        close_to(transform.forward_val(val), forward_f(val), tol)
 
 
-def check_vector_transform_identity(transform, domain):
-    return check_transform_identity(transform, domain, tt.dvector, test=np.array([0, 0]))
+def check_vector_transform(transform, domain):
+    return check_transform(transform, domain, tt.dvector, test=np.array([0, 0]))
 
 
 def get_values(transform, domain=R, constructor=tt.dscalar, test=0):
@@ -35,9 +39,9 @@ def get_values(transform, domain=R, constructor=tt.dscalar, test=0):
 
 
 def test_simplex():
-    check_vector_transform_identity(tr.stick_breaking, Simplex(2))
-    check_vector_transform_identity(tr.stick_breaking, Simplex(4))
-    check_transform_identity(tr.stick_breaking, MultiSimplex(
+    check_vector_transform(tr.stick_breaking, Simplex(2))
+    check_vector_transform(tr.stick_breaking, Simplex(4))
+    check_transform(tr.stick_breaking, MultiSimplex(
         3, 2), constructor=tt.dmatrix, test=np.zeros((2, 2)))
 
 
@@ -56,8 +60,8 @@ def test_simplex_jacobian_det():
 
 
 def test_sum_to_1():
-    check_vector_transform_identity(tr.sum_to_1, Simplex(2))
-    check_vector_transform_identity(tr.sum_to_1, Simplex(4))
+    check_vector_transform(tr.sum_to_1, Simplex(2))
+    check_vector_transform(tr.sum_to_1, Simplex(4))
 
 
 def test_sum_to_1_jacobian_det():
@@ -95,7 +99,7 @@ def check_jacobian_det(transform, domain,
 
 
 def test_log():
-    check_transform_identity(tr.log, Rplusbig)
+    check_transform(tr.log, Rplusbig)
     check_jacobian_det(tr.log, Rplusbig, elemwise=True)
     check_jacobian_det(tr.log, Vector(Rplusbig, 2),
                        tt.dvector, [0, 0], elemwise=True)
@@ -105,7 +109,7 @@ def test_log():
 
 
 def test_log_exp_m1():
-    check_transform_identity(tr.log_exp_m1, Rplusbig)
+    check_transform(tr.log_exp_m1, Rplusbig)
     check_jacobian_det(tr.log_exp_m1, Rplusbig, elemwise=True)
     check_jacobian_det(tr.log_exp_m1, Vector(Rplusbig, 2),
                        tt.dvector, [0, 0], elemwise=True)
@@ -115,7 +119,7 @@ def test_log_exp_m1():
 
 
 def test_logodds():
-    check_transform_identity(tr.logodds, Unit)
+    check_transform(tr.logodds, Unit)
     check_jacobian_det(tr.logodds, Unit, elemwise=True)
     check_jacobian_det(tr.logodds, Vector(Unit, 2),
                        tt.dvector, [.5, .5], elemwise=True)
@@ -127,7 +131,7 @@ def test_logodds():
 
 def test_lowerbound():
     trans = tr.lowerbound(0.0)
-    check_transform_identity(trans, Rplusbig)
+    check_transform(trans, Rplusbig)
     check_jacobian_det(trans, Rplusbig, elemwise=True)
     check_jacobian_det(trans, Vector(Rplusbig, 2),
                        tt.dvector, [0, 0], elemwise=True)
@@ -138,7 +142,7 @@ def test_lowerbound():
 
 def test_upperbound():
     trans = tr.upperbound(0.0)
-    check_transform_identity(trans, Rminusbig)
+    check_transform(trans, Rminusbig)
     check_jacobian_det(trans, Rminusbig, elemwise=True)
     check_jacobian_det(trans, Vector(Rminusbig, 2),
                        tt.dvector, [-1, -1], elemwise=True)
@@ -151,7 +155,7 @@ def test_interval():
     for a, b in [(-4, 5.5), (.1, .7), (-10, 4.3)]:
         domain = Unit * np.float64(b - a) + np.float64(a)
         trans = tr.interval(a, b)
-        check_transform_identity(trans, domain)
+        check_transform(trans, domain)
         check_jacobian_det(trans, domain, elemwise=True)
 
         vals = get_values(trans)
@@ -161,7 +165,7 @@ def test_interval():
 
 def test_circular():
     trans = tr.circular
-    check_transform_identity(trans, Circ)
+    check_transform(trans, Circ)
     check_jacobian_det(trans, Circ)
 
     vals = get_values(trans)

pymc3/util.py

@@ -143,7 +143,7 @@ def update_start_vals(a, b, model):
                         d.transformation for d in model.deterministics if d.name == name]
                     if transform_func:
                         b[tname] = transform_func[0].forward_val(
-                            a[name], point=b).eval()
+                            a[name], point=b)
 
     a.update({k: v for k, v in b.items() if k not in a})