Use random number generators, not global seeds (#94)

* Use random number generators, not global seeds * Use rng in notebooks and allow it to be passed as input to random funcs * lint * docstring update * Add release note * use rng varname * Fix varnames * Fix seed
Qiskit · Oct 29, 2024 · 6f5660c · 6f5660c
1 parent 0b7dae5
commit 6f5660c
Show file tree

Hide file tree

Showing 8 changed files with 81 additions and 69 deletions.
diff --git a/docs/how_tos/choose_subspace_dimension.ipynb b/docs/how_tos/choose_subspace_dimension.ipynb
diff --git a/docs/how_tos/integrate_dice_solver.ipynb b/docs/how_tos/integrate_dice_solver.ipynb
@@ -47,10 +47,10 @@
     "nuclear_repulsion_energy = mf_as.mol.energy_nuc()\n",
     "\n",
     "# Create a seed to control randomness throughout this workflow\n",
-    "rand_seed = 42\n",
+    "rng = np.random.default_rng(24)\n",
     "\n",
     "# Generate random samples\n",
-    "counts_dict = generate_counts_uniform(10_000, num_orbitals * 2, rand_seed=rand_seed)\n",
+    "counts_dict = generate_counts_uniform(10_000, num_orbitals * 2, rand_seed=rng)\n",
     "\n",
     "# Convert counts into bitstring and probability arrays\n",
     "bitstring_matrix_full, probs_arr_full = counts_to_arrays(counts_dict)\n",
@@ -81,7 +81,7 @@
     "            occupancies_bitwise,\n",
     "            num_elec_a,\n",
     "            num_elec_b,\n",
-    "            rand_seed=rand_seed,\n",
+    "            rand_seed=rng,\n",
     "        )\n",
     "\n",
     "    # Throw out samples with incorrect hamming weight and create batches of subsamples.\n",
@@ -92,7 +92,7 @@
     "        hamming_left=num_elec_b,\n",
     "        samples_per_batch=samples_per_batch,\n",
     "        num_batches=n_batches,\n",
-    "        rand_seed=rand_seed,\n",
+    "        rand_seed=rng,\n",
     "    )\n",
     "    # Run eigenstate solvers in a loop. This loop should be parallelized for larger problems.\n",
     "    int_e = np.zeros(n_batches)\n",

diff --git a/docs/how_tos/use_oo_to_optimize_hamiltonian_basis.ipynb b/docs/how_tos/use_oo_to_optimize_hamiltonian_basis.ipynb
diff --git a/docs/tutorials/01_chemistry_hamiltonian.ipynb b/docs/tutorials/01_chemistry_hamiltonian.ipynb
diff --git a/qiskit_addon_sqd/configuration_recovery.py b/qiskit_addon_sqd/configuration_recovery.py
@@ -54,7 +54,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 +82,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 seed for controlling randomness
 
     Returns:
         A refined bitstring matrix and an updated probability array.
@@ -92,6 +92,8 @@ def recover_configurations(
              arXiv:2405.05068 [quant-ph].
 
     """
+    rng = 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 +106,7 @@ def recover_configurations(
             avg_occupancies,
             num_elec_a,
             num_elec_b,
-            rand_seed=rand_seed,
+            rng=rng,
         )
         bs_str = "".join("1" if bit else "0" for bit in bs_corrected)
         corrected_dict[bs_str] += freq
@@ -183,7 +185,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 +194,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 +203,6 @@ def _bipartite_bitstring_correcting(
     # This function must not mutate the input arrays.
     bit_array = bit_array.copy()
 
-    np.random.seed(rand_seed)
-
     # The number of bits should be even
     num_bits = bit_array.shape[0]
     partition_size = num_bits // 2
@@ -246,7 +246,7 @@ def _bipartite_bitstring_correcting(
             probs_left[bit_array[:partition_size]]
         )
         # Correct the hamming by probabilistically flipping some bits to flip to 0
-        indices_to_flip = np.random.choice(
+        indices_to_flip = rng.choice(
             indices_occupied, size=round(n_diff), replace=False, p=p_choice
         )
         bit_array[:partition_size][indices_to_flip] = False
@@ -259,7 +259,7 @@ def _bipartite_bitstring_correcting(
             probs_left[np.logical_not(bit_array[:partition_size])]
         )
         # Correct the hamming by probabilistically flipping some bits to flip to 1
-        indices_to_flip = np.random.choice(
+        indices_to_flip = rng.choice(
             indices_empty, size=round(np.abs(n_diff)), replace=False, p=p_choice
         )
         bit_array[:partition_size][indices_to_flip] = np.logical_not(
@@ -280,7 +280,7 @@ def _bipartite_bitstring_correcting(
             probs_right[bit_array[partition_size:]]
         )
         # Correct the hamming by probabilistically flipping some bits to flip to 0
-        indices_to_flip = np.random.choice(
+        indices_to_flip = rng.choice(
             indices_occupied, size=round(n_diff), replace=False, p=p_choice
         )
         bit_array[partition_size:][indices_to_flip] = np.logical_not(
@@ -295,7 +295,7 @@ def _bipartite_bitstring_correcting(
             probs_right[np.logical_not(bit_array[partition_size:])]
         )
         # Correct the hamming by probabilistically flipping some bits to flip to 1
-        indices_to_flip = np.random.choice(
+        indices_to_flip = rng.choice(
             indices_empty, size=round(np.abs(n_diff)), replace=False, p=p_choice
         )
         bit_array[partition_size:][indices_to_flip] = np.logical_not(

diff --git a/qiskit_addon_sqd/counts.py b/qiskit_addon_sqd/counts.py
@@ -41,7 +41,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 +62,13 @@ 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.")
-    np.random.seed(rand_seed)
+
+    rng = 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 = np.random.choice([0, 1], size=(num_samples, num_bits))
+    bts_matrix = rng.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 +83,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,17 +114,13 @@ def generate_counts_bipartite_hamming(
     if hamming_left < 0 or hamming_right < 0:
         raise ValueError("Hamming weights must be specified as non-negative integers.")
 
-    np.random.seed(rand_seed)
+    rng = 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 = np.random.choice(np.arange(num_bits // 2), hamming_right, replace=False).astype(
-            "int"
-        )
-        dn_flips = np.random.choice(np.arange(num_bits // 2), hamming_left, replace=False).astype(
-            "int"
-        )
+        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")
 
         # Create a bitstring with the chosen bits flipped
         bts_arr = np.zeros(num_bits)

diff --git a/qiskit_addon_sqd/subsampling.py b/qiskit_addon_sqd/subsampling.py
@@ -28,7 +28,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 +68,8 @@ def postselect_and_subsample(
     if hamming_left < 0 or hamming_right < 0:
         raise ValueError("Hamming weight must be specified with a non-negative integer.")
 
+    rng = 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 +81,15 @@ 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=rng)
 
 
 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,7 +124,8 @@ def subsample(
     if num_batches < 1:
         raise ValueError("The number of batches must be specified with a positive integer.")
 
-    np.random.seed(rand_seed)
+    rng = np.random.default_rng(rand_seed)
+
     num_bitstrings = bitstring_matrix.shape[0]
 
     # If the number of requested samples is >= the number of bitstrings, return
@@ -136,7 +139,7 @@ def subsample(
     batches = []
     for _ in range(num_batches):
         if randomly_sample:
-            indices = np.random.choice(
+            indices = rng.choice(
                 np.arange(num_bitstrings).astype("int"),
                 samples_per_batch,
                 replace=False,

diff --git a/releasenotes/notes/rng-620825ec3215e415.yaml b/releasenotes/notes/rng-620825ec3215e415.yaml
@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    All functions which take a ``rand_seed`` argument now also accept a ``numpy.random.Generator`` instance as the ``rand_seed``.