Subset simulation reproducibility and testing

matflow/data/scripts/uq/collate_results.py

@@ -19,9 +19,12 @@ def estimate_cov(indicator, p_i: float) -> float:
             r_k += r_k_i
         r[k] = (r_k / (N - (k + 1) * num_chains)) - p_i**2
 
-    r_0 = np.sum(indicator**2) / N - p_i**2  # autocovariance at lag zero (exact)
     r_0 = p_i * (1 - p_i)
 
+    if np.isclose(r_0, 0.0):
+        # i.e. p_i is 1.0
+        return 0.0
+
     rho = r / r_0
 
     gamma = 2 * sum(
@@ -47,8 +50,7 @@ def collate_results(g, x, p_0, all_g, all_x, all_accept, level_cov):
 
     if all_g:
         # from multiple Markov chains:
-        g_2D = np.array([i[:] for i in all_g])  # (num_chains, num_states)
-        g_unsrt = np.concatenate(g_2D)
+        g_unsrt = np.concatenate([i[:] for i in all_g])
         accept = np.vstack([i[:] for i in all_accept])
         x = np.vstack([i[:] for i in all_x])
         accept_rate = np.mean(accept)
@@ -87,14 +89,17 @@ def collate_results(g, x, p_0, all_g, all_x, all_accept, level_cov):
 
     threshold = (g[num_chains - 1] + g[num_chains]) / 2
 
+    # indicator function:
+    g_2D = np.reshape(g_unsrt, (num_chains, num_states))
+    indicator = (g_2D > np.minimum(threshold, 0)).astype(int)
+
     # failure probability at this level:
-    fail_bool = g > 0
-    level_pf = np.mean(fail_bool) if threshold > 0 else p_0
+    level_pf = np.mean(indicator)
 
     chain_seeds = x[:num_chains]
     chain_g = g[:num_chains]
 
-    is_finished = (num_failed / num_samples) >= p_0
+    is_finished = (threshold > 0).item()
 
     pf = p_0**level_idx * num_failed / num_samples
 
@@ -106,9 +111,6 @@ def collate_results(g, x, p_0, all_g, all_x, all_accept, level_cov):
 
     if all_g:
         # from multiple Markov chains:
-        indicator = np.reshape(
-            g_2D > np.minimum(threshold, 0), (num_chains, num_states)
-        ).astype(int)
         level_cov = estimate_cov(indicator, level_pf)
     else:
         # from initial direct Monte Carlo samples:

matflow/data/scripts/uq/generate_next_state.py

@@ -1,4 +1,5 @@
 import logging
+import os
 
 import numpy as np
 from numpy.typing import NDArray
@@ -17,33 +18,70 @@ def set_up_logger():
     return logger
 
 
-def generate_next_state(x, prop_std):
+def _init_rng(chain_index, loop_idx):
+    """Initialise the random number generator, which must happen once for each Markov
+    chain."""
+    random_seed = int(os.environ["MATFLOW_RUN_RANDOM_SEED"])
+    spawn_key_str = os.environ["MATFLOW_RUN_RNG_SPAWN_KEY"]
+    task_iID = int(os.environ["MATFLOW_TASK_INSERT_ID"])
+    spawn_key = tuple(int(i) for i in spawn_key_str.split(",")) if spawn_key_str else ()
+    spawn_key = tuple([*spawn_key, task_iID, loop_idx["levels"], chain_index])
+    rng = np.random.default_rng(
+        np.random.SeedSequence(
+            random_seed,
+            spawn_key=spawn_key,
+        )
+    )
+    return rng
+
+
+def generate_next_state(x, prop_std, 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.
+    rng
+        Random number generator to be used in this function.
+    chain_index
+        Index of the Markov chain within the subset simulation level loop.
 
     Returns
     -------
     dict:
         x:
             Generated candidate state.
+        rng:
+            Random number generator to be used in the next invocation of this function,
+            for this chain.
     """
 
+    loop_idx = {
+        loop_name: int(loop_idx)
+        for loop_name, loop_idx in (
+            item.split("=")
+            for item in os.environ["MATFLOW_ELEMENT_ITER_LOOP_IDX"].split(";")
+        )
+    }
+
+    if loop_idx["markov_chain_state"] == 0:
+        # new chain, so want a new RNG:
+        rng = _init_rng(chain_index, loop_idx)
+
     x = x[:]  # convert to numpy array
     dim = len(x)
     current_state = x
-    rng = np.random.default_rng()
     xi = np.empty(dim)
 
     proposal = norm(loc=current_state, scale=prop_std)
-    xi_hat = np.atleast_1d(proposal.rvs())
+    xi_hat = np.atleast_1d(proposal.rvs(random_state=rng))
     accept_ratios = np.divide(*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}
+    return {"x": xi, "rng": rng}

matflow/data/scripts/uq/generate_next_state_ACS.py

@@ -1,26 +1,67 @@
+import os
+
 import numpy as np
 from numpy.typing import NDArray
 from scipy.stats import norm
 
 
-def generate_next_state_ACS(x, prop_std, lambda_):
+def _init_rng(chain_index, loop_idx):
+    """Initialise the random number generator, which must happen once for each Markov
+    chain."""
+    random_seed = int(os.environ["MATFLOW_RUN_RANDOM_SEED"])
+    spawn_key_str = os.environ["MATFLOW_RUN_RNG_SPAWN_KEY"]
+    task_iID = int(os.environ["MATFLOW_TASK_INSERT_ID"])
+    spawn_key = tuple(int(i) for i in spawn_key_str.split(",")) if spawn_key_str else ()
+    spawn_key = tuple([*spawn_key, task_iID, loop_idx["levels"], chain_index])
+    rng = np.random.default_rng(
+        np.random.SeedSequence(
+            random_seed,
+            spawn_key=spawn_key,
+        )
+    )
+    return rng
+
+
+def generate_next_state_ACS(x, prop_std, lambda_, rng, chain_index):
     """Generate the next candidate state using adaptive conditional sampling (a.k.a
     adaptive subset-infinity)
 
     Parameters
     ----------
     x
         Current state on which the candidate state will depend.
+    prop_std
+        Proposal distribution standard deviation.
+    lambda_
+        Scaling parameter.
+    rng
+        Random number generator to be used in this function.
+    chain_index
+        Index of the Markov chain within the subset simulation level loop.
 
     Returns
     -------
     dict:
         x:
             Generated candidate state.
+        rng:
+            Random number generator to be used in the next invocation of this function,
+            for this chain.
     """
 
+    loop_idx = {
+        loop_name: int(loop_idx)
+        for loop_name, loop_idx in (
+            item.split("=")
+            for item in os.environ["MATFLOW_ELEMENT_ITER_LOOP_IDX"].split(";")
+        )
+    }
+
+    if loop_idx["markov_chain_state"] == 0:
+        # new chain, so want a new RNG:
+        rng = _init_rng(chain_index, loop_idx)
+
     x = x[:]  # convert to numpy array
     sigma = np.minimum(1, lambda_ * prop_std)
     rho = np.sqrt(1 - sigma**2)
-
-    return {"x": norm.rvs(loc=x * rho, scale=sigma)}
+    return {"x": norm.rvs(loc=x * rho, scale=sigma, random_state=rng), "rng": rng}

matflow/data/scripts/uq/generate_next_state_CS.py

@@ -1,25 +1,64 @@
+import os
+
 import numpy as np
 from numpy.typing import NDArray
 from scipy.stats import norm
 
 
-def generate_next_state_CS(x, prop_std):
+def _init_rng(chain_index, loop_idx):
+    """Initialise the random number generator, which must happen once for each Markov
+    chain."""
+    random_seed = int(os.environ["MATFLOW_RUN_RANDOM_SEED"])
+    spawn_key_str = os.environ["MATFLOW_RUN_RNG_SPAWN_KEY"]
+    task_iID = int(os.environ["MATFLOW_TASK_INSERT_ID"])
+    spawn_key = tuple(int(i) for i in spawn_key_str.split(",")) if spawn_key_str else ()
+    spawn_key = tuple([*spawn_key, task_iID, loop_idx["levels"], chain_index])
+    rng = np.random.default_rng(
+        np.random.SeedSequence(
+            random_seed,
+            spawn_key=spawn_key,
+        )
+    )
+    return rng
+
+
+def generate_next_state_CS(x, prop_std, rng, chain_index):
     """Generate the next candidate state using conditional sampling (a.k.a
     subset-infinity)
 
     Parameters
     ----------
     x
         Current state on which the candidate state will depend.
+    prop_std
+        Proposal distribution standard deviation.
+    rng
+        Random number generator to be used in this function.
+    chain_index
+        Index of the Markov chain within the subset simulation level loop.
 
     Returns
     -------
     dict:
         x:
             Generated candidate state.
+        rng:
+            Random number generator to be used in the next invocation of this function,
+            for this chain.
     """
 
+    loop_idx = {
+        loop_name: int(loop_idx)
+        for loop_name, loop_idx in (
+            item.split("=")
+            for item in os.environ["MATFLOW_ELEMENT_ITER_LOOP_IDX"].split(";")
+        )
+    }
+
+    if loop_idx["markov_chain_state"] == 0:
+        # new chain, so want a new RNG:
+        rng = _init_rng(chain_index, loop_idx)
+
     x = x[:]  # convert to numpy array
     rho = np.sqrt(1 - prop_std**2)
-
-    return {"x": norm.rvs(loc=x * rho, scale=prop_std)}
+    return {"x": norm.rvs(loc=x * rho, scale=prop_std, random_state=rng), "rng": rng}

matflow/data/scripts/uq/sample_direct_MC.py

@@ -1,10 +1,21 @@
+import os
+
 import numpy as np
 from scipy.stats import multivariate_normal
 
 
 def sample_direct_MC(dimension):
+    # RNG:
+    random_seed = int(os.environ["MATFLOW_RUN_RANDOM_SEED"])
+    sample_idx = int(os.environ["MATFLOW_ELEMENT_IDX"])  # TODO: use repeats idx?
+    task_iID = int(os.environ["MATFLOW_TASK_INSERT_ID"])
+    spawn_key_str = os.environ["MATFLOW_RUN_RNG_SPAWN_KEY"]
+    spawn_key = tuple(int(i) for i in spawn_key_str.split(",")) if spawn_key_str else ()
+    spawn_key = tuple([*spawn_key, task_iID, sample_idx])
+    rng = np.random.default_rng(np.random.SeedSequence(random_seed, spawn_key=spawn_key))
+
     # sample in standard normal space, variates are transformed before they are passed to
     # the performance function:
-    pi = multivariate_normal(mean=np.zeros(dimension), cov=None)
+    pi = multivariate_normal(mean=np.zeros(dimension), cov=None, seed=rng)
     x = np.atleast_1d(pi.rvs())
     return {"x": x}

matflow/data/template_components/task_schemas.yaml

@@ -1867,8 +1867,12 @@
   inputs:
     - parameter: x
     - parameter: prop_std
+    - parameter: chain_index
+    - parameter: rng
+      default_value: null
   outputs:
     - parameter: x
+    - parameter: rng
   actions:
     - script: <<script:uq/generate_next_state.py>>
       script_data_in: direct
@@ -1884,8 +1888,12 @@
   inputs:
     - parameter: x
     - parameter: prop_std
+    - parameter: chain_index
+    - parameter: rng
+      default_value: null    
   outputs:
     - parameter: x
+    - parameter: rng
   actions:
     - script: <<script:uq/generate_next_state_CS.py>>
       script_data_in: direct
@@ -1902,8 +1910,12 @@
     - parameter: x
     - parameter: prop_std
     - parameter: lambda_
+    - parameter: chain_index
+    - parameter: rng
+      default_value: null    
   outputs:
     - parameter: x
+    - parameter: rng
   actions:
     - script: <<script:uq/generate_next_state_ACS.py>>
       script_data_in: direct

matflow/data/workflows/subset_simulation_toy_model.yaml

@@ -13,12 +13,15 @@ loops:
 resources:
   any:
     combine_scripts: true
+    random_seed: 1234
 
 tasks:
   - schema: sample_direct_MC
     inputs:
       dimension: 200
-    repeats: 100 # num_samples
+    repeats:
+      number: 100 # num_samples
+      generate_seed_sequence: false
     groups:
       - name: all
   - schema: system_analysis_toy_model

matflow/data/workflows/subset_simulation_toy_model_ACS.yaml

@@ -18,12 +18,15 @@ loops:
 resources:
   any:
     combine_scripts: true
+    random_seed: 1234
 
 tasks:
   - schema: sample_direct_MC
     inputs:
       dimension: 200
-    repeats: 1000 # num_samples
+    repeats:
+      number: 100 # num_samples
+      generate_seed_sequence: false
     groups:
       - name: all
   - schema: system_analysis_toy_model

matflow/data/workflows/subset_simulation_toy_model_CS.yaml

@@ -13,12 +13,15 @@ loops:
 resources:
   any:
     combine_scripts: true
+    random_seed: 1234    
 
 tasks:
   - schema: sample_direct_MC
     inputs:
       dimension: 200
-    repeats: 100 # num_samples
+    repeats:
+      number: 100 # num_samples
+      generate_seed_sequence: false
     groups:
       - name: all
   - schema: system_analysis_toy_model