caught a goof in optimize that was causing bad results in my experiments and, inspired by what it showed me, rewrote the optimizer unit tests to hopefully catch things like that in the future

Author: pavelkomarov

examples/4_performance_analysis.ipynb

@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "5e5ea3c3-aed3-4b8d-a7c7-1db05fc73f3c",
    "metadata": {},
    "outputs": [],
@@ -26,7 +26,7 @@
     "from tqdm import tqdm\n",
     "from collections import defaultdict\n",
     "import multiprocessing as mp\n",
-    "mp.set_start_method(\"fork\", force=True)\n",
+    "#mp.set_start_method(\"fork\", force=True)\n",
     "\n",
     "from pynumdiff.utils.simulate import sine, triangle, pop_dyn, linear_autonomous, pi_cruise_control, lorenz_x\n",
     "from pynumdiff.utils.evaluate import rmse, error_correlation\n",
@@ -277,8 +277,8 @@
     "\tlegend1 = ax[0].legend(ncol=2, columnspacing=0.5, handletextpad=0, loc='upper left', fontsize=15)\n",
     "\tax[0].add_artist(legend1)\n",
     "\tfor handle in legend1.legend_handles:\n",
-    "\t    handle.set_edgecolor('dimgray')\n",
-    "\t    if len(handle.get_facecolor()) == 1: handle.set_facecolor('dimgray') # for those that are filled\n",
+    "\t\thandle.set_edgecolor('dimgray')\n",
+    "\t\tif len(handle.get_facecolor()) == 1: handle.set_facecolor('dimgray') # for those that are filled\n",
     "\tsim_patches = [mpatches.Patch(color=colors[j], label=sname) for j,(sim,sname) in enumerate(sims)]\n",
     "\tlegend2 = ax[0].legend(handles=sim_patches, loc='upper left', fontsize=15, bbox_to_anchor=(0.175, 1.0))\n",
     "\n",

pynumdiff/optimize/_optimize.py

@@ -230,7 +230,7 @@ def optimize(func, x, dt, dxdt_truth=None, tvgamma=1e-2, search_space_updates={}
                 categorical_params=categorical_combo, search_space_types=search_space_types, dxdt_truth=dxdt_truth,
                 metric=metric, tvgamma=tvgamma, padding=padding, cache=cache)
             _minimize = partial(scipy.optimize.minimize, _obj_fun, method=opt_method, bounds=bounds, options={'maxiter':maxiter})
-            results = [_minimize(p) for p in starting_points]
+            results += [_minimize(p) for p in starting_points]
 
     opt_idx = np.nanargmin([r.fun for r in results])
     opt_point = results[opt_idx].x

pynumdiff/tests/test_optimize.py

@@ -2,59 +2,46 @@
 from pytest import skip
 
 from ..finite_difference import first_order as iterated_finite_difference
-from ..smooth_finite_difference import mediandiff, meandiff, gaussiandiff, friedrichsdiff, butterdiff
+from ..smooth_finite_difference import butterdiff
 from ..basis_fit import spectraldiff
 from ..polynomial_fit import polydiff, savgoldiff, splinediff
 from ..total_variation_regularization import velocity, acceleration, iterative_velocity
 from ..kalman_smooth import rtsdiff
 from ..optimize import optimize
 from ..utils.simulate import pi_cruise_control
+from ..utils.evaluate import rmse
 
 
-# simulation
 dt = 0.01
 x, x_truth, dxdt_truth = pi_cruise_control(duration=2, noise_type='normal', noise_parameters=[0, 0.01], dt=dt)
-cutoff_frequency = 10 # in Hz
-log_gamma = -1.6 * np.log(cutoff_frequency) - 0.71 * np.log(dt) - 5.1
-tvgamma = np.exp(log_gamma)
+cutoff_frequency = 3 # in Hz
+tvgamma = np.exp(-1.6 * np.log(cutoff_frequency) - 0.71 * np.log(dt) - 5.1)
 
 
-def test_finite_difference():
-    params1, val1 = optimize(iterated_finite_difference, x, dt, dxdt_truth=dxdt_truth, padding='auto')
-    params2, val2 = optimize(iterated_finite_difference, x, dt, tvgamma=tvgamma, dxdt_truth=None, padding='auto')
-    assert params1['num_iterations'] == 5
-    assert params2['num_iterations'] == 1
+def test_parallel_same_as_serial():
+    """Ensure running optimize across several processes returns the same result as running in a single process"""
+    params_parallel, val_parallel = optimize(rtsdiff, x, dt, tvgamma=tvgamma, parallel=True)
+    params_serial, val_serial = optimize(rtsdiff, x, dt, tvgamma=tvgamma, parallel=False)
 
-def test_iterative_velocity():
-    params1, val1 = optimize(iterative_velocity, x, dt, dxdt_truth=dxdt_truth, search_space_updates={'num_iterations':1}, padding='auto')
-    params2, val2 = optimize(iterative_velocity, x, dt, tvgamma=tvgamma, search_space_updates={'num_iterations':1}, padding='auto')
+    assert np.allclose(val_serial, val_parallel)
+    assert params_serial == params_parallel
+
+
+def test_targeting_rmse_vs_tvgamma_loss():
+    """Ensure optimization properly targets different metrics"""
+    params_rmse, val_rmse = optimize(splinediff, x, dt, dxdt_truth=dxdt_truth)
+    params_loss, val_loss = optimize(splinediff, x, dt, tvgamma=tvgamma)
 
-    np.testing.assert_almost_equal(params1['gamma'], 0.0001, decimal=4)
-    np.testing.assert_almost_equal(params2['gamma'], 0.0001, decimal=4)
-
-def test_savgoldiff():
-    params1, val1 = optimize(savgoldiff, x, dt, dxdt_truth=dxdt_truth, padding='auto')
-    params2, val2 = optimize(savgoldiff, x, dt, tvgamma=tvgamma, padding='auto')
-    assert (params1['degree'], params1['window_size'], params1['smoothing_win']) == (7, 41, 3)
-    assert (params2['degree'], params2['window_size'], params2['smoothing_win']) == (3, 3, 5)
-
-def test_spectraldiff():
-    params1, val1 = optimize(spectraldiff, x, dt, dxdt_truth=dxdt_truth, padding='auto')
-    params2, val2 = optimize(spectraldiff, x, dt, tvgamma=tvgamma, padding='auto')
-    np.testing.assert_almost_equal(params1['high_freq_cutoff'], 0.18, decimal=2)
-    np.testing.assert_almost_equal(params2['high_freq_cutoff'], 0.155, decimal=2)
-
-def test_polydiff():
-    params1, val1 = optimize(polydiff, x, dt, dxdt_truth=dxdt_truth, search_space_updates={'step_size':1}, padding='auto')
-    params2, val2 = optimize(polydiff, x, dt, tvgamma=tvgamma, search_space_updates={'step_size':1}, padding='auto')
-    assert (params1['degree'], params1['window_size'], params1['kernel']) == (6, 50, 'friedrichs')
-    assert (params2['degree'], params2['window_size'], params2['kernel']) == (3, 10, 'gaussian')
-
-# This test runs in a reasonable amount of time locally but for some reason takes forever in CI
-# def test_rtsdiff_with_irregular_step():
-#     t = np.arange(len(x))*dt; np.random.seed(7) # seed so the test can't randomly fail
-#     t_irreg = t + np.random.uniform(-dt/10, dt/10, *t.shape) # add jostle
-#     params1, val1 = optimize(rtsdiff, x, t, dxdt_truth=dxdt_truth)
-#     params2, val2 = optimize(rtsdiff, x, t_irreg, dxdt_truth=dxdt_truth)
-#     assert val2 < 1.15*val1 # optimization works and comes out similar, since jostle is small
-#     assert params1['qr_ratio']*0.85 < params2['qr_ratio'] < params1['qr_ratio']*1.15
+    x_hat, dxdt_hat = splinediff(x, dt, **params_loss)
+    loss_rmse = rmse(dxdt_truth, dxdt_hat)
+
+    assert val_rmse < loss_rmse < 1.1*val_rmse # This exact bound might break if using a different diff method or data series, but the point is they should be close
+
+
+def test_search_space_updates_applied():
+    """Ensure search space updates are used in optimization"""
+    params2, _ = optimize(butterdiff, x, dt, search_space_updates={'filter_order':2}, tvgamma=tvgamma)
+    params3, _ = optimize(butterdiff, x, dt, search_space_updates={'filter_order':3}, tvgamma=tvgamma)
+
+    assert params2['filter_order'] == 2
+    assert params3['filter_order'] == 3