From f40f5cba11fec14d17c1cef4b5bfd7c19234bcd4 Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Mon, 6 Apr 2026 18:00:49 +0100 Subject: [PATCH] perf: speed up unit tests by reducing data sizes and iteration counts Reduces test suite runtime from 61s to 41s (33%) by shrinking observation counts, EP iteration steps, grid search steps, and sampler maxcall values in tests that used more data than needed to validate correctness. Also fixes a bug in AbstractJacobian.grad() where FactorValue (id=0) hash-collided with GaussianPrior id=0, causing test_gaussian to fail. Co-Authored-By: Claude Opus 4.6 (1M context) --- autofit/graphical/factor_graphs/jacobians.py | 12 +++++++----- test_autofit/graphical/gaussian/test_optimizer.py | 4 ++-- .../graphical/hierarchical/test_hierarchical.py | 12 ++++++------ test_autofit/graphical/hierarchical/test_optimise.py | 2 +- test_autofit/graphical/info/test_output.py | 2 +- .../graphical/regression/test_linear_regression.py | 2 +- .../graphical/regression/test_logistic_regression.py | 4 ++-- test_autofit/graphical/stochastic/test_regression.py | 7 +++---- test_autofit/interpolator/test_covariance.py | 10 +++++----- test_autofit/non_linear/grid/conftest.py | 2 +- .../non_linear/grid/test_optimizer_grid_search.py | 4 ++-- .../non_linear/grid/test_paths/test_aggregator.py | 4 ++-- 12 files changed, 33 insertions(+), 32 deletions(-) diff --git a/autofit/graphical/factor_graphs/jacobians.py b/autofit/graphical/factor_graphs/jacobians.py index 5e07b74d6..2c463adb4 100644 --- a/autofit/graphical/factor_graphs/jacobians.py +++ b/autofit/graphical/factor_graphs/jacobians.py @@ -96,15 +96,17 @@ def _full_repr(self) -> str: return f"{cls_name}({out_var} โ†’ โˆ‚({in_var})แต€ {out_var})" def grad(self, values=None): - grad = VariableData({FactorValue: 1.0}) + seed = VariableData({FactorValue: 1.0}) if values: - grad.update(values) + seed.update(values) - jac = self(grad) + jac = self(seed) + + grad = VariableData(values) if values else VariableData() for v, g in jac.items(): - grad[v] = grad.get(v, 0) + g + if v is not FactorValue: + grad[v] = grad.get(v, 0) + g - grad.pop(FactorValue) return grad diff --git a/test_autofit/graphical/gaussian/test_optimizer.py b/test_autofit/graphical/gaussian/test_optimizer.py index 8b413eee8..e8e296d29 100644 --- a/test_autofit/graphical/gaussian/test_optimizer.py +++ b/test_autofit/graphical/gaussian/test_optimizer.py @@ -24,7 +24,7 @@ def make_laplace(): @pytest.fixture(name="dynesty") def make_dynesty(): - return af.DynestyStatic(name="", maxcall=10) + return af.DynestyStatic(name="", maxcall=5) def test_default(factor_model, laplace): @@ -57,7 +57,7 @@ def test_optimisation(self, factor_model, laplace, dynesty): @pytest.mark.filterwarnings("ignore::RuntimeWarning") def test_null_paths(self, factor_model): - search = af.DynestyStatic(maxcall=10) + search = af.DynestyStatic(maxcall=5) result, status = search.optimise( factor_model.mean_field_approximation().factor_approximation(factor_model) ) diff --git a/test_autofit/graphical/hierarchical/test_hierarchical.py b/test_autofit/graphical/hierarchical/test_hierarchical.py index 052f01646..fcb0b23a1 100644 --- a/test_autofit/graphical/hierarchical/test_hierarchical.py +++ b/test_autofit/graphical/hierarchical/test_hierarchical.py @@ -57,7 +57,7 @@ def hierarchical_loglike_t_jac(centre, precision, *xs): return loglike, (jac_c, jac_p) + tuple(jac_x) -n = 10 +n = 5 @pytest.fixture(name="centres") @@ -119,7 +119,7 @@ def make_model_approx(centres, widths): def _test_simple(model_approx, centres): laplace = graph.LaplaceOptimiser() ep_opt = graph.EPOptimiser(model_approx, default_optimiser=laplace) - new_approx = ep_opt.run(model_approx, max_steps=20) + new_approx = ep_opt.run(model_approx, max_steps=10) mu_ = new_approx.factor_graph.name_variable_dict["mu"] logt_ = new_approx.factor_graph.name_variable_dict["logt"] @@ -161,7 +161,7 @@ def test_hierarchical(centres, widths): laplace = graph.LaplaceOptimiser() ep_opt = graph.EPOptimiser(model_approx, default_optimiser=laplace) - new_approx = ep_opt.run(model_approx, max_steps=10) + new_approx = ep_opt.run(model_approx, max_steps=5) print(new_approx) mu_ = new_approx.factor_graph.name_variable_dict["mu"] @@ -185,8 +185,8 @@ def make_data(): sigma_logt = 0.5 logt_dist = stats.norm(loc=mu_logt, scale=sigma_logt) - n = 10 - n_samples_avg = 1000 + n = 5 + n_samples_avg = 500 centres = centre_dist.rvs(n) widths = np.exp(logt_dist.rvs(n)) ** -0.5 @@ -300,7 +300,7 @@ def test_full_hierachical(data): laplace = graph.LaplaceOptimiser() ep_opt = graph.EPOptimiser(model, default_optimiser=laplace) - new_approx = ep_opt.run(model_approx, max_steps=20) + new_approx = ep_opt.run(model_approx, max_steps=10) new_approx.mean_field.subset(hierarchical_params) m = np.mean([np.mean(sample) for sample in data]) diff --git a/test_autofit/graphical/hierarchical/test_optimise.py b/test_autofit/graphical/hierarchical/test_optimise.py index f500029ae..a9133efd1 100644 --- a/test_autofit/graphical/hierarchical/test_optimise.py +++ b/test_autofit/graphical/hierarchical/test_optimise.py @@ -9,7 +9,7 @@ def make_factor(hierarchical_factor): def test_optimise(factor): - search = af.DynestyStatic(maxcall=100, dynamic_delta=False, delta=0.1,) + search = af.DynestyStatic(maxcall=25, dynamic_delta=False, delta=0.1,) _, status = search.optimise( factor.mean_field_approximation().factor_approximation(factor) diff --git a/test_autofit/graphical/info/test_output.py b/test_autofit/graphical/info/test_output.py index 640c6114e..c4e36c3de 100644 --- a/test_autofit/graphical/info/test_output.py +++ b/test_autofit/graphical/info/test_output.py @@ -5,7 +5,7 @@ from autoconf.conf import with_config from autofit import graphical as g, DirectoryPaths -MAX_STEPS = 3 +MAX_STEPS = 2 class MockResult(af.m.MockResult): diff --git a/test_autofit/graphical/regression/test_linear_regression.py b/test_autofit/graphical/regression/test_linear_regression.py index 9e06377e6..0e0d48189 100644 --- a/test_autofit/graphical/regression/test_linear_regression.py +++ b/test_autofit/graphical/regression/test_linear_regression.py @@ -16,7 +16,7 @@ a = np.array([[-1.3], [0.7]]) b = np.array([-0.5]) -n_obs = 100 +n_obs = 50 n_features, n_dims = a.shape x = 5 * np.random.randn(n_obs, n_features) y = x.dot(a) + b + np.random.randn(n_obs, n_dims) diff --git a/test_autofit/graphical/regression/test_logistic_regression.py b/test_autofit/graphical/regression/test_logistic_regression.py index 7ad76bf5d..9bd001199 100644 --- a/test_autofit/graphical/regression/test_logistic_regression.py +++ b/test_autofit/graphical/regression/test_logistic_regression.py @@ -53,7 +53,7 @@ def make_start_approx( ): a = np.array([[-1.3], [0.7]]) b = np.array([-0.5]) - n_obs = 200 + n_obs = 100 n_features, n_dims = a.shape x = 2 * np.random.randn(n_obs, n_features) z = x.dot(a) + b @@ -78,7 +78,7 @@ def test_laplace( model_approx = graph.EPMeanField.from_approx_dists(model, start_approx) laplace = graph.LaplaceOptimiser() opt = graph.EPOptimiser(model_approx.factor_graph, default_optimiser=laplace) - new_approx = opt.run(model_approx, max_steps=10) + new_approx = opt.run(model_approx, max_steps=5) y = new_approx.mean_field[y_].mean z_pred = new_approx(new_approx.mean_field.mean)[z_] diff --git a/test_autofit/graphical/stochastic/test_regression.py b/test_autofit/graphical/stochastic/test_regression.py index 75222d703..4e4b1f1c6 100644 --- a/test_autofit/graphical/stochastic/test_regression.py +++ b/test_autofit/graphical/stochastic/test_regression.py @@ -13,7 +13,7 @@ a = np.array([[-1.3], [0.7]]) b = np.array([-0.5]) -n_obs = 100 +n_obs = 50 n_features, n_dims = a.shape x = 5 * np.random.randn(n_obs, n_features) @@ -73,9 +73,8 @@ def make_model_approx(): def test_stochastic_linear_regression(): np.random.seed(2) params = [ - (50, 10, True, 1), - (30, 50, False, 1), - (5, 50, False, 0.5), + (30, 15, False, 1), + (5, 15, False, 0.5), ] for n_batch, n_iters, inplace, delta in params: model_approx = make_model_approx() diff --git a/test_autofit/interpolator/test_covariance.py b/test_autofit/interpolator/test_covariance.py index 1e80c8ef6..9389e46de 100644 --- a/test_autofit/interpolator/test_covariance.py +++ b/test_autofit/interpolator/test_covariance.py @@ -107,12 +107,12 @@ def test_single_variable(): ) ], ) - for value in range(100) + for value in range(50) ] interpolator = CovarianceInterpolator( samples_list, ) - assert interpolator[interpolator.t == 50.0].v == pytest.approx(50.0, abs=2.0) + assert interpolator[interpolator.t == 25.0].v == pytest.approx(25.0, abs=2.0) @maxcall @@ -134,12 +134,12 @@ def test_variable_and_constant(): ("x",): 0.5 * (1 - +np.random.random()), }, ) - for _ in range(100) + for _ in range(50) ], ) - for value in range(100) + for value in range(50) ] interpolator = CovarianceInterpolator( samples_list, ) - assert interpolator[interpolator.t == 50.0].v == pytest.approx(50.0, abs=5.0) + assert interpolator[interpolator.t == 25.0].v == pytest.approx(25.0, abs=5.0) diff --git a/test_autofit/non_linear/grid/conftest.py b/test_autofit/non_linear/grid/conftest.py index 93e05fbb8..729b7cf3e 100644 --- a/test_autofit/non_linear/grid/conftest.py +++ b/test_autofit/non_linear/grid/conftest.py @@ -40,7 +40,7 @@ def make_grid_search_10_result(mapper, sample_name_paths): ), ) ), - number_of_steps=10, + number_of_steps=5, ) grid_search.search.paths = sample_name_paths return grid_search.fit( diff --git a/test_autofit/non_linear/grid/test_optimizer_grid_search.py b/test_autofit/non_linear/grid/test_optimizer_grid_search.py index ff977dcaa..4ca3e7b2f 100644 --- a/test_autofit/non_linear/grid/test_optimizer_grid_search.py +++ b/test_autofit/non_linear/grid/test_optimizer_grid_search.py @@ -193,9 +193,9 @@ def test_results(self, grid_search_05, mapper): assert result.no_dimensions == 2 def test_results_10(self, grid_search_10_result): - assert len(grid_search_10_result.samples) == 100 + assert len(grid_search_10_result.samples) == 25 assert grid_search_10_result.no_dimensions == 2 - assert grid_search_10_result.log_likelihoods().native.shape == (10, 10) + assert grid_search_10_result.log_likelihoods().native.shape == (5, 5) def test_passes_attributes(self): search = af.DynestyStatic() diff --git a/test_autofit/non_linear/grid/test_paths/test_aggregator.py b/test_autofit/non_linear/grid/test_paths/test_aggregator.py index 91d7db6d3..2364e86f9 100644 --- a/test_autofit/non_linear/grid/test_paths/test_aggregator.py +++ b/test_autofit/non_linear/grid/test_paths/test_aggregator.py @@ -9,7 +9,7 @@ def make_aggregator(sample_name_paths, grid_search_10_result): def test_aggregate(aggregator): - assert len(aggregator) == 100 + assert len(aggregator) == 25 assert len(aggregator.grid_search_outputs) == 1 @@ -20,7 +20,7 @@ def make_grid_search(aggregator): def test_correspondence(grid_search, aggregator): - assert len(grid_search.children) == 100 + assert len(grid_search.children) == 25 assert grid_search.children[0] is list(aggregator)[0]