store MAP even when self._ci = True

convoys/regression.py

@@ -184,11 +184,11 @@ def fit(self, X, B, T, W=None, fix_k=None, fix_p=None):
             jac=autograd.grad(lambda x: -generalized_gamma_LL(x, *args)),
             method='SLSQP',
         )
-        x0 = res.x
+        result = {'map': res.x}
 
         # Let's sample from the posterior to compute uncertainties
         if self._ci:
-            dim, = x0.shape
+            dim, = res.x.shape
             nwalkers = 5*dim
             sampler = emcee.EnsembleSampler(
                 nwalkers=nwalkers,
@@ -197,42 +197,43 @@ def fit(self, X, B, T, W=None, fix_k=None, fix_p=None):
                 args=args
             )
             mcmc_initial_noise = 1e-3
-            p0 = [x0 + mcmc_initial_noise * numpy.random.randn(dim)
+            p0 = [result['map'] + mcmc_initial_noise * numpy.random.randn(dim)
                   for i in range(nwalkers)]
             nburnin = 20
             nsteps = numpy.ceil(1000. / nwalkers)
             print('\nStarting MCMC with %d walkers and %d steps:' % (
                     nwalkers, nburnin+nsteps))
             sampler.run_mcmc(p0, nburnin+nsteps)
             print('\n')
-            data = sampler.chain[:, nburnin:, :].reshape((-1, dim)).T
-        else:
-            data = x0
+            result['samples'] = sampler.chain[:, nburnin:, :] \
+                                       .reshape((-1, dim)).T
 
         # The `data` array is either 1D (for MAP) or 2D (for MCMC)
-        self.params = {
+        self.params = {k: {
             'k': exp(data[0]),
             'p': exp(data[1]),
             'a': data[4],
             'b': data[5],
             'alpha': data[6:6+n_features].T,
             'beta': data[6+n_features:6+2*n_features].T,
-            }
+        } for k, data in result.items()}
 
     def cdf(self, x, t, ci=None):
         x = numpy.array(x)
         t = numpy.array(t)
-        lambd = exp(dot(x, self.params['alpha'].T) + self.params['a'])
-        c = expit(dot(x, self.params['beta'].T) + self.params['b'])
+        if ci is None:
+            params = self.params['map']
+        else:
+            assert self._ci
+            params = self.params['samples']
+        lambd = exp(dot(x, params['alpha'].T) + params['a'])
+        c = expit(dot(x, params['beta'].T) + params['b'])
         M = c * gammainc(
-            self.params['k'],
-            numpy.multiply.outer(t, lambd)**self.params['p'])
+            params['k'],
+            numpy.multiply.outer(t, lambd)**params['p'])
 
-        if not self._ci:
-            assert not ci
+        if not ci:
             return M
-        elif ci is None:
-            return numpy.mean(M, axis=-1)
         else:
             # Replace the last axis with a 3-element vector
             y = numpy.mean(M, axis=-1)
@@ -250,12 +251,13 @@ def rvs(self, x, n_curves=1, n_samples=1, T=None):
             assert T.shape == (n_curves, n_samples)
         B = numpy.zeros((n_curves, n_samples), dtype=numpy.bool)
         C = numpy.zeros((n_curves, n_samples))
-        for i, j in enumerate(numpy.random.randint(len(self.params['k']),
+        params = self.params['samples']
+        for i, j in enumerate(numpy.random.randint(len(params['k']),
                                                    size=n_curves)):
-            k = self.params['k'][j]
-            p = self.params['p'][j]
-            lambd = exp(dot(x, self.params['alpha'][j]) + self.params['a'][j])
-            c = expit(dot(x, self.params['beta'][j]) + self.params['b'][j])
+            k = params['k'][j]
+            p = params['p'][j]
+            lambd = exp(dot(x, params['alpha'][j]) + params['a'][j])
+            c = expit(dot(x, params['beta'][j]) + params['b'][j])
             z = numpy.random.uniform(size=(n_samples,))
             cdf_now = c * gammainc(
                 k,

test_convoys.py

@@ -147,7 +147,7 @@ def test_gamma_regression_model(c=0.3, lambd=0.1, k=3.0, n=10000):
     model = convoys.regression.Gamma()
     model.fit(X, B, T)
     assert 0.80*c < model.cdf([1], float('inf')) < 1.30*c
-    assert 0.80*k < numpy.mean(model.params['k']) < 1.30*k
+    assert 0.80*k < numpy.mean(model.params['map']['k']) < 1.30*k
 
 
 def _generate_dataframe(cs=[0.3, 0.5, 0.7], k=0.5, lambd=0.1, n=1000):