Switching to Bridgestan and JAX

.readthedocs.yml

@@ -5,6 +5,11 @@
 # Required
 version: 2
 
+# Set the OS, Python version and other tools you might need
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
 # Build documentation in the docs/ directory with Sphinx
 sphinx:
   configuration: docs/source/conf.py
@@ -19,7 +24,6 @@ sphinx:
 
 # Optionally set the version of Python and requirements required to build your docs
 python:
-  version: 3.8
   install:
     - method: pip
       path: .[docs]

MANIFEST.in

@@ -4,3 +4,5 @@ recursive-exclude * __pycache__
 recursive-exclude * *.py[co]
 
 include viabel/data/*.stan
+include viabel/data/*.json
+

requirements-dev.txt

@@ -5,7 +5,7 @@
 codecov
 coverage
 pytest
-pystan==2.19.1.1
+pystan>=3.1.0
 
 # lint
 autoflake

requirements-docs.txt

@@ -9,3 +9,4 @@ sphinx_rtd_theme
 ipykernel
 nbsphinx
 nbstripout
+bridgestan

requirements.txt

@@ -1,10 +1,11 @@
 numpy>=1.13
-scipy~=1.7.3
-autograd~=1.3
+scipy>=1.7.3
+jax>=0.4.1
 tqdm~=4.64.1
-paragami~=0.42
-
-pystan~=2.19.1.1
+bridgestan>=2.0.0
+pystan>=3.1.0
 pytest~=7.1.2
 viabel~=0.5.1
-setuptools~=65.5.0
+setuptools~=65.5.0
+jaxlib>=0.4.1
+nest_asyncio>=1.5.8

setup.py

@@ -44,7 +44,7 @@ def get_long_description():
     extras_require={ 'docs' : get_requirements_docs(),
                      'dev' : get_requirements_dev()
                      },
-    python_requires='>=3.8',
+    python_requires='>=3.9',
     classifiers=['Programming Language :: Python :: 3',
                  'Natural Language :: English',
                  'License :: OSI Approved :: MIT License',

viabel/__init__.py

@@ -4,3 +4,5 @@
 from viabel.models import *
 from viabel.objectives import *
 from viabel.optimization import *
+from viabel.patterns import *
+from viabel.function_patterns import *

viabel/_distributions.py

@@ -1,6 +1,6 @@
-import autograd.numpy as np
-from autograd.numpy import linalg
-from autograd.scipy import special, stats
+import jax.numpy as np
+from jax.numpy import linalg
+from jax.scipy import special, stats
 
 
 # See: https://github.com/scipy/scipy/blob/master/scipy/stats/_multivariate.py

viabel/_mc_diagnostics.py

@@ -1,6 +1,6 @@
 import warnings
 
-import autograd.numpy as np
+import jax.numpy as np
 from scipy.fftpack import next_fast_len
 
 
@@ -64,30 +64,30 @@ def ess(samples):
 
     rho_hat_t = np.zeros(n_draw)
     rho_hat_even = 1.0
-    rho_hat_t[0] = rho_hat_even
+    rho_hat_t =rho_hat_t.at[0].set(rho_hat_even)
     rho_hat_odd = 1.0 - (mean_var - np.mean(acov[:, 1])) / var_plus
-    rho_hat_t[1] = rho_hat_odd
+    rho_hat_t = rho_hat_t.at[1].set(rho_hat_odd)
 
     # Geyer's initial positive sequence
     t = 1
     while t < (n_draw - 3) and (rho_hat_even + rho_hat_odd) > 0.0:
         rho_hat_even = 1.0 - (mean_var - np.mean(acov[:, t + 1])) / var_plus
         rho_hat_odd = 1.0 - (mean_var - np.mean(acov[:, t + 2])) / var_plus
         if (rho_hat_even + rho_hat_odd) >= 0:
-            rho_hat_t[t + 1] = rho_hat_even
-            rho_hat_t[t + 2] = rho_hat_odd
+            rho_hat_t = rho_hat_t.at[t+1].set(rho_hat_even)
+            rho_hat_t = rho_hat_t.at[t + 2].set(rho_hat_odd)
         t += 2
 
     max_t = t - 2
     # improve estimation
     if rho_hat_even > 0:
-        rho_hat_t[max_t + 1] = rho_hat_even
+        rho_hat_t = rho_hat_t.at[max_t + 1].set(rho_hat_even)
     # Geyer's initial monotone sequence
     t = 1
     while t <= max_t - 2:
         if (rho_hat_t[t + 1] + rho_hat_t[t + 2]) > (rho_hat_t[t - 1] + rho_hat_t[t]):
-            rho_hat_t[t + 1] = (rho_hat_t[t - 1] + rho_hat_t[t]) / 2.0
-            rho_hat_t[t + 2] = rho_hat_t[t + 1]
+            rho_hat_t = rho_hat_t.at[t + 1].set((rho_hat_t[t - 1] + rho_hat_t[t]) / 2.0)
+            rho_hat_t = rho_hat_t.at[t + 2].set(rho_hat_t[t + 1])
         t += 2
 
     ess = n_chain * n_draw
@@ -117,6 +117,7 @@ def MCSE(sample):
     n_iters, d = sample.shape
     sd_dev = np.sqrt(np.var(sample, ddof=1, axis=0))
     eff_samp = [ess(sample[:, i].reshape(1, n_iters)) for i in range(d)]
+    eff_samp = np.asarray(eff_samp)
     mcse = sd_dev / np.sqrt(eff_samp)
     return eff_samp, mcse
 
@@ -178,7 +179,8 @@ def R_hat_convergence_check(samples, windows, Rhat_threshold=1.1):
     best_W: `int`
         Best window size
     """
-    R_hat_array = [np.max(compute_R_hat(np.array(samples[-window:]), 0)) for window in windows]
+    R_hat_array = [np.max(compute_R_hat(np.array(samples[-window:]), 0)) for window in windows]  
+    R_hat_array = np.asarray(R_hat_array)
     best_R_hat_ind = np.argmin(R_hat_array)
     success = R_hat_array[best_R_hat_ind] <= Rhat_threshold
     return success, windows[best_R_hat_ind]

viabel/_psis.py

@@ -134,6 +134,7 @@ def psislw(lw, Reff=1.0, overwrite_lw=False):
         Pareto tail indices
 
     """
+    lw = np.array(lw)
     if lw.ndim == 2:
         n, m = lw.shape
     elif lw.ndim == 1:

viabel/_utils.py

@@ -1,11 +1,5 @@
-import os
-import pickle
-import shutil
 import time
-from hashlib import md5
-
-import autograd.numpy as np
-import pystan
+import numpy as np
 
 
 def vectorize_if_needed(f, a, axis=-1):
@@ -36,47 +30,4 @@ def __exit__(self, *args):
         self.interval = self.end - self.start
 
 
-def _data_file_path(filename):
-    """Returns the path to an internal file"""
-    return os.path.abspath(os.path.join(__file__, '../stan_models', filename))
-
-
-def _stan_model_cache_dir():
-    return _data_file_path('cached-stan-models')
-
-
-def clear_stan_model_cache():
-    stan_model_dir = _stan_model_cache_dir()
-    if os.path.exists(stan_model_dir):
-        shutil.rmtree(stan_model_dir)
-
-
-def StanModel_cache(model_code=None, model_name=None, **kwargs):
-    """Use just as you would `StanModel`"""
-    if model_code is None:
-        if model_name is None:
-            raise ValueError('Either model_code or model_name must be provided')
-        model_file = _data_file_path(model_name + '.stan')
-        if not os.path.isfile(model_file):
-            raise ValueError('invalid model "{}"'.format(model_name))
-        with open(model_file) as f:
-            model_code = f.read()
-    stan_model_dir = _stan_model_cache_dir()
-    os.makedirs(stan_model_dir, exist_ok=True)
-    code_hash = md5(model_code.encode('ascii')).hexdigest()
-    if model_name is None:
-        cache_fn = 'cached-model-{}.pck'.format(code_hash)
-    else:
-        cache_fn = 'cached-{}-{}.pck'.format(model_name, code_hash)
-    cache_file = os.path.join(stan_model_dir, cache_fn)
-    if os.path.exists(cache_file):
-        print('Using cached StanModel{}'.format('' if model_name is None
-                                                else ' for ' + model_name))
-        with open(cache_file, 'rb') as f:
-            sm = pickle.load(f)
-    else:
-        sm = pystan.StanModel(model_code=model_code, model_name=model_name)
-        with open(cache_file, 'wb') as f:
-            pickle.dump(sm, f)
 
-    return sm

viabel/approximations.py

@@ -1,16 +1,15 @@
 from abc import ABC, abstractmethod
+import scipy
 
-import autograd.numpy as np
-import autograd.numpy.random as npr
-import autograd.scipy.stats.norm as norm
-import autograd.scipy.stats.t as t_dist
-from autograd import elementwise_grad
-from autograd.scipy.linalg import sqrtm
-from paragami import (
-    FlattenFunctionInput, NumericArrayPattern, NumericVectorPattern, PatternDict,
-    PSDSymmetricMatrixPattern)
+import jax.numpy as np
+import numpy.random as npr
+import jax.scipy.stats.norm as norm
+import jax.scipy.stats.t as t_dist
+from jax import jvp
+from viabel.function_patterns import FlattenFunctionInput
+from viabel.patterns import NumericArrayPattern, NumericVectorPattern, PatternDict, PSDSymmetricMatrixPattern
 
-from ._distributions import multivariate_t_logpdf
+from viabel._distributions import multivariate_t_logpdf
 
 __all__ = [
     'ApproximationFamily',
@@ -318,7 +317,10 @@ def _get_mu_sigma_pattern(dim):
     ms_pattern['Sigma'] = PSDSymmetricMatrixPattern(size=dim)
     return ms_pattern
 
-
+def sqrtm(matrix):
+    L = scipy.linalg.cholesky(matrix)
+    return L
+
 class MultivariateT(ApproximationFamily):
     """A full-rank multivariate t approximation family."""
 
@@ -415,17 +417,19 @@ def __init__(self, layers_shapes, nonlinearity=np.tanh, last=np.tanh,
 
     def forward(self, var_param, x):
         log_det_J = np.zeros(x.shape[0])
-        derivative = elementwise_grad(self._nonlinearity)
-        derivative_last = elementwise_grad(self._last)
         for layer_id in range(self._layers):
             W = var_param[str(layer_id)]
             b = var_param[str(layer_id) + "_b"]
             if layer_id + 1 == self._layers:
                 x = self._last(np.dot(x, W) + b)
-                log_det_J += np.log(np.abs(np.dot(derivative_last(x), W.T).sum(axis=1)))
+                _, elementwise_gradient_last = jvp(self._last, (x,), (np.ones_like(x),))
+
+                log_det_J += np.log(np.abs(np.dot(elementwise_gradient_last, W.T).sum(axis=1)))
             else:
                 x = self._nonlinearity(np.dot(x, W) + b)
-                log_det_J += np.log(np.abs(np.dot(derivative(x), W.T).sum(axis=1)))
+                _, elementwise_gradient = jvp(self._nonlinearity, (x,), (np.ones_like(x),))
+
+                log_det_J += np.log(np.abs(np.dot(elementwise_gradient, W.T).sum(axis=1)))
         return x, log_det_J
 
     def sample(self, var_param, n_samples):

viabel/convenience.py

@@ -3,7 +3,7 @@
 from viabel.diagnostics import all_diagnostics
 from viabel.models import Model, StanModel
 from viabel.objectives import ExclusiveKL
-from viabel.optimization import RAABBVI, FASO, RMSProp
+from viabel.optimization import RAABBVI, FASO, RMSProp, AveragedRMSProp
 
 all = [
     'bbvi',
@@ -80,7 +80,7 @@ def bbvi(dimension, *, n_iters=10000, num_mc_samples=10, log_density=None,
         objective = ExclusiveKL(approx, model, num_mc_samples)
     if init_var_param is None:
         init_var_param = approx.init_param()
-    base_opt = RMSProp(learning_rate, diagnostics=True, **RMS_kwargs)
+    base_opt = AveragedRMSProp(learning_rate, diagnostics=True, **RMS_kwargs)
     if adaptive and not fixed_lr:
         opt = RAABBVI(base_opt, **RAABBVI_kwargs)
     elif adaptive and fixed_lr: