Improvements to subset simulation

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,
     }

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,
+    }

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,
+    }

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:uq/collate_results.py>>
       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:uq/initialise_markov_chains.py>>
       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:uq/generate_next_state.py>>
       script_data_in: direct

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:

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:

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
 

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

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:

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:

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:

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