Use rng in notebooks and allow it to be passed as input to random funcs

docs/how_tos/integrate_dice_solver.ipynb

@@ -47,7 +47,7 @@
     "nuclear_repulsion_energy = mf_as.mol.energy_nuc()\n",
     "\n",
     "# Create a seed to control randomness throughout this workflow\n",
-    "rand_seed = 42\n",
+    "rand_seed = np.random.default_rng(2**24)\n",
     "\n",
     "# Generate random samples\n",
     "counts_dict = generate_counts_uniform(10_000, num_orbitals * 2, rand_seed=rand_seed)\n",

qiskit_addon_sqd/configuration_recovery.py

@@ -15,6 +15,7 @@
 
 from __future__ import annotations
 
+import warnings
 from collections import defaultdict
 from collections.abc import Sequence
 
@@ -54,7 +55,7 @@ def recover_configurations(
     avg_occupancies: np.ndarray,
     num_elec_a: int,
     num_elec_b: int,
-    rand_seed: int | None = None,
+    rand_seed: np.random.Generator | int | None = None,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Refine bitstrings based on average orbital occupancy and a target hamming weight.
 
@@ -82,7 +83,7 @@ def recover_configurations(
             ``i`` in ``bitstring_matrix``.
         num_elec_a: The number of spin-up electrons in the system.
         num_elec_b: The number of spin-down electrons in the system.
-        rand_seed: A seed to control random behavior
+        rand_seed: A random number generator
 
     Returns:
         A refined bitstring matrix and an updated probability array.
@@ -92,6 +93,9 @@ def recover_configurations(
              arXiv:2405.05068 [quant-ph].
 
     """
+    if isinstance(rand_seed, int) or rand_seed is None:
+        rand_seed = np.random.default_rng(rand_seed)
+
     if num_elec_a < 0 or num_elec_b < 0:
         raise ValueError("The numbers of electrons must be specified as non-negative integers.")
 
@@ -104,7 +108,7 @@ def recover_configurations(
             avg_occupancies,
             num_elec_a,
             num_elec_b,
-            rand_seed=rand_seed,
+            rng=rand_seed,
         )
         bs_str = "".join("1" if bit else "0" for bit in bs_corrected)
         corrected_dict[bs_str] += freq
@@ -183,7 +187,7 @@ def _bipartite_bitstring_correcting(
     avg_occupancies: np.ndarray,
     hamming_right: int,
     hamming_left: int,
-    rand_seed: int | None = None,
+    rng: np.random.Generator,
 ) -> np.ndarray:
     """Use occupancy information and target hamming weight to correct a bitstring.
 
@@ -192,7 +196,7 @@ def _bipartite_bitstring_correcting(
         avg_occupancies: A 1D array containing the mean occupancy of each orbital.
         hamming_right: The target hamming weight used for the right half of the bitstring
         hamming_left: The target hamming weight used for the left half of the bitstring
-        rand_seed: A seed to control random behavior
+        rng: A random number generator
 
     Returns:
         A corrected bitstring
@@ -201,8 +205,6 @@ def _bipartite_bitstring_correcting(
     # This function must not mutate the input arrays.
     bit_array = bit_array.copy()
 
-    rng = np.random.default_rng(rand_seed)
-
     # The number of bits should be even
     num_bits = bit_array.shape[0]
     partition_size = num_bits // 2

qiskit_addon_sqd/counts.py

@@ -15,6 +15,8 @@
 
 from __future__ import annotations
 
+import warnings
+
 import numpy as np
 
 
@@ -41,7 +43,7 @@ def counts_to_arrays(counts: dict[str, float | int]) -> tuple[np.ndarray, np.nda
 
 
 def generate_counts_uniform(
-    num_samples: int, num_bits: int, rand_seed: None | int = None
+    num_samples: int, num_bits: int, rand_seed: np.random.Generator | int | None = None
 ) -> dict[str, int]:
     """Generate a bitstring counts dictionary of samples drawn from the uniform distribution.
 
@@ -62,11 +64,14 @@ def generate_counts_uniform(
         raise ValueError("The number of samples must be specified with a positive integer.")
     if num_bits < 1:
         raise ValueError("The number of bits must be specified with a positive integer.")
-    rng = np.random.default_rng(rand_seed)
+
+    if isinstance(rand_seed, int) or rand_seed is None:
+        rand_seed = np.random.default_rng(rand_seed)
+
     sample_dict: dict[str, int] = {}
     # Use numpy to generate a random matrix of bit values and
     # convert it to a dictionary of bitstring samples
-    bts_matrix = rng.choice([0, 1], size=(num_samples, num_bits))
+    bts_matrix = rand_seed.choice([0, 1], size=(num_samples, num_bits))
     for i in range(num_samples):
         bts_arr = bts_matrix[i, :].astype("int")
         bts = "".join("1" if bit else "0" for bit in bts_arr)
@@ -81,7 +86,7 @@ def generate_counts_bipartite_hamming(
     *,
     hamming_right: int,
     hamming_left: int,
-    rand_seed: None | int = None,
+    rand_seed: np.random.Generator | int | None = None,
 ) -> dict[str, int]:
     """Generate a bitstring counts dictionary with specified bipartite hamming weight.
 
@@ -112,13 +117,18 @@ def generate_counts_bipartite_hamming(
     if hamming_left < 0 or hamming_right < 0:
         raise ValueError("Hamming weights must be specified as non-negative integers.")
 
-    rng = np.random.default_rng(rand_seed)
+    if isinstance(rand_seed, int) or rand_seed is None:
+        rand_seed = np.random.default_rng(rand_seed)
 
     sample_dict: dict[str, int] = {}
     for _ in range(num_samples):
         # Pick random bits to flip such that the left and right hamming weights are correct
-        up_flips = rng.choice(np.arange(num_bits // 2), hamming_right, replace=False).astype("int")
-        dn_flips = rng.choice(np.arange(num_bits // 2), hamming_left, replace=False).astype("int")
+        up_flips = rand_seed.choice(np.arange(num_bits // 2), hamming_right, replace=False).astype(
+            "int"
+        )
+        dn_flips = rand_seed.choice(np.arange(num_bits // 2), hamming_left, replace=False).astype(
+            "int"
+        )
 
         # Create a bitstring with the chosen bits flipped
         bts_arr = np.zeros(num_bits)

qiskit_addon_sqd/subsampling.py

@@ -15,6 +15,8 @@
 
 from __future__ import annotations
 
+import warnings
+
 import numpy as np
 
 from .configuration_recovery import post_select_by_hamming_weight
@@ -28,7 +30,7 @@ def postselect_and_subsample(
     hamming_left: int,
     samples_per_batch: int,
     num_batches: int,
-    rand_seed: int | None = None,
+    rand_seed: np.random.Generator | int | None = None,
 ) -> list[np.ndarray]:
     """Subsample batches of bit arrays with correct hamming weight from an input ``bitstring_matrix``.
 
@@ -68,6 +70,9 @@ def postselect_and_subsample(
     if hamming_left < 0 or hamming_right < 0:
         raise ValueError("Hamming weight must be specified with a non-negative integer.")
 
+    if isinstance(rand_seed, int) or rand_seed is None:
+        rand_seed = np.random.default_rng(rand_seed)
+
     # Post-select only bitstrings with correct hamming weight
     mask_postsel = post_select_by_hamming_weight(
         bitstring_matrix, hamming_right=hamming_right, hamming_left=hamming_left
@@ -79,15 +84,17 @@ def postselect_and_subsample(
     if len(probs_postsel) == 0:
         return [np.array([])] * num_batches
 
-    return subsample(bs_mat_postsel, probs_postsel, samples_per_batch, num_batches, rand_seed)
+    return subsample(
+        bs_mat_postsel, probs_postsel, samples_per_batch, num_batches, rand_seed=rand_seed
+    )
 
 
 def subsample(
     bitstring_matrix: np.ndarray,
     probabilities: np.ndarray,
     samples_per_batch: int,
     num_batches: int,
-    rand_seed: int | None = None,
+    rand_seed: np.random.Generator | int | None = None,
 ) -> list[np.ndarray]:
     """Subsample batches of bit arrays from an input ``bitstring_matrix``.
 
@@ -122,6 +129,9 @@ def subsample(
     if num_batches < 1:
         raise ValueError("The number of batches must be specified with a positive integer.")
 
+    if isinstance(rand_seed, int) or rand_seed is None:
+        rand_seed = np.random.default_rng(rand_seed)
+
     num_bitstrings = bitstring_matrix.shape[0]
 
     # If the number of requested samples is >= the number of bitstrings, return
@@ -135,8 +145,7 @@ def subsample(
     batches = []
     for _ in range(num_batches):
         if randomly_sample:
-            rng = np.random.default_rng(rand_seed)
-            indices = rng.choice(
+            indices = rand_seed.choice(
                 np.arange(num_bitstrings).astype("int"),
                 samples_per_batch,
                 replace=False,