From a20c84b7d801786ce7a0653f9bc55dd6595203df Mon Sep 17 00:00:00 2001 From: Adam Plowman Date: Mon, 14 Sep 2026 10:19:40 +0100 Subject: [PATCH 1/4] refactor: remove initial attempts at SuS-DA --- matflow/tests/subset_simulation.py | 379 +---------------------------- 1 file changed, 4 insertions(+), 375 deletions(-) diff --git a/matflow/tests/subset_simulation.py b/matflow/tests/subset_simulation.py index 6c00f19a..481fc0f1 100644 --- a/matflow/tests/subset_simulation.py +++ b/matflow/tests/subset_simulation.py @@ -434,149 +434,6 @@ def generate_next_level_samples_ACS( } -def generate_next_level_samples_MLDA_incorrect( - performance, - performance_coarse, - num_chains, - num_states, - dimension, - chain_seeds, - chain_g, - all_x, - all_g, - level_idx, - master_seed, - threshold, - proposal, - num_coarse_states, - transformation: Callable | None = None, - threshold_coarse=None, - debug: bool = False, -): - """ - Incorrect! Multilevel delayed-acceptance (MLDA) modified Metropolis algorithm. - """ - if threshold_coarse is None: - threshold_coarse = threshold - - subset_accept_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - mcmc_accept_arr = np.zeros((num_chains, num_states - 1)) - fine_eval_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - coarse_eval_count_arr = np.zeros((num_chains, num_states - 1)).astype(int) - - debug_chain_states = [] - debug_trial_x = [] - debug_current_x = [] - debug_indices = [] - for chain_index in range(num_chains): - - debug_trial_x_chain_i = [] - debug_current_x_chain_i = [] - debug_indices_chain_i = [] - # proceed this Markov chain until all states have been generated - all_x[chain_index, 0] = chain_seeds[chain_index] - all_g[chain_index, 0] = chain_g[chain_index] - - chain_rng = None - for state_idx in range(1, num_states): - - # RNG seed sequence for Markov chains: - if state_idx == 1: - # spawn key to match the task ID in the matflow workflow - spawn_key = (4, level_idx, chain_index) - chain_rng = np.random.default_rng( - np.random.SeedSequence(master_seed, spawn_key=spawn_key) - ) - - x = all_x[chain_index, state_idx - 1] - g = all_g[chain_index, state_idx - 1] - - x_inner = [x] - g_inner = [g] - for inner_state_idx in range(num_coarse_states): - - current_x_inner = x_inner[-1] - current_g_inner = g_inner[-1] - - trial_x_inner, mcmc_accept_rate = generate_next_state( - x=current_x_inner, - proposal=proposal, - rng=chain_rng, - ) - - if debug: - debug_indices_chain_i.append( - { - "chain_index": chain_index, - "level_idx": level_idx, - "state_idx": state_idx, - "inner_state_idx": inner_state_idx, - } - ) - debug_current_x_chain_i.append(current_x_inner) - debug_trial_x_chain_i.append(trial_x_inner) - - trial_x_inner_t = ( - transformation(trial_x_inner) if transformation else trial_x_inner - ) - trial_g_inner = performance_coarse(trial_x_inner_t) - is_ss_accept_inner = trial_g_inner > threshold_coarse - coarse_eval_count_arr[chain_index, state_idx - 1] += 1 - - new_x_inner = trial_x_inner if is_ss_accept_inner else current_x_inner - new_g_inner = trial_g_inner if is_ss_accept_inner else current_g_inner - - x_inner.append(new_x_inner) - g_inner.append(new_g_inner) - - trial_x = x_inner[-1] - - mcmc_accept_arr[chain_index, state_idx - 1] = mcmc_accept_rate - - current_x = x - current_g = g - fine_eval_arr[chain_index, state_idx - 1] = True - trial_x_t = transformation(trial_x) if transformation else trial_x - trial_g = performance(trial_x_t) - is_ss_accept = trial_g > threshold - - subset_accept_arr[chain_index, state_idx - 1] = is_ss_accept - new_x = trial_x if is_ss_accept else current_x - new_g = trial_g if is_ss_accept else current_g - - all_x[chain_index, state_idx] = new_x - all_g[chain_index, state_idx] = new_g - - if debug: - debug_chain_states.append(chain_rng.bit_generator.state["state"]["state"]) - debug_indices.append(debug_indices_chain_i) - debug_trial_x.append(debug_trial_x_chain_i) - debug_current_x.append(debug_current_x_chain_i) - - subset_accept = np.mean(subset_accept_arr).item() - mcmc_accept = np.mean(mcmc_accept_arr).item() - out = { - "subset_accept": subset_accept, - "mcmc_accept": mcmc_accept, - "fine_eval_rate": fine_eval_arr.mean().item(), - "num_fine_evals": int(fine_eval_arr.sum()), - "num_coarse_evals": int(coarse_eval_count_arr.sum()), - } - if debug: - out.update( - { - "debug_chain_states": debug_chain_states, - # early_stop makes the number of inner hops per state vary, so the - # per-chain trial/current-x lists are ragged -- use dtype=object - # rather than assuming a uniform (num_states, num_coarse_states) shape. - "debug_trial_x": np.array(debug_trial_x, dtype=object), - "debug_current_x": np.array(debug_current_x, dtype=object), - "debug_indices": debug_indices, - } - ) - return out - - def weakest_link_coarse_gradient_xt(x_t, group_idx): """Analytic gradient of `weakest_link_performance_coarse` (i.e. of `max(block_average(x_t)) - y_star`) with respect to `x_t` (the *transformed*/ @@ -612,238 +469,6 @@ def cumsum_transformation_adjoint(grad_xt): return np.cumsum(grad_xt[::-1])[::-1] -def estimate_conservative_threshold_coarse( - performance, - performance_coarse, - chain_seeds, - proposal, - threshold, - transformation: Callable | None = None, - num_pilot_trials: int = 500, - quantile: float = 0.99, - window_width_std: float = 1.0, - min_window_count: int = 30, - rng=None, -): - """Calibrate a conservative coarse-stage threshold via pilot sampling. - - Draws pilot MCMC trial proposals from the *current level's* chain seeds - (i.e. the actual states the real kernel will propose from), evaluates - both fine and coarse performance, and sets - - threshold_coarse = threshold - margin - - where margin is a high quantile of the gap g_fine - g_coarse, restricted - to trials whose coarse response falls near the tail boundary -- the - region where the screening decision is actually consequential. - - Because g_coarse <= g_fine always (block-averaging can only smooth, per - weakest_link_performance_coarse's docstring), this gap is >= 0, and a - sufficiently large margin makes the false-rejection event - {g_coarse <= threshold_coarse and g_fine > threshold} arbitrarily rare - -- at the cost of accepting more coarse trials to the fine stage (lower - screening efficiency). - - Returns - ------- - threshold_coarse : float - margin : float - diagnostics : dict with 'gap', 'g_fine', 'g_coarse', 'window_mask' - """ - if rng is None: - rng = np.random.default_rng() - - num_chains = len(chain_seeds) - reps = int(np.ceil(num_pilot_trials / num_chains)) - - trial_xs = [] - for _ in range(reps): - for seed in chain_seeds: - trial_x, _ = generate_next_state(x=seed, proposal=proposal, rng=rng) - trial_xs.append(trial_x) - trial_xs = np.asarray(trial_xs[:num_pilot_trials]) - - trial_x_t = transformation(trial_xs) if transformation else trial_xs - g_fine = performance(trial_x_t) - g_coarse = performance_coarse(trial_x_t) - gap = g_fine - g_coarse # >= 0 by construction - - # restrict to trials near the tail boundary, where the coarse - # accept/reject decision is actually live. Trials far below threshold - # are uninformative (they'll be correctly rejected either way, and - # including them would let a heavy low-gap bulk dilute the quantile). - window_mask = np.abs(g_coarse - threshold) < window_width_std * np.std(g_coarse) - if (window_count := window_mask.sum()) < min_window_count: - window_mask = np.ones_like(gap, dtype=bool) # fallback: use all pilots - - margin = np.quantile(gap[window_mask], quantile) - threshold_coarse = threshold - margin - - return ( - threshold_coarse, - margin, - { - "gap": gap, - "g_fine": g_fine, - "g_coarse": g_coarse, - "window_count": window_count, - "window_mask": window_mask, - }, - ) - - -def generate_next_level_samples_DA_threshold_calibration( - performance, - performance_coarse, - num_chains, - num_states, - dimension, - chain_seeds, - chain_g, - all_x, - all_g, - level_idx, - master_seed, - threshold, - proposal, - transformation: Callable | None = None, - threshold_coarse=None, - previous_threshold_coarse_margin=None, - num_pilot_trials: int = 0, - threshold_coarse_quantile: float = 0.99, - debug: bool = False, -): - """Delayed-acceptance modified Metropolis algorithm for subset simulation.""" - - add_fine_evals = 0 - add_coarse_evals = 0 - calib_diag = None - if threshold_coarse is None: - # estimate the coarse threshold only for the first subset level: - if level_idx == 0: - if not num_pilot_trials: - threshold_coarse = threshold - margin = 0 - else: - ( - threshold_coarse, - margin, - calib_diag, - ) = estimate_conservative_threshold_coarse( - performance=performance, - performance_coarse=performance_coarse, - chain_seeds=chain_seeds, - proposal=proposal, - threshold=threshold, - transformation=transformation, - num_pilot_trials=num_pilot_trials, - quantile=threshold_coarse_quantile, - ) - add_fine_evals += num_pilot_trials - add_coarse_evals += num_pilot_trials - else: - threshold_coarse = threshold - previous_threshold_coarse_margin - margin = previous_threshold_coarse_margin - # print(f"{threshold=!r} => {threshold_coarse=!r}") - - subset_accept_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - debug_subset_accept_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - coarse_accept_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - mcmc_accept_arr = np.zeros((num_chains, num_states - 1)) - fine_eval_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - - for chain_index in range(num_chains): - - # proceed this Markov chain until all states have been generated - all_x[chain_index, 0] = chain_seeds[chain_index] - all_g[chain_index, 0] = chain_g[chain_index] - - chain_rng = None - for state_idx in range(1, num_states): - - # RNG seed sequence for Markov chains: - if state_idx == 1: - # spawn key to match the task ID in the matflow workflow - spawn_key = (4, level_idx, chain_index) - chain_rng = np.random.default_rng( - np.random.SeedSequence(master_seed, spawn_key=spawn_key) - ) - - x = all_x[chain_index, state_idx - 1] - g = all_g[chain_index, state_idx - 1] - - trial_x, mcmc_accept_rate = generate_next_state( - x=x, - proposal=proposal, - rng=chain_rng, - ) - mcmc_accept_arr[chain_index, state_idx - 1] = mcmc_accept_rate - - # stage 1: cheap coarse-model screening -- only proceed to the expensive - # fine model if the trial looks promising according to the coarse model. - trial_x_t = transformation(trial_x) if transformation else trial_x - trial_g_coarse = performance_coarse(trial_x_t) - is_coarse_accept = trial_g_coarse > threshold_coarse - coarse_accept_arr[chain_index, state_idx - 1] = is_coarse_accept - - current_x = x - current_g = g - - if is_coarse_accept: - # stage 2: expensive fine-model evaluation - fine_eval_arr[chain_index, state_idx - 1] = True - trial_x_t = transformation(trial_x) if transformation else trial_x - trial_g = performance(trial_x_t) - is_ss_accept = trial_g > threshold - new_x = trial_x if is_ss_accept else current_x - new_g = trial_g if is_ss_accept else current_g - debug_is_ss_accept = is_ss_accept - else: - is_ss_accept = False - debug_is_ss_accept = is_ss_accept - new_x = current_x - new_g = current_g - if debug: - # run anyway to see if it would be rejected: - trial_x_t = transformation(trial_x) if transformation else trial_x - trial_g = performance(trial_x_t) - debug_is_ss_accept = trial_g > threshold - - subset_accept_arr[chain_index, state_idx - 1] = is_ss_accept - debug_subset_accept_arr[chain_index, state_idx - 1] = debug_is_ss_accept - - all_x[chain_index, state_idx] = new_x - all_g[chain_index, state_idx] = new_g - - subset_accept = np.mean(subset_accept_arr).item() - coarse_accept = np.mean(coarse_accept_arr).item() - mcmc_accept = np.mean(mcmc_accept_arr).item() - fine_eval_rate = np.mean(fine_eval_arr).item() - num_trials = num_chains * (num_states - 1) - - false_coarse_rejection_rate = None - if debug: - false_coarse_rejection_rate = np.mean( - np.logical_and(~coarse_accept_arr, debug_subset_accept_arr) - ) - - return { - "subset_accept": subset_accept, - "mcmc_accept": mcmc_accept, - "subset_accept_arr": subset_accept_arr, - "coarse_accept_arr": coarse_accept_arr, - "coarse_accept": coarse_accept, - "fine_eval_rate": fine_eval_rate, - "num_fine_evals": int(fine_eval_arr.sum()) + add_fine_evals, - "num_coarse_evals": num_trials + add_coarse_evals, - "threshold_coarse": threshold_coarse, - "threshold_coarse_margin": margin, - "threshold_coarse_debug": calib_diag, - "debug_subset_accept_arr": debug_subset_accept_arr, - "false_coarse_rejection_rate": false_coarse_rejection_rate, - } - - def log_surrogate_weight(g_coarse, threshold, temperature): """Note temperature should be of a similar order of magnitude to `g_coarse` and `threshold`.""" @@ -1250,6 +875,9 @@ def generate_next_level_samples_DA( else np.nan ) + # overall acceptance rate for each of the outer trials + subset_accept = n_fine_accepts / num_outer_trials + # Number of coarse model evaluations. # # This assumes generate_coarse_subchain evaluates the coarse @@ -1280,6 +908,7 @@ def generate_next_level_samples_DA( # Fine correction diagnostics # -------------------------------------------------------- "fine_subset_pass_rate": fine_subset_pass_rate, + "subset_accept": subset_accept, "fine_correction_accept_rate": fine_correction_accept_rate, "outer_move_rate": outer_move_rate, "fine_eval_arr": fine_eval_arr, From 73f6cd8f7ebe7a56414b9a348ad1a8e3726cb320 Mon Sep 17 00:00:00 2001 From: Adam Plowman Date: Mon, 14 Sep 2026 13:20:13 +0100 Subject: [PATCH 2/4] refactor: define a `SubsetSimulationResult` object --- .../demo_workflows/test_demo_workflows.py | 11 +- matflow/tests/subset_simulation.py | 565 +++++++++--------- matflow/tests/subset_simulation_result.py | 220 +++++++ 3 files changed, 510 insertions(+), 286 deletions(-) create mode 100644 matflow/tests/subset_simulation_result.py diff --git a/matflow/tests/demo_workflows/test_demo_workflows.py b/matflow/tests/demo_workflows/test_demo_workflows.py index b42aadfc..47ed2f51 100644 --- a/matflow/tests/demo_workflows/test_demo_workflows.py +++ b/matflow/tests/demo_workflows/test_demo_workflows.py @@ -11,7 +11,6 @@ from matflow.tests.subset_simulation import ( log_surrogate_weight, generate_next_level_samples_DA, - generate_next_level_samples_MLDA_incorrect, get_approx_y_star_random_walk, make_voxel_grouping, subset_simulation, @@ -67,7 +66,7 @@ def test_damask_input_files(tmp_path, save_fig, reference_array_data): @pytest.mark.demo_workflows -@pytest.mark.skip(reason="takes too long") +# @pytest.mark.skip(reason="takes too long") def test_subset_simulation_toy_model_prediction(tmp_path): """Validate the MatFlow subset simulation implementation for a toy model. @@ -91,7 +90,7 @@ def test_subset_simulation_toy_model_prediction(tmp_path): performance = partial(system_analysis_toy_model, dimension=200, target_pf=1e-4) # run via single function implementation: - pf_sf, cov_sf, sus_acc_sf, mcmc_acc_sf = subset_simulation( + result = subset_simulation( dimension=200, performance=performance, p_0=0.1, @@ -114,9 +113,9 @@ def test_subset_simulation_toy_model_prediction(tmp_path): # TODO: also verify same result with `subset_simulation_toy_model_external`, once # that can be submitted without a ridiculous number of processes. - assert pf == pf_sf - assert cov == cov_sf - assert np.allclose(sus_acc, sus_acc_sf) + assert pf == result.pf + assert cov == result.cov + assert np.allclose(sus_acc, result.outer_move_rates) # TODO: store and check mcmc_accept diff --git a/matflow/tests/subset_simulation.py b/matflow/tests/subset_simulation.py index 481fc0f1..3906b73c 100644 --- a/matflow/tests/subset_simulation.py +++ b/matflow/tests/subset_simulation.py @@ -16,6 +16,14 @@ import matflow as mf +from .subset_simulation_result import ( + LevelSamplingResult, + DALevelSamplingResult, + ACSLevelSamplingResult, + SubsetSimulationResult, + rms_jump_distances, +) + @TimeIt.decorator def sample_direct_MC( @@ -208,10 +216,10 @@ def generate_next_level_samples( proposal, transformation: Callable | None = None, debug: bool = False, -): +) -> LevelSamplingResult: subset_accept_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - mcmc_accept_arr = np.zeros((num_chains, num_states - 1)) + component_accept_arr = np.zeros((num_chains, num_states - 1)) for chain_index in range(num_chains): @@ -238,7 +246,7 @@ def generate_next_level_samples( proposal=proposal, rng=chain_rng, ) - mcmc_accept_arr[chain_index, state_idx - 1] = mcmc_accept_rate + component_accept_arr[chain_index, state_idx - 1] = mcmc_accept_rate trial_x_t = transformation(trial_x) if transformation else trial_x trial_g = performance(trial_x_t) @@ -253,14 +261,22 @@ def generate_next_level_samples( all_g[chain_index, state_idx] = new_g subset_accept = np.mean(subset_accept_arr).item() - mcmc_accept = np.mean(mcmc_accept_arr).item() num_fine_evals = num_chains * (num_states - 1) - return { - "subset_accept": subset_accept, - "mcmc_accept": mcmc_accept, - "num_fine_evals": num_fine_evals, - "num_coarse_evals": 0, - } + component_acceptance_rate = np.mean(component_accept_arr).item() + jump_distances = rms_jump_distances(all_x) + mean_jump_distance = np.mean(jump_distances).item() + outer_move_rate = np.mean(jump_distances > 0).item() + + return LevelSamplingResult( + x=all_x, + g=all_g, + component_acceptance_rate=component_acceptance_rate, + subset_acceptance_rate=subset_accept, + mean_jump_distance=mean_jump_distance, + jump_distances=jump_distances, + outer_move_rate=outer_move_rate, + num_fine_evals=num_fine_evals, + ) def generate_next_level_samples_CS( @@ -278,7 +294,7 @@ def generate_next_level_samples_CS( prop_std, transformation: Callable | None = None, debug: bool = False, -): +) -> LevelSamplingResult: """Conditional sampling algorithm for generating states in the subset simulation level (aka subset infinity). @@ -325,11 +341,20 @@ def generate_next_level_samples_CS( subset_accept = np.mean(subset_accept_arr).item() num_fine_evals = num_chains * (num_states - 1) - return { - "subset_accept": subset_accept, - "num_fine_evals": num_fine_evals, - "num_coarse_evals": 0, - } + jump_distances = rms_jump_distances(all_x) + mean_jump_distance = np.mean(jump_distances).item() + outer_move_rate = np.mean(jump_distances > 0).item() + + return LevelSamplingResult( + x=all_x, + g=all_g, + component_acceptance_rate=1.0, + subset_acceptance_rate=subset_accept, + mean_jump_distance=mean_jump_distance, + jump_distances=jump_distances, + outer_move_rate=outer_move_rate, + num_fine_evals=num_fine_evals, + ) def generate_next_level_samples_ACS( @@ -349,7 +374,7 @@ def generate_next_level_samples_ACS( prop_std=1.0, lambda_=1.0, debug: bool = False, -): +) -> ACSLevelSamplingResult: """Adaptive conditional sampling algorithm for generating states in the subset simulation level (aka adaptive subset infinity). @@ -426,12 +451,20 @@ def generate_next_level_samples_ACS( subset_accept = np.mean(batch_avgs).item() num_fine_evals = num_chains * (num_states - 1) - return { - "lambda_": lambda_, - "subset_accept": subset_accept, - "num_fine_evals": num_fine_evals, - "num_coarse_evals": 0, - } + + jump_distances = rms_jump_distances(all_x) + mean_jump_distance = np.mean(jump_distances).item() + outer_move_rate = np.mean(jump_distances > 0).item() + + return ACSLevelSamplingResult( + lambda_=lambda_, + mcmc_acceptance_rate=1.0, + subset_acceptance_rate=subset_accept, + mean_jump_distance=mean_jump_distance, + jump_distances=jump_distances, + outer_move_rate=outer_move_rate, + num_fine_evals=num_fine_evals, + ) def weakest_link_coarse_gradient_xt(x_t, group_idx): @@ -494,6 +527,7 @@ def generate_coarse_subchain( current_sub_chain_gc = gc inner_accepts = 0 + mmh_component_acceptance_sum = 0.0 debug_data = {} if debug: @@ -507,9 +541,10 @@ def generate_coarse_subchain( if debug: debug_data["rng_states"].append(rng.bit_generator.state["state"]) - trial_x, _ = generate_next_state( + trial_x, mmh_component_acceptance = generate_next_state( x=current_sub_chain_x, proposal=proposal, rng=rng ) + mmh_component_acceptance_sum += mmh_component_acceptance if debug: debug_data["current_sub_chain_x"].append(current_sub_chain_x) @@ -536,7 +571,13 @@ def generate_coarse_subchain( current_sub_chain_gc = trial_gc inner_accepts += 1 - return current_sub_chain_x, current_sub_chain_gc, inner_accepts, debug_data + return ( + current_sub_chain_x, + current_sub_chain_gc, + inner_accepts, + mmh_component_acceptance_sum, + debug_data, + ) def generate_next_level_samples_DA( @@ -559,7 +600,7 @@ def generate_next_level_samples_DA( num_inner_states: int = 1, spawn_key: tuple[int] | None = None, debug: bool = False, -): +) -> DALevelSamplingResult: """Fixed-length subchain surrogate transition for Subset Simulation. This is the randomised-length subchain surrogate transition (RST) algorithm but with a @@ -627,6 +668,7 @@ def generate_next_level_samples_DA( # ------------------------------------------------------------ all_gc = np.full((num_chains, num_states), np.nan, dtype=float) + total_mmh_component_acceptance = 0.0 debug_data = {} if debug: @@ -687,6 +729,7 @@ def generate_next_level_samples_DA( psi, psi_gc, n_inner_accepts, + mmh_component_acceptance_sum, sub_chain_debug_data, ) = generate_coarse_subchain( x=current_x, @@ -702,6 +745,7 @@ def generate_next_level_samples_DA( chain_idx=chain_index, debug=debug, ) + total_mmh_component_acceptance += mmh_component_acceptance_sum if debug: debug_dat_cs_ij["generate_coarse_subchain_data"] = sub_chain_debug_data @@ -829,45 +873,47 @@ def generate_next_level_samples_DA( # ============================================================ n_inner_proposals = num_chains * (num_states - 1) * num_inner_states + n_inner_accepts = int(inner_accept_count_arr.sum()) n_fine_evals = int(fine_eval_arr.sum()) n_fine_subset_pass = int(fine_subset_pass_arr.sum()) n_fine_accepts = int(fine_accept_arr.sum()) n_endpoint_moves = int(endpoint_move_arr.sum()) - # Mean acceptance probability of individual coarse MH steps. - inner_accept_rate = ( + component_acceptance_rate = ( + total_mmh_component_acceptance / n_inner_proposals + if n_inner_proposals > 0 + else np.nan + ).item() + + coarse_acceptance_rate = ( n_inner_accepts / n_inner_proposals if n_inner_proposals > 0 else np.nan ) - # Fraction of outer transitions for which the coarse subchain - # produced an endpoint different from the current state. endpoint_move_rate = ( n_endpoint_moves / num_outer_trials if num_outer_trials > 0 else np.nan ) - # Fraction of outer transitions requiring an expensive evaluation. fine_eval_rate = n_fine_evals / num_outer_trials if num_outer_trials > 0 else np.nan - # Of the endpoints that were evaluated by the fine model, - # how many were actually in the fine subset? + # Of the endpoints actually evaluated with the fine model, how many satisfy the fine + # subset condition? fine_subset_pass_rate = ( n_fine_subset_pass / n_fine_evals if n_fine_evals > 0 else np.nan ) - # Of the endpoints passing the fine subset, how many passed - # the final RST correction? - fine_correction_accept_rate = ( + # Of the endpoints in the fine subset, how many pass the final delayed-acceptance + # correction? + fine_correction_acceptance_rate = ( n_fine_accepts / n_fine_subset_pass if n_fine_subset_pass > 0 else np.nan ) - # Probability of an actual outer-chain move. + # Fraction of outer transitions that actually change the stored state. outer_move_rate = ( n_fine_accepts / num_outer_trials if num_outer_trials > 0 else np.nan ) - # Among fine-evaluated endpoints, fraction also above the - # coarse threshold. + # Among endpoints in the fine subset, how many are also above the coarse threshold? coarse_given_fine = ( np.sum(endpoint_coarse_subset_pass_arr & fine_subset_pass_arr) / n_fine_subset_pass @@ -875,57 +921,44 @@ def generate_next_level_samples_DA( else np.nan ) - # overall acceptance rate for each of the outer trials - subset_accept = n_fine_accepts / num_outer_trials + jump_distances = rms_jump_distances(all_x) + mean_jump_distance = np.mean(jump_distances).item() - # Number of coarse model evaluations. - # - # This assumes generate_coarse_subchain evaluates the coarse - # model once per inner step. If it skips evaluation when the - # MMH proposal makes no move, adjust this using a counter - # returned by generate_coarse_subchain. + # One coarse evaluation for each chain seed plus one per# inner proposal. num_coarse_evals = num_chains + n_inner_proposals - return { - # -------------------------------------------------------- - # Main computational quantities - # -------------------------------------------------------- - "num_fine_evals": n_fine_evals, - "num_coarse_evals": num_coarse_evals, - "fine_eval_rate": fine_eval_rate, - # -------------------------------------------------------- - # Inner coarse-MH diagnostics - # -------------------------------------------------------- - "inner_accept_rate": inner_accept_rate, - "inner_accept_count_arr": inner_accept_count_arr, - "inner_accept_rate_arr": inner_accept_rate_arr, - # -------------------------------------------------------- - # Endpoint diagnostics - # -------------------------------------------------------- - "endpoint_move_rate": endpoint_move_rate, - "endpoint_move_arr": endpoint_move_arr, - # -------------------------------------------------------- - # Fine correction diagnostics - # -------------------------------------------------------- - "fine_subset_pass_rate": fine_subset_pass_rate, - "subset_accept": subset_accept, - "fine_correction_accept_rate": fine_correction_accept_rate, - "outer_move_rate": outer_move_rate, - "fine_eval_arr": fine_eval_arr, - "fine_subset_pass_arr": fine_subset_pass_arr, - "fine_accept_arr": fine_accept_arr, - "fine_log_alpha_arr": fine_log_alpha_arr, - # -------------------------------------------------------- - # Coarse-vs-fine diagnostic at endpoints - # -------------------------------------------------------- - "coarse_given_fine": coarse_given_fine, - "endpoint_coarse_subset_pass_arr": endpoint_coarse_subset_pass_arr, - # -------------------------------------------------------- - # Optional debugging arrays - # -------------------------------------------------------- - "all_gc": all_gc, - "debug_data": debug_data, - } + outer_move_rate_2 = np.mean(jump_distances > 0).item() + assert outer_move_rate == outer_move_rate_2 + + return DALevelSamplingResult( + x=all_x, + g=all_g, + component_acceptance_rate=component_acceptance_rate, + subset_acceptance_rate=fine_subset_pass_rate, + mean_jump_distance=mean_jump_distance, + num_fine_evals=n_fine_evals, + num_coarse_evals=num_coarse_evals, + jump_distances=jump_distances if debug else None, + debug_data=debug_data if debug else None, + coarse_acceptance_rate=coarse_acceptance_rate, + endpoint_move_rate=endpoint_move_rate, + fine_eval_rate=fine_eval_rate, + fine_subset_pass_rate=fine_subset_pass_rate, + fine_correction_acceptance_rate=fine_correction_acceptance_rate, + outer_move_rate=outer_move_rate, + coarse_given_fine=coarse_given_fine, + all_gc=all_gc, + coarse_acceptance_rates=(inner_accept_rate_arr if debug else None), + coarse_acceptance_counts=(inner_accept_count_arr if debug else None), + endpoint_moves=(endpoint_move_arr if debug else None), + fine_evals=(fine_eval_arr if debug else None), + fine_subset_passes=(fine_subset_pass_arr if debug else None), + fine_correction_accepts=(fine_accept_arr if debug else None), + fine_log_alpha=(fine_log_alpha_arr if debug else None), + endpoint_coarse_subset_passes=( + endpoint_coarse_subset_pass_arr if debug else None + ), + ) def generate_next_level_samples_DA_single_inner( @@ -1098,137 +1131,165 @@ def subset_simulation( transformation: Callable | None = None, mimic_matflow: bool = False, debug: bool = False, -): +) -> SubsetSimulationResult: + """Estimate failure probability using Subset Simulation. + + ``num_levels`` includes the initial direct-Monte-Carlo level. Therefore, + ``num_levels=1`` performs direct Monte Carlo only and does not invoke + ``sampling_method``. + + The ``level_idx`` passed to ``sampling_method`` identifies the transition + from level ``level_idx`` to level ``level_idx + 1``. + """ + + if num_levels < 1: + raise ValueError("num_levels must be at least 1") + + # ------------------------------------------------------------------ + # Level 0: direct Monte Carlo + # ------------------------------------------------------------------ x = sample_direct_MC( dimension, num_samples, seed=master_seed, - spawn_key=(0,), # spawn key to match the task ID in the matflow workflow + spawn_key=(0,), mimic_matflow=mimic_matflow, ) x_t = transformation(x) if transformation else x g = performance(x_t) - sampling_method_kwargs = copy.deepcopy(sampling_method_kwargs) - # pass the performance function on to the sampling method: - if "performance" not in sampling_method_kwargs: - sampling_method_kwargs["performance"] = performance + x_original = x.copy() if debug else None + + sampling_method_kwargs = copy.deepcopy(sampling_method_kwargs or {}) + sampling_method_kwargs.setdefault("performance", performance) + + # ------------------------------------------------------------------ + # Simulation-level results + # ------------------------------------------------------------------ + + thresholds: list[float] = [] + level_covs: list[float] = [] + levels: list[LevelSamplingResult] = [] + num_failed_per_level: list[int] = [] - if debug: - x_original = x.copy() - - level_covs = [] - subset_accepts = [] - coarse_accepts = [] - subset_accepts_arr = [] - coarse_accepts_arr = [] - mcmc_accepts = [] - fine_eval_rates = [] - thresholds = [] - thresholds_coarse = [] - threshold_coarse_debugs = [] - p_coarse = [] - p_fine = [] - n_fine_pass = [] - n_both_pass = [] - p_coarse_given_fine = [] - num_failed_all = [] - - # for DA: includes fine-acceptance even if coarse rejected: - debug_subset_accepts_arr = [] - false_coarse_rejection_rates = [] - - debug_data = {"level_data": []} - - ret = None - all_x = None - all_g = None - # the initial direct-MC draw is always evaluated with the fine model num_fine_evals_total = num_samples num_coarse_evals_total = 0 - for level_idx in range(num_levels): - if debug: - debug_data["level_data"].append({}) + debug_data = {"level_data": []} if debug else None + + chain_seeds = None + chain_g = None + all_x = None + all_g = None + pf = np.nan + + # ------------------------------------------------------------------ + # Analyse each available simulation level + # ------------------------------------------------------------------ + + for level_idx in range(num_levels): num_failed = int(np.sum(g > 0)) - num_failed_all.append(num_failed) + num_failed_per_level.append(num_failed) + num_chains = int(len(g) * p_0) num_states = int(num_samples / num_chains) - g_unsrt = g.copy() - # sort responses - srt_idx = np.argsort(g)[::-1] # sort by closest-to-failure first - g = g[srt_idx] - x = x[srt_idx, :] + g_unsorted = g.copy() + + # Sort with the points closest to failure first. + sort_idx = np.argsort(g)[::-1] + g = g[sort_idx] + x = x[sort_idx, :] - threshold = (g[num_chains - 1] + g[num_chains]) / 2 + threshold = float((g[num_chains - 1] + g[num_chains]) / 2) thresholds.append(threshold) - # failure probability at this level: indicator = np.reshape( - g_unsrt > np.minimum(threshold, 0), (num_chains, num_states) + g_unsorted > np.minimum(threshold, 0), + (num_chains, num_states), ).astype(int) - level_pf = np.mean(indicator) + + level_pf = np.mean(indicator).item() chain_seeds = x[:num_chains] chain_g = g[:num_chains] pf = p_0**level_idx * num_failed / num_samples + if level_idx == 0: - level_cov = np.sqrt((1 - level_pf) / (num_samples * level_pf)) + level_cov = np.sqrt((1 - level_pf) / (num_samples * level_pf)).item() else: - level_cov = estimate_cov(indicator, level_pf) - level_covs.append(level_cov) + level_cov = estimate_cov( + indicator, + level_pf, + ).item() - # TODO: if final level, break here? so num_levels=1 just gives direct MC result? + level_covs.append(level_cov) - if is_finished := threshold > 0: - cov = np.sqrt(sum(np.pow(level_covs, 2))).item() - if debug: - return { - "pf": pf, - "cov": cov, - "subset_accepts": subset_accepts, - "coarse_accepts": coarse_accepts, - "subset_accepts_arr": subset_accepts_arr, - "coarse_accepts_arr": coarse_accepts_arr, - "mcmc_accepts": mcmc_accepts, - "x_original": x_original, - "thresholds": thresholds, - "thresholds_coarse": thresholds_coarse, - "debug_chain_states": (ret or {}).get("debug_chain_states"), - "debug_current_x": (ret or {}).get("debug_current_x"), - "debug_trial_x": (ret or {}).get("debug_trial_x"), - "debug_indices": (ret or {}).get("debug_indices"), - "chain_seeds": chain_seeds, - "chain_g": chain_g, - "all_x": all_x, - "all_g": all_g, - "num_failed": num_failed, - "num_failed_all": num_failed_all, + if debug: + debug_data["level_data"].append( + { + "level_idx": level_idx, "threshold": threshold, "level_pf": level_pf, "level_cov": level_cov, - "fine_eval_rates": fine_eval_rates, - "num_fine_evals": num_fine_evals_total, - "num_coarse_evals": num_coarse_evals_total, - "threshold_coarse_debugs": threshold_coarse_debugs, - "debug_subset_accepts_arr": debug_subset_accepts_arr, - "false_coarse_rejection_rates": false_coarse_rejection_rates, - "p_coarse": p_coarse, - "p_fine": p_fine, - "n_fine_pass": n_fine_pass, - "n_both_pass": n_both_pass, - "p_coarse_given_fine": p_coarse_given_fine, - "debug_data": debug_data, + "num_failed": num_failed, + "chain_seeds": chain_seeds, + "chain_g": chain_g, } - return pf, cov, subset_accepts, mcmc_accepts + ) + + # -------------------------------------------------------------- + # The current level has reached the failure domain. + # -------------------------------------------------------------- + + if threshold > 0: + cov = np.sqrt(np.sum(np.square(level_covs))).item() + + return SubsetSimulationResult( + pf=pf, + cov=cov, + converged=True, + thresholds=np.array(thresholds), + level_covs=np.array(level_covs), + levels=levels, + num_fine_evals=num_fine_evals_total, + num_coarse_evals=num_coarse_evals_total, + num_failed_per_level=np.array(num_failed_per_level), + x_original=x_original, + final_chain_seeds=(chain_seeds if debug else None), + final_chain_g=(chain_g if debug else None), + final_all_x=(all_x if debug else None), + final_all_g=(all_g if debug else None), + debug_data=debug_data, + ) + + # -------------------------------------------------------------- + # No further level is permitted. + # + # In particular, num_levels == 1 reaches this point after + # analysing only the direct-MC samples. + # -------------------------------------------------------------- - all_x = np.ones((num_chains, num_states, dimension)) * np.nan - all_g = np.ones((num_chains, num_states)) * np.nan - ret = sampling_method( + if level_idx == num_levels - 1: + break + + # -------------------------------------------------------------- + # Generate simulation level level_idx + 1. + # -------------------------------------------------------------- + + all_x = np.full( + (num_chains, num_states, dimension), + np.nan, + ) + all_g = np.full( + (num_chains, num_states), + np.nan, + ) + + level_result = sampling_method( num_chains=num_chains, num_states=num_states, dimension=dimension, @@ -1243,115 +1304,59 @@ def subset_simulation( transformation=transformation, **sampling_method_kwargs, ) - if "debug_data" in (ret or {}): - debug_data["level_data"][level_idx]["sampling_method_data"] = ret[ - "debug_data" - ] - - if "subset_accept" in (ret or {}): - subset_accepts.append(ret["subset_accept"]) - - if "coarse_accept" in (ret or {}): - coarse_accepts.append(ret["coarse_accept"]) - - if "coarse_accept_arr" in (ret or {}): - coarse_accepts_arr.append(ret["coarse_accept_arr"]) - - if "subset_accept_arr" in (ret or {}): - subset_accepts_arr.append(ret["subset_accept_arr"]) - - if "mcmc_accept" in (ret or {}): - mcmc_accepts.append(ret["mcmc_accept"]) - - if "lambda_" in (ret or {}): - sampling_method_kwargs["lambda_"] = ret["lambda_"] - - if "fine_eval_rate" in (ret or {}): - fine_eval_rates.append(ret["fine_eval_rate"]) - if "threshold_coarse" in (ret or {}): - thresholds_coarse.append(ret["threshold_coarse"]) - sampling_method_kwargs["previous_threshold_coarse_margin"] = ret[ - "threshold_coarse_margin" - ] - - if "threshold_coarse_debug" in (ret or {}): - threshold_coarse_debugs.append(ret["threshold_coarse_debug"]) - - if "debug_subset_accept_arr" in (ret or {}): - debug_subset_accepts_arr.append(ret["debug_subset_accept_arr"]) - - if "false_coarse_rejection_rate" in (ret or {}): - false_coarse_rejection_rates.append(ret["false_coarse_rejection_rate"]) - - if "p_coarse" in (ret or {}): - p_coarse.append(ret["p_coarse"]) - - if "p_fine" in (ret or {}): - p_fine.append(ret["p_fine"]) - - if "n_fine_pass" in (ret or {}): - n_fine_pass.append(ret["n_fine_pass"]) + if not isinstance(level_result, LevelSamplingResult): + raise TypeError( + "sampling_method must return a " + "LevelSamplingResult instance; " + f"got {type(level_result).__name__}" + ) - if "n_both_pass" in (ret or {}): - n_both_pass.append(ret["n_both_pass"]) + levels.append(level_result) - if "p_coarse_given_fine" in (ret or {}): - p_coarse_given_fine.append(ret["p_coarse_given_fine"]) + # Carry ACS adaptation state into the next transition. + if isinstance(level_result, ACSLevelSamplingResult): + sampling_method_kwargs["lambda_"] = level_result.lambda_ - num_fine_evals_total += ret.get("num_fine_evals", 0) - num_coarse_evals_total += ret.get("num_coarse_evals", 0) + num_fine_evals_total += level_result.num_fine_evals + num_coarse_evals_total += level_result.num_coarse_evals if debug: - debug_data["level_data"][level_idx]["all_x"] = all_x - debug_data["level_data"][level_idx]["all_g"] = all_g + debug_data["level_data"][level_idx][ + "sampling_method_data" + ] = level_result.debug_data - g = all_g.reshape((num_samples)) - x = all_x.reshape((num_samples, dimension)) + # Use the returned result as the authoritative sample state. + all_x = level_result.x + all_g = level_result.g - if debug: - print( - f"Failed to estimate in {num_levels} levels. Try increasing. Debug info " - f"is returned." + x = level_result.x.reshape( + num_samples, + dimension, ) - return { - "pf": pf, - "subset_accepts": subset_accepts, - "coarse_accepts": coarse_accepts, - "subset_accepts_arr": subset_accepts_arr, - "coarse_accepts_arr": coarse_accepts_arr, - "mcmc_accepts": mcmc_accepts, - "x_original": x_original, - "thresholds": thresholds, - "thresholds_coarse": thresholds_coarse, - "debug_chain_states": (ret or {}).get("debug_chain_states"), - "debug_current_x": (ret or {}).get("debug_current_x"), - "debug_trial_x": (ret or {}).get("debug_trial_x"), - "debug_indices": (ret or {}).get("debug_indices"), - "chain_seeds": chain_seeds, - "chain_g": chain_g, - "all_x": all_x, - "all_g": all_g, - "num_failed": num_failed, - "num_failed_all": num_failed_all, - "threshold": threshold, - "level_pf": level_pf, - "level_cov": level_cov, - "fine_eval_rates": fine_eval_rates, - "num_fine_evals": num_fine_evals_total, - "num_coarse_evals": num_coarse_evals_total, - "threshold_coarse_debugs": threshold_coarse_debugs, - "debug_subset_accepts_arr": debug_subset_accepts_arr, - "false_coarse_rejection_rates": false_coarse_rejection_rates, - "p_coarse": p_coarse, - "p_fine": p_fine, - "n_fine_pass": n_fine_pass, - "n_both_pass": n_both_pass, - "p_coarse_given_fine": p_coarse_given_fine, - "debug_data": debug_data, - } - else: - raise RuntimeError(f"Failed to estimate in {num_levels} levels. Try increasing.") + g = level_result.g.reshape(num_samples) + + # ------------------------------------------------------------------ + # The maximum permitted number of levels was reached. + # ------------------------------------------------------------------ + + return SubsetSimulationResult( + pf=pf, + cov=None, + converged=False, + thresholds=np.array(thresholds), + level_covs=np.array(level_covs), + levels=levels, + num_fine_evals=num_fine_evals_total, + num_coarse_evals=num_coarse_evals_total, + num_failed_per_level=np.array(num_failed_per_level), + x_original=x_original, + final_chain_seeds=(chain_seeds if debug else None), + final_chain_g=(chain_g if debug else None), + final_all_x=(all_x if debug else None), + final_all_g=(all_g if debug else None), + debug_data=debug_data, + ) def get_stats( diff --git a/matflow/tests/subset_simulation_result.py b/matflow/tests/subset_simulation_result.py new file mode 100644 index 00000000..31161fde --- /dev/null +++ b/matflow/tests/subset_simulation_result.py @@ -0,0 +1,220 @@ +from dataclasses import dataclass, fields, is_dataclass +from typing import Any +import numpy as np +from numpy.typing import NDArray + + +def rms_jump_distances(all_x: np.ndarray) -> np.ndarray: + """RMS jump distance between consecutive stored chain states.""" + dx = np.diff(all_x, axis=1) + return np.linalg.norm(dx, axis=-1) / np.sqrt(dx.shape[-1]) + + +@dataclass +class LevelSamplingResult: + """Result of generating one conditional Subset Simulation level.""" + + #: Sampled states. + x: np.ndarray + + #: Performance function evaluations associated with each of the sampled states. + g: np.ndarray + + #: Average component acceptance rate from Modified Metropolis Hastings, inclusive of + #: proposals that subsequently fail the subset condition. + component_acceptance_rate: float + + #: Fraction of stored outer-chain transitions that change state. + outer_move_rate: float + + #: Fraction of candidate states that pass the fine-model subset condition. + subset_acceptance_rate: float + + #: Average distance between adjacent states in the outer chain (a subset rejection + #: produces a zero jump distance). + mean_jump_distance: float + + #: Number of (fine) evaluations. + num_fine_evals: int + + #: Number of coarse evaluations, if any were performed. + num_coarse_evals: int = 0 + + # optional detailed diagnostics. + jump_distances: np.ndarray | None = None + debug_data: dict[str, Any] | None = None + + +@dataclass +class ACSLevelSamplingResult(LevelSamplingResult): + lambda_: float = 1.0 + + +@dataclass +class DALevelSamplingResult(LevelSamplingResult): + """Result of generating one delayed-acceptance conditional level.""" + + #: Fraction of individual coarse-MH proposals accepted by the + #: coarse surrogate target. + coarse_acceptance_rate: float = np.nan + + #: Fraction of outer transitions for which the coarse subchain endpoint + #: differs from the current outer-chain state. + endpoint_move_rate: float = np.nan + + #: Fraction of outer transitions requiring a fine-model evaluation. + fine_eval_rate: float = np.nan + + #: Fraction of fine-evaluated candidate endpoints that pass the + #: fine-model subset condition. + fine_subset_pass_rate: float = np.nan + + #: Fraction of candidate endpoints that pass the fine subset condition + #: and are subsequently accepted by the final DA correction. + fine_correction_acceptance_rate: float = np.nan + + #: Among endpoints in the fine subset, fraction that are also above the + #: coarse threshold. + coarse_given_fine: float = np.nan + + #: Coarse performance-function values associated with the stored + #: outer-chain states. + all_gc: np.ndarray | None = None + + # Optional detailed diagnostics. + coarse_acceptance_rates: np.ndarray | None = None + coarse_acceptance_counts: np.ndarray | None = None + endpoint_moves: np.ndarray | None = None + fine_evals: np.ndarray | None = None + fine_subset_passes: np.ndarray | None = None + fine_correction_accepts: np.ndarray | None = None + fine_log_alpha: np.ndarray | None = None + endpoint_coarse_subset_passes: np.ndarray | None = None + + +@dataclass(repr=False) +class SubsetSimulationResult: + pf: float + cov: float | None + converged: bool + + thresholds: NDArray + level_covs: NDArray + levels: list[LevelSamplingResult] + + num_fine_evals: int + num_coarse_evals: int + num_failed_per_level: NDArray + + # Debugging/state data omitted from repr. + x_original: NDArray | None = None + final_chain_seeds: NDArray | None = None + final_chain_g: NDArray | None = None + final_all_x: NDArray | None = None + final_all_g: NDArray | None = None + debug_data: dict[str, Any] | None = None + + @property + def component_acceptance_rates(self) -> NDArray: + return np.array([level.component_acceptance_rate for level in self.levels]) + + @property + def subset_acceptance_rates(self) -> NDArray: + return np.array([level.subset_acceptance_rate for level in self.levels]) + + @property + def mean_jump_distances(self) -> NDArray: + return np.array([level.mean_jump_distance for level in self.levels]) + + @property + def outer_move_rates(self) -> NDArray: + return np.array([level.outer_move_rate for level in self.levels]) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, SubsetSimulationResult): + return NotImplemented + + if type(self) is not type(other): + return False + + return all( + _values_equal( + getattr(self, field.name), + getattr(other, field.name), + ) + for field in fields(self) + ) + + def __repr__(self) -> str: + cls_name = type(self).__name__ + cov_str = f"{self.cov:.4f}" if self.cov is not None else "None" + return ( + f"{cls_name}(\n" + f" pf={self.pf:.4e},\n" + f" cov={cov_str},\n" + f" converged={self.converged!r},\n" + f" thresholds={self.thresholds!r},\n" + f" level_covs={self.level_covs!r},\n" + f" num_conditional_levels={len(self.levels)},\n" + f" component_acceptance_rates={self.component_acceptance_rates!r},\n" + f" subset_acceptance_rates={self.subset_acceptance_rates!r},\n" + f" outer_move_rates={self.outer_move_rates!r},\n" + f" mean_jump_distances={self.mean_jump_distances!r},\n" + f" num_fine_evals={self.num_fine_evals!r},\n" + f" num_coarse_evals={self.num_coarse_evals!r},\n" + f" num_failed_per_level={self.num_failed_per_level!r},\n" + f")" + ) + + +def _values_equal(left: Any, right: Any) -> bool: + """Compare nested result data, including NumPy arrays and NaNs.""" + + if isinstance(left, np.ndarray) or isinstance(right, np.ndarray): + if not (isinstance(left, np.ndarray) and isinstance(right, np.ndarray)): + return False + + return np.array_equal( + left, + right, + equal_nan=True, + ) + + if is_dataclass(left) or is_dataclass(right): + if not (is_dataclass(left) and is_dataclass(right) and type(left) is type(right)): + return False + + return all( + _values_equal( + getattr(left, field.name), + getattr(right, field.name), + ) + for field in fields(left) + ) + + if isinstance(left, dict) or isinstance(right, dict): + if not ( + isinstance(left, dict) + and isinstance(right, dict) + and left.keys() == right.keys() + ): + return False + + return all(_values_equal(left[key], right[key]) for key in left) + + if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): + if type(left) is not type(right) or len(left) != len(right): + return False + + return all(_values_equal(a, b) for a, b in zip(left, right)) + + # Treat NaN diagnostics as equal. + if ( + isinstance(left, (float, np.floating)) + and isinstance(right, (float, np.floating)) + and np.isnan(left) + and np.isnan(right) + ): + return True + + return bool(left == right) From 8182661c163a05d953990a093ace22d61ff505bc Mon Sep 17 00:00:00 2001 From: Adam Plowman Date: Mon, 14 Sep 2026 17:09:11 +0100 Subject: [PATCH 3/4] refactor: encapsulate some behaviour in subset simulation related objects --- matflow/subset_simulation/markov_chain.py | 419 ++++++++++++ .../subset_simulation.py | 607 ++++++------------ .../subset_simulation_result.py | 326 ++++++++++ .../demo_workflows/test_demo_workflows.py | 12 +- matflow/tests/subset_simulation_result.py | 220 ------- matflow/tests/test_subset_simulation.py | 85 +++ matflow/tests/utils.py | 57 +- 7 files changed, 1085 insertions(+), 641 deletions(-) create mode 100644 matflow/subset_simulation/markov_chain.py rename matflow/{tests => subset_simulation}/subset_simulation.py (71%) create mode 100644 matflow/subset_simulation/subset_simulation_result.py delete mode 100644 matflow/tests/subset_simulation_result.py create mode 100644 matflow/tests/test_subset_simulation.py diff --git a/matflow/subset_simulation/markov_chain.py b/matflow/subset_simulation/markov_chain.py new file mode 100644 index 00000000..b678fdeb --- /dev/null +++ b/matflow/subset_simulation/markov_chain.py @@ -0,0 +1,419 @@ +from __future__ import annotations + +from dataclasses import dataclass, fields + +import numpy as np + +from matflow.tests.utils import _values_equal + + +def rms_jump_distances(all_x: np.ndarray) -> np.ndarray: + """RMS jump distance between consecutive stored chain states.""" + dx = np.diff(all_x, axis=1) + return np.linalg.norm(dx, axis=-1) / np.sqrt(dx.shape[-1]) + + +@dataclass +class _ChainAccumulator: + """Mutable internal storage used while generating one chain.""" + + x: np.ndarray + g: np.ndarray + next_state_idx: int = 1 + + num_components_accepted: int = 0 + num_components_proposed: int = 0 + + num_subset_trials: int = 0 + num_subset_accepts: int = 0 + num_outer_moves: int = 0 + + jump_distance_sum: float = 0.0 + num_fine_evals: int = 0 + num_coarse_evals: int = 0 + + @property + def current_x(self) -> np.ndarray: + """Most recently stored chain state.""" + return self.x[self.next_state_idx - 1] + + @property + def current_g(self) -> float: + """Fine performance value of the current state.""" + return float(self.g[self.next_state_idx - 1]) + + @property + def is_complete(self) -> bool: + """Whether the target number of states has been generated.""" + return self.next_state_idx == self.x.shape[0] + + @property + def num_transitions(self) -> int: + """Number of outer-chain transitions generated so far.""" + return self.next_state_idx - 1 + + def append_transition( + self, + trial_x: np.ndarray, + trial_g: float, + threshold: float, + *, + num_components_accepted: int = 0, + num_components_proposed: int = 0, + fine_evaluated: bool = True, + ) -> None: + """Apply and store one proposed transition.""" + + if self.is_complete: + raise ValueError("Cannot append to a completed chain.") + + current_x = self.current_x + current_g = self.current_g + + subset_accepted = trial_g > threshold + + if subset_accepted: + new_x = trial_x + new_g = trial_g + else: + new_x = current_x + new_g = current_g + + outer_moved = not np.array_equal(new_x, current_x) + + jump_distance = np.linalg.norm(new_x - current_x) / np.sqrt(current_x.size) + + self.x[self.next_state_idx] = new_x + self.g[self.next_state_idx] = new_g + self.next_state_idx += 1 + + self.num_components_accepted += num_components_accepted + self.num_components_proposed += num_components_proposed + self.num_subset_accepts += int(subset_accepted) + self.num_subset_trials += 1 + self.num_outer_moves += int(outer_moved) + self.jump_distance_sum += jump_distance + self.num_fine_evals += int(fine_evaluated) + + def finalise(self, *, retain_jump_distances: bool = False) -> MarkovChainResult: + """Create a result from the completed chain.""" + + if not self.is_complete: + raise RuntimeError("Cannot finalise an incomplete chain.") + + jump_distances = rms_jump_distances(self.x[np.newaxis, ...])[0] + + return MarkovChainResult( + x=self.x, + g=self.g, + num_components_accepted=(self.num_components_accepted), + num_components_proposed=(self.num_components_proposed), + num_subset_accepts=(self.num_subset_accepts), + num_subset_trials=self.num_transitions, + num_outer_moves=self.num_outer_moves, + jump_distance_sum=self.jump_distance_sum, + num_fine_evals=self.num_fine_evals, + num_coarse_evals=self.num_coarse_evals, + jump_distances=(jump_distances if retain_jump_distances else None), + ) + + +@dataclass +class _DAChainAccumulator(_ChainAccumulator): + """Mutable internal storage for one delayed-acceptance chain.""" + + gc: np.ndarray | None = None + + num_coarse_accepts: int = 0 + num_coarse_proposals: int = 0 + num_endpoint_moves: int = 0 + num_fine_subset_passes: int = 0 + num_fine_correction_accepts: int = 0 + num_coarse_and_fine_subset_passes: int = 0 + + def append_DA_transition( + self, + *, + new_x, + new_g, + new_gc, + num_components_accepted, + num_components_proposed, + num_coarse_accepts, + num_coarse_proposals, + endpoint_moved, + fine_evaluated, + fine_subset_pass, + fine_correction_accept, + coarse_subset_pass, + ) -> None: + """Store one completed delayed-acceptance transition.""" + + if self.is_complete: + raise ValueError("Cannot append to a completed chain.") + + current_x = self.current_x + + outer_moved = not np.array_equal( + new_x, + current_x, + ) + jump_distance = np.linalg.norm(new_x - current_x) / np.sqrt(current_x.size) + + self.x[self.next_state_idx] = new_x + self.g[self.next_state_idx] = new_g + self.gc[self.next_state_idx] = new_gc + self.next_state_idx += 1 + + self.num_components_accepted += num_components_accepted + self.num_components_proposed += num_components_proposed + + self.num_subset_accepts += int(fine_subset_pass) + self.num_subset_trials += int(fine_evaluated) + self.num_outer_moves += int(outer_moved) + self.jump_distance_sum += jump_distance + + self.num_fine_evals += int(fine_evaluated) + self.num_coarse_evals += num_coarse_proposals + self.num_coarse_accepts += num_coarse_accepts + self.num_coarse_proposals += num_coarse_proposals + self.num_endpoint_moves += int(endpoint_moved) + self.num_fine_subset_passes += int(fine_subset_pass) + self.num_fine_correction_accepts += int(fine_correction_accept) + self.num_coarse_and_fine_subset_passes += int( + coarse_subset_pass and fine_subset_pass + ) + + def finalise( + self, + *, + retain_jump_distances=False, + ) -> DAMarkovChainResult: + """Create a result from a completed DA chain.""" + + if not self.is_complete: + raise RuntimeError("Cannot finalise an incomplete chain.") + + jump_distances = rms_jump_distances(self.x[np.newaxis, ...])[0] + + return DAMarkovChainResult( + x=self.x, + g=self.g, + gc=self.gc, + num_components_accepted=(self.num_components_accepted), + num_components_proposed=(self.num_components_proposed), + num_subset_accepts=(self.num_subset_accepts), + num_subset_trials=self.num_subset_trials, + num_outer_moves=self.num_outer_moves, + jump_distance_sum=self.jump_distance_sum, + num_fine_evals=self.num_fine_evals, + num_coarse_evals=self.num_coarse_evals, + num_coarse_accepts=self.num_coarse_accepts, + num_coarse_proposals=(self.num_coarse_proposals), + num_endpoint_moves=self.num_endpoint_moves, + num_fine_subset_passes=(self.num_fine_subset_passes), + num_fine_correction_accepts=(self.num_fine_correction_accepts), + num_coarse_and_fine_subset_passes=(self.num_coarse_and_fine_subset_passes), + jump_distances=(jump_distances if retain_jump_distances else None), + ) + + +@dataclass(eq=False, repr=False) +class MarkovChainResult: + """Result of generating one completed conditional Markov chain.""" + + #: Stored chain states, with shape ``(states, dimensions)``. + x: np.ndarray + + #: Fine performance values associated with the stored states. + g: np.ndarray + + #: Total number of accepted MMH component proposals. + num_components_accepted: int + + #: Total number of attempted MMH component proposals. + num_components_proposed: int + + #: Number of candidates passing the fine-model subset condition. + num_subset_accepts: int + + #: Number of candidates tested against the fine-model subset condition. + num_subset_trials: int + + #: Number of stored outer-chain transitions that change state. + num_outer_moves: int + + #: Sum of RMS distances between adjacent stored chain states. + jump_distance_sum: float + + #: Number of fine performance-function evaluations. + num_fine_evals: int + + #: Number of coarse performance-function evaluations. + num_coarse_evals: int = 0 + + #: Per-transition RMS jump distances, if retained. + jump_distances: np.ndarray | None = None + + def __post_init__(self) -> None: + if self.x.ndim != 2: + raise ValueError( + "x must have shape (states, dimensions); " f"got shape {self.x.shape}." + ) + + if self.g.ndim != 1: + raise ValueError("g must have shape (states,); " f"got shape {self.g.shape}.") + + if self.x.shape[0] != self.g.shape[0]: + raise ValueError( + "x and g must contain the same number of states; " + f"got {self.x.shape[0]} and {self.g.shape[0]}." + ) + + if self.jump_distances is not None and self.jump_distances.shape != ( + self.num_transitions, + ): + raise ValueError( + "jump_distances must have shape " + f"({self.num_transitions},); got " + f"{self.jump_distances.shape}." + ) + + @property + def num_states(self) -> int: + """Number of stored states, including the initial chain seed.""" + return self.x.shape[0] + + @property + def dimension(self) -> int: + """Dimension of each Markov-chain state.""" + return self.x.shape[1] + + @property + def num_transitions(self) -> int: + """Number of attempted outer-chain transitions.""" + return max(self.num_states - 1, 0) + + @property + def component_acceptance_rate(self) -> float: + """Return the average MMH component acceptance rate. + + The rate includes component proposals belonging to candidate states + that subsequently fail the Subset Simulation condition. + """ + if self.num_components_proposed == 0: + return np.nan + + return self.num_components_accepted / self.num_components_proposed + + @property + def subset_acceptance_rate(self) -> float: + """Return the fraction of candidates passing the fine subset condition.""" + if self.num_subset_trials == 0: + return np.nan + + return self.num_subset_accepts / self.num_subset_trials + + @property + def outer_move_rate(self) -> float: + """Return the fraction of stored outer transitions that change state.""" + if self.num_transitions == 0: + return np.nan + + return self.num_outer_moves / self.num_transitions + + @property + def mean_jump_distance(self) -> float: + """Return the mean RMS distance between adjacent stored states. + + Rejected outer transitions contribute a zero jump distance. + """ + if self.num_transitions == 0: + return np.nan + + return self.jump_distance_sum / self.num_transitions + + def __eq__(self, other: object) -> bool: + """Compare chain results, including their NumPy arrays.""" + if not isinstance(other, MarkovChainResult): + return NotImplemented + + if type(self) is not type(other): + return False + + return all( + _values_equal(getattr(self, field.name), getattr(other, field.name)) + for field in fields(self) + ) + + def __repr__(self) -> str: + """Return a compact representation that omits stored arrays.""" + return ( + f"{type(self).__name__}(\n" + f" num_states={self.num_states},\n" + f" dimension={self.dimension},\n" + f" num_transitions={self.num_transitions},\n" + f" component_acceptance_rate={self.component_acceptance_rate!r},\n" + f" subset_acceptance_rate={self.subset_acceptance_rate!r},\n" + f" outer_move_rate={self.outer_move_rate!r},\n" + f" mean_jump_distance={self.mean_jump_distance!r},\n" + f" num_fine_evals={self.num_fine_evals},\n" + f" num_coarse_evals={self.num_coarse_evals},\n" + f")" + ) + + +@dataclass(eq=False, repr=False) +class DAMarkovChainResult(MarkovChainResult): + """Result of generating one delayed-acceptance Markov chain.""" + + gc: np.ndarray | None = None + + num_coarse_accepts: int = 0 + num_coarse_proposals: int = 0 + num_endpoint_moves: int = 0 + num_fine_subset_passes: int = 0 + num_fine_correction_accepts: int = 0 + num_coarse_and_fine_subset_passes: int = 0 + + @property + def coarse_acceptance_rate(self) -> float: + """Fraction of inner coarse-MH proposals accepted.""" + if self.num_coarse_proposals == 0: + return np.nan + return self.num_coarse_accepts / self.num_coarse_proposals + + @property + def endpoint_move_rate(self) -> float: + """Fraction of coarse subchains whose endpoint moved.""" + if self.num_transitions == 0: + return np.nan + return self.num_endpoint_moves / self.num_transitions + + @property + def fine_eval_rate(self) -> float: + """Fraction of outer transitions requiring a fine evaluation.""" + if self.num_transitions == 0: + return np.nan + return self.num_fine_evals / self.num_transitions + + @property + def fine_subset_pass_rate(self) -> float: + """Fraction of fine-evaluated endpoints in the fine subset.""" + if self.num_fine_evals == 0: + return np.nan + return self.num_fine_subset_passes / self.num_fine_evals + + @property + def fine_correction_acceptance_rate(self) -> float: + """Fraction of fine-subset endpoints accepted by DA correction.""" + if self.num_fine_subset_passes == 0: + return np.nan + return self.num_fine_correction_accepts / self.num_fine_subset_passes + + @property + def coarse_given_fine(self) -> float: + """Fraction of fine-subset endpoints also in the coarse subset.""" + if self.num_fine_subset_passes == 0: + return np.nan + return self.num_coarse_and_fine_subset_passes / self.num_fine_subset_passes diff --git a/matflow/tests/subset_simulation.py b/matflow/subset_simulation/subset_simulation.py similarity index 71% rename from matflow/tests/subset_simulation.py rename to matflow/subset_simulation/subset_simulation.py index 3906b73c..8163139e 100644 --- a/matflow/tests/subset_simulation.py +++ b/matflow/subset_simulation/subset_simulation.py @@ -1,7 +1,10 @@ """Module containing functions to run a subset simulation on a simple toy model, used as a validation of the MatFlow implementation.""" +from __future__ import annotations + import copy + from datetime import datetime import pickle from pathlib import Path @@ -15,13 +18,19 @@ from hpcflow.sdk.log import TimeIt import matflow as mf +from matflow.subset_simulation.markov_chain import ( + _ChainAccumulator, + _DAChainAccumulator, + DAMarkovChainResult, + MarkovChainResult, +) + from .subset_simulation_result import ( LevelSamplingResult, DALevelSamplingResult, ACSLevelSamplingResult, SubsetSimulationResult, - rms_jump_distances, ) @@ -208,8 +217,6 @@ def generate_next_level_samples( dimension, chain_seeds, chain_g, - all_x, - all_g, level_idx, master_seed, threshold, @@ -217,66 +224,56 @@ def generate_next_level_samples( transformation: Callable | None = None, debug: bool = False, ) -> LevelSamplingResult: + """Generate one conditional Subset Simulation level using MMH.""" - subset_accept_arr = np.zeros((num_chains, num_states - 1)).astype(bool) - component_accept_arr = np.zeros((num_chains, num_states - 1)) + chain_results: list[MarkovChainResult] = [] for chain_index in range(num_chains): - # proceed this Markov chain until all states have been generated - all_x[chain_index, 0] = chain_seeds[chain_index] - all_g[chain_index, 0] = chain_g[chain_index] + chain_x = np.full((num_states, dimension), np.nan) + chain_g_values = np.full(num_states, np.nan) - chain_rng = None - for state_idx in range(1, num_states): + chain_x[0] = chain_seeds[chain_index] + chain_g_values[0] = chain_g[chain_index] - # RNG seed sequence for Markov chains: - if state_idx == 1: - # spawn key to match the task ID in the matflow workflow - spawn_key = (4, level_idx, chain_index) - chain_rng = np.random.default_rng( - np.random.SeedSequence(master_seed, spawn_key=spawn_key) - ) + chain = _ChainAccumulator(x=chain_x, g=chain_g_values) - x = all_x[chain_index, state_idx - 1] - g = all_g[chain_index, state_idx - 1] + # Preserve the existing MatFlow-compatible RNG mapping. + spawn_key = ( + 4, + level_idx, + chain_index, + ) + chain_rng = np.random.default_rng( + np.random.SeedSequence(master_seed, spawn_key=spawn_key) + ) - trial_x, mcmc_accept_rate = generate_next_state( - x=x, + while not chain.is_complete: + current_x = chain.current_x + + (trial_x, component_acceptance_rate,) = generate_next_state( + x=current_x, proposal=proposal, rng=chain_rng, ) - component_accept_arr[chain_index, state_idx - 1] = mcmc_accept_rate + + # Preferably generate_next_state() would return this + # integer count directly. + num_components_accepted = int(np.rint(component_acceptance_rate * dimension)) + trial_x_t = transformation(trial_x) if transformation else trial_x trial_g = performance(trial_x_t) - current_x = x - current_g = g - is_ss_accept = trial_g > threshold - subset_accept_arr[chain_index, state_idx - 1] = is_ss_accept - new_x = trial_x if is_ss_accept else current_x - new_g = trial_g if is_ss_accept else current_g + chain.append_transition( + trial_x=trial_x, + trial_g=trial_g, + threshold=threshold, + num_components_accepted=(num_components_accepted), + ) - all_x[chain_index, state_idx] = new_x - all_g[chain_index, state_idx] = new_g + chain_results.append(chain.finalise(retain_jump_distances=debug)) - subset_accept = np.mean(subset_accept_arr).item() - num_fine_evals = num_chains * (num_states - 1) - component_acceptance_rate = np.mean(component_accept_arr).item() - jump_distances = rms_jump_distances(all_x) - mean_jump_distance = np.mean(jump_distances).item() - outer_move_rate = np.mean(jump_distances > 0).item() - - return LevelSamplingResult( - x=all_x, - g=all_g, - component_acceptance_rate=component_acceptance_rate, - subset_acceptance_rate=subset_accept, - mean_jump_distance=mean_jump_distance, - jump_distances=jump_distances, - outer_move_rate=outer_move_rate, - num_fine_evals=num_fine_evals, - ) + return LevelSamplingResult(chains=tuple(chain_results)) def generate_next_level_samples_CS( @@ -286,8 +283,6 @@ def generate_next_level_samples_CS( dimension, chain_seeds, chain_g, - all_x, - all_g, level_idx, master_seed, threshold, @@ -295,66 +290,44 @@ def generate_next_level_samples_CS( transformation: Callable | None = None, debug: bool = False, ) -> LevelSamplingResult: - """Conditional sampling algorithm for generating states in the subset simulation level - (aka subset infinity). + """Generate one conditional level using conditional sampling (aka subset infinity).""" + + chain_results: list[MarkovChainResult] = [] - """ - subset_accept_arr = np.zeros((num_chains, num_states - 1)).astype(bool) for chain_index in range(num_chains): + chain_x = np.full((num_states, dimension), np.nan) + chain_g_values = np.full(num_states, np.nan) - # proceed this Markov chain until all states have been generated - all_x[chain_index, 0] = chain_seeds[chain_index] - all_g[chain_index, 0] = chain_g[chain_index] + chain_x[0] = chain_seeds[chain_index] + chain_g_values[0] = chain_g[chain_index] - chain_rng = None - for state_idx in range(1, num_states): + chain = _ChainAccumulator(x=chain_x, g=chain_g_values) - # RNG seed sequence for Markov chains: - if state_idx == 1: - # spawn key to match the task ID in the matflow workflow - spawn_key = (4, level_idx, chain_index) - chain_rng = np.random.default_rng( - np.random.SeedSequence(master_seed, spawn_key=spawn_key) - ) - - x = all_x[chain_index, state_idx - 1] - g = all_g[chain_index, state_idx - 1] + spawn_key = (4, level_idx, chain_index) + chain_rng = np.random.default_rng( + np.random.SeedSequence(master_seed, spawn_key=spawn_key) + ) + while not chain.is_complete: trial_x = generate_next_state_CS( - x=x, - prop_std=prop_std, - rng=chain_rng, + x=chain.current_x, prop_std=prop_std, rng=chain_rng ) + trial_x_t = transformation(trial_x) if transformation else trial_x trial_g = performance(trial_x_t) - current_x = x - current_g = g + chain.append_transition( + trial_x=trial_x, + trial_g=trial_g, + threshold=threshold, + num_components_accepted=0, # CS has no component-wise MMH stage. + num_components_proposed=0, + fine_evaluated=True, + ) - is_accept = trial_g > threshold - subset_accept_arr[chain_index, state_idx - 1] = is_accept - new_x = trial_x if is_accept else current_x - new_g = trial_g if is_accept else current_g + chain_results.append(chain.finalise(retain_jump_distances=debug)) - all_x[chain_index, state_idx] = new_x - all_g[chain_index, state_idx] = new_g - - subset_accept = np.mean(subset_accept_arr).item() - num_fine_evals = num_chains * (num_states - 1) - jump_distances = rms_jump_distances(all_x) - mean_jump_distance = np.mean(jump_distances).item() - outer_move_rate = np.mean(jump_distances > 0).item() - - return LevelSamplingResult( - x=all_x, - g=all_g, - component_acceptance_rate=1.0, - subset_acceptance_rate=subset_accept, - mean_jump_distance=mean_jump_distance, - jump_distances=jump_distances, - outer_move_rate=outer_move_rate, - num_fine_evals=num_fine_evals, - ) + return LevelSamplingResult(chains=tuple(chain_results)) def generate_next_level_samples_ACS( @@ -364,8 +337,6 @@ def generate_next_level_samples_ACS( dimension, chain_seeds, chain_g, - all_x, - all_g, level_idx, master_seed, threshold, @@ -375,8 +346,8 @@ def generate_next_level_samples_ACS( lambda_=1.0, debug: bool = False, ) -> ACSLevelSamplingResult: - """Adaptive conditional sampling algorithm for generating states in the subset - simulation level (aka adaptive subset infinity). + """Generate one conditional level using adaptive conditional sampling (aka adaptive + subset infinity). Parameters ---------- @@ -391,80 +362,79 @@ def generate_next_level_samples_ACS( """ - A_STAR = 0.44 + target_acceptance_rate = 0.44 num_chains_per_update = int(chains_per_update * num_chains) - assert float(num_chains_per_update) == chains_per_update * num_chains - num_batches = int(num_chains / num_chains_per_update) - batch_avgs = [] # mean acceptance for each batch + if num_chains_per_update != chains_per_update * num_chains: + raise ValueError("chains_per_update * num_chains must be an integer.") + + if num_chains_per_update < 1: + raise ValueError("chains_per_update must select at least one chain.") + + if num_chains % num_chains_per_update: + raise ValueError("num_chains must be divisible by num_chains_per_update.") + + num_batches = num_chains // num_chains_per_update + + chain_results: list[MarkovChainResult] = [] for batch_idx in range(num_batches): + sigma = np.minimum(1.0, lambda_ * prop_std) + rho = np.sqrt(1.0 - sigma**2) - sigma = np.minimum(1, lambda_ * prop_std) - rho = np.sqrt(1 - sigma**2) + batch_results: list[MarkovChainResult] = [] - is_accept_arr = np.zeros((num_chains_per_update, num_states - 1)).astype(bool) for batch_chain_idx in range(num_chains_per_update): + chain_index = batch_idx * num_chains_per_update + batch_chain_idx - chain_index = batch_chain_idx + (batch_idx * num_chains_per_update) + chain_x = np.full((num_states, dimension), np.nan) + chain_g_values = np.full(num_states, np.nan) - # proceed this Markov chain until all states have been generated - all_x[chain_index, 0] = chain_seeds[chain_index] - all_g[chain_index, 0] = chain_g[chain_index] + chain_x[0] = chain_seeds[chain_index] + chain_g_values[0] = chain_g[chain_index] - chain_rng = None - for state_idx in range(1, num_states): + chain = _ChainAccumulator(x=chain_x, g=chain_g_values) - # RNG seed sequence for Markov chains: - if state_idx == 1: - # spawn key to match the task ID in the matflow workflow - spawn_key = (4, level_idx, chain_index) - chain_rng = np.random.default_rng( - np.random.SeedSequence(master_seed, spawn_key=spawn_key) - ) + spawn_key = (4, level_idx, chain_index) + chain_rng = np.random.default_rng( + np.random.SeedSequence(master_seed, spawn_key=spawn_key) + ) - x = all_x[chain_index, state_idx - 1] - g = all_g[chain_index, state_idx - 1] + while not chain.is_complete: + trial_x = norm.rvs( + loc=chain.current_x * rho, + scale=sigma, + random_state=chain_rng, + ) - trial_x = norm.rvs(loc=x * rho, scale=sigma, random_state=chain_rng) trial_x_t = transformation(trial_x) if transformation else trial_x trial_g = performance(trial_x_t) - current_x = x - current_g = g - - is_accept = trial_g > threshold - is_accept_arr[batch_chain_idx, state_idx - 1] = is_accept - - new_x = trial_x if is_accept else current_x - new_g = trial_g if is_accept else current_g - - all_x[chain_index, state_idx] = new_x - all_g[chain_index, state_idx] = new_g + chain.append_transition( + trial_x=trial_x, + trial_g=trial_g, + threshold=threshold, + num_components_accepted=0, + num_components_proposed=0, + fine_evaluated=True, + ) - accept_batch_avg = np.mean(is_accept_arr) - batch_avgs.append(accept_batch_avg) + chain_result = chain.finalise(retain_jump_distances=debug) + batch_results.append(chain_result) + chain_results.append(chain_result) - zeta = 1 / np.sqrt(batch_idx + 1) - lambda_ *= np.exp(zeta * (accept_batch_avg - A_STAR)) + num_batch_accepts = sum(chain.num_subset_accepts for chain in batch_results) + num_batch_trials = sum(chain.num_subset_trials for chain in batch_results) - subset_accept = np.mean(batch_avgs).item() - num_fine_evals = num_chains * (num_states - 1) + accept_batch_avg = ( + num_batch_accepts / num_batch_trials if num_batch_trials > 0 else np.nan + ) - jump_distances = rms_jump_distances(all_x) - mean_jump_distance = np.mean(jump_distances).item() - outer_move_rate = np.mean(jump_distances > 0).item() + zeta = 1.0 / np.sqrt(batch_idx + 1) + lambda_ *= np.exp(zeta * (accept_batch_avg - target_acceptance_rate)) - return ACSLevelSamplingResult( - lambda_=lambda_, - mcmc_acceptance_rate=1.0, - subset_acceptance_rate=subset_accept, - mean_jump_distance=mean_jump_distance, - jump_distances=jump_distances, - outer_move_rate=outer_move_rate, - num_fine_evals=num_fine_evals, - ) + return ACSLevelSamplingResult(chains=tuple(chain_results), lambda_=lambda_) def weakest_link_coarse_gradient_xt(x_t, group_idx): @@ -588,8 +558,6 @@ def generate_next_level_samples_DA( dimension, chain_seeds, chain_g, - all_x, - all_g, level_idx, master_seed, threshold, @@ -598,10 +566,10 @@ def generate_next_level_samples_DA( transformation: Callable | None = None, temperature=1.0, num_inner_states: int = 1, - spawn_key: tuple[int] | None = None, + spawn_key: tuple[int, ...] | None = None, debug: bool = False, ) -> DALevelSamplingResult: - """Fixed-length subchain surrogate transition for Subset Simulation. + """Generate one conditional level using fixed-subchain RST. This is the randomised-length subchain surrogate transition (RST) algorithm but with a fixed length subchain. For num_inner_states=1, this should be identical to @@ -633,71 +601,30 @@ def generate_next_level_samples_DA( The endpoint is then corrected to the fine target. """ - num_outer_trials = num_chains * (num_states - 1) - - # ------------------------------------------------------------ - # Per-outer-transition diagnostics - # ------------------------------------------------------------ - - # Fraction of inner coarse-MH proposals accepted. - inner_accept_rate_arr = np.zeros((num_chains, num_states - 1), dtype=float) - - # Number of accepted coarse moves within each inner subchain. - inner_accept_count_arr = np.zeros((num_chains, num_states - 1), dtype=int) - - # Whether the final endpoint differs from the outer current state. - endpoint_move_arr = np.zeros((num_chains, num_states - 1), dtype=bool) - - # Whether the expensive fine model was evaluated. - fine_eval_arr = np.zeros((num_chains, num_states - 1), dtype=bool) - - # Whether the endpoint passed the fine subset condition. - fine_subset_pass_arr = np.zeros((num_chains, num_states - 1), dtype=bool) - - # Whether the final fine correction accepted the endpoint. - fine_accept_arr = np.zeros((num_chains, num_states - 1), dtype=bool) - - # Whether the accepted endpoint is above the coarse threshold. - endpoint_coarse_subset_pass_arr = np.zeros((num_chains, num_states - 1), dtype=bool) - - # Final coarse surrogate acceptance log probability. - fine_log_alpha_arr = np.full((num_chains, num_states - 1), np.nan, dtype=float) - - # ------------------------------------------------------------ - # Store coarse performance at outer states - # ------------------------------------------------------------ - - all_gc = np.full((num_chains, num_states), np.nan, dtype=float) - total_mmh_component_acceptance = 0.0 - - debug_data = {} - if debug: - debug_data["chain_data"] = [] + chain_results: list[DAMarkovChainResult] = [] + debug_data = {"chain_data": []} if debug else None for chain_index in range(num_chains): - if debug: - debug_data["chain_data"].append({"state_data": []}) + seed_x = np.asarray(chain_seeds[chain_index]) - # -------------------------------------------------------- - # Initial state - # -------------------------------------------------------- + chain_x = np.full((num_states, dimension), np.nan) + chain_g_values = np.full(num_states, np.nan) + chain_gc = np.full(num_states, np.nan) - seed_x = chain_seeds[chain_index] - - all_x[chain_index, 0] = seed_x - all_g[chain_index, 0] = chain_g[chain_index] + chain_x[0] = seed_x + chain_g_values[0] = chain_g[chain_index] seed_x_t = transformation(seed_x) if transformation else seed_x + chain_gc[0] = performance_coarse(seed_x_t) - seed_gc = performance_coarse(seed_x_t) - all_gc[chain_index, 0] = seed_gc - - # -------------------------------------------------------- - # RNG for this chain - # -------------------------------------------------------- + chain = _DAChainAccumulator( + x=chain_x, + g=chain_g_values, + gc=chain_gc, + num_coarse_evals=1, # initial coarse evaluation of the seed. + ) - # `4` is usually the task insert ID of the generate_next_state task: - spawn_key_ = tuple([*(spawn_key or (4,)), level_idx, chain_index]) + spawn_key_ = (*(spawn_key or (4,)), level_idx, chain_index) chain_rng = np.random.default_rng( np.random.SeedSequence( master_seed, @@ -705,32 +632,19 @@ def generate_next_level_samples_DA( ) ) - # -------------------------------------------------------- - # Generate outer states - # -------------------------------------------------------- + chain_debug_data = {"state_data": []} if debug else None - for state_idx in range(1, num_states): - - if debug: - debug_data["chain_data"][chain_index]["state_data"].append({}) - debug_dat_cs_ij = debug_data["chain_data"][chain_index]["state_data"][ - state_idx - 1 - ] - - current_x = all_x[chain_index, state_idx - 1] - current_g = all_g[chain_index, state_idx - 1] - current_gc = all_gc[chain_index, state_idx - 1] - - # ---------------------------------------------------- - # Run the coarse subchain - # ---------------------------------------------------- + while not chain.is_complete: + current_x = chain.current_x + current_g = chain.current_g + current_gc = float(chain.gc[chain.next_state_idx - 1]) ( psi, psi_gc, - n_inner_accepts, - mmh_component_acceptance_sum, - sub_chain_debug_data, + num_inner_accepts, + component_acceptance_sum, + subchain_debug_data, ) = generate_coarse_subchain( x=current_x, gc=current_gc, @@ -745,108 +659,54 @@ def generate_next_level_samples_DA( chain_idx=chain_index, debug=debug, ) - total_mmh_component_acceptance += mmh_component_acceptance_sum - - if debug: - debug_dat_cs_ij["generate_coarse_subchain_data"] = sub_chain_debug_data - debug_dat_cs_ij["psi"] = psi - inner_accept_count_arr[chain_index, state_idx - 1] = n_inner_accepts - - inner_accept_rate_arr[chain_index, state_idx - 1] = ( - n_inner_accepts / num_inner_states if num_inner_states > 0 else np.nan - ) - - # ---------------------------------------------------- - # Did the coarse subchain actually move? - # ---------------------------------------------------- + num_components_proposed = num_inner_states * dimension + num_components_accepted = int(np.rint(component_acceptance_sum * dimension)) endpoint_moved = not np.array_equal(psi, current_x) - endpoint_move_arr[chain_index, state_idx - 1] = endpoint_moved + fine_evaluated = False + fine_subset_pass = False + fine_correction_accept = False + coarse_subset_pass = False + fine_log_alpha = np.nan + psi_g = None if not endpoint_moved: - - # The subchain ended where it started. - # No fine evaluation is necessary. new_x = current_x new_g = current_g new_gc = current_gc - if debug: - debug_dat_cs_ij["psi_g"] = None - else: - - # ------------------------------------------------ - # Endpoint coarse diagnostics - # ------------------------------------------------ - - endpoint_coarse_subset_pass_arr[chain_index, state_idx - 1] = ( - psi_gc > threshold - ) - - # ------------------------------------------------ - # Fine evaluation - # ------------------------------------------------ - - fine_eval_arr[chain_index, state_idx - 1] = True + coarse_subset_pass = psi_gc > threshold + fine_evaluated = True psi_t = transformation(psi) if transformation else psi - psi_g = performance(psi_t) - - if debug: - debug_dat_cs_ij["psi_g"] = psi_g - - # ------------------------------------------------ - # Fine subset test - # ------------------------------------------------ - fine_subset_pass = psi_g > threshold - fine_subset_pass_arr[chain_index, state_idx - 1] = fine_subset_pass if not fine_subset_pass: - - # The fine target is zero here. new_x = current_x new_g = current_g new_gc = current_gc else: - - # ------------------------------------------------ - # Final RST correction - # - # alpha_F = - # min(1, pi_C(current_x) / pi_C(psi)) - # - # The phi terms cancel, leaving: - # - # alpha_F = - # min(1, s(current_x) / s(psi)) - # ------------------------------------------------ - log_s_current = log_surrogate_weight( current_gc, threshold, temperature, ) - log_s_psi = log_surrogate_weight( psi_gc, threshold, temperature, ) - log_alpha_fine = min(0.0, log_s_current - log_s_psi) - random_num = chain_rng.random() - fine_log_alpha_arr[chain_index, state_idx - 1] = log_alpha_fine - fine_accept = np.log(random_num) < log_alpha_fine + fine_log_alpha = min(0.0, log_s_current - log_s_psi) - fine_accept_arr[chain_index, state_idx - 1] = fine_accept + fine_correction_accept = np.log(chain_rng.random()) < fine_log_alpha - if fine_accept: + if fine_correction_accept: new_x = psi new_g = psi_g new_gc = psi_gc @@ -855,110 +715,41 @@ def generate_next_level_samples_DA( new_g = current_g new_gc = current_gc - if debug: - debug_dat_cs_ij["new_x"] = new_x - debug_dat_cs_ij["new_g"] = new_g - debug_dat_cs_ij["new_gc"] = new_gc - - # ---------------------------------------------------- - # Store outer state - # ---------------------------------------------------- - - all_x[chain_index, state_idx] = new_x - all_g[chain_index, state_idx] = new_g - all_gc[chain_index, state_idx] = new_gc - - # ============================================================ - # Aggregate diagnostics - # ============================================================ - - n_inner_proposals = num_chains * (num_states - 1) * num_inner_states - - n_inner_accepts = int(inner_accept_count_arr.sum()) - n_fine_evals = int(fine_eval_arr.sum()) - n_fine_subset_pass = int(fine_subset_pass_arr.sum()) - n_fine_accepts = int(fine_accept_arr.sum()) - n_endpoint_moves = int(endpoint_move_arr.sum()) - - component_acceptance_rate = ( - total_mmh_component_acceptance / n_inner_proposals - if n_inner_proposals > 0 - else np.nan - ).item() - - coarse_acceptance_rate = ( - n_inner_accepts / n_inner_proposals if n_inner_proposals > 0 else np.nan - ) - - endpoint_move_rate = ( - n_endpoint_moves / num_outer_trials if num_outer_trials > 0 else np.nan - ) - - fine_eval_rate = n_fine_evals / num_outer_trials if num_outer_trials > 0 else np.nan - - # Of the endpoints actually evaluated with the fine model, how many satisfy the fine - # subset condition? - fine_subset_pass_rate = ( - n_fine_subset_pass / n_fine_evals if n_fine_evals > 0 else np.nan - ) + chain.append_DA_transition( + new_x=new_x, + new_g=new_g, + new_gc=new_gc, + num_components_accepted=(num_components_accepted), + num_components_proposed=(num_components_proposed), + num_coarse_accepts=num_inner_accepts, + num_coarse_proposals=num_inner_states, + endpoint_moved=endpoint_moved, + fine_evaluated=fine_evaluated, + fine_subset_pass=fine_subset_pass, + fine_correction_accept=(fine_correction_accept), + coarse_subset_pass=coarse_subset_pass, + ) - # Of the endpoints in the fine subset, how many pass the final delayed-acceptance - # correction? - fine_correction_acceptance_rate = ( - n_fine_accepts / n_fine_subset_pass if n_fine_subset_pass > 0 else np.nan - ) + if debug: + chain_debug_data["state_data"].append( + { + "psi": psi, + "psi_gc": psi_gc, + "psi_g": psi_g, + "endpoint_moved": endpoint_moved, + "fine_subset_pass": fine_subset_pass, + "fine_correction_accept": (fine_correction_accept), + "fine_log_alpha": fine_log_alpha, + "subchain_data": (subchain_debug_data), + } + ) - # Fraction of outer transitions that actually change the stored state. - outer_move_rate = ( - n_fine_accepts / num_outer_trials if num_outer_trials > 0 else np.nan - ) + chain_results.append(chain.finalise(retain_jump_distances=debug)) - # Among endpoints in the fine subset, how many are also above the coarse threshold? - coarse_given_fine = ( - np.sum(endpoint_coarse_subset_pass_arr & fine_subset_pass_arr) - / n_fine_subset_pass - if n_fine_subset_pass > 0 - else np.nan - ) + if debug: + debug_data["chain_data"].append(chain_debug_data) - jump_distances = rms_jump_distances(all_x) - mean_jump_distance = np.mean(jump_distances).item() - - # One coarse evaluation for each chain seed plus one per# inner proposal. - num_coarse_evals = num_chains + n_inner_proposals - - outer_move_rate_2 = np.mean(jump_distances > 0).item() - assert outer_move_rate == outer_move_rate_2 - - return DALevelSamplingResult( - x=all_x, - g=all_g, - component_acceptance_rate=component_acceptance_rate, - subset_acceptance_rate=fine_subset_pass_rate, - mean_jump_distance=mean_jump_distance, - num_fine_evals=n_fine_evals, - num_coarse_evals=num_coarse_evals, - jump_distances=jump_distances if debug else None, - debug_data=debug_data if debug else None, - coarse_acceptance_rate=coarse_acceptance_rate, - endpoint_move_rate=endpoint_move_rate, - fine_eval_rate=fine_eval_rate, - fine_subset_pass_rate=fine_subset_pass_rate, - fine_correction_acceptance_rate=fine_correction_acceptance_rate, - outer_move_rate=outer_move_rate, - coarse_given_fine=coarse_given_fine, - all_gc=all_gc, - coarse_acceptance_rates=(inner_accept_rate_arr if debug else None), - coarse_acceptance_counts=(inner_accept_count_arr if debug else None), - endpoint_moves=(endpoint_move_arr if debug else None), - fine_evals=(fine_eval_arr if debug else None), - fine_subset_passes=(fine_subset_pass_arr if debug else None), - fine_correction_accepts=(fine_accept_arr if debug else None), - fine_log_alpha=(fine_log_alpha_arr if debug else None), - endpoint_coarse_subset_passes=( - endpoint_coarse_subset_pass_arr if debug else None - ), - ) + return DALevelSamplingResult(chains=tuple(chain_results), debug_data=debug_data) def generate_next_level_samples_DA_single_inner( @@ -1280,23 +1071,12 @@ def subset_simulation( # Generate simulation level level_idx + 1. # -------------------------------------------------------------- - all_x = np.full( - (num_chains, num_states, dimension), - np.nan, - ) - all_g = np.full( - (num_chains, num_states), - np.nan, - ) - level_result = sampling_method( num_chains=num_chains, num_states=num_states, dimension=dimension, chain_seeds=chain_seeds, chain_g=chain_g, - all_x=all_x, - all_g=all_g, level_idx=level_idx, master_seed=master_seed, threshold=threshold, @@ -1307,8 +1087,7 @@ def subset_simulation( if not isinstance(level_result, LevelSamplingResult): raise TypeError( - "sampling_method must return a " - "LevelSamplingResult instance; " + "sampling_method must return a LevelSamplingResult instance; " f"got {type(level_result).__name__}" ) diff --git a/matflow/subset_simulation/subset_simulation_result.py b/matflow/subset_simulation/subset_simulation_result.py new file mode 100644 index 00000000..53ed85f5 --- /dev/null +++ b/matflow/subset_simulation/subset_simulation_result.py @@ -0,0 +1,326 @@ +from __future__ import annotations + +from dataclasses import dataclass, fields, is_dataclass +from typing import Any, ClassVar, TYPE_CHECKING + +import numpy as np +from numpy.typing import NDArray +from hpcflow.sdk.core.parameters import ParameterValue + +from matflow.tests.utils import _values_equal + +if TYPE_CHECKING: + from matflow.subset_simulation.subset_simulation import MarkovChainResult + + +from dataclasses import dataclass +from functools import cached_property +from typing import Any + +import numpy as np + + +@dataclass(repr=False) +class LevelSamplingResult: + """Result of generating one conditional Subset Simulation level.""" + + #: Results for the individual Markov chains used to generate this level. + chains: tuple[MarkovChainResult, ...] + + #: Optional algorithm-specific debugging information. + debug_data: dict[str, Any] | None = None + + def __post_init__(self) -> None: + if not self.chains: + raise ValueError("LevelSamplingResult requires at least one chain.") + + @cached_property + def x(self) -> np.ndarray: + """Sampled states, arranged as ``(chains, states, dimensions)``.""" + return np.stack([chain.x for chain in self.chains], axis=0) + + @cached_property + def g(self) -> np.ndarray: + """Fine performance values associated with the sampled states.""" + return np.stack([chain.g for chain in self.chains], axis=0) + + @property + def num_chains(self) -> int: + """Number of Markov chains used to generate the level.""" + return len(self.chains) + + @property + def num_states(self) -> int: + """Total number of stored states across all chains.""" + return sum(chain.num_states for chain in self.chains) + + @property + def num_transitions(self) -> int: + """Total number of attempted outer-chain transitions.""" + return sum(chain.num_transitions for chain in self.chains) + + @property + def num_components_accepted(self) -> int: + """Total number of MMH component proposals accepted.""" + return sum(chain.num_components_accepted for chain in self.chains) + + @property + def num_components_proposed(self) -> int: + """Total number of MMH component proposals attempted.""" + return sum(chain.num_components_proposed for chain in self.chains) + + @property + def component_acceptance_rate(self) -> float: + """Return the aggregate MMH component acceptance rate. + + The rate includes MMH component proposals belonging to candidate + states that subsequently fail the Subset Simulation condition. + """ + if self.num_components_proposed == 0: + return np.nan + return self.num_components_accepted / self.num_components_proposed + + @property + def num_subset_accepts(self) -> int: + """Total number of candidates passing the fine subset condition.""" + return sum(chain.num_subset_accepts for chain in self.chains) + + @property + def num_subset_trials(self) -> int: + """Total number of candidates tested against the fine subset condition.""" + return sum(chain.num_subset_trials for chain in self.chains) + + @property + def subset_acceptance_rate(self) -> float: + """Return the fraction of candidates passing the fine subset condition.""" + if self.num_subset_trials == 0: + return np.nan + return self.num_subset_accepts / self.num_subset_trials + + @property + def num_outer_moves(self) -> int: + """Total number of stored outer-chain transitions that change state.""" + return sum(chain.num_outer_moves for chain in self.chains) + + @property + def outer_move_rate(self) -> float: + """Return the fraction of stored outer-chain transitions that move.""" + if self.num_transitions == 0: + return np.nan + return self.num_outer_moves / self.num_transitions + + @property + def jump_distance_sum(self) -> float: + """Sum of RMS jump distances over all outer-chain transitions.""" + return sum(chain.jump_distance_sum for chain in self.chains) + + @property + def mean_jump_distance(self) -> float: + """Return the mean RMS distance between adjacent stored states. + + Rejected outer transitions contribute a zero jump distance. + """ + if self.num_transitions == 0: + return np.nan + return self.jump_distance_sum / self.num_transitions + + @cached_property + def jump_distances(self) -> np.ndarray | None: + """Return per-transition RMS jump distances, if retained. + + The resulting array has shape ``(chains, transitions)``. The value is + ``None`` if detailed jump distances were not retained for every chain. + """ + if any(chain.jump_distances is None for chain in self.chains): + return None + return np.stack([chain.jump_distances for chain in self.chains], axis=0) + + @property + def num_fine_evals(self) -> int: + """Total number of fine performance-function evaluations.""" + return sum(chain.num_fine_evals for chain in self.chains) + + @property + def num_coarse_evals(self) -> int: + """Total number of coarse performance-function evaluations.""" + return sum(chain.num_coarse_evals for chain in self.chains) + + def __repr__(self) -> str: + return ( + f"{type(self).__name__}(\n" + f" num_chains={self.num_chains},\n" + f" num_states={self.num_states},\n" + f" num_transitions={self.num_transitions},\n" + f" component_acceptance_rate={self.component_acceptance_rate!r},\n" + f" subset_acceptance_rate={self.subset_acceptance_rate!r},\n" + f" outer_move_rate={self.outer_move_rate!r},\n" + f" mean_jump_distance={self.mean_jump_distance!r},\n" + f" num_fine_evals={self.num_fine_evals},\n" + f" num_coarse_evals={self.num_coarse_evals},\n" + f")" + ) + + +@dataclass(repr=False) +class ACSLevelSamplingResult(LevelSamplingResult): + """Result of one adaptive conditional sampling level.""" + + #: Updated proposal scaling parameter. + lambda_: float = 1.0 + + +@dataclass(repr=False) +class DALevelSamplingResult(LevelSamplingResult): + """Result of one delayed-acceptance conditional level.""" + + @property + def coarse_acceptance_rate(self) -> float: + """Fraction of all inner coarse-MH proposals accepted.""" + proposed = sum(chain.num_coarse_proposals for chain in self.chains) + if proposed == 0: + return np.nan + + return sum(chain.num_coarse_accepts for chain in self.chains) / proposed + + @property + def endpoint_move_rate(self) -> float: + """Fraction of coarse subchains whose endpoints moved.""" + if self.num_transitions == 0: + return np.nan + + return ( + sum(chain.num_endpoint_moves for chain in self.chains) / self.num_transitions + ) + + @property + def fine_eval_rate(self) -> float: + """Fraction of outer transitions requiring fine evaluation.""" + if self.num_transitions == 0: + return np.nan + + return self.num_fine_evals / self.num_transitions + + @property + def fine_subset_pass_rate(self) -> float: + """Fraction of fine-evaluated endpoints in the fine subset.""" + if self.num_fine_evals == 0: + return np.nan + + return ( + sum(chain.num_fine_subset_passes for chain in self.chains) + / self.num_fine_evals + ) + + @property + def fine_correction_acceptance_rate(self) -> float: + """Fraction of fine-subset endpoints accepted by correction.""" + num_subset_passes = sum(chain.num_fine_subset_passes for chain in self.chains) + if num_subset_passes == 0: + return np.nan + + return ( + sum(chain.num_fine_correction_accepts for chain in self.chains) + / num_subset_passes + ) + + @property + def coarse_given_fine(self) -> float: + """Fraction of fine-subset endpoints also in the coarse subset.""" + num_subset_passes = sum(chain.num_fine_subset_passes for chain in self.chains) + if num_subset_passes == 0: + return np.nan + + return ( + sum(chain.num_coarse_and_fine_subset_passes for chain in self.chains) + / num_subset_passes + ) + + @cached_property + def gc(self) -> np.ndarray: + """Stored coarse values with shape ``(chains, states)``.""" + return np.stack( + [chain.gc for chain in self.chains], + axis=0, + ) + + +@dataclass(repr=False) +class SubsetSimulationResult(ParameterValue): + + _typ: ClassVar[str] = "subset_simulation_result" + + pf: float + cov: float | None + converged: bool + + thresholds: NDArray + level_covs: NDArray + levels: list[LevelSamplingResult] + + num_fine_evals: int + num_coarse_evals: int + num_failed_per_level: NDArray + + # Debugging/state data omitted from repr. + x_original: NDArray | None = None + final_chain_seeds: NDArray | None = None + final_chain_g: NDArray | None = None + final_all_x: NDArray | None = None + final_all_g: NDArray | None = None + debug_data: dict[str, Any] | None = None + + @property + def component_acceptance_rates(self) -> NDArray: + return np.array([level.component_acceptance_rate for level in self.levels]) + + @property + def subset_acceptance_rates(self) -> NDArray: + return np.array([level.subset_acceptance_rate for level in self.levels]) + + @property + def mean_jump_distances(self) -> NDArray: + return np.array([level.mean_jump_distance for level in self.levels]) + + @property + def outer_move_rates(self) -> NDArray: + return np.array([level.outer_move_rate for level in self.levels]) + + @property + def num_conditional_levels(self) -> int: + return len(self.levels) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, SubsetSimulationResult): + return NotImplemented + + if type(self) is not type(other): + return False + + return all( + _values_equal( + getattr(self, field.name), + getattr(other, field.name), + ) + for field in fields(self) + ) + + def __repr__(self) -> str: + cls_name = type(self).__name__ + cov_str = f"{self.cov:.4f}" if self.cov is not None else "None" + return ( + f"{cls_name}(\n" + f" pf={self.pf:.4e},\n" + f" cov={cov_str},\n" + f" converged={self.converged!r},\n" + f" thresholds={self.thresholds!r},\n" + f" level_covs={self.level_covs!r},\n" + f" num_conditional_levels={self.num_conditional_levels},\n" + f" component_acceptance_rates={self.component_acceptance_rates!r},\n" + f" subset_acceptance_rates={self.subset_acceptance_rates!r},\n" + f" outer_move_rates={self.outer_move_rates!r},\n" + f" mean_jump_distances={self.mean_jump_distances!r},\n" + f" num_fine_evals={self.num_fine_evals!r},\n" + f" num_coarse_evals={self.num_coarse_evals!r},\n" + f" num_failed_per_level={self.num_failed_per_level!r},\n" + f")" + ) diff --git a/matflow/tests/demo_workflows/test_demo_workflows.py b/matflow/tests/demo_workflows/test_demo_workflows.py index 47ed2f51..efbaa21d 100644 --- a/matflow/tests/demo_workflows/test_demo_workflows.py +++ b/matflow/tests/demo_workflows/test_demo_workflows.py @@ -8,7 +8,7 @@ import numpy as np import matflow as mf -from matflow.tests.subset_simulation import ( +from matflow.subset_simulation.subset_simulation import ( log_surrogate_weight, generate_next_level_samples_DA, get_approx_y_star_random_walk, @@ -66,7 +66,7 @@ def test_damask_input_files(tmp_path, save_fig, reference_array_data): @pytest.mark.demo_workflows -# @pytest.mark.skip(reason="takes too long") +@pytest.mark.skip(reason="takes too long") def test_subset_simulation_toy_model_prediction(tmp_path): """Validate the MatFlow subset simulation implementation for a toy model. @@ -161,7 +161,7 @@ def test_subset_simulation_toy_model_DA_prediction(tmp_path): weakest_link_performance_coarse, y_star=y_star, group_idx=group_idx ) - debug = subset_simulation( + result = subset_simulation( performance=performance, dimension=dimension, p_0=0.1, @@ -186,6 +186,6 @@ def test_subset_simulation_toy_model_DA_prediction(tmp_path): wk.wait() iter_i = wk.tasks.collate_results.elements[0].iterations[-1] - assert iter_i.get("outputs.pf") == debug["pf"] - assert iter_i.get("outputs.cov") == debug["cov"] - assert iter_i.get("outputs.threshold") == debug["thresholds"][-1] + assert iter_i.get("outputs.pf") == result.pf + assert iter_i.get("outputs.cov") == result.cov + assert iter_i.get("outputs.threshold") == result.thresholds[-1] diff --git a/matflow/tests/subset_simulation_result.py b/matflow/tests/subset_simulation_result.py deleted file mode 100644 index 31161fde..00000000 --- a/matflow/tests/subset_simulation_result.py +++ /dev/null @@ -1,220 +0,0 @@ -from dataclasses import dataclass, fields, is_dataclass -from typing import Any -import numpy as np -from numpy.typing import NDArray - - -def rms_jump_distances(all_x: np.ndarray) -> np.ndarray: - """RMS jump distance between consecutive stored chain states.""" - dx = np.diff(all_x, axis=1) - return np.linalg.norm(dx, axis=-1) / np.sqrt(dx.shape[-1]) - - -@dataclass -class LevelSamplingResult: - """Result of generating one conditional Subset Simulation level.""" - - #: Sampled states. - x: np.ndarray - - #: Performance function evaluations associated with each of the sampled states. - g: np.ndarray - - #: Average component acceptance rate from Modified Metropolis Hastings, inclusive of - #: proposals that subsequently fail the subset condition. - component_acceptance_rate: float - - #: Fraction of stored outer-chain transitions that change state. - outer_move_rate: float - - #: Fraction of candidate states that pass the fine-model subset condition. - subset_acceptance_rate: float - - #: Average distance between adjacent states in the outer chain (a subset rejection - #: produces a zero jump distance). - mean_jump_distance: float - - #: Number of (fine) evaluations. - num_fine_evals: int - - #: Number of coarse evaluations, if any were performed. - num_coarse_evals: int = 0 - - # optional detailed diagnostics. - jump_distances: np.ndarray | None = None - debug_data: dict[str, Any] | None = None - - -@dataclass -class ACSLevelSamplingResult(LevelSamplingResult): - lambda_: float = 1.0 - - -@dataclass -class DALevelSamplingResult(LevelSamplingResult): - """Result of generating one delayed-acceptance conditional level.""" - - #: Fraction of individual coarse-MH proposals accepted by the - #: coarse surrogate target. - coarse_acceptance_rate: float = np.nan - - #: Fraction of outer transitions for which the coarse subchain endpoint - #: differs from the current outer-chain state. - endpoint_move_rate: float = np.nan - - #: Fraction of outer transitions requiring a fine-model evaluation. - fine_eval_rate: float = np.nan - - #: Fraction of fine-evaluated candidate endpoints that pass the - #: fine-model subset condition. - fine_subset_pass_rate: float = np.nan - - #: Fraction of candidate endpoints that pass the fine subset condition - #: and are subsequently accepted by the final DA correction. - fine_correction_acceptance_rate: float = np.nan - - #: Among endpoints in the fine subset, fraction that are also above the - #: coarse threshold. - coarse_given_fine: float = np.nan - - #: Coarse performance-function values associated with the stored - #: outer-chain states. - all_gc: np.ndarray | None = None - - # Optional detailed diagnostics. - coarse_acceptance_rates: np.ndarray | None = None - coarse_acceptance_counts: np.ndarray | None = None - endpoint_moves: np.ndarray | None = None - fine_evals: np.ndarray | None = None - fine_subset_passes: np.ndarray | None = None - fine_correction_accepts: np.ndarray | None = None - fine_log_alpha: np.ndarray | None = None - endpoint_coarse_subset_passes: np.ndarray | None = None - - -@dataclass(repr=False) -class SubsetSimulationResult: - pf: float - cov: float | None - converged: bool - - thresholds: NDArray - level_covs: NDArray - levels: list[LevelSamplingResult] - - num_fine_evals: int - num_coarse_evals: int - num_failed_per_level: NDArray - - # Debugging/state data omitted from repr. - x_original: NDArray | None = None - final_chain_seeds: NDArray | None = None - final_chain_g: NDArray | None = None - final_all_x: NDArray | None = None - final_all_g: NDArray | None = None - debug_data: dict[str, Any] | None = None - - @property - def component_acceptance_rates(self) -> NDArray: - return np.array([level.component_acceptance_rate for level in self.levels]) - - @property - def subset_acceptance_rates(self) -> NDArray: - return np.array([level.subset_acceptance_rate for level in self.levels]) - - @property - def mean_jump_distances(self) -> NDArray: - return np.array([level.mean_jump_distance for level in self.levels]) - - @property - def outer_move_rates(self) -> NDArray: - return np.array([level.outer_move_rate for level in self.levels]) - - def __eq__(self, other: object) -> bool: - if not isinstance(other, SubsetSimulationResult): - return NotImplemented - - if type(self) is not type(other): - return False - - return all( - _values_equal( - getattr(self, field.name), - getattr(other, field.name), - ) - for field in fields(self) - ) - - def __repr__(self) -> str: - cls_name = type(self).__name__ - cov_str = f"{self.cov:.4f}" if self.cov is not None else "None" - return ( - f"{cls_name}(\n" - f" pf={self.pf:.4e},\n" - f" cov={cov_str},\n" - f" converged={self.converged!r},\n" - f" thresholds={self.thresholds!r},\n" - f" level_covs={self.level_covs!r},\n" - f" num_conditional_levels={len(self.levels)},\n" - f" component_acceptance_rates={self.component_acceptance_rates!r},\n" - f" subset_acceptance_rates={self.subset_acceptance_rates!r},\n" - f" outer_move_rates={self.outer_move_rates!r},\n" - f" mean_jump_distances={self.mean_jump_distances!r},\n" - f" num_fine_evals={self.num_fine_evals!r},\n" - f" num_coarse_evals={self.num_coarse_evals!r},\n" - f" num_failed_per_level={self.num_failed_per_level!r},\n" - f")" - ) - - -def _values_equal(left: Any, right: Any) -> bool: - """Compare nested result data, including NumPy arrays and NaNs.""" - - if isinstance(left, np.ndarray) or isinstance(right, np.ndarray): - if not (isinstance(left, np.ndarray) and isinstance(right, np.ndarray)): - return False - - return np.array_equal( - left, - right, - equal_nan=True, - ) - - if is_dataclass(left) or is_dataclass(right): - if not (is_dataclass(left) and is_dataclass(right) and type(left) is type(right)): - return False - - return all( - _values_equal( - getattr(left, field.name), - getattr(right, field.name), - ) - for field in fields(left) - ) - - if isinstance(left, dict) or isinstance(right, dict): - if not ( - isinstance(left, dict) - and isinstance(right, dict) - and left.keys() == right.keys() - ): - return False - - return all(_values_equal(left[key], right[key]) for key in left) - - if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): - if type(left) is not type(right) or len(left) != len(right): - return False - - return all(_values_equal(a, b) for a, b in zip(left, right)) - - # Treat NaN diagnostics as equal. - if ( - isinstance(left, (float, np.floating)) - and isinstance(right, (float, np.floating)) - and np.isnan(left) - and np.isnan(right) - ): - return True - - return bool(left == right) diff --git a/matflow/tests/test_subset_simulation.py b/matflow/tests/test_subset_simulation.py new file mode 100644 index 00000000..3e613b72 --- /dev/null +++ b/matflow/tests/test_subset_simulation.py @@ -0,0 +1,85 @@ +from functools import partial + +import numpy as np +from scipy.stats import norm + +from matflow.subset_simulation.subset_simulation import ( + generate_next_level_samples, + generate_next_level_samples_DA, + get_approx_y_star_random_walk, + log_surrogate_weight, + make_voxel_grouping, + subset_simulation, + system_analysis_toy_model, + weakest_link_performance_coarse, + weakest_link_performance_fine, +) + + +def test_subset_simulation_toy_model(): + seed = 1234 + performance = partial(system_analysis_toy_model, dimension=200, target_pf=1e-4) + result = subset_simulation( + dimension=200, + performance=performance, + p_0=0.1, + num_samples=100, + num_levels=7, + sampling_method=generate_next_level_samples, + sampling_method_kwargs={ + "proposal": norm(scale=1.0), + }, + master_seed=seed, + mimic_matflow=True, + ) + assert result.converged + assert result.num_conditional_levels == 3 + assert np.isclose(result.pf, 1.3000e-04) + assert np.isclose(result.cov, 0.9559503335693211) + assert np.allclose(result.outer_move_rates, [0.51111111, 0.3, 0.25555556]) + + +def test_subset_simulation_toy_model_DA(): + seed = 123 + + dimension = 200 + block_size = 10 + target_pf = 1e-4 + + y_star = get_approx_y_star_random_walk( + sigma=1, dimension=dimension, target_pf=target_pf + ) + + group_idx = make_voxel_grouping(dimension, block_size) + + performance = partial(weakest_link_performance_fine, y_star=y_star) + performance_coarse = partial( + weakest_link_performance_coarse, y_star=y_star, group_idx=group_idx + ) + + result = subset_simulation( + performance=performance, + dimension=dimension, + p_0=0.1, + num_samples=100, + num_levels=4, + master_seed=seed, + sampling_method=generate_next_level_samples_DA, + sampling_method_kwargs={ + "proposal": norm(), + "temperature": 1, # should be of a similar order of magnitude to threshold + "performance_coarse": performance_coarse, + "log_surrogate_weight": log_surrogate_weight, + # "log_surrogate_weight": lambda *args, **kwargs: 1, + "num_inner_states": 3, + "spawn_key": (5,), + }, + transformation=lambda x: np.cumsum(x, axis=-1), # random walk model + debug=True, + mimic_matflow=True, + ) + assert result.converged + assert result.num_conditional_levels == 3 + assert np.isclose(result.pf, 1.1000e-04) + assert np.isclose(result.cov, 0.6454545454545455) + assert np.allclose(result.outer_move_rates, [0.64444444, 0.42222222, 0.35555556]) diff --git a/matflow/tests/utils.py b/matflow/tests/utils.py index 46fe0b80..fe65a0b9 100644 --- a/matflow/tests/utils.py +++ b/matflow/tests/utils.py @@ -1,15 +1,17 @@ """Utility functions to assist testing.""" +from dataclasses import fields, is_dataclass from functools import partial +from typing import Any from hpcflow.sdk.core.test_utils import ( make_test_data_YAML_workflow, make_test_data_YAML_workflow_template, ) +import numpy as np import matflow as mf - make_test_data_YAML_workflow = partial( make_test_data_YAML_workflow, app=mf, pkg="matflow.tests.data" ) @@ -17,3 +19,56 @@ make_test_data_YAML_workflow_template = partial( make_test_data_YAML_workflow_template, app=mf, pkg="matflow.tests.data" ) + + +def _values_equal(left: Any, right: Any) -> bool: + """Compare nested result data, including NumPy arrays and NaNs.""" + + if isinstance(left, np.ndarray) or isinstance(right, np.ndarray): + if not (isinstance(left, np.ndarray) and isinstance(right, np.ndarray)): + return False + + return np.array_equal( + left, + right, + equal_nan=True, + ) + + if is_dataclass(left) or is_dataclass(right): + if not (is_dataclass(left) and is_dataclass(right) and type(left) is type(right)): + return False + + return all( + _values_equal( + getattr(left, field.name), + getattr(right, field.name), + ) + for field in fields(left) + ) + + if isinstance(left, dict) or isinstance(right, dict): + if not ( + isinstance(left, dict) + and isinstance(right, dict) + and left.keys() == right.keys() + ): + return False + + return all(_values_equal(left[key], right[key]) for key in left) + + if isinstance(left, (list, tuple)) or isinstance(right, (list, tuple)): + if type(left) is not type(right) or len(left) != len(right): + return False + + return all(_values_equal(a, b) for a, b in zip(left, right)) + + # Treat NaN diagnostics as equal. + if ( + isinstance(left, (float, np.floating)) + and isinstance(right, (float, np.floating)) + and np.isnan(left) + and np.isnan(right) + ): + return True + + return bool(left == right) From 3f250c71a72ee04f73151c5d0c716aa8c7142dc9 Mon Sep 17 00:00:00 2001 From: Adam Plowman Date: Mon, 14 Sep 2026 23:31:42 +0100 Subject: [PATCH 4/4] feat: return mean MMH component acceptance rate and mean jump distance --- matflow/data/scripts/uq/collate_results.py | 37 ++++++++++++++++--- .../data/scripts/uq/generate_next_state.py | 18 +++++++-- .../scripts/uq/initialise_markov_chains.py | 10 ++++- .../template_components/task_schemas.yaml | 21 +++++++++-- .../subset_simulation_DAMASK_Mg.yaml | 2 + .../subset_simulation_DAMASK_Mg_DA.yaml | 2 + ...subset_simulation_DAMASK_Mg_two_level.yaml | 3 ++ .../subset_simulation_toy_model.yaml | 2 + .../subset_simulation_toy_model_DA.yaml | 2 + .../subset_simulation_toy_model_external.yaml | 2 + ...subset_simulation_toy_model_two_level.yaml | 2 + ...mulation_toy_model_two_level_external.yaml | 2 + .../subset_simulation/subset_simulation.py | 34 ++++++++--------- .../demo_workflows/test_demo_workflows.py | 27 ++++++++++---- matflow/tests/test_subset_simulation.py | 6 +++ 15 files changed, 133 insertions(+), 37 deletions(-) diff --git a/matflow/data/scripts/uq/collate_results.py b/matflow/data/scripts/uq/collate_results.py index 31f35a62..5003e218 100644 --- a/matflow/data/scripts/uq/collate_results.py +++ b/matflow/data/scripts/uq/collate_results.py @@ -36,6 +36,12 @@ def estimate_cov(indicator, p_i: float) -> float: return delta +def rms_jump_distances(all_x: np.ndarray, diff_axis) -> np.ndarray: + """RMS jump distance between consecutive stored chain states.""" + dx = np.diff(all_x, axis=diff_axis) + return np.linalg.norm(dx, axis=-1) / np.sqrt(dx.shape[-1]) + + def collate_results( g, x, @@ -49,6 +55,8 @@ def collate_results( fine_eval_count, coarse_eval_count, fine_eval_rates, + num_components_accepted, + num_components_proposed, ): # all iterations of g are passed just to get the level index: @@ -73,8 +81,21 @@ def collate_results( for iter_dat in all_accept.values(): if iter_dat["value"]: all_all_accept.append(np.vstack([i[:] for i in iter_dat["value"]])) - accept_rate = np.mean(all_all_accept, axis=(1, 2)) - x = np.vstack([i[:] for i in all_x]) + outer_move_rate = np.mean(all_all_accept, axis=(1, 2)) + x = [] + jump_distances = [] + for i in all_x: + x.append(i[:]) + jump_distances.extend(rms_jump_distances(i, diff_axis=0)) + x = np.vstack(x) + + # sum over chains: + num_components_accepted = sum(num_components_accepted) + num_components_proposed = sum(num_components_proposed) + component_acceptance_rate = num_components_accepted / num_components_proposed + + jump_distances = np.array(jump_distances) + mean_jump_distance = np.mean(jump_distances).item() # sum across chains: fine_eval_count = np.sum(fine_eval_count) @@ -92,9 +113,11 @@ def collate_results( # from initial direct Monte Carlo samples: g_unsrt = np.array(g) x = np.vstack([i[:] for i in x]) - accept_rate = None + outer_move_rate = None total_states = num_samples fine_eval_count = g_unsrt.size + mean_jump_distance = None + component_acceptance_rate = None fine_eval_count = int(fine_eval_count) total_fine_eval_count += fine_eval_count @@ -170,7 +193,7 @@ def collate_results( f"num_samples: {num_samples!r}\n" f"num_chains: {num_chains!r}\n" f"num_failed: {num_failed!r}\n" - f"accept_rate: {accept_rate[:] if accept_rate is not None else '-'}\n" + f"outer_move_rate: {outer_move_rate if outer_move_rate is not None else '-'}\n" f"total_fine_eval_count: {total_fine_eval_count!r}\n" f"total_coarse_eval_count: {total_coarse_eval_count!r}\n" f"level_pf: {level_pf.item()!r}\n" @@ -179,6 +202,8 @@ def collate_results( f"level_fine_eval_rate: {fine_eval_rate!r}\n" f"level_coarse_eval_count: {coarse_eval_count if coarse_eval_count is not None else '-'}\n" f"fine_eval_rates: {fine_eval_rates!r}\n" + f"component_acceptance_rate: {component_acceptance_rate if component_acceptance_rate is not None else '-'}\n" + f"mean_jump_distance: {mean_jump_distance if mean_jump_distance is not None else '-'}\n" "\n", ) @@ -196,8 +221,10 @@ def collate_results( "fine_eval_rates": fine_eval_rates, "pf": pf, "is_finished": is_finished, - "accept_rate": accept_rate, + "outer_move_rate": outer_move_rate, "cov": cov, "total_fine_eval_count": total_fine_eval_count, "total_coarse_eval_count": total_coarse_eval_count, + "component_acceptance_rate": component_acceptance_rate, + "mean_jump_distance": mean_jump_distance, } diff --git a/matflow/data/scripts/uq/generate_next_state.py b/matflow/data/scripts/uq/generate_next_state.py index eaf8dd88..a3b4594b 100644 --- a/matflow/data/scripts/uq/generate_next_state.py +++ b/matflow/data/scripts/uq/generate_next_state.py @@ -34,7 +34,9 @@ def _init_rng(chain_index, loop_idx): return rng -def generate_next_state(x, proposal, rng, chain_index): +def generate_next_state( + x, proposal, rng, chain_index, num_components_accepted, num_components_proposed +): """Generate the next candidate state in a modified Metropolis algorithm. Parameters @@ -48,6 +50,10 @@ def generate_next_state(x, proposal, rng, chain_index): Random number generator to be used in this function. chain_index Index of the Markov chain within the subset simulation level loop. + num_components_accepted + Running total of the number of accepted MMH components. + num_components_proposed + Running total of the number of proposed MMH components. Returns ------- @@ -98,6 +104,12 @@ def generate_next_state(x, proposal, rng, chain_index): xi[accept_idx] = xi_hat[accept_idx] xi[~accept_idx] = current_state[~accept_idx] - mcmc_accept_rate = np.mean(accept_idx) + num_components_accepted += np.sum(accept_idx).item() + num_components_proposed += current_state.size - return {"x": xi, "mcmc_accept_rate": mcmc_accept_rate, "rng": rng} + return { + "x": xi, + "num_components_accepted": num_components_accepted, + "num_components_proposed": num_components_proposed, + "rng": rng, + } diff --git a/matflow/data/scripts/uq/initialise_markov_chains.py b/matflow/data/scripts/uq/initialise_markov_chains.py index 79bac875..5abb6392 100644 --- a/matflow/data/scripts/uq/initialise_markov_chains.py +++ b/matflow/data/scripts/uq/initialise_markov_chains.py @@ -8,4 +8,12 @@ def initialise_markov_chains(chain_index, chain_seeds, chain_g): all_x = np.array(x)[None] all_g = np.array([g]) all_accept = np.array([], dtype=bool) - return {"x": x, "g": g, "all_x": all_x, "all_g": all_g, "all_accept": all_accept} + return { + "x": x, + "g": g, + "all_x": all_x, + "all_g": all_g, + "all_accept": all_accept, + "num_components_accepted": 0, + "num_components_proposed": 0, + } diff --git a/matflow/data/template_components/task_schemas.yaml b/matflow/data/template_components/task_schemas.yaml index ba20bb79..fba427a9 100644 --- a/matflow/data/template_components/task_schemas.yaml +++ b/matflow/data/template_components/task_schemas.yaml @@ -1819,6 +1819,12 @@ - parameter: coarse_eval_count default_value: null group: all + - parameter: num_components_accepted + group: all + default_value: null + - parameter: num_components_proposed + group: all + default_value: null - parameter: fine_eval_rates default_value: null outputs: @@ -1836,9 +1842,11 @@ - parameter: total_coarse_eval_count - parameter: pf - parameter: is_finished - - parameter: accept_rate + - parameter: outer_move_rate - parameter: cov - - parameter: fine_eval_rates + - parameter: fine_eval_rates + - parameter: component_acceptance_rate + - parameter: mean_jump_distance actions: - script: <> script_data_in: @@ -1864,6 +1872,8 @@ - parameter: all_x - parameter: all_g - parameter: all_accept + - parameter: num_components_accepted + - parameter: num_components_proposed actions: - script: <> script_data_in: direct @@ -1909,10 +1919,15 @@ - parameter: chain_index - parameter: rng default_value: null + - parameter: num_components_accepted + default_value: 0 + - parameter: num_components_proposed + default_value: 1 outputs: - parameter: x - - parameter: mcmc_accept_rate - parameter: rng + - parameter: num_components_accepted + - parameter: num_components_proposed actions: - script: <> script_data_in: direct diff --git a/matflow/data/workflows/subset_simulation_DAMASK_Mg.yaml b/matflow/data/workflows/subset_simulation_DAMASK_Mg.yaml index cf64ba37..52361d88 100644 --- a/matflow/data/workflows/subset_simulation_DAMASK_Mg.yaml +++ b/matflow/data/workflows/subset_simulation_DAMASK_Mg.yaml @@ -259,6 +259,8 @@ tasks: proposal: type: norm scale: 1.0 + groups: + - name: all - schema: system_analysis # [inner loop] - schema: increment_chain # [inner loop] groups: diff --git a/matflow/data/workflows/subset_simulation_DAMASK_Mg_DA.yaml b/matflow/data/workflows/subset_simulation_DAMASK_Mg_DA.yaml index a6c3f1f0..043bb56f 100644 --- a/matflow/data/workflows/subset_simulation_DAMASK_Mg_DA.yaml +++ b/matflow/data/workflows/subset_simulation_DAMASK_Mg_DA.yaml @@ -510,6 +510,8 @@ tasks: proposal: type: norm scale: 1.0 + groups: + - name: all - schema: system_analysis_16 resources: diff --git a/matflow/data/workflows/subset_simulation_DAMASK_Mg_two_level.yaml b/matflow/data/workflows/subset_simulation_DAMASK_Mg_two_level.yaml index e0282b1a..bf619950 100644 --- a/matflow/data/workflows/subset_simulation_DAMASK_Mg_two_level.yaml +++ b/matflow/data/workflows/subset_simulation_DAMASK_Mg_two_level.yaml @@ -479,6 +479,9 @@ tasks: proposal: type: norm scale: 1.0 + groups: + - name: all + - schema: system_analysis_16 # [inner MC loop; outer MC loop] # 7,8,9 - schema: increment_chain_inner # [inner MC loop; outer MC loop] # outputs x, like `generate_next_state` does # 10 diff --git a/matflow/data/workflows/subset_simulation_toy_model.yaml b/matflow/data/workflows/subset_simulation_toy_model.yaml index 072841b4..87be009f 100644 --- a/matflow/data/workflows/subset_simulation_toy_model.yaml +++ b/matflow/data/workflows/subset_simulation_toy_model.yaml @@ -44,6 +44,8 @@ tasks: proposal: type: norm scale: 1.0 + groups: + - name: all - schema: system_analysis_toy_model # [inner loop] inputs: dimension: 200 diff --git a/matflow/data/workflows/subset_simulation_toy_model_DA.yaml b/matflow/data/workflows/subset_simulation_toy_model_DA.yaml index e20386ec..b92fe3f4 100644 --- a/matflow/data/workflows/subset_simulation_toy_model_DA.yaml +++ b/matflow/data/workflows/subset_simulation_toy_model_DA.yaml @@ -81,6 +81,8 @@ tasks: proposal: type: norm scale: 1.0 + groups: + - name: all - schema: system_analysis_toy_model_max_random_walk_coarse resources: diff --git a/matflow/data/workflows/subset_simulation_toy_model_external.yaml b/matflow/data/workflows/subset_simulation_toy_model_external.yaml index 6e9ed157..0e396e76 100644 --- a/matflow/data/workflows/subset_simulation_toy_model_external.yaml +++ b/matflow/data/workflows/subset_simulation_toy_model_external.yaml @@ -56,6 +56,8 @@ tasks: proposal: type: norm scale: 0.5 + groups: + - name: all - schema: system_analysis # [inner loop] - schema: increment_chain # [inner loop] groups: diff --git a/matflow/data/workflows/subset_simulation_toy_model_two_level.yaml b/matflow/data/workflows/subset_simulation_toy_model_two_level.yaml index b8e9b443..27fffb97 100644 --- a/matflow/data/workflows/subset_simulation_toy_model_two_level.yaml +++ b/matflow/data/workflows/subset_simulation_toy_model_two_level.yaml @@ -53,6 +53,8 @@ tasks: proposal: type: norm scale: 1.0 + groups: + - name: all - schema: system_analysis_toy_model # [inner MC loop; outer MC loop] inputs: diff --git a/matflow/data/workflows/subset_simulation_toy_model_two_level_external.yaml b/matflow/data/workflows/subset_simulation_toy_model_two_level_external.yaml index e991937b..8dd6ad11 100644 --- a/matflow/data/workflows/subset_simulation_toy_model_two_level_external.yaml +++ b/matflow/data/workflows/subset_simulation_toy_model_two_level_external.yaml @@ -62,6 +62,8 @@ tasks: proposal: type: norm scale: 1.0 + groups: + - name: all - schema: system_analysis # [inner MC loop; outer MC loop] # 7,8,9 - schema: increment_chain_inner # [inner MC loop; outer MC loop] # outputs x, like `generate_next_state` does # 10 diff --git a/matflow/subset_simulation/subset_simulation.py b/matflow/subset_simulation/subset_simulation.py index 8163139e..8b55e04e 100644 --- a/matflow/subset_simulation/subset_simulation.py +++ b/matflow/subset_simulation/subset_simulation.py @@ -180,6 +180,9 @@ def estimate_cov(indicator, p_i: float) -> float: def generate_next_state(x, proposal, rng): """ Proposal must be a symmetric distribution centred on zero. + + Returns the new candidate state and the number of new components in that state (i.e. + number of accepted components). """ dim = len(x) @@ -193,9 +196,8 @@ def generate_next_state(x, proposal, rng): xi[accept_idx] = xi_hat[accept_idx] xi[~accept_idx] = current_state[~accept_idx] - mcmc_accept_rate = np.mean(accept_idx) - return xi, mcmc_accept_rate + return xi, np.sum(accept_idx) def generate_next_state_CS(x, prop_std, rng): @@ -251,16 +253,12 @@ def generate_next_level_samples( while not chain.is_complete: current_x = chain.current_x - (trial_x, component_acceptance_rate,) = generate_next_state( + trial_x, num_components_accepted = generate_next_state( x=current_x, proposal=proposal, rng=chain_rng, ) - # Preferably generate_next_state() would return this - # integer count directly. - num_components_accepted = int(np.rint(component_acceptance_rate * dimension)) - trial_x_t = transformation(trial_x) if transformation else trial_x trial_g = performance(trial_x_t) @@ -268,7 +266,8 @@ def generate_next_level_samples( trial_x=trial_x, trial_g=trial_g, threshold=threshold, - num_components_accepted=(num_components_accepted), + num_components_accepted=num_components_accepted, + num_components_proposed=chain.current_x.size, ) chain_results.append(chain.finalise(retain_jump_distances=debug)) @@ -497,7 +496,7 @@ def generate_coarse_subchain( current_sub_chain_gc = gc inner_accepts = 0 - mmh_component_acceptance_sum = 0.0 + num_components_accepted_sum = 0.0 debug_data = {} if debug: @@ -511,10 +510,10 @@ def generate_coarse_subchain( if debug: debug_data["rng_states"].append(rng.bit_generator.state["state"]) - trial_x, mmh_component_acceptance = generate_next_state( + trial_x, num_components_accepted = generate_next_state( x=current_sub_chain_x, proposal=proposal, rng=rng ) - mmh_component_acceptance_sum += mmh_component_acceptance + num_components_accepted_sum += num_components_accepted if debug: debug_data["current_sub_chain_x"].append(current_sub_chain_x) @@ -545,7 +544,7 @@ def generate_coarse_subchain( current_sub_chain_x, current_sub_chain_gc, inner_accepts, - mmh_component_acceptance_sum, + num_components_accepted_sum, debug_data, ) @@ -643,7 +642,7 @@ def generate_next_level_samples_DA( psi, psi_gc, num_inner_accepts, - component_acceptance_sum, + num_components_accepted, subchain_debug_data, ) = generate_coarse_subchain( x=current_x, @@ -660,8 +659,7 @@ def generate_next_level_samples_DA( debug=debug, ) - num_components_proposed = num_inner_states * dimension - num_components_accepted = int(np.rint(component_acceptance_sum * dimension)) + num_components_proposed = num_inner_states * current_x.size endpoint_moved = not np.array_equal(psi, current_x) @@ -719,14 +717,14 @@ def generate_next_level_samples_DA( new_x=new_x, new_g=new_g, new_gc=new_gc, - num_components_accepted=(num_components_accepted), - num_components_proposed=(num_components_proposed), + num_components_accepted=num_components_accepted, + num_components_proposed=num_components_proposed, num_coarse_accepts=num_inner_accepts, num_coarse_proposals=num_inner_states, endpoint_moved=endpoint_moved, fine_evaluated=fine_evaluated, fine_subset_pass=fine_subset_pass, - fine_correction_accept=(fine_correction_accept), + fine_correction_accept=fine_correction_accept, coarse_subset_pass=coarse_subset_pass, ) diff --git a/matflow/tests/demo_workflows/test_demo_workflows.py b/matflow/tests/demo_workflows/test_demo_workflows.py index efbaa21d..1ccab1db 100644 --- a/matflow/tests/demo_workflows/test_demo_workflows.py +++ b/matflow/tests/demo_workflows/test_demo_workflows.py @@ -108,15 +108,19 @@ def test_subset_simulation_toy_model_prediction(tmp_path): final_iter = wk.tasks.collate_results.elements[0].latest_iteration_non_skipped pf = final_iter.get("outputs.pf") cov = final_iter.get("outputs.cov") - sus_acc = final_iter.get("outputs.accept_rate")[:] + outer_move_rate = final_iter.get("outputs.outer_move_rate") + mean_jump = final_iter.get("outputs.mean_jump_distance") + comp_accept = final_iter.get("outputs.component_acceptance_rate") # TODO: also verify same result with `subset_simulation_toy_model_external`, once # that can be submitted without a ridiculous number of processes. + # note the actual values are tested in ``test_subset_simulation`` assert pf == result.pf assert cov == result.cov - assert np.allclose(sus_acc, result.outer_move_rates) - # TODO: store and check mcmc_accept + assert np.allclose(outer_move_rate, result.outer_move_rates) + assert np.isclose(mean_jump, result.mean_jump_distances[-1]) + assert np.isclose(comp_accept, result.component_acceptance_rates[-1]) @pytest.mark.demo_workflows @@ -185,7 +189,16 @@ def test_subset_simulation_toy_model_DA_prediction(tmp_path): wk.wait() - iter_i = wk.tasks.collate_results.elements[0].iterations[-1] - assert iter_i.get("outputs.pf") == result.pf - assert iter_i.get("outputs.cov") == result.cov - assert iter_i.get("outputs.threshold") == result.thresholds[-1] + final_iter = wk.tasks.collate_results.elements[0].iterations[-1] + pf = final_iter.get("outputs.pf") + cov = final_iter.get("outputs.cov") + outer_move_rate = final_iter.get("outputs.outer_move_rate") + mean_jump = final_iter.get("outputs.mean_jump_distance") + comp_accept = final_iter.get("outputs.component_acceptance_rate") + + # note the actual values are tested in ``test_subset_simulation`` + assert pf == result.pf + assert cov == result.cov + assert np.allclose(outer_move_rate, result.outer_move_rates) + assert np.isclose(mean_jump, result.mean_jump_distances[-1]) + assert np.isclose(comp_accept, result.component_acceptance_rates[-1]) diff --git a/matflow/tests/test_subset_simulation.py b/matflow/tests/test_subset_simulation.py index 3e613b72..fcdbcd9b 100644 --- a/matflow/tests/test_subset_simulation.py +++ b/matflow/tests/test_subset_simulation.py @@ -37,6 +37,8 @@ def test_subset_simulation_toy_model(): assert np.isclose(result.pf, 1.3000e-04) assert np.isclose(result.cov, 0.9559503335693211) assert np.allclose(result.outer_move_rates, [0.51111111, 0.3, 0.25555556]) + assert np.allclose(result.mean_jump_distances, [0.33655373, 0.19935043, 0.16935632]) + assert np.allclose(result.component_acceptance_rates, [0.704, 0.7025, 0.69927778]) def test_subset_simulation_toy_model_DA(): @@ -83,3 +85,7 @@ def test_subset_simulation_toy_model_DA(): assert np.isclose(result.pf, 1.1000e-04) assert np.isclose(result.cov, 0.6454545454545455) assert np.allclose(result.outer_move_rates, [0.64444444, 0.42222222, 0.35555556]) + assert np.allclose(result.mean_jump_distances, [0.53359356, 0.31473844, 0.24871597]) + assert np.allclose( + result.component_acceptance_rates, [0.70409259, 0.70640741, 0.7047037] + )