Repository: eigenfoo/littlemcmc

Issue:
s/size/model_ndim/g
Currently we have the `size` parameter in `lmc.sample`. However, `model_ndim` might be a more descriptive variable name (and also mirrors what PyMC3 uses).

Comments:

Patch:
diff --git a/docs/_static/notebooks/quickstart.ipynb b/docs/_static/notebooks/quickstart.ipynb
index 3ea0e3d..c171dbc 100644
--- a/docs/_static/notebooks/quickstart.ipynb
+++ b/docs/_static/notebooks/quickstart.ipynb
@@ -72,10 +72,6 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/home/george/littlemcmc/venv/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: divide by zero encountered in log\n",
-      "  \n",
-      "/home/george/littlemcmc/venv/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: divide by zero encountered in log\n",
-      "  \n",
       "/home/george/littlemcmc/venv/lib/python3.6/site-packages/ipykernel_launcher.py:2: RuntimeWarning: divide by zero encountered in log\n",
       "  \n"
      ]
@@ -84,7 +80,7 @@
    "source": [
     "trace, stats = lmc.sample(\n",
     "    logp_dlogp_func=logp_dlogp_func,\n",
-    "    size=1,\n",
+    "    model_ndim=1,\n",
     "    draws=1000,\n",
     "    tune=500,\n",
     "    progressbar=None,  # HTML progress bars don't render well in RST.\n",
@@ -106,8 +102,8 @@
     {
      "data": {
       "text/plain": [
-       "array([[ 0.08150886, -0.91618983, -1.20928858, ..., -0.06641805,\n",
-       "        -1.00700313,  1.09983883]])"
+       "array([[ 0.91989736,  0.89910128,  0.88961585, ...,  1.31528136,\n",
+       "        -1.52068354, -0.61179308]])"
       ]
      },
      "execution_count": 4,
@@ -147,24 +143,24 @@
     {
      "data": {
       "text/plain": [
-       "{'depth': array([2, 2, 1, ..., 1, 2, 2]),\n",
-       " 'step_size': array([1.17740792, 1.17740792, 1.17740792, ..., 1.2966478 , 1.2966478 ,\n",
-       "        1.2966478 ]),\n",
+       "{'depth': array([2, 2, 1, ..., 2, 2, 1]),\n",
+       " 'step_size': array([1.15894778, 1.15894778, 1.15894778, ..., 1.50735926, 1.50735926,\n",
+       "        1.50735926]),\n",
        " 'tune': array([False, False, False, ..., False, False, False]),\n",
-       " 'mean_tree_accept': array([0.87820821, 0.46292512, 0.87717804, ..., 1.        , 0.86187468,\n",
-       "        0.98578587]),\n",
-       " 'step_size_bar': array([1.20720148, 1.20720148, 1.20720148, ..., 1.24048036, 1.24048036,\n",
-       "        1.24048036]),\n",
-       " 'tree_size': array([3., 3., 1., ..., 1., 3., 3.]),\n",
+       " 'mean_tree_accept': array([0.75240493, 1.        , 1.        , ..., 0.96866415, 0.96184355,\n",
+       "        1.        ]),\n",
+       " 'step_size_bar': array([1.24138909, 1.24138909, 1.24138909, ..., 1.38460359, 1.38460359,\n",
+       "        1.38460359]),\n",
+       " 'tree_size': array([3., 3., 1., ..., 3., 3., 1.]),\n",
        " 'diverging': array([False, False, False, ..., False, False, False]),\n",
-       " 'energy_error': array([ 0.        ,  0.17517448,  0.13104529, ..., -0.00202233,\n",
-       "         0.22495278,  0.04357828]),\n",
-       " 'energy': array([1.18523089, 3.17544136, 1.8060492 , ..., 0.9244717 , 1.42648794,\n",
-       "        1.55020885]),\n",
-       " 'max_energy_error': array([ 0.18993171,  1.45544483,  0.13104529, ..., -0.00202233,\n",
-       "         0.22495278, -0.21811027]),\n",
-       " 'model_logp': array([-0.92226038, -1.33864044, -1.65012797, ..., -0.92114421,\n",
-       "        -1.42596619, -1.52376126])}"
+       " 'energy_error': array([ 0.13713966, -0.00798758, -0.00358263, ...,  0.09872431,\n",
+       "         0.1215682 , -0.40449202]),\n",
+       " 'energy': array([1.913862  , 1.5033014 , 1.48600433, ..., 1.82498413, 2.11625971,\n",
+       "        1.70961724]),\n",
+       " 'max_energy_error': array([ 0.45459613, -0.09927514, -0.00358263, ..., -0.25742807,\n",
+       "        -0.35807476, -0.40449202]),\n",
+       " 'model_logp': array([-1.34204411, -1.32313009, -1.31464672, ..., -1.78392107,\n",
+       "        -2.07517774, -1.10608392])}"
       ]
      },
      "execution_count": 6,
@@ -218,7 +214,7 @@
    "source": [
     "trace, stats = lmc.sample(\n",
     "    logp_dlogp_func=logp_dlogp_func,\n",
-    "    size=1,\n",
+    "    model_ndim=1,\n",
     "    draws=1000,\n",
     "    tune=500,\n",
     "    progressbar=None,\n",
@@ -232,10 +228,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=1, target_accept=0.9)\n",
+    "step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, model_ndim=1, target_accept=0.9)\n",
     "trace, stats = lmc.sample(\n",
     "    logp_dlogp_func=logp_dlogp_func,\n",
-    "    size=1,\n",
+    "    model_ndim=1,\n",
     "    step=step,\n",
     "    draws=1000,\n",
     "    tune=500,\n",
diff --git a/littlemcmc/base_hmc.py b/littlemcmc/base_hmc.py
index 4946ad2..1d08bd4 100644
--- a/littlemcmc/base_hmc.py
+++ b/littlemcmc/base_hmc.py
@@ -32,7 +32,7 @@ class BaseHMC:
     def __init__(
         self,
         logp_dlogp_func: Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]],
-        size: int,
+        model_ndim: int,
         scaling: Optional[np.ndarray],
         is_cov: bool,
         potential: QuadPotential,
@@ -52,7 +52,7 @@ class BaseHMC:
         logp_dlogp_func : Python callable
             Python callable that returns the log-probability and derivative of
             the log-probability, respectively.
-        size : int
+        model_ndim : int
             Total number of parameters. Dimensionality of the output of
             ``logp_dlogp_func``.
         scaling : 1 or 2-dimensional array-like
@@ -98,8 +98,8 @@ class BaseHMC:
         self.adapt_step_size = adapt_step_size
         self.Emax = Emax
         self.iter_count = 0
-        self.size = size
-        self.step_size = step_scale / (size ** 0.25)
+        self.model_ndim = model_ndim
+        self.step_size = step_scale / (model_ndim ** 0.25)
         self.target_accept = target_accept
         self.step_adapt = step_sizes.DualAverageAdaptation(
             self.step_size, target_accept, gamma, k, t0
@@ -108,9 +108,9 @@ class BaseHMC:
 
         if scaling is None and potential is None:
             # Default to diagonal quadpotential
-            mean = np.zeros(size)
-            var = np.ones(size)
-            potential = QuadPotentialDiagAdapt(size, mean, var, 10)
+            mean = np.zeros(model_ndim)
+            var = np.ones(model_ndim)
+            potential = QuadPotentialDiagAdapt(model_ndim, mean, var, 10)
 
         if scaling is not None and potential is not None:
             raise ValueError("Cannot specify both `potential` and `scaling`.")
diff --git a/littlemcmc/hmc.py b/littlemcmc/hmc.py
index 926e172..2c01f69 100644
--- a/littlemcmc/hmc.py
+++ b/littlemcmc/hmc.py
@@ -52,7 +52,7 @@ class HamiltonianMC(BaseHMC):
     def __init__(
         self,
         logp_dlogp_func: Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]],
-        size: int,
+        model_ndim: int,
         scaling: Optional[np.ndarray] = None,
         is_cov: bool = False,
         potential=None,
@@ -74,7 +74,7 @@ class HamiltonianMC(BaseHMC):
         logp_dlogp_func : Python callable
             Python callable that returns the log-probability and derivative of
             the log-probability, respectively.
-        size : int
+        model_ndim : int
             Total number of parameters. Dimensionality of the output of
             ``logp_dlogp_func``.
         scaling : 1 or 2-dimensional array-like
@@ -124,7 +124,7 @@ class HamiltonianMC(BaseHMC):
             step_scale=step_scale,
             is_cov=is_cov,
             logp_dlogp_func=logp_dlogp_func,
-            size=size,
+            model_ndim=model_ndim,
             potential=potential,
             Emax=Emax,
             target_accept=target_accept,
diff --git a/littlemcmc/nuts.py b/littlemcmc/nuts.py
index 2c9aa33..d3123bf 100644
--- a/littlemcmc/nuts.py
+++ b/littlemcmc/nuts.py
@@ -109,7 +109,7 @@ class NUTS(BaseHMC):
     def __init__(
         self,
         logp_dlogp_func: Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]],
-        size: int,
+        model_ndim: int,
         scaling: Optional[np.ndarray] = None,
         is_cov: bool = False,
         potential=None,
@@ -132,7 +132,7 @@ class NUTS(BaseHMC):
         logp_dlogp_func : Python callable
             Python callable that returns the log-probability and derivative of
             the log-probability, respectively.
-        size : int
+        model_ndim : int
             Total number of parameters. Dimensionality of the output of
             ``logp_dlogp_func``.
         scaling : 1 or 2-dimensional array-like
@@ -188,7 +188,7 @@ class NUTS(BaseHMC):
         """
         super(NUTS, self).__init__(
             logp_dlogp_func=logp_dlogp_func,
-            size=size,
+            model_ndim=model_ndim,
             scaling=scaling,
             is_cov=is_cov,
             potential=potential,
diff --git a/littlemcmc/sampling.py b/littlemcmc/sampling.py
index 87d046a..50a3c96 100644
--- a/littlemcmc/sampling.py
+++ b/littlemcmc/sampling.py
@@ -31,7 +31,7 @@ _log = logging.getLogger("littlemcmc")
 
 def _sample_one_chain(
     logp_dlogp_func: Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]],
-    size: int,
+    model_ndim: int,
     draws: int,
     tune: int,
     step: Union[NUTS, HamiltonianMC],
@@ -49,7 +49,7 @@ def _sample_one_chain(
         progressbar_position = 0
 
     q = start
-    trace = np.zeros([size, tune + draws])
+    trace = np.zeros([model_ndim, tune + draws])
     stats: List[SamplerWarning] = []
 
     if progressbar == "notebook":
@@ -75,7 +75,7 @@ def _sample_one_chain(
 
 def sample(
     logp_dlogp_func: Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]],
-    size: int,
+    model_ndim: int,
     draws: int = 1000,
     tune: int = 1000,
     step: Union[NUTS, HamiltonianMC] = None,
@@ -96,7 +96,7 @@ def sample(
     logp_dlogp_func: Python callable
         Python callable that returns a tuple of the model joint log probability and its
         derivative, in that order.
-    size: int
+    model_ndim: int
         The number of parameters of the model.
     draws: int
         The number of samples to draw. Defaults to 1000. The number of tuned samples are
@@ -183,7 +183,7 @@ def sample(
     if step is None or start is None:
         start_, step_ = init_nuts(
             logp_dlogp_func=logp_dlogp_func,
-            size=size,
+            model_ndim=model_ndim,
             init=init,
             random_seed=random_seed,
             **kwargs,
@@ -197,7 +197,7 @@ def sample(
     results = Parallel(n_jobs=cores, backend="multiprocessing")(
         delayed(_sample_one_chain)(
             logp_dlogp_func=logp_dlogp_func,
-            size=size,
+            model_ndim=model_ndim,
             draws=draws,
             tune=tune,
             step=step,
@@ -225,7 +225,7 @@ def sample(
 
 def init_nuts(
     logp_dlogp_func: Callable[[np.ndarray], Tuple[np.ndarray, np.ndarray]],
-    size: int,
+    model_ndim: int,
     init: str = "auto",
     random_seed: Union[None, int, List[int]] = None,
     **kwargs,
@@ -278,28 +278,30 @@ def init_nuts(
         np.random.seed(random_seed)
 
     if init == "adapt_diag":
-        start = np.zeros(size)
+        start = np.zeros(model_ndim)
         mean = start
-        var = np.ones(size)
-        potential: QuadPotential = QuadPotentialDiagAdapt(size, mean, var, 10)
+        var = np.ones(model_ndim)
+        potential: QuadPotential = QuadPotentialDiagAdapt(model_ndim, mean, var, 10)
     elif init == "jitter+adapt_diag":
-        start = 2 * np.random.rand(size) - 1
+        start = 2 * np.random.rand(model_ndim) - 1
         mean = start
-        var = np.ones(size)
-        potential = QuadPotentialDiagAdapt(size, mean, var, 10)
+        var = np.ones(model_ndim)
+        potential = QuadPotentialDiagAdapt(model_ndim, mean, var, 10)
     elif init == "adapt_full":
-        start = np.zeros(size)
+        start = np.zeros(model_ndim)
         mean = start
-        cov = np.eye(size)
-        potential = QuadPotentialFullAdapt(size, mean, cov, 10)
+        cov = np.eye(model_ndim)
+        potential = QuadPotentialFullAdapt(model_ndim, mean, cov, 10)
     elif init == "jitter+adapt_full":
-        start = 2 * np.random.rand(size) - 1
+        start = 2 * np.random.rand(model_ndim) - 1
         mean = start
-        cov = np.eye(size)
-        potential = QuadPotentialFullAdapt(size, mean, cov, 10)
+        cov = np.eye(model_ndim)
+        potential = QuadPotentialFullAdapt(model_ndim, mean, cov, 10)
     else:
         raise ValueError("Unknown initializer: {}.".format(init))
 
-    step = NUTS(logp_dlogp_func=logp_dlogp_func, size=size, potential=potential, **kwargs)
+    step = NUTS(
+        logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim, potential=potential, **kwargs
+    )
 
     return start, step


Test patch:
diff --git a/tests/test_hmc.py b/tests/test_hmc.py
index 3fd96c9..375ea0c 100644
--- a/tests/test_hmc.py
+++ b/tests/test_hmc.py
@@ -22,11 +22,11 @@ from test_utils import logp_dlogp_func
 
 def test_leapfrog_reversible():
     np.random.seed(42)
-    size = 1
-    scaling = np.random.rand(size)
-    step = HamiltonianMC(logp_dlogp_func=logp_dlogp_func, size=size, scaling=scaling)
+    model_ndim = 1
+    scaling = np.random.rand(model_ndim)
+    step = HamiltonianMC(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim, scaling=scaling)
     p = step.potential.random()
-    q = np.random.randn(size)
+    q = np.random.randn(model_ndim)
     start = step.integrator.compute_state(p, q)
 
     for epsilon in [0.01, 0.1]:
@@ -41,14 +41,14 @@ def test_leapfrog_reversible():
 
 
 def test_nuts_tuning():
-    size = 1
+    model_ndim = 1
     draws = 5
     tune = 5
-    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=size)
+    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
     chains = 1
     cores = 1
     trace, stats = lmc.sample(
-        logp_dlogp_func, size, draws, tune, step=step, chains=chains, cores=cores
+        logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
 
     assert not step.tune
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
index d47e49f..19fca98 100644
--- a/tests/test_sampling.py
+++ b/tests/test_sampling.py
@@ -22,59 +22,59 @@ from test_utils import logp_dlogp_func
     "method", ["adapt_diag", "jitter+adapt_diag", "adapt_full", "jitter+adapt_full",],
 )
 def test_init_nuts(method):
-    start, step = lmc.init_nuts(logp_dlogp_func=logp_dlogp_func, size=1, init=method)
+    start, step = lmc.init_nuts(logp_dlogp_func=logp_dlogp_func, model_ndim=1, init=method)
     assert isinstance(start, np.ndarray)
     assert len(start) == 1
     assert isinstance(step, lmc.NUTS)
 
 
 def test_hmc_sampling_runs():
-    size = 1
-    step = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, size=size)
+    model_ndim = 1
+    step = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
     draws = 2
     tune = 1
     chains = 1
     cores = 1
     trace, stats = lmc.sample(
-        logp_dlogp_func, size, draws, tune, step=step, chains=chains, cores=cores
+        logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
     assert trace.shape == (1, 2)
 
 
 def test_nuts_sampling_runs():
-    size = 1
-    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=size)
+    model_ndim = 1
+    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
     draws = 2
     tune = 1
     chains = 1
     cores = 1
     trace, stats = lmc.sample(
-        logp_dlogp_func, size, draws, tune, step=step, chains=chains, cores=cores
+        logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
     assert trace.shape == (1, 2)
 
 
 def test_multiprocess_sampling_runs():
-    size = 1
-    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=size)
+    model_ndim = 1
+    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
     draws = 1
     tune = 1
     chains = None
     cores = None
     trace, stats = lmc.sample(
-        logp_dlogp_func, size, draws, tune, step=step, chains=chains, cores=cores
+        logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
 
 
 def test_hmc_recovers_1d_normal():
-    size = 1
-    step = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, size=size)
+    model_ndim = 1
+    step = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
     draws = 1000
     tune = 1000
     chains = 1
     cores = 1
     trace, stats = lmc.sample(
-        logp_dlogp_func, size, draws, tune, step=step, chains=chains, cores=cores
+        logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
 
     assert np.allclose(np.mean(trace), 0, atol=1)
@@ -82,14 +82,14 @@ def test_hmc_recovers_1d_normal():
 
 
 def test_nuts_recovers_1d_normal():
-    size = 1
-    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=size)
+    model_ndim = 1
+    step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
     draws = 1000
     tune = 1000
     chains = 1
     cores = 1
     trace, stats = lmc.sample(
-        logp_dlogp_func, size, draws, tune, step=step, chains=chains, cores=cores
+        logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
 
     assert np.allclose(np.mean(trace), 0, atol=1)