new training set sampling method

src/qibo/models/error_mitigation.py

@@ -222,12 +222,7 @@ def sample_training_circuit(
     # Compute the scores
     prob = np.exp(-(distance**2) / sigma**2)
     # Sample which of the RZ found to substitute
-    index = np.random.choice(
-        range(len(gates_to_replace)),
-        size=min(int(len(gates_to_replace) / 2), 50),
-        replace=False,
-        p=prob.sum(-1) / prob.sum(),
-    )
+    index = range(len(gates_to_replace))
     gates_to_replace = np.array([gates_to_replace[i] for i in index])
     prob = [prob[i] for i in index]
     # Sample which replacement gate to substitute with
@@ -248,6 +243,55 @@ def sample_training_circuit(
     return sampled_circuit
 
 
+def sample_training_set(
+    circuit,
+    observable,
+    nshots,
+    n_training_samples,
+    obs_frequencies={"-1": 0.1, "0": 0.8, "1": 0.1},
+):
+    """
+    Sample CDR training set.
+    """
+
+    n_subsets = len(obs_frequencies)
+    obs_values = list(obs_frequencies.keys())
+    residual = int(
+        n_training_samples
+        - sum(obs_frequencies[val] * n_training_samples for val in obs_frequencies)
+    )
+    print(residual)
+
+    training_circuits = [[] for _ in range(n_subsets)]
+
+    sanity_check = True
+
+    while sanity_check:
+        c = sample_training_circuit(circuit)
+        expv = round(c(nshots=nshots).expectation_from_samples(observable))
+        idx = obs_values.index(str(expv))
+
+        if len(training_circuits[idx]) < int(
+            obs_frequencies[str(expv)] * n_training_samples
+        ):
+            training_circuits[idx].append(c)
+
+        # sanity is checked if each expected value is represented by subset_dim circuits
+        if sum(len(training_circuits[i]) for i in range(n_subsets)) == (
+            n_training_samples - residual
+        ):
+            sanity_check = False
+
+    circuits = []
+    for circuit_type in training_circuits:
+        circuits.extend(circuit_type)
+
+    for _ in range(residual):
+        circuits.append(sample_training_circuit(circuit))
+
+    return circuits
+
+
 def CDR(
     circuit,
     observable,
@@ -291,9 +335,30 @@ def CDR(
     if backend is None:  # pragma: no cover
         backend = GlobalBackend()
     # Sample the training set
-    training_circuits = [
-        sample_training_circuit(circuit) for n in range(n_training_samples)
-    ]
+    training_circuits = sample_training_set(
+        circuit=circuit,
+        observable=observable,
+        nshots=nshots,
+        n_training_samples=n_training_samples,
+    )
+
+    sanity_check = True
+    while sanity_check:
+        training_circuits = []
+        expected_values = []
+        different_values = 0
+        for _ in range(n_training_samples):
+            c = sample_training_circuit(circuit)
+            expv = round(c(nshots=nshots).expectation_from_samples(observable))
+            if expv not in expected_values:
+                different_values += 1
+            training_circuits.append(c)
+            expected_values.append(expv)
+
+        # exit if we get more than one point for the fit
+        if different_values > 1:
+            sanity_check = False
+
     # Run the sampled circuits
     train_val = {"noise-free": [], "noisy": []}
     for c in training_circuits:
@@ -396,6 +461,7 @@ def vnCDR(
     training_circuits = [
         sample_training_circuit(circuit) for n in range(n_training_samples)
     ]
+
     train_val = {"noise-free": [], "noisy": []}
 
     # Add the different noise levels and run the circuits