1import unittest
2
3import numpy as np
4
5from pystan import StanModel, stan
6from pystan._compat import PY2
7from pystan.tests.helper import get_model
8
9# REF: rstan/tests/unitTests/runit.test.stanfit.R
10
11class TestStanfit(unittest.TestCase):
12
13    def test_init_zero_exception_inf_grad(self):
14        code = """
15        parameters {
16            real x;
17        }
18        model {
19            target += 1 / log(x);
20        }
21        """
22        sm = StanModel(model_code=code)
23        with self.assertRaises(RuntimeError):
24            sm.sampling(init='0', iter=1, chains=1)
25
26    def test_grad_log(self):
27        y = np.array([0.70,  -0.16,  0.77, -1.37, -1.99,  1.35, 0.08,
28                      0.02,  -1.48, -0.08,  0.34,  0.03, -0.42, 0.87,
29                      -1.36,  1.43,  0.80, -0.48, -1.61, -1.27])
30
31        code = '''
32        data {
33            int N;
34            real y[N];
35        }
36        parameters {
37            real mu;
38            real<lower=0> sigma;
39        }
40        model {
41            y ~ normal(mu, sigma);
42        }'''
43
44        def log_prob_fun(mu, log_sigma, adjust=True):
45            sigma = np.exp(log_sigma)
46            lp = -1 * np.sum((y - mu)**2) / (2 * (sigma**2)) - len(y) * np.log(sigma)
47            if adjust:
48                lp = lp + np.log(sigma)
49            return lp
50
51        def log_prob_grad_fun(mu, log_sigma, adjust=True):
52            sigma = np.exp(log_sigma)
53            g_lsigma = np.sum((y - mu)**2) * sigma**(-2) - len(y)
54            if adjust:
55                g_lsigma = g_lsigma + 1
56            g_mu = np.sum(y - mu) * sigma**(-2)
57            return (g_mu, g_lsigma)
58
59        sm = get_model("normal_mu_sigma_model", code)
60        sf = sm.sampling(data=dict(y=y, N=20), iter=200)
61        mu = 0.1
62        sigma = 2
63        self.assertEqual(sf.log_prob(sf.unconstrain_pars(dict(mu=mu, sigma=sigma))),
64                         log_prob_fun(mu, np.log(sigma)))
65        self.assertEqual(sf.log_prob(sf.unconstrain_pars(dict(mu=mu, sigma=sigma)), False),
66                         log_prob_fun(mu, np.log(sigma), adjust=False))
67        g1 = sf.grad_log_prob(sf.unconstrain_pars(dict(mu=mu, sigma=sigma)), False)
68        np.testing.assert_allclose(g1, log_prob_grad_fun(mu, np.log(sigma), adjust=False))
69
70    def test_specify_args(self):
71        y = (0.70,  -0.16,  0.77, -1.37, -1.99,  1.35, 0.08,
72                0.02,  -1.48, -0.08,  0.34,  0.03, -0.42, 0.87,
73                -1.36,  1.43,  0.80, -0.48, -1.61, -1.27)
74        code = """
75            data {
76                int N;
77                real y[N];
78            }
79            parameters {
80                real mu;
81                real<lower=0> sigma;
82            }
83            model {
84                y ~ normal(mu, sigma);
85            }"""
86        stepsize0 = 0.15
87        sm = get_model("normal_mu_sigma_model", code)
88        sf = sm.sampling(data=dict(y=y, N=20), iter=200,
89                         control=dict(adapt_engaged=False, stepsize=stepsize0))
90        self.assertEqual(sf.get_sampler_params()[0]['stepsize__'][0], stepsize0)
91        sf2 = stan(fit=sf, iter=20, algorithm='HMC', data=dict(y=y, N=20),
92                    control=dict(adapt_engaged=False, stepsize=stepsize0))
93        self.assertEqual(sf2.get_sampler_params()[0]['stepsize__'][0], stepsize0)
94        sf3 = stan(fit=sf, iter=1, data=dict(y=y, N=20), init=0, chains=1)
95        i_u = sf3.unconstrain_pars(sf3.get_inits()[0])
96        np.testing.assert_equal(i_u, [0, 0])
97