Fix slowdown when using periodic_save_freq in MCMC

orbitize/results.py

@@ -58,18 +58,24 @@ def __init__(
             self.param_idx = self.system.param_idx
             self.standard_param_idx = self.system.basis.standard_basis_idx
 
-    def add_samples(self, orbital_params, lnlikes, curr_pos=None): 
+    def add_samples(self, orbital_params, lnlikes, curr_pos=None):
         """
-        Add accepted orbits, their likelihoods, and the orbitize version number 
+        Add accepted orbits, their likelihoods, and the orbitize version number
         to the results
 
         Args:
-            orbital_params (np.array): add sets of orbital params (could be multiple) 
+            orbital_params (np.array): add sets of orbital params (could be multiple)
                 to results
             lnlike (np.array): add corresponding lnlike values to results
-            curr_pos (np.array of float): for MCMC only. A multi-D array of the 
+            curr_pos (np.array of float): for MCMC only. A multi-D array of the
                 current walker positions
 
+        .. Note:: ``post``/``lnlike`` are backed by an internal buffer that's
+            over-allocated and grown by doubling, so repeated calls (e.g. from
+            ``periodic_save_freq`` during MCMC) append in amortized O(1) time
+            per row instead of reallocating and copying the full accumulated
+            array on every call.
+
         Written: Henry Ngo, 2018
 
         API Update: Sarah Blunt, 2021

pyproject.toml

@@ -14,7 +14,6 @@ dependencies = [
     "numpy",
     "astropy>=4",
     "scipy",
-    "emcee>=3",
     "ptemcee_for_orbitize",
     "matplotlib",
     "corner",

requirements.txt

@@ -1,7 +1,6 @@
 astropy>=4
 numpy
 scipy
-emcee>=3
 ptemcee_for_orbitize
 matplotlib
 corner

tests/test_mcmc.py

@@ -19,13 +19,11 @@
     'sma1': 0, 'ecc1':1, 'inc1':2, 'aop1':3, 'pan1':4, 'tau1':5, 'plx':6, 'mtot':7
 }
 
-def do_mcmc_runs(num_temps=0, num_threads=1, make_corner_plot=False):
+def do_mcmc_runs(num_temps=1, num_threads=1, make_corner_plot=False):
     """
     Tests the MCMC sampler by making sure it runs
     Args:
         num_temps (int): Number of temperatures to use
-            Uses Parallel Tempering MCMC (ptemcee) if > 1,
-            otherwises, uses Affine-Invariant Ensemble Sampler (emcee)
         num_threads (int): number of threads to run
         make_corner_plot (bool): if True, make a corner plot of the results
     """
@@ -106,8 +104,6 @@ def do_examine_chop_chains(num_temps=0, num_threads=1):
     Tests the MCMC sampler's examine_chains and chop_chains methods
     Args:
         num_temps: Number of temperatures to use
-            Uses Parallel Tempering MCMC (ptemcee) if > 1,
-            otherwises, uses Affine-Invariant Ensemble Sampler (emcee)
         num_threads: number of threads to run
     """
 
@@ -181,7 +177,7 @@ def test_mcmc_param_idx():
     # construct Driver with fixed mass and plx
     n_walkers = 100
     myDriver = Driver(input_file, 'MCMC', 1, 1, 0.01,
-                      mcmc_kwargs={'num_temps': 0, 'num_threads': 1,
+                      mcmc_kwargs={'num_temps': 1, 'num_threads': 1,
                                    'num_walkers': n_walkers}
                       )
 
@@ -190,7 +186,7 @@ def test_mcmc_param_idx():
 
     # construct Driver with no fixed params
     myDriver = Driver(input_file, 'MCMC', 1, 1, 0.01, mass_err=0.1, plx_err=0.2,
-                      mcmc_kwargs={'num_temps': 0, 'num_threads': 1,
+                      mcmc_kwargs={'num_temps': 1, 'num_threads': 1,
                                    'num_walkers': n_walkers}
                       )
 
@@ -201,13 +197,13 @@ def test_mcmc_runs():
     do_mcmc_runs(num_temps=2, num_threads=1, make_corner_plot=True)
     do_mcmc_runs(num_temps=2, num_threads=4)
     # Ensemble MCMC tests
-    do_mcmc_runs(num_temps=0, num_threads=1)
-    do_mcmc_runs(num_temps=0, num_threads=8)
+    do_mcmc_runs(num_temps=1, num_threads=1)
+    do_mcmc_runs(num_temps=1, num_threads=8)
 
 def test_chop_chains():
     # Test examine/chop chains
     do_examine_chop_chains(num_temps=5)  # PT
-    do_examine_chop_chains(num_temps=0)  # Ensemble
+    do_examine_chop_chains(num_temps=1)  # Ensemble
 
 
 if __name__ == "__main__":

tests/test_multiplanet.py

@@ -226,16 +226,16 @@ def test_fit_selfconsist():
     # start walkers near the true location for the orbital parameters
     np.random.seed(123)
     # planet b
-    samp.curr_pos[:, 0] = np.random.normal(b_params[0], 0.01, n_walkers)  # sma
-    samp.curr_pos[:, 1] = np.random.normal(b_params[-1], 0.01, n_walkers)  # tau
+    samp.curr_pos[0, :, 0] = np.random.normal(b_params[0], 0.01, n_walkers)  # sma
+    samp.curr_pos[0, :, 1] = np.random.normal(b_params[-1], 0.01, n_walkers)  # tau
     # planet c
-    samp.curr_pos[:, 2] = np.random.normal(c_params[0], 0.01, n_walkers)  # sma
-    samp.curr_pos[:, 3] = np.random.normal(c_params[-1], 0.01, n_walkers)  # tau
+    samp.curr_pos[0, :, 2] = np.random.normal(c_params[0], 0.01, n_walkers)  # sma
+    samp.curr_pos[0, :, 3] = np.random.normal(c_params[-1], 0.01, n_walkers)  # tau
     # we will make a fairly broad mass starting position
-    samp.curr_pos[:, 4] = np.random.uniform(mass_b * 0.25, mass_b * 4, n_walkers)
-    samp.curr_pos[:, 5] = np.random.uniform(mass_c * 0.25, mass_c * 4, n_walkers)
-    samp.curr_pos[0, 4] = mass_b
-    samp.curr_pos[0, 5] = mass_c
+    samp.curr_pos[0, :, 4] = np.random.uniform(mass_b * 0.25, mass_b * 4, n_walkers)
+    samp.curr_pos[0, :, 5] = np.random.uniform(mass_c * 0.25, mass_c * 4, n_walkers)
+    samp.curr_pos[0, 0, 4] = mass_b
+    samp.curr_pos[0, 0, 5] = mass_c
 
     samp.run_sampler(n_walkers * 50, burn_steps=50)
 
@@ -259,5 +259,5 @@ def test_fit_selfconsist():
 
 
 if __name__ == "__main__":
-    # test_compute_model()
+    test_compute_model()
     test_fit_selfconsist()