4
4
from . import models
5
5
from pymc3 .step_methods .hmc .base_hmc import BaseHMC
6
6
import pymc3
7
+ from pymc3 .theanof import floatX
7
8
from .checks import close_to
9
+ from .helpers import select_by_precision
8
10
9
11
10
12
def test_leapfrog_reversible ():
@@ -13,17 +15,16 @@ def test_leapfrog_reversible():
13
15
step = BaseHMC (vars = model .vars , model = model )
14
16
bij = DictToArrayBijection (step .ordering , start )
15
17
q0 = bij .map (start )
16
- p0 = np .ones (n ) * .05
17
-
18
+ p0 = floatX ( np .ones (n ) * .05 )
19
+ precision = select_by_precision ( float64 = 1E-8 , float32 = 1E-5 )
18
20
for epsilon in [.01 , .1 , 1.2 ]:
19
21
for n_steps in [1 , 2 , 3 , 4 , 20 ]:
20
22
21
23
q , p = q0 , p0
22
- q , p , _ = step .leapfrog (q , p , np .array (epsilon ), np .array (n_steps , dtype = 'int32' ))
23
- q , p , _ = step .leapfrog (q , - p , np .array (epsilon ), np .array (n_steps , dtype = 'int32' ))
24
-
25
- close_to (q , q0 , 1e-8 , str ((n_steps , epsilon )))
26
- close_to (- p , p0 , 1e-8 , str ((n_steps , epsilon )))
24
+ q , p , _ = step .leapfrog (q , p , floatX (np .array (epsilon )), np .array (n_steps , dtype = 'int32' ))
25
+ q , p , _ = step .leapfrog (q , - p , floatX (np .array (epsilon )), np .array (n_steps , dtype = 'int32' ))
26
+ close_to (q , q0 , precision , str ((n_steps , epsilon )))
27
+ close_to (- p , p0 , precision , str ((n_steps , epsilon )))
27
28
28
29
29
30
def test_leapfrog_reversible_single ():
@@ -36,7 +37,8 @@ def test_leapfrog_reversible_single():
36
37
for method , step in zip (integrators , steps ):
37
38
bij = DictToArrayBijection (step .ordering , start )
38
39
q0 = bij .map (start )
39
- p0 = np .ones (n ) * .05
40
+ p0 = floatX (np .ones (n ) * .05 )
41
+ precision = select_by_precision (float64 = 1E-8 , float32 = 1E-5 )
40
42
for epsilon in [0.01 , 0.1 , 1.2 ]:
41
43
for n_steps in [1 , 2 , 3 , 4 , 20 ]:
42
44
dlogp0 = step .dlogp (q0 )
@@ -46,13 +48,13 @@ def test_leapfrog_reversible_single():
46
48
47
49
energy = step .compute_energy (q , p )
48
50
for _ in range (n_steps ):
49
- q , p , v , dlogp , _ = step .leapfrog (q , p , dlogp , np .array (epsilon ))
51
+ q , p , v , dlogp , _ = step .leapfrog (q , p , dlogp , floatX ( np .array (epsilon ) ))
50
52
p = - p
51
53
for _ in range (n_steps ):
52
- q , p , v , dlogp , _ = step .leapfrog (q , p , dlogp , np .array (epsilon ))
54
+ q , p , v , dlogp , _ = step .leapfrog (q , p , dlogp , floatX ( np .array (epsilon ) ))
53
55
54
- close_to (q , q0 , 1e-8 , str (('q' , method , n_steps , epsilon )))
55
- close_to (- p , p0 , 1e-8 , str (('p' , method , n_steps , epsilon )))
56
+ close_to (q , q0 , precision , str (('q' , method , n_steps , epsilon )))
57
+ close_to (- p , p0 , precision , str (('p' , method , n_steps , epsilon )))
56
58
57
59
58
60
def test_nuts_tuning ():
0 commit comments