diff --git a/searches_minimal/nautilus_jax_mlp.py b/searches_minimal/nautilus_jax_mlp.py deleted file mode 100644 index a75b0f9..0000000 --- a/searches_minimal/nautilus_jax_mlp.py +++ /dev/null @@ -1,275 +0,0 @@ -""" -Nautilus JAX-MLP Emulator Prototype ------------------------------------ - -Phase 3 of the nautilus-nn-bottleneck research task -(https://github.com/Jammy2211/autofit_workspace_developer/issues/18). - -Replaces nautilus's sklearn `MLPRegressor` ensemble (`nautilus/neural.py`) -with a JAX implementation and measures the end-to-end effect on the same -scenarios as nautilus_profile.py. Design choices, and why: - - - The whole ensemble trains as ONE jitted program: `jax.vmap` over the - n_networks parameter sets, full-batch Adam for a fixed number of steps. - sklearn trains the 4 MLPs serially, each pinned to a single thread. - - Training inputs are padded to power-of-two bucket sizes with masked - sample weights, so the JIT cache is reused across bounds instead of - recompiling for every new training-set size. - - `predict` is also jitted + padded — it sits on the sampling hot path - (`NeuralBound.contains` calls it for every proposal batch). - - Same architecture and normalisation as nautilus's default - (hidden layers (100, 50, 20), relu, inputs standardised), so the - comparison isolates the training engine, not the model. - - Fixed-step training replaces sklearn's early stopping - (`tol=0, n_iter_no_change=10`); N_STEPS below was chosen so the - emulator matches sklearn's end-to-end sampling quality (log Z / ESS), - which the benchmark verifies. - -This is a *prototype*: HDF5 checkpoint write/read of the emulator is not -implemented, so it must not be used with `filepath=`. Promotion paths -(PyAutoFit wrapper hook vs upstream nautilus PR) are the Phase 4 verdict. - -On this laptop JAX is CPU-only; on a GPU node the ensemble training is -expected to shrink further. The benchmark prints the measured numbers only. -""" - -import time -from pathlib import Path - -import numpy as np - -import jax -import jax.numpy as jnp -import optax - -from searches_minimal.nautilus_profile import ( - TIMINGS, - Sampler, - MLTracker, - NeuralNetworkEmulator, - _train_original, - gaussian_prior_transform, - gaussian_log_likelihood, - fast_prior_transform, - fast_log_likelihood, - model, - N_DIM_FAST, -) - -HIDDEN_LAYERS = (100, 50, 20) -N_STEPS = 600 -LEARNING_RATE = 3e-3 - - -def _pad_size(n, minimum=64): - size = minimum - while size < n: - size *= 2 - return size - - -def _init_params(key, n_dim): - sizes = (n_dim,) + HIDDEN_LAYERS + (1,) - params = [] - for n_in, n_out in zip(sizes[:-1], sizes[1:]): - key, subkey = jax.random.split(key) - w = jax.random.normal(subkey, (n_in, n_out)) * jnp.sqrt(2.0 / n_in) - params.append((w, jnp.zeros(n_out))) - return params - - -def _forward(params, x): - for w, b in params[:-1]: - x = jax.nn.relu(x @ w + b) - w, b = params[-1] - return (x @ w + b)[..., 0] - - -@jax.jit -def _train_ensemble(params_ensemble, x, y, weights): - optimizer = optax.adam(LEARNING_RATE) - - def loss_fn(params): - pred = _forward(params, x) - return jnp.sum(weights * (pred - y) ** 2) / jnp.sum(weights) - - def train_one(params): - opt_state = optimizer.init(params) - - def step(carry, _): - params, opt_state = carry - grads = jax.grad(loss_fn)(params) - updates, opt_state = optimizer.update(grads, opt_state) - params = optax.apply_updates(params, updates) - return (params, opt_state), None - - (params, _), _ = jax.lax.scan(step, (params, opt_state), None, length=N_STEPS) - return params - - return jax.vmap(train_one)(params_ensemble) - - -@jax.jit -def _predict_ensemble(params_ensemble, x): - return jnp.mean(jax.vmap(_forward, in_axes=(0, None))(params_ensemble, x), axis=0) - - -class JaxNeuralNetworkEmulator: - """Drop-in for nautilus.neural.NeuralNetworkEmulator (train/predict only).""" - - _seed_counter = 0 - - @classmethod - def train(cls, x, y, n_networks=4, neural_network_kwargs={}, pool=None): - emulator = cls() - emulator.mean = np.mean(x, axis=0) - emulator.scale = np.std(x, axis=0) - - x_t = (x - emulator.mean) / emulator.scale - n, n_dim = x_t.shape - n_pad = _pad_size(n) - x_pad = np.zeros((n_pad, n_dim)) - x_pad[:n] = x_t - y_pad = np.zeros(n_pad) - y_pad[:n] = y - weights = np.zeros(n_pad) - weights[:n] = 1.0 - - cls._seed_counter += 1 - keys = jax.random.split(jax.random.PRNGKey(cls._seed_counter), n_networks) - params = jax.vmap(lambda k: _init_params(k, n_dim))(keys) - - emulator.params = _train_ensemble( - params, jnp.array(x_pad), jnp.array(y_pad), jnp.array(weights) - ) - jax.block_until_ready(emulator.params) - return emulator - - def predict(self, x): - x_t = (x - self.mean) / self.scale - n = x_t.shape[0] - n_pad = _pad_size(n) - x_pad = np.zeros((n_pad, x_t.shape[1])) - x_pad[:n] = x_t - result = _predict_ensemble(self.params, jnp.array(x_pad)) - return np.asarray(result)[:n] - - -def _use_jax_emulator(): - def train_timed(cls, *args, **kwargs): - t0 = time.time() - result = JaxNeuralNetworkEmulator.train(*args, **kwargs) - TIMINGS.add("nn_train", time.time() - t0) - return result - - NeuralNetworkEmulator.train = classmethod(train_timed) - - -def _use_sklearn_emulator(): - def train_timed(cls, *args, **kwargs): - t0 = time.time() - result = _train_original(cls, *args, **kwargs) - TIMINGS.add("nn_train", time.time() - t0) - return result - - NeuralNetworkEmulator.train = classmethod(train_timed) - - -def run_config(name, prior_transform, log_likelihood, n_dim, n_live=200): - TIMINGS.totals.clear() - TIMINGS.counts.clear() - TIMINGS.per_call.clear() - - tracker = MLTracker() - - def wrapped(params): - log_l = log_likelihood(params) - tracker.record(log_l) - return log_l - - sampler = Sampler( - prior=prior_transform, - likelihood=wrapped, - n_dim=n_dim, - n_live=n_live, - seed=1, - ) - t_start = time.time() - sampler.run(verbose=False) - wall = time.time() - t_start - - points, log_w, log_l = sampler.posterior() - evals_to_ml, _ = tracker.finalise(max_log_l=float(np.max(log_l)), tolerance=1.0) - - nn = TIMINGS.totals.get("nn_train", 0.0) - bound = TIMINGS.totals.get("bound_compute", 0.0) - result = dict( - name=name, - wall=wall, - nn=nn, - geometry=bound - nn, - shell=TIMINGS.totals.get("sample_shell", 0.0), - n_evals=int(sampler.n_like), - log_z=float(sampler.log_z), - ess=float(sampler.n_eff), - evals_to_ml=evals_to_ml, - ) - print( - f" {name:<14s} wall={wall:6.2f} s nn={nn:6.2f} s " - f"shell={result['shell']:5.2f} s logZ={result['log_z']:8.3f} " - f"ESS={result['ess']:7.1f} evals={result['n_evals']}" - ) - return result - - -def format_rows(title, results): - lines = [ - f"--- {title} ---", - f"{'config':<14s} {'wall s':>7s} {'nn s':>6s} {'geom s':>6s} " - f"{'shell s':>7s} {'evals':>7s} {'log Z':>9s} {'ESS':>7s} {'ev->ML':>7s}", - ] - for r in results: - lines.append( - f"{r['name']:<14s} {r['wall']:>7.2f} {r['nn']:>6.2f} " - f"{r['geometry']:>6.2f} {r['shell']:>7.2f} {r['n_evals']:>7d} " - f"{r['log_z']:>9.3f} {r['ess']:>7.1f} " - f"{r['evals_to_ml'] if r['evals_to_ml'] is not None else 'n/a':>7}" - ) - return "\n".join(lines) - - -if __name__ == "__main__": - scenarios = [ - ("fast_gauss_10d", fast_prior_transform, fast_log_likelihood, N_DIM_FAST), - ( - "gaussian_1d_data", - gaussian_prior_transform, - gaussian_log_likelihood, - model.prior_count, - ), - ] - - tables = [] - for title, prior_transform, log_likelihood, n_dim in scenarios: - print(f"== {title} ==") - results = [] - _use_sklearn_emulator() - results.append(run_config("sklearn", prior_transform, log_likelihood, n_dim)) - _use_jax_emulator() - results.append(run_config("jax_mlp", prior_transform, log_likelihood, n_dim)) - # Second JAX run reuses the padded-shape JIT cache — the steady-state - # cost a long fit would see. - results.append( - run_config("jax_mlp_warm", prior_transform, log_likelihood, n_dim) - ) - tables.append(format_rows(title, results)) - - output = "\n\n".join(tables) + "\n" - print() - print(output) - - output_dir = Path(__file__).parent / "output" - output_dir.mkdir(parents=True, exist_ok=True) - summary_path = output_dir / f"{Path(__file__).stem}_summary.txt" - summary_path.write_text(output) - print(f"Summary written to: {summary_path}") diff --git a/searches_minimal/nautilus_sweep.py b/searches_minimal/nautilus_sweep.py deleted file mode 100644 index da22b68..0000000 --- a/searches_minimal/nautilus_sweep.py +++ /dev/null @@ -1,195 +0,0 @@ -""" -Nautilus Zero-Code Tuning Sweep -------------------------------- - -Phase 2 of the nautilus-nn-bottleneck research task -(https://github.com/Jammy2211/autofit_workspace_developer/issues/18). - -Sweeps the tuning surface nautilus 1.0.5 already exposes — no sampler code -changes — and measures wall time against sampling quality, using the same -instrumentation and scenarios as nautilus_profile.py: - - baseline n_networks=4, serial training (the PyAutoFit JAX-path - default: fit_x1_cpu passes pool=None) - nets_2 halve the MLP ensemble - nets_0 no NN at all — pure ellipsoid bounds; sampling efficiency - drops but each extra eval is ~free in the fast regime - pool_s_4 pool=(None, 4): parallel MLP training via a sampler pool, - likelihood stays serial/vectorized (what fit_x1_cpu could - pass but doesn't) - small_mlp one hidden layer (32,) instead of (100, 50, 20) - n_update_2x retrain bounds half as often (n_update = 2 * n_live) - vectorized batched likelihood callback (numpy-vectorized here; the - stand-in for a vmapped GPU likelihood) - -Quality is judged on log Z / ESS / evals-to-ML agreement with the baseline -(reference log Z for the 3-parameter Gaussian problem: ~-57.5). -""" - -import time -from pathlib import Path - -import numpy as np - -from searches_minimal.nautilus_profile import ( - TIMINGS, - Sampler, - MLTracker, - gaussian_prior_transform, - gaussian_log_likelihood, - fast_prior_transform, - fast_log_likelihood, - model, - _mean, - _cov_inv, - N_DIM_FAST, -) - - -def fast_log_likelihood_vectorized(params): - """Batched variant of the fast 10-d Gaussian: (n_batch, n_dim) -> (n_batch,).""" - d = params - _mean - return -0.5 * np.einsum("bi,ij,bj->b", d, _cov_inv, d) - - -def run_config( - name, prior_transform, log_likelihood, n_dim, n_live=200, **sampler_kwargs -): - TIMINGS.totals.clear() - TIMINGS.counts.clear() - TIMINGS.per_call.clear() - - tracker = MLTracker() - vectorized = sampler_kwargs.get("vectorized", False) - - if vectorized: - - def wrapped(params): - log_l = log_likelihood(params) - for value in np.atleast_1d(log_l): - tracker.record(float(value)) - return log_l - - else: - - def wrapped(params): - log_l = log_likelihood(params) - tracker.record(log_l) - return log_l - - sampler = Sampler( - prior=prior_transform, - likelihood=wrapped, - n_dim=n_dim, - n_live=n_live, - seed=1, - **sampler_kwargs, - ) - - t_start = time.time() - sampler.run(verbose=False) - wall = time.time() - t_start - - points, log_w, log_l = sampler.posterior() - max_logl = float(np.max(log_l)) - evals_to_ml, _ = tracker.finalise(max_log_l=max_logl, tolerance=1.0) - - nn = TIMINGS.totals.get("nn_train", 0.0) - bound = TIMINGS.totals.get("bound_compute", 0.0) - result = dict( - name=name, - wall=wall, - nn=nn, - geometry=bound - nn, - shell=TIMINGS.totals.get("sample_shell", 0.0), - like=TIMINGS.totals.get("evaluate_likelihood", 0.0), - n_bounds=TIMINGS.counts.get("bound_compute", 0), - n_evals=int(sampler.n_like), - log_z=float(sampler.log_z), - ess=float(sampler.n_eff), - evals_to_ml=evals_to_ml, - ) - - for pool_attr in ("pool_l", "pool_s"): - pool = getattr(sampler, pool_attr, None) - if pool is not None: - try: - pool.close() - pool.join() - except Exception: - pass - - print( - f" {name:<12s} wall={wall:6.2f} s nn={nn:6.2f} s " - f"logZ={result['log_z']:8.3f} ESS={result['ess']:7.1f} " - f"evals={result['n_evals']}" - ) - return result - - -def format_table(title, results, baseline): - lines = [ - f"--- {title} ---", - f"{'config':<12s} {'wall s':>7s} {'x':>5s} {'nn s':>6s} {'geom s':>6s} " - f"{'shell s':>7s} {'like s':>6s} {'bounds':>6s} {'evals':>7s} " - f"{'log Z':>9s} {'ESS':>7s} {'ev->ML':>7s}", - ] - for r in results: - speedup = baseline["wall"] / r["wall"] - lines.append( - f"{r['name']:<12s} {r['wall']:>7.2f} {speedup:>4.1f}x {r['nn']:>6.2f} " - f"{r['geometry']:>6.2f} {r['shell']:>7.2f} {r['like']:>6.2f} " - f"{r['n_bounds']:>6d} {r['n_evals']:>7d} {r['log_z']:>9.3f} " - f"{r['ess']:>7.1f} " - f"{r['evals_to_ml'] if r['evals_to_ml'] is not None else 'n/a':>7}" - ) - return "\n".join(lines) - - -if __name__ == "__main__": - small_mlp = dict(hidden_layer_sizes=(32,), n_iter_no_change=5) - - print("== fast_gauss_10d (fast-GPU-likelihood regime) ==") - fast_args = (fast_prior_transform, fast_log_likelihood, N_DIM_FAST) - fast_results = [ - run_config("baseline", *fast_args), - run_config("nets_2", *fast_args, n_networks=2), - run_config("nets_0", *fast_args, n_networks=0), - run_config("pool_s_4", *fast_args, pool=(None, 4)), - run_config("small_mlp", *fast_args, neural_network_kwargs=small_mlp), - run_config("n_update_2x", *fast_args, n_update=400), - run_config( - "vectorized", - fast_prior_transform, - fast_log_likelihood_vectorized, - N_DIM_FAST, - vectorized=True, - ), - ] - - print("== gaussian_1d_data (3 params; reference log Z ~ -57.5) ==") - gauss_args = ( - gaussian_prior_transform, - gaussian_log_likelihood, - model.prior_count, - ) - gauss_results = [ - run_config("baseline", *gauss_args), - run_config("nets_0", *gauss_args, n_networks=0), - run_config("pool_s_4", *gauss_args, pool=(None, 4)), - run_config("small_mlp", *gauss_args, neural_network_kwargs=small_mlp), - ] - - tables = [ - format_table("fast_gauss_10d", fast_results, fast_results[0]), - format_table("gaussian_1d_data", gauss_results, gauss_results[0]), - ] - output = "\n\n".join(tables) + "\n" - print() - print(output) - - output_dir = Path(__file__).parent / "output" - output_dir.mkdir(parents=True, exist_ok=True) - summary_path = output_dir / f"{Path(__file__).stem}_summary.txt" - summary_path.write_text(output) - print(f"Summary written to: {summary_path}") diff --git a/searches_minimal/output/nautilus_profile_summary.txt b/searches_minimal/output/nautilus_profile_summary.txt deleted file mode 100644 index f3657de..0000000 --- a/searches_minimal/output/nautilus_profile_summary.txt +++ /dev/null @@ -1,39 +0,0 @@ ---- Nautilus Overhead Profile: gaussian_1d_data (3 params) --- -n_dim / n_live: 3 / 200 -Wall time: 48.67 s -Likelihood evals: 16400 -Raw likelihood: 0.256 ms/call (timed outside sampler, 1000 calls) - ---- Wall-time breakdown --- -nn_train 25.18 s 51.7 % (26 emulator trainings, 26 ellipsoid-ensembles x 4 MLPs) -bound_geometry 12.26 s 25.2 % (26 bounds; ellipsoids/bound: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]) -sample_shell 2.53 s 5.2 % (164 calls) -evaluate_likelihood 6.95 s 14.3 % (164 batched calls) -residue 1.75 s 3.6 % (bookkeeping, shell updates) - ---- Quality (MLTracker contract) --- -Max log L: -38.1725 -Log evidence: -57.4552 -ESS: 10059.8 -Evals to ML: 4278 -Time to ML: 25.96 s - ---- Nautilus Overhead Profile: fast_gauss_10d --- -n_dim / n_live: 10 / 200 -Wall time: 15.52 s -Likelihood evals: 10200 -Raw likelihood: 0.005 ms/call (timed outside sampler, 1000 calls) - ---- Wall-time breakdown --- -nn_train 6.24 s 40.2 % (5 emulator trainings, 5 ellipsoid-ensembles x 4 MLPs) -bound_geometry 5.93 s 38.2 % (5 bounds; ellipsoids/bound: [1, 1, 1, 1, 1]) -sample_shell 2.80 s 18.0 % (102 calls) -evaluate_likelihood 0.09 s 0.6 % (102 batched calls) -residue 0.47 s 3.0 % (bookkeeping, shell updates) - ---- Quality (MLTracker contract) --- -Max log L: -0.0039 -Log evidence: -0.0102 -ESS: 10048.5 -Evals to ML: 1 -Time to ML: 0.01 s