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) <noreply@anthropic.com>

Author: 2 people

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
 
 

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)
         )

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])

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)

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):

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)

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_]

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()

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)

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(