Merge pull request #550 from hpcflow/fix/sus-acs

Fix adaptive conditional sampling variant of the toy model subset simulation workflow

Author: aplowman

matflow/data/scripts/uq/collate_results.py

@@ -51,9 +51,12 @@ def collate_results(g, x, p_0, all_g, all_x, all_accept, level_cov):
     if all_g:
         # from multiple Markov chains:
         g_unsrt = np.concatenate([i[:] for i in all_g])
-        accept = np.vstack([i[:] for i in all_accept])
+        all_all_accept = []
+        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])
-        accept_rate = np.mean(accept)
     else:
         # from initial direct Monte Carlo samples:
         g_unsrt = np.array(g)

matflow/data/scripts/uq/generate_next_state.py

@@ -2,8 +2,7 @@
 import os
 
 import numpy as np
-from numpy.typing import NDArray
-from scipy.stats import norm
+import scipy.stats
 
 
 def set_up_logger():
@@ -35,15 +34,16 @@ def _init_rng(chain_index, loop_idx):
     return rng
 
 
-def generate_next_state(x, prop_std, rng, chain_index):
+def generate_next_state(x, proposal, rng, chain_index):
     """Generate the next candidate state in a modified Metropolis algorithm.
 
     Parameters
     ----------
     x
         Current state on which the candidate state will depend.
-    prop_std
-        Proposal distribution standard deviation.
+    proposal
+        Type and arguments to a Scipy distribution that should be used as the proposal.
+        The proposal must be a symmetrical distribution, centred on zero.
     rng
         Random number generator to be used in this function.
     chain_index
@@ -59,6 +59,9 @@ def generate_next_state(x, prop_std, rng, chain_index):
             for this chain.
     """
 
+    dist_type = proposal.pop("type")
+    prop_dist = getattr(scipy.stats, dist_type)(**proposal)
+
     loop_idx = {
         loop_name: int(loop_idx)
         for loop_name, loop_idx in (
@@ -76,12 +79,13 @@ def generate_next_state(x, prop_std, rng, chain_index):
     current_state = x
     xi = np.empty(dim)
 
-    proposal = norm(loc=current_state, scale=prop_std)
-    xi_hat = np.atleast_1d(proposal.rvs(random_state=rng))
-    accept_ratios = np.divide(*norm.pdf([xi_hat, current_state]))
+    xi_hat = np.atleast_1d(current_state + prop_dist.rvs(size=dim, random_state=rng))
+    accept_ratios = np.divide(*scipy.stats.norm.pdf([xi_hat, current_state]))
     accept_idx = rng.random(len(accept_ratios)) < np.minimum(1, accept_ratios)
 
     xi[accept_idx] = xi_hat[accept_idx]
     xi[~accept_idx] = current_state[~accept_idx]
 
-    return {"x": xi, "rng": rng}
+    mcmc_accept_rate = np.mean(accept_idx)
+
+    return {"x": xi, "mcmc_accept_rate": mcmc_accept_rate, "rng": rng}

matflow/data/template_components/task_schemas.yaml

@@ -1807,7 +1807,7 @@
       group: all
       default_value: null
     - parameter: level_cov
-      default_value: null      
+      default_value: null
     - parameter: p_0
   outputs:
     - parameter: chain_seeds
@@ -1826,6 +1826,7 @@
       script_data_in:
         g: { format: direct, all_iterations: true }
         level_cov: { format: direct, all_iterations: true }
+        all_accept: { format: direct, all_iterations: true }
         "*": direct
       script_data_out: direct
       script_exe: python_script
@@ -1886,12 +1887,13 @@
 - objective: generate_next_state
   inputs:
     - parameter: x
-    - parameter: prop_std
+    - parameter: proposal
     - parameter: chain_index
     - parameter: rng
       default_value: null
   outputs:
     - parameter: x
+    - parameter: mcmc_accept_rate
     - parameter: rng
   actions:
     - script: <<script:uq/generate_next_state.py>>
@@ -1910,7 +1912,7 @@
     - parameter: prop_std
     - parameter: chain_index
     - parameter: rng
-      default_value: null    
+      default_value: null
   outputs:
     - parameter: x
     - parameter: rng
@@ -1932,7 +1934,7 @@
     - parameter: lambda_
     - parameter: chain_index
     - parameter: rng
-      default_value: null    
+      default_value: null
   outputs:
     - parameter: x
     - parameter: rng

matflow/data/workflows/subset_simulation_DAMASK_Al.yaml

@@ -117,7 +117,9 @@ tasks:
           step: 1
   - schema: generate_next_state # [inner loop]
     inputs:
-      prop_std: 0.5
+      proposal:
+        type: norm
+        scale: 0.5
   - schema: system_analysis # [inner loop]
   - schema: increment_chain # [inner loop]
     groups:

matflow/data/workflows/subset_simulation_DAMASK_Mg.yaml

@@ -256,7 +256,9 @@ tasks:
           step: 1
   - schema: generate_next_state # [inner loop]
     inputs:
-      prop_std: 1.0
+      proposal:
+        type: norm
+        scale: 1.0
   - schema: system_analysis # [inner loop]
   - schema: increment_chain # [inner loop]
     groups:

matflow/data/workflows/subset_simulation_DAMASK_Mg_two_level.yaml

@@ -476,7 +476,9 @@ tasks:
 
   - schema: generate_next_state # [inner MC loop; outer MC loop] # 6
     inputs:
-      prop_std: 1.0
+      proposal:
+        type: norm
+        scale: 1.0
   - 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

@@ -1,9 +1,9 @@
 loops:
   - name: markov_chain_state # [inner loop]
-    tasks: [4, 5, 6]
+    tasks: [ 4, 5, 6 ]
     num_iterations: 9 # num_states - 1: (1 / p_0) - 1
   - name: levels # [outer loop]
-    tasks: [2, 3, 4, 5, 6]
+    tasks: [ 2, 3, 4, 5, 6 ]
     num_iterations: 7
     termination_task: 2
     termination:
@@ -26,7 +26,7 @@ tasks:
   - schema: system_analysis_toy_model
     inputs:
       dimension: 200
-      target_pf: 1.0e-4
+      target_pf: 1e-4
     groups:
       - name: all
   - schema: collate_results # [outer loop]
@@ -41,11 +41,13 @@ tasks:
           step: 1
   - schema: generate_next_state # [inner loop]
     inputs:
-      prop_std: 1.0
+      proposal:
+        type: norm
+        scale: 1.0
   - schema: system_analysis_toy_model # [inner loop]
     inputs:
       dimension: 200
-      target_pf: 1.0e-4
+      target_pf: 1e-4
   - schema: increment_chain # [inner loop]
     groups:
       - name: all

matflow/data/workflows/subset_simulation_toy_model_delayed_acceptance.yaml

@@ -47,7 +47,9 @@ tasks:
 
   - schema: generate_next_state # [inner MC loop; outer MC loop]
     inputs:
-      prop_std: 1.0
+      proposal:
+        type: norm
+        scale: 1.0
 
   - schema: system_analysis_DA_coarse # [inner MC loop; outer MC loop]
     resources:

matflow/data/workflows/subset_simulation_toy_model_external.yaml

@@ -53,7 +53,9 @@ tasks:
           step: 1
   - schema: generate_next_state # [inner loop]
     inputs:
-      prop_std: 0.5
+      proposal:
+        type: norm
+        scale: 0.5
   - schema: system_analysis # [inner loop]
   - schema: increment_chain # [inner loop]
     groups:

matflow/data/workflows/subset_simulation_toy_model_two_level.yaml

@@ -45,7 +45,9 @@ tasks:
 
   - schema: generate_next_state # [inner MC loop; outer MC loop]
     inputs:
-      prop_std: 1.0
+      proposal:
+        type: norm
+        scale: 1.0
   - schema: system_analysis_toy_model # [inner MC loop; outer MC loop]
     inputs:
       dimension: 200